DMuon: Efficient Distributed Muon Training with Near-Adam Overhead
作者: Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang
分类: cs.DC, cs.LG
发布日期: 2026-06-25
💡 一句话要点
提出DMuon以解决分布式Muon优化器效率问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 分布式训练 优化器 深度学习 矩阵正交化 高效算法 模型训练 开源实现
📋 核心要点
- 现有的分布式训练基础设施主要围绕元素级优化器设计，无法有效支持矩阵级优化器如Muon，导致性能低下。
- DMuon是一个开源的分布式Muon实现，能够无缝集成到现有训练管道中，优化了矩阵级更新的效率。
- 在基础模型和大型语言模型的训练中，DMuon实现了显著的加速，优化器步骤时间提升达到163倍，极大提高了训练效率。
📝 摘要（中文）
基于矩阵正交化的优化器，如Muon，已在现代深度学习任务中展示出强大的收敛性。然而，现有的分布式训练基础设施主要围绕元素级优化器构建，无法有效支持矩阵级优化器如Muon，导致其实现的前向和反向传递成本超过2倍。为了解决这一问题，本文提出了DMuon，一个开源的分布式Muon实现，能够无缝集成到现有训练管道中。DMuon在基础模型和大型语言模型的训练中，实现了1.48倍至3.01倍的端到端步骤时间加速，以及6.85倍至163.00倍的优化器步骤时间加速，使每步延迟接近AdamW水平，从而有效提升模型训练的可扩展性。
🔬 方法详解
问题定义：本文旨在解决现有分布式训练基础设施与矩阵级优化器（如Muon）之间的不匹配问题。传统的元素级优化器假设导致Muon实现的前向和反向传递成本过高，影响训练效率。
核心思路：DMuon通过设计一个开源的分布式实现，能够无缝集成到现有训练管道中，优化了矩阵级更新的过程，减少了对昂贵的Newton-Schulz迭代的依赖。
技术框架：DMuon的整体架构包括数据并行和模型并行的结合，利用矩阵感知更新机制，支持高效的分布式训练。主要模块包括数据分发、梯度聚合和优化器步骤执行。
关键创新：DMuon的核心创新在于其能够在不修改框架的情况下，提升矩阵级优化器的效率，使其在分布式环境中表现出色，显著降低了每步延迟。
关键设计：DMuon在参数设置上进行了优化，采用了适应性的学习率调整策略，并在损失函数中引入了矩阵正交化的约束，以确保收敛性和稳定性。
🖼️ 关键图片
📊 实验亮点
DMuon在基础模型和大型语言模型的训练中，分别实现了1.48倍至3.01倍的端到端步骤时间加速，以及6.85倍至163.00倍的优化器步骤时间加速，显示出其在性能上的显著提升，接近于AdamW的水平。
🎯 应用场景
DMuon的研究成果在深度学习模型训练中具有广泛的应用潜力，尤其是在需要高效分布式训练的场景，如大规模自然语言处理和计算机视觉任务。其高效的优化器步骤时间和可扩展性将推动更复杂模型的训练，提升实际应用的效率和效果。
📄 摘要（原文）
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.