SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales
Mikail Khona, Aditya Vavre, Boxiang Wang, Deyu Fu, Hao Wu, Mike Chrzanowski, Bryan Catanzaro, Dheevatsa Mudigere, Jeff Pool, Michael Lightstone, Mohammad Shoeybi, Mostofa Patwary, Nima Tajbakhsh, Tijmen Blankevoort
cs.LG, cs.AI
2026-07-14
NVIDIA stabilizes SOAP and Muon for billion-parameter, trillion-token pretraining via per-step QR orthogonalization and a layer-wise distributed optimizer; at large batch AdamW degrades while Muon holds MMLU above 74 on 3T tokens.
AdamW is the default optimizer for LLM pretraining, but a family of higher-order optimizers (Muon, SOAP) promises faster convergence. Two things block them at scale: computing the preconditioner costs more than AdamW, and they go numerically unstable at very large batch, so no one had pushed them cleanly to multi-billion parameters and trillions of tokens. This NVIDIA paper runs SOAP and Muon through large-scale LLM pretraining and patches the engineering along the way.
The team first pins down why SOAP diverges at large batch: the reference implementation refreshes eigenbases only every 10 steps and excludes the current gradient, so the preconditioner is always stale. Weight-gradient norms oscillate and the language-modeling loss spikes. They call it slingshot instability.
Fixes come in layers. Per-step QR orthogonalization recomputes the QR decomposition of the preconditioner matrices every step, with the current gradient included, killing the lag. KL-Shampoo covariance estimation replaces standard Kronecker-factor accumulation with a KL-regularized estimate, cutting the condition number from κ to √κ and reducing numerical error in the eigendecomposition. Fair comparison is itself a contribution: optimizers move weights by very different amounts, so transferring learning rates directly is unfair. Update-RMS matching matches the RMS norm of parameter updates and applies a roughly 0.2 correction to Muon to offset EMA-momentum damping, putting SOAP, Muon, and AdamW on the same ruler.
The last piece is a layer-wise distributed optimizer for Megatron-LM: parameter matrices are distributed round-robin by size, each GPU updates its own layers, communication is an async all-gatherv bucketed by model execution order, and computation overlaps fetching the next bucket. No approximation of the optimizer math, so convergence is preserved.
Across an 8B dense GPT, Qwen3-30B-A3B, Nemotron-3-Nano-30B, and a 72B-A8B, trained to trillions of tokens with batch sizes up to 100M tokens. The headline: at large batch AdamW degrades while SOAP and Muon hold. On Nemotron-3-Nano at 3T tokens:
| Setting | MMLU | HumanEval | Math500 |
| AdamW 1× batch | 73.38 | 63.26 | 71.50 |
| Muon 1× batch | 73.71 | 62.56 | 75.50 |
| Muon 2× batch | 74.80 | 64.79 | 74.05 |
| Muon 3× batch | 74.00 | 62.22 | 73.25 |
Doubling or tripling batch barely dents Muon's downstream quality, while AdamW starts breaking under the same scaling. On Qwen3-30B, KL-SOAP holds a small but consistent edge over Muon in cross-entropy loss across 1×/2×/4× batch. As a side note, Muon's Newton-Schulz approximate orthogonalization matches exact polar decomposition (MOP) in loss, so the approximation is good enough.
For teams pretraining big models, this is solid evidence for "can we switch off AdamW once it hits a wall." The answer is that Muon and SOAP stay stable as batch scales up, and the layer-wise distributed implementation makes them runnable on big clusters. The code is open-sourced as Emerging-Optimizers and is ready to try.
The ε numerical-stability term was not tuned systematically (1e-8 for SOAP, 1e-7 for Muon), and how it should scale with model size is left to future work. Conv1D filters in Mamba-style architectures are geometrically ill-motivated for orthogonalization, so the method's boundary does not cover them. The 100M-token batch is the tested upper bound; whether new instabilities appear beyond it is not addressed.