一句话总结

NVIDIA 研究人员通过每步 $QR$QR 正交化和改进的预调节策略来解决 SOAP 在大批次大小下的不稳定性，并通过使用更新 $RMS$RMS 匹配的统一实证研究，表明 SOAP 和 Muon 在数万亿 token 训练的多十亿参数模型上始终优于 AdamW，同时还引入了一个与 Megatron-LM 兼容的逐层分布式优化器，该优化器无需近似即可高效扩展训练。

核心贡献

• 对 SOAP 的算法修改，包括每步 QR 正交化和改进的预调节策略，消除了损失尖峰，使得在大批次大小下训练稳定。
• 使用更新 RMS 匹配进行公平学习率迁移的统一实证研究表明，SOAP 和 Muon 在数万亿 token 训练的多十亿参数模型上始终优于 AdamW，并在高达 1 亿 token 的批次大小下保持稳定，而 AdamW 则性能下降。
• 一款与 Megatron-LM 兼容的逐层分布式优化器在避免近似的同时平衡了内存并隐藏了通信，并嵌入了系统级改进以加速大规模训练。

引言

优化器的选择对大规模模型训练至关重要，直接塑造了内存分片、通信开销和收敛速度。虽然像 AdamW 这样的逐元素优化器易于在分布式系统中扩展，但它们忽略了梯度张量的相关结构。捕捉损失曲面曲率的高阶方法可以更快收敛并使用更大的批次大小，但其复杂性历来阻碍了它们扩展到前沿模型。最近，诸如 SOAP 和 Muon 这样的张量级优化器作为折中方案出现，以可管理的开销提供了二阶优势，但它们在大规模混合专家架构上仍会遇到稳定性和可扩展性障碍。作者通过让 Muon 和 SOAP 在全局批次大小高达 1 亿 token 时保持 token 高效的收敛，识别并修复了 SOAP 在大批次训练下预调节器计算中的一个关键不稳定性，并在 Megatron-LM 中提供了一个逐层分布式实现，以及一个开源优化器研究库，从而解决了这些差距。

数据集

作者从 Nemotron-3 数据集 中提取训练数据，并使用两个固定大小的子集进行所有预训练实验：

• 1 万亿 token 子集
• 3 万亿 token 子集

两个子集均直接从 Nemotron-3 完整语料库中提取；论文未详细说明数据集的原始组成、领域划分或过滤规则。

训练中的使用 这些子集用于预训练多种模型架构：80 亿参数的密集 GPT 模型，3B 激活/30B 总参数 MoE Transformer 模型，以及一个 8B 激活/72B 总参数的混合 Mamba-Transformer MoE。实验将子集视为可互换的训练语料库；最终选择（1T 或 3T）取决于进行的具体优化器对比。

流水线中的处理

• 所有数据被 token 化并打包成 8192 个 token 的序列。
• 全局批次大小固定为 3072 个样本（相当于每个优化步 2500 万 token）。
• 除了 Nemotron-3 的原始预处理外，未提及额外的裁剪、去重或元数据构建。

方法

作者引入了一个面向生产的训练框架，将矩阵感知优化器，特别是 Muon 和 SOAP，集成到 Megatron-LM 生态系统中。核心挑战在于，传统的数据并行切分策略（例如 ZeRO-1、FSDP）将权重矩阵分割成逐元素切片，这破坏了需要完整 2D 参数张量的预调节和正交化操作。解决方案是逐层分布式优化器，结合对张量并行和批次大小缩放的仔细处理，使 Muon 和 SOAP 能够在大规模下运行，而不会产生过多的通信开销。

逐层优化器将整个参数矩阵分配给不同的 GPU，而不是分割单个张量。通过对所有参数矩阵按大小排序，并以轮询方式将它们分配到数据并行秩上，实现负载均衡。这保证了每个 GPU 持有大致相等的内存占用，同时保留了 Muon 的 Newton-Schulz 迭代和 SOAP 的 Kronecker 因子更新所需的矩阵结构。每个 GPU 独立更新其指定的层，然后将结果展平到一个缓冲区中，供下一次前向传播使用。

为了隐藏通信，系统将参数同步与计算重叠。参数矩阵被打包成与模型执行顺序对齐的顺序桶。对于每个桶，一个异步 all-gather 集合操作从其所属 GPU 获取更新后的参数，同时前向传播正在处理当前桶的权重。由于逐层分配导致各 GPU 的参数数量略有不同，因此使用了可变大小的 all-gather-V 原语，而不是标准的均匀集合操作，从而避免了填充或进一步碎片化。

对于张量并行设置，当预调节应用于整个层的权重时，有两种模式可保留完整矩阵结构。在复制模式下，张量并行域内的所有 GPU 首先执行 all-gather 获取权重；然后每个 GPU 独立运行 Newton-Schulz 迭代。对于通信占主导的较小层，此模式是首选。在分布式模式下，每次 Newton-Schulz 迭代期间的中间矩阵乘法结果会进行 all-reduce，这适用于计算成为瓶颈的较大层。两种模式都使用整个层权重的统计信息对输入进行归一化，确保在数学上与非张量并行执行等效。

