Dion3：全端正交更新

摘要

Muon 最佳化器因其三次方時間的 Newton-Schulz 正交化步驟而產生顯著的額外成本。當權重被分片時，通訊開銷更會加劇此計算成本，在許多情況下侵蝕了 Muon 的優勢。我們提出 Dion3，這是 Muon 的修訂版本，旨在堆疊的各個層級解決此開銷。我們的 Gram Newton-Schulz 演算法降低了正交化的 FLOP 成本，我們的 CuteDSL 核心透過利用對稱性來加速運算，而我們的巨型批次（megabatching）策略則減少了通訊開銷。此外，我們對更新規則提出了一項簡單的修改，進一步降低成本：在每一步僅選取動量矩陣的一部分列進行正交化。此更新規則在速度和效能上均優於 Dion（另一種 Muon 的「壓縮」版本）。整體而言，Dion3 在損失上與 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.