SOAP和Muon优化器在大规模语言模型训练中表现优于AdamW，且可稳定处理超大批次。
SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales
- 改进SOAP算法，通过每步正交化与预条件策略解决大批次训练中的不稳定性。
- 在万亿级令牌、百亿参数模型上，SOAP和Muon在超大批次下仍保持稳定且性能超越AdamW。
- 设计分层分布式优化器，兼容Megatron-LM，实现高效通信与计算，适合大规模训练场景。
高阶优化器如Muon和SOAP相比AdamW具有更快的收敛速度，但其计算开销和数值稳定性问题限制了在大规模语言模型预训练中的应用。本文通过改进预条件梯度方法，克服了大规模训练中的实际挑战。我们发现SOAP在大批次下存在不稳定性，提出每步QR正交化与优化的预条件策略，消除损失尖峰并实现稳定训练。进一步，通过更新均方根匹配，对SOAP、Muon和AdamW进行统一实证比较。在多百亿参数模型、万亿级令牌数据上的实验表明，SOAP和Muon在测试规模下持续优于AdamW。尤其在高达1亿令牌批次的下一词预测任务中，二者保持训练稳定与高质量，而AdamW性能下降。为此，我们提出一种与Megatron-LM兼容的分层分布式优化器，在不近似优化计算的前提下平衡内存与通信开销，保留其收敛优势。此外，我们识别并实施特定系统优化以加速分层实现。为支持研究社区，我们开源了包含新兴优化算法的代码库：https://github.com/NVIDIA-NeMo/Emerging-Optimizers。
原文摘要 · Abstract (English)
Higher-order optimizers such as Muon and SOAP offer faster convergence than AdamW, but their computational cost and numerical stability challenges have limited adoption at scale. In this work, we adapt and enhance preconditioned gradient methods to overcome the practical challenges of large-scale LLM pretraining. We first identify instabilities in SOAP at large batch sizes and propose algorithmic modifications including per-step QR orthogonalization and improved preconditioning strategies that eliminate loss spikes and enable stable training in these regimes. We then present a unified empirical study of SOAP, Muon, and AdamW using update-RMS matching to ensure fair learning rate transfer across optimizers. As part of this analysis, we empirically evaluate the orthogonalization quality of Muon. Our experiments on multi-billion-parameter models trained on trillions of tokens reveal that SOAP and Muon consistently outperform AdamW at the scales we tested. Notably, at batch sizes of up to 100M tokens for next-token prediction, these optimizers maintain training stability and quality while AdamW degrades. To enable efficient training at large scale, we introduce a layer-wise distributed optimizer compatible with Megatron-LM. Our implementation balances memory and hides communication while avoiding approximations to the optimizer computations, thus retaining their convergence benefits. Additionally, we identify and build specific system-level improvements to further accelerate our layer-wise implementation. To support the research community, we release a codebase that contains emerging algorithms for optimization: https://github.com/NVIDIA-NeMo/Emerging-Optimizers
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。