8. 优化器：AdamW 与 Muon

训练循环算出梯度后，需要有人决定“参数应该改多少”。这个角色就是优化器。

nanochat 使用一个组合优化器：MuonAdamW。它不是两个优化器各跑一遍，而是在同一个类里，根据参数类型分别使用 Muon 或 AdamW。

读组合优化器说明nanochat/optim.py第 184–259 行↗

为什么需要两种优化器

模型里的参数不是同一种东西：

• 二维矩阵适合用 Muon 更新。
• 词嵌入、输出层和少量标量参数不适合 Muon，继续用 AdamW。

代码中的分组大致是：

Transformer 块里的矩阵参数  → Muon
词嵌入、输出层、值嵌入      → AdamW
残差缩放、x0 缩放等标量      → AdamW

图 1：参数分组。Muon 只处理 Transformer 块中的二维矩阵，其余参数交给 AdamW。

AdamW 在做什么

AdamW 可以拆成三个动作：

1. 动量：记录梯度的滑动平均，像滚下坡的球保留一点惯性。
2. 二阶矩：记录梯度平方的滑动平均，让更新步长根据每个参数的历史变化自动调整。
3. 解耦权重衰减：在每个步骤中直接把参数往零拉一点，而不是把衰减混进梯度。

核心更新可以写成：

m = β1 × m + (1 - β1) × 梯度
v = β2 × v + (1 - β2) × 梯度²
m̂ = m / (1 - β1^t)
v̂ = v / (1 - β2^t)
参数 = 参数 - 学习率 × m̂ / (√v̂ + ε)

nanochat 把这些操作融合进一个编译后的函数，减少 Python 和 GPU 之间的调度开销。数学意义没有改变。

读 AdamW 更新nanochat/optim.py第 23–63 行↗

Muon 在做什么

Muon 只用于二维矩阵。它的更新分几步：

1. 用带动量的梯度得到初步更新方向。
2. 用一个迭代过程把更新矩阵“近似正交化”。
3. 做方差缩减，让不同行的更新尺度更均衡。
4. 用谨慎权重衰减更新参数。

“正交化”听起来抽象，可以理解为：不要让更新只集中在少数几个方向上，而要把各个方向拉得更均匀。

标准梯度矩阵可能在某些方向上很大，在另一些方向上很小。Muon 把这个矩阵调整得更接近一个正交矩阵，让更新在多个方向上都有合理的作用。这样往往能加快小模型的训练。

代码使用 Polar Express 迭代来完成近似正交化，而不是精确的奇异值分解。精确分解太慢，近似迭代在 GPU 上更实用。

读 Muon 更新nanochat/optim.py第 111–180 行↗

初次阅读时，不需要推导每一步的系数。先记住三个边界：

• Muon 只处理二维矩阵。词嵌入、输出层和 0 维、1 维参数继续用 AdamW。
• 同一个 Muon 组里的参数必须形状相同，代码会把它们堆叠起来一起计算。
• 分布式训练时，优化器状态按 rank 分片，采用 ZeRO-2 风格的思路，减少每张卡的显存占用。

学习率不是固定值

训练过程中，学习率会按计划变化。base_train.py 使用三段式：

线性预热 → 保持 → 线性下降

读学习率、动量和权重衰减计划scripts/base_train.py第 359–386 行↗

• 预热：前 40 步从较小值升到目标值，避免一开始更新过猛。
• 保持：在中段用相对稳定的学习率。
• 下降：最后 65% 的步数逐步降到初始值的 5%。

Muon 的动量也会从 0.85 逐渐升到 0.97，再在训练末尾降到 0.90。权重衰减则按余弦曲线逐渐降到零。

这些计划不是装饰，它们直接影响训练稳定性。恢复训练时，这些状态也会被保存，否则学习率会突然跳回初始值，破坏训练。

哪些参数用哪个学习率

模型配置里把参数分成了几组，并给每组设置不同的学习率：

参数组	优化器	默认学习率（示意）
输出层 lm_head	AdamW	0.004
词嵌入 wte	AdamW	0.2
值嵌入	AdamW	词嵌入学习率的一半
残差缩放、x0 缩放	AdamW	0.5 或更小
Transformer 块矩阵	Muon	0.02

这些默认值不是普适常数，而是针对 d12 参考模型调过的。模型宽度变化时，AdamW 的学习率还会按 1/√(d_model/768) 缩放。

读参数组和学习率设置nanochat/gpt.py第 419–457 行↗

这也解释了一个常见误解：Muon 不是“替换 AdamW 的新优化器”。nanochat 并没有把全部参数都交给 Muon，而是让两种优化器各自处理擅长的部分。

分布式训练时发生了什么

在多卡环境中，每张卡先算出自己那份数据的梯度。如果直接各更新各的，模型会分叉。所以优化器需要：

1. 同步梯度，让每张卡看到全局梯度。
2. 把参数的更新任务分配到不同卡上。
3. 把更新后的参数收集回来。

nanochat 把这部分放在优化器内部完成，而不是使用 PyTorch 的 DDP 包装。这样做的好处是能把通信和计算重叠，并按 ZeRO-2 的方式分片优化器状态，减少每张卡的显存占用。

第一遍不需要读懂全部通信代码。先理解目标：多卡算出的梯度必须合在一起，参数必须保持一致。

亲手检查参数组

加载或创建一个模型后，可以打印优化器的参数组：

model.setup_optimizer()
for i, group in enumerate(model.optimizer.param_groups):
    print(i, group["kind"], len(group["params"]), group["lr"])

真实代码中 setup_optimizer() 会返回优化器，而不是把它挂在模型上，所以更准确的写法是：

optimizer = model.setup_optimizer()
for i, group in enumerate(optimizer.param_groups):
    print(i, group["kind"], len(group["params"]), group["lr"])

你会看到多个 adamw 组和多个 muon 组。Muon 组按矩阵形状拆分，因为不同形状的矩阵不能直接堆叠。

本章要记住的三件事

1. AdamW 用一阶矩和二阶矩自适应地更新参数，适合词嵌入、输出层和标量。
2. Muon 对二维矩阵的更新做近似正交化，让更新方向更均衡。
3. 学习率、动量和权重衰减都随训练进度变化，恢复训练时要一起恢复。

下一章看怎样判断训练结果好坏：评估。