Dion3: Full-Stack Orthogonal Updates

Microsoft Research (senior author John Langford) with Princeton (Tri Dao, Jack Zhang, Berlin Chen), NVIDIA, NYU and Yale attack the reason Muon is expensive at scale: its cubic-time Newton-Schulz orthogonalization, whose compute and FSDP communication overhead the paper argues Kimi only escaped through a lucky alignment of flat-buffer sharding, fine-grained MoE (small matrices), and PP+EP parallelism. Dion3 is a full-stack revision with four stacked contributions: Gram Newton-Schulz (iterate on the small symmetric Gram matrix XXᵀ instead of X — mathematically identical output, two rectangular matmuls total), symmetric GEMM kernels in CuTeDSL for Hopper/Blackwell that compute only the lower triangle (~half the FLOPs, beating cuBLAS across sizes), a fractional update rule that orthogonalizes only a selected fraction of momentum rows per step with error feedback on the rest (reducing exactly to Muon at f=1; superseding the unpublished "Dion2"), and megabatching, which cuts FSDP all-to-all rounds per optimizer step to a small constant.

Quality is preserved or slightly improved: 1B dense transformers (GQA + sliding window, MXFP8, FSDP) on 100B ClimbMix tokens reach slightly better validation loss than a tuned NorMuon baseline, with larger-scale runs up to 14B and a 7B case study on four GH200s. Speed: symmetric kernels + Gram Newton-Schulz give ≥1.5× over standard Muon, and with f=0.25 row selection the stack reaches a 6× optimizer-step speedup for larger models (measured via CUDA events on GH200, single-GPU and 4-GPU FSDP), while matching or improving Muon's loss. Ships as the MIT-licensed pip-installable dion package (microsoft/dion, ~540 stars) with FSDP2/DDP/mixed-sharding backends that also implement Muon, NorMuon and standard Newton-Schulz — a drop-in replacement aimed at making orthogonal optimizers practical beyond the architectures that happen to suit them.