SOAP, Muon, and Beyond: Pushing LLM Pretraining Scales
Mikail Khona, 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
cs.LG, cs.AI
2026-07-14
NVIDIA 用逐步 QR 正交化把 SOAP/Muon 推到万亿 token;大 batch 下 AdamW 退化、Muon 在 3T token 训练中 MMLU 守住 74 以上。
预训练 LLM 默认用 AdamW,但一批「高阶优化器」(Muon、SOAP)号称收敛更快。问题是它们在大规模上有两道坎:算 preconditioner 比 AdamW 贵,而且在超大 batch 下数值不稳,此前没人把它们稳定推到几十亿参数、上万亿 token。NVIDIA 这篇把 SOAP/Muon 真正放到大规模 LLM 预训练里跑通,顺手补上了一系列工程化的修法。
团队先定位了 SOAP 在大 batch 下发散的根因:参考实现每 10 步才刷新一次特征基,还把当前步的梯度排除在外,导致 preconditioner 总是「过时」。权重梯度范数来回振荡,language modeling loss 跟着飙,他们把它叫 slingshot 不稳定。
修法分几层。逐步 QR 正交化:每个训练步都对 preconditioner 矩阵做 QR 分解,把当前梯度算进去,消灭滞后。KL-Shampoo 协方差估计:用 KL 散度正则替换标准的 Kronecker 因子累积,把条件数从 κ 压到 √κ,减小特征分解里的数值误差。
公平比较本身也是贡献。不同优化器的更新幅度差很多,直接迁学习率不公平。作者用 update-RMS matching:匹配参数更新的均方根范数,给 Muon 乘一个约 0.2 的校正因子来抵消 EMA 动量的衰减,这样 SOAP、Muon、AdamW 才在同一把尺子上比。
最后是逐层分布式优化器:把参数矩阵按大小轮流分到各 GPU,每块卡独立更新自己那层,通信用按模型执行顺序分桶的异步 all-gatherv,计算和取下一桶参数重叠。没有对优化器计算做近似,收敛性保住。
横跨 8B 稠密 GPT、Qwen3-30B-A3B、Nemotron-3-Nano-30B 与 72B-A8B,训练到万亿级 token,batch 推到 1 亿 token(100M)。核心结论:大 batch 下 AdamW 退化,SOPA/Muon 稳。以 Nemotron-3-Nano 训 3T token 为例:
| 设置 | MMLU | HumanEval | Math500 |
| AdamW 1× batch | 73.38 | 63.26 | 71.50 |
| Muon 1× batch | 73.71 | 62.56 | 75.50 |
| Muon 2× batch | 74.80 | 64.79 | 74.05 |
| Muon 3× batch | 74.00 | 62.22 | 73.25 |
batch 翻倍到三倍,Muon 的下游质量基本不掉,AdamW 在同等放大下开始出问题。在 Qwen3-30B 上 KL-SOAP 相对 Muon 在 CE loss 上有「持续但很小」的优势,跨 1×/2×/4× batch。顺带一提,Muon 用的 Newton-Schulz 近似正交化和精确极分解(MOP)对比 loss 几乎一样,说明近似质量够用。
对正在预训练大模型的团队,这是「AdamW 撞墙之后能不能换」的一份扎实证据。结论是 Muon/SOAP 在 batch 放大时比 AdamW 稳,逐层分布式实现让它们在大集群上跑得动。代码以 Emerging-Optimizers 名义开源,工程上可直接拿来试。
ε 这个数值稳定项没有系统调过(SOAP 用 1e-8,Muon 用 1e-7),它怎么随模型规模变化是作者点名留给后续的工作。Mamba 这类架构里的 Conv1D 滤波器在几何上不适合做正交化,方法的适用边界没覆盖到。1 亿 token batch 是上界,再往上有没有新的不稳,论文没给。