使用 NVIDIA Megatron 推进新兴优化器以加速 LLM 训练

至少十年来，Shampoo 等高阶优化算法已被有效应用于神经网络训练。近年来，这些方法在应用于领先的 LLM 时取得了显著成功。特别是，Muon（MomentUm Orthogonalized by Newton-Schulz）被用于训练当今一些最优秀的开源模型，包括 Kimi K2 和 GLM-5。

本文介绍 NVIDIA 如何为 Muon 及其他前沿新兴优化器，以及使它们能够训练大规模模型的相关技术提供全面支持。

NVIDIA GB300 NVL72 上的 Muon 训练性能

表 1 总结了在 NVIDIA GB300 NVL72 系统上，使用 Muon 和 AdamW 优化器训练 Kimi K2 与 Qwen3 30B 模型的吞吐量。借助下一节将介绍的技术，结果显示，与 AdamW 相比，使用 Muon 优化器造成的训练性能损失非常小。在计入 Newton-Schulz 迭代中矩阵乘法的 FLOPs 时，Muon 的模型 FLOPs 利用率（MFU）更高。

这些测量结果是使用 NVIDIA NeMo Megatron Bridge 26.02 获得的；它是 NeMo Framework 中的一个 PyTorch 原生库，为热门 LLM 和 VLM 模型提供预训练、SFT 和 LoRA。NVIDIA 团队在 Kimi K2 训练中使用了 256 块 NVIDIA GB300 GPU，并采用 PP4DP64EP64；在 Qwen3 30B-A3B 训练中使用了 8 块 NVIDIA GB300 GPU，并采用 DP8EP8。

有关更详细的实验设置和复现该测量的步骤，请参见下方的 Quickstart 部分。

支持大规模 Muon 训练的技术

尽管 Muon 等高阶优化器有望提升 LLM 训练效率，但其大规模实际部署面临若干障碍。使用 Newton-Schulz 迭代或特征分解的预条件处理步骤会显著增加计算成本和内存消耗。此外，在混合精度训练和梯度累积过程中可能出现数值不稳定。要在数千个 GPU 之间高效分发同步且正交化的更新，可能会引入通信瓶颈，从而削弱整体效率收益。

本节介绍用于克服这些障碍并扩展到数千个 GPU 的技术。所选技术在通用性、吞吐量和实现复杂度之间取得平衡。我们倾向于选择可泛化到 Muon、SOAP（ShampoO with Adam in the Preconditioner’s eigenbasis）以及其他复杂优化器的技术，而不是仅关注 Muon；由于 Muon 的简单性，针对它还可以应用更多优化。

逐层分布式优化器

传统的逐元素分布式优化器（通常大规模应用于 AdamW 等逐元素优化器）提供以下功能：

• 状态分片：优化器状态不会由每个 GPU 全部存储，而是被均匀切分到所有可用 GPU 上。
• Reduce-scatter 梯度：对所有梯度执行 reduce-scatter 操作，每个 GPU 获取与其“拥有”的参数相对应的一部分梯度。
• 本地更新：每个 GPU 只更新其“拥有”的模型参数的特定部分。
• AllGather 参数：更新后，GPU 之间进行通信，以确保所有设备都拥有完整模型的更新后权重，用于下一次前向传播。

更高级的版本会拆分各项操作，并将通信与计算重叠执行。

逐元素分布式优化器与 Adam 优化器配合一直表现良好。然而，包括 Muon 在内的许多新兴优化器，都需要整个层的梯度来计算该层权重的更新。如果权重和优化器状态在数据并行（DP）rank 之间均匀分布，就无法基于每个 GPU 上可用的数据来计算更新。为收集数据以计算完整更新，将需要额外的通信。

引入了一种逐层分布式优化器，以支持这类优化器。不同层的参数被分布到不同的 DP rank。每个 GPU 都拥有完整层的参数，从而可以计算预条件器。

逐层分布带来的一个变化是可变大小通信。由于整个层的大小各不相同，每个 GPU 都需要通过 all_gatherv 从不同 GPU 收集大小不同的参数更新。

这种逐层分布式优化器已完全集成到 NVIDIA Megatron Core 中，这是一个用于构建和训练大规模模型的开源库，支持高级并行、混合精度和优化的 GPU 内核。实现可在 layer_wise_optimizer.py 中找到。

分布式 Newton-Schulz

大规模 LLM 训练通常会将张量并行（TP）与其他并行方式结合使用，这带来了独特的挑战。TP 会沿特定维度将单个权重矩阵拆分到多个 GPU 上，这意味着没有任何单个设备持有完整的参数张量。Muon 优化器的关键正交化步骤需要访问整个动量缓冲矩阵才能进行计算。当该动量缓冲在设备之间分片时，需要额外通信来处理分片后的动量和权重。

TensorParallelMuon 提供了两种方法来处理此类场景：duplicated mode 和 distributed mode。

正文：Duplicated mode

首先在 TP 域内对 momentum 进行 all-gather，使每个 GPU 都拥有用于 Newton-Schulz 迭代的全部数据。然后，每个 GPU 执行完整的 Newton-Schulz 迭代，并使用相应的正交化 momentum 更新其拥有的权重。该模式以每个 GPU 执行相同的 NS 迭代为代价，对网络延迟进行了优化。对于每次更新，无论使用多少次 NS 迭代（通常为五次），都只需要预先进行一次 all-gather 通信。

