NorMuon: Making Muon more efficient and scalable — part 2/3

Cleaned-text code points [15963, 31551). Bibliography omitted here; source.md includes it.

We observe that after the row-normalization the resulting direction has a much larger norm. Hence, during the update, we add a learning rate scaling following ( Jordan et al., 2024b ) to keep a similar RMS norm to match Adam’s RMS norm (line 10).

We would like to note that in the idealized case where the 𝐎 t \mathbf{O}_{t} is strictly orthogonalized (i.e., not approximated by NS5), the per-neuron norm would be strictly 1 for full-rank matrix with m ≤ n m\leq n . On these matrices, the neuron-wise normalization would not be beneficial. However, since orthogonalization is approximated in practice, we observe that this normalization remains necessary and helpful even for m ≤ n m\leq n matrix (validated in Section 4.1.3 ).

### 3.2 Distributed NorMuon

As LLM training scales larger, distributed training becomes essential for both memory constraints and computational efficiency. We develop a distributed version of NorMuon compatible with the FSDP2 framework ( Feng et al., 2022 ) , which employs ZeRO-3 style ( Rajbhandari et al., 2020 ) sharding to partition optimizer states, parameters, and gradients across multiple devices.

While coordinate-wise optimizers like Adam naturally extend to distributed settings, NorMuon presents unique challenges due to Muon’s orthogonalization step, which requires access to complete momentum matrices. An existing distributed implementation of Muon ( Liu et al., 2025a ) gathers the full momentum on all devices and replicates the orthogonalization computation. We avoid such replicated costs by near-uniformly assign parameters to different devices.

Algorithm 2 presents our distributed implementation. The key modifications from Algorithm 1 are:

∙ \bullet Efficient Orthogonalization Distribution (line 5-9): Rather than having all devices compute orthogonalization for all parameters, we first sort the parameter list by matrix size (line 2) to ensure uniform work assignment, where Numel ​ ( ⋅ ) \text{Numel}(\cdot) counts the number of elements in each matrix. We then assign each parameter tensor to a specific device using a round-robin scheme. Only the assigned device gathers the full momentum matrix via all-gather communication and performs the Newton-Schulz orthogonalization (lines 5-8), before scattering the result back to all devices (line 9). This approach eliminates redundant computation while maintaining load balance across devices.

∙ \bullet Shard-Local Row Normalization (lines 10-12): A key advantage of our design is that the row-wise normalization operates entirely on local shards without additional communication. This is possible because FSDP2 employs row-wise sharding, ensuring that each device holds complete rows of the weight matrix. The computation of row statistics and normalization thus proceed independently.

### 3.3 Overhead Analysis

Memory Overhead . NorMuon maintains Muon’s memory efficiency. For a weight matrix W ∈ R m × n W\in\mathbb{R}^{m\times n} , the memory consumption of optimizer states for each optimizer is: (1) Adam: 2 ​ m ​ n 2mn (first and second-order momentum); (2) Muon: m ​ n mn (first-order momentum only); (3) NorMuon: m ⁡ ( n + 1 ) m(n+1) (first-order momentum + per-neuron second-order statistics). The additional memory overhead of NorMuon compared to Muon is negligible (1/n factor), while remaining approximately 50% more memory-efficient than Adam.

Communication Overhead . NorMuon introduces moderate additional communication compared to standard FSDP training. Under FP32 training with AdamW, the per-parameter communication cost is: 4 bytes (forward all-gather) + 4 bytes (backward all-gather) + 4 bytes (gradient reduce-scatter) = 12 bytes. With NS5 iteration computed in BF16 precision, NorMuon requires: 12 bytes (standard FSDP communication) + 2 bytes (momentum gather, BF16) + 2 bytes (update scatter, BF16) = 16 bytes.

This represents a 33% increase in communication volume. When parameters use BF16, the relative overhead increases to 50%. However, this communication can be overlapped with orthogonalization computation to minimize latency impact. In our experiments (Section 4.1.4 ), we demonstrate that the per-iteration latency of NorMuon is only 3% higher than AdamW, while achieving significantly better convergence efficiency.

## 4 Experiments

In this section, we conduct pretraining experiments across four different model scales to validate the effectiveness of NorMuon: 124M, 350M, 1.1B, and 5.4B parameters. For the larger models (1.1B and 5.4B), we adopt the experimental setup from previous work on architecture scaling ( Ren et al., 2025 ) , with results and configurations presented in Section 4.1 . For the smaller models (124M and 350M), we follow the experimental setting of Modded-NanoGPT ( Jordan et al., 2024a ) , with results and settings provided in Section 4.2 . We include extensive ablation studies that justify our design choices, along with detailed efficiency analyses (Section 4.1.3 and 4.1.4 ).

### 4.1 Experiments on 1.1B and 5.4B Models

#### 4.1.1 Setup

