Beyond AdamW

Higher-order optimizers have promised faster convergence than AdamW for a while, with computational cost and numerical stability as the standing objections. This NVIDIA work adapts them for large-scale pretraining, identifying instabilities in SOAP at large batch sizes and eliminating the loss spikes with per-step QR orthogonalization and improved preconditioning, then running a unified study of SOAP, Muon, and AdamW under update-RMS matching for fair learning rate transfer. On multi-billion-parameter models trained over trillions of tokens, SOAP and Muon consistently beat AdamW, and at batch sizes up to 100M tokens for next-token prediction they hold stability and quality while AdamW degrades. A layer-wise distributed optimizer compatible with Megatron-LM balances memory and hides communication without approximating the optimizer math.