Paper deep dive
Dion3: Full-Stack Orthogonal Updates
Noah Amsel, Jack Zhang, Kwangjun Ahn, Ali Naeimi, Austin Feng, Berlin Chen, Tri Dao, John Langford
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:The Muon optimizer incurs a significant overhead cost due to its cubic-time Newton-Schulz orthogonalization step. When weights are sharded, communication overhead compounds this computational cost, eroding the benefits of Muon in many settings. We present Dion3, a revision of Muon that targets this overhead at every level of the stack. Our Gram Newton-Schulz algorithm reduces the FLOP cost of orthogonalization, our CuteDSL kernels accelerate it by exploiting symmetry, and our megabatching strategy reduces communication overhead. Moreover, we propose a simple change to the update rule that cuts costs even further: selecting only a fraction of the momentum matrix's rows to orthogonalize at each step. This update rule improves on Dion (another "compressed" version of Muon), in both speed and performance. Overall, Dion3 matches or improves on the loss achieved by Muon but reduces optimizer step time by up to 6x. Dion3 is available via the dion package (this https URL) as a drop-in replacement for Muon.
Tags
Links
- Source: https://arxiv.org/abs/2608.11612v1
- Canonical: https://arxiv.org/abs/2608.11612v1
Trouble viewing inline? Open PDF directly →
Full Text
134,174 characters extracted from source content.
Expand or collapse full text
Dion3: Full-stack orthogonal updates Noah Amsel Work performed while at Microsoft Research. New York University Jack Zhang Princeton University Kwangjun Ahn11footnotemark: 1 NVIDIA Ali Naeimi Independent Researcher Austin Feng11footnotemark: 1 Yale University Berlin Chen Princeton University Tri Dao Princeton University John Langford Microsoft Research Abstract The Muon optimizer incurs a significant overhead cost due to its cubic-time Newton-Schulz orthogonalization step. When weights are sharded, communication overhead compounds this computational cost, eroding the benefits of Muon in many settings. We present Dion3, a revision of Muon that targets this overhead at every level of the stack. Our Gram Newton-Schulz algorithm reduces the FLOP cost of orthogonalization, our CuteDSL kernels accelerate it by exploiting symmetry, and our megabatching strategy reduces communication overhead. Moreover, we propose a simple change to the update rule that cuts costs even further: selecting only a fraction of the momentum matrix’s rows to orthogonalize at each step. This update rule improves on Dion (another “compressed” version of Muon), in both speed and performance [2]. Overall, Dion3 matches or improves on the loss achieved by Muon but reduces optimizer step time by up to 6×6×. Dion3 is available via the dion package111https://github.com/microsoft/dion as a drop-in replacement for Muon. Figure 1: Optimizer step time of Muon (excluding forward/backward pass) relative to AdamW for a 7B-parameter language model trained on four GH200s. Each colored bar adds one of our contributions on top of the previous one. Together they reduce Muon’s cost from 26×26× AdamW to just 4×4×. Contents 1 Introduction 2 The Challenge of Scaling Muon 2.1 Muon and NorMuon Recap 2.2 Standard Newton-Schulz 2.3 Scalability of Muon 3 Comparison with Related Work 3.1 Improving Newton-Schulz 3.2 Orthogonalizing a Smaller Matrix 4 Gram Newton-Schulz 4.1 Runtime of Naive Gram Newton-Schulz 4.2 Stabilizing Gram Newton-Schulz 5 Symmetric GEMM Kernels in CuteDSL 6 The Dion3 Update Rule 7 Megabatching and Communication 8 Experiments 8.1 Model Quality Is Preserved 8.2 Dion3 Accelerates the Optimizer 9 Conclusion References A Alternative Experiments on Gram Newton-Schulz A.1 Setup A.2 Kernelized Gram Newton-Schulz preserves quality A.3 Kernelized Gram Newton-Schulz speeds up the optimizer step B Stability of Gram Newton-Schulz B.1 Instability of Naive Gram Newton-Schulz B.2 Stabilizing Gram Newton-Schulz by Restarting C Kernel Implementation Details C.1 Symmetric GEMM Kernel Details C.2 Implementation Strategy in Code C.3 Kernel Optimizations for Standard Newton-Schulz D Additional Experiments D.1 Architecture and Optimization Setup D.2 Finer-Grained Timing Metrics D.3 Benchmarking all-to-all communications D.4 Ablations E Case Studies of End-to-End Training Time 1 Introduction Muon is becoming the method of choice for training frontier LLMs like Kimi K2 and GLM-5 [21, 13]. Compared to AdamW, Muon needs fewer optimizer steps to reach a given loss, but each step is more expensive. This overhead is due to Muon’s Newton-Schulz orthogonalization procedure, a cubic-time matrix operation not present in older optimizers. As model size increases, the overhead of computing each Muon step grows rapidly. Moreover, Muon is more difficult than traditional optimizers to parallelize, introducing additional communication overhead in distributed settings. Thus, scaling Muon presents a range of algorithmic and systems-level challenges that diminish its effectiveness in the most demanding settings and hinder its adoption to new problems. This paper gives a comprehensive answer to these challenges. We provide a full-stack solution that performs well across a wide range of model sizes, architectures, cluster sizes, and parallelism strategies. We package our contributions into a single optimizer, Dion3, built on four improvements that each reduce the cost of Muon’s orthogonalization step and compound when used together: 1. Gram Newton-Schulz, a mathematically equivalent reformulation of Newton-Schulz that iterates on the small symmetric Gram matrix and cuts the FLOP cost of each orthogonalization dramatically (Section˜4). 2. Custom GPU kernels for symmetric matrix multiplication written in CuteDSL, which speed up Gram Newton-Schulz and take maximal advantage of its symmetric structure (Section˜5). 3. A new optimizer update rule that subsamples the rows or columns of the momentum matrix before orthogonalizing it to make each step faster (Section˜6). Our update rule is simpler and faster than Dion [2] while matching or improving the optimization quality of Muon and NorMuon. 4. Megabatched communication, which eliminates a significant source of overhead by reducing the number of rounds of communication per optimizer step to a small constant. Our improvements combine to accelerate the optimizer step at every level of the stack. The GPU kernels at the lowest level speed up Gram Newton-Schulz, which in turn reduces the FLOP cost of executing our optimizer’s update rule. At the top level, the compression achieved by our update rule joins with the megabatch strategy to limit the communication cost. Each of our contributions is of independent value, but combining them gives Dion3 maximum flexibility to handle different parallelism strategies, architectures, and model sizes efficiently. Moreover, Dion3 benefits from synergies between them that further reduce computational cost. Gram Newton-Schulz uses more symmetric multiplications than standard Newton-Schulz does, compounding the benefits of our CuteDSL kernels. Likewise, Gram Newton-Schulz is especially fast for matrices with highly asymmetric shapes—like those produced by our update rule’s subsampling strategy. Surprisingly, Dion3 appears to benefit the update quality as well as the computational footprint in some settings, even with a strong optimized baseline. To realize these gains in practice, we provide open-source implementations in the form of two interoperable packages: gram-newton-schulz222https://github.com/Dao-AILab/gram-newton-schulz is a drop-in replacement for Muon’s Newton-Schulz routine that also incorporates our CuteDSL kernels, while the dion package implements our optimizer update rule (along with baselines like Muon and NorMuon), and handles parallelism in the distributed setting (megabatching, FSDP, DDP, etc.). In all, Dion3 allows practitioners working in a wide range of settings to realize the benefits of Muon and related optimizers with almost no overhead. 2 The Challenge of Scaling Muon 2.1 Muon and NorMuon Recap The Muon optimizer [19] is best described as steepest-direction descent with respect to the spectral norm [5]. At a given training step, let ∈ℝn×m W ^n× m be a weight matrix and let G be the gradient of the loss with respect to W. The Muon update rule is ←μ+←−ηpolar() split M&←μ M+ G\\ W&← W- ( M) split (1) where μ is the momentum coefficient, η is the learning rate, and M is the momentum matrix (with 0:= M_0:= 0). In many ways, Muon resembles basic stochastic gradient descent (SGD) with momentum. Its key innovation is the polarpolar operation, which is defined as follows: Definition 1 (Polar Decomposition). If =⊤ X= U V is the singular value decomposition (SVD) of a matrix, then polar()=⊤polar( X)= U V . This operation—the “orthogonalization step”—balances the spectrum of the update, ensuring that it has full numerical rank. NorMuon [23] is a popular variant of Muon that aims to combine its strengths with those of Adam. While Muon’s update has a balanced spectrum, it does not have balanced row norms; each step could greatly affect some neurons while almost neglecting others. Following Adam, NorMuon uses the second moment of the (orthogonalized) gradients to adaptively adjust the stepsize for each neuron, helping all neurons to learn at every step. A final rescaling ensures that the overall magnitude of the update (as measured in the Frobenius norm) matches that of Muon: ←μ+,←polar()i←β2i+(1−β2)⋅1m∑jij2,^ij←iji+ϵ←−η‖^‖ split M&←μ M+ G, O ( M)\\ v_i&← _2 v_i+(1- _2)· 1mΣ _j O_ij^2, O_ij← O_ij v_i+ε\\ W&← W-η \| O\|_ F\| O\|_ F O split (2) The cost of NorMuon’s extra normalization steps is negligible compared to that of computing polar()polar( M), so NorMuon’s benefits come almost for free. Our methods, which speed up the computation of polar()polar( M), apply to both Muon and NorMuon alike. 2.2 Standard Newton-Schulz The main challenge to scaling Muon is the need to compute polar()polar( M) for each weight matrix at every step. Since polar(⋅)polar(·) is expensive to compute exactly, Muon uses the Newton-Schulz method to approximate it. Newton-Schulz is an iterative method based on matrix polynomials. Beginning with 0 X_0, each iteration improves the approximation t≈polar(0) X_t ( X_0) according to the update rule t+1=att+bttt⊤t+ct(tt⊤)2t. X_t+1=a_t X_t+b_t X_t X_t X_t+c_t ( X_t X_t )^2 X_t. We can interpret Newton-Schulz by understanding how it affects the singular value decomposition. Let 0=⊤ X_0= U V be the SVD, where ⊤=⊤= U U= V V= I and is diagonal with positive entries called the singular values. A direct computation shows 1=(a1+b13+c15)⊤=p1()⊤ X_1= U (a_1 +b_1 ^3+c_1 ^5 ) V = Up_1( ) V where p1(x):=a1x+b1x3+c1x5p_1(x):=a_1x+b_1x^3+c_1x^5. Since U and V have orthonormal columns and p1()p_1( ) is diagonal, the right-hand side of this equation must be the SVD of 1 X_1. By extension, T X_T also has the same singular vectors U and V as 0 X_0, and its singular values have been transformed according to the composition of polynomials (pT∘⋯∘p1)()(p_T ·s p_1)( ). If we normalize the input matrix 0=/‖ X_0= X/\| X\|_ F, then all singular values of 0 X_0 lie in [0,1][0,1]. [19] identified a sequence of degree-5 odd polynomials for which (pT∘⋯∘p1)(x)≈1(p_T ·s p_1)(x)≈ 1 on [0,1][0,1]. Therefore, T=(pT∘⋯∘p1)()⊤≈⊤=:polar(0) X_T= U(p_T ·s p_1)( ) V ≈ U V =:polar( X_0) Algorithm 1 Standard Newton-Schulz 1:∈ℝn×m X ^n× m with n≤mn≤ m, coefficients (at,bt,ct)t=15\(a_t,b_t,c_t)\_t=1^5 2:←/(‖+ϵ) X← X\,/\,(\| X\|_ F+ε) ⊳ Normalize sing. vals. to [0,1][0,1]; ϵ=10−7ε=10^-7 3:←bfloat16() X← bfloat16( X) ⊳ Cast to half precision for speed 4:for t=1,…,5t=1,…,5 do ⊳ Apply pt()p_t( X) 5: ←⊤ A← X X 6: ←bt+ct2 B← b_t A+c_t A^2 7: ←at+ X← a_t X+ B X 8:end for 9:return X Algorithm˜1 gives the standard implementation of Newton-Schulz. We now analyze its runtime in FLOPs to help us understand its performance bottlenecks. We count only the cubic-time matrix multiplication operations, ignoring the lower-order scalar multiplications and matrix additions. For clarity, we let T denote the number of iterations, remembering that Muon typically sets T=5T=5. We also assume without loss of generality that n≤mn≤ m and define the aspect ratio α=m/n≥1α=m/n≥ 1. Intuitively, α measures how asymmetric the shape of the matrix is, with α=1α=1 being square and α≫1α 1 being very asymmetric. Each iteration has three steps. Each step contains a single matrix multiplication (⊤ X X , 2 A^2, B X) costing, respectively, 2mn22mn^2, 2n32n^3, and 2mn22mn^2 FLOPs for a total cost of T(4mn2+2n3)=2T(2α+1)n3T(4mn^2+2n^3)=2T(2α+1)n^3 FLOPs. When T=5T=5, the cost is (20α+10)n3(20α+10)n^3 spread across 15 matrix-matrix multiplications (GEMMs). 2.3 Scalability of Muon While Muon comfortably outperforms traditional optimizers like Adam on a per-step basis, two factors immediately make it more difficult to scale: • Super-linear complexity. SGD and Adam perform only inexpensive element-wise operations like addition and scalar multiplication; therefore, the computational complexity of these methods scales linearly with the number of parameters. Orthogonalization is substantially more expensive; for an n×n× n weight matrix, it requires O(n3)O(n^3) time. Although sparse MoE architectures keep most weight matrices smaller, they also reduce overall model FLOPs, increasing the relative cost of the optimizer step [11] compared to the forward and backward passes. Moreover, MoEs still include large weight matrices in their dense layers [9]. • Distributed training. There are additional challenges when weights are sharded across GPUs, as is standard in distributed training. Element-wise optimizers like Adam can update each shard separately, but Muon must gather the shards together to compute polarpolar on the full matrix. Early implementations of Muon duplicated the orthogonalization step across devices, but this increases its already considerable cost. [11] propose using all-to-all communication along the sharding dimension, with each device processing a different weight in parallel. [2] implemented this strategy in PyTorch FSDP2, and [24] examined its compute-communication overlap characteristics. While this approach makes the overhead of distributed Newton-Schulz manageable in some settings, it does not fully resolve the scalability limitations. In Section˜6, we introduce a variant of Muon that addresses these two challenges by subselecting (that is, compressing) the momentum matrix before orthogonalizing it. Section˜7 describes a megabatching strategy and a flexible, pip-installable package that gracefully handle the distributed setting. Our other contributions are inspired by the analysis in the previous section, which reveals two shortcomings of the standard Newton-Schulz algorithm: • Symmetric Matrix Multiplication. The matrices =⊤ A= X X and =bt+ct2 B=b_t A+c_t A^2 computed at each iteration of Newton-Schulz are symmetric by definition, but standard Newton-Schulz does not exploit this structure. We can compute the lower triangular part of these matrices in the usual way and simply copy the results to the upper triangular part. This technique halves the cost of computing ⊤ X X and 2 A^2, giving an overall total of T(3α+1)n3T(3α+1)n^3 FLOPs. Section˜5 describes custom CuteDSL kernels that implement this technique. • Dependence on Aspect Ratio. Newton-Schulz’s runtime is dominated by the large rectangular matrix multiplications needed to compute ⊤ X X and B X, which together cost 3αn33α n^3 FLOPs per iteration even when using symmetric matrix multiplications. A typical implementation with T=5T=5 requires 10 of these expensive rectangular multiplications. This strong dependence on α is unfortunate. Most of the weight matrices in transformer architectures are rectangular, including the MLP weights, MoE weights, and attention projection weights when using GQA or MLA. Furthermore, we observe that the latest MoE architectures are trending towards finer-grained, sparser experts, meaning that the aspect ratios of their hidden dimensions to intermediate dimensions are increasing as well [21, 16, 34, 31]. Thus, at large scales, pretraining time would benefit greatly from an algorithm that uses fewer rectangular multiplications and more small symmetric ones. We develop such an algorithm in Section˜4. 2.3.1 Kimi’s Success Muon has been scaled successfully by Moonshot AI [26, 21]. Their success was enabled by an alignment of several factors: 1. Earlier releases of PyTorch and Megatron-LM used DP-sharding strategies for optimizer states that were, by chance, favorable for Muon [27]. Model and optimizer states were stored in large contiguous flat buffers, and data-parallel shards were produced by splitting this buffer. As a result, only tensors crossing a DP boundary required an additional gather. These advantageous strategies, however, have been deprecated in more recent releases. 2. They adopt a fine-grained MoE architecture with only a single dense layer (even fewer than in DeepSeek-V3 [9]), so most matrices remain small even at the one-trillion-parameter scale, keeping the Newton-Schulz overhead manageable. 3. Their main distributed training strategy combines pipeline parallelism with expert parallelism, which naturally distributes Muon’s computation across devices with minimal communication. In sum, Muon has been successfully scaled, but its success relies on a subtle alignment of architectural, parallelism, and framework factors. For Muon to serve as a general-purpose replacement for Adam, it would benefit from cheaper, more flexible scaling properties that relax these constraints. This is the goal of the present work. 3 Comparison with Related Work As Muon has gained popularity, successive work has sought to improve it in several ways. Most of these proposals (e.g. NorMuon) modify Muon’s update rule so as to reach a given loss in fewer training steps; however, they use the same Newton-Schulz routine described above and generally suffer from the same scaling challenges. 3.1 Improving Newton-Schulz A few papers have attempted to improve the Newton-Schulz step itself—e.g., by optimizing the sequence of polynomials (at,bt,ct)(a_t,b_t,c_t) or the normalization step [3, 15, 6]—but they retain the form of Algorithm˜1. In contrast, our Gram Newton-Schulz algorithm departs from this form. Since its output is still mathematically identical to the standard version, it remains compatible with nearly all varieties of Muon, including prior improvements to Newton-Schulz. Gram Newton-Schulz is closely akin to a method proposed in Appendix J of [3]. Both aim to reduce the FLOP cost of Newton-Schulz, both form the Gram matrix to reduce the number of m×nm× n matrix multiplications, and both are mathematically identical to standard Newton-Schulz. However, our work supersedes [3] in several ways. First, the precise formulas of Gram Newton-Schulz are different and, we believe, more stable. Second, we use symmetric matrix multiplication kernels; the opportunity to use these kernels more is an essential advantage of Gram Newton-Schulz not studied previously. Third, we undertake a thorough stability analysis and provide practical recommendations that allow Gram Newton-Schulz to be used in practice with minimal ad-hoc hyperparameter tuning. The iterative part of the method, which approximates the inverse square root of the Gram matrix T≈0−1/2:=(⊤)−1/2 Q_T≈ R_0^-1/2:=( X X )^-1/2, generalizes and speeds up the method of [22] (see also [17, Eq. 7.18 ]), but Lakić does not consider the symmetric kernels, the stability issues, or the application to polar()polar( X). The idea of exploiting symmetry to reduce the arithmetic cost of computing ⊤ A A has appeared before, both in standard linear algebra packages like BLAS and in the context of Muon’s Newton-Schulz routine [30, 33, 25]. However, our results show that the true potential of this trick is only realized when combined with Gram Newton-Schulz, which requires more general symmetric operations like +β A B+β C. Furthermore, our symmetric matrix multiplication kernels are written in CuteDSL, exploiting the advanced features of NVIDIA’s Hopper and Blackwell GPUs. 3.2 Orthogonalizing a Smaller Matrix Our update rule is part of a second line of work that reduces the cost of the Newton-Schulz step by shrinking the size of its input, an approach pioneered by Dion [2]. Dion constructs a low-rank approximation of the momentum matrix and orthogonalizes only this approximation. If ≈^⋅^⊤ M≈ M V· V is a rank-k approximation with ^⊤^=∈ℝk×k V V= I ^k× k, then polar()≈polar(^^⊤)=polar(^)^⊤.polar( M) ( M V V )=polar( M V) V . (3) Because the dimensions of M V are much smaller than those of M, this approximation greatly reduces the runtime of Newton-Schulz. In Dion, V is found by a warm-started power-iteration procedure. If V spans the top-k right singular vectors of M, then this approximation gives the optimal rank-k approximation of polar()polar( M). [2] found that approximating V from only a single warm-started power-iteration suffices for good downstream performance. A key element of Dion’s success is error feedback, a technique that helps offset the error introduced by low-rank approximation in future iterations. We adopt error feedback and describe it fully in Section˜6. Trion [29] adopts the main ideas of Dion, but changes how the low-rank approximation is computed. Instead of constructing V via power iteration, it selects k columns from the discrete cosine transform matrix. Surprisingly, this simpler approximation performs even better than Dion, suggesting that finding a good low-rank approximation is not essential; the error-feedback mechanism compensates even for large differences between M and ^^⊤ M V V . We push this simplification further by using the simplest possible “low-rank approximation” of M: we select k of its rows or columns and do not operate on the others. This achieves the same savings in the runtime of Newton-Schulz as Dion and Trion, but makes all the other steps of the algorithm (construction of V, error feedback, etc.) much easier to implement, especially in the distributed setting. Besides low-rank approximation, block orthogonalization is an alternative way to shrink the size of the input to Newton-Schulz. This approach partitions the momentum matrix into blocks—each typically corresponding to a shard in weight-sharded settings—and orthogonalizes each block separately. Several papers have proposed versions of this approach [7, 20, 32]. In particular, [20] found success by periodically alternating block-wise steps with full steps, a method they call MuonBP. Both low-rank and block orthogonalization methods use Newton-Schulz as a black box. Therefore, they are trivial to combine with improvements to Newton-Schulz. 4 Gram Newton-Schulz Our first contribution is an orthogonalization algorithm that uses fewer rectangular matrix multiplications than standard Newton-Schulz. Instead of iterating directly on the large input matrix X, we iterate on the small symmetric Gram matrix ⊤ X X . The output of our algorithm is mathematically identical to that of standard Newton-Schulz, but it is significantly cheaper to compute. At a high level, our strategy is based on the following formula. If ∈ℝn×m X ^n× m with n≤mn≤ m, then polar()=(⊤)−1/2polar( X)=( X X )^-1/2 X. Rather than use an iterative method to approximate T≈polar() X_T ( X) directly, we instead: 1. Compute the n×n× n Gram matrix ⊤ X X 2. Use an iterative method to approximate T≈(⊤)−1/2 Q_T≈( X X )^-1/2 3. Output T Q_T X Step 2—which comprises almost all of the algorithm’s wall clock runtime and FLOP cost—works entirely with small n×n× n symmetric matrices. This version uses just two rectangular matrix multiplications: ⊤ X X in the beginning and T Q_T X at the end. It also synergizes well with our symmetric GEMM kernels (Section˜5). Because we use more symmetric multiplications than before, these kernels provide an even greater speedup. Since our algorithm works on the n×n× n Gram matrix of X, we call it “Gram Newton-Schulz”. How can we ensure that the output matches that of standard Newton-Schulz? Recall that Newton-Schulz outputs (pT∘⋯∘p1)()(p_T ·s p_1)( X), where each ptp_t is an odd polynomial p(x)=ax+bx3+cx5p(x)=ax+bx^3+cx^5. Any odd polynomial can be rewritten in the form p(x)=xh(x2)p(x)=xh(x^2), where h is a lower-degree polynomial with the same coefficients, like h(x)=a+bx+cx2h(x)=a+bx+cx^2. Intuitively, if p(x)≈1p(x)≈ 1, then h(y)=p(y1/2)y−1/2≈y−1/2h(y)=p(y^1/2)y^-1/2≈ y^-1/2, so the Newton-Schulz polynomials implicitly provide a way to approximate inverse square roots. If we use this approximation in step 2 of Gram Newton-Schulz, then its output will match that of standard Newton-Schulz. Formally, Gram Newton-Schulz is based on the following theorem. In effect, it shows how to compute T X_T from 0 X_0 without ever constructing the intermediate values 1,…,T−1 X_1,…, X_T-1: Theorem 2. If pt(x)=xht(x2)p_t(x)=xh_t(x^2) for all t∈1,…,Tt∈\1,…,T\, then (pT∘⋯∘p1)(x)=qTx(p_T ·s p_1)(x)=q_Tx, where qTq_T is defined by the iteration r0=x2,q0=1,r_0=x^2,\,\,q_0=1, and zt=ht(rt−1),rt=rt−1zt2,qt=qt−1zt z_t=h_t(r_t-1), r_t=r_t-1z_t^2, q_t=q_t-1z_t for all t∈1,…,Tt∈\1,…,T\. Proof. Define x0=x_0=x and xt=pt(xt−1)x_t=p_t(x_t-1) for t∈1,…,Tt∈\1,…,T\. We will show by induction that rt=xt2r_t=x_t^2 and qt=xt/x0q_t=x_t/x_0 for all t. The base case t=0t=0 holds by the definition r0=x2,q0=1r_0=x^2,q_0=1. Now assume the hypothesis holds for t−1t-1. By assumption, xt=pt(xt−1)=xt−1ht(xt−12)x_t=p_t(x_t-1)=x_t-1h_t(x_t-1^2) By the inductive hypothesis, ht(xt−12)=ht(rt−1)=zth_t(x_t-1^2)=h_t(r_t-1)=z_t, so xt=xt−1ztx_t=x_t-1z_t. Squaring both sides, xt2=xt−12zt2=rt−1zt2=rtx_t^2=x_t-1^2z_t^2=r_t-1z_t^2=r_t If we instead divide both sides by x0x_0 and apply the other part of the inductive hypothesis, we get xtx0=xt−1x0zt=qt−1zt=qt x_tx_0= x_t-1x_0z_t=q_t-1z_t=q_t Thus, both parts of the hypothesis hold for t. Finally, conclude (pT∘⋯∘p1)(x)=xT=qTx0(p_T ·s p_1)(x)=x_T=q_Tx_0. ∎ Note that, as an immediate corollary of the proof, qt=xt/x0→1/x0=(x02)−1/2q_t=x_t/x_0→ 1/x_0= (x_0^2 )^-1/2. In effect, this shows that T→(⊤)−1/2 Q_T→( X X )^-1/2. To obtain our initial version of Gram Newton-Schulz, we simply lift the iteration from Theorem˜2 to matrices. As in standard Newton-Schulz, each matrix operation preserves singular vectors. Therefore, each singular value of t R_t, t Q_t, and t Z_t evolves independently of the others according to the scalar iteration described above. Note that while this algorithm is mathematically equivalent to standard Newton-Schulz, it is not yet practical due to numerical instability. The only difference between our proposed method (Algorithm˜3) and this naive version is the presence of what we call a “restart” at the beginning of iteration 3 of the loop. We will motivate this change below (Section˜4.2). Algorithm 2 Naive Gram Newton-Schulz 1:∈ℝn×m X ^n× m with n≤mn≤ m, coefficients (at,bt,ct)t=15\(a_t,b_t,c_t)\_t=1^5 2:←/(∥+ϵ) X← X\,/\,( X _ F+ε) ⊳ Normalize sing vals to [0,1][0,1], ϵ=10−7ε=10^-7 3:0=⊤ R_0= X X 4:0= Q_0= I 5:for t=1,…,5t=1,…,5 do 6: t←at+btt−1+ctt−12 Z_t← a_t I+b_t R_t-1+c_t R_t-1^2 ⊳ Apply ht(t−1)h_t( R_t-1) 7: t←t−1t Q_t← Q_t-1 Z_t 8: t←tt−1t R_t← Z_t R_t-1 Z_t 9:end for 10:return 5 Q_5 X 4.1 Runtime of Naive Gram Newton-Schulz We now calculate the FLOP count of this new algorithm to show how its runtime improves on standard Newton-Schulz. There are four matrix multiplications per iteration; if we use our symmetric GEMM kernel, these cost n3n^3 FLOPs each. The initialization (⊤ X X ) and output (5 Q_5 X) steps cost mn2mn^2 and 2mn22mn^2, respectively, since 5 Q_5 X is not symmetric. Computing 1=01 Q_1= Q_0 Z_1 is free since 0= Q_0= I, and we can skip computing 5=545 R_5= Z_5 R_4 Z_5; together this saves 3n33n^3 FLOPs. Thus, the total cost is T⋅4n3+3mn2−3n3=(4T+3α−3)n3T· 4n^3+3mn^2-3n^3=(4T+3α-3)n^3 FLOPs. Compare this to standard Newton-Schulz’s T(3α+1)n3T(3α+1)n^3 FLOPs when using symmetric GEMMs. When α=1α=1, they are equal. When α>1α>1, Gram Newton-Schulz is cheaper, often significantly so. For a typical Muon application (T=5,α=4T=5,α=4333Transformers’ MLP blocks typically have an intermediate dimension 4×4× the model’s hidden dimension.), it saves 55% of the FLOPs used by standard Newton-Schulz with symmetric GEMMs, or 68% compared to a typical implementation without symmetric GEMMs. For larger T and α, the savings are even greater, as the leading order term is O((T+α)n3)O((T+α)n^3) instead of O(Tαn3)O(Tα n^3). In practice, when α=1α=1, we fall back to standard Newton-Schulz with our symmetric GEMMs (Section˜C.3), since it launches fewer GEMMs and has a faster wall clock time. 4.2 Stabilizing Gram Newton-Schulz As written, Algorithm˜2 destabilizes Muon and can even diverge. The leading cause of this instability is the introduction of spurious negative eigenvalues in the Gram matrix ⊤ X X due to large rounding errors in half-precision arithmetic. These negative eigenvalues theoretically should not exist in a Gram matrix like ⊤ X X . The main loop of Gram Newton-Schulz approximates inverse square root of positive numbers, but it diverges for negative inputs. We provide a thorough theoretical analysis and numerical experiments showing the impact of spurious negative eigenvalues in the Gram matrix in Section˜B.1. In practice, we can fully mitigate this instability using a restarting strategy. Instead of running all five iterations of the loop in a single pass and outputting 5 Q_5 X, we run only the first two iterations and compute 2=2 X_2= Q_2 X. We then “restart” the algorithm treating 2 X_2 as the new input; we construct the Gram matrix 22⊤ X_2 X_2 , initialize 2= Q_2= I, and proceed with the remaining three iterations of the loop. This resets any spurious negative eigenvalues to near zero and restores commutativity of X, Q, and R at the cost of 3(α−1)n33(α-1)n^3 FLOPs. We find this strategy is sufficient to preserve training quality. To select the best iteration after which to restart, we sweep for the restart location that provides the best bound on the condition number of t Q_t, assuming that forming the Gram matrix introduces spurious eigenvalues as negative as −4⋅10−4-4· 10^-4. We walk through an example of how to automate this sweep in Section˜B.2. Algorithm˜3 presents our stabilized, training-ready version of Gram Newton-Schulz, which uses the restart strategy as well as a slight reformulation of the intermediate polynomials and a switch to float16 instead of bfloat16. The latter two changes are motivated in Sections˜B.2.3 and B.2.2. Our algorithm is implemented in a pip-installable package called gram-newton-schulz, which uses the Polar Express coefficients [3]. Algorithm 3 Stabilized Gram Newton-Schulz 1:∈ℝn×m X ^n× m with n≤mn≤ m, coefficients (at,bt,ct)t=15\(a_t,b_t,c_t)\_t=1^5 2:←/(∥+ϵ) X← X\,/\,( X _ F+ε) ⊳ Normalize sing vals to [0,1][0,1], ϵ=10−7ε=10^-7 3:←float16() X← float16( X) ⊳ Cast to half precision for speed 4:0←⊤ R_0← X X 5:0← Q_0← I 6:for t=1,…,5t=1,…,5 do 7: if t=3t=3 then ⊳ Restart to stabilize 8: ←2 X← Q_2 X 9: 2←⊤ R_2← X X 10: 2← Q_2← I 11: end if 12: t←btt−1+ctt−12 Z_t← b_t R_t-1+c_t R_t-1^2 13: t←t−1t+att−1 Q_t← Q_t-1 Z_t+a_t Q_t-1 ⊳ t=t−1ht(t−1) Q_t= Q_t-1h_t( R_t-1) 14: ()t←t−1t+att−1(RZ)_t← R_t-1 Z_t+a_t R_t-1 15: t←t()t+at()t R_t← Z_t(RZ)_t+a_t(RZ)_t ⊳ t=t−1ht(t−1)2 R_t= R_t-1h_t( R_t-1)^2 16:end for 17:←5 X← Q_5 X 18:return X 5 Symmetric GEMM Kernels in CuteDSL Our second contribution is a set of custom GPU kernels for the operations A B and α+βα A B+β C that assume A B and C are symmetric. These symmetric GEMM kernels compute the lower triangle of the output matrix in the usual way and copy the results to the upper triangle (Figure˜2, left), saving about half the floating point operations used by standard matrix multiplication routines. Our kernels allow us to take advantage of the symmetric structure of Gram Newton-Schulz. As noted above, they also accelerate standard Newton-Schulz, but their impact is far greater when combined with Gram Newton-Schulz. We target the Hopper and Blackwell GPU architectures. Figure˜2 (right) shows that our kernels achieve superb performance on both architectures, significantly outperforming the standard GEMM routine from cuBLAS across a range of matrix sizes. Figure 2: Left: Symmetric GEMM computes 256×256256× 256 tiles from the lower triangle and main diagonal, then transposes and copies each lower tile to the corresponding upper tile. Right: Our CuteDSL symmetric GEMM kernels benchmarked against cuBLAS GEMM kernels on Hopper and Blackwell GPUs. Input matrices ,, A, B, C have dimensions n×n× n. For large enough n, our kernels achieve a ∼2× 2× speedup over cuBLAS, both with and without an epilogue addition of C. Most GEMM kernels have the following form: 1. Scheduler: The output matrix is divided into tiles. A schedule is created that assigns each tile to a group of workers. 2. Computing each output tile of A B or α+βα A B+β C requires three steps: (a) Prologue: The rows of A and columns of B needed for the current tile are loaded in from general memory (high-bandwidth memory) to shared memory (SRAM). (b) Matrix-Multiply Accumulate (MMA): The rows and columns are multiplied and written to the register file (Hopper) or tensor memory (Blackwell). (c) Epilogue: Additional tensors needed for fused operations (e.g. C, α, β) are loaded, fused operations are executed, and the final output is written from the register file to shared memory and then to general memory. Our symmetric GEMM kernels differ from the standard one only in their schedulers and epilogues. Triangular Scheduler In the standard GEMM, the entire output matrix is partitioned into work tiles that are load balanced and evenly partitioned amongst clusters of thread blocks, where thread blocks in the same cluster can access the same shared memory and are therefore scheduled to run together. Each cluster then computes its assigned work tiles in succession. Our tile scheduler in the symmetric GEMM is almost identical, except that it only partitions work tiles from the lower triangle of the matrix (including the main diagonal); work tiles in the upper triangle are unassigned. This triangular scheduler ensures that clusters are load balanced and no unnecessary work is performed. Epilogue: Writing to the Transposed Tile In the GEMM epilogue, when the computed values of the lower triangle (excluding the main diagonal) are written to their assigned tile in general memory (HBM), they are also written to their transposed tile location in the upper triangle. The left panel of Figure˜2 illustrates this process. We describe the details of our symmetric kernels and of our kernelized implementation of standard Newton-Schulz in Appendix˜C. 6 The Dion3 Update Rule Our third contribution is a variant of Muon that avoids orthogonalizing the full momentum matrix, making the optimizer even faster. At each step, we select only a fraction of the rows from the momentum matrix and perform the Muon update (including the orthogonalization) as if the other rows did not exist. In words, the main steps of the Dion3 update rule are as follows: 1. Select a fraction of the rows (or columns) of the momentum matrix. 2. Orthogonalize the selected submatrix using Gram Newton-Schulz. 3. Update the selected rows of the weight matrix using the output of the previous step. Do not update the other rows. 4. Error Feedback Decay only the selected rows of the momentum matrix by a multiplicative factor. Unlike our other contributions, this procedure changes the optimization trajectory, but experiments in Section˜8.1 show that this change is benign. When all rows are selected (i.e., the fraction equals 1), our update rule reduces to Muon in exact arithmetic. Note that an earlier version of this update rule appeared in an unpublished report under the name Dion2. As it was not formally published, we drop the old name and present it here as part of Dion3. We now describe each step of the optimizer in detail. Algorithm˜4 gives the full pseudocode. Algorithm 4 Dion3 update rule (single weight matrix) 1:Weight ∈ℝn×m W ^n× m, gradient G, momentum buffer M, row-wise 2nd moment buffer ∈ℝn v ^n, fraction f∈(0,1]f∈(0,1], momentum μ, decay β2 _2, learning rate η. 2:←+ M← M+ G 3:←S← indices of the k=⌈fn⌉k= fn rows of M with largest ℓ1 _1 norm ⊳ Selection 4:←polar([,:]) O ( M[S,:]) ⊳ Orthogonalize k×mk× m submatrix 5:i←β2i+(1−β2)1m∑jij2 v_i← _2 v_i+(1- _2) 1m _j O_ij^2 for i∈i ⊳ Optional: NorMuon steps… 6:^ij←ij/(i+ϵ) O_ij← O_ij/( v_i+ε) 7:^←^∥/∥^∥ O← O\, O _ F/ O _ F 8:←(1−η⋅wd) W←(1-η· wd) W ⊳ Optional: weight decay 9:ij←ij−η^ij W_ij← W_ij-η\, O_ij for i∈i ⊳ Update selected rows 10:ij←μij M_ij←μ\, M_ij for i∈i ⊳ Error Feedback Selection Our update rule introduces a new hyperparameter f∈(0,1]f∈(0,1] that controls the fraction of rows (or columns) to select. We think of f as a compression factor; f=1f=1 corresponds to Muon, and decreasing f speeds up the algorithm. We recommend f=1/4f= 14 or f=1/8f= 18. We could subselect either the rows or the columns. When the weight matrix is sharded across one of these dimensions, we subselect along that same dimension. Otherwise, we pick whichever dimension is smaller. For brevity, we refer only to row selection throughout this section. We use a simple selection strategy: we pick the k=⌈fn⌉k= fn rows of largest ℓ1 _1 norm. Initial experiments showed that optimization quality when using random selection was not much worse, so other selection strategies may perform well. In the distributed setting, selecting the top rows across all shards would require an extra synchronization round and a ragged all-to-all, so we simply select the top-f fraction of rows from each shard. True global selection is available as an option in our package. Orthogonalization The benefit of our selection strategy is realized in this step. For an input of size n×mn× m, each matrix multiplication in Newton-Schulz costs either mn2mn^2 or n3n^3 FLOPs with our symmetric kernels. Selection reduces the dimensions to fn×mfn× m, slashing the cost of these multiplications by a factor of at least 1/f21/f^2. This benefit is compounded by using Gram Newton-Schulz, which is faster than standard Newton-Schulz precisely when the aspect ratio α=m/nα=m/n is large. Happily, selection increases α by a factor of 1/f1/f. Our optimizer also benefits the communication cost of Muon in the distributed setting. Rather than assemble the entire momentum matrix on a single device to perform the orthogonalization, we only assemble the selected rows, reducing the communication volume by a factor of 1/f1/f. We describe our communication strategy in detail in Section˜7. Weight Update We experiment with two flavors of the update rule: Muon and NorMuon (corresponding to (1) and (2), respectively). For the Muon flavor, we use each row of the orthogonalized submatrix =polar([,:]) O=polar( M[S,:]) to update the matching row of the weight matrix: [,:]←[,:]−η W[S,:]← W[S,:]-η O. The NorMuon flavor is similar, but with extra steps that update the row-wise second moments and use them to rescale O; see Algorithm˜4. In both flavors, if weight decay is enabled, it is applied to all rows of W before updating the selected rows. Our initial implementation of Dion3 suffered from a nefarious numerical issue that caused it to underperform Muon and NorMuon in final loss. Normally, the compiler fuses the weight update into a single operation ←(1−η⋅wd)−η W←(1-η· wd) W-η O. Tensors are cast from bfloat16 to float32, the arithmetic is performed, and the result is cast back to bfloat16. When we introduced the row selection operator, it broke this automatic fusion, creating multiple rounds of upcasting and downcasting. To restore the original numerical behavior, we implemented a custom Triton kernel for ˜8 and 9 of Algorithm˜4. Furthermore, we found that NorMuon’s normalization step should be performed in float32 to avoid similar issues. Error Feedback We adapt the error feedback mechanism from Dion [2] to our simpler selection-based strategy for approximating M. The usual rule for updating the momentum matrix is ←+;←μ M← M+ G;\, M←μ M, which ensures that M is an exponential moving average of the gradients across iterations. The main idea of error feedback is to decay only the component of M that was captured by our approximation. Decompose M into two parts, the selected part M—which matches M at the selected rows and contains zeros in all other rows—and the unselected part − M- M. (In Dion, M would be a low-rank approximation of M.) Instead of the usual rule ←μ M←μ M, error feedback does ←μ^+(−^) M←μ M+( M- M). That is, it decays the selected rows by a factor of μ but leaves the other rows alone. In effect, error feedback adds an extra (1−μ)(−^)(1-μ)( M- M) to the momentum. By boosting the residual component, this extra term nudges future iterations to select rows that were ignored by the current iteration. Even if a given row of G is consistently small, error feedback allows it to build up in M over many iterations, so it eventually gets selected and participates in the update. When the approximation is exact (= M= M), we recover standard momentum. Finite Precision At f=1f=1 every row is selected each step, and Dion3 reduces to Muon/NorMuon up to four implementation details: the selection step permutes the rows, momentum damping is applied after forming the update444That is, Newton-Schulz sees + M+ G rather than μ+μ M+ G. Since M is initialized to zero, these matrices differ by a constant factor of μ, which is removed by orthogonalization anyway., the NorMuon normalization steps run in float32, and the update step uses our custom Triton kernel. We confirmed that, despite these differences, the convergence curve of our optimizer matches that of Muon / NorMuon almost exactly (see Section˜D.4). Kernel launch minimization Dion3’s row selection process necessarily requires more kernel launches than standard Muon does. At small scale, the cost of these extra kernel launches can dominate. Since the kernel dependency structure is the same from round to round, we use a standard CUDA graph capture and replay to avoid this overhead and realize the full potential of our optimizer on smaller-scale models. 7 Megabatching and Communication Our final contribution concerns the distributed setting. Under fully-sharded data parallelism (FSDP), each momentum matrix is partitioned across GPUs, so it must first be assembled onto a single device before Newton-Schulz can run. This requires an all-to-all communication along the sharded dimension and a second all-to-all to scatter back the result. As discussed in Section˜2.3, the cost of these operations is a major obstacle to scaling Muon [11, 2]. To help users easily adopt Dion3 in a wide range of scenarios, we implement it in dion, a pip-installable package that gracefully manages this communication for FSDP2, DDP, and mixed sharding strategies with minimal overhead. The package supports our kernelized implementation of Gram Newton-Schulz and both flavors of the Dion3 update rule as well as Muon, NorMuon, and standard Newton-Schulz, all of which share the same distributed backend. We believe dion provides a full-stack, production-ready solution for distributed orthogonal optimization. A key feature of our implementation is megabatching. While the update rule of Section˜6 reduces the communication volume (the number of bits that must be assembled and scattered), megabatching reduces the number of rounds of communication. A naive implementation of Muon performs the update in batches of world_size matrices at a time; each matrix in the batch is assembled and orthogonalized in parallel on a different GPU without any GPU sitting idle. For a model with N orthogonalized matrices, this requires O(N/world_size)O(N/ world\_size) rounds of communication (all-to-alls) per optimizer step. Unfortunately, launching and synchronizing each round incurs overhead. Since this overhead is independent of the communication volume, our update rule does not mitigate it. Furthermore, each peer-to-peer message is small—just one shard of one parameter. These messages may fail to saturate the interconnect, so bandwidth utilization is poor. We demonstrate these effects in Section˜D.3. Our solution is megabatching; we group all matrices of the same shape into a single batch. For a given shape, the local momentum shards are packed into one all-to-all, assembled, orthogonalized as a batch, and scattered back together. Transformers contain only a handful of distinct weight shapes, so megabatching reduces the number of communication rounds to O(1)O(1), independent of model depth. We implement this megabatching strategy in the dion package, where it powers Muon, NorMuon, and Dion3 alike. Note that megabatching only changes when data moves; the assembly, momentum-state layout, checkpoint format, and orthogonalization math are all the same. Benchmarking We compare megabatching against the previous communication strategy by measuring the wall-clock time of the optimizer step. We hold the model, data, and optimizer (Muon) fixed, and time both strategies back-to-back on the same GPUs. We test two model sizes (1B and 14B parameters) and two FSDP configurations (8 or 32 GPUs). Results are shown in Table˜1. For the 14B model, the computational cost of Newton-Schulz dominates, so reducing communication overhead has little impact. For the 1B model with 32 shards, each rank holds only one or two matrices of each shape anyway (see Table˜4), so the batch size is about the same for both strategies. However, for the 1B model with 8 shards, the optimizer is communication-bound and each rank holds many matrices. Here, megabatching has a major impact, reducing step time by 35%35\%. Table 1: Effect of megabatched communication on the per-GPU optimizer step time (Muon, median over steps). “Batching” orthogonalizes matrices in groups whose size equals the shard count; “megabatch” uses a single batch per shape group. When the weights are small (so Newton-Schulz is cheap) and each rank holds several matrices, megabatching substantially reduces optimizer time. Each node has four GH200s. Model size Nodes FSDP Shards Batching (ms) Megabatching (ms) Change 1B 1 8 80.780.7 52.152.1 −35%-35\% 1B 4 32 61.961.9 59.359.3 −4%-4\% 14B 1 8 144.0144.0 140.8140.8 −2%-2\% 14B 4 32 94.794.7 89.189.1 −6%-6\% Compressed Data Parallelism There is a prospect for further savings at larger scales, where data parallelism spans multiple pods in a data-center network (DCN) or hybrid sharding is used [35]. A common deployment combines model parallelism within an inter-chip interconnect (ICI) with pure data parallelism across pods [4]. Because DCN bandwidth is typically far smaller than ICI bandwidth, shrinking the volume of data-parallel communication is especially valuable in this regime. Here, Dion3 allows for compressed data-parallel synchronizations analogous to those of Dion [2, §3.3]. Ordinarily, gradients are averaged across data-parallel replicas at every step so that the momentum buffers, and hence the weights, stay consistent. With Dion3, however, we need not synchronize the entire momentum matrix at every step, only the selected submatrix [,:] M[S,:]. When the rows S can be selected without first synchronizing M (e.g., by picking S at random), this strategy performs an identical update at a fraction of the communication cost. 8 Experiments 8.1 Model Quality Is Preserved Of our four contributions, only our update rule explicitly changes the optimizer trajectory. However, Gram Newton-Schulz and our symmetric GEMM kernels introduce numerical differences due to finite precision arithmetic. Experiments in Section˜A.2 show that these differences do not affect training. We also confirmed that, as expected, Dion3 with selection fraction f=1f=1 matches Muon/NorMuon almost exactly (see Figure˜22). In contrast, Dion3 with f<1f<1 is a genuinely new optimizer. In this section, we demonstrate that Dion3 achieves a slightly better loss than the baseline. We did not set out to improve training quality, so this result is a small but pleasant surprise. At a minimum, it shows that our subselection strategy does no harm to the loss. We train 1B-parameter dense transformers on 100B tokens of the ClimbMix dataset [10], using fully-sharded data parallelism (FSDP) with MXFP8 weights. (Below, we also train at larger scales.) As always, Newton-Schulz runs in half precision arithmetic. The models use grouped-query and sliding-window attention, but are not mixtures of experts. We report the cross-entropy loss (CE) of next-token prediction on a held-out validation split of ClimbMix. For further details about our setup, see Section˜D.1. For a fair comparison, we begin by tuning the baseline. In this section, we compare NorMuon to the NorMuon flavor of Dion3, as initial experiments gave it a slight edge over Muon. We find that the optimal learning rate and momentum coefficient for NorMuon are η=0.01η=0.01 and μ=0.9μ=0.9, respectively. Learning-rate transfer rule Since Dion3 updates only a fraction of the rows in each step, it reduces the effective step size. This in turn changes the optimal learning rate. To understand this effect, we sweep both the fraction f∈1/2,1/4,1/8,1/16f∈\ 12, 14, 18, 116\ and the learning rate η of Dion3, along with our earlier sweep for NorMuon (f=1f=1). The results in Figure˜3 show a clear trend: the optimal learning rate for a given f scales as η∝1/fη 1/f. We can explain this transfer rule as follows. For Muon, the size of the update is ‖η⋅polar()‖=ηn\|η·polar( M)\|_ F=η n, where n is the smaller dimension. For Dion3, it is ‖η′⋅polar([,:])‖=η′fn\|η ·polar( M[S,:])\|_ F=η fn. For these to match, we must set η′=η/fη =η/ f. We do not bother retuning the momentum coefficient; we simply copy the optimal setting from NorMuon (μ=0.9μ=0.9) to Dion3. Figure 3: Left: Final validation loss for 1B-parameter models trained on 100B tokens of ClimbMix with NorMuon or Dion3, as a function of row-selection fraction f and learning rate η. NorMuon is f=1f=1. Bolded cell in each row shows best η. Optimal η track the line ηf=0.01η f=0.01, shown in red on a log-log scale. Not shown: f=1,η=0.03f=1,η=0.03 (loss = 2.2372.237) and f=1/32,η=0.04f= 132,η=0.04 (loss =2.191=2.191). Right: The same data reduced to two dimensions: validation loss as a function of ηfη f. Final loss is minimized when ηf≈0.01η f≈ 0.01. Dion3 improves the loss Figure˜3 contains another remarkable finding: Dion3 with f<1f<1 actually outperforms NorMuon when tuned correctly. Indeed, the lowest loss is achieved at f=1/8f= 18. Figure˜4 shows the validation loss curve for each optimizer (excluding f=1/4f= 14) at its best learning rate. It shows a clear gap throughout training between NorMuon and the Dion3 variants, which confirms this finding. Figure 4: Validation loss curves for 1B-parameter models trained on 100B tokens of ClimbMix with NorMuon or Dion3 at f=1/2,1/8f= 12, 18, or 1/16 116. Each optimizer uses its best learning rate from Figure˜3. Parentheses show final loss. All Dion3 variants track below the fully-tuned NorMuon baseline throughout training, finishing about 0.010.01 points lower. To strengthen this finding, we scale up the model size to between 3B and 14B parameters. For reasons of cost, we now train on 10B tokens instead of 100B. We compare Dion3 with f=1/4f= 14 against NorMuon, each using their best learning rate as identified above. For each scale and optimizer, Table˜2 reports both validation loss on ClimbMix and downstream accuracy on a suite of 12 standard benchmarks. In validation loss, Dion3 outperforms NorMuon at every scale, achieving the largest improvement (−0.027-0.027) at the largest scale (14B). Figure˜5 shows that, as for the 1B parameter model, validation loss is lower throughout training when using Dion3. In downstream accuracy, Dion3 wins at three of four scales, with an improvement of 0.7 percentage points at 14B. Table 2: Dion3 (f=1/4f\!=\! 14, η=0.02η\!=\!0.02) versus NorMuon (η=0.01η\!=\!0.01) across model sizes trained on 10B tokens of ClimbMix: final validation loss (cross-entropy; lower is better) and downstream accuracy (macro-average over 12 standard benchmarks: ARC easy/challenge, BoolQ, COPA, HellaSwag, LAMBADA, MMLU, OpenBookQA, PIQA, RTE, TruthfulQA, WinoGrande; higher is better). The winner of each comparison is bolded. Model Size Validation Loss Downstream Accuracy (%) NorMuon Dion3 (f=1/4f\!=\! 14) Δ NorMuon Dion3 (f=1/4f\!=\! 14) Δ 3B 2.269 2.257 −0.012-0.012 53.9 54.9 +1.0+1.0 4B 2.243 2.232 −0.011-0.011 55.2 54.9 −0.3-0.3 7B 2.220 2.206 −0.014-0.014 56.0 56.1 +0.1+0.1 14B 2.189 2.162 −0.027-0.027 57.4 58.1 +0.7+0.7 Figure 5: Validation loss over training at 14B for Dion3 (f=1/4f\!=\! 14, η=0.02η\!=\!0.02) versus NorMuon (η=0.01η\!=\!0.01), trained on 10B tokens of ClimbMix. Dion3 tracks below NorMuon throughout. 8.2 Dion3 Accelerates the Optimizer We now measure how much Dion3 reduces the cost of Muon’s optimizer step. Starting with standard Muon, we successively add our improvements—symmetric kernels, Gram Newton-Schulz, and fractional updates with f=0.5f=0.5 or f=0.25f=0.25—stacking each on top of the previous ones. (All versions use megabatching.) We also compare against plain AdamW. We run each version of the optimizer on models of various sizes either on 1 GPU or using FSDP on 4 GPUs. We use CUDA events to measure the GPU time of the optimizer step. For fidelity, we run full training steps with forward and backward passes (on synthetic data) but these are excluded from the recorded timings. Figure 6: Optimizer step time (excluding forward/backward pass) relative to standard Muon across model scales when training on 1 GH200 (left) and FSDP over 4 GH200s (right). Each colored line adds one of our contributions on top of the previous one; AdamW is shown for reference. Lines show median over 25 steps and bands show interquartile range. Adding each of our contributions consistently reduces runtime; together they achieve a 6×6× speedup for larger models. Results are shown in Figure˜6. A subset of these results appears in Figure˜1, and the corresponding timings for the NorMuon family are given in Section˜D.4 (Figure˜23). Each part of Dion3 yields a consistent speedup across scales and parallelism configurations. Our symmetric kernels and Gram Newton-Schulz give a combined speedup of 1.5×1.5× or more over standard Muon. The fractional update rule cuts the runtime dramatically. As expected, its impact grows with model size, as the cubic computational cost of Newton-Schulz increasingly dominates the other operations. For sufficiently large models, setting f=1/2f= 12 and f=1/4f= 14 gives further 2×2× and 3.7×3.7× reductions respectively, for overall speedups of 3.6×3.6× and 6.5×6.5×. Thus, while Dion3 is still slower than AdamW, the gap is significantly smaller. In Section˜D.2, we report CPU time and communication volume in addition to GPU time, showing that CPU overhead is small and that fractional updates reduce communication by a factor of 1/f1/f. The communication savings of our fractional update rule can be crucial in some settings, though the setting of Figure˜6 is compute-bound: neither communication nor host time is exposed. In Appendix˜A, we conduct an additional suite of benchmarks on a wider range of architectures, including open source models like Gemma and mixtures-of-experts. On architectures like these, whose weight matrices have higher aspect ratios (α=8α=8 instead of 44), Gram Newton-Schulz and the symmetric kernels alone achieve a speedup of 2×2× (see Figures˜9 and 10). 9 Conclusion For LLMs trained with Muon, estimates suggest that optimizer step time accounts for between 1% (e.g., [1]) and 17% of total training time (see Appendix˜E). While Muon has been scaled successfully, doing so has demanded particular choices of architecture and parallelism strategy, along with significant engineering effort. As interest in Muon continues to spread, this paradigm is too brittle; a flexible, general-purpose remedy for Muon’s overhead is needed. Dion3 allows practitioners to easily realize the benefits of Muon without paying the high cost of its orthogonalization step—even for large-scale models in highly distributed settings. Our Gram Newton-Schulz algorithm and CuteDSL kernels speed up Muon by 1.5×1.5× for dense models and 2×2× for MoEs, a rare case of free lunch performance. Megabatching has an equally large effect in certain distributed settings. Dion3’s fractional update rule (f=1/4f= 14) provides an additional 3.7×3.7× speedup. Though this update rule changes the optimizer trajectory, we find that it actually improves training quality in our setting. This improvement is unexpected, as we designed Dion3 to cheaply approximate Muon rather than improve upon it. Further work is needed to determine how widely this improvement generalizes, but we find support in a recent result similar to our own: [18] showed that randomly masking blocks of the update improves the trajectory of SGD with momentum. Overall, Dion3 as implemented in our dion and gram-newton-schulz packages provides the tools to make orthogonal optimizers practical and accessible in a wide range of settings. Acknowledgements The authors thank Zichong Li for contributing an implementation of NorMuon to the dion repo. NA is supported by NSF award 2234660. References [1] J. Abadji, M. Abdin, C. Adams, E. Alcaide, M. Altun, M. Artoni, J. Bao, U. Barar, V. Bekiaris, A. Bessonov, B. Bütikofer, J. Chang, Y. Chen, D. Chernenkov, Y. Chi, F. Christianos, F. Christopoulou, R. Ciocoiu, T. Cohen, Y. Coppel, D. Emelianenko, B. Fergerson, B. Fitzgerald, M. Gallé, A. Golonzovskyi, G. Grigorev, Y. Hao, C. Hensel, J. Huenermann, Y. Ji, S. Joshi, E. Kant, K. Khandpur, S. Kim, V. Kirichenko, U. Kocasarac, I. Kochik, I. Komarov, C. Kong, A. Koul, F. Lacroix, S. Laktionov, W. Long, Q. Malartic, V. Markovtsev, A. Marques, R. McHardy, C. Mocholí, D. Monakhov, A. Morris, M. Muller, C. Mürtz, R. Nabel, T. Nguyen, R. Novosel, S. Ozog, A. Patankar, A. Petrov, A. Piché, A. Pignet, T. Poncu, P. Potter, A. Rakowski, P. Ritschard, J. Roberts, J. Rowell, P. Sarna, P. Savalle, U. Sazanovich, N. Shapovalov, A. Shevchenko, M. Shilkov, A. Sokol, M. Soliman, J. Stephenson, V. Storchan, D. Tantaru, A. Tyurin, A. Wälchli, P. Wang, J. Yang, R. Zayashnikov, A. Z. Martin, N. Zinov, C. Bercier, J. Caldeira, M. Garcia, T. George, K. Gharzai, G. Hitchcock, C. Klingenberg, I. Pinto, V. Randery, N. Smith, A. Sugako, and J. Warner (2026) Laguna m.1/xs.2 technical report. External Links: 2605.27605, Link Cited by: §9. [2] K. Ahn, B. Xu, N. Abreu, Y. Fan, G. Magakyan, P. Sharma, Z. Zhan, and J. Langford (2025) Dion: distributed orthonormalized updates. External Links: 2504.05295, Link Cited by: item 3, 2nd item, §3.2, §3.2, §6, §7, §7. [3] N. Amsel, D. Persson, C. Musco, and R. M. Gower (2026) The polar express: optimal matrix sign methods and their application to the muon algorithm. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §A.2, §B.2.1, §3.1, §3.1, §4.2. [4] J. Austin, S. Douglas, R. Frostig, A. Levskaya, C. Chen, S. Vikram, F. Lebron, P. Choy, V. Ramasesh, A. Webson, and R. Pope (2025) How to scale your model. Google DeepMind. External Links: Link Cited by: §7. [5] J. Bernstein (2025) Deriving Muon. External Links: Link Cited by: §2.1. [6] T. Boissin, T. Massena, F. Mamalet, and M. Serrurier (2025) Turbo-muon: accelerating orthogonality-based optimization with pre-conditioning. External Links: 2512.04632, Link Cited by: §3.1. [7] V. Boreiko, Z. Bu, and S. Zha (2025) Towards understanding of orthogonalization in muon. In Tiny Titans: The next wave of On-Device Learning for Foundation Models (TTODLer-FM), External Links: Link Cited by: §3.2. [8] F. L. Cesista, J. You, and K. Jordan (2025) Squeezing 1-2% efficiency gains out of muon by optimizing the newton-schulz coefficients. External Links: Link Cited by: §A.2, §A.2. [9] DeepSeek-AI, A. Liu, B. Feng, B. Xue, B. Wang, B. Wu, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, D. Dai, D. Guo, D. Yang, D. Chen, D. Ji, E. Li, F. Lin, F. Dai, F. Luo, G. Hao, G. Chen, G. Li, H. Zhang, H. Bao, H. Xu, H. Wang, H. Zhang, H. Ding, H. Xin, H. Gao, H. Li, H. Qu, J. L. Cai, J. Liang, J. Guo, J. Ni, J. Li, J. Wang, J. Chen, J. Chen, J. Yuan, J. Qiu, J. Li, J. Song, K. Dong, K. Hu, K. Gao, K. Guan, K. Huang, K. Yu, L. Wang, L. Zhang, L. Xu, L. Xia, L. Zhao, L. Wang, L. Zhang, M. Li, M. Wang, M. Zhang, M. Zhang, M. Tang, M. Li, N. Tian, P. Huang, P. Wang, P. Zhang, Q. Wang, Q. Zhu, Q. Chen, Q. Du, R. J. Chen, R. L. Jin, R. Ge, R. Zhang, R. Pan, R. Wang, R. Xu, R. Zhang, R. Chen, S. S. Li, S. Lu, S. Zhou, S. Chen, S. Wu, S. Ye, S. Ye, S. Ma, S. Wang, S. Zhou, S. Yu, S. Zhou, S. Pan, T. Wang, T. Yun, T. Pei, T. Sun, W. L. Xiao, W. Zeng, W. Zhao, W. An, W. Liu, W. Liang, W. Gao, W. Yu, W. Zhang, X. Q. Li, X. Jin, X. Wang, X. Bi, X. Liu, X. Wang, X. Shen, X. Chen, X. Zhang, X. Chen, X. Nie, X. Sun, X. Wang, X. Cheng, X. Liu, X. Xie, X. Liu, X. Yu, X. Song, X. Shan, X. Zhou, X. Yang, X. Li, X. Su, X. Lin, Y. K. Li, Y. Q. Wang, Y. X. Wei, Y. X. Zhu, Y. Zhang, Y. Xu, Y. Xu, Y. Huang, Y. Li, Y. Zhao, Y. Sun, Y. Li, Y. Wang, Y. Yu, Y. Zheng, Y. Zhang, Y. Shi, Y. Xiong, Y. He, Y. Tang, Y. Piao, Y. Wang, Y. Tan, Y. Ma, Y. Liu, Y. Guo, Y. Wu, Y. Ou, Y. Zhu, Y. Wang, Y. Gong, Y. Zou, Y. He, Y. Zha, Y. Xiong, Y. Ma, Y. Yan, Y. Luo, Y. You, Y. Liu, Y. Zhou, Z. F. Wu, Z. Z. Ren, Z. Ren, Z. Sha, Z. Fu, Z. Xu, Z. Huang, Z. Zhang, Z. Xie, Z. Zhang, Z. Hao, Z. Gou, Z. Ma, Z. Yan, Z. Shao, Z. Xu, Z. Wu, Z. Zhang, Z. Li, Z. Gu, Z. Zhu, Z. Liu, Z. Li, Z. Xie, Z. Song, Z. Gao, and Z. Pan (2025) DeepSeek-v3 technical report. External Links: 2412.19437, Link Cited by: item 1, Appendix E, 1st item, item 2. [10] S. Diao, Y. Yang, Y. Fu, X. Dong, D. Su, M. Kliegl, Z. Chen, P. Belcak, Y. Suhara, H. Yin, M. Patwary, Y. C. Lin, J. Kautz, and P. Molchanov (2025) Nemotron-CLIMB: clustering-based iterative data mixture bootstrapping for language model pre-training. In The Thirty-ninth Annual Conference on Neural Information Processing Systems Datasets and Benchmarks Track, External Links: Link Cited by: §8.1. [11] Essential AI (2025) Layer Sharding for Large-Scale Training with Muon. External Links: Link Cited by: 1st item, 2nd item, §7. [12] Gemma Team, A. Kamath, J. Ferret, S. Pathak, N. Vieillard, R. Merhej, S. Perrin, T. Matejovicova, A. Ramé, M. Rivière, L. Rouillard, T. Mesnard, G. Cideron, J. Grill, S. Ramos, E. Yvinec, M. Casbon, E. Pot, I. Penchev, G. Liu, F. Visin, K. Kenealy, L. Beyer, X. Zhai, A. Tsitsulin, R. Busa-Fekete, A. Feng, N. Sachdeva, B. Coleman, Y. Gao, B. Mustafa, I. Barr, E. Parisotto, D. Tian, M. Eyal, C. Cherry, J. Peter, D. Sinopalnikov, S. Bhupatiraju, R. Agarwal, M. Kazemi, D. Malkin, R. Kumar, D. Vilar, I. Brusilovsky, J. Luo, A. Steiner, A. Friesen, A. Sharma, A. Sharma, A. M. Gilady, A. Goedeckemeyer, A. Saade, A. Feng, A. Kolesnikov, A. Bendebury, A. Abdagic, A. Vadi, A. György, A. S. Pinto, A. Das, A. Bapna, A. Miech, A. Yang, A. Paterson, A. Shenoy, A. Chakrabarti, B. Piot, B. Wu, B. Shahriari, B. Petrini, C. Chen, C. L. Lan, C. A. Choquette-Choo, C. Carey, C. Brick, D. Deutsch, D. Eisenbud, D. Cattle, D. Cheng, D. Paparas, D. S. Sreepathihalli, D. Reid, D. Tran, D. Zelle, E. Noland, E. Huizenga, E. Kharitonov, F. Liu, G. Amirkhanyan, G. Cameron, H. Hashemi, H. Klimczak-Plucińska, H. Singh, H. Mehta, H. T. Lehri, H. Hazimeh, I. Ballantyne, I. Szpektor, I. Nardini, J. Pouget-Abadie, J. Chan, J. Stanton, J. Wieting, J. Lai, J. Orbay, J. Fernandez, J. Newlan, J. Ji, J. Singh, K. Black, K. Yu, K. Hui, K. Vodrahalli, K. Greff, L. Qiu, M. Valentine, M. Coelho, M. Ritter, M. Hoffman, M. Watson, M. Chaturvedi, M. Moynihan, M. Ma, N. Babar, N. Noy, N. Byrd, N. Roy, N. Momchev, N. Chauhan, N. Sachdeva, O. Bunyan, P. Botarda, P. Caron, P. K. Rubenstein, P. Culliton, P. Schmid, P. G. Sessa, P. Xu, P. Stanczyk, P. Tafti, R. Shivanna, R. Wu, R. Pan, R. Rokni, R. Willoughby, R. Vallu, R. Mullins, S. Jerome, S. Smoot, S. Girgin, S. Iqbal, S. Reddy, S. Sheth, S. Põder, S. Bhatnagar, S. R. Panyam, S. Eiger, S. Zhang, T. Liu, T. Yacovone, T. Liechty, U. Kalra, U. Evci, V. Misra, V. Roseberry, V. Feinberg, V. Kolesnikov, W. Han, W. Kwon, X. Chen, Y. Chow, Y. Zhu, Z. Wei, Z. Egyed, V. Cotruta, M. Giang, P. Kirk, A. Rao, K. Black, N. Babar, J. Lo, E. Moreira, L. G. Martins, O. Sanseviero, L. Gonzalez, Z. Gleicher, T. Warkentin, V. Mirrokni, E. Senter, E. Collins, J. Barral, Z. Ghahramani, R. Hadsell, Y. Matias, D. Sculley, S. Petrov, N. Fiedel, N. Shazeer, O. Vinyals, J. Dean, D. Hassabis, K. Kavukcuoglu, C. Farabet, E. Buchatskaya, J. Alayrac, R. Anil, Dmitry, Lepikhin, S. Borgeaud, O. Bachem, A. Joulin, A. Andreev, C. Hardin, R. Dadashi, and L. Hussenot (2025) Gemma 3 technical report. External Links: 2503.19786, Link Cited by: §A.1. [13] GLM-5 Team, A. Zeng, X. Lv, Z. Hou, Z. Du, Q. Zheng, B. Chen, D. Yin, C. Ge, C. Huang, C. Xie, C. Zhu, C. Yin, C. Wang, G. Pan, H. Zeng, H. Zhang, H. Wang, H. Chen, J. Zhang, J. Jiao, J. Guo, J. Wang, J. Du, J. Wu, K. Wang, L. Li, L. Fan, L. Zhong, M. Liu, M. Zhao, P. Du, Q. Dong, R. Lu, Shuang-Li, S. Cao, S. Liu, T. Jiang, X. Chen, X. Zhang, X. Huang, X. Dong, Y. Xu, Y. Wei, Y. An, Y. Niu, Y. Zhu, Y. Wen, Y. Cen, Y. Bai, Z. Qiao, Z. Wang, Z. Wang, Z. Zhu, Z. Liu, Z. Li, B. Wang, B. Wen, C. Huang, C. Cai, C. Yu, C. Li, C. Hu, C. Zhang, D. Zhang, D. Lin, D. Yang, D. Wang, D. Ai, E. Zhu, F. Yi, F. Chen, G. Wen, H. Sun, H. Zhao, H. Hu, H. Zhang, H. Liu, H. Zhang, H. Peng, H. Tai, H. Zhang, H. Liu, H. Wang, H. Yan, H. Ge, H. Liu, H. Chu, J. Zhao, J. Wang, J. Zhao, J. Ren, J. Wang, J. Zhang, J. Gui, J. Zhao, J. Li, J. An, J. Li, J. Yuan, J. Du, J. Liu, J. Zhi, J. Duan, K. Zhou, K. Wei, K. Wang, K. Luo, L. Zhang, L. Sha, L. Xu, L. Wu, L. Ding, L. Chen, M. Li, N. Lin, P. Ta, Q. Zou, R. Song, R. Yang, S. Tu, S. Yang, S. Wu, S. Zhang, S. Li, S. Li, S. Fan, W. Qin, W. Tian, W. Zhang, W. Yu, W. Liang, X. Kuang, X. Cheng, X. Li, X. Yan, X. Hu, X. Ling, X. Fan, X. Xia, X. Zhang, X. Zhang, X. Pan, X. Zou, X. Zhang, Y. Liu, Y. Wu, Y. Li, Y. Wang, Y. Zhu, Y. Tan, Y. Zhou, Y. Pan, Y. Zhang, Y. Su, Y. Geng, Y. Yan, Y. Tan, Y. Bi, Y. Shen, Y. Yang, Y. Li, Y. Liu, Y. Wang, Y. Li, Y. Wu, Y. Zhang, Y. Duan, Y. Zhang, Z. Liu, Z. Jiang, Z. Yan, Z. Zhang, Z. Wei, Z. Chen, Z. Feng, Z. Yao, Z. Chai, Z. Wang, Z. Zhang, B. Xu, M. Huang, H. Wang, J. Li, Y. Dong, and J. Tang (2026) GLM-5: from vibe coding to agentic engineering. External Links: 2602.15763, Link Cited by: §A.1.1, §1. [14] A. Grattafiori, A. Dubey, A. Jauhri, A. Pandey, A. Kadian, A. Al-Dahle, A. Letman, A. Mathur, A. Schelten, A. Vaughan, A. Yang, A. Fan, A. Goyal, A. Hartshorn, A. Yang, A. Mitra, A. Sravankumar, A. Korenev, A. Hinsvark, A. Rao, A. Zhang, A. Rodriguez, A. Gregerson, A. Spataru, B. Roziere, B. Biron, B. Tang, B. Chern, C. Caucheteux, C. Nayak, C. Bi, C. Marra, C. McConnell, C. Keller, C. Touret, C. Wu, C. Wong, C. C. Ferrer, C. Nikolaidis, D. Allonsius, D. Song, D. Pintz, D. Livshits, D. Wyatt, D. Esiobu, D. Choudhary, D. Mahajan, D. Garcia-Olano, D. Perino, D. Hupkes, E. Lakomkin, E. AlBadawy, E. Lobanova, E. Dinan, E. M. Smith, F. Radenovic, F. Guzmán, F. Zhang, G. Synnaeve, G. Lee, G. L. Anderson, G. Thattai, G. Nail, G. Mialon, G. Pang, G. Cucurell, H. Nguyen, H. Korevaar, H. Xu, H. Touvron, I. Zarov, I. A. Ibarra, I. Kloumann, I. Misra, I. Evtimov, J. Zhang, J. Copet, J. Lee, J. Geffert, J. Vranes, J. Park, J. Mahadeokar, J. Shah, J. van der Linde, J. Billock, J. Hong, J. Lee, J. Fu, J. Chi, J. Huang, J. Liu, J. Wang, J. Yu, J. Bitton, J. Spisak, J. Park, J. Rocca, J. Johnstun, J. Saxe, J. Jia, K. V. Alwala, K. Prasad, K. Upasani, K. Plawiak, K. Li, K. Heafield, K. Stone, K. El-Arini, K. Iyer, K. Malik, K. Chiu, K. Bhalla, K. Lakhotia, L. Rantala-Yeary, L. van der Maaten, L. Chen, L. Tan, L. Jenkins, L. Martin, L. Madaan, L. Malo, L. Blecher, L. Landzaat, L. de Oliveira, M. Muzzi, M. Pasupuleti, M. Singh, M. Paluri, M. Kardas, M. Tsimpoukelli, M. Oldham, M. Rita, M. Pavlova, M. Kambadur, M. Lewis, M. Si, M. K. Singh, M. Hassan, N. Goyal, N. Torabi, N. Bashlykov, N. Bogoychev, N. Chatterji, N. Zhang, O. Duchenne, O. Çelebi, P. Alrassy, P. Zhang, P. Li, P. Vasic, P. Weng, P. Bhargava, P. Dubal, P. Krishnan, P. S. Koura, P. Xu, Q. He, Q. Dong, R. Srinivasan, R. Ganapathy, R. Calderer, R. S. Cabral, R. Stojnic, R. Raileanu, R. Maheswari, R. Girdhar, R. Patel, R. Sauvestre, R. Polidoro, R. Sumbaly, R. Taylor, R. Silva, R. Hou, R. Wang, S. Hosseini, S. Chennabasappa, S. Singh, S. Bell, S. S. Kim, S. Edunov, S. Nie, S. Narang, S. Raparthy, S. Shen, S. Wan, S. Bhosale, S. Zhang, S. Vandenhende, S. Batra, S. Whitman, S. Sootla, S. Collot, S. Gururangan, S. Borodinsky, T. Herman, T. Fowler, T. Sheasha, T. Georgiou, T. Scialom, T. Speckbacher, T. Mihaylov, T. Xiao, U. Karn, V. Goswami, V. Gupta, V. Ramanathan, V. Kerkez, V. Gonguet, V. Do, V. Vogeti, V. Albiero, V. Petrovic, W. Chu, W. Xiong, W. Fu, W. Meers, X. Martinet, X. Wang, X. Wang, X. E. Tan, X. Xia, X. Xie, X. Jia, X. Wang, Y. Goldschlag, Y. Gaur, Y. Babaei, Y. Wen, Y. Song, Y. Zhang, Y. Li, Y. Mao, Z. D. Coudert, Z. Yan, Z. Chen, Z. Papakipos, A. Singh, A. Srivastava, A. Jain, A. Kelsey, A. Shajnfeld, A. Gangidi, A. Victoria, A. Goldstand, A. Menon, A. Sharma, A. Boesenberg, A. Baevski, A. Feinstein, A. Kallet, A. Sangani, A. Teo, A. Yunus, A. Lupu, A. Alvarado, A. Caples, A. Gu, A. Ho, A. Poulton, A. Ryan, A. Ramchandani, A. Dong, A. Franco, A. Goyal, A. Saraf, A. Chowdhury, A. Gabriel, A. Bharambe, A. Eisenman, A. Yazdan, B. James, B. Maurer, B. Leonhardi, B. Huang, B. Loyd, B. D. Paola, B. Paranjape, B. Liu, B. Wu, B. Ni, B. Hancock, B. Wasti, B. Spence, B. Stojkovic, B. Gamido, B. Montalvo, C. Parker, C. Burton, C. Mejia, C. Liu, C. Wang, C. Kim, C. Zhou, C. Hu, C. Chu, C. Cai, C. Tindal, C. Feichtenhofer, C. Gao, D. Civin, D. Beaty, D. Kreymer, D. Li, D. Adkins, D. Xu, D. Testuggine, D. David, D. Parikh, D. Liskovich, D. Foss, D. Wang, D. Le, D. Holland, E. Dowling, E. Jamil, E. Montgomery, E. Presani, E. Hahn, E. Wood, E. Le, E. Brinkman, E. Arcaute, E. Dunbar, E. Smothers, F. Sun, F. Kreuk, F. Tian, F. Kokkinos, F. Ozgenel, F. Caggioni, F. Kanayet, F. Seide, G. M. Florez, G. Schwarz, G. Badeer, G. Swee, G. Halpern, G. Herman, G. Sizov, Guangyi, Zhang, G. Lakshminarayanan, H. Inan, H. Shojanazeri, H. Zou, H. Wang, H. Zha, H. Habeeb, H. Rudolph, H. Suk, H. Aspegren, H. Goldman, H. Zhan, I. Damlaj, I. Molybog, I. Tufanov, I. Leontiadis, I. Veliche, I. Gat, J. Weissman, J. Geboski, J. Kohli, J. Lam, J. Asher, J. Gaya, J. Marcus, J. Tang, J. Chan, J. Zhen, J. Reizenstein, J. Teboul, J. Zhong, J. Jin, J. Yang, J. Cummings, J. Carvill, J. Shepard, J. McPhie, J. Torres, J. Ginsburg, J. Wang, K. Wu, K. H. U, K. Saxena, K. Khandelwal, K. Zand, K. Matosich, K. Veeraraghavan, K. Michelena, K. Li, K. Jagadeesh, K. Huang, K. Chawla, K. Huang, L. Chen, L. Garg, L. A, L. Silva, L. Bell, L. Zhang, L. Guo, L. Yu, L. Moshkovich, L. Wehrstedt, M. Khabsa, M. Avalani, M. Bhatt, M. Mankus, M. Hasson, M. Lennie, M. Reso, M. Groshev, M. Naumov, M. Lathi, M. Keneally, M. Liu, M. L. Seltzer, M. Valko, M. Restrepo, M. Patel, M. Vyatskov, M. Samvelyan, M. Clark, M. Macey, M. Wang, M. J. Hermoso, M. Metanat, M. Rastegari, M. Bansal, N. Santhanam, N. Parks, N. White, N. Bawa, N. Singhal, N. Egebo, N. Usunier, N. Mehta, N. P. Laptev, N. Dong, N. Cheng, O. Chernoguz, O. Hart, O. Salpekar, O. Kalinli, P. Kent, P. Parekh, P. Saab, P. Balaji, P. Rittner, P. Bontrager, P. Roux, P. Dollar, P. Zvyagina, P. Ratanchandani, P. Yuvraj, Q. Liang, R. Alao, R. Rodriguez, R. Ayub, R. Murthy, R. Nayani, R. Mitra, R. Parthasarathy, R. Li, R. Hogan, R. Battey, R. Wang, R. Howes, R. Rinott, S. Mehta, S. Siby, S. J. Bondu, S. Datta, S. Chugh, S. Hunt, S. Dhillon, S. Sidorov, S. Pan, S. Mahajan, S. Verma, S. Yamamoto, S. Ramaswamy, S. Lindsay, S. Lindsay, S. Feng, S. Lin, S. C. Zha, S. Patil, S. Shankar, S. Zhang, S. Zhang, S. Wang, S. Agarwal, S. Sajuyigbe, S. Chintala, S. Max, S. Chen, S. Kehoe, S. Satterfield, S. Govindaprasad, S. Gupta, S. Deng, S. Cho, S. Virk, S. Subramanian, S. Choudhury, S. Goldman, T. Remez, T. Glaser, T. Best, T. Koehler, T. Robinson, T. Li, T. Zhang, T. Matthews, T. Chou, T. Shaked, V. Vontimitta, V. Ajayi, V. Montanez, V. Mohan, V. S. Kumar, V. Mangla, V. Ionescu, V. Poenaru, V. T. Mihailescu, V. Ivanov, W. Li, W. Wang, W. Jiang, W. Bouaziz, W. Constable, X. Tang, X. Wu, X. Wang, X. Wu, X. Gao, Y. Kleinman, Y. Chen, Y. Hu, Y. Jia, Y. Qi, Y. Li, Y. Zhang, Y. Zhang, Y. Adi, Y. Nam, Yu, Wang, Y. Zhao, Y. Hao, Y. Qian, Y. Li, Y. He, Z. Rait, Z. DeVito, Z. Rosnbrick, Z. Wen, Z. Yang, Z. Zhao, and Z. Ma (2024) The llama 3 herd of models. External Links: 2407.21783, Link Cited by: §A.1, Appendix E. [15] E. Grishina, M. Smirnov, and M. Rakhuba (2026) Accelerating newton-schulz iteration for orthogonalization via chebyshev-type polynomials. External Links: 2506.10935, Link Cited by: §3.1. [16] W. Guo, M. Mishra, X. Cheng, I. Stoica, and T. Dao (2026) SonicMoE: accelerating moe with io and tile-aware optimizations. External Links: 2512.14080, Link Cited by: 2nd item. [17] N. J. Higham (2008) Functions of matrices: Theory and Computation. edition, Society for Industrial and Applied Mathematics, . External Links: Document, Link Cited by: §3.1. [18] T. Joo, W. Xia, C. Kim, M. Zhang, and E. Ie (2026) On surprising effectiveness of masking updates in adaptive optimizers. External Links: 2602.15322, Link Cited by: §9. [19] K. Jordan, Y. Jin, V. Boza, J. You, F. Cesista, L. Newhouse, and J. Bernstein (2024) Muon: an optimizer for hidden layers in neural networks. External Links: Link Cited by: §2.1, §2.2. [20] A. Khaled, K. Ozkara, T. Yu, M. Hong, and Y. Park (2026) MuonBP: faster muon via block-periodic orthogonalization. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §3.2. [21] Kimi Team, Y. Bai, Y. Bao, Y. Charles, C. Chen, G. Chen, H. Chen, H. Chen, J. Chen, N. Chen, R. Chen, Y. Chen, Y. Chen, Y. Chen, Z. Chen, J. Cui, H. Ding, M. Dong, A. Du, C. Du, D. Du, Y. Du, Y. Fan, Y. Feng, K. Fu, B. Gao, C. Gao, H. Gao, P. Gao, T. Gao, Y. Ge, S. Geng, Q. Gu, X. Gu, L. Guan, H. Guo, J. Guo, X. Hao, T. He, W. He, W. He, Y. He, C. Hong, H. Hu, Y. Hu, Z. Hu, W. Huang, Z. Huang, Z. Huang, T. Jiang, Z. Jiang, X. Jin, Y. Kang, G. Lai, C. Li, F. Li, H. Li, M. Li, W. Li, Y. Li, Y. Li, Y. Li, Z. Li, Z. Li, H. Lin, X. Lin, Z. Lin, C. Liu, C. Liu, H. Liu, J. Liu, J. Liu, L. Liu, S. Liu, T. Y. Liu, T. Liu, W. Liu, Y. Liu, Y. Liu, Y. Liu, Y. Liu, Z. Liu, E. Lu, H. Lu, L. Lu, Y. Luo, S. Ma, X. Ma, Y. Ma, S. Mao, J. Mei, X. Men, Y. Miao, S. Pan, Y. Peng, R. Qin, Z. Qin, B. Qu, Z. Shang, L. Shi, S. Shi, F. Song, J. Su, Z. Su, L. Sui, X. Sun, F. Sung, Y. Tai, H. Tang, J. Tao, Q. Teng, C. Tian, C. Wang, D. Wang, F. Wang, H. Wang, H. Wang, J. Wang, J. Wang, J. Wang, S. Wang, S. Wang, S. Wang, X. Wang, Y. Wang, Y. Wang, Y. Wang, Y. Wang, Y. Wang, Z. Wang, Z. Wang, Z. Wang, Z. Wang, C. Wei, Q. Wei, H. Wu, W. Wu, X. Wu, Y. Wu, C. Xiao, J. Xie, X. Xie, W. Xiong, B. Xu, J. Xu, L. H. Xu, L. Xu, S. Xu, W. Xu, X. Xu, Y. Xu, Z. Xu, J. Xu, J. Xu, J. Yan, Y. Yan, H. Yang, X. Yang, Y. Yang, Y. Yang, Z. Yang, Z. Yang, Z. Yang, H. Yao, X. Yao, W. Ye, Z. Ye, B. Yin, L. Yu, E. Yuan, H. Yuan, M. Yuan, S. Yuan, H. Zhan, D. Zhang, H. Zhang, W. Zhang, X. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Y. Zhang, Z. Zhang, H. Zhao, Y. Zhao, Z. Zhao, H. Zheng, S. Zheng, L. Zhong, J. Zhou, X. Zhou, Z. Zhou, J. Zhu, Z. Zhu, W. Zhuang, and X. Zu (2025) Kimi k2: open agentic intelligence. External Links: 2507.20534, Link Cited by: Appendix E, §1, 2nd item, §2.3.1. [22] S. Lakić (1998) On the computation of the matrix k-th root. Zeitschrift für Angewandte Mathematik und Mechanik 78 (3), p. 167–172. External Links: Link Cited by: §3.1. [23] Z. Li, L. Liu, C. Liang, W. Chen, and T. Zhao (2026) NorMuon: making muon more efficient and scalable. In Forty-third International Conference on Machine Learning, External Links: Link Cited by: §2.1. [24] J. Lim, S. Lee, D. Kim, T. Kim, E. Park, J. Lee, J. Lee, J. Lee, W. T. Cheung, D. Choi, J. Her, J. Huh, H. Jung, C. Kang, B. Kim, M. Kim, T. Kim, Y. Kim, H. Kweon, H. Lee, K. Lee, D. Oh, Y. Park, B. Ryu, and D. Weon (2025) Motif 2 12.7B technical report. External Links: 2511.07464, Link Cited by: 2nd item. [25] T. Lin (2025) Flash-muon: an efficient implementation of muon optimizer. External Links: Link Cited by: §3.1. [26] J. Liu, J. Su, X. Yao, Z. Jiang, G. Lai, Y. Du, Y. Qin, W. Xu, E. Lu, J. Yan, Y. Chen, H. Zheng, Y. Liu, S. Liu, B. Yin, W. He, H. Zhu, Y. Wang, J. Wang, M. Dong, Z. Zhang, Y. Kang, H. Zhang, X. Xu, Y. Zhang, Y. Wu, X. Zhou, and Z. Yang (2025) Muon is scalable for llm training. External Links: 2502.16982, Link Cited by: §A.1, §2.3.1. [27] J. Liu (2025) A proof of concept for Distributed Muon. GitHub. Note: Pull Request #1428, NVIDIA/Megatron-LMAccessed: 2026-07-10 External Links: Link Cited by: item 1. [28] A. Lozhkov, L. Ben Allal, L. von Werra, and T. Wolf (2024) FineWeb-edu: the finest collection of educational content. Hugging Face. External Links: Link, Document Cited by: §A.1. [29] I. Modoranu, M. Safaryan, E. Schultheis, M. Ryabinin, A. Chumachenko, and D. Alistarh (2026) Trion: FFT-based dynamic subspace selection for low-rank adaptive optimization of LLMs. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §3.2. [30] L. Newhouse, D. Goldberg, and R. Ruiz (2024) Faster symmetric matrix multiplication with ThunderKittens. External Links: Link Cited by: §3.1. [31] OpenAI, S. Agarwal, L. Ahmad, J. Ai, S. Altman, A. Applebaum, E. Arbus, R. K. Arora, Y. Bai, B. Baker, H. Bao, B. Barak, A. Bennett, T. Bertao, N. Brett, E. Brevdo, G. Brockman, S. Bubeck, C. Chang, K. Chen, M. Chen, E. Cheung, A. Clark, D. Cook, M. Dukhan, C. Dvorak, K. Fives, V. Fomenko, T. Garipov, K. Georgiev, M. Glaese, T. Gogineni, A. Goucher, L. Gross, K. G. Guzman, J. Hallman, J. Hehir, J. Heidecke, A. Helyar, H. Hu, R. Huet, J. Huh, S. Jain, Z. Johnson, C. Koch, I. Kofman, D. Kundel, J. Kwon, V. Kyrylov, E. Y. Le, G. Leclerc, J. P. Lennon, S. Lessans, M. Lezcano-Casado, Y. Li, Z. Li, J. Lin, J. Liss, L. (. Liu, J. Liu, K. Lu, C. Lu, Z. Martinovic, L. McCallum, J. McGrath, S. McKinney, A. McLaughlin, S. Mei, S. Mostovoy, T. Mu, G. Myles, A. Neitz, A. Nichol, J. Pachocki, A. Paino, D. Palmie, A. Pantuliano, G. Parascandolo, J. Park, L. Pathak, C. Paz, L. Peran, D. Pimenov, M. Pokrass, E. Proehl, H. Qiu, G. Raila, F. Raso, H. Ren, K. Richardson, D. Robinson, B. Rotsted, H. Salman, S. Sanjeev, M. Schwarzer, D. Sculley, H. Sikchi, K. Simon, K. Singhal, Y. Song, D. Stuckey, Z. Sun, P. Tillet, S. Toizer, F. Tsimpourlas, N. Vyas, E. Wallace, X. Wang, M. Wang, O. Watkins, K. Weil, A. Wendling, K. Whinnery, C. Whitney, H. Wong, L. Yang, Y. Yang, M. Yasunaga, K. Ying, W. Zaremba, W. Zhan, C. Zhang, B. Zhang, E. Zhang, and S. Zhao (2025) Gpt-oss-120b & gpt-oss-20b model card. External Links: 2508.10925, Link Cited by: 2nd item. [32] Z. Tang, T. Xu, Y. Saad, and Y. Xi (2026) Hierarchical muon: tiled newton-schulz updates for efficient muon optimization. External Links: 2606.27216, Link Cited by: §3.2. [33] B. Xu (2025) Transpose one of the MLP matrices + add Triton kernel for symmetric matmul. GitHub. Note: Pull Request #109, KellerJordan/modded-nanogptAccessed: 2026-07-10 External Links: Link Cited by: §3.1. [34] A. Yang, A. Li, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Gao, C. Huang, C. Lv, C. Zheng, D. Liu, F. Zhou, F. Huang, F. Hu, H. Ge, H. Wei, H. Lin, J. Tang, J. Yang, J. Tu, J. Zhang, J. Yang, J. Yang, J. Zhou, J. Zhou, J. Lin, K. Dang, K. Bao, K. Yang, L. Yu, L. Deng, M. Li, M. Xue, M. Li, P. Zhang, P. Wang, Q. Zhu, R. Men, R. Gao, S. Liu, S. Luo, T. Li, T. Tang, W. Yin, X. Ren, X. Wang, X. Zhang, X. Ren, Y. Fan, Y. Su, Y. Zhang, Y. Zhang, Y. Wan, Y. Liu, Z. Wang, Z. Cui, Z. Zhang, Z. Zhou, and Z. Qiu (2025) Qwen3 technical report. External Links: 2505.09388, Link Cited by: §A.1, 2nd item. [35] Y. Zhao, A. Gu, R. Varma, L. Luo, C. Huang, M. Xu, L. Wright, H. Shojanazeri, M. Ott, S. Shleifer, et al. (2023) PyTorch FSDP: experiences on scaling fully sharded data parallel. Proceedings of the VLDB Endowment 16 (12), p. 3848–3860. Cited by: §7. Appendix A Alternative Experiments on Gram Newton-Schulz This section presents additional experimental results about Gram Newton-Schulz and the symmetric GEMM kernels. Complementing Section˜8, these experiments study a wider range of architectures but at smaller scale, using only one GPU at a time. A.1 Setup We train four architectures: Llama-430M, Qwen-600M, Gemma-1B, and a custom MoE architecture with 1B parameters, of which ∼20% 20\% are active [14, 34, 12]. We train on FineWeb-Edu [28]. The number of training tokens for each dense model is given by the Chinchilla scaling law; for MoE-1B it is twice the Chinchilla scaling law with respect to the active parameters. We use a cosine learning rate scheduler with the base learning rates given in Table˜3. We use Muon on most matrix parameters—the attention layer’s projection matrices (Q,K,V W_Q, W_K, W_V), the projection following attention (O W_O), the SwiGLU MLP weights (MLP,UP W_MLP,UP, MLP,GATE W_MLP,GATE, MLP,DOWN W_MLP,DOWN), and the token router of the MoE (router W_router)—but exclude the embedding and unembedding layers. Although we are not in the distributed setting, we use megabatching (Section˜7), as batching Newton-Schulz’s GEMM operations makes them faster. As is standard for Muon, we adjust the learning rate by scaling the update for each weight matrix based on its dimensions. For these experiments, we found that Moonshot AI’s strategy of scaling by 0.2max(fan_out,fan_in)0.2 (fan\_out,fan\_in) yields the best loss curves [26, §2.2]. Table 3: Architecture, base learning rate, and training-token budget of each model. Model Hidden dim Layers MLP dim Learning rate Training tokens Llama-430M 10241024 2424 40964096 3×10−33× 10^-3 1010B Qwen-600M 10241024 2828 30723072 1.5×10−31.5× 10^-3 1212B Gemma-1B 20482048 88 1638416384 3×10−43× 10^-4 2222B MoE-1B 10241024 99 (per-expert) 256256 2.5×10−32.5× 10^-3 1111B A.1.1 Splitting the Weights We draw special attention to the fact that we split MLP,UP W_MLP,UP from MLP,GATE W_MLP,GATE and orthogonalize them separately. Ordinary implementations of SwiGLU MLPs concatenate these into a single weight matrix; however, their contributions to the activation are fundamentally different, so gradients are different too. We find that orthogonalizing them separately improves the final loss; for example, in Llama-430M, we observe an improvement of ≈0.2≈ 0.2 in perplexity. For MoE architectures—whose intermediate size is typically smaller than the hidden size—separating them also reduces the FLOP cost of orthogonalization by a factor of two for standard Newton-Schulz and even more for Gram Newton-Schulz. Likewise, while earlier implementations of Muon orthogonalized the combined matrix [Q|K|V] bmatrix W_Q\,|\, W_K\,|\, W_V bmatrix, we orthogonalize each piece separately. We are also aware that in some settings, including pretraining of GLM-5, Muon benefits from splitting the Multi-Latent Attention weights (UQ W^UQ, UK W^UK, and UV W^UV) by attention head before orthogonalizing [13]. This choice is principled, since the actual function computed by attention treats each head separately; it does not see the concatenated weights. Inspired by this, we experimented with splitting Q W_Q, K W_K, V W_V, and O W_O by attention head to form H matrices each of size dH×d dH× d, where d is the embedding dimension and H is the number of heads. However, this design led to higher losses throughout training, so we did not adopt it. Still, we believe that there are other settings like GLM-5 where this strategy works well. Such cases would benefit immensely from Gram Newton-Schulz, since the aspect ratio of these weight matrices would be the number of heads H. For a standard attention weight like Q W_Q with H=16H=16 and T=5T=5, Gram Newton-Schulz on the little matrices would use ×80× fewer FLOPs than orthogonalizing the big matrix! A.2 Kernelized Gram Newton-Schulz preserves quality We first verify that switching from standard Newton-Schulz to our kernelized implementation of Gram Newton-Schulz does not affect training quality. We try each implementation with both the Polar Express coefficients [3] and those of [8]. As Figure˜7 shows, the loss curves are identical in all cases, and the final validation perplexity is preserved to within 0.010.01. Figure˜7 uses a Hopper GPU, but we got the same results on Blackwell. We see loss preserved as follows, when both using the Polar Express coefficients and the coefficients derived by [8]: Figure 7: When training with Muon on a Hopper GPU, switching from standard Newton-Schulz to Gram Newton-Schulz preserves the validation perplexity throughout training (up to a 0.01 difference in GNS’s favor). Setup is described in Section˜A.1. A.3 Kernelized Gram Newton-Schulz speeds up the optimizer step Newton-Schulz Performance We observe that our method speeds up the runtime of the Newton-Schulz step by 1.51.5–2×2×. Figure˜8 reports these speed-ups for each model, benchmarked on both H100 and B300 GPUs. As expected, savings are greatest for weights that are highly rectangular, like Gemma’s MLP weights (α=8α=8) and MoE-1B’s expert weights (α=4α=4). Note that these experiments use standard Newton-Schulz as the fallback when m=nm=n. Figure 8: Newton-Schulz time per model weight for (1) standard Newton-Schulz implemented in pure PyTorch, (2) standard Newton-Schulz with our symmetric GEMM kernels, and (3) Gram Newton-Schulz with our kernels. We test four architectures on a single Hopper (top) or Blackwell (bottom) GPU. Square weights (Q W_Q or K W_K) benefit from the kernels. Rectangular weights—especially those with a high aspect ratio (e.g. Up/Gate in Gemma-1B)—additionally benefit from Gram Newton-Schulz. End-to-End Optimizer Performance Figure˜9 shows the end-to-end wall clock time of the Muon optimizer step for each version of Newton-Schulz, along with that of AdamW. For Muon, these timings include updating the momentum matrix, learning rate scaling, applying the weight update, and the AdamW updates for weights not assigned to Muon (such as the embedding layer and the vector-valued weights). Our method yields a 1.31.3–2×2× speedup, with both Gram Newton-Schulz and the kernels contributing significantly. As before, the impact of Gram Newton-Schulz is greatest for architectures with highly rectangular weights. Due to its MLP’s higher 8×8× aspect ratio, Gemma-1B sees the largest speedup. Figure 9: Optimizer step time for AdamW and Muon with three different Newton-Schulz routines on a single H100. Gram Newton-Schulz with our symmetric kernels gives a 1.31.3–2×2× speedup over standard Newton-Schulz. Timings include matrix splitting and recombination for QKV and MLP, learning rate scaling, weight updates, and the scalar optimizer (AdamW) step for non-2D weights. Estimating Gram Newton-Schulz time in Kimi K2 Kimi K2 is a trillion parameter sparse, fine-grained MoE model with 384384 experts per layer, a hidden size of 71687168, and a small expert intermediate dimension of 20482048. Since models are trending towards finer-grained MoE architectures and Kimi K2 was trained with Muon, this is a perfect setting to benchmark Gram Newton-Schulz. Due to Kimi’s sophisticated pipeline parallel strategy, the optimizer steps of many of its weights are completely hidden behind the backward pass of the next pipeline stage. We estimate that orthogonalization steps of following weights are exposed: 216216 expert up/gate/down weights of shape 2048×71682048× 7168 and 11 dense up/gate/down weight of shape 7168×184327168× 18432. Therefore, we estimate the exposed Newton-Schulz time by orthogonalizing weights of these shapes on a single GPU. Figure˜10 shows the results: our method yields a 2×2× speedup. As before, both our kernels and the Gram Newton-Schulz algorithm contribute significantly. Figure 10: Estimated exposed Newton-Schulz time for one step of Kimi K2 with pipeline parallelism. We measure the runtime of the exposed operations on a single Hopper (top) or Blackwell (bottom) GPU. Gram Newton-Schulz with kernels is 2×2× faster than standard Newton-Schulz. Appendix B Stability of Gram Newton-Schulz B.1 Instability of Naive Gram Newton-Schulz We trained Llama-430M with Muon using Naive Gram Newton-Schulz (Algorithm˜2). The loss curve is shown in Figure˜11. Clearly, training is very unstable. Not only do we observe loss spikes, but eventually, the outputs of Gram Newton-Schulz become full of Infs. The problem is due to floating point arithmetic. While Gram Newton-Schulz is mathematically equivalent to standard Newton-Schulz in exact arithmetic, it behaves differently in finite precision, especially in half precision. As half precision is essential for good performance, this makes Naive Gram Newton-Schulz unworkable in practice. We now analyze the source of the instability in order to motivate our solution. A Jupyter notebook reproducing the experiments in this section is available here. Figure 11: Naive Gram Newton-Schulz used to train Llama-430M in half precision. Numerical instability wrecks training. B.1.1 Tracking Eigenvalues of Intermediate Matrices To understand how the matrices in Naive Gram Newton-Schulz evolve and why they diverge, we track their eigenvalues and singular values. Recall that the entries of any matrix are upper bounded by its largest singular value, so if we can bound the singular values, we can prevent blowups. As a baseline, we start by running Naive Gram Newton-Schulz in full float64 precision for 88 steps to simulate its behavior in exact arithmetic. We use a synthetic 128×512128× 512 matrix with an exponentially decaying spectrum. To make our plots more readable, the experiments in this section use the coefficients (at,bt,ct)=(158,−108,38)(a_t,b_t,c_t)=( 158,- 108, 38) at every iteration, but our conclusions will generalize to the coefficients used in practice. If 0=⊤ X_0= U V is the SVD of the input matrix, then the intermediate matrices of Algorithm˜2 (t R_t, t Q_t, t Z_t) are square symmetric with eigenvectors U. We therefore plot the diagonal entries of ⊤t U R_t U and ⊤t U Q_t U against the corresponding singular values in to track how each evolves according to the polynomial update rules—or diverges from them. Even though Gram Newton-Schulz does not need to compute 1,…,T−1 X_1,…, X_T-1, we do so here for demonstration using the formula t=t0 X_t= Q_t X_0 and plot ⊤t U X_t V against . The top panel of Figure˜12 shows that the eigenvalues evolve as predicted by Theorem˜2. Initially, x0∈[0,1]x_0∈[0,1], r0=x02∈[0,1]r_0=x_0^2∈[0,1] and q0=1q_0=1. As the algorithm progresses, rt→1r_t→ 1, qt=xt/x0→1/x0q_t=x_t/x_0→ 1/x_0, and xt→1x_t→ 1. For x0≈1x_0≈ 1, convergence is faster; for x0≪1x_0 1, it is slower. Figure 12: Evolution of eigenvalues of t R_t, t Q_t, and t X_t in Naive Gram Newton-Schulz with (at,bt,ct)=(15/8,10/8,3/8)(a_t,b_t,c_t)=( 158, 108, 38) as a function of the corresponding singular value of 0 X_0. Top: float64. Bottom: bfloat16. Now we rerun the experiment in bfloat16 arithmetic (Figure˜12, bottom). The first few iterations look correct, but by step 7, we see noisy and unexpected behavior (rt<0r_t<0, xt>1x_t>1). By step 8, the spectrum diverges completely, and by step 10 the algorithm returns Infs. We now describe two key causes of divergence: spurious negative eigenvalues in the Gram matrix ⊤ X X , and eigenvector drift. B.1.2 Spurious Negative Eigenvalues The leading cause of divergence is the presence of negative eigenvalues in the Gram matrix due to half-precision arithmetic. These negative eigenvalues blow up after too many iterations of Gram Newton-Schulz. By construction, rt=xt2≥0r_t=x_t^2≥ 0, so t R_t should be positive semidefinite (Theorem˜2), but Figure˜12 shows that is has negative eigenvalues. In fact, even 0 R_0 has tiny negative eigenvalues introduced in the first matrix multiplication 00⊤ X_0 X_0 . Because 0 X_0 has many singular values that are nearly zero (as is the case in Muon), 0=00⊤ R_0= X_0 X_0 has many eigenvalues that are numerically equal to zero. In bfloat16, even a slightly negative number can be numerically equal to zero. Later computations can introduce additional negative eigenvalues into t R_t as well. These eigenvalues represent nothing about the original problem; they are just an artifact of floating point arithmetic. Therefore, we call them “spurious eigenvalues”. These spurious negative eigenvalues start small, but Figure˜12 shows that their magnitude grows quickly. Recall the update rule: rt=rt−1zt2=rt−1ht(rt−1)2r_t=r_t-1z_t^2=r_t-1h_t(r_t-1)^2 (4) For our choice of coefficients, ht(x)=158−108x+38x2h_t(x)= 158- 108x+ 38x^2. Figure˜13 plots this update rule. As it shows, rt<(158)2rt−1r_t< ( 158 )^2r_t-1. Thus, if any rt<0r_t<0, the spurious negative eigenvalues grow exponentially, diverging to −∞-∞ ! This sets off a chain reaction that causes t Q_t and t X_t to diverge as well. This problem cannot be fixed by choosing different polynomials; while the main loop of Gram Newton-Schulz approximates the inverse square root of r0>0r_0>0, it diverges for negative inputs. Figure 13: When t R_t evolves according to (4), negative eigenvalues diverge to −∞-∞. To prove that the tiny spurious negative eigenvalues of 0 R_0 suffice to cause divergence, we rerun the experiment in float64 precision, but cast 0 R_0 from float64 to bfloat16 and then back to float64 to induce a small floating point error. As Figure˜14 shows, this is enough to cause a blowup. Figure 14: Evolution of eigenvalues of t R_t, t Q_t, and t X_t when all operations use float64 except 0=00⊤ R_0= X_0 X_0 , which uses bfloat16. B.1.3 Eigenvector Drift Spurious negative eigenvalues are not the only source of numerical instability. If we take the input matrix 0 X_0 to have no small singular values (i.e., all ≥0.017≥ 0.017), then we do not observe any negative eigenvalues in t R_t, but t X_t still fails to converge. The culprit seems to be “eigenvector drift”. In exact arithmetic, the eigenvectors of all intermediate matrices match U, the left singular vectors of 0 X_0, but in finite precision they do not. We demonstrate this effect by measuring how far ⊤t U R_t U, ⊤t U Q_t U, and ⊤t U X_t V are from being diagonal matrices. That is, we take the Frobenius norm of the off-diagonal entries as a fraction of the total Frobenius norm. Figure˜15 shows that after several iterations, the eigenvectors / singular vectors of t Q_t and t X_t have drifted significantly from those of 0 X_0. At the same time, the eigenvalues of t Q_t (and by extension, those of t X_t) diverge from where they would be in exact arithmetic. The growing eigenvalues of t Q_t seem to spill into one another. The strength of this effect is less consistent than that of negative eigenvalues, but it is still harmful. Figure 15: As the eigenvectors drift, the spectral norms of t R_t, t Q_t, and t X_t diverge. B.2 Stabilizing Gram Newton-Schulz by Restarting As we have seen, if we run Gram Newton-Schulz for more than a few iterations, the spurious negative eigenvalues of t R_t diverge to negative infinity and t Q_t blows up. Our solution is simple: run Gram Newton-Schulz for only a few iterations. For instance, rather than using Gram Newton-Schulz to compute T X_T directly, we can use it to compute 5 X_5 in a stable manner. While 5 X_5 is not a good approximation to polar(0)polar( X_0), it is closer than where we started. Now we apply Gram Newton-Schulz a second time on the input 5 X_5 to compute 10 X_10 stably. This process can be repeated to reach any desired T X_T. This restarting technique sacrifices some of the performance gains of Gram Newton-Schulz, but it still offers a significant speedup over standard Newton-Schulz. Figure˜16 repeats the experiment from above with a restart every five iterations. While t R_t develops some negative eigenvalues, unlike before, the growth of these eigenvalues is controlled. Each time we restart, we re-initialize t=tt⊤ R_t= X_t X_t , eliminating any large negative eigenvalues. By restarting regularly, we prevent any from growing too large. As expected, t Q_t resets to the identity at iterations 5,10,15,205,10,15,20, and 2525. Therefore, the eigenvalues of t Q_t never grow beyond ≈12≈ 12, despite the negative eigenvalues in t R_t. Since the eigenvalues of t Q_t remain controlled, those of t=tt−5 X_t= Q_t X_t-5 stay at or below 11. Figure 16: Restarting prevents the divergence of t R_t in half-precision. Restarting also helps control eigenvector drift. We repeat the experiment from Figure˜15 on the same matrix (with all singular values >0.017>0.017) with a restart after step 55. Diagonalization error always remains ≤0.05≤ 0.05, and the maximum eigenvalues now closely track their predicted values. Note that we always measure eigenvector drift relative to the original input 0 X_0, not the restarted 5 X_5. Figure 17: Restarting curbs eigenvector drift. B.2.1 When to Restart: Polar Express Coefficients for Muon We now derive the optimal restarting schedule for a given total number of iterations T. To prevent t=t0 X_t= Q_t X_0 from blowing up, we must control the condition number of t Q_t, even when 0 R_0 has spurious negative eigenvalues. (Because qt→1/x0≥1q_t→ 1/x_0≥ 1, this also controls the maximum eigenvalue of t Q_t.) The growth of t Q_t in turn depends on the size of the spurious negative eigenvalues and the specific sequence of polynomials defined by (at,bt,ct)t=1T\(a_t,b_t,c_t)\_t=1^T. Fixing a desired number of restarts, we sweep over all possible restart schedules and pick the one that minimizes the maximum condition number of t Q_t across tts. For the application to Muon, we analyze five iterations of the Polar Express coefficients [3]. In the experiments of the previous section, we observe that the most negative spurious eigenvalue of 0 R_0 is about −4⋅10−4-4· 10^-4. Therefore, we simulate Gram Newton-Schulz in full precision with Polar Express coefficients and track how the eigenvalues of t R_t and t Q_t evolve when 0 R_0 has eigenvalues in the range [−4⋅10−4,1][-4· 10^-4,1]. The left panels of Figure˜18 show that without restarting, they blow up. We then repeat the simulation with one restart, sweeping all possible choices of when to restart. Every time we restart and form =⊤ R= X X , we subtract 4⋅10−44· 10^-4 I to simulate potentially dangerous corruption of the eigenvalues due to floating point error. As the right panel shows, restarting after the second iteration provides the best bounds, ensuring that the eigenvalues of t R_t stay well above −0.4-0.4 and those of t Q_t stay below ≈100≈ 100 for all iterations. Figure˜19 shows that for Gram Newton-Schulz with Polar Express coefficients and one restart after the second iteration, all eigenvalues of t X_t converge stably to 11 . Figure 18: Left: Min/max eigenvalue of t R_t and t Q_t without restarts. 0 R_0 starts with a negative eigenvalue at −4⋅10−4-4· 10^-4. Right: Minimum eigenvalue of t R_t and condition number of t Q_t if restart is placed after iteration 11, 22, 33, or 44. Restarting after iteration 22 best controls the size of t Q_t. Figure 19: Gram Newton-Schulz with Polar Express coefficients and a restart after 22 iterations converges stably. B.2.2 Further Precautions While restarting greatly improves stability, it is not absolutely foolproof. A second or third restart may be required if running for more than five iterations or using a numerically sensitive set of coefficients. Moreover, the usual numerical precautions for standard Newton-Schulz still apply. Safety Factors Most choices of Newton-Schulz polynomials are designed to converge only when the input lies in [0,1][0,1]; any singular values larger than 11 may diverge rapidly, as Figure˜20 shows. Even when 0 X_0 is properly normalized, singular values greater than 11 can arise due to numerical error. This problem affects standard Newton-Schulz too, so the Polar Express polynomials are typically adjusted according to the formula p~t(x)=pt(x/1.02) p_t(x)=p_t(x/1.02). This ensures convergence even for singular values as large as 1.021.02. When using Gram Newton-Schulz, roundoff errors like this can worsen due to computations like ⊤ X X , which do not have such numerical buffers; however, we have never seen this happen when using our recommended setup (float16 arithmetic with restarting after 22 iterations). It is wise to be conservative in the choice of safety factor, for instance, by replacing 1.021.02 with 1.051.05. Figure 20: Theoretical behavior of both standard and Gram Newton-Schulz (Polar Express coefficients) when 0 X_0 has singular values slightly above one. Float16 vs BFloat16 in Newton-Schulz In addition, we argue for using float16 instead of bfloat16 to implement Newton-Schulz. Both use 1616 bits and thus run with the same performance; however, the distribution of the 1616 bits across the mantissa and exponent bits differs. Compared to bfloat16, float16 represents values from a narrower range, but it has greater precision within that range. For Newton-Schulz, however, the range of float16 (roughly 6.1⋅10−56.1· 10^-5 to 6.5⋅1046.5· 10^4) suffices because the magnitudes of its intermediate matrices are not too large. On certain test matrices, we see more accurate polar()polar( X) approximations with float16, but in practice, we have not found a case where the training quality is meaningfully different between float16 and bfloat16. Still, we default to float16. High Accuracy Setting In the application to Muon, we do not need to compute polar()polar( X) to high accuracy, and Section˜A.2 shows that Muon with Gram Newton-Schulz yields effectively identical results to Muon with standard Newton-Schulz in terms of training quality. However, when high accuracy is desired, the usual warnings about forming the Gram matrix apply. Since forming ⊤ X X immediately squares the condition number, Gram Newton-Schulz may not be appropriate in these cases. B.2.3 Computing Matrix Quadratics A key step in Gram Newton-Schulz is computing the matrix quadratic t←at+btt−1+ctt−12 Z_t← a_t I+b_t R_t-1+c_t R_t-1^2. Standard Newton-Schulz implicitly computes matrix quadratics too, but PyTorch implementations typically avoid adding the identity explicitly. Rather, they compute (at+bt+ct2) X(a_t I+b_t A+c_t A^2) in two steps, each with a single GEMM: 1. ←bt+ct2 B← b_t A+c_t A^2 2. ←at+ X← a_t X+ B X Our symmetric GEMM kernel can compute matrix quadratics in one launch, fusing the addition of ata_t I by adding ata_t to all diagonal entries of the output when they are at the register level. This optimization completely obviates any I/O operations needed for the ata_t I addition, typically outspeeding gemm_symmetric(A, B, C, c_t, b_t) + a_t * I, which would require loading I from general memory to shared memory to registers. Once t Z_t is assembled, Gram Newton-Schulz can perform its three subsequent multiplications without the need for further fused additions. However, our tests show that adding ata_t I explicitly can be less stable than distributing it into those later three multiplications. If we stress-test our method by ignoring some of our own advice—restarting after three iterations instead of two, using a Polar Express safety factor of 1.021.02 instead of 1.051.05, and computing the quadratic with ata_t I explicitly—we observe instability. Interestingly, this instability disappears if we use non-symmetric GEMMs (either from PyTorch or Quack) instead of our symmetric kernels; however, if we force symmetry after calling standard PyTorch GEMMs, we see instability again. We conclude that fusing +at+a_t I into a symmetric GEMM is numerically unfavorable; this is not just a kernel bug. We believe this effect can be explained as follows. While the fused kernel computes at+btt−1+ctt−12a_t I+b_t R_t-1+c_t R_t-1^2 in float32 arithmetic under the hood, the result t Z_t is rounded back down to float16 at the end of the GEMM. Future computations like tt Q_t Z_t suffer from this loss of precision in ata_t. In contrast, if the ata_t I term is handled implicitly, all arithmetic involving ata_t takes place in float32. Therefore, it is more stable to compute att+t(btt−1+ctt−12)a_t Q_t+ Q_t (b_t R_t-1+c_t R_t-1^2 ) than t(at+btt−1+ctt−12) Q_t (a_t I+b_t R_t-1+c_t R_t-1^2 ). We reiterate that in all our experiments, this instability can be avoided entirely by restarting correctly or using a safety factor of 1.051.05. Out of an abundance of caution, we rearrange Naive Gram Newton-Schulz (Algorithm˜2) to avoid adding ata_t I explicitly. We change 1. t←at+btt−1+ctt−12 Z_t← a_t I+b_t R_t-1+c_t R_t-1^2 // Apply ht(t−1)h_t( R_t-1) 2. t←t−1t Q_t← Q_t-1 Z_t 3. t←tt−1t R_t← Z_t R_t-1 Z_t to 1. t←btt−1+ctt−12 Z_t← b_t R_t-1+c_t R_t-1^2 2. t←t−1t+att−1 Q_t← Q_t-1 Z_t+a_t Q_t-1 3. ()t←t−1t+att−1(RZ)_t← R_t-1 Z_t+a_t R_t-1 4. t←t()t+at()t R_t← Z_t(RZ)_t+a_t(RZ)_t This change fixes all collected training examples in which symmetric GEMMs were less stable than non-symmetric GEMMs. Appendix C Kernel Implementation Details C.1 Symmetric GEMM Kernel Details We implement all of our symmetric GEMM kernels with square cluster work tiles. Hopper uses cluster size (2,1)(2,1) and thread block tile size (128,256)(128,256), and Blackwell uses cluster size (2,1)(2,1) and 2-CTA collaboration, in which the 2 thread blocks in the cluster collaborate on the same big (256,256)(256,256) tile. Notably, highly optimized custom GEMM kernels on Hopper typically use Ping Pong Scheduling, in which the MMA of tile i and the epilogue of tile i−1i-1 are overlapped in two consumer warp groups.555https://pytorch.org/blog/cutlass-ping-pong-gemm-kernel/ However, Ping Pong Scheduling uses more registers at once, and (128,256)(128,256) is too large of a tile size for Ping Pong Scheduling, leading to register spillage. This is much slower than standard single producer warp, single consumer warp scheduling. Thus, our Hopper symmetric kernels do not use Ping Pong Scheduling. Blackwell GEMM kernels have no explicit conception of Ping Pong Scheduling, since by default in cuBLAS and most kernel libraries, two accumulators are kept in the new tensor memory hierarchy, and MMA is computed on one accumulator while the epilogue is computed on the other. As a small implementation detail, note that the main diagonal of 256×256256× 256 cluster work tiles is part of the work assigned by the triangular scheduler. Since their transposed locations are identical to their current locations, we only write those values to general memory once—writing twice can cause inaccurate values or NaNs. C.2 Implementation Strategy in Code There are only two differences between the symmetric GEMM kernel and the standard GEMM kernel: the triangular scheduler and the transposed tile write in the epilogue. We design our symmetric kernel such that it is abstracted around the standard GEMM kernel to enable lightweight but maximally performant GEMM epilogue fusions. Using these abstractions, we are able to implement the symmetric GEMM kernel for both Hopper and Blackwell in just 160 lines, while achieving state-of-the-art performance. We override the standard tile scheduler with our triangular scheduler and wrap the symmetric GEMM class around the GEMM with activation class. GEMM with activation itself is a wrapper around the Blackwell and Hopper default GEMMs. It supports writing two output tensors—the standard GEMM output (the preactivation) and the standard GEMM output with an activation function such as SwiGLU or ReLU applied (the postactivation). We define the activation function to be the identity and the postactivation tensor to be the inplace transpose of the preactivation tensor. Then, when the GEMM with activation class writes to the postactivation, it is really writing to the upper triangle with a transposed layout—this is exactly the intent of the symmetric GEMM kernel. We override the epilogue of GEMM with activation just to ensure we do not write twice to the diagonal tiles, for the correctness reasons mentioned previously. C.3 Kernel Optimizations for Standard Newton-Schulz Using just optimized CuteDSL kernels, we can accelerate standard Newton-Schulz with two changes. Symmetric Matrix Multiplication As discussed above, the matrices =⊤ A= X X and =bt+ct2 B=b_t A+c_t A^2 computed at each iteration of Newton-Schulz are symmetric by definition. Therefore, we use our symmetric GEMM kernels for these operations, reducing their FLOP cost by half. Fused GEMM + Add The typical way to implement the non-symmetric multiplication ←at+ X← a_t X+ B X is to use torch.baddbmm, which calls cuBLAS under the hood. However, we offer a much faster implementation of this “Fused GEMM + Add” operation for Hopper. Unlike cuBLAS, our “Fused GEMM + Add” supports Ping Pong Scheduling for Hopper, which better hides the epilogue addition of ata_t X. Appendix D Additional Experiments D.1 Architecture and Optimization Setup We now describe details of the experiments in Section˜8.1. (The benchmarks in Section˜8.2 are similar, but with minor architectural differences, synthetic data, and GH200s.) Architecture. The models are decoder-only dense transformers (not mixtures of experts). They use grouped-query attention, sliding-window attention (window 20482048), rotary position embeddings, and RMSNorm, at sequence length 81928192 over a vocabulary of 100,352100,352 tokens. Weights are stored and updated in MXFP8. Table˜4 lists the per-scale dimensions. Each run uses a single node of eight B200s. Table 4: Model dimensions by scale. All share sequence length 81928192, vocabulary 100,352100,352, sliding window 20482048, grouped-query attention, and rotary position embeddings. Scale Hidden dim Layers Heads KV heads MLP hidden dim 1B 15361536 2424 1616 88 66566656 3B 23042304 3232 1818 66 98569856 4B 25602560 3636 3232 88 1075210752 7B 40964096 3232 3232 88 1100811008 14B 51205120 4040 4040 88 1459214592 Optimization. Matrix parameters are updated by the orthogonalizing optimizer under study—Muon, NorMuon, or their filtered variants. For matrix parameters, we multiply the base learning rate by dout/din d_out/d_in. Embeddings, the language-model head, norm scales, and biases are updated by AdamW (β1=0.9 _1=0.9, β2=0.95 _2=0.95) at the base learning rate, without the dout/din d_out/d_in scaling applied to the matrix parameters (the language-model head, in particular, also uses the base rate). We apply weight decay 0.010.01 to the matrix parameters only; the AdamW (non-matrix) groups use no weight decay. The sequence length is 81928192. The learning-rate schedule is warmup–stable–cooldown: a 200200-step linear warmup from 2×10−32× 10^-3 to the peak learning rate, a constant phase, and a cooldown over the final 25%25\% of training back down to 2×10−32× 10^-3. All reported held-out cross-entropies are taken at the end of cooldown. D.2 Finer-Grained Timing Metrics The main text reports GPU (device) time, the metric of interest for the orthogonalization compute. Here we report three per-step metrics side by side: • GPU (device) time, read from CUDA events between two stream markers around optimizer.step(); • CPU (host) time, from time.perf_counter() stopped immediately after step() returns and before any device sync — cost of kernel launch, Python, and any host-side syncs the optimizer performs; • Communication volume, the number of bytes moved by the megabatch collectives (all_to_all, all_gather, reduce_scatter) at each step, measured by tallying the receive-side payload per-GPU. As always, Newton-Schulz runs in half precision arithmetic. Results for a 14B model sharded over eight B200s and a 7B model on four GH200s appear in Table˜5 and Table˜6, respectively. Unlike in Section˜8.2, we show all four combinations of kernel and Newton-Schulz variant. We see that Gram Newton-Schulz without symmetric kernels achieves up to a 1.2×1.2× speedup. Otherwise, GPU timings are consistent with the results presented in Section˜8.2. The NorMuon family has higher costs overall due to the overhead of its extra normalization steps. Our contributions still yield similar speedups: up to 1.5×1.5× for symmetric kernels and Gram Newton-Schulz, 3.3×3.3× for f=1/2f= 12, and 5.6×5.6× for f=1/4f= 14. Table˜5 shows that the CPU cost is negligible—just ∼0.1 0.1 ms across configurations. This near-constant host cost is a direct consequence of CUDA-graph capture. Row selection, error feedback, and Gram Newton-Schulz issue many small operations and kernel launches. Without capture, the host would dispatch each one individually, and this per-launch overhead would dominate, especially for the filtered configurations, whose device work is small. Capture and replay instead collapse the entire step into a single graph launch, so the host issues one call regardless of how many kernels the step contains or how much they compute. The reported host time is therefore just the fixed cost of that one launch (∼0.1 0.1 ms), essentially independent of the configuration. Finally, communication volume is directly proportional to the fraction f, as expected. Table 5: Fine-grained per-step metrics for the distributed 14B / 8-GPU benchmark: GPU (device) time from CUDA events, CPU (host) time from perf_counter stopped before the device sync, and per-GPU communication volume of one step. We distinguish the GPU time between the Muon and NorMuon families of optimizers; CPU time and communication volume are reported once because neither depends on the family: host time is the fixed CUDA-graph launch cost, and communication volume is set by the parameter shapes and the fraction f. GPU (ms) Configuration Kernels Algorithm NS frac. Muon NorMuon CPU (ms) Comm (MB) Standard Newton-Schulz PyTorch Standard 11 153.8 174.1 0.1 5479 + symmetric kernels CuteDSL Standard 11 132.7 153.4 0.1 5479 + Gram NS PyTorch Gram 11 137.6 158.5 0.1 5479 + kernels + Gram NS CuteDSL Gram 11 106.8 127.0 0.1 5479 + filtering (f=1/2f= 12) CuteDSL Gram 1/2 12 66.4 73.8 0.1 2740 + filtering (f=1/4f= 14) CuteDSL Gram 1/4 14 34.8 38.1 0.1 1370 Table 6: Replication of Table˜5 for a 7B-parameter model sharded over a node of four GH200s with NVLink. GPU (ms) Configuration Kernels Algorithm NS frac. Muon NorMuon Comm (MB) Standard Newton-Schulz PyTorch Standard 11 360.6 385.4 6442 + symmetric kernels CuteDSL Standard 11 274.4 315.0 6442 + Gram NS PyTorch Gram 11 296.6 321.1 6442 + kernels + Gram NS CuteDSL Gram 11 217.7 252.1 6442 + filtering (f=1/2f= 12) CuteDSL Gram 1/2 12 101.5 116.7 3221 + filtering (f=1/4f= 14) CuteDSL Gram 1/4 14 59.6 68.6 1610 D.3 Benchmarking all-to-all communications In Section˜7, we claim that reducing the number of communication rounds via megabatching is beneficial in two ways: (1) each round adds a fixed overhead not dependent on the size of the payload, and (2) small messages fail to saturate the interconnect. Here, we verify both claims directly with a standalone microbenchmark of NCCL all_to_all over NVLink, independent of any optimizer. On a single node of four H100s (SXM, NVLink-4 mesh, NCCL 2.29), we sweep the size of the payload per rank from 11 KiB to 11 GiB. We time each all-to-all with CUDA events (10 warmup, 50 timed iterations), taking the per-iteration time to be the maximum over ranks, and measure the bandwidth. We confirmed that the transport is a pure NVLink peer-to-peer communication, with no PCIe or network fallback and that the measured bandwidths agree with the nccl-tests alltoall_perf tool to within 1%1\%. Figure˜21 shows the results. Both effects discussed in Section˜7 are clearly visible. The latency floor is ∼25μ 25\, (left panel), of which ∼15 15–17μ17\, is host-side dispatch (as measured separately with perf_counter and no device synchronization). Without megabatching, we pay this cost over and over. Bandwidth climbs by more than an order of magnitude as the per-link payload grows (right panel), so coalescing many small all-to-alls into one megabatch increases the effective bandwidth of each transfer. Figure 21: NCCL all_to_all over NVLink on 4×4×H100. Left: median all-to-all time versus payload per rank (sender) over 50 trials. For small payloads, runtime is dominated by a fixed latency. Right: bus bandwidth versus payload per link (sender → receiver). Bandwidth is near zero for small messages and does not reach 8080–90%90\% of its peak until the per-link payload is ∼16 16–3232 MiB. At world_size=4 world\_size=4, a 256256 KiB per-link payload achieves only ∼10% 10\% of the peak bandwidth. D.4 Ablations Tuning the baseline The experiments in Section˜8.1 compare Dion3 to a NorMuon baseline. Table˜7, below, demonstrates that this baseline is properly tuned for a fair comparison. Table 7: Tuning hyperparameters (learning rate η, momentum μ) of NorMuon at 1B. Final validation loss on 100B tokens of ClimbMix. Winner: η=0.01η=0.01, μ=0.9μ=0.9 (bold). Sweep η (at μ=0.9μ=0.9) η CE 0.0050.005 2.2202.220 0.010.01 2.1942.194 0.020.02 2.2152.215 0.030.03 2.2372.237 Sweep μ (at η=0.01η=0.01) μ CE 0.90.9 2.1942.194 0.9490.949 2.2012.201 0.9740.974 2.2172.217 0.9870.987 2.2382.238 Dion3 with no compression As a control, we run Dion3 with f=1f=1, which should reduce to plain NorMuon. Figure˜22 confirms this at 1B; the two curves coincide throughout training, and their final validation losses differ by 0.00050.0005. Figure 22: Validation loss at 1B for NorMuon and for Dion3 at f≈1f≈ 1, on 100B tokens of ClimbMix. As expected, Dion3 tracks NorMuon throughout, ending at 2.21412.2141 versus 2.21462.2146. Timings with NorMuon family Section˜8.2 compares the runtime of Muon to the Muon version of Dion3. Here, in Figure˜23, we repeat this benchmark with the NorMuon family. Results are quite alike, but due to the extra overhead of NorMuon’s normalization steps, it benefits slightly less from our contributions, which target the orthogonalization step. Figure 23: Optimizer step time (excluding forward/backward pass) relative to standard Muon across model scales when training on 1 GH200 (left) and FSDP over 4 GH200s (right). Each colored line adds one of our contributions on top of the previous one; AdamW is shown for reference. Lines show median over 25 steps and bands show interquartile range. Compare to Figure˜6. Appendix E Case Studies of End-to-End Training Time The share of end-to-end training time taken up by Newton-Schulz can vary widely depending on the training setup. To explain this variability, we analyze two idealized scenarios. In one, standard Newton-Schulz takes 2% of training time; in the other it takes 17%. Case Study 1: Standard Newton-Schulz takes 2% of Kimi K2 training time The following analysis gives a very optimistic estimate of the optimizer’s wall clock time. We assume an efficient training infrastructure with highly optimized pipeline parallelism. Moreover, we assume that the optimizer step of each pipeline stage is completely hidden behind the backward pass of the next pipeline stage. Kimi K2 Thinking is a 1.11.1 trillion parameter model with 3232 billion active parameters. It has 11 dense layer followed by 6060 MoE layers [21]. It is pretrained with 256256-GPU model parallel groups, 1616-way pipeline parallelism, 1616-way expert parallelism within each pipeline stage, and a huge batch size of 6767 million tokens. We use a single H100 to approximate the share of each training step’s runtime occupied by Newton-Schulz in this setting under the following assumptions: 1. The training cluster consists of 256256 nodes of eight H100s each (20482048 GPUs in total), connected with NDR 400 Gb/s InfiniBand inter-node (8 NICs per node, 1:1 NIC-to-GPU ratio) and NVLink 4.0 intra-node. This is the size of the cluster used to train DeepSeek-V3, with upgraded hardware [9]. 2. Training in bfloat16 attains 40%40\% model flop utilization (MFU), which is typical for MoEs at this scale on H100s. 3. The only non-overlapped optimizer time is that of the last pipeline stage to complete its backward pass (i.e., pipeline stage 11 of 1616). The optimizer steps of pipeline stages 22 to 1616 are fully hidden behind the backwards of stages 11 to 1515. 4. Pipeline stage 11 contains the dense layer and 33 MoE layers. Under these assumptions, the optimal way to partition the Newton-Schulz work of pipeline stage 11 is as follows. Each of the 1616 GPUs in pipeline stage 11’s expert parallel group gets 384 experts/layer×3 MoE layers16 GPUs=72 experts/GPU=216 expert up-gate-down/GPU. 384 experts/layer× 3 MoE layers16 GPUs=72 experts/GPU=216 expert up-gate-down/GPU. Each of the 1616 GPUs has its own unique expert weights, so no communication is needed. Pipeline stage 11 also contains the dense MLP’s three 7168×184327168× 18432 weights (up/gate/down) and three shared experts. Orthogonalizing the dense MLP weights is the dominant cost, so they are sent to three different GPUs; the shared experts are split amongst the remaining 1313 GPUs. Thus, the orthogonalization time for pipeline stage 11 is the time it takes for a single GPU to orthogonalize 216216 expert up/gate/down weights and 11 dense up/gate/down weight. As benchmarked in Figure˜10, standard Newton-Schulz implemented in PyTorch takes 315315 ms to do this. Per our assumption, stage 11’s orthogonalization step is the only one that is not overlapped. Having estimated the Newton-Schulz time, we now estimate the end-to-end wall clock time of an entire Kimi K2 global training step. We use the standard estimate of 6NB6NB FLOPs for the forward and backward pass, where N is the number of active parameters and B is the global batch size. Given • Active parameters: N=32⋅109N=32· 10^9 • H100 peak: P=989⋅1012P=989· 10^12 FLOP/s • Model flop utilization: MFU=40%MFU=40\% • Cluster size: G=2048G=2048 GPUs • Global batch size: B=67⋅106B=67· 10^6 tokens we have sec/batch=6N×BP×MFU×G=15.9sec/batch= 6N× BP×MFU× G=15.9 Thus, Newton-Schulz takes approximately 315 ms15900 ms+315 ms=1.9% 315 ms15900 ms+315 ms=1.9\% of total pretraining wall clock time in this setting. Case Study 2: Standard Newton-Schulz takes 17% of Llama3-70B SFT time Llama3-70B is an 80-layer dense model with hidden size 81928192, intermediate size 2867228672, and grouped query attention with K,V W_K, W_V of size 1024×81921024× 8192 and Q,O W_Q, W_O of size 8192×81928192× 8192 [14]. Supervised finetuning (SFT) typically uses small batch sizes, ranging from 3232 to 256256 sequences [9]. We analyze the following SFT case: 1. Training uses 3232 H100s across 44 nodes (8 GPUs per node). 2. Training in bfloat16 hits 40%40\% MFU. 3. Weights are sharded evenly across GPUs using FSDP, and the exposed Newton-Schulz time is that of 80 layers/32 GPUs≈3 layers80 layers/32 GPUs≈ 3 layers. Each layer has 33 MLP weights (up, gate, down), and the attention weights Q W_Q, K W_K, V W_V, and O W_O. According to our benchmarking, standard Newton-Schulz of • nine 8192×286728192× 28672 weights takes 739739 ms, • six 8192×81928192× 8192 weights takes 156156 ms, • six 1024×81921024× 8192 weights takes 2.322.32 ms, totaling 897897 ms. Given • Parameters: N=70⋅109N=70· 10^9 • H100 peak: P=989⋅1012P=989· 10^12 FLOP/s • Model flop utilization: MFU=40%MFU=40\% • Cluster size: G=32G=32 GPUs • Global batch size: B=64 sequences×2048 tokens/sequence=131,072B=64 sequences× 2048 tokens/sequence=131,072 tokens we have sec/batch=6N×BP×MFU×G=4.35 ssec/batch= 6N× BP×MFU× G=4.35 s Thus, Newton-Schulz takes approximately 897 ms4350 ms+897 ms=17% 897 ms4350 ms+897 ms=17\% of total SFT wall clock time in this setting.