跳转至

Muon Optimizer

导言

muon you

Muon 相对 AdamW 的核心优势不是“单步更便宜”，而是“到达同等 validation loss 所需 token/FLOPs 更少 + optimizer state 显存更少”。在大模型预训练里，这通常比单步 optimizer 算子开销更重要。

截至 2026-05-06 的公开资料看：Kimi K2 使用的是 MuonClip，即 Muon + QK-Clip；DeepSeek-V4 模型卡也称使用 Muon 来获得更快收敛和更好稳定性。注意这通常不是“纯 Muon 替代所有 AdamW”，而是 Muon 优化 2D 矩阵权重，AdamW 仍处理 embedding、norm、LM head、1D 参数。

1. Loss 下降：主要优势是 token/FLOP efficiency

机制差异

AdamW 是逐元素自适应优化：

\[ m_t = \beta_1 m_{t-1} + (1-\beta_1)g_t \]
\[ v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2 \]
\[ \Delta W \propto \frac{m_t}{\sqrt{v_t}+\epsilon} \]

它的预条件是 diagonal / element-wise，对矩阵结构利用有限。

Muon 对矩阵参数先做 momentum，再对 momentum 矩阵做近似正交化：

\[ M_t = \beta M_{t-1} + g_t \]
\[ U_t \approx \text{orthogonalize}(M_t) \]

其中 orthogonalization 通常通过 Newton–Schulz 迭代近似实现。直观上，它把更新矩阵的奇异值拉得更均衡，避免更新能量过度集中到少数主方向。Moonshot/UCLA 的 Muon 论文也用 SVD entropy 观察到，Muon 训练出的权重矩阵在多数分组上比 AdamW 有更高的谱多样性。

实验量级

Moonshot/UCLA 的 Muon is Scalable for LLM Training 报告中，scaling-law 实验显示：Muon 达到 AdamW 同等 loss 大约只需要 52% 的训练 FLOPs。换句话说，如果 AdamW 需要 100 单位训练计算量到某个 validation loss，Muon 大约 52 单位就能到。来源：Muon is Scalable for LLM Training 1。

Kimi K2 技术报告进一步说明，K2 使用 MuonClip，在 15.5T token 预训练中实现了无 loss spike 的稳定训练；报告也强调 Muon 在相同 compute/model/data 下相对 AdamW 有更好的 token efficiency。来源：Kimi K2: Open Agentic Intelligence 2。

所以在 loss 维度可以概括为：

指标 AdamW Muon / MuonClip
同等 token/FLOPs baseline validation loss 通常更低
同等目标 loss 需要更多 token/FLOPs 约 0.5x AdamW 训练计算量，取决于实现和任务
稳定性 成熟稳定 裸 Muon 大规模可能有 attention logit 爆炸，Kimi 用 QK-Clip 解决

2. 显存：optimizer state 约省一半，但总显存不会省一半

AdamW 对每个参数通常维护两个状态：

  • first moment：m
  • second moment：v

Muon 对矩阵参数主要维护一个 momentum buffer，没有 AdamW 的 second moment。

所以只看 optimizer states：

\[ \text{AdamW states} \approx 2P \]
\[ \text{Muon states} \approx 1P \]

即 矩阵参数部分 optimizer state 显存约为 AdamW 的 50%。Muon 论文也明确指出，Distributed Muon 的额外 optimizer state memory 约为 Distributed AdamW 的一半。来源：Muon is Scalable for LLM Training 1。

但工程上要注意三点：

  1. 不是所有参数都用 Muon NVIDIA NeMo 的 Muon 文档也提醒：Muon 要求传入 2D 参数，不应用于 embedding、final FC、1D 参数，这些通常仍用 AdamW。来源：NVIDIA NeMo Emerging Optimizers - Muon 3。

  2. 总训练显存不只 optimizer state 总显存还包括：

  3. 参数
  4. 梯度
  5. optimizer states
  6. activation
  7. communication buffer
  8. temporary tensor

所以 optimizer state 省 50%，不代表总 HBM 省 50%。

  1. 混合优化器下的实际节省取决于矩阵参数占比 设可用 Muon 的矩阵参数占比为 \(f\)，则 optimizer-state 节省约为：
\[ \frac{f}{2} \]

如果 90% 参数可用 Muon，则 optimizer-state 约省 45%。 如果再算 FP32 master weight、activation 等，总显存下降会更低，可能是 10%～30% 量级，具体看 ZeRO、activation checkpoint、sequence length 和 MoE 并行策略。

3. 耗时：单步未必更快，但 time-to-loss 明显更优

这是最容易误解的点。

单步耗时

AdamW 单步是逐元素操作，理论上很便宜。Muon 额外做 Newton–Schulz 正交化，涉及矩阵乘法，因此 裸看 optimizer step，Muon 比 AdamW 更复杂。

