让矩阵优化器高效分布式训练，速度逼近AdamW。
DMuon: Efficient Distributed Muon Training with Near-Adam Overhead
- 用矩阵正交化方法改进优化器，避免逐元素更新的瓶颈。
- 在大模型训练中，单步耗时提速1.48倍至3.01倍，优化器步骤快6.85到163倍。
- 无需修改框架，可直接接入现有训练流程，适合大规模模型训练者。
基于矩阵正交化的优化器（如Muon）在多种现代深度学习任务中展现出优异的收敛性能。其矩阵级更新机制为传统逐元素优化提供了有力替代，尤其适用于日益规模化与异构化的模型架构。然而，当前分布式训练系统多基于逐元素优化器设计，难以匹配像Muon这类需耦合整个权重矩阵并执行昂贵Newton-Schulz迭代的矩阵级优化器，导致原生Muon实现的开销超过前向与反向传播的2倍以上。为此，我们提出DMuon——一个开源的分布式Muon实现，可作为即插即用模块集成至现有训练流程，无需框架级修改。在具身基础模型与大型语言模型（LLM）训练任务中，DMuon实现了1.48x–3.01x的端到端步骤加速，以及6.85x–163.00x的优化器步骤加速，使每步延迟接近AdamW水平，支持高效扩展。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。