PaperlayerView exact revision on arXiv

Source-linked AI summary

Dion3: Full-Stack Orthogonal Updates

Noah Amsel, Jack Zhang, Kwangjun Ahn, Ali Naeimi, Austin Feng, Berlin Chen, Tri Dao, John Langford

arXiv:2608.11612v1cs.LGcs.AI

TL;DR

Muon’s cubic-time orthogonalization and distributed communication costs make scaling difficult. Dion3 addresses these costs with algorithmic, kernel, update-rule, and communication changes, achieving substantial speedups while preserving or improving training quality in the reported setting.

  • Problem

    Muon’s cubic-time orthogonalization and additional distributed communication create scaling challenges that limit its flexibility as a general-purpose optimizer.

  • Method

    Dion3 combines Gram Newton-Schulz, symmetry-exploiting CuteDSL kernels, fractional momentum-row updates, and megabatched communication to reduce Muon’s optimization costs.

  • Results

    Dion3’s Gram Newton-Schulz and CuteDSL components speed up Muon by 1.5× for dense models and 2× for MoEs, while the fractional update rule with f = 1/4 adds a 3.7× speedup and improves training quality in the reported setting.

  • Takeaways & Limitations

    Dion3 provides practical open-source packages that make orthogonal optimizers more accessible across large-scale and distributed settings.

  • Takeaways & Limitations

    Gram Newton-Schulz can diverge when half-precision rounding introduces spurious negative Gram-matrix eigenvalues, and the fractional update rule’s quality improvement may not generalize widely.

Abstract

