🤖 machine learning

DMuon: Efficient Distributed Muon Training with Near-Adam Overhead

该论文介绍了 DMuon，这是一个 Muon 优化器的开源分布式实现，它能够无缝集成到现有的训练流水线中，从而大幅降低优化器步进延迟并达到接近 AdamW 的效率，进而实现基于矩阵正交化的优化在大型基础模型中的规模化应用。

原作者： Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

发布于 2026-06-26

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Vincent Chen, Starrick Liu, Regis Cheng, Dance Yang, Shalfun Li, Ryan Yu, Lucy Liang, Hang Su, Roy Gan, Hao Wang, Qian Wang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在训练一个巨大的机器人大脑（大型语言模型或具身智能）来学习。为了教它，你需要一个“优化器”——一个观察机器人的错误并调整其内部连接（权重）以使其变得更聪明的教练。

多年来，标准的教练一直是 AdamW。它就像一支由成千上万名微型工人组成的团队，每人负责修理机器人上的一个微小螺丝。它快速且高效，但它把每一个螺丝都视为独立的个体。

最近，人们发现了一个更聪明的教练，名叫 Muon。Muon 不再是一个个修理螺丝，而是观察整个机器人的“面板”（完整的数据矩阵）并同时调整它们。这种“团体治疗”的方法让机器人学习得更快，性能也更高。然而，它有一个代价：在分布式系统（由许多计算机组成的集群）上运行起来非常缓慢。

问题：“全员会议”瓶颈
在分布式训练设置中，机器人的大脑被拆分到了许多台计算机（GPU）上。

AdamW 就像一个远程团队：计算机 A 修理它的螺丝，计算机 B 修理它的螺丝，它们之间不需要进行太多交流。
Muon 就像一个委员会：为了修理一个面板，每台计算机都必须先看到整个面板。

在一种朴素的实现方式中，每一台计算机都必须停止手头的工作，下载来自所有人的完整数据，执行复杂的数学运算，然后将结果传回。这就像是在举行一场大规模会议，每个人都要在做任何决定之前先读完一份 1000 页的报告。这耗时太长了，以至于“会议”（优化步骤）所花费的时间比实际工作（学习步骤）还要长！事实上，论文指出这可能会导致 2 倍于 学习本身的时间成本！

解决方案：DMuon（分布式 Muon）
X Square 机器人团队 开发了 DMuon，这是一个让 Muon 运行速度几乎接近 AdamW 的新系统。他们通过改变团队的工作方式，而不是改变教练的数学原理，解决了这个瓶颈。

他们通过以下三个主要技巧解决了瓶颈：

1. “指定专家”策略（以所有者为中心的执行）

DMuon 不再让每个人都尝试修理每个面板，而是为每个面板分配 一台特定的计算机 作为“所有者”（Owner）。

之前： 所有人聚集数据，所有人进行计算，所有人都在浪费时间。
现在： 只有“所有者”计算机负责收集数据并进行繁重的数学运算。其他计算机则处于等待状态或进行其他工作。
类比： 想象一个建筑工地。与其让每个工人尝试建造整栋房子，不如指派一名专家负责屋顶，另一名专家负责管道。其他人不去尝试建造屋顶，他们只是向专家递送所需的工具。这避免了重复劳动。

2. “传送带”系统（通信重叠）

当“所有者”正在进行数学运算时，其他计算机并不是闲置的。DMuon 创建了一个流水线：

前向传播（Forward Pass）： 当机器人在“思考”（处理数据）时，系统会在后台悄悄地将下一组工具发送给专家。
反向传播（Backward Pass）： 当机器人在“从错误中学习”时，系统已经在为专家收集下一批数据了。
类比： 这就像餐厅的厨房。厨师（所有者）正在切菜。与其等待厨师切完菜再让服务员送下一份订单，不如在厨师切菜的同时，服务员就把下一个订单送过来。等待的时间被隐藏在了切菜的时间里。

3. “智能装配线”（批处理与调优）

Muon 进行的数学运算非常复杂。有时处理的部件很大（如一面巨大的墙），有时又很小（如一块小瓷砖）。

问题： 如果你试图用一台巨大的机器去处理一块微小的瓷砖，机器就会处于闲置状态。如果你处理的是一面巨大的墙，则会耗费很长时间。
解决方法： DMuon 将许多微小的瓷砖组合成一个单一的“批次”（Batch），使机器保持忙碌。它还使用了一个“调优器”，能够自动找到处理每种特定形状瓷砖的最快方式。
类比： DMuon 不会为了一个信封就派一辆快递车，而是等待攒够一整车的信封后再统一发送。它还会根据目的地是城市还是农场，来改变卡车的路线。

