Distributed Implementation of Muon and Emerging Optimizers in Megatron-Core

傅德禹, NVIDIA GPU 加速计算专家团队 | Al Open Day | Nov 07, 2025

议程

• Muon 与新兴优化器
• 实现与 Mcore 集成
• 逐层分布式优化器
• 部署与结果
• 结论与未来路线图

新兴优化器的涌现

为何是现在？

• 大规模语言模型（LLM）训练的规模和成本达到了前所未有的水平，强调了对更高效、更先进算法的需求。
• 优化器的改进与其他优化方法是正交的，能够惠及训练的各个阶段。
• AdamW 作为标准已经有相当长一段时间，但二阶方法应该能更快收敛。

• 已有足够的先前工作表明追求新优化器是有益的：

  • Shampoo: https://arxiv.org/abs/1802.09568
  • SOAP: https://arxiv.org/abs/2409.11321
  • Muon: https://kellerjordan.github.io/posts/muon/

• 最佳的开源基础模型（如 Kimi-K2, GLM-4.5）使用 Muon 进行训练，证明了其在大规模训练中的有效性。

基于预处理（Pre-Condition）的优化器

近似二阶方法

• 二阶方法: 利用二阶偏导数（Hessian 矩阵 H）和牛顿方向 p = -H⁻¹∇f(xₜ)。

• 预处理方法: 使用多种预处理器来“模拟”Hessian矩阵的信息，以加速一阶（基于梯度）方法，即找到一个“更好”的方向。
• 以 Shampoo 为例，对于给定的权重层 W 和梯度 G：
  • 更新预处理器：
    Lₜ = Lₜ₋₁ + GₜGₜᵀ
Rₜ = Rₜ₋₁ + GₜᵀGₜ
  • 更新权重：
    Wₜ₊₁ = Wₜ - ηLₜ⁻¹ᐟ⁴ Gₜ Rₜ⁻¹ᐟ⁴

对 Shampoo 的推广

• 泛化 Shampoo:

• SOAP:

  • Kronecker 因子的特征向量作为预处理器，接近但不完全精确，其中 m = n = -1/2 且 β 是一个参数。
  • 选择使用特征分解或 QR 分解来更新特征基。

• Muon:

  • 实际上等同于设置 β = 0, m = -1/2, n = 0。
  • 使用 Newton-Schulz 方法来近似矩阵的逆。

对 Muon 的浓厚兴趣

为何选择 Muon？

• 始终表现出优于 AdamW 的结果。

• 已被开源基础模型验证。

  • 实现更快的收敛，减少所需的训练步数。
  • 通过更好的优化环境条件来提高最终模型的质量。
  • 由于对学习率的敏感性降低，使得超参数调整更加高效。

• 所有操作都映射为矩阵乘法，易于实现。

实现与 Mcore 集成

独立的 GitHub 代码库

地址: https://github.com/NVIDIA-NeMo/Emerging-Optimizers

• 这是一个旨在全面理解新兴优化器性能影响的研究项目。

• 关注方面包括：

  • 优化器的计算成本、实现和实际速度。
  • 收敛速度和成本效益（基于挂钟时间）。
  • 最终解的质量。
  • 内存消耗。
  • 可扩展性。
  • 低精度下的数值稳定性。
  • 超参数。

• 提供包括 Shampoo、SOAP、Muon 等优化器的 GPU 优化实现。

• Muon 是目前需求最大的，也将是后续工作的重点。

Muon 算法

• 在应用更新前，使用 Newton-Schulz 方法进行正交化 SGD 更新。
• 作用于 2D 权重，而非逐点操作。
• Newton-Schulz 是一个迭代且计算密集的过程。

Muon 优化

优化 NS (Newton-Schulz) 步骤

• 通过 SYRK（对称矩阵的秩-k 更新）加速 Newton-Schulz，其中对称矩阵 C 通过通用矩阵 A 更新。

  • NS 操作在对称矩阵上进行。
  • 每一步可以调用两次 GEMM 上的 SYRK。
  • 理论上只需要一半的浮点运算，并且可以比相同大小的 GEMM 快 2 倍。
  • 提供一个由 Triton 生成的实现，更快的核即将推出。

• 分布式 Newton-Schulz：提供 3 种不同模式来处理部分梯度。

  • 分布式：梯度可以因模型并行（TP）或 ZeRO 优化而分布。
  • 复制：allgather 权重，每个 rank 对复制的常规 NS 步骤执行。
  • 分布式：各 rank 协同执行 NS 步骤与并行 GEMM。
  • 分块：每个 rank 仅用本地梯度执行 NS 步骤（数学上不同）。

Megatron-Core 集成

易用性

• 代码库位置: https://github.com/NVIDIA/Megatron-LM/blob/dev/megatron/core/optimizer/muon.py
• 支持以下功能：
  • 所有模型并行方式 (TP, PP, EP)。
  • get_megatron_muon_optimizer() 函数用于在训练循环中替换 megatron get()。
  • 自动为 AdamW + Muon 混合使用（Muon 仅用于线性层权重）拆分参数组。
  • QKV 融合权重的 Muon 更新。

逐层分布式优化器

Mcore 分布式优化器 (ZeRO-1)

• 与 ZeRO-1 相似，但不完全相同。

  • 优化器和其状态在数据并行（DP）的 ranks 之间分片。
  • Reduce-scatter 梯度和 all-gather 参数。
  • 主要使用 fp32 参数和混合梯度。

• 结果是，优化器更新是分布式的。

