SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales
本文提出SOAP与Muon优化器，通过算法改进实现大规模LLM预训练中的稳定性与效率提升。
核心发现
方法论
研究采用改进的预条件梯度方法，结合QR正交化和谱预条件策略，解决SOAP在大批次训练中的不稳定性。通过update-RMS匹配确保不同优化器间学习率的公平转移，评估Muon的正交化质量。实验证明在亿级参数模型上，SOAP与Muon在训练稳定性和模型性能上优于AdamW，尤其在批次达1亿tokens时表现出明显优势。系统层面，开发了兼容Megatron-LM的分层分布式优化器，优化内存与通信，保持收敛性。
关键结果
- 在多亿参数模型上，SOAP和Muon在批次规模达100M tokens时，训练稳定性和模型质量优于AdamW，后者在大批次下出现性能下降。
- 通过引入逐步QR正交化和KL-divergence估算，成功消除SOAP在大批次中的损失波峰，提升训练稳定性。
- 在不同模型（如8B Dense GPT、3B MoE、72B混合模型）上，优化器性能持续优越，尤其在超大批次训练中，Muon表现出更强的鲁棒性和收敛速度。
研究意义
该研究突破了高阶优化器在大规模预训练中的应用瓶颈，为训练亿级参数模型提供了更稳定、更高效的算法方案。推动了深度学习在模型规模扩展中的技术边界，特别是在超大批次训练和系统架构优化方面，具有重要的理论与实践意义。这不仅改善了模型训练的收敛速度，也为未来大规模模型的部署提供了技术基础，助力AI行业实现更大规模、更高性能的模型开发。
技术贡献
技术创新包括引入逐步QR正交化策略，结合KL-divergence的协方差估算，解决SOAP在大批次中的数值不稳定问题。提出层级分布式优化器，兼容Megatron-LM架构，优化内存与通信效率，保持高阶优化器的收敛优势。系统层面，开发了多项加速技术，支持高效大规模训练。实验验证显示，SOAP与Muon在亿级参数模型中，超越AdamW，显著提升训练稳定性与模型性能。
新颖性
首次系统性评估SOAP与Muon在超大规模模型中的表现，提出解决SOAP大批次不稳定的算法改进，结合系统优化实现高效分布式训练。创新点在于算法与系统的深度结合，突破了高阶优化器在实际大规模训练中的应用限制，推动了高阶优化在工业界的落地。
局限性
- 算法在极端大批次（超过100M tokens）时仍存在数值敏感性，需进一步优化正交化与预条件策略。
- 系统实现复杂，依赖特定硬件与通信架构，迁移到不同平台可能面临挑战。
- 实验主要集中在特定模型和数据集，泛化到其他任务和架构仍需验证。
未来方向
未来将探索更高效的预条件策略，降低系统复杂度，拓展优化器在多任务、多模态模型中的应用。计划结合自动调参技术，进一步提升训练稳定性与性能。还将研究多GPU异步通信机制，减少通信瓶颈，推动高阶优化器在更大规模模型中的应用落地。
AI 总览摘要
随着深度学习模型规模不断扩大，传统的优化算法如AdamW逐渐暴露出在大批次训练中的稳定性与效率瓶颈。本文引入两种高阶优化器——SOAP与Muon，结合算法改进与系统工程，显著提升了亿级参数模型的训练表现。通过引入逐步QR正交化和谱预条件策略，有效解决了SOAP在超大批次下的数值不稳定问题。实验结果显示，在批次达1亿tokens时，SOAP与Muon不仅保持训练稳定，还在模型性能上优于AdamW，验证了其在大规模预训练中的潜力。此外，研究还开发了与Megatron-LM兼容的分层分布式优化器，优化内存与通信，确保高阶优化器的高效实现。这些技术创新为未来超大规模模型的训练提供了坚实的基础，推动深度学习向更高的规模与性能迈进。尽管如此，算法在极端条件下仍面临数值敏感和系统复杂的挑战，未来工作将集中在算法优化与硬件适配上，以实现更广泛的应用推广。
深度分析
研究背景
近年来，深度学习模型规模不断突破，从数亿到万亿参数，推动了自然语言处理、计算机视觉等领域的飞跃。优化算法作为模型训练的核心，经历了从SGD到AdamW的演变，解决了梯度调整与正则化问题。高阶优化器如Shampoo、SOAP和Muon's出现，旨在利用二阶信息提升收敛速度与模型质量，但在大规模训练中面临数值不稳定和计算复杂的挑战。系统层面，分布式训练架构不断优化，支持更大规模模型的训练需求。尽管如此，如何在保证稳定性和效率的同时，充分利用高阶优化器的优势，仍是当前研究的热点。
核心问题
在超大规模模型训练中，批次规模不断扩大带来数值不稳定、收敛缓慢等问题。SOAP作为高阶优化器，虽具潜力，但在大批次训练时出现损失波峰，影响模型性能。AdamW虽易于实现，但在超大批次下表现出性能下降和不稳定。系统层面，如何设计兼容高阶优化器的分布式架构，减少通信与内存开销，确保算法的可扩展性，也是亟待解决的问题。这些挑战限制了高阶优化器在工业界的应用推广。
核心创新
本研究提出了多项创新：一是引入逐步QR正交化策略，解决SOAP在大批次中的数值不稳定问题；二是结合KL-divergence的协方差估算，提升预条件器的更新稳定性；三是设计层级分布式优化器，兼容Megatron-LM架构，优化内存与通信效率。这些创新使得高阶优化器在亿级参数模型中实现稳定训练，突破了传统方法的局限。系统方面，开发了多项加速技术，支持高效大规模训练，为工业界大模型训练提供了新工具。
方法详解
- �� 识别SOAP在大批次中的不稳定性，分析预条件器与梯度统计的滞后问题。• 引入逐步QR正交化，确保每步梯度旋转的数值稳定性。• 结合KL-divergence估算协方差矩阵，动态更新Kronecker因子，消除损失波峰。• 设计层级分布式优化器，将参数层划分到不同GPU，支持矩阵级预条件计算。• 实现兼容Megatron-LM的分层通信策略，优化内存与通信开销。• 采用update-RMS匹配，确保不同优化器间学习率的公平转移。• 在多亿参数模型上进行大规模训练实验，验证算法效果。
实验设计
使用8B Dense GPT、3B MoE和72B混合模型，训练数据为1T和3T tokens，采用特定超参数配置。比较AdamW、SOAP和Muon在不同批次规模下的性能，重点关注批次达1亿tokens时的稳定性与模型质量。通过AB测试和指标如训练损失、验证准确率，验证算法的鲁棒性。实验还包括正交化质量评估和系统性能分析，确保算法在实际大规模环境中的可行性。
结果分析
在亿级参数模型中，SOAP与Muon在100M tokens批次下，训练稳定性优于AdamW，模型性能提升约3-5%。在超大批次（超过50M tokens）时，AdamW表现出明显的性能下降和训练不稳定，而Muon保持低损失和高质量输出。引入算法改进后，SOAP的损失波峰显著减少，训练过程更平滑。系统优化实现了高效的通信与内存管理，支持更大规模模型的训练，验证了算法与系统的结合优势。
应用场景
该技术适用于训练超大规模语言模型、图像识别模型等，特别是在需要高效利用大规模分布式硬件资源的场景。企业和研究机构可借助此方法，加快模型训练速度，提升模型性能，降低训练成本。未来，结合自动调参和异步通信，将进一步推动工业界大模型的快速部署。
局限与展望
算法在极端大批次（如超过100M tokens）时仍存在数值敏感问题，需进一步优化正交化策略。系统实现复杂，依赖特定硬件架构，迁移难度较大。实验主要集中在特定模型和数据集，泛化到其他任务和架构仍需验证。未来需简化算法复杂度，提升系统的适应性。
通俗解读 非专业人士也能看懂
想象你在厨房做饭，锅里放满了各种食材。传统的炒菜方法就像用普通锅铲，简单快速，但可能会炒得不均匀或焦掉。而高阶优化器就像用一把智能的厨具，能根据食材的状态自动调整火候和翻炒方式，让菜炒得更均匀、更美味。SOAP和Muon就是这样的智能厨具，它们能更好地理解食材的结构（模型参数的关系），在大批量炒菜时保持火候稳定，不会出现焦糊或生硬。通过算法的改进和系统的优化，厨师（训练者）可以用更少的时间做出更好吃的菜（模型），而且在大厨房（大模型）中也能保持高效。这就像用高科技厨具，让厨房变得更智能、更高效，做出更美味的饭菜。
简单解释 像给14岁少年讲一样
想象你在学校的食堂里准备大批量的饭菜，平时用普通锅铲炒菜，速度快但容易糊锅或不均匀。现在，厨师们发明了一种新型的智能厨具，能根据食材的状态自动调整火候和翻炒方式，确保每一份饭都炒得刚刚好。这就像论文中的SOAP和Muon优化器，它们能在训练大模型时，像智能厨具一样，保持火候（训练稳定）并让菜（模型）变得更好吃（性能更优）。这些新工具不仅让炒菜（训练）变得更快、更稳定，还能在更大的厨房（更大规模模型）中工作。虽然还需要一些改进，但未来用这种智能厨具做饭，学校的饭菜一定会越来越好吃！
原文摘要
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