{
    "git_commit": "b1fea531a1bb067c5975282a447929a416793831",
    "wandb_run_id": "ef3f1rcw",
    "config": {
        "seed": 42,
        "notes": null,
        "trainer": {
            "run_dir": "local_data/runs/ablate-scales",
            "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
                },
                "fast": {
                    "_target_": "optax.contrib.muon",
                    "lr": 0.0005,
                    "beta": 0.9,
                    "adam_b1": 0.9,
                    "adam_b2": 0.95,
                    "eps": 1e-06,
                    "weight_decay": 0.01,
                    "ns_steps": 5,
                    "consistent_rms": 0.2,
                    "nesterov": true,
                    "warmup_steps": 100,
                    "stable_steps": 899900,
                    "decay_steps": 100000,
                    "min_lr_ratio": 0.0
                }
            },
            "_target_": "trainers.forte_trainer.ForteTrainer"
        },
        "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": "ablate-scales",
        "initialization": "aklein4/forte-init-instruct",
        "initialization_step": 0,
        "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.forte.ForteModel",
            "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",
            "num_fast_vectors": 512,
            "num_fast_heads": 16,
            "base_lr": 0.0001,
            "pool_size": 512,
            "grad_rms_eps": 1e-12
        }
    },
    "step": 200,
    "checkpoint_items": [
        "params"
    ]
}