{
    "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"
        ]
    }
}