from arXiv · show

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 (https://github.com/microsoft/dion) as a drop-in replacement for Muon.

1 Introduction

Dion3 addresses Muon’s computational and communication overhead with a full-stack optimizer that combines four complementary improvements. Together, these changes target orthogonalization cost across algorithms, kernels, update rules, and distributed execution.

  • Dion3 combines four improvements that reduce Muon’s orthogonalization cost and compound when used together.The components are Gram Newton-Schulz, symmetric CuteDSL kernels, fractional momentum subsampling, and megabatched communication.
  • Gram Newton-Schulz reformulates orthogonalization around a small symmetric Gram matrix, reducing FLOPs while preserving mathematical equivalence.
  • CuteDSL kernels exploit symmetric matrix multiplication to accelerate Gram Newton-Schulz.
  • The fractional update rule subsamples momentum rows or columns before orthogonalization, making Dion3 simpler and faster than Dion while matching or improving optimization quality.
  • Megabatched communication reduces optimizer-step communication rounds to a small constant, while the package supports distributed settings with little adoption overhead.

2 The Challenge of Scaling Muon

Muon’s polar-update approximation reduces training steps but is costly to compute and difficult to distribute. Its Newton-Schulz routine has super-linear matrix complexity, while sharding introduces additional communication and scaling constraints.

  • Muon’s update: Muon applies polar decomposition to momentum matrices, balancing the update spectrum and giving it full numerical rank.
  • Newton-Schulz: Newton-Schulz approximates polar(M) iteratively with matrix polynomials, preserving singular vectors while transforming singular values toward one.
  • Newton-Schulz: When T = 5, standard Newton-Schulz costs (20α + 10)n^3 FLOPs across 15 matrix-matrix multiplications.
  • Scalability: Orthogonalization scales as O(n^3), making Muon’s optimizer step increasingly expensive relative to linear-cost optimizers such as SGD and Adam.
  • Scalability: Sharded training adds all-to-all communication to assemble and scatter momentum matrices, creating a major obstacle to scaling Muon.
  • Scalability: Muon’s successful scaling has depended on a subtle alignment of architecture, parallelism, and framework factors, motivating more flexible scaling properties.

3 Comparison with Related Work

Dion3 differs from prior Muon accelerations by changing both the orthogonalization computation and the optimizer input reduction. Its row or column subsampling simplifies compressed updates while remaining compatible with Newton-Schulz improvements.

  • Most Muon variants improve step efficiency but retain Newton-Schulz and therefore generally retain its scaling challenges.
  • Gram Newton-Schulz departs from standard Newton-Schulz while remaining mathematically identical and compatible with nearly all Muon variants.
  • Compared with related Gram-based work, Dion3 adds different formulas, symmetric kernels, stability analysis, and practical recommendations for deployment.
  • Dion3 selects k rows or columns of the momentum matrix instead of constructing a low-rank approximation, preserving runtime savings while simplifying implementation and distributed handling.
  • Unlike low-rank and block orthogonalization methods, Dion3’s compressed update can be combined directly with improvements to Newton-Schulz.

4 Gram Newton-Schulz

Gram Newton-Schulz computes the same orthogonalization result as standard Newton-Schulz by iterating on the smaller symmetric Gram matrix. Its FLOP savings are substantial for asymmetric matrices, but numerical instability requires restarting.

  • Gram Newton-Schulz iterates on XX⊤ rather than directly on X, producing an output mathematically identical to standard Newton-Schulz.
  • The algorithm uses small n × n symmetric matrices for most runtime and requires only two rectangular multiplications: XX⊤ initially and Q_T X finally.
  • Gram Newton-Schulz rewrites odd Newton-Schulz polynomials as xh(x^2), enabling equivalent iteration through inverse-square-root approximations of the Gram matrix.
  • Runtime: For T = 5 and α = 43, Gram Newton-Schulz saves 55% of FLOPs versus symmetric-GEMM Newton-Schulz and 68% versus a typical nonsymmetric implementation.
  • Stability: The naive algorithm can diverge because half-precision rounding introduces spurious negative eigenvalues into XX⊤, where inverse-square-root iteration is unstable.
  • Stability: Restarting after two iterations resets spurious negative eigenvalues near zero and preserves training quality, at an added cost of 3(α−1)n^3 FLOPs.
  • Implementation: The training-ready implementation combines restarting, polynomial reformulation, float16 arithmetic, and Polar Express coefficients.

5 Symmetric GEMM Kernels in CuteDSL

Dion3’s symmetric GEMM kernels exploit output symmetry by computing only the lower triangle and copying its transpose, reducing work while accelerating Gram Newton-Schulz.

  • Symmetric GEMM design: The kernels save about half the floating-point operations used by standard matrix multiplication by exploiting symmetric outputs.They support AB and αAB + βC when AB and C are symmetric.
  • Symmetric GEMM design: Symmetric GEMM computes lower-triangle and diagonal 256 × 256 tiles, then transposes and copies each lower tile to the upper triangle.The triangular scheduler assigns only lower-triangle tiles, while the epilogue writes computed lower-triangle values to transposed upper-tile locations.
  • Performance: For large enough n, CuteDSL kernels achieve a ∼2× speedup over cuBLAS on Hopper and Blackwell, with or without epilogue addition of C.The comparison uses n × n input matrices A, B, and C.
  • Implementation: The symmetric kernels differ from standard GEMM in their schedulers and epilogues, while retaining the usual tiled computation pipeline.Standard GEMM loads operands, performs matrix-multiply accumulation, fuses epilogue operations, and writes the output.

6 The Dion3 Update Rule

Dion3 accelerates Muon by orthogonalizing only selected momentum rows, applying the update to those rows, and using error feedback to preserve ignored information across iterations.

  • Update rule: Dion3 selects a fraction of momentum rows, orthogonalizes the selected submatrix with Gram Newton-Schulz, and updates only the selected weight rows.The remaining weight rows are not updated in that step.
  • Selection: The compression factor f ∈(0, 1] controls selection; f = 1 recovers Muon, while f = 1/4 or f = 1/8 is recommended.Rows or columns are selected according to sharding, or the smaller dimension when unsharded.
  • Selection: Dion3 selects the k = ⌈fn⌉ rows with largest ℓ1 norm, while random selection produced optimization quality that was not much worse in initial experiments.Distributed selection normally chooses the top-f fraction independently on each shard to avoid extra synchronization and ragged all-to-all communication.
  • Efficiency: Selection reduces Newton-Schulz multiplication cost by at least 1/f 2 and reduces distributed communication volume by 1/f.Gram Newton-Schulz adds greater benefit when the aspect ratio α = m/n is large.
  • Error feedback: Error feedback decays selected momentum rows by µ while leaving unselected rows unchanged, allowing ignored rows to accumulate and become selected later.When M = c M, the approximation is exact and standard momentum is recovered.
  • Validation: At f = 1, Dion3’s convergence curve matches Muon or NorMuon almost exactly despite implementation differences.Those differences include row permutation, damping placement, float32 NorMuon normalization, and a custom Triton update kernel.
  • Implementation: CUDA graph capture and replay offsets the extra kernel-launch overhead introduced by row selection on smaller-scale models.The repeated dependency structure makes the optimizer suitable for graph replay.

7 Megabatching and Communication

Dion3 addresses distributed Muon overhead by assembling only selected momentum rows and grouping same-shaped matrices into megabatches, reducing communication volume and rounds.

  • Distributed communication: Under FSDP, Muon requires an all-to-all to assemble each momentum matrix and another to scatter the orthogonalized result.Dion3’s dion package manages these operations for FSDP2, DDP, and mixed sharding strategies.
  • Megabatching: Megabatching groups all same-shaped matrices into one batch, packing their shards into one all-to-all and processing them together.Because Transformers have only a handful of distinct weight shapes, communication rounds become O(1), independent of model depth.
  • Benchmark: 35% lower optimizer step time is achieved by megabatching for the 1B model with 32 shards.The benchmark measures Muon’s median per-GPU optimizer step time on fixed models, data, optimizers, and GPUs.
  • Benchmark: Megabatching helps most when weights are small, Newton-Schulz is cheap, and each rank holds several matrices; its impact is limited when Newton-Schulz dominates computation.Table 1 compares shard-count-sized batching with one megabatch per shape group on 1B and 14B models using 8 or 32 GPUs.
  • Distributed communication: Dion3 communicates only the selected submatrix M[S, :], reducing communication relative to synchronizing the entire momentum matrix.When rows can be selected without synchronizing M, the update is identical at a fraction of the communication cost.

8 Experiments

Experiments show that Dion3 preserves or improves model quality while substantially reducing optimizer-step cost across model sizes and parallelism settings. Fractional updates require learning-rate adjustment, and their runtime and communication benefits grow with scale.

  • 8.1 Model Quality Is Preserved: Dion3 with f = 1 matches Muon/NorMuon almost exactly, while fractional updates with f < 1 slightly improve loss.The authors report that subselection does no harm to loss and can improve it when tuned correctly.
  • 8.1 Model Quality Is Preserved: The lowest 1B-model loss occurs at f = 1/8, with Dion3 variants maintaining a clear validation-loss advantage over NorMuon throughout training.Figure 4 reports the Dion3 variants finishing about 0.01 loss points below the fully tuned NorMuon baseline.
  • 8.1 Model Quality Is Preserved: At 3B–14B parameters, Dion3 outperforms NorMuon in validation loss at every scale, with the largest improvement of −0.027 at 14B.Dion3 also wins downstream accuracy at three of four scales, including a 0.7-percentage-point improvement at 14B.
  • 8.1 Model Quality Is Preserved: Optimal learning rates follow η√f ≈ 0.01 as the row-selection fraction f decreases.The effective step size shrinks when only a fraction of momentum rows is updated, motivating η′ = η/√f.
  • 8.2 Dion3 Accelerates the Optimizer: Symmetric kernels and Gram Newton-Schulz provide a combined speedup of 1.5× or more over standard Muon, while fractional updates cut runtime further at larger scales.The fractional update has increasing impact because Newton-Schulz’s cubic cost increasingly dominates other operations.
  • 8.2 Dion3 Accelerates the Optimizer: Fractional updates reduce communication by a factor of 1/f, although the Figure 6 setting is compute-bound and exposes neither communication nor host time.Dion3 remains slower than AdamW, but the gap is significantly smaller.
  • 8.2 Dion3 Accelerates the Optimizer: On higher-aspect-ratio architectures such as the evaluated MoE models, Gram Newton-Schulz and symmetric kernels alone achieve a 2× speedup.The appendix evaluates a wider range of architectures, including Gemma and mixtures-of-experts.

9 Conclusion

Dion3 addresses Muon’s computational and distributed overhead with a full-stack optimizer design. The combined system preserves or improves loss while achieving large step-time reductions across dense and MoE settings.

  • 9 Conclusion: Muon’s cubic-time Newton-Schulz orthogonalization and distributed communication overhead make scaling increasingly costly, motivating a flexible remedy.Optimizer-step time can account for 1% to 17% of total training time in LLMs trained with Muon.
  • 9 Conclusion: Dion3 combines Gram Newton-Schulz, CuteDSL symmetric kernels, fractional updates, and megabatching to reduce orthogonalization and communication costs.The contributions target different stack levels and compound when used together.
  • 9 Conclusion: Gram Newton-Schulz and CuteDSL kernels provide 1.5× speedups for dense models and 2× for mixtures-of-experts, with f = 1/4 adding 3.7×.The authors describe the fractional-update quality improvement as unexpected and requiring further study of generalization.
  • 9 Conclusion: The dion and gram-newton-schulz packages make orthogonal optimizers practical across a wide range of settings.The implementations are provided as interoperable, pip-installable tools for distributed and nondistributed use.
  • 9 Conclusion: Together, Dion3’s contributions achieve up to a 6× optimizer-step speedup for larger models relative to standard Muon.The speed measurements exclude forward and backward passes and cover one GH200 and FSDP over four GH200s.

Appendices

The appendices test Dion3 components across architectures and hardware, showing preserved training quality and architecture-dependent speedups. They also examine weight-splitting choices that affect both quality and orthogonalization cost.

  • A.1 Additional experiments: Appendix experiments cover Llama-430M, Qwen-600M, Gemma-1B, and a 1B-parameter MoE trained on FineWeb-Edu.These experiments use one GPU at a time and include megabatching.
  • A.1.1 Splitting the Weights: Separating SwiGLU up-projection and gate weights before orthogonalization improves final loss, yielding approximately 0.2 lower perplexity in Llama-430M.The separation also reduces orthogonalization FLOPs for MoE architectures with smaller intermediate dimensions.
  • A.1.1 Splitting the Weights: Splitting attention projection weights by head produced higher losses throughout training, so the authors did not adopt that design.The authors note that the strategy may still work in settings such as GLM-5.
  • A.1.1 Splitting the Weights: For H = 16 heads and T = 5, applying Gram Newton-Schulz to separate attention matrices uses 80× fewer FLOPs than orthogonalizing the combined matrix.The potential saving follows from the aspect ratio of the smaller matrices.
  • A.1 Additional experiments: Kernelized Gram Newton-Schulz preserves validation perplexity within 0.01 of standard Newton-Schulz across both tested coefficient sets and on Hopper and Blackwell GPUs.The reported loss curves are identical in the tested comparisons.

A.3 Kernelized Gram Newton-Schulz speeds up the optimizer step

Gram Newton-Schulz with optimized kernels reduces orthogonalization and end-to-end Muon step time, especially for rectangular weights, but naive half-precision execution is unstable.

  • Kernelized Gram Newton-Schulz: 1.5–2× speedups reduce Newton-Schulz runtime, with the largest savings for highly rectangular weights.The gains are greatest for Gemma’s MLP weights and MoE-1B expert weights.
  • Kernelized Gram Newton-Schulz: 1.3–2× speedups reduce end-to-end Muon optimizer-step time when Gram Newton-Schulz and symmetric kernels are used.Timings include momentum updates, learning-rate scaling, weight updates, and AdamW updates for non-2D weights.
  • Stability: Naive Gram Newton-Schulz becomes unstable in half precision, producing loss spikes, divergent spectra, and Infs.The instability arises from finite-precision effects despite mathematical equivalence to standard Newton-Schulz in exact arithmetic.
  • Kernelized Gram Newton-Schulz: Gram Newton-Schulz with kernels is 2× faster than standard Newton-Schulz for exposed Kimi K2 operations.The estimate covers exposed Newton-Schulz time on Hopper and Blackwell GPUs with pipeline parallelism.
  • Stability: Spurious negative Gram-matrix eigenvalues can grow exponentially, driving Rt toward −∞ and causing Qt and Xt to diverge.Eigenvector drift is a separate finite-precision instability when inputs lack small singular values.

B.2 Stabilizing Gram Newton-Schulz by Restarting

Restarting Gram Newton-Schulz limits finite-precision instability by repeatedly resetting the Gram computation, while preserving a significant speed advantage over standard Newton-Schulz.

  • Restarting: Restarting every few iterations controls negative eigenvalue growth and keeps Qt and Xt bounded in half precision.With five-iteration restarts, Qt eigenvalues remain below approximately 12 and Xt eigenvalues stay at or below 1.
  • Restarting: Restarting limits eigenvector drift, keeping diagonalization error at or below 0.05 in the tested setting.The measurement is relative to the original input X0 rather than the restarted matrix.
  • Restart schedule: The restart schedule is selected by minimizing the maximum condition number of Qt across iterations.The analysis sweeps possible schedules for a fixed number of restarts and uses the resulting Qt conditioning to control Xt.
  • Restart schedule: A restart after two iterations best controls Qt for the tested Polar Express coefficients and initial eigenvalue range.The corresponding stable configuration is shown for Gram Newton-Schulz with Polar Express coefficients.
  • Further precautions: High-accuracy polar computation can be unsuitable for Gram Newton-Schulz because forming XX^T squares the condition number.For Muon, Gram Newton-Schulz produces effectively identical training quality to standard Newton-Schulz because high polar accuracy is unnecessary.
  • Further precautions: Fusing +atI into symmetric GEMMs is numerically unfavorable, while distributing the addition into later multiplications avoids the observed instability.The instability appeared under a stress test with three-iteration restarts and a 1.02 safety factor, but disappeared with non-symmetric GEMMs.
  • End-to-end performance: Dion3 timing experiments show up to 1.5× speedup from symmetric kernels and Gram Newton-Schulz, 3.3× at f = 1/2, and 5.6× at f = 1/4.The measurements use half precision on sharded 14B and 7B models; CUDA-graph capture keeps CPU cost near 0.1 ms.

D.3 Benchmarking all-to-all communications

All-to-all communication has a substantial fixed latency and poor bandwidth for small payloads, supporting megabatching as a way to coalesce transfers and improve effective bandwidth.

  • Communication benchmark: ∼25 µs is the measured all-to-all latency floor, including ∼15–17 µs of host-side dispatch.Without megabatching, this fixed cost is paid repeatedly across communication rounds.
  • Communication benchmark: Megabatching increases effective bandwidth by coalescing many small all-to-alls into larger transfers.The benchmark independently verifies both fixed-latency and small-message bandwidth effects on four H100s over NVLink.
  • Communication benchmark: Small per-link payloads achieve poor bandwidth, while bandwidth does not reach 80–90% of peak until payloads reach ∼16–32 MiB.At world_size = 4, a 256 KiB per-link payload reaches only ∼10% of peak bandwidth.
  • Control comparison: Dion3 with f = 1 matches plain NorMuon, with final validation losses differing by 0.0005.The two validation curves coincide throughout training at 1B on 100B tokens of ClimbMix.
  • NorMuon ablation: NorMuon benefits slightly less from Dion3’s contributions because its additional normalization steps add overhead beyond orthogonalization.The timing comparison reports similar overall results but lower benefit for the NorMuon family.

E Case Studies of End-to-End Training Time

The case studies estimate Newton-Schulz’s share of end-to-end training time under large-scale MoE pretraining and dense-model SFT settings. Under the stated assumptions, its share ranges from 1.9% to 17%.

  • Cross-case comparison: Newton-Schulz’s training-time share varies substantially with the training setup, reaching 2% in one idealized scenario and 17% in another.These scenarios motivate evaluating optimizer overhead relative to complete forward/backward training steps.
  • Kimi K2 pretraining: Newton-Schulz work is partitioned across 16 expert-parallel GPUs, with dense MLP weights assigned to three GPUs and shared experts split across the remaining 13.Each GPU receives 216 expert up/gate/down weights, while dense MLP orthogonalization remains the dominant cost.
  • Kimi K2 pretraining: 1.9% of total pretraining wall-clock time is attributed to Newton-Schulz for Kimi K2 under the stated 2048-GPU assumptions.The estimate uses 315 ms of Newton-Schulz time over a 15.9-second forward/backward batch time.
  • Llama3-70B SFT: The Llama3-70B SFT estimate exposes approximately three layers per GPU after evenly sharding 80 layers across 32 GPUs.The exposed work includes three MLP weights and four attention weights per layer.
  • Llama3-70B SFT: 17% of total SFT wall-clock time is attributed to Newton-Schulz for Llama3-70B under 32-GPU FSDP assumptions.The estimate combines 897 ms of Newton-Schulz time with a 4.35-second forward/backward batch time.
Loading 2608.11612v1…