SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales
This paper introduces SOAP and Muon optimizers, with algorithmic improvements enabling stable, efficient large-scale LLM pretraining at billion-parameter scales.
Key Findings
Methodology
The study employs enhanced preconditioned gradient algorithms, integrating per-step QR orthogonalization and spectral preconditioning to address SOAP's instability at large batch sizes. Update-RMS matching ensures fair learning rate transfer across optimizers, while empirical evaluation of Muon's orthogonalization quality demonstrates superior stability. Experiments on models up to 72B parameters trained on trillions of tokens show SOAP and Muon outperform AdamW, particularly at batch sizes up to 100M tokens, where AdamW's performance degrades. A layer-wise distributed optimizer compatible with Megatron-LM is developed, balancing memory and communication, preserving convergence benefits. System-level accelerations further improve scalability.
Key Results
- In training billion-parameter models, SOAP and Muon maintain stability and outperform AdamW at batch sizes up to 100M tokens, with improvements in loss reduction and model quality. AdamW exhibits instability beyond critical batch sizes, with loss spikes and degraded performance.
- Algorithmic modifications, including per-step eigenbasis updates and KL-divergence-based covariance estimation, effectively eliminate loss spikes, ensuring stable training at large scales.
- The distributed layer-wise optimizer implementation achieves efficient communication and memory management, enabling training of models with hundreds of billions of parameters while retaining optimizer convergence advantages.
Significance
This work advances the practical application of second-order and spectral optimizers in large-scale language model training, addressing longstanding stability and efficiency challenges. It bridges the gap between theoretical benefits and real-world deployment, enabling faster convergence, better model quality, and scalable training infrastructure. The innovations open pathways for training even larger models, reducing training costs and accelerating AI development. The open-source release fosters community engagement, promoting further research into emerging optimization algorithms and system architectures, thus shaping the future of large-scale AI systems.
Technical Contribution
Key contributions include the development of a robust algorithmic framework combining QR orthogonalization and spectral covariance estimation to stabilize SOAP at large batch sizes. The design of a layer-wise distributed optimizer compatible with Megatron-LM architecture allows efficient large-scale deployment without sacrificing optimizer benefits. The integration of update-RMS matching ensures fair hyperparameter transfer, facilitating comprehensive comparative studies. Experimental validation on models up to 72B parameters demonstrates consistent outperformance over AdamW, establishing new benchmarks for optimizer scalability and stability.
Novelty
This is the first comprehensive evaluation of SOAP and Muon in billion-parameter regimes, with novel algorithmic modifications addressing SOAP's instability in large-batch training. The proposed layer-wise distributed optimizer uniquely balances memory, communication, and convergence, enabling high-order optimizers to be practically applied at unprecedented scales. The combination of spectral preconditioning, orthogonalization, and system-level engineering represents a significant step forward in optimizer research, bridging theoretical advances with scalable system implementation.
Limitations
- Despite improvements, the algorithms still face numerical sensitivity at extreme batch sizes (>100M tokens), requiring further refinement of orthogonalization procedures.
- System complexity and hardware dependency may limit portability across different hardware architectures and distributed frameworks.
- Experiments are primarily on specific models and datasets; generalization to other architectures, tasks, or modalities remains to be validated.
Future Work
Future directions include refining the orthogonalization and covariance estimation methods for enhanced stability, reducing system complexity for broader applicability, and extending the framework to multi-task and multimodal models. Exploring asynchronous communication strategies and automated hyperparameter tuning could further improve scalability and robustness. Additionally, integrating these optimizers into mainstream training frameworks will facilitate wider adoption and accelerate large-scale AI research.
AI Executive Summary
As the scale of language models continues to grow exponentially, the limitations of traditional first-order optimizers like AdamW become increasingly apparent. These methods, while computationally simple, struggle with stability and convergence at billion-parameter scales, especially under large batch training regimes. To address this, the paper introduces SOAP and Muon, two advanced second-order and spectral optimizers, which leverage curvature-aware preconditioning to accelerate training and improve stability.
The core innovation lies in algorithmic modifications—particularly, per-step QR orthogonalization and covariance estimation via KL-divergence—that eliminate loss spikes and stabilize training at batch sizes up to 100 million tokens. These improvements enable the optimizers to maintain high training quality, outperforming AdamW in both convergence speed and final model performance across models ranging from 8B to 72B parameters.
Complementing the algorithmic advances, the authors develop a layer-wise distributed optimizer compatible with Megatron-LM, balancing memory and communication to support efficient large-scale training. This system-level engineering ensures that the benefits of high-order optimizers are preserved in practical, industrial-scale environments.
Experimentally, the results demonstrate that SOAP and Muon consistently outperform AdamW at large batch sizes, with Muon showing particular robustness due to spectral orthogonalization. These findings mark a significant step toward scalable, stable training of massive models, opening new avenues for AI research and deployment. Despite some remaining challenges—such as numerical sensitivity at extreme scales—the work sets a new benchmark for optimizer design, system integration, and large-scale model training.
Deep Analysis
Background
近年来，深度学习模型规模不断扩大，从数亿参数到万亿参数级别，带动了自然语言处理、计算机视觉等领域的飞跃。优化算法作为模型训练的核心，从SGD到AdamW的演变，解决了梯度调整与正则化问题。高阶优化器如Shampoo、SOAP和Muon的出现，旨在利用二阶信息提升收敛速度和模型性能，但在大规模训练中面临数值不稳定和计算复杂的挑战。系统架构不断优化，支持更大规模模型的训练需求，但如何在保证稳定性和效率的同时充分利用高阶优化器，仍是研究热点。
Core Problem
在超大规模模型训练中，批次规模不断扩大带来数值不稳定、收敛缓慢等问题。SOAP虽具潜力，但在大批次训练时出现损失波峰，影响模型性能。AdamW虽易于实现，但在超大批次下表现出性能下降和不稳定。系统设计难以兼顾高阶优化器的矩阵结构需求，通信与内存瓶颈限制了其应用。解决这些问题，成为推动大模型训练的关键。
Innovation
本研究提出多项创新：一是引入逐步QR正交化策略，解决SOAP在大批次中的不稳定性；二是结合KL-divergence的协方差估算，提升预条件器的稳定性；三是设计层级分布式优化器，支持矩阵级预条件计算，兼容Megatron-LM架构。系统层面，开发了多项加速技术，支持高效大规模训练。这些创新使得高阶优化器在亿级参数模型中实现稳定训练，突破传统限制，推动深度学习模型的规模极限。
Methodology
- �� 识别SOAP在大批次中的不稳定性，分析预条件器与梯度统计滞后问题。• 引入逐步QR正交化，确保每步梯度旋转的数值稳定性。• 结合KL-divergence估算协方差矩阵，动态更新Kronecker因子，消除损失波峰。• 设计层级分布式优化器，将参数层划分到不同GPU，支持矩阵级预条件计算。• 实现兼容Megatron-LM的分层通信策略，优化内存与通信开销。• 采用update-RMS匹配，确保不同优化器间学习率的公平转移。• 在多亿参数模型上进行大规模训练实验，验证算法效果。
Experiments
使用8B Dense GPT、3B MoE和72B混合模型，训练数据为1T和3T tokens，设定特定超参数。比较AdamW、SOAP和Muon在不同批次规模下的性能，重点关注100M tokens批次的稳定性。通过训练损失、验证指标评估算法鲁棒性，结合正交化质量和系统性能分析，验证算法在实际大规模环境中的适用性。
Results
在亿级参数模型中，SOAP与Muon在100M tokens批次下，训练稳定性优于AdamW，模型损失降低约3-5%。超大批次（超50M tokens）时，AdamW表现出明显性能下降和训练不稳定，而Muon保持低损失和高性能。算法改进显著减少损失波峰，训练过程更平滑。系统优化实现高效通信和内存管理，支持大模型训练，验证了算法与系统的结合优势。
Applications
该技术适用于训练超大规模语言模型、图像识别模型等，特别是在需要高效利用大规模分布式硬件资源的场景。企业和研究机构可借助此方法，加快模型训练速度，提升模型性能，降低成本。未来结合自动调参和异步通信，将推动工业界大模型的快速部署。
Limitations & Outlook
算法在极端大批次（如超过100M tokens）仍存在数值敏感问题，需进一步优化正交化策略。系统实现复杂，依赖硬件架构，迁移难度大。实验主要在特定模型和数据集上，泛化到其他任务和架构仍需验证。未来需简化算法复杂度，提升系统适应性。
Plain Language Accessible to non-experts
想象你在厨房做饭，锅里放满了各种食材。传统用普通锅铲炒菜，速度快但容易炒得不均或焦掉。而高阶优化器就像用一把智能厨具，能根据食材状态自动调节火候和翻炒方式，让菜炒得更均匀、更美味。SOAP和Muon就像这样的智能厨具，能更好理解食材的结构（模型参数关系），在大批量炒菜时保持火候稳定，不会出现焦糊或生硬。通过算法改进和系统优化，厨师（训练者）用更少时间做出更好吃的菜（模型），在大厨房（大模型）中也能保持高效。这就像用高科技厨具，让厨房更智能、更高效，做出更美味的饭菜。
ELI14 Explained like you're 14
想象你在学校食堂准备大批量的饭菜，用普通锅铲炒菜，速度快但容易糊锅或不均。现在，厨师们发明了一种新型智能厨具，能根据食材状态自动调节火候和翻炒，让每份饭都炒得刚刚好。这就像论文中的SOAP和Muon优化器，它们能在训练大模型时，像智能厨具一样，保持火候（训练稳定）并让模型变得更好。用这些新工具，不仅炒菜更快、更稳，还能在更大的厨房里工作。虽然还需改进，但未来用这种厨具做饭，饭菜一定更好吃！
Abstract
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