第 8 章 大模型优化器：AdamW、Muon 与训练配方¶
开场问题：同样的模型、同样的数据、同样的卡，换一个优化器能省下三成算力——这件事在 2025 年之后从论文变成了工业实践。
8.1 优化器同时管三件事¶
在大模型语境下，选优化器不只是选收敛速度：
- 收敛效率：达到同样 loss 需要多少 token 和多少算力。这是最直接的钱。
- 稳定性：能不能在几万步、上千卡的规模下不出 loss 尖刺。
- 显存：优化器状态是训练显存的最大单项（第 1 章算过，AdamW 占每参数 8 字节的矩状态，加 fp32 主权重是 12 字节）。
这三件事互相牵制，所以"哪个优化器最好"没有脱离规模的答案。
8.2 AdamW 为什么当了十年默认¶
Adam 为每个参数维护梯度的一阶矩（动量）和二阶矩（梯度平方的滑动平均），更新时用一阶矩除以二阶矩的平方根。这相当于给每个参数一个自适应的步长：梯度长期很大的参数步子迈小一点，梯度稀疏的参数步子迈大一点。
这个性质对 Transformer 特别重要，因为不同层、不同模块（attention 与 FFN、embedding 与 norm）的梯度量级差异可能有几个数量级，单一学习率根本不work。AdamW 的改进是把权重衰减从梯度里解耦出来，直接作用在参数上，让正则强度不受二阶矩缩放的影响。
代价就是那两份状态：每个参数 8 字节（fp32 的 m 和 v），8B 模型就是 64 GB，比模型本身还大。
8.3 Muon：把"逐参数"换成"逐矩阵"¶
Muon 的出发点是一个观察：神经网络里绝大部分参数是二维矩阵，而 Adam 把它们当成一堆互不相干的标量来处理，丢掉了矩阵结构本身的信息。
它的做法分两步。第一步和 Adam 一样做动量。第二步是关键：对动量矩阵做一次正交化，粗略地说，就是把这个更新矩阵的所有奇异值都拉到 1 附近，只保留它的"方向"，抹掉各个方向上量级的悬殊差异。严格做需要 SVD，太慢；实际用 Newton-Schulz 迭代做近似，只需要几次矩阵乘法，可以完全在 GPU 上跑，开销相对整个训练步很小。
直觉上的差别可以这样理解：Adam 控制的是"每个元素变化多大"，Muon 控制的是"这次更新能把输入向量最多放大多少倍"（谱范数意义上的步长）。后者更接近我们真正想限制的东西——一次更新对模型输出的扰动。实践中它表现为：同样的 loss 用更少的 token 和算力达到。Moonshot 在 Moonlight 的实验里报告了约两倍的算力效率——这是在他们那套设定下的结果，不是普适常数，换模型形状和数据配比都会变。
三个必须知道的工程事实：
- 它只用于二维矩阵参数。embedding、lm_head、LayerNorm 的 scale、bias 这些一维或性质特殊的参数仍然用 AdamW。所以真实训练里跑的是一个混合优化器，不是纯 Muon。
- 它更省显存。只需要动量一份状态，比 Adam 少一份二阶矩。
- 分布式实现是主要落地成本。Newton-Schulz 需要完整的矩阵，而 FSDP/ZeRO 下参数是分片的，必须先把矩阵聚合起来或者设计分布式版本（Moonshot 为此专门做了 distributed Muon）。这是它从"小规模很香"走到"千卡可用"要跨的那道坎。
8.4 大规模下的稳定性：MuonClip 这类补丁¶
换优化器不是无痛的。Muon 在大规模训练上遇到的一个具体问题是 attention logits 爆炸——第一册 4.4 节讲 QK-Norm 时描述过同一个现象。Moonshot 的做法是在 Muon 基础上加 qk-clip，直接约束 query 和 key 的规模，合称 MuonClip，用它完成了万亿参数 MoE 的预训练且全程无 loss 尖刺。
这件事的普适教训比技术细节更重要：优化器、模型结构的稳定性设计、学习率方案是一套耦合的系统。换掉其中一个，另外两个大概率要跟着调。
8.5 其他值得知道名字的¶
- Shampoo / SOAP：更彻底的二阶方法，为每层维护预条件矩阵。效果好但矩阵求逆开销和显存都重，SOAP 是把 Adam 放到 Shampoo 的特征基上的简化版本。
- Adafactor：把二阶矩做低秩分解，显存从 O(参数量) 降到 O(行+列)。显存紧张时的经典选择，代价是效果略逊。
- 8-bit 优化器：把状态量化到 8 bit 存储，用时反量化。省下大半优化器显存，对效果影响很小，中小规模训练里性价比很高。
- Lion：只保留动量，用符号更新。显存省、实现简单，但对超参更敏感，大规模上的表现不如前几者稳定。
8.6 判断一个新优化器值不值得换¶
这个领域论文很多，绝大多数经不起以下三问：
- 对比是否公平。很多"超越 AdamW"的结果赢在对方没调好。合格的对比要求双方在同一算力预算下各自调到最优。
- 验证规模够不够。小规模上赢、大规模上输的情况极其常见。要看它在你的目标规模附近有没有证据。
- 工程代价算没算。分布式实现、checkpoint 兼容、与现有并行策略的配合、以及超参要不要全部重调——这些成本经常超过收益。
对大多数团队，合理的姿态是：AdamW 仍是安全默认；Muon 已经有大规模成功案例，值得在新项目上认真评估；其余的先观望。
8.7 学习率与 batch：和优化器绑在一起的两个旋钮¶
学习率调度现在主流是两种：cosine（经典，但要求提前知道总步数）和 WSD（warmup-stable-decay，中间保持恒定学习率，最后才退火）。WSD 的工程优势很大——中途可以随时从稳定段切出一个分支做退火，得到一个可用模型，而不必重跑；也方便中途加数据或改配比。
需要强调的是：换优化器之后，最优学习率通常完全不同，Muon 和 AdamW 的合适取值不在一个量级上。第一册附录 F 讲的 muP 那套超参迁移结论也是针对特定参数化的，不能直接照搬到新优化器上。
核心问答¶
Q：Muon 和 AdamW 的本质区别用一句话怎么说？ AdamW 逐元素归一化更新，Muon 逐矩阵归一化更新——前者管每个数字变多少，后者管整个矩阵对输出的放大倍数。
Q：为什么 Muon 不能用在 embedding 上？ 它的正交化操作依赖"这是一个有意义的二维变换矩阵"这个前提。embedding 表和 lm_head 更像是一堆独立的向量，LayerNorm 的 scale 则是一维的，都不适用。实际训练用混合优化器。
Q：Muon 省显存吗？ 省一份状态（没有二阶矩），每参数少 4 字节。对 8B 模型是 32 GB，不算小，但它的主要卖点仍是收敛效率。
Q：换优化器要重调哪些超参？ 至少是学习率和权重衰减，通常 warmup 和调度形状也要跟着变。把旧超参直接搬过去做对比，是最常见的错误评测方式。
Q：为什么优化器放在第二册而不是第一册？ 它不改变模型结构，只改变训练过程——和并行、精度、数据管线一样属于"怎么把它训出来"。第一册只关心模型长什么样。