SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales论文地址https://arxiv.org/abs/2607.20548项目页面https://github.com/NVIDIA-NeMo/Emerging-Optimizers学术交流922230617目录1. 引言2. 优化器概览AdamW、Muon 与 SOAP2.1 AdamW简单但受限2.2 Shampoo结构感知的 Kronecker 因子分解2.3 SOAP在预处理特征基中执行 Adam2.4 Muon直接进行谱正交化3. 大规模 Batch 训练的挑战与解决方案3.1 MoE 模型中的有效 Batch Size3.2 学习率缩放平方根规则3.3 分布式策略全矩阵 vs. 分片4. 背景与相关工作5. 预训练实验Muon vs. SOAP vs. AdamW5.1 模型与数据5.2 公平对比update-RMS 匹配5.3 Muon 在大 Batch 下显著优于 AdamW5.4 解决 SOAP 的训练不稳定性5.4.1 “弹弓” 失稳现象5.4.2 修复每步 QR 包含当前梯度5.4.3 引入 KL-散度协方差估计5.5 Muon vs. SOAP 正面比较6. 系统实现层间分布式优化器7. 结论与未来方向1. 引言在大模型训练中优化器不仅仅是 “调参工具”——它决定了分布式系统的内存布局优化器状态常比模型参数更占显存、通信开销batch size 极限以及最终的数据效率和泛化能力。长期以来AdamW 凭借其元素级自适应的简单性和可扩展性成为事实标准。但这种 “无视矩阵结构” 的做法忽略了梯度间的相关性也放弃了利用曲率信息加速收敛的机会。高阶优化器如 Shampoo、SOAP、Muon通过Tensor 级别的预处理Kronecker 因子或正交化来逼近二阶信息理论上能走更大步长、更快收敛。然而它们的计算开销、数值稳定性、分布式实现复杂性一直阻碍其在千亿级模型上落地。NVIDIA 团队在最新技术报告中系统性地解决了这些问题并在数万亿 token 的预训练任务中验证了 SOAP 和 Muon 的优越性。核心贡献包括大 batch 规模下的稳定性MoE 模型 batch 达 100M token 时Muon/SOAP 仍优于 AdamW修复 SOAP 的 “失稳”通过每步 QR 正交化和 KL 散度协方差估计消除 loss 尖峰公平对比协议基于 update-RMS 匹配实现优化器间学习率公平迁移可扩展的分布式实现在 Megatron-LM 中引入层间分布式优化器隐藏通信开销。2. 优化器概览AdamW、Muon 与 SOAP现代 LLM 预训练优化器遵循一条从 “逐元素自适应” 到 “结构感知” 的演进路线。为了理解 Muon 和 SOAP 的优势首先需要理清它们与奠基性方法 Shampoo 以及最通用的 AdamW 之间的血缘关系。2.1 AdamW简单但受限AdamW 是目前最通用的基准优化器。它通过维护梯度的一阶矩动量 m_t​和二阶矩方差 v_t​的指数移动平均EMA为每个参数元素独立调整更新步长一阶矩二阶矩其预处理更新方向为这种对角预处理使得 AdamW 内存友好且易于分片但它完全忽略了权重矩阵的行列结构及相关性。2.2 Shampoo结构感知的 Kronecker 因子分解为了捕捉梯度的结构性关联Shampoo 应运而生。它不将梯度展平为一维向量而是保留其矩阵形状 G_t​ ∈ R^{m×n}并分别维护行和列的协方差矩阵其更新方向通过 Kronecker 积形成的全矩阵预处理生成Shampoo 通过 Kronecker 分解近似了二阶曲率信息但计算矩阵逆平方根R^{−1/4}的代价较高且内存占用显著大于 AdamW。2.3 SOAP在预处理特征基中执行 AdamSOAP 直接建立在 Shampoo 的数学框架之上但做了一个巧妙的 “嫁接”它利用 Shampoo 维护的特征向量矩阵Q_L​, Q_R​将梯度旋转到预处理器的对角化特征基中然后在该基中应用 Adam 的逐元素自适应更新最后再旋转回原空间这种设计使得 SOAP 既保留了 Shampoo 的结构感知能力大曲率方向被抑制又继承了 AdamW 平滑的自适应特性理论上更易于超参数的迁移。但其代价是需要维护完整的 Kronecker 因子、特征基且对数值稳定性极为敏感。2.4 Muon直接进行谱正交化与 SOAP 的 “复杂嫁接” 不同Muon 走出了另一条截然不同的高效路径。它完全不计算协方差矩阵或预处理器而是先像 AdamW 一样计算动量矩阵 M_t​随后直接通过牛顿-舒尔茨Newton-Schulz迭代近似计算该动量的极分解Polar Decomposition即最近的正交矩阵Muon 的更新本质上是谱更新它不调整每个元素的大小而是改变更新方向使特征值分布更加均衡。相比 SOAPMuon 省去了二阶矩内存但要求矩阵运算在完整的 2D 张量上进行无法像 AdamW 那样按元素任意分片。小结Shampoo 通过全矩阵预处理奠定了结构感知的基础SOAP 在 Shampoo 的基础上嫁接了 Adam 的自适应更新表达能力最强但开销最大Muon 则通过极分解简化了预处理以最小的额外内存开销实现了高效的谱更新。3. 大规模 Batch 训练的挑战与解决方案3.1 MoE 模型中的有效 Batch Size对于混合专家Mixture-of-ExpertsMoE模型由于稀疏路由机制每个专家实际处理的 token 数远小于全局 batch size。假设路由器负载均衡理想任何单个专家所看到的有效 batch size 由以下公式给出其中k 为 Top-K 路由数量N 为专家总数。这意味着增加全局 batch size 主要对稠密参数如注意力层构成压力而稀疏专家参数仍处于较易优化的低 batch regime。3.2 学习率缩放平方根规则当将批量大小从初始大小 B 调整为目标大小 B 时为保证不同 batch size 下更新方差的一致性求解得采用平方根缩放规则Square Root Scaling Rule这使得在扩大 batch size 时参数更新的随机波动保持相似水平。3.3 分布式策略全矩阵 vs. 分片AdamW 的 element-wise 状态可任意分片但 Muon/SOAP 需要完整 2D 矩阵才能计算正交化或 Kronecker 因子。为此文中设计了层间分布式优化器见第 6 节而非传统的 ZeRO/FSDP 按元素切分。4. 背景与相关工作Shampoo 变体Eigen-Shampoo、KL-Shampoo、DASH 等改善预处理估计。谱优化Muon、Scion 及后续 Muon、AdaMuon、NorMuon 等。系统层面veScale-FSDP 支持非对称切分Canzona 实现异步矩阵收集。5. 预训练实验Muon vs. SOAP vs. AdamW5.1 模型与数据模型8B 稠密 GPT3B/30B 纯 Transformer MoE8B/72B 混合 Mamba-Transformer MoE。训练数据1T~3T token 的 Nemotron-3 数据集。配置序列长度 8192全局 batch 25M token基线学习率采用 WSDWarmup-Stable-Decay计划。5.2 公平对比update-RMS 匹配由于旋转矩阵不改变 Frobenius 范数SOAP 与 AdamW 的更新 RMS 自然一致。而 Muon 需要乘以因子来匹配更新尺度。文中明确采用了 Kimi-Moonshot 团队的 RMS 匹配协议。5.3 Muon 在大 Batch 下显著优于 AdamW实验设计batch size 从 25M 逐步升至 50M、75M带 ramp-up学习率按平方根同步调整。结果图1、图2表5AdamW 在 500M 时 loss 明显退化而 Muon 保持降低 loss。在 30B-A3B 和 72B-A8B MoE 上Muon 在编程和常识推理任务上 gains 最显著。注意Mamba 的 Conv1D 滤波器不适合 Muon回退至 AdamW 效果更好。5.4 解决 SOAP 的训练不稳定性5.4.1 “弹弓” 失稳现象原 SOAP 实现为省计算每 10 步才更新特征基且不包含当前步梯度。在大 batch 早期loss 变化剧烈陈旧预处理导致梯度范数振荡进而引发 loss 尖峰图4。在 8B 模型上直接导致训练发散图5。5.4.2 修复每步 QR 包含当前梯度修改为每步重新计算特征基QR 分解并纳入当前梯度即可完全消除尖峰图7。QR 比 eig 更快且效果无差别附录 E 图13。5.4.3 引入 KL-散度协方差估计令在时间步 处梯度的奇异值分解为 G UΣVᵀ。假设 Kronecker 因子 S 和 S 与梯度的特征基大致对齐使得其中 Λ 和 Λ 分别表示 Kronecker 因子所捕获的特征值。将原 Shampoo 的 G·G^T 累加替换为 KL-Shampoo 的耦合更新这使 Kronecker 因子的条件数从 (σ_max ⁡/ σ_min)^2⁡ 降为 σ_max ⁡/ σ_min⁡​平方根改善极大增强数值稳定性图6推导见论文附录。文中推荐KL-SOAP即加了 KL 更新的 SOAP为最强版本。5.5 Muon vs. SOAP 正面比较控制变量统一不拆分 QKV统一 blocking 策略。结果图7、图8在大 batch1x、2x、4x下KL-SOAP 的 CE loss 始终略低于 Muon差距稳定但微小。使用精确 SVD 极分解的 MOPMuon 变体也略优于原 Muon说明正交化精度有提升空间。ϵ 超参数差异SOAP 的 ϵ 类似 Adam作为二阶估计分母的软下界。Muon 的 ϵ 用于归一化矩阵 Frobenius 范数下限。因此将 视为特定优化器的超参数对于SOAP它控制预处理自适应矩和因子反演中的数值下限而对于Muon则控制牛顿-舒尔茨迭代中进入的最小归一化尺度。局限文中未系统调优将作为后续研究。6. 系统实现层间分布式优化器为实现 Muon/SOAP 的高效分布式训练文中在 Megatron-LM 中构建了专门方案负载均衡将整个参数矩阵不分片按大小排序后 round-robin 分配到各 DP rank确保每卡内存均匀。独立更新每卡只更新自己负责的完整矩阵计算完成后 flatten 到 buffer。异步 All-Gather 重叠利用 Megatron 的 DDP bucket按层顺序发出可变长 All-Gather-V 集体通信在计算前向的同时拉取下一 bucket 的更新参数有效隐藏通信延迟。该实现不牺牲数学等价性全矩阵预处理。7. 结论与未来方向结论。Muon 和 SOAP 在大 batch≥50M token下明显优于 AdamW尤其适合 MoE 的密集层。KL-SOAP综合表现最佳若显存充裕优先推荐 KL-SOAP。未来工作。SOAP 的 TP 支持与 QKV 分裂进一步适配 Tensor Parallelism。优化器驱动的内存布局让 DDP buffer 按矩阵形状分配减少碎片和通信开销。架构协同设计研究 MLA、LoRA 等结构与全秩预处理的匹配性设计混合策略。更精确的正交化大模型中小奇异值易受浮点噪声影响需探索截断/正则化正交化。Batch size 标度律将系统约束最大化 GPU 利用率与理论标度律结合动态优化 batch 计划。