DMuon: Efficient Distributed Muon Training with Near-Adam Overhead
For practitioners using matrix-orthogonalization-based optimizers like Muon, DMuon eliminates the major performance bottleneck, making these optimizers practical for large-scale distributed training.
DMuon reduces the overhead of Muon optimizer from >2x to near-AdamW levels, achieving 1.48x-3.01x speedup in end-to-end step time and 6.85x-163.00x speedup in optimizer-step time across embodied and LLM training workloads.
Matrix-orthogonalization-based optimizers, exemplified by Muon, have demonstrated strong convergence behavior across a wide range of modern deep learning workloads. The matrix-aware updates offer a compelling alternative to conventional element-wise optimization, particularly as model architectures continue to grow in scale and heterogeneity. Yet contemporary distributed training infrastructure built around the assumption of element-wise optimizers is poorly matched to matrix-level optimizers such as Muon, whose updates couple entire weight matrices and require costly Newton-Schulz iterations. Vanilla Muon implementations incur more than 2x the cost of forward and backward passes. To close this gap, we present DMuon, an open-source distributed Muon implementation that integrates into existing training pipelines as a drop-in module, with no framework-level modifications. Across both embodied foundation model and large language model (LLM) training workloads, DMuon achieves a 1.48x-3.01x speedup in end-to-end step time and a 6.85x-163.00x speedup in optimizer-step time, bringing per-step latency to near-AdamW levels and enabling efficient scaling in our model training.