Muon is Scalable for LLM Training
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng, Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang, Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, Zhilin Yang
cs.LG, cs.AI, cs.CL
2025-02-24
Moonshot adds weight decay and per-matrix update scaling to Muon, matching AdamW loss at ~52% of the FLOPs in compute-optimal runs, and trains Moonlight on 5.7T tokens.
Muon, from Jordan et al. 2024, updates matrix-shaped weights by accumulating gradient momentum and approximately orthogonalizing it with a Newton-Schulz iteration. Singular values of the update are pushed near 1, so the step does not collapse onto a few dominant directions. That orthogonalization is steepest descent under the spectral norm. Weights act as operators on their inputs, and a spectral-norm ball fits that role better than Adam's moving max-of-max constraint.
On small language models this beat AdamW. The gain thinned once parameter counts and token budgets grew. Two mechanisms show up in the runs. Without weight decay, weight RMS and layer-output RMS climb until they leave the accurate range of bf16. Update size is also tied to shape. For a full-rank matrix [A, B], the theoretical Muon update RMS is 1/√max(A,B). Wide MLP matrices then take steps that are too small to use the model's capacity. Splitting KV heads in GQA or MLA makes those matrices small, and the steps get large enough to destabilize training. Distributed training adds a further constraint: orthogonalization needs the whole gradient matrix, so an elementwise ZeRO-1 Adam step does not carry over.
Two changes let Muon reuse a learning rate and a weight decay already tuned for AdamW.
Weight decay has the same form as in AdamW: each step subtracts λ times the current weight along with the orthogonalized update. On an 800M model trained for 100B tokens, about 5× the compute-optimal token budget, vanilla Muon leads early, by about 0.023 validation loss at step 24,000. Later the weights grow. By step 66,000 the decayed run is ahead by about 0.017, and it is also ahead of AdamW.
The shape fix multiplies the orthogonalized matrix by 0.2·√max(A,B). The square root cancels the 1/√max(A,B) shrinkage. The 0.2 factor puts update RMS in the 0.2 to 0.4 band typical of AdamW. Matrix parameters use Muon. RMSNorm, embeddings, and the LM head stay on AdamW. Both sides share η and λ.
An 800M ablation replaces SwiGLU with a two-layer MLP of shape [H, 4H] and trains 4B tokens of a 20B schedule. Scaling every update by 0.2·√H, which is the right factor for square matrices, gives validation loss 2.812. Forcing update RMS to exactly 0.2, or multiplying by √max(A,B), both reach 2.789. The shape-based rule is cheaper, and later runs use it.
Ten Newton-Schulz steps orthogonalize more tightly than five and do not lower loss, so N stays at 5. Momentum is 0.95. The polynomial coefficients stay at the original 3.4445, −4.7750, and 2.0315. Weight decay is 0.1.
Distributed Muon shards optimizer state across the data-parallel group in the ZeRO-1 style, gathers a full matrix for Newton-Schulz, and keeps only the local shard for the parameter update. Muon stores one momentum buffer against AdamW's two, so the extra optimizer memory is about half. Communication lands between 1× and 1.25× a distributed AdamW step; the upper bound is an fp32 reduce-scatter, an fp32 all-gather, and one extra bf16 gather. Tensor parallelism adds another bf16 gather on the TP group. The optimizer is usually 1% to 3% of forward-plus-backward time. On the cluster used here, end-to-end latency is not distinguishable from AdamW.
The scaling laws use dense Llama-style models from 399M to 1.5B, compute-optimal token budgets, and 8K context. AdamW gets a grid search. Muon reuses that learning rate and batch size. Fitted loss is 2.506×C^(−0.052) for Muon and 2.608×C^(−0.054) for AdamW. Under compute-optimal training, matching AdamW's loss takes about 52% of the FLOPs. That is the 2× efficiency quoted in the abstract.
Moonlight is the large run: a DeepSeek-V3-Small-like MoE with 2.24B activated and 15.29B total parameters excluding embeddings (about 3B/16B including them), trained on 5.7T tokens at 8K context. The learning rate rises to 4.2e-4, cosine-decays to 4.2e-5, then a 500B-token cooldown on math, code, and reasoning data takes it to zero.
The matched optimizer comparison stops at 1.2T, before cooldown. Against Moonlight-A, the same recipe on AdamW, MMLU is essentially tied, code and math move up, and BBH moves down. HumanEval and MBPP are pass@1.
| Method | Metric | Result |
| Moonlight (Muon, 1.2T) | MMLU / BBH / HumanEval / MATH | 60.4 / 43.2 / 37.2 / 19.8 |
| Moonlight-A (AdamW, 1.2T) | MMLU / BBH / HumanEval / MATH | 60.2 / 45.3 / 29.3 / 16.1 |
| DeepSeek-V3-Small (AdamW, 1.33T) | MMLU / BBH / HumanEval / MATH | 53.3 / 41.4 / 26.8 / 10.7 |
V3-Small sits well below both. Moonlight-A already scores 60.2 MMLU, so the recipe gap is larger than the optimizer gap on that benchmark.
At 5.7T, Moonlight scores MMLU 70.0, MMLU-Pro 42.4, BBH 65.2, HumanEval 48.1, GSM8K 77.4, and MATH 45.3. DeepSeek-V2-Lite, also listed as 2.4B/16B and trained for 5.7T with AdamW, scores 58.3, 25.5, 44.1, 29.9, 41.1, and 17.1. Qwen2.5-3B, a dense 3B trained for 18T tokens, scores MMLU 65.6 and GSM8K 79.1. Moonlight is higher on MMLU and slightly lower on GSM8K. Figure 1 places Moonlight on the Pareto frontier of performance versus training FLOPs. The jump from 60.4 at 1.2T to 70.0 at 5.7T includes the remaining tokens and the cooldown. There is no AdamW twin at 5.7T.
Across the 1.2T pretraining, mean SVD entropy stays higher for Muon in all six weight groups. At the 1.2T checkpoint, more than 90% of individual matrices also have higher SVD entropy under Muon. Router weights show the widest gap.
SFT works best with the same optimizer used in pretraining. The 1.2T checkpoints are finetuned for two epochs on tulu-3. MMLU in this table is 0-shot chain-of-thought, a different protocol from the 5-shot pretraining number. Muon then Muon scores MMLU 55.7, HumanEval 57.3, GSM8K 68.0. AdamW then AdamW scores 52.0, 53.1, 64.6. Muon pretraining followed by AdamW SFT drops MMLU to 50.2. On Qwen2.5-7B, already pretrained with AdamW, using Muon only for SFT is slightly worse on every reported metric: MMLU 70.8 vs 71.4, HumanEval 77.4 vs 79.3, MBPP 71.6 vs 71.9, GSM8K 85.8 vs 89.8.
For a pretraining run that starts from scratch, the offer is specific. Keep the AdamW learning rate, set weight decay to 0.1, and the scaling curve puts the same loss at about 52% of the FLOPs. Optimizer state drops by one momentum buffer. The extra communication did not show up as end-to-end slowdown on the reported cluster. The distributed implementation, pretrained weights, instruction-tuned weights, and intermediate checkpoints are released with the paper.
Checkpoints that already finished pretraining under AdamW are a different story. Swapping only the SFT optimizer shows no reliable gain. In the matched 1.2T comparison, the gains sit in HumanEval and MATH. MMLU moves by 0.2.
Three items are left open. RMSNorm, embeddings, and the LM head still use AdamW. Extending the spectral-norm view to a general Schatten norm is future work. Mixing optimizers between pretraining and finetuning hurts, which blocks a drop-in swap onto the existing pile of AdamW checkpoints.
The 52% figure is a fit on dense models from 399M to 1.5B at compute-optimal token counts. The exponents are −0.052 and −0.054, so most of the gap is the leading coefficient. Moonlight is a 2.24B-active MoE trained for 5.7T tokens, far past compute-optimal. The matched AdamW run stops at 1.2T, where MMLU differs by 0.2 and BBH is 2.1 points worse for Muon. The double-digit gaps versus V2-Lite mix in data and architecture changes. The 0.2 scale factor comes from the observed AdamW update-RMS band, and the theory does not single it out. The claim of no visible latency comes without a wall-clock number at Moonlight scale.