正文：Distributed mode

顾名思义，NS 迭代的计算分布在 TP 域中的各个 GPU 上。每次 NS 交互会执行三次矩阵乘法，在每次迭代的第一次矩阵乘法之后都需要一次 all-reduce。该模式以更频繁的通信为代价，对计算进行了优化。

此外，还支持一种 blockwise 模式，该模式仅使用每个 GPU 所拥有的动量和权重进行正交化和更新。与另外两种模式相比，它的计算成本更低，并且不需要任何通信。然而，它并不等同于对整个动量矩阵一起进行正交化。从数学上看，它在概念上是一种块正交化（因此得名 blockwise），类似于 Scalable Second Order Optimization for Deep Learning 中首次引入的分块方法。

其他优化

其他优化可以进一步提升训练吞吐量。本节将简要讨论其中一些重要优化，并将通过 NVIDIA Emerging Optimizers 研究项目逐步介绍。

通信隐藏

在 NVIDIA Emerging Optimizers 项目的当前版本中，参数会在优化器步骤之后立即被聚合并完全暴露，在某些情况下这可能成为瓶颈。逐元素分布式优化器中使用的相同技术也可以应用于此：参数聚合可以延迟到下一批次的前向步骤，并与计算重叠。

负载均衡

由于注意力和 MLP 权重矩阵的大小各不相同，优化器因预处理而产生的计算成本在各层之间也会有所不同。即使在理论上，也无法实现完美的负载均衡。目前，在同一个 DP 域内，层以轮询方式分布到各个 GPU 上。也有一些方法可以基于对计算和通信成本的估算，实现更好的负载均衡。

SYRK 和融合的 all-reduce

NS 交互中的三次矩阵乘法里，前两次可以映射为 SYRK（SYmmetric Rank-K update，对称秩 K 更新），从而节省接近一半的浮点运算。节省量略低于一半，因为沿对角线的 tile 需要完整计算。当前版本提供了 Triton kernel。

SYRK 输出的是完整矩阵，而不是以压缩格式表示的上三角（或下三角）矩阵。在分布式模式下，all-reduce 仍需要通信整个矩阵，而不是一半矩阵。更高级的版本可以将通信融合到 SYRK kernel 中，这不仅节省带宽，还可以以更细粒度（即 tile 级别）隐藏通信。计划在未来版本中提供 CuTe DSL 实现。

NVIDIA 还支持哪些其他用于研究的优化器？

除了 Muon 之外，NVIDIA 还支持许多其他优化器，供研究社区探索，包括：

• 正交化优化器 MOP（Momentum Orthogonalized by Polar decomposition，基于极分解的动量正交化）的最终形式
• 一种高级 SOAP 变体，在 REKLS 中通过特征分解并结合 KL 修正逐步更新特征基

复现 Muon 训练结果

本节说明如何复现表 1 中的吞吐量结果，以及 Muon 如何集成到 NVIDIA 软件栈中。

快速开始

请遵循 NVIDIA-NeMo/Megatron-Bridge GitHub 仓库中性能方案的说明。

例如，要在 256 个 NVIDIA GB300 GPU 上运行 Kimi，请从 Megatron Bridge 仓库根目录开始：

CONTAINER="nvcr.io/nvidia/nemo:26.04"

python scripts/performance/setup_experiment.py 
  --account <slurm_account> \
  -i ${CONTAINER} \
  --partition <slurm_partition> \
  -m kimi \
  -mr kimi_k2\
  --log_dir <result_dir> \
  --num_gpus 256 \
  --gpus_per_node 4 \
  -t "00:15:00" \
  -g gb300 \
  -c fp8_mx \
  -hf <HF_TOKEN>

集成

Muon 已集成到 Megatron Core 中，包括一个 Muon 类和按层分布式包装器。

请注意，对于 Muon，use_distributed_optimizer 会自动将其分派到按层分布式优化器，而不是逐元素对应的优化器。

开始使用用于 LLM 训练的新兴优化器

像 Muon 这样的高阶优化器正被证明对于突破 LLM 训练效率的边界至关重要。通过逐层分布式优化、定制化的分布式 Newton-Schulz 迭代，以及在通信隐藏和融合 SYRK/all-reduce 内核方面的持续工作，这些强大的方法现在可以在超大规模场景中有效部署。借助 Megatron Core，你现在就可以开始使用它们。

• Muon 和其他新兴优化器已集成到 Megatron Core 中。
• MOP 和 REKLS 等研究型优化器可在 NVIDIA-NeMo/Emerging-Optimizers GitHub 仓库中用于实验。

选择配置时，请牢记以下注意事项：

• 分布式 NS 模式：当网络延迟占主导时，使用 duplicated 模式。当计算成为瓶颈时，使用 distributed 模式。Blockwise 模式完全避免通信，但会近似完整正交化。
• 规模：逐层分布式优化器使 Muon 能够以极低开销扩展到大规模 DP。表 1 显示，在 NVIDIA GB300 上，其吞吐量与 AdamW 接近持平。
• 未来改进：通信隐藏、更好的负载均衡以及融合的 SYRK/all-reduce 内核将在即将发布的版本中进一步缩小任何剩余的吞吐量差距。

准备好开始了吗？请查看 Megatron Bridge 性能配方。