Dion3, a revision of Muon that targets this overhead at every level of the stack with 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.
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 (https://github.com/microsoft/dion) as a drop-in replacement for Muon.
Periodic Row-wise Muon is introduced, which performs a full NS5 spectral update once every \(K\) steps and applies a low compute and communication cost row-wise constrained update based on the current momentum at the remaining steps to preserve Muon's generative quality advantage while translating it into end-to-end training efficiency for large DiTs.
Chenghao Li, Xiao Han, Xin-Xin Huang et al.· 0 citations
FlashBoot is presented, a hardware-friendly, framework-workflow co-designed weight-loading subsystem built on SGLang that accelerates single-node weight loading by up to 50x and concurrent rack-level weight loading by>270x and scales poorly to concurrent multi-node bring-up.
Issac Zhu, Hscos Zhang, Ke Jiang et al.· 0 citations
The limited-memory BFGS (L-BFGS) Hessian update scheme is the critical kernel in many quasi-Newton optimization algorithms. The most common approach to implementing L-BFGS uses $2m$ sequential rank-1 updates as part of solving a linear system when there are $m$ history steps. The performance of this approach suffers when the latency of synchronization is significant, and its poor temporal locality increases the memory traffic when vectors do not fit in cache. The compact dense representation of L-BFGS results in an approach that has minimal synchronization latency and better temporal locality, but it requires an additional pass over the basis vectors and an additional basis that must be recomputed when the $B_0$ matrix changes as in variable-metric methods. In the Portable Extensible Toolkit for Scientific Computation and the Toolkit for Advanced Optimization (PETSc/TAO), we have implemented an intermediate dense formulation of BFGS that retains most of the good characteristics of both the recursive and compact dense approaches. We report single-node performance tests of these implementations on the U.S. Department of Energy's Polaris and Frontier machines, testing both GPU-based and CPU-based computations.
Spectral-Aware Muon is introduced, which holds the head at the Muon scale and amplifies the bulk using a static spectral prior, and both variants outperform tuned AdamW and Muon (Scion implementation) baselines in all evaluated model-scale and batch-size configurations.
Xiaodong Wu, Wenyi Yu, Chao Zhang et al.· 1 citation
BaKron is an efficient solver that combines anti-diagonal parallelism with a recursive divide-and-conquer construction that matches the cubic scaling of GPTQ while exploiting richer curvature information.
In training Mixture-of-Experts (MoE) models, sharded data parallelism partitions each expert's parameters across GPUs, requiring an Allgather operation to reconstruct the full weight matrix before each layer executes. This communication often dominates iteration time. Prior work often reduces this overhead using lossy compression methods that sacrifice numerical fidelity, while existing lossless methods do not exploit cross-iteration exponent stability. In this paper, we propose Cross-Iteration Exponent Reuse Allgather (CIERA), a lossless, system-aware communication method for sharded MoE training. We observe that after a brief warmup phase, the exponent values of most weights remain unchanged across iterations. Based on this, we cache exponents locally and transmit only the sign and mantissa when exponents are unchanged. The receiver reconstructs the original weights exactly by combining the cached exponents with the received data. Since parameter matrices vary in shape across layers, compression is applied only when it yields net time savings. Moreover, the compression operation is overlapped with both Allgather communication and computation. Our real experiments and large-scale trace-driven simulator show that on OLMoE-1B-7B at 16 GPUs, CIERA achieves a 3.70x speedup over the lossless baseline and 3.68x over the lossy baseline, projected to reach 4.28x and 4.42x respectively at 128 GPUs, while preserving bitwise-exact parameter reconstruction in all evaluated runs.
Unknown authors· 0 citations
We use cookies to run the site and, with your consent, for analytics and to show ads.
See our Cookie Policy.
