Muon is Scalable for LLM Training
Jingyuan Liu, Jianlin Su, Xingcheng Yao, Zhejun Jiang, Guokun Lai, Yulun Du, Yidao Qin, Weixin Xu, Enzhe Lu, Junjie Yan, Yanru Chen, Huabin Zheng, Yibo Liu, Shaowei Liu, Bohong Yin, Weiran He, Han Zhu, Yuzhi Wang, Jianzhou Wang, Mengnan Dong, Zheng Zhang, Yongsheng Kang, Hao Zhang, Xinran Xu, Yutao Zhang, Yuxin Wu, Xinyu Zhou, Zhilin Yang
cs.LG, cs.AI, cs.CL
2025-02-24
Moonshot 给 Muon 加上权重衰减、按矩阵形状补齐更新后，缩放实验里对齐同一损失大约只需 AdamW 52% 的训练算力，并据此训出 5.7T token 的 Moonlight。
Muon 是 Jordan 等人 2024 年的矩阵优化器。动量累积之后，用 Newton-Schulz 迭代把动量矩阵近似正交化，让更新的奇异值靠近 1，免得权重只沿少数主导方向长。正交化对应谱范数约束下的最速下降。网络权重是作用在输入上的算子，谱范数这种算子范数比 Adam 那种随训练变动的约束更贴这个角色。小模型上它比 AdamW 快，放到十亿级参数、万亿级 token，优势会缩掉。
缩掉来自两件具体的事。不加权重衰减时，权重和层输出的 RMS 一直涨，会冲出 bf16 的高精度区间。更新幅度还被矩阵形状绑死：满秩矩阵 [A, B] 的理论更新 RMS 是 1/√max(A,B)。又大的 MLP 矩阵步子太小，容量用不满；GQA、MLA 里把单个 KV head 拆开时矩阵偏小，步子又太大，训练不稳。分布式还有一道坎，正交化要整块梯度，按元素更新的 ZeRO-1 套不上。
两处改动之后，学习率和权重衰减可以直接拿 AdamW 已经搜过的那套。
权重衰减与 AdamW 同形，每步多减掉 λ 倍当前权重。800M 模型、100B token，大约最优 token 量的 5 倍。不加衰减的 Muon 前段更低，24000 步验证损失领先约 0.023；后段权重胀开，66000 步时带衰减的版本反超约 0.017，也低于 AdamW。
形状补偿写进同一步更新：正交化结果乘上 0.2·√max(A,B)。根号项抵消上面的缩小，0.2 把更新 RMS 拉到 AdamW 常见的 0.2 到 0.4。矩阵走 Muon，RMSNorm、embedding、LM head 仍走 AdamW，两边共用 η 和 λ。
三种缩放在 800M 上比过。实验把 SwiGLU 换成 [H, 4H] 的两层 MLP，只训 4B/20B token。只乘 0.2·√H 的基线，验证损失 2.812；把更新硬归一到 RMS=0.2，以及按形状乘 √max(A,B)，验证损失都是 2.789。后一种更省，后面都用它。
Newton-Schulz 做 10 步比 5 步更接近正交，损失没有更好，N 取 5。动量 0.95，多项式系数沿用原作 3.4445、−4.7750、2.0315。权重衰减 0.1。
分布式版本把优化器状态按 ZeRO-1 切在数据并行组上，多一次组内 gather，拼出完整矩阵做 Newton-Schulz，再只留本卡分片去更新。动量缓冲只有 AdamW 的一半，优化器额外显存大约少一半。通信量约是分布式 AdamW 的 1 到 1.25 倍，上界对应 fp32 的 reduce-scatter、all-gather，再加一次 bf16 gather。开了张量并行，还要在 TP 组上多一次 bf16 gather。优化器通常只占前向加反向的 1% 到 3%。大规模集群上，端到端延迟相对 AdamW 看不出来。
缩放实验是 Llama 结构的稠密模型，399M 到 1.5B，token 按算力最优配，上下文 8K。AdamW 做了网格搜索，Muon 复用同一学习率和批次。拟合是 Muon 2.506×C^(−0.052)，AdamW 2.608×C^(−0.054)。算力最优设定下，对齐同一损失大约只要 AdamW 52% 的训练 FLOPs。
用这套优化器从零训出的模型叫 Moonlight，骨架靠近 DeepSeek-V3-Small：激活 2.24B，总参数 15.29B（不含 embedding，含 embedding 约 3B/16B），5.7T token，上下文 8K。学习率升到 4.2e-4 再余弦降到 4.2e-5，最后 500B token 做 cooldown，数据偏数学、代码和推理。
干净对照停在 1.2T，还没进入 cooldown。和同配方 AdamW 基线 Moonlight-A 比，MMLU 几乎打平，代码和数学明显更高，BBH 反而是 AdamW 更高。HumanEval 与 MBPP 都按 pass@1 计。
| 方法 | 指标 | 结果 |
| Moonlight（Muon，1.2T） | MMLU / BBH / HumanEval / MATH | 60.4 / 43.2 / 37.2 / 19.8 |
| Moonlight-A（AdamW，1.2T） | MMLU / BBH / HumanEval / MATH | 60.2 / 45.3 / 29.3 / 16.1 |
| DeepSeek-V3-Small（AdamW，1.33T） | MMLU / BBH / HumanEval / MATH | 53.3 / 41.4 / 26.8 / 10.7 |
公开的 V3-Small 更弱，但 Moonlight-A 的 MMLU 已经有 60.2，配方差距大于优化器差距。
训到 5.7T，Moonlight 的 MMLU 70.0，MMLU-Pro 42.4，BBH 65.2，HumanEval 48.1，GSM8K 77.4，MATH 45.3。同为 2.4B/16B、同为 5.7T 的 DeepSeek-V2-Lite 是 58.3、25.5、44.1、29.9、41.1、17.1。Qwen2.5-3B 是稠密 3B、18T token，MMLU 65.6，GSM8K 79.1，Moonlight 的 MMLU 更高，GSM8K 略低。图 1 把它标在性能对训练 FLOPs 的帕累托前沿上。从 1.2T 的 60.4 到满训 70.0，多出来的 token 和 cooldown 都在里面，没有同期的 AdamW 孪生模型。
1.2T 这段预训练里，六个权重组的平均 SVD 熵一直是 Muon 更高；到 1.2T 检查点，超过 90% 的单块矩阵也是 Muon 更高，路由矩阵拉开最多。
SFT 和预训练最好用同一个优化器。1.2T 检查点在 tulu-3 上微调两个 epoch，这里的 MMLU 是 0-shot CoT，和预训练的 5-shot 不是同一套协议。预训练和 SFT 都用 Muon 时，MMLU 55.7，HumanEval 57.3，GSM8K 68.0；都用 AdamW 时是 52.0、53.1、64.6。Muon 预训练再改 AdamW 做 SFT，MMLU 掉到 50.2。Qwen2.5-7B 这种 AdamW 预训练模型只在 SFT 换 Muon，四项都略差：MMLU 70.8 对 71.4，HumanEval 77.4 对 79.3，MBPP 71.6 对 71.9，GSM8K 85.8 对 89.8。
预训练团队若从零开训，有机会沿用 AdamW 的学习率，把权重衰减放在 0.1，换来缩放曲线上大约 52% 的训练 FLOPs。动量状态少一半，通信增量没有变成可见的端到端变慢。分布式实现、预训练权重、指令微调权重和中间检查点一并放出。
已经用 AdamW 训完的模型，只在 SFT 换 Muon，目前看不到稳定收益。1.2T 的同配方对照里，涨点集中在 HumanEval 和 MATH，MMLU 只动了 0.2。
讨论里留了三件事。RMSNorm、embedding、LM head 仍用 AdamW。谱范数推广到一般 Schatten 范数还没做。预训练和微调优化器不一致时变差，存量 AdamW 检查点不好直接接上 Muon。
52% 这个数字来自 399M 到 1.5B 的稠密模型、算力最优 token 下的拟合，两条曲线指数是 −0.052 和 −0.054，差别主要在系数。Moonlight 是激活 2.24B 的 MoE，训了 5.7T，远超算力最优。同配方 AdamW 只对照到 1.2T，MMLU 差 0.2，BBH 还倒退 2.1。和 V2-Lite 的十几分里混着数据和结构差异。0.2 这个对齐系数取自 AdamW 更新 RMS 的经验区间，理论没有规定唯一取值。端到端延迟看不出来，没有配上 Moonlight 规模的墙钟。