File size: 1,616 Bytes
13d8176
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
{
    "type": "oloop.OLoopModel",
    "pretrained_url": "aklein4/Llama-3.2-1B-TPU",
    "pretrained_step": 0,
    "pretrained_strict": false,
    "torch_dtype": "float32",
    "vocab_size": 128256,
    "bos_token_id": 128000,
    "eos_token_id": 128001,
    "pad_token_id": -100,
    "hidden_size": 2048,
    "num_hidden_layers": 16,
    "num_attention_heads": 32,
    "num_key_value_heads": 8,
    "intermediate_size": 8192,
    "hidden_act": "silu",
    "max_position_embeddings": 131072,
    "rope_theta": 500000.0,
    "rope_scaling": {
        "factor": 32.0,
        "high_freq_factor": 4.0,
        "low_freq_factor": 1.0,
        "original_context_len": 8192
    },
    "initializer_range": 0.02,
    "attention_dropout": false,
    "attention_bias": false,
    "rms_norm_eps": 1e-05,
    "attention_kernel": "flash_attention",
    "pure_modules": [],
    "fast_weight_size": 2048,
    "base_lr": 0.001,
    "momentum_beta": 0.9,
    "momentum_dtype": "bfloat16",
    "state_dtype": "float32",
    "sharding": {
        "model.layers.*": [
            [
                "data",
                "fsdp"
            ],
            null,
            null
        ],
        "lm_head": [
            [
                "data",
                "fsdp"
            ],
            null,
            null
        ]
    },
    "remat": {
        "activation_checkpoint_layers": [
            "LlamaDecoderLayer"
        ],
        "optimization_barrier_layers": [
            "LlamaDecoderLayer"
        ],
        "scan_layers": "model.layers",
        "offload_tensors": [
            "decoder_input"
        ]
    }
}