Preprint 1 Mention
SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales
Mikail Khona2026
Aditya VavreBoxiang Wang
Low Citations
0 citations · Artificial Intelligence
Open Access

TLDR

Some new training methods help big language models learn faster and more reliably than older methods, especially when working with huge amounts of data.

Summary

1 Study Aim

The paper aims to improve and test advanced training methods, called higher-order optimizers, for large language models (LLMs). The authors focus on making these optimizers, specifically SOAP (Second-Order Adaptive Preconditioner) and Muon, more stable and practical for use with very large datasets and models. They also compare these methods to the widely used AdamW optimizer to see which works best at large scales. The study wants to find out if new training methods can make big language models learn faster and more reliably than older ones.

2 Study Design

The researchers first identify problems with the SOAP optimizer when used with very large batch sizes (large groups of data processed at once). They then make changes to the algorithm, such as adding per-step QR orthogonalization (a mathematical process to keep calculations stable) and better preconditioning strategies (ways to adjust the optimizer for smoother learning). The team runs experiments on models with billions of parameters, training them on trillions of words. They compare SOAP, Muon, and AdamW using a method called update-RMS matching to ensure a fair comparison. They also develop a new distributed optimizer that works with Megatron-LM (a system for training large models), making it efficient for large-scale use. The researchers test new and old training methods on very large language models to see which ones work best and stay stable.

3 Findings

The study reveals that both SOAP and Muon optimizers consistently outperform AdamW when training very large language models, especially at batch sizes up to 100 million tokens. The authors show that their improvements to SOAP remove instability issues, such as sudden spikes in loss, and allow for stable training even with massive datasets. AdamW, in contrast, becomes less stable and effective at these scales. The new layer-wise distributed optimizer enables efficient and accurate training without sacrificing the benefits of the advanced optimizers. The authors also provide a public codebase to help others use and build on these methods. They recommend using SOAP or Muon for large-scale language model training to achieve better results and stability. The new training methods help big language models learn better and more reliably than older methods, especially when using huge amounts of data.

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

Referenced In