Source-linked AI summary
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
TL;DR
The paper studies whether Muon, previously demonstrated mainly on small language models, can scale to large-scale LLM training. It adds weight decay and adjusts per-parameter update scales, finding comparable performance to AdamW at approximately 52% of the training FLOPs and applying the optimizer to Moonlight. The authors release distributed Muon code, models, and checkpoints for further research.
Problem
Muon showed promise on small language models, but its scalability to larger models remained unproven.
Method
The paper analyzes Muon’s scaling behavior, adds weight decay and parameter-wise update-scale adjustments, and develops a distributed implementation.
Results
Muon achieves comparable performance to AdamW-trained counterparts while requiring approximately 52% of the training FLOPs, and is used to train the 3B/16B-parameter Moonlight MoE model on 5.7 trillion tokens.
Takeaways & Limitations
Muon can serve as an optimizer for large-scale LLM training, with released implementations, models, and intermediate checkpoints supporting further research.
Takeaways & Limitations
The report identifies extending Muon beyond the spectral norm to general Schatten norms as an unresolved direction for potentially broader optimization capabilities.
Abstract
from arXiv · showhide
Recently, the Muon optimizer based on matrix orthogonalization has demonstrated strong results in training small-scale language models, but the scalability to larger models has not been proven. We identify two crucial techniques for scaling up Muon: (1) adding weight decay and (2) carefully adjusting the per-parameter update scale. These techniques allow Muon to work out-of-the-box on large-scale training without the need of hyper-parameter tuning. Scaling law experiments indicate that Muon achieves $\sim\!2\times$ computational efficiency compared to AdamW with compute optimal training. Based on these improvements, we introduce Moonlight, a 3B/16B-parameter Mixture-of-Expert (MoE) model trained with 5.7T tokens using Muon. Our model improves the current Pareto frontier, achieving better performance with much fewer training FLOPs compared to prior models. We open-source our distributed Muon implementation that is memory optimal and communication efficient. We also release the pretrained, instruction-tuned, and intermediate checkpoints to support future research.
1 Introduction
The report addresses whether Muon can scale from small language models to large-scale LLM training. It identifies scaling techniques, validates efficiency against AdamW, and applies Muon in the Moonlight MoE model.
- Motivation: Muon uses orthogonalized gradient momentum for matrix parameters, but its scalability to billion-parameter models trained on trillions of tokens remained unresolved.The report also identifies distributed orthogonalization and cross-stage generalization as open challenges.
- Scaling Muon: Weight decay and parameter-wise update-scale adjustments improve Muon’s stability and allow out-of-the-box large-scale training without hyper-parameter tuning.The report identifies weight decay as crucial to scalability and proposes adjusted update rules.
- Scaling-law validation: Muon matches AdamW-trained counterparts while requiring approximately 52% of the training FLOPs.This result comes from scaling-law comparisons against strong AdamW baselines.
- Results: The authors report that Muon can replace AdamW for large-scale LLM training, improving both training efficiency and model performance.They support this claim with comprehensive experiments and introduce Moonlight as a large-scale application.
- Results: Moonlight is a 16B-parameter MoE model trained using Muon, and the authors release its implementation and intermediate checkpoints for further research.The released materials are intended to facilitate research on scalable optimization techniques for LLMs.
2 Methods
Muon scales to large-model training through weight decay, parameter-wise update scaling, and a distributed ZeRO-1-style implementation. These changes address growing weights, shape-dependent update RMS, and Muon’s full-matrix gradient requirement while retaining favorable memory and communication properties.
- 2.1 Background: Muon orthogonalizes matrix-valued momentum with Newton-Schulz iterations before applying the parameter update.The update uses O_t = Newton-Schulz(M_t) and W_t = W_{t−1} − η_tO_t; Nesterov-style momentum may be used in practice.
- 2.2 Scaling Up Muon: Weight decay prevents weights and layer-output RMS from growing too large during long, large-model training.In an 800M-parameter, 100B-token experiment, Muon with weight decay achieved lower validation loss than vanilla Muon and AdamW in the over-train regime.
- 2.2 Scaling Up Muon: Muon’s theoretical update RMS is 1/max(A, B) for a full-rank [A, B] matrix, so matrix shape can make updates too small or too large.Large max(A, B) can limit representational capacity, while small max(A, B) can cause training instabilities.
- 2.2 Scaling Up Muon: The revised update scales each matrix update to maintain consistent RMS across shapes and match AdamW’s typical 0.2 to 0.4 update-RMS range.The adjustment allows Muon to reuse learning-rate and weight-decay hyperparameters tuned for AdamW.
- 2.3 Distributed Muon: Muon uses one momentum buffer instead of AdamW’s two, while its communication workload is reported as (1, 1.25] times Distributed AdamW.Distributed Muon has larger end-to-end latency, but optimizer latency is usually negligible relative to the forward-backward pass, at 1% to 3%.
3 Experiments
Experiments evaluate Muon’s update scaling, compute efficiency, large-scale MoE pretraining, and optimization dynamics. Muon uses adjusted update scaling and weight decay, matches AdamW with fewer FLOPs, and shows advantages in model performance and update-spectrum diversity.
- 3.1 Consistent Update RMS: Adjusted LR scales each matrix’s learning rate by 0.2 · p max(A, B), while Update Norm directly normalizes Newton–Schulz updates to RMS 0.2.The experiments compare both methods against a baseline that maintains consistent RMS with AdamW.
- 3.1 Consistent Update RMS: Muon’s inconsistent update RMS becomes more problematic as matrix dimensions diverge, especially for matrices with highly unequal shapes.The analysis modifies MLP matrix shapes to expose this effect during early training.
- 3.1 Consistent Update RMS: Both Update Norm and Adjusted LR outperform the baseline, but Adjusted LR is selected for future experiments because it has lower cost.For [H, 4H] MLP weights, both methods produce roughly doubled weight RMS relative to the baseline; for [H, H] query weights, Adjusted LR preserves similar weight RMS.
- 3.2 Scaling Law of Muon: Muon requires about 52% training FLOPs to match AdamW performance under compute-optimal training.The comparison uses scaling-law experiments on dense Llama-architecture models with AdamW hyperparameter search and matched Muon update RMS.
- 3.3 Pretraining with Muon: Moonlight uses a DeepSeek-V3-small MoE architecture with 2.24B activated and 15.29B total parameters, or 3B and 16B including embeddings.The model is pretrained with a maximum context length of 8K and staged learning-rate, batch-size, and cooldown schedules.
- 3.3 Pretraining with Muon: Moonlight outperforms comparable models trained on equivalent tokens and lies on the Pareto frontier of model performance versus training budget.At 1.2T tokens, Moonlight also significantly outperforms the AdamW-trained Moonlight-A baseline; Muon especially excels on math- and code-related tasks.
- 3.4 Dynamics of Singular Spectrum: Muon produces higher SVD entropy than AdamW across checkpoints and weight-matrix groups, with the largest discrepancy in router weights.At the 1.2T-token checkpoint, over 90% of weight matrices have higher SVD entropy under Muon.
- 3.5 Supervised Finetuning (SFT) with Muon: Muon’s benefits persist in SFT when the model is both Muon-pretrained and Muon-finetuned, but not when the SFT optimizer differs from pretraining.The authors identify further exploration of optimizer consistency across stages as future work.
4 Discussions
The report identifies remaining research directions for Muon, including covering all parameters, extending beyond the spectral norm, and resolving optimizer mismatch between pretraining and fine-tuning.
- Muon currently uses Adam for some non-matrix parameters, leaving exclusive optimization of all parameters within Muon as an open research direction.
- Muon can be interpreted as steepest descent under the spectral norm, motivating extensions to general Schatten norms.
- Extending Muon to general Schatten norms may provide additional optimization capabilities and potentially improve results over the spectral norm implementation.
- Models pretrained with AdamW and fine-tuned with Muon, or vice versa, can show suboptimal performance, creating a barrier to using AdamW-pretrained checkpoints.
5 Conclusions
The report concludes that systematic improvements make Muon applicable to large-scale LLM training, including a 3B/16B-parameter MoE trained on 5.7 trillion tokens. It releases implementations, Moonlight, and intermediate checkpoints for further research.
- Muon was applied to a 3B/16B-parameter MoE model trained on 5.7 trillion tokens.
- Muon can effectively replace AdamW as the standard optimizer for large-scale LLM training, with advantages in training efficiency and model performance.
- The authors open-source the Muon implementation, Moonlight model, and intermediate training checkpoints to facilitate research on scalable optimization.
A Update RMS
The update-RMS analysis derives the RMS of a rank-r matrix update and evaluates Muon update scales against AdamW, finding similar performance for selected RMS settings.
- RMS(X) = √(r/mn) for the Muon update format X = U[:,:r]V[:r,:] with orthogonal factors.
- For full-rank matrices where r = m, the update RMS simplifies to √(1/n).
- Muon update RMS was tested at [0.05, 0.1, 0.2, 0.4, 0.8] against AdamW as the baseline.
- 0.2 RMS and 0.4 RMS performed similarly and substantially better than the other tested settings.
B AdamW Baseline Scaling Law
The AdamW baseline study searches scaling-law hyper-parameters under fixed FLOPs budgets and visualizes their loss landscapes. Each budget has a distinct minimum representing an optimal configuration.
- AdamW scaling-law experiments determine optimal model size N, training tokens D, learning rate η, and batch size B under computational budget C.
- The search varies model size N while adjusting token count D inversely to preserve C = 6ND.
- Figure 5 depicts loss landscapes over training tokens, learning rate, and batch size across varying FLOPs budgets.
- Each FLOPs-specific bowl-shaped loss curve has a distinct global minimum corresponding to an optimal hyper-parameter configuration.
C Model Architecture
Moonlight uses a DeepSeek-V3-Small-like architecture with several targeted modifications. These include omitting MTP layers, changing the auxfree bias update, and selecting a gate scaling factor to match dense-model output RMS.
- Muon is architecture-agnostic, and Moonlight is based on a model similar to DeepSeek-V3-Small.
- Moonlight omits multi-token prediction layers because MTP showed no significant pretraining benefit in the authors’ experiments.
- Auxfree Bias Update: The auxfree bias update subtracts the average sign of expert violating ratios to control bias magnitude without changing topk selection.
- Gate Scaling Factor: Moonlight uses a gate scaling factor of 2.446 to produce output RMS similar to dense models.DeepSeek-V2-Lite used no factor, while DeepSeek-V3 used 2.5.
D Training Stability
Moonlight training was smooth in loss and gradient norm, but maximum attention logits rose above 100 in some layers early in training. Extreme logits remained sparse and later declined, while RMSNorm gamma decay supports stability by limiting output RMS.
- Moonlight training showed no loss or gradient norm spikes.The authors describe the training process as very smooth.
- Maximum Attention Logit: Maximum attention logits exceeded 100 in specific layers during the initial training phase despite stable loss and gradient norm.
- Large Attention Logits Ratio: The proportion of attention logits above 100 stayed near 10^-4, indicating that extreme values were sparse.
- Large Attention Logits Ratio: Maximum attention logits gradually decreased as training progressed, suggesting healthier optimization dynamics.
- RMSNorm Gamma Weight Decay: Applying weight decay to RMSNorm gamma is crucial for stability because it prevents excessively high layer output RMS values.
E Comparison with More Expensive Models
Moonlight is compared with publicly available models trained using greater computational resources, including Llama3.1-8B, Gemma-9B, and Qwen2.5-7B. The comparison covers multiple benchmarks, with GSM8k shown explicitly.
- Moonlight is compared against Llama3.1-8B, Gemma-9B, and Qwen2.5-7B, which were trained with greater computational resources.
- The comparison evaluates the models across various benchmarks.
- GSM8k Comparison: Figure 8 specifically presents GSM8k performance comparisons between Moonlight and comparable models.
F Singular Value Distributions of Weight Matrices
The paper visualizes normalized singular-value spectra for attention and feed-forward weight matrices. Muon generally produces flatter spectra than AdamW, supporting a more diverse spectrum of updates.
- Muon- and AdamW-optimized weight matrices are compared using singular values plotted in descending order and normalized by each matrix’s largest value.
- For most weight matrices, Muon produces flatter singular-value distributions than AdamW.
- The flatter distributions further confirm the hypothesis that Muon provides a more diverse spectrum of updates.
- Figures 9 and 10 cover weight matrices in attention and feed-forward network layers, respectively.The captions identify the matrices by their roles in projections, latent spaces, keys, queries, and expert models.