但大模型预训练中，主要时间在 forward/backward 的 GEMM 和 attention，而不是 optimizer step。Muon 论文中提到，分布式 Muon 的额外 latency 通常相对 forward/backward 很小，optimizer 端延迟大约只占 1%～3%，工程上通过 BF16、tensor core、overlap communication 可以做到几乎无感。来源：Muon is Scalable for LLM Training 1。

所以：

维度 结论
单步 optimizer 算子 Muon 通常比 AdamW 重
端到端 step time 优化良好时接近 AdamW，可能略慢
到达同等 loss 的总时间 Muon 通常显著更短，约 1.8x～2x 加速可能成立

简单估算

如果 AdamW 到目标 loss 需要 100 GPU-days，Muon 需要 52% FLOPs：

  • 若 Muon 单步开销近似相同： [ 100 / 52 \approx 1.92\times ]

  • 若 Muon 单步慢 5%： [ 100 / (52 \times 1.05) \approx 1.83\times ]

  • 若实现很差，单步慢 30%： [ 100 / (52 \times 1.3) \approx 1.48\times ]

所以 Muon 的耗时优势来自“更少步数 / 更少 token / 更少 FLOPs 到达目标 loss”，不是来自单步更便宜。

4. 为什么 Kimi/DeepSeek 这类 MoE 前沿模型特别适合 Muon？

主要原因有三类：

  1. 矩阵参数占比极高 Transformer 的 Q/K/V/O projection、MLP up/down/gate、MoE experts、router 都是矩阵。Muon 正好作用于这些核心权重。

  2. MoE 对训练效率极敏感 万亿参数 MoE 的训练瓶颈不是单纯参数量，而是 token efficiency、通信、显存、稳定性。Muon 如果能用 0.5x FLOPs 达到同等 loss，收益非常大。

  3. 高质量 token 越来越稀缺 Kimi K2 报告明确强调 token efficiency 的重要性，并采用 MuonClip 来提升每个 token 的学习效率。来源：Kimi K2 Technical Report 2。

DeepSeek-V4 模型卡也称其采用 Muon optimizer 以获得更快收敛和更好训练稳定性。来源：DeepSeek-V4-Flash Model Card 4。

5. 但 Muon 不是无脑替换 AdamW

工程上要注意这些坑：

  1. 不能直接给所有参数用 Muon 常规做法是：
  2. 2D hidden matrix：Muon

  3. embedding / LM head / norm / bias / 1D 参数：AdamW

  4. 需要 weight decay 和 update RMS scaling Moonshot 的 scalable Muon 版本里，weight decay 和 per-parameter update scale 是关键改动；否则大规模训练可能出现权重 RMS 增长、稳定性问题。

  5. 裸 Muon 在超大模型上可能有 attention logit 爆炸 Kimi K2 报告指出，Muon 比 AdamW 更容易出现 attention logits explosion，所以他们提出 QK-Clip，形成 MuonClip。来源：Kimi K2: Open Agentic Intelligence 2。

4. fine-tuning 不一定收益明显

Muon 的优势最明确是在 pretraining。对于已经用 AdamW 预训练好的公开模型，直接用 Muon 做 SFT 不一定比 AdamW 好，甚至存在 optimizer mismatch。

Qwen 3.5 35B SFT 时，Muon 20步的loss下降没有 Adaw快

Qwen checkpoint 可能是 AdamW/Adam 系预训练，SFT 阶段切 Muon 存在 optimizer mismatch； Muon 的主要优势在预训练 token efficiency，不保证 SFT 每步 train loss 更快； AdamW 的逐元素自适应更新更适合短程 SFT 快速压低 CE loss； Muon 对 lr、update scale、weight decay、参数分组非常敏感；

总结

维度 相对 AdamW 的优势 注意事项
Loss 下降 同等 compute 下 validation loss 更低；同等目标 loss 约只需 52% FLOPs 主要在大规模预训练显著；SFT 不一定
显存 矩阵参数 optimizer state 约省 50% 总显存不会省 50%；非 2D 参数仍用 AdamW
耗时 time-to-target-loss 可接近 2x 加速 单步不一定更快，依赖高质量实现
稳定性 配合 weight decay、RMS scaling、QK-Clip 后可大规模稳定 裸 Muon 可能有 attention logit 爆炸
适用场景 LLM / MoE / Transformer pretraining 小模型、数据受限、普通 fine-tuning 收益不确定

一句话：Muon 的价值在于把 AdamW 的“可靠但逐元素”的更新，换成更符合矩阵权重几何结构的正交化更新；它用一点 optimizer 计算开销，换来更好的 loss/FLOP 曲线和更低 optimizer-state 显存。对于万亿级 MoE 预训练，这个 trade-off 非常划算。

评论