config.json

{
  "run_name": "model2_50m_ctx2n",
  "results_tag": "canonical",
  "hf_repo_id": "FAIRC/token-averaging-model2_50m_ctx2n",
  "checkpoint_format": {
    "keys": [
      "step",
      "tokens_seen",
      "cumulative_flops",
      "model",
      "optimizer",
      "scheduler"
    ],
    "note": "Raw torch.save dict from experiments/chinchilla/train.py. Load with torch.load(..., map_location='cpu', weights_only=False) and take state['model']. Not a transformers AutoModel checkpoint."
  },
  "model_config": {
    "name": "model2_50m_ctx2n",
    "d_model": 512,
    "n_heads": 8,
    "n_layers": 8,
    "context_len": 2048,
    "averaging_k": 1,
    "method_name": null,
    "multi_token_phase_ratio": 0.0,
    "tie_embeddings": true,
    "grad_checkpoint": true,
    "color": "#f0a500",
    "label": "~50M standard (2n=2048)",
    "target_tokens": 1000000000,
    "lr": 0.0002,
    "warmup_steps": 2000,
    "n_params_approx": 50897408
  }
}