Dion3：全栈正交更新

摘要

Muon优化器因其三次时间复杂度的Newton-Schulz正交化步骤而产生显著的开销成本。当权重进行分片时，通信开销进一步加剧了这一计算成本，在许多场景下削弱了Muon的优势。我们提出Dion3，这是对Muon的改进版本，从技术栈的每一层入手解决这一开销问题。我们的Gram Newton-Schulz算法降低了正交化的FLOP成本，我们的CuteDSL内核通过利用对称性加速正交化过程，而我们的大规模批处理策略则减少了通信开销。此外，我们对更新规则提出了一项简单改进以进一步削减成本：每步仅选取动量矩阵的一部分行进行正交化。这一更新规则在速度和性能两方面均优于Dion（Muon的另一种"压缩"版本）。总体而言，Dion3在损失（loss）上与Muon相当或更优，同时将优化器步进时间最多降低6倍。Dion3可通过dion包（https://github.com/microsoft/dion）获取，可作为Muon的直接替代品。

English

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.