SOAP、Muon 及更多：推动LLM预训练规模扩展

# SOAP、Muon 及其他：推动大语言模型预训练规模边界
来源：https://arxiv.org/html/2607.20548
Aditya Vavre, Boxiang Wang, Deyu Fu, Hao Wu, Mike Chrzanowski, Bryan Catanzaro, Dheevatsa Mudigere, Jeff Pool, Michael Lightstone, Mohammad Shoeybi, Mostofa Patwary, Nima Tajbakhsh, Tijmen Blankevoort NVIDIA \{mkhona, avavre, boxiangw, deyuf, skyw, mchrzanowski, bcatanzaro, dmudigere, jpool, mlightstone, mshoeybi, mpatwary, ntajbakhsh, tblankevoort\}@nvidia.com

###### 摘要

诸如 Muon 和 SOAP 等高阶优化器比 AdamW 收敛更快，但其计算成本和数值稳定性挑战限制了在大规模场景下的采用。在本工作中，我们调整并增强预处理梯度方法，以克服大规模大语言模型预训练中的实际挑战。

我们首先识别 SOAP 在大量批次训练中的不稳定性，并提出算法修改，包括逐步骤 QR 正交化和改进的预处理策略，以消除损失尖峰并确保在此类设置中的稳定训练。然后，我们通过更新均方根匹配对 SOAP、Muon 和 AdamW 进行统一实证研究，以确保优化器之间公平的学习率迁移。作为此项分析的一部分，我们实证评估了 Muon 的正交化质量。我们针对训练数万亿 token 的数十亿参数模型的实验表明，SOAP 和 Muon 在我们测试的规模下始终优于 AdamW。值得注意的是，对于下一 token 预测，在高达 1 亿 token 的批次大小下，这些优化器保持训练稳定性和质量，而 AdamW 性能下降。

为了实现大规模高效训练，我们引入了一种与 Megatron-LM 兼容的层级分布式优化器。我们的实现平衡了内存开销并隐藏了通信，同时避免了对优化器计算的近似，从而保留了其收敛优势。此外，我们识别并构建了特定的系统级改进，以进一步加速我们的层级实现。为支持研究社区，我们发布了一个包含优化新兴算法的代码库：

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

## 1 引言

优化算法的选择依然是前沿模型训练框架的核心。实际上，它作为所有大规模模型训练的心脏，将系统工程与算法和模型收敛性连接起来。从系统角度看，优化器决定了分布式执行的基本约束：优化器状态通常比模型参数本身占用更多内存，严重影响模型分片和内存管理策略。此外，优化器在极端批次大小下保持稳定性的能力决定了训练能在多大程度上扩展到大型加速器集群，而不会因同步训练工作负载的通信开销成为瓶颈。从算法角度看，优化器控制着数据效率、收敛速度和模型的泛化能力。

尽管具有这种双重影响，深度学习优化的发展一直由一种实际紧张关系定义：元素级标量优化器（如 AdamW [kingma2014adam, loshchilov2017decoupled]、RMSProp [hinton2012neural] 和 LaProp [ziyin2020laprop]）的计算简单性和可扩展性在很大程度上掩盖了关注曲率的高阶方法。虽然像 AdamW 这样的一阶方法因其易用性占据主导地位，但它们本质上独立处理单个参数元素更新，忽略了梯度之间的相关结构以及相应神经网络权重的算子性质。相反，高阶优化器（近似二阶信息）捕捉损失景观的几何形状 [george2018fast, martens2015optimizing, bollapragada2018progressive]，理论上允许更大的步长、更快的收敛以及可能更高的精度，但由于其复杂性而缺乏可扩展性。

最近，张量级自适应方法，如 Shampoo 优化器和现代变体 SOAP [vyas2024soap]、特征校正 Shampoo [eschenhagen2025purifying] 和 KL-Shampoo [lin2025understanding]，以及谱优化器如 Muon 和 Scion [pethick2025training]，已作为折衷方案出现，提供了二阶方法的好处，同时计算和内存开销可控。然而，当这些方法应用于如细粒度 MoE 等大型前沿模型时，面临可扩展性障碍。

本文其余部分围绕扩展这些高阶优化器所需的公式、实证评估和系统工程进行组织。我们的核心贡献如下：

