SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales

Empirical and systems study of higher-order preconditioned optimizers at production pretraining scale. Diagnoses SOAP's large-batch "slingshot" instability — stale preconditioner eigenbases cause oscillations that a 621M model damps out but that catastrophically diverge an 8B dense model — and eliminates it with per-step QR orthogonalization, real-time eigenbasis updates, and gradient covariance estimation via KL-divergence minimization (the paper's KL-SOAP variant). Optimizers are compared under update-RMS matching so learning rates transfer fairly across AdamW, Muon, and SOAP.

On multi-billion-parameter dense (8B) and MoE configurations — including Nemotron-3-Nano-30B-A3B (128 experts, top-6) and a Nemotron-3-72B-A8B config (512 experts, top-6) — trained on ~3T-token horizons, SOAP and Muon consistently outperform AdamW and stay stable at global batch sizes up to 100M tokens, past AdamW's critical batch size where its quality degrades. The systems contribution is a layer-wise distributed optimizer compatible with Megatron-LM that balances memory and hides communication without approximating the optimizer computation, released in the Apache-2.0 Emerging-Optimizers codebase. Senior authors include Bryan Catanzaro and Mohammad Shoeybi. Extends the head-to-head-optimizer line of Stanford's Fantastic Pretraining Optimizers and Moonshot AI's Muon scaling work from academic scales to frontier batch sizes and token budgets.