One-sentence Summary

NVIDIA researchers address SOAP's instabilities at large batch sizes with per-step $QR$QR orthogonalization and improved preconditioning, and through a unified empirical study using update-$RMS$RMS matching, show that SOAP and Muon consistently outperform AdamW on multi-billion-parameter models trained on trillions of tokens, while also introducing a layer-wise distributed optimizer compatible with Megatron-LM that efficiently scales training without approximations.

Key Contributions

• Algorithmic modifications to SOAP, including per-step QR orthogonalization and improved preconditioning strategies, eliminate loss spikes and enable stable training at large batch sizes.
• A unified empirical study of SOAP, Muon, and AdamW using update-RMS matching for fair learning rate transfer shows that SOAP and Muon consistently outperform AdamW on multi-billion-parameter models trained on trillions of tokens, maintaining stability at batch sizes up to 100M tokens where AdamW degrades.
• A layer-wise distributed optimizer compatible with Megatron-LM balances memory and hides communication while avoiding approximations, and incorporates system-level improvements to accelerate large-scale training.

Introduction

The choice of optimizer is central to large-scale model training, directly shaping memory sharding, communication overhead, and convergence speed. While element-wise optimizers like AdamW are simple to scale across distributed systems, they ignore the correlational structure of gradient tensors. Higher-order methods that capture loss-landscape curvature can converge faster and use larger batch sizes, but their complexity has historically prevented them from scaling to frontier models. Recently, tensor-wise optimizers such as SOAP and Muon have emerged as a compromise, offering second-order benefits with manageable overhead, yet they still encounter stability and scalability hurdles on large mixture-of-experts architectures. The authors address these gaps by enabling Muon and SOAP to maintain token-efficient convergence at global batch sizes up to 100M tokens, identifying and fixing a critical instability in SOAP’s preconditioner computation under large-batch training, and providing a layer-wise distributed implementation within Megatron-LM along with an open-source optimizer research library.

Dataset

The authors draw training data from the Nemotron-3 dataset and use two fixed-size subsets for all pretraining experiments:

• 1‑trillion‑token subset
• 3‑trillion‑token subset

Both subsets are directly extracted from the full Nemotron-3 corpus; the paper does not detail the dataset’s original composition, domain breakdown, or filtering rules.

Usage in training The subsets feed the pretraining of multiple model architectures: 8B‑parameter dense GPT models, 3B‑active/30B‑total MoE transformer models, and an 8B‑active/72B‑total hybrid Mamba‑Transformer MoE. The experiments treat the subsets as interchangeable training corpora; the final selection (1T or 3T) depends on the specific optimizer comparison being run.

Processing in the pipeline

• All data is tokenized and packed into sequences of 8192 tokens.
• The global batch size is fixed at 3072 samples (equivalent to 25 million tokens per optimization step).
• No additional cropping, deduplication, or metadata construction is mentioned beyond the original preprocessing of Nemotron-3.

Method

The authors introduce a production-oriented training framework that integrates matrix‑aware optimizers, specifically Muon and SOAP, into the Megatron‑LM ecosystem. The core challenge is that conventional data‑parallel sharding strategies (e.g., ZeRO‑1, FSDP) fragment weight matrices into element‑wise slices, which breaks the preconditioning and orthogonalization operations that require full 2D parameter tensors. The solution is a layer‑wise distributed optimizer combined with careful handling of tensor parallelism and batch‑size scaling, enabling Muon and SOAP to run at scale without excessive communication overhead.

The layer‑wise optimizer assigns entire parameter matrices to distinct GPUs rather than splitting individual tensors. Load balancing is achieved by sorting all parameter matrices by size and distributing them across data‑parallel ranks in a round‑robin sequence. This guarantees that each GPU holds a roughly equal memory footprint while preserving the matrix structure needed for Muon’s Newton‑Schulz iterations and SOAP’s Kronecker‑factor updates. Each GPU independently updates its assigned layers and then flattens the results into a buffer for the next forward pass.

To hide communication, the system overlaps parameter synchronization with computation. Parameter matrices are grouped into sequential buckets that align with the model’s execution order. For each bucket, an asynchronous all‑gather collective fetches the updated parameters from their owning GPU while the forward pass is processing the current bucket’s weights. Because the layer‑wise assignment leads to slightly uneven parameter counts across GPUs, a variable‑sized all‑gather‑V primitive is used instead of standard uniform collectives, avoiding padding or further fragmentation.

For tensor‑parallel settings, two modes preserve the full‑matrix structure when preconditioning is applied to the entire layer’s weights. In the duplicated mode, all GPUs inside the tensor‑parallel domain first all‑gather the weights; each GPU then runs the Newton‑Schulz iteration independently. This mode is preferred for smaller layers where communication dominates. In the distributed mode, intermediate matrix‑multiplication results are all‑reduced during each Newton‑Schulz iteration, making it suitable for larger layers where computation is the bottleneck. Both modes normalize the inputs using the statistics of the whole layer’s weight, ensuring mathematical equivalence to a non‑tensor‑parallel execution.

