m21/config.json

{
  "model_id": "m21",
  "run_name": "m21-gqa-rope-dense-muonv2",
  "data_dir": "data/fineweb-edu",
  "vocab_size": 50257,
  "block_size": 1024,
  "d_model": 512,
  "num_layers": 12,
  "num_heads": 8,
  "ffn_mult": 4,
  "dropout": 0.0,
  "bias": false,
  "norm_type": "rmsnorm",
  "tie_embeddings": true,
  "attention_type": "gqa",
  "gqa_num_kv_heads": 2,
  "mla_latent_dim": 64,
  "pos_type": "rope",
  "rope_base": 10000,
  "ffn_type": "dense",
  "moe_num_experts": 8,
  "moe_top_k": 2,
  "optimizer": "muon_v2",
  "lr": 0.0006,
  "min_lr": 6e-05,
  "weight_decay": 0.1,
  "beta1": 0.9,
  "beta2": 0.95,
  "grad_clip": 1.0,
  "muon_lr": 0.02,
  "muon_momentum": 0.95,
  "training_tokens": 1000000000,
  "warmup_tokens": 6000000,
  "lr_schedule": "cosine",
  "batch_size": 4,
  "grad_accum_steps": 16,
  "device": "cuda",
  "dtype": "bfloat16",
  "compile": true,
  "seed": 1337,
  "eval_interval": 250,
  "eval_iters": 100,
  "log_interval": 10,
  "checkpoint_interval": 1000,
  "results_dir": "results",
  "mlflow_experiment": "slm-ablation",
  "mlflow_tracking_uri": "sqlite:///mlflow.db",
  "mlflow_artifact_location": "./mlartifacts"
}