• 更重要的是，这减少了优化器的状态并节省了内存。
• reduce-scatter 和 all-gather 作为性能优化可以重叠。
• 内存消耗对比: 假设我们使用混合精度训练和 Adam 优化器。比较使用和不使用 ZeRO-1 的总内存使用量。ψ 表示模型大小（参数数量），Nd 表示 DP 的度。
  • 不使用 ZeRO-1 的内存消耗是 (2 + 2 + 3 * 4) * ψ = 16ψ。
  • 使用 ZeRO-1 的内存消耗是 2ψ + 2ψ + 12ψ/Nd。

规模化挑战

Muon 无法与 Mcore 分布式优化器协同工作

• 在规模化场景下，分布式优化器（或 FSDP）是必需的。
• 预处理器（preconditioner）的计算需要整个层的梯度。传统的 uniform RS -> update -> AG 流程无法工作。

下图展示了梯度缓冲区分片（全局、本地、参数）的情况，其中每个 DP rank 只持有部分参数的梯度。

我们的解决方案

逐层分解 (Layer wise decomposition)

• 将不同的层分发到不同的 GPU，而不是将扁平化的完整模型均匀地分发。
• 这种方法需要解决一些问题：
  • 通信将变为非均匀的，需要可变大小的 reduce-scatter 和 all-gather。
  • 工作量和内存消耗可能不均衡。
  • 随着 DP ranks 数量的增加，需要扩展到更多的层。

下图展示了两种不同的分发策略。

替代算法

由 Kimi 使用

• 各个 rank 在层权重边界与邻居交换梯度，然后执行常规更新并丢弃“多余”部分。
• 一个 rank 最多与2个邻居通信并收集2个层的数据。
• 大层跨越多个 rank 可能会导致问题（下图黄色参数部分）。
• 更大的问题：新版本的 Distopt 和 FSDP 使用更细粒度的权重分片，例如，所有权重均匀地分片到所有 rank。这会导致大量的通信和重复更新，从而导致性能不佳。

实现与优化

再次强调，优先考虑易用性

• 可在以下地址获取：https://github.com/NVIDIA/Megatron-LM/blob/dev/megatron/core/optimizer/layer_wise_optimizer.py
• 实现为 Mcore ChainedOptimizer，为 MLM 训练循环保持功能对等。
• 支持 DP 和 EP 权重，这需要不同的通信组。
• 复用 DDP 的 all-reduce，而不是可变大小的 reduce-scatter，以获得良好性能和重叠。
• 在分布式更新后，使用单个 all-gather 来同步参数。
• 通过根据大小和形状分布层来平衡工作和内存。
• 保持简单和可维护性。与分布式优化器相比，代码行数为 165 行 vs 2600 行。

部署与结果

通过 MLM 参数更改进行简单部署

• Megatron-LM 用户可以指定 "--optimizer muon" 来为线性层使用 muon，其余部分默认使用 AdamW。

• 指定 "--optimizer dist_muon" 来使用 Muon + Layerwise 分布式优化器。

  • 确保 "--use-distributed-optimizer" 和相关选项已关闭。

• 可以从 Mcore 优化器包装器直接调用。

• 不与 Muon 耦合，可以与其他优化器一起使用，包括 AdamW（如下一张幻灯片所示）。
  • 此选项当前未作为 megatron-lm 参数公开。

DS-v2 代理模型，GPU == DP SIZE == 64

内存使用单位为 GB，时间单位为 ms

* 仅供技术讨论和参考。性能可能因不同产品组合而异。

性能结果

• 端到端（E2E）性能由于重叠更能代表优化器性能，但方差较大。此处的数字是第10步到第100步的90个步骤的平均值。注意，通常优化器时间不到端到端时间的5%。
• 当使用 AdamW 作为内部优化器时，Layerwise 的性能介于有/无 all gather 重叠的 distopt 之间。
• Muon 的计算密集度更高，但在分片到 DP ranks 后，其相较于 AdamW 的额外成本是可以接受的。
• Layerwise + Muon 在内存消耗和速度上与 Mcore dist-opt 相匹配，尽管 Muon 的成本更高。
• 我们将在未来几周内分享在更大模型和规模上的更多结果。

收敛性

超参数（HP）建议

• 拆分 QKV 能产生更好的结果。拆分头（heads）则不能。
• 不使用 Nesterov 能产生更好的结果。
• 需要额外的 0.2 缩放因子来使 Muon 的更新 RMS 与 AdamW 相似。
• 复用其他 AdamW 的超参数。

Qwen-30B 收敛结果

与 AdamW 相比，Muon 最终获得了更好的损失，并且尖峰更少。

结论与未来路线图

结论

• 新兴优化器为 LLM 训练提供了简便的成本节约方案，值得进一步探索。
• Muon 技术成熟，并被证明比 AdamW 收敛更快、更稳定。
• Mcore Muon 集成在速度和准确性方面均达到生产标准，同时易于使用和维护。

未来路线图

• 近期（数周内）

  • 首先登陆 Mcore Dev 分支！
  • 实现 MuonClip，无需重新计算 logits。
  • 拆分通用的融合权重（包括 MLA）。
  • 在各种开源模型上进行更多基准测试。
  • 使用 Muon 训练更大的模型（进行中）。

• 长期

  • 更好的核函数和分布式 NS 实现。
  • 为 Layerwise 分布式优化器提供更快的通信和重叠。
  • 低精度训练支持（目前仅 bf16 经过良好测试）。
  • FSDP 和更多的检查点格式。

• 欢迎贡献和反馈！