Underpinning the optimizer strategy is a principled approach to batch‑size scaling. When the global batch size is changed, the square‑root scaling rule adjusts the learning rate proportionally to $\sqrt{B'/B}$B′/B​ to preserve the variance of parameter updates. For mixture‑of‑experts models, the effective batch size seen by an expert is $B_{\mathrm{Global}} \times k/N$BGlobal​×k/N, which is dramatically smaller than the global batch due to sparse top‑$k$k routing. Consequently, the dense and shared parameters absorb the full large‑batch burden, making the stability of Muon and SOAP on these dense components critical. Together, the layer‑wise sharding, asynchronous communication, and batch‑size scaling rules form a system that unlocks the convergence benefits of higher‑order optimizers at the largest training scales.

Experiment

The evaluation compares Muon and SOAP against AdamW on 8B dense and up to 8B‑active/72B‑total MoE models using 1T‑3T token subsets of Nemotron‑3, with optimizer learning rates aligned via update RMS matching. Muon scales to large batch sizes without the instability that plagues AdamW, maintaining lower loss and fewer spikes, while SOAP requires per‑step eigenbasis updates including the current gradient and KL‑regularized covariance estimation to avoid “slingshot” divergence. In controlled comparisons, both Muon and SOAP outperform AdamW, with KL‑SOAP holding a slight advantage in cross‑entropy loss, though systematic epsilon tuning remains an open direction.

All evaluated architectures share 32 attention heads with a head dimension of 128 and Grouped-Query Attention, but differ in the number of key-value groups. Mixture-of-Experts (MoE) models use deeper networks with smaller expert hidden sizes, while the dense baseline has a wider feed-forward network and fewer layers. Across MoE variants, scaling behaviors differ: Nemotron models increase hidden size and expert size while keeping layer count constant, whereas Qwen3-30B-A3B adopts a narrower hidden dimension and smaller experts. All models use 32 attention heads, head dimension 128, and Grouped-Query Attention; the number of key-value groups ranges from 2 to 8. MoE architectures have 48 to 52 layers and expert hidden sizes between 768 and 2688, while the dense 8B model has 32 layers and a FFN hidden size of 21504. The Nemotron-3-Nano-30B-A3B and Nemotron-3-72B-A8B both use 52 layers, scaling hidden size from 2688 to 4096 and expert size from 1856 to 2688.

The Qwen-3-30B-A3B architecture, chosen for the Muon versus SOAP optimizer comparison, uses 128 total experts and activates 8 per token, with no shared expert. In contrast, both Nemotron models incorporate a shared expert and activate only 6 routed experts. Qwen-3-30B-A3B has 128 total experts, 8 active experts, and no shared expert. Nemotron-3-Nano-30B-A3B has 128 experts, 6 active, and a shared expert of size 3712. Nemotron-3-72B-A8B has 512 experts, 6 active, and a shared expert of size 5376. The optimizer comparison study between Muon and SOAP was conducted on the Qwen-3-30B-A3B architecture.

All models share a common WSD schedule, global batch size of 3072, and weight decay of 0.1, differing only in their learning rate extremes. The Qwen3-30B-A3B uses the highest values (max 1.6e-3, min 1.6e-5), the 8B Dense GPT and Nemotron-3-72B-A8B share the lowest (max 8e-4, min 8e-6), and Nemotron-3-Nano-30B-A3B sits in between (max 1e-3, min 1e-5). The Qwen3-30B-A3B model trains with a maximum learning rate twice as large as that of the 8B Dense GPT and Nemotron-3-72B-A8B architectures. Across all models, the minimum learning rate is consistently set to 1% of the maximum learning rate.

Muon is configured with a momentum coefficient of 0.9, decoupled weight decay of 0.1, 16 Newton-Schulz iterations, and an epsilon of 10^{-7} that floors the Frobenius norm used for normalization. Under controlled structural conditions, KL-SOAP maintains a slight but consistent edge in cross-entropy loss over Muon. The epsilon hyperparameter plays distinct roles: in Muon it bounds the normalization scale, while in SOAP it stabilizes the denominator of the second-moment estimate in the preconditioned basis. Muon's epsilon (10^{-7}) lower-bounds the Frobenius norm for normalization, whereas SOAP's epsilon floors the second-moment estimate after rotation into the eigenbasis. In a controlled comparison with identical tensor parallelism and fused QKV projection, KL-SOAP exhibits a slight but consistent advantage over Muon in cross-entropy loss. Muon uses 16 Newton-Schulz iterations with coefficients from the PolarExpress method for orthogonalization, while the MOP variant uses exact SVD-based polar decomposition. The momentum EMA coefficient for Muon is set to 0.9, and the decoupled weight decay strength is 0.1.

Across model sizes, Muon allows larger batch sizes while preserving or improving downstream task performance. The 8B model shows Muon at a doubled batch size outperforming AdamW at the baseline batch size on all reported metrics, with especially large gains in coding. Nano-V3 results indicate that Muon's coding accuracy peaks at 2× batch size, though reasoning tasks are more sensitive to batch size. On the 8B Hybrid MTP model, Muon at 2× batch size scores 65.79 on HumanEval versus AdamW's 57.62 at 1× batch size, a clear advantage in code generation. Muon’s downstream metrics remain robust as batch size scales up; for Nano-V3, Muon at 2× batch size achieves the highest HumanEval and HumanEval+ scores, while MMLU PRO CoT is best with the 1× batch size.

The experiments compare MoE and dense architectures with different expert counts, shared expert usage, and learning rate schedules, including a dedicated Muon versus SOAP optimizer study on the Qwen3-30B-A3B model. While KL-SOAP achieves a slight cross-entropy advantage over Muon, Muon permits larger batch sizes and delivers robust downstream performance gains, especially in code generation, with no degradation as batch size scales. The findings highlight Muon's scalability and effectiveness across model sizes and tasks.