结果

该论文在真实世界的模型上测试了 DMuon，包括一个机器人训练模型（Wall-OSS）和一个语言模型（Qwen2.5）。

速度： DMuon 使“优化步骤”（即会议）比旧的朴素方式快了 6 到 163 倍。
效率： 现在训练一个步骤的总时间仅比标准的 AdamW 方法慢 2%。
结论： DMuon 让团队能够在不支付巨大时间惩罚的情况下，使用更优越的“团体治疗”式学习风格。它将一个缓慢、笨拙的过程变成了一个流畅、高效的过程，使其能够用于大规模 AI 模型的生产环境。

简而言之，DMuon 是一个交通控制系统，它让最聪明的教练（Muon）能够在计算机高速公路上运行而不会造成交通拥堵，终于使这种方式在训练下一代 AI 时变得切实可行。

技术摘要：DMuon —— 高效分布式 Muon 训练，实现接近 Adam 的开销

1. 问题陈述

矩阵感知优化器（特别是 Muon）通过对权重矩阵的动量聚合梯度应用 Newton–Schulz (NS) 迭代，展现出了优于元素级优化器（如 AdamW）的收敛特性。这种方法将更新的奇异值推向均匀谱，在预训练大型模型（如 LLM 和具身基础模型）时提供了显著的计算效率提升。

然而，在现代分布式训练环境中部署 Muon 面临着关键的系统瓶颈。当代分布式训练栈（如 ZeRO、FSDP）运行在一种隐含的元素级优化契约之下，即更新是在参数分片上独立计算的。Muon 违反了这一契约，因为其核心 NS 迭代要求单个设备上必须存在完整的权重矩阵，以计算必要的矩阵乘积（具体为 $XX^\top$ ）。

在朴素的分布式实现（“先聚合后计算”）中，每个 rank 都必须：

All-gather 所有分片的完整缩减梯度矩阵。
在本地执行 NS 迭代。
丢弃它不拥有的分片结果。

这种方法引入了两个主要成本：

通信开销： 每一步重建完整矩阵都会引入集体通信，其规模随模型大小和分布式宽度增加。
冗余计算： 每个 rank 都会执行相同的昂贵正交化操作，这使得计算成本乘以了 rank 的数量 ( $D$ )。

在实践中，这些成本导致分布式 Muon 的优化器步进时间足以媲美甚至超过前向和反向传播的总壁钟时间，从而削弱了 Muon 的算法优势。对于具身模型训练（如 VLA）而言，这种惩罚尤为严重，因为相比于标准的 LLM 预训练，其前向-反向传播的计算量相对较轻。

2. 方法论：DMuon 系统设计

DMuon 是一种开源的分布式实现，旨在缩小 Muon 与 AdamW 之间的性能差距。它通过一个协同设计的系统架构来实现这一点，该架构集成了通信优化、负载均衡和内核级效率，且无需修改底层训练框架（如 PyTorch/FSDP2）。

2.1 以所有者为中心的执行策略

DMuon 不在所有 rank 之间复制计算，而是为每个矩阵参数分配一个单一的所有者 rank (owner rank)。

所有者角色： 所有者维护权威的优化器状态，并为每个矩阵精确执行一次 Muon 更新（Newton–Schulz 迭代）。
非所有者角色： 其他 rank 仅在进行前向/反向计算时才实例化参数，并且不执行正交化。
通信流：
- 前向： 所有者将更新后的参数异步发布给非所有者，并与层执行过程进行流水线化。
- 反向： 梯度在打包的缓冲区中累积，并被归约（reduce）至所有者。
- 重叠： DMuon 采用两级层次结构（节点内和节点间）以及细粒度的权重布局，以实现通信与计算的重叠，防止集体通信冲突。

2.1 高效 Gram Newton–Schulz 迭代

为了降低所有者侧更新的计算成本，DMuon 采用了 Gram Newton–Schulz 公式：