支撑优化器策略的是批次大小缩放的原则性方法。当全局批次大小改变时，平方根缩放规则根据 $\sqrt{B'/B}$B′/B​ 调整学习率，以保持参数更新的方差。对于混合专家模型，专家看到的有效批次大小为 $B_{\mathrm{Global}} \times k/N$BGlobal​×k/N，由于稀疏的 top-$k$k 路由，这远小于全局批次。因此，密集和共享参数承受了全部大批次负载，使得 Muon 和 SOAP 在这些密集组件上的稳定性至关重要。逐层切分、异步通信和批次大小缩放规则共同构成了一个系统，解锁了在最大训练规模下高阶优化器的收敛收益。

实验

评估在 80 亿参数密集模型和高达 8B 激活/72B 总参数的 MoE 模型上，使用 Nemotron-3 的 1T-3T token 子集，将 Muon 和 SOAP 与 AdamW 进行对比，优化器学习率通过更新 RMS 匹配对齐。Muon 能够扩展到大批次大小，而不会出现困扰 AdamW 的不稳定性，保持更低的损失和更少的尖峰，而 SOAP 则需要每步更新特征基，包括当前梯度和 KL 正则化的协方差估计，以避免“弹弓”发散。在受控对比中，Muon 和 SOAP 均优于 AdamW，KL-SOAP 在交叉熵损失上略有优势，但系统性的 epsilon 调优仍是一个开放方向。

所有评估架构共享 32 个注意力头，头维度为 128，并采用 Grouped-Query Attention，但键值组数量不同。混合专家（MoE）模型使用更深的网络，专家隐藏尺寸更小，而密集基线具有更宽的前馈网络和更少的层数。在 MoE 变体中，缩放行为不同：Nemotron 模型增加隐藏尺寸和专家尺寸，同时保持层数不变，而 Qwen3-30B-A3B 采用更窄的隐藏维度和更小的专家。所有模型都使用 32 个注意力头，头维度 128，以及 Grouped-Query Attention；键值组数量从 2 到 8 不等。MoE 架构有 48 到 52 层，专家隐藏尺寸在 768 到 2688 之间，而密集 8B 模型有 32 层，FFN 隐藏尺寸为 21504。Nemotron-3-Nano-30B-A3B 和 Nemotron-3-72B-A8B 均使用 52 层，隐藏尺寸从 2688 增加到 4096，专家尺寸从 1856 增加到 2688。

用于 Muon 与 SOAP 优化器对比的 Qwen-3-30B-A3B 架构使用了 128 个总专家，每个 token 激活 8 个，没有共享专家。相比之下，两个 Nemotron 模型都包含一个共享专家，并且只激活 6 个路由专家。Qwen-3-30B-A3B 有 128 个总专家，8 个活跃专家，无共享专家。Nemotron-3-Nano-30B-A3B 有 128 个专家，6 个活跃，共享专家大小为 3712。Nemotron-3-72B-A8B 有 512 个专家，6 个活跃，共享专家大小为 5376。Muon 与 SOAP 的优化器对比研究是在 Qwen-3-30B-A3B 架构上进行的。

所有模型共享相同的 WSD schedule，全局批次大小为 3072，权重衰减为 0.1，仅学习率极值不同。Qwen3-30B-A3B 使用最高值（最大 1.6e-3，最小 1.6e-5），8B 密集 GPT 和 Nemotron-3-72B-A8B 共享最低值（最大 8e-4，最小 8e-6），Nemotron-3-Nano-30B-A3B 介于中间（最大 1e-3，最小 1e-5）。Qwen3-30B-A3B 模型训练的最大学习率是 8B 密集 GPT 和 Nemotron-3-72B-A8B 架构的两倍。在所有模型中，最小学习率始终保持为最大学习率的 1%。

Muon 配置了动量系数 0.9，解耦权重衰减 0.1，16 次 Newton-Schulz 迭代，以及 epsilon 为 10^{-7}，用于对归一化所用的 Frobenius 范数进行下限取值。在受控结构条件下，KL-SOAP 在交叉熵损失上相对于 Muon 保持着微小但持续的优势。epsilon 超参数扮演着不同的角色：在 Muon 中，它限制了归一化尺度，而在 SOAP 中，它稳定了预条件基中二阶矩估计的分母。Muon 的 epsilon（10^{-7}）对 Frobenius 范数进行下界，而 SOAP 的 epsilon 在旋转到特征基后对二阶矩估计进行下限取值。在相同张量并行和融合 QKV 投影的受控比较中，KL-SOAP 在交叉熵损失上表现出对 Muon 的微小但持续的优势。Muon 使用 16 次 Newton-Schulz 迭代，系数来自 PolarExpress 方法进行正交化，而 MOP 变体使用基于精确 SVD 的极分解。Muon 的动量 EMA 系数设为 0.9，解耦权重衰减强度为 0.1。

在不同模型尺寸上，Muon 允许更大的批次大小，同时保持或提升下游任务性能。8B 模型显示，在批次大小加倍的情况下，Muon 在所有报告的指标上均优于基线批次大小的 AdamW，在编码任务上增益尤为显著。Nano-V3 结果表明，Muon 的编码准确率在 2 倍批次大小时达到峰值，尽管推理任务对批次大小更敏感。在 8B Hybrid MTP 模型上，2 倍批次大小的 Muon 在 HumanEval 上得分为 65.79，而 1 倍批次大小的 AdamW 为 57.62，在代码生成方面具有明显优势。随着批次大小增加，Muon 的下游指标保持稳健；对于 Nano-V3，2 倍批次大小的 Muon 取得了最高的 HumanEval 和 HumanEval+ 分数，而 MMLU PRO CoT 在 1 倍批次大小下表现最佳。

实验比较了具有不同专家数量、共享专家使用方式和学习率策略的 MoE 和密集架构，包括在 Qwen3-30B-A3B 模型上专门进行的 Muon 与 SOAP 优化器研究。虽然 KL-SOAP 在交叉熵损失上对 Muon 有微小优势，但 Muon 允许更大的批次大小，并提供稳健的下游性能提升，特别是在代码生成方面，且批次大小扩展时性能不下降。这些发现突显了 Muon 在不同模型规模和任务上的可扩展性和有效性。