final/config.json

{
  "step": 0,
  "epoch": 0,
  "total_tokens": 0,
  "best_val": Infinity,
  "val_loss": 11.654257669448853,
  "mtp_val": 11.198893718719482,
  "fp8_mode": "torchao_rowwise",
  "seq_len": 4096,
  "config": {
    "size_name": "7b-moe",
    "vocab_size": 50304,
    "d_model": 1536,
    "n_layers": 30,
    "n_heads": 12,
    "n_kv_heads": 2,
    "head_dim": 128,
    "ffn_hidden": 6144,
    "use_moe": true,
    "n_dense_layers": 1,
    "n_routed_experts": 64,
    "n_shared_experts": 2,
    "moe_top_k": 6,
    "expert_hidden": 768,
    "routed_scale": 1.0,
    "bias_update_rate": 0.001,
    "aux_loss_weight": 0.001,
    "router_z_weight": 0.0001,
    "global_every": 4,
    "window": 1024,
    "rope_local": 10000.0,
    "rope_global": 500000.0,
    "use_qk_norm": true,
    "doc_masking": true,
    "tie_embeddings": false,
    "z_loss": 0.0001,
    "n_ce_chunks": 8,
    "use_mtp": true,
    "mtp_weight": 0.3,
    "resid_dropout": 0.0,
    "attn_dropout": 0.0,
    "use_fim": true,
    "fim_rate": 0.6,
    "fim_spm_rate": 0.5,
    "eos_id": 50268,
    "pad_id": 50272,
    "fim_pre_id": 50269,
    "fim_suf_id": 50270,
    "fim_mid_id": 50271,
    "fim_pad_id": -1,
    "param_dtype_str": "bf16",
    "stochastic_rounding": true,
    "ckpt_attn": true,
    "ckpt_ffn": true,
    "muon_lr": 0.0022,
    "adam_lr": 0.00045,
    "muon_momentum": 0.95,
    "muon_ns_steps": 5,
    "weight_decay": 0.1,
    "beta1": 0.9,
    "beta2": 0.95,
    "grad_clip": 1.0,
    "min_lr_ratio": 0.02,
    "warmup_tokens": 100000000,
    "decay_frac": 0.2,
    "target_total_tokens": 10000000000,
    "max_epochs": 4.0,
    "seq_len": 4096,
    "long_ctx_len": 16384,
    "long_ctx_frac": 0.12,
    "max_len": 16384,
    "micro_batch": 16,
    "grad_accum": 16,
    "batch_ramp_tokens": 2000000000,
    "log_every": 10,
    "eval_every_steps": 500,
    "eval_batches": 40,
    "ckpt_every_min": 45.0,
    "max_session_hours": 11.3,
    "overfit_patience": 5,
    "seed": 1337
  }
}