Gram 空间递归： 迭代不是在原始的 $m \times n$ 空间中进行，而是在 $m \times m$ 的 Gram 空间（ $G = XX^\top$ ）中进行。当 $m < n$ 时，这把复杂度从 $O(m^2n)$ 降低到了 $O(m^3)$ 。
对称性利用： 由于 $G$ 是对称的，DMuon 使用 SYRK 风格的执行路径来仅计算下三角部分，从而使主矩阵乘法的算术工作量减半。
一旦解决 GPU 占用率问题，DMuon 通过将形状相似的矩阵进行批处理，并联合执行 Gram NS 递归，从而分摊内核调度开销。

2.3 计算感知型负载均衡

由于 Muon 更新成本随矩阵形状的变化而显著不同，朴素的分配方式（如轮询）会导致 straggler rank（掉队者）。DMuon 将所有者分配建模为一个混合整数线性规划 (MILP) 问题：

剖析 (Profiling)： 在初始化期间，系统会对各种批次大小和内核调度针对每个不同矩阵形状的执行时间进行基准测试。
优化： MILP 求解器最小化所有者 rank 之间的预测最大完成时间（makespan），确保通过分配最重的矩阵来平衡工作负载。
回退机制： 对于极其庞大的搜索空间，使用贪心启发式算法来限制初始化时间。

2.4 形状自适应内核自动调优

DMuon 利用 DSL 驱动的自动调优工作流（使用 TileLang 和 CUTE DSL）来生成高性能内核。

它针对目标硬件和矩阵形状，在 tile 形状、块大小和软件流水线配置上进行搜索。
最优配置会被缓存并重复使用，以确保内核调度始终接近 Transformer 训练中所遇到的特定问题维度的最优解。

3. 核心贡献

细粒度通信优化： 一个以所有者为中心的运行时，通过层次化调度和细粒度权重布局最小化通信开销，同时保持精确的优化器语义。
形状自适应执行栈： 一个结合了批处理 Gram Newton–Schulz、对称感知内核和自动调优的高性能引擎，以高效处理异构矩阵工作负载。
计算感知型负载均衡： 一种基于实际运行时剖析（而非理论 FLOPs）的测量成本优化方法 (MILP)，用于分配矩阵所有者。
即插即用集成： 一个可以集成到现有训练流水线（如 FSDP2）中的模块，只需极小的努力，无需对框架进行底层修改。
生产级验证： 在机器人（具身）和语言模型训练工作负载上成功部署。

4. 实验结果

作者在四个生产级工作负载上评估了 DMuon：Wall-OSS-0.5、Pi0、Wall-WM 和 Qwen2.5-7B，规模从 8 到 256 块 A800 GPU。

优化器步进时间： 与朴素的“先聚合后计算”分布式 Muon 实现相比，DMuon 在优化器步进时间上实现了 6.85 倍至 163.00 倍的加速。
端到端步进时间： DMuon 在总步进时间上实现了 1.48 倍至 3.01 倍的加速。
相对于 AdamW 的开销： DMuon 的每步延迟被降至接近 AdamW 的水平，在评估的工作负载中，平均开销仅比 AdamW 高出 +2%。
扩展性： 在 Wall-OSS 上，随着 GPU 数量增加，DMuon 保持了接近 AdamW 的吞吐量，而朴素的 Muon 则会遭受严重的通信瓶颈。

加速分解：

对称 Gram 内核： 贡献了 48% 的加速（减少了主乘法中的算术量）。
所有者调度与负载均衡： 贡献了 32%（消除了 $D$ 倍的冗余计算）。
自动调优与批处理： 贡献了 16%（减少了小矩阵的内核启动开销）。

5. 重要性与主张

论文声称 DMuon 成功解决了矩阵感知优化器与分布式训练基础设施之间的系统失配问题。通过消除冗余计算并优化通信，DMuon 使 Muon 的收敛优势（如更高的计算效率和更好的优化几何特性）能够转化为实际的壁钟时间加速。

作者将 DMuon 定位为一种必要的系统赋能工具，而非一种新算法，它使 Muon 成为大规模生产训练中可以作为 AdamW 的即插即用替代品。这项工作证明了可以将矩阵感知优化的“惩罚”降低到可以忽略不计的水平，从而为大型语言模型和具身基础模型 (VLA) 实现高效扩展。

论文明确指出，DMuon 不会改变 Muon 的更新规则或收敛特性；它严格解决的是系统开销问题。因此，其收益在分布式设置中最为显著，而在单 GPU 训练中，由于缺乏需要消除的分布式冗余，提升相对较小（优化器速度提升约 2 倍）。