microsoft/dion — distributed Muon, Dion2, Dion, and NorMuon orthonormal optimizers

TL;DR

Microsoft Research’s reference implementation of four orthonormal optimizers — Muon, Dion, Dion2, and NorMuon — built specifically for modern distributed PyTorch (≥2.7, DTensor / FSDP2 / TP). Dion and Dion2 are the lab’s own contributions: they shrink the momentum matrix before the Newton–Schulz orthonormalization step (Dion via power iteration with error feedback, Dion2 via simple submatrix selection), reducing both compute and the communication-intensive full-matrix reconstruction that Muon under FSDP needs. The package ships as pip install git+... with mature parallelism support (DDP/FSDP2/HSDP, TP for Dion only), per-head Newton–Schulz for fused QKV projections, compressed data-parallel gradient sync via low-rank states (PowerSGD-style), Triton kernels for the NS step, and a mixed-precision config for bf16 optimizer states.

Key claims

The package implements four orthonormal optimizers — Muon, Dion, Dion2, NorMuon — under one API; Dion2 is the currently recommended algorithm and Dion is marked legacy [README “Welcome”, “Optimizers” tables].
Dion/Dion2’s communication-efficiency story: orthonormalize a shrunk matrix instead of the full sharded weight, avoiding Muon’s all-gather-to-reconstruct pattern; Dion uses power iteration + error feedback, Dion2 uses submatrix selection + error feedback [README “Introduction”].
Distributed support matrix: single-device, DDP, FSDP2 supported for all four; 2D sharding (FSDP2 + TP) is supported only for Dion (the legacy implementation), not Dion2/Muon/NorMuon [README “Optimizers” table].
Per-parameter-group setup is mandatory and non-trivial: weight matrices use Dion/Muon; biases, norms, embeddings, and the LM head must be split into separate AdamW/Lion groups, with the LM head additionally scaled by lr / sqrt(d_in) — the optimizer cannot detect these from tensor shape alone [README “Building Parameter Groups”, table].
Per-head Newton–Schulz: fused QKV nn.Linear of shape (num_heads*head_dim, in_features) can be orthonormalized per head by passing num_heads=... on the param group, treating the weight as a batch of head matrices internally; sharding dim 0 on head boundaries avoids the all-to-all under FSDP [README “Per-Head Newton-Schulz for Attention Projections”].
Compressed data-parallel gradient sync: in DDP/HSDP, Dion can skip the full-gradient all-reduce and instead all-reduce only its low-rank states, yielding identical model weights with 1/rank_fraction less DP-sync bandwidth (PowerSGD-derived); requires replicate_mesh_grad_sync=True, no_sync() under DDP, and an explicit synchronize_for_checkpoint() call before saving because optimizer states are intentionally decoupled across DP replicas [README “Compressed Data-Parallel Gradient Sync”].
Default LR scaling for matrix params is sqrt(d_out/d_in) for both Dion and Muon — chosen to produce consistent activation-vector change so LR transfers across model size; Muon additionally supports Moonshot AI’s 0.2 * sqrt(max(d_in, d_out)) scaling [README “Best Practices”].
Empirical rank-fraction guidance: at 3B params, rank_fraction=0.25 matches full-rank Dion; the authors expect 1/8 or 1/16 to work at 10B+ [README “Best Practices”].
The Muon NS step ships a custom Triton kernel that exploits the symmetry of A @ A.T to compute only half the output and mirror; enabled via use_triton=True [README “Triton Kernels for Muon Newton-Schulz”].
The repo points to its two backing papers — Dion (arXiv:2504.05295, Ahn/Xu/Abreu/Langford 2025) and Dion2 (arXiv:2512.16928, Ahn/Amsel/Langford 2025) — and explicitly cites Kimi K2 and GLM-4.5 as the production-scale validations of orthonormal optimizers more broadly [README “Introduction”, “Citation”].

Method

The two new algorithms factor Muon into “shrink, then orthonormalize.” Muon orthonormalizes the momentum matrix M ∈ R^{d_out × d_in} directly via Newton–Schulz iteration; under FSDP this requires reconstructing the full M from its shards before each NS step, which is communication-bound. Dion computes a rank-r approximation Q ≈ M via power iteration with an error-feedback buffer that tracks M − Q @ Q.T @ M, then runs NS on the small r × d_in factor — the reconstruction call becomes a much cheaper all-gather of a low-rank matrix. Dion2 simplifies further: it picks an α-fraction submatrix of the momentum, NS-orthonormalizes only that submatrix, and uses error feedback to absorb the residual. The submatrix path lets Dion2 use an all-to-all communication pattern that simultaneously unshards and distributes only the selected fraction across devices.

NorMuon is a third variant that adds neuron-wise normalization on top of Muon for additional stability (separate arXiv: 2510.05491).

Beyond the algorithms, the package’s distinguishing engineering: a DeviceMesh argument (distributed_mesh=...) tells each optimizer how the model is sharded so it can choose between (i) all-to-all unsharding when parameters are FSDP-sharded along the optimizer’s mesh, (ii) work distribution + all-gather when not, and (iii) for Dion only, a separate outer_shard_mesh / inner_shard_mesh split that lets FSDP and TP shards both live inside the orthonormalization step. The compressed-DP-sync mode (replicate_mesh_grad_sync=True) requires the user to keep FSDP’s device mesh 1D (sharded sub-mesh only) so FSDP doesn’t double-reduce gradients that Dion is already synchronizing in compressed form.

Results

The README does not report training-loss numbers — it is a tooling / reference-implementation repo, not a benchmark paper. The empirical claims it relays are second-hand: Muon’s effectiveness was first shown in the modded-NanoGPT speedrun and validated at scale by Kimi K2 (1T-param-class) and GLM-4.5. The repo’s own positioning is communication and compute efficiency: orthonormalizing a low-rank or sub-block factor instead of the full matrix; rank-fraction 0.25 reportedly matches full-rank at 3B; rank-fraction is expected to fall further with scale. Triton-kernel speedup for the symmetric NS step is asserted but not quantified in the README.

Why it’s interesting

This is the missing piece in the wiki’s optimizer thread. The Training stability at scale page already tracks the architecture-vs-optimizer fork — SSO (Controlled LLM Training on Spectral Sphere) argues for spectral-sphere optimizers, The Newton-Muon Optimizer iterates on Muon’s NS recipe, and Addition of Muon optimizer to torch.optim — PyTorch team welcomes PR (Issue #148819) flagged that torch.optim.Muon shipped without FSDP/HSDP support, leaving a real infrastructure gap. microsoft/dion is one of the candidate fillers of that gap: it ships distributed Muon (FSDP2 + DDP today, TP-not-yet) and introduces a low-rank/sub-block variant designed to make the orthonormalization itself cheaper to distribute. Complements Aurora: A Leverage-Aware Optimizer for Rectangular Matrices (leverage-aware step sizing for rectangular matrices) and QK-Clip: Taking Muon Further on the Scaleup Journey (QK-clip for Muon stability) — same lineage, different bottleneck. The compressed-DP-sync trick (PowerSGD reuse for orthonormal optimizers) is the under-discussed angle: it folds gradient-compression into the optimizer state for free at non-trivial rank fractions. The Triton symmetric-NS kernel is the same hardware-level trick Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon pursues from the Gram-matrix side — exploit A @ A.T symmetry to write only half the output — packaged here for the rectangular-Muon path.