Models. We follow a simple linear rule from prior works ( Kaplan et al., 2020 ; Ren et al., 2025 ) for scaling the architectural shape. Specifically, for a model with depth d d layers, we configure the architecture as follows: hidden dimension α ​ d \alpha d , number of attention query heads d d , number of attention key-value heads d / 4 d/4 , and MLP intermediate dimension 4 ​ α ​ d 4\alpha d , where α = 128 \alpha=128 . The α \alpha and ratios are derived relative to Llama-3-8B ( Dubey et al., 2024 ) . Our 1.1B and 5.4B parameter models correspond to depths of d = 16 d=16 and d = 28 d=28 layers, respectively.

Dataset. We conduct pretraining on the SlimPajama dataset ( Soboleva et al., 2023 ) and train our models on 50B tokens.

Hyperparameters. We employ Depth- μ \mu p ( Yang et al., 2023 ) to scale the learning rate inversely proportional to d \sqrt{d} based on model depth. The base learning rate is set to 4 × 10 − 4 4\times 10^{-4} with a base model depth of 16. The learning rate schedule consists of a 1B token warmup phase followed by linear decay to 0. We apply 0.1 weight decay for 2D parameters in hidden layers and zero weight decay for others to enhance training stability ( Ren et al., 2025 ) . The batch size is fixed at 2M tokens with a sequence length of 4096 tokens. For optimization, we use the following configurations: Adam optimizer with ( β 1 , β 2 ) = ( 0.9 , 0.95 ) (\beta_{1},\beta_{2})=(0.9,0.95) , Muon optimizer with β 1 = 0.95 \beta_{1}=0.95 following Jordan et al. (2024b) , and our proposed NorMuon optimizer with ( β 1 , β 2 ) = ( 0.95 , 0.95 ) (\beta_{1},\beta_{2})=(0.95,0.95) .

Baselines. We compare NorMuon against three established optimizers: AdamW ( Loshchilov and Hutter, 2017 ) , the standard adaptive optimizer with decoupled weight decay; Muon ( Jordan et al., 2024b ) , which applies orthogonalization to update directions; and Dion ( Ahn et al., 2025 ) , a scalable orthogonalization-based method that uses low-rank power iteration. For orthogonalization-based optimizers, we apply orthogonalization to the concatenated QKV matrix rather than separately for a fair comparison, as we observe no performance improvement with separate application for Muon.

#### 4.1.2 Main results

Figure 2 presents the validation loss curve across different model scales. NorMuon demonstrates consistent and substantial improvements over all baseline optimizers. While orthogonalization-based optimizers (Muon and Dion) already outperform AdamW, NorMuon amplifies this advantage through the integration of our proposed neuron-wise adaptive learning rate.

To quantify the performance gains, Table 1 reports the percentage reduction in training steps required for each optimizer to achieve the same final validation loss as Adam. NorMuon achieves the best efficiency gains of 21.74% and 13.91% for the 1.1B and 5.4B models, respectively.

#### 4.1.3 Ablation studies.

To validate our design choices in NorMuon, we conduct ablation studies that examine three key aspects: the granularity of adaptive learning rates and the positioning of normalization relative to orthogonalization, and the impact of applying normalization universally versus selectively based on matrix dimensions. The results on pretraining 1.1B model are presented in Figure 3 .

Adaptive Learning Rate Granularity . We compare our neuron-wise adaptive approach against coordinate-wise adaptation in the “Muon+Adam” baseline, which applies Adam’s coordinate-wise normalization after Muon’s orthogonalization. This approach is similarly explored in concurrent work ( Si et al., 2025 ) , though we removed the sign stabilization component as we found it performs better without it. While Muon+Adam demonstrates improvements over vanilla Muon, it underperforms NorMuon across the training trajectory. Importantly, this approach requires maintaining full second-order momentum of orthogonalized updates, doubling the memory overhead (Section 4.1.4 ).

Normalization Positioning . We further investigate whether the positioning of adaptive normalization matters by testing “NorMuon (Front)”, which applies neuron-wise adaptive learning rates before orthogonalization rather than after. This variant still improves upon Muon, but slightly underperforms NorMuon.

Universal vs. Selective Normalization : To test whether normalization is beneficial to m ≤ n m\leq n matrices as discussed in Section 3.1 , we evaluate ”NorMuon ( m > n m>n only)”, which applies normalization only to m > n m>n matrices. We can see that this selective approach underperforms the full NorMuon, demonstrating that applying normalization to those with m ≤ n m\leq n is helpful.

#### 4.1.4 Computational and Memory Overhead Analysis

In this section, we demonstrate that the overheads incur by NorMuon are manageable and do not diminish the practical benefits.

Wall-Clock Performance . Figure 4 presents the validation loss as a function of wall-clock training time. Despite the additional computation required for orthogonalization and neuron-wise normalization, NorMuon maintains substantial performance advantages over AdamW.