- • 使用 Muon 进行 MoE 的大批量规模扩展：我们使用更新均方根匹配框架来公平地迁移学习率，并比较 AdamW、Muon 和 SOAP 优化器。我们在第 3 节 (https://arxiv.org/html/2607.20548#S3) 中阐述了 MoE 大批量设置，并在第 5.3 节 (https://arxiv.org/html/2607.20548#S5.SS3) 和第 5.5 节 (https://arxiv.org/html/2607.20548#S5.SS5) 中展示了实证结果。我们在高达 720 亿参数的 MoE 模型上的评估表明，虽然 AdamW 在超出临界批次大小后性能下降，但 Muon 和 SOAP 在高达 1 亿 token 的全局批次大小下仍保持 token 高效的收敛。
- • 修复 SOAP 预处理器计算的不稳定性：我们识别了标准 SOAP111https://github.com/nikhilvyas/SOAP 实现中在大量批次设置中特有的关键不稳定性，该不稳定性由预处理器与当前梯度统计之间的滞后导致。第 5.4 节 (https://arxiv.org/html/2607.20548#S5.SS4) 表明，通过强制使用当前梯度进行逐步骤特征基更新，并集成最近开发的基于 KL 散度的协方差估计方法来累积 SOAP 的 Kronecker 因子 [lin2025understanding]，可消除所有训练损失尖峰并提高训练稳定性，从而更有信心在生产中扩展到更大的批次和模型大小。
- • 比较 SOAP 和 Muon：我们比较了这两种优化器在大批量设置中的准确性和稳定性，发现 Muon 和 SOAP 均优于 AdamW，且 KL-SOAP 相对于 Muon 略有优势。我们还讨论了结论的局限性。
- • 可扩展的并行化：我们提出了一种高效的层级分布式实现，旨在在 Megatron-LM 框架内实现可扩展性。我们在第 3.3 节 (https://arxiv.org/html/2607.20548#S3.SS3) 中简要讨论了分布式优化器的约束，并在第 6 节 (https://arxiv.org/html/2607.20548#S6) 中展示了实现。
- • Emerging-Optimizers：我们贡献了一个用于优化器研究的开源库，包含本报告中考虑的所有优化器实现以及许多实验变体。

## 2 Adam、Muon 和 SOAP 入门

Adam 风格优化器是现代大语言模型预训练中最常用的优化器，它们将随机梯度下降 (SGD) 扩展为坐标级自适应学习率和解耦权重衰减 [kingma2014adam, loshchilov2017decoupled]。这里我们简要概述这些方法及其关系。

令 \(g_t \in \mathbb{R}^{d=mn}\) 为步骤 \(t\) 参数张量 \(G_t \in \mathbb{R}^{m \times n}\) 的展平梯度向量。AdamW 独立地为每个参数元素维护梯度和平方梯度的指数移动平均 (EMA)，其中 \(\beta_1, \beta_2\) 为 EMA 时间尺度系数。

\[
m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t
\]
\[
v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t \odot g_t
\]
其中 \(\odot\) 表示逐元素（Hadamard）乘积。AdamW（为简单起见忽略一阶动量 \(m_t\) 和偏差校正）的预处理更新方向对梯度应用对角预处理器：

\[
u_t = m_t \circ \text{diag}\left(\dfrac{1}{(v_t)^{1/2} + \epsilon}\right)
\]

这种逐元素结构使得 AdamW 在内存上规则，并且易于在大型分布式系统中进行分片。然而，这也意味着 AdamW 忽略了神经网络权重的矩阵或张量结构：每个坐标单独重新缩放，但未显式建模跨行、列、头、专家或投影子空间的相关性。

另一类优化器试图通过预处理来利用这种结构。Shampoo [gupta2018shampoo] 通过为每个张量维度维护 Kronecker 因子化的梯度统计来近似二阶信息，允许使用曲率感知变换更新矩阵或张量形状的参数，而无需形成全密集预处理器。Shampoo 为行和列维护两个独立的对称协方差矩阵：

\[
L_t = \beta_2 L_{t-1} + (1-\beta_2) G_t G_t^\top
\]
\[
R_t = \beta_2 R_{t-1} + (1-\beta_2) G_t^\top G_t
\]
然后，Shampoo 的更新方向通过非对角的 Kronecker 因子化预处理器产生，该预处理器源自 Kronecker 因子的逆幂：

\[
u_t = (R_t^{-1/4} \otimes L_t^{-1/4}) g_t
\]

SOAP（算法 2 (https://arxiv.org/html/2607.20548#alg2)）建立在 Shampoo 之上，将 Shampoo 风格的预处理与预处理器特征基中的 Adam 式逐元素自适应更新相结合 [vyas2024soap]：

\[
u_t = Q_L \texttt{Adam}(Q_L^T m_t Q_R) Q_R^T
\]
(1)

直观上，SOAP 首先将梯度旋转到其行和列相关性近似对角化的基（即 \(R_t\) 和 \(L_t\) 的特征基）中，在该基中应用自适应的逐元素更新，然后将更新旋转回预处理器特征基。这使得 SOAP 比经典 Shampoo 与 AdamW 的联系更紧密，因为对于单位特征向量，我们正好恢复 AdamW，同时保留利用权重矩阵内结构的能力，可能允许更平滑的超参数迁移。然而，Shampoo 和 SOAP 都维护全精度预处理器和 Kronecker 因子矩阵，所需内存显著大于 AdamW。

Muon（算法 1 (https://arxiv.org/html/2607.20548#alg1)）跳过计算和存储预处理器，遵循不同的策略。Muon 不估算预处理器并在预处理基中计算完整的自适应二阶矩统计，而是像 AdamW 和 Shampoo 一样（公式 2 (https://arxiv.org/html/2607.20548#S2.Ex1)）应用动量，然后对其进行正交化。为了正交化 \(M_t\) 并近似其极因子（最接近的正交矩阵，等价于奇异值分解 \(M_t = U \Sigma V^\top\) 中的 \(UV^\top\)），Muon 采用 Newton-Schulz 迭代 [pethick2025training, liu2025muon, shah2025practical, amsel2025polar, jordanmuon]，该方法使用矩阵多项式迭代来近似计算极因子。

这可以被视为谱更新，因为对于二维权重矩阵，Muon 改变更新方向以控制其奇异值，而不是像 AdamW 那样应用独立的坐标级重新缩放。简而言之，SOAP/Shampoo 风格的白化和 Muon 风格的正交化密切相关：两者都减少大奇异方向的主导地位，并产生更好地尊重线性层算子结构的更新。值得注意的是，在 Shampoo 中关闭 EMA（设置 \(\beta_1, \beta_2 = 0\)）在数学上将其预处理更新直接简化为极因子 \(UV^\top\)，因为 \((\mathbf{GG^T})^{-1/4} \mathbf{G} (\mathbf{G^T G})^{-1/4} = UV^\top\)，从而恢复 Muon 更新。详见附录 A.1 (https://arxiv.org/html/2607.20548#A1.SS1) 中关于 SOAP 预处理的 SVD 解释以及三种优化器之间的直观关系。

这些优化器在规模使用时存在实际差异。这里我们总结定性差异。AdamW 是最易于扩展的优化器，因为其状态和计算是逐元素的，因此无限可分片。SOAP 更具表达力，但必须维护和更新 Kronecker 因子和特征基，使其对数值稳定性、预处理器新鲜度更敏感，具有更大的内存占用，以及用于分片的复杂矩阵分布式张量布局。Muon 的占用空间小于 AdamW，因为它不维护二阶矩，但仍需要完整的二维矩阵更新，面临与 SOAP 相同的基于矩阵的分片问题。Muon 在优化器步骤中使用更多计算，其开销取决于正交化的质量和成本，而正交化的成本随 Newton-Schulz 迭代次数缩放。

## 3 大规模、大批量预训练优化入门

### 3.1 混合专家模型的批量大小扩展

对于混合专家（MoE）模型，大批量扩展尤其微妙，因为模型的大部分参数位于专家线性层中（通常占参数总数的 \(>90\%\)）。与密集模型不同，由于稀疏的 top-k 路由，这些层观察到显著更低的**有效**批量大小。考虑一个全局批量大小为 \(B_{\text{Global}}\) 的 token，采用 top-\(k\) 路由策略在 \(N\) 个总专家上训练。假设负载均衡的理想路由器，任何单个专家观测到的有效批量大小 \(B^{\text{expert}}_{\text{eff}}\) 由下式给出：

\[
B^{\text{expert}}_{\text{eff}} = B_{\text{Global}} \times \frac{k}{N}
\]
(2)

由于在典型的细粒度稀疏 MoE 配置中 \(k \ll N\)（例如，top-8 门控与 256 个专家 [liu2024deepseek]），导致 \(B^{\text{expert}}_{\text{eff}} / B_{\text{Global}} = 0.03125\)，每个专家仅看到全局批量的一小部分。因此，当扩展到非常大的全局批量大小时，稀疏专家参数保持在较低有效批量范围内，通常更容易优化。相比之下，密集和共享参数必须处理整个全局批量。因此，增加全局批量大小主要对密集组件的大量批处理容忍性提出挑战。这使得 MoE 训练高度依赖于能够为这些在极端批量大小下运行的密集参数保持稳定性和 token 效率的优化方法 [sun2024hunyuan]。

### 3.2 批量大小–学习率缩放与更新方差

由于我们的比较将全局批量大小改变最多 4 倍，我们需要一个学习率迁移规则，使得在将差异归因于优化器本身之前，更新尺度保持可比。在 SGD 背景下，修改小批量大小需要相应调整学习率，以保持稳定且可迁移的优化动态。考虑一个简单的 SGD 更新，其中梯度估计器 \(\hat{g}_B\) 在大小为 \(B\) 的小批量上计算如下：

\[
\hat{g}_B = \frac{1}{B} \sum_{i=1}^B \nabla L_i(\theta)
\]
参数更新定义为 \(\Delta \theta = \eta \hat{g}_B\)，其中 \(\eta\) 表示学习率。假设单个样本梯度独立，梯度估计器的协方差与批量大小成反比。