{
    "git_commit": "a34fdf8bdceb4b644a9106819b9d6f7d6d0020e0",
    "wandb_run_id": "fnn7p0ul",
    "config": {
        "seed": 42,
        "notes": null,
        "trainer": {
            "run_dir": "local_data/runs/baseline-fixed",
            "checkpoint": {
                "load_state": false,
                "save_optimizer": false,
                "interval": 50
            },
            "wandb": {
                "project": "oloop-instruct",
                "enabled": true
            },
            "huggingface": {
                "private": false,
                "enabled": true
            },
            "global_batch_size": 128,
            "max_steps": 1000000,
            "num_logit_iterations": 4,
            "aux_loss_weight": 0.1,
            "max_grad_norm": 1.0,
            "debug": false,
            "groups": {
                "slow": {
                    "_target_": "optax.adamw",
                    "lr": 3e-05,
                    "b1": 0.9,
                    "b2": 0.95,
                    "eps": 1e-06,
                    "weight_decay": 0.0,
                    "warmup_steps": 100,
                    "stable_steps": 899900,
                    "decay_steps": 100000,
                    "min_lr_ratio": 0.0
                }
            },
            "_target_": "trainers.horizon_lm_trainer.HorizonLMTrainer"
        },
        "data": {
            "skip_batches": 0,
            "dataset": {
                "path": "aklein4/300k-horizons",
                "split": "train",
                "streaming": true
            },
            "collator": {
                "_target_": "collators.horizon.HorizonCollator",
                "tokenizer_url": "meta-llama/Llama-3.2-1B-Instruct",
                "max_length": 1024,
                "cluster_length": 64
            }
        },
        "distributed": {
            "initialize": "auto"
        },
        "mesh": {
            "fsdp_axis_size": -1
        },
        "name": "baseline-fixed",
        "initialization": "meta-llama/Llama-3.2-1B-Instruct",
        "initialization_step": null,
        "model": {
            "sharding": {
                "default": [],
                "rules": [
                    {
                        "pattern": "^embed_tokens/embedding$",
                        "spec": [
                            "fsdp",
                            null
                        ]
                    },
                    {
                        "pattern": "layers/mlp/(gate_proj|up_proj|down_proj)",
                        "spec": [
                            null,
                            null,
                            "fsdp"
                        ]
                    }
                ]
            },
            "_target_": "models.llama.LlamaForCausalLM",
            "vocab_size": 128256,
            "pad_token_id": -100,
            "hidden_size": 2048,
            "intermediate_size": 8192,
            "num_hidden_layers": 16,
            "num_attention_heads": 32,
            "num_key_value_heads": 8,
            "head_dim": 64,
            "max_position_embeddings": 131072,
            "rope": {
                "theta": 500000.0,
                "factor": 32.0,
                "high_freq_factor": 4.0,
                "low_freq_factor": 1.0,
                "original_context_len": 8192
            },
            "rms_norm_eps": 1e-05,
            "compute_dtype": "bfloat16",
            "gradient_checkpointing": true,
            "attention_backend": "auto"
        }
    },
    "step": 950,
    "checkpoint_items": [
        "params"
    ]
}