Memory and Computational Overhead . Table 2 provides a breakdown of the computational and memory requirements for each optimizer. NorMuon achieves comparable memory efficiency to Muon with nearly a 50% reduction compared to AdamW or Muon + Adam. In terms of computational cost, NorMuon introduces only a 2.9% increase in training step time compared to AdamW. The neuron-wise norm computation adds minimal cost relative to orthogonalization operations. Our efficient orthogonalization distribution strategy across GPUs proves crucial for maintaining low overhead. Without proper work distribution, the optimizer step time increases to approximately 2.7×. Notably, this strategy can and has been applied to Muon for fair comparison.

### 4.2 Experiments on Modded-NanoGPT

To further verify the advantages of NorMuon over Muon, we conduct experiments using Muon’s original experimental setting on Modded-NanoGPT ( Jordan et al., 2024a ) . Detailed experimental configurations and more ablation results are provided in Appendix C and A.2 , respectively.

Main Results . Figure 5 presents the comparison between NorMuon and Muon on 124M and 350M parameter models. NorMuon consistently outperforms Muon across both model sizes. Since Muon’s improvements over Adam have been extensively demonstrated in Jordan et al. (2024b) , we omit those baseline comparisons here.

To quantify the computational benefits of NorMuon, we conduct an additional analysis where Muon is trained with the same learning rate schedule but for a longer total number of iterations, until it reaches the same validation loss as NorMuon (denoted as “Muon-rescale” in Figure 5 ). We observe that on the 124M model, Muon requires 6% more iterations than NorMuon to achieve the same validation loss. On the 350M model, this efficiency gap increases substantially to 15%, demonstrating the broad advantage of NorMuon over Muon.

## 5 Conclusion

In this work, we introduced NorMuon , a simple yet effective optimizer that integrates Muon’s orthogonalization with neuron-wise adaptive learning rates. To make NorMuon practical for large-scale training, we developed an efficient distributed implementation under the FSDP2 framework, carefully orchestrating momentum gathering and orthogonalization to eliminate redundant computation and communication overhead. Our experiment results shows notable improvement over Muon, demonstrating that orthogonalization and adaptive scaling need not be mutually exclusive, but can be complementary, when combined, lead to superior optimization dynamics.

## Acknowledgments

We thank Liliang Ren for valuable discussions and feedback regarding the code and its implementation.

## Appendix A Additional Experiments

### A.1 Adam-mini’s results on optimization geometry analysis

### A.2 Ablation Experiments on Modded-NanoGPT

To further verify the effectiveness of NorMuon, we conducted several ablation experiments under the setting of training a 350M Modded-NanoGPT on FineWeb, and show the results in Figure 7 :

(1) Standard NorMuon, denoted as ”NorMuon” in Figure 7 .

(2) Standard Muon used in original setting of Modded-NanoGPT ( Jordan et al., 2024b ) , denoted as ”Muon” in Figure 7 .

(3) Applying normalization directly to Muon’s update such that the update is strictly m × n \sqrt{m\times n} , denoted as ”Muon w/ normalization” in Figure 7 .

(4) applying NorMuon only to weight matrices with m > n m>n , while using the normalized muon mentioned in (3) for all other weight matrices, denoted as ”NorMuon ablation” in Figure 7 .

We can see that although Muon with normalization performs slightly better than Muon in the early stages, it is eventually surpassed by Muon, indicating that the effectiveness of NorMuon cannot be attributed to normalization. Furthermore, since weight matrices with m > n m>n correspond only to the MLP up-projection matrices, which constitute only a small portion of the model, applying NorMuon only to this subset of parameters greatly diminishes the effect of NorMuon, resulting in only a marginal improvement over Muon.

## Appendix B Implementation Details

For experiments involving 1.1B and 5.4B parameter models, we conducted training on 2 nodes, each equipped with 8 A100 GPUs (80GB) connected via NVLink for optimized inter-GPU communication. Training duration was approximately 2 days for the 1.1B model and 7 days for the 5.4B model.

All experiments using Modded-NanoGPT, including both 124M and 350M parameter models, were performed on a single node with 8 A100 GPUs.

## Appendix C Experiment Setup of Modded-NanoGPT

We strictly follow the experimental setup of Muon ( Jordan et al., 2024b ) , with details below:

Models : The model architecture is consistent with GPT-2 ( Radford et al., 2019 ) , with 124M and 350M parameter configurations obtained by adjusting width and depth.

Dataset : We train all models on the FineWeb dataset ( Penedo et al., 2024 ) . The 124M model is trained on approximately 3.2B tokens, while the 350M model uses approximately 4B tokens.

Hyperparameters : Since Muon has already performed extensive hyperparameter tuning in this setting ( Jordan et al., 2024b ) , we adopt their optimized configurations except for β 1 \beta_{1} , which we slightly tune. We use a batch size of 512, sequence length of 1024, and the Warmup-Stable-Decay (WSD) learning rate schedule. Training iterations are set to 6,200 for the 124M model and 7,500 for the 350M model.