Search papers, labs, and topics across Lattice.
Dion3 is a refined version of the Muon optimizer that addresses the significant overhead costs associated with its cubic-time Newton-Schulz orthogonalization step, particularly when weights are sharded. By introducing a Gram Newton-Schulz algorithm, CuteDSL kernels, and a megabatching strategy, Dion3 effectively reduces the FLOP cost and communication overhead, while a novel update rule further enhances efficiency by selectively orthogonalizing only a fraction of the momentum matrix's rows. As a result, Dion3 achieves comparable or improved loss performance to Muon while decreasing optimizer step time by up to 6x, making it a highly efficient alternative for practitioners.
Dion3 slashes optimizer step time by up to 6x while maintaining or improving loss performance compared to its predecessor, Muon.
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.
Microsoft Research