02. 一步训练发生了什么

上一篇的地图里，预训练要跑上百万个 GPU 小时，资料助手的 SFT 只要一张卡几小时。规模差了几个数量级，最小的重复单位却是同一个：训练 step，也就是一次参数更新。模型训练入门用三条运费记录写过最小的训练循环：算预测、算损失、求梯度、按学习率更新。这一篇把那个循环换成语言模型，看它多出了哪些零件：损失怎样从“猜下一个 token”算出来，优化器为什么用 AdamW 而不是简单地减去梯度，batch、step、epoch、梯度累积怎样换算，学习率为什么先升后降，以及混合精度和梯度裁剪各防住了什么。读完以后，你能看懂一份训练配置里的大多数超参数，并按样本数算出要训多少步。

一个 step 的全貌

资料助手的 SFT 配置是这样的：每次往 GPU 里送 4 条问答（一个 micro-batch），累积 8 次再更新一次参数。于是一个 step 里，前向和反向各做 8 次，参数只改 1 次：

上排是每个 micro-batch 都要走一遍的计算，下排是每个 step 只做一次的收尾。和运费例子相比，结构没有变：前向算预测，损失打分，反向求梯度，优化器更新。变化在每个零件内部，下面按图中编号逐个拆开。

一句制度原文变成几道选择题

训练样本里有一句制度原文“年假需提前三天申请”。为了便于手算，假设分词器把它切成 5 个 token：年假 / 需 / 提前 / 三天 / 申请（真实分词器的切法会不同，第 06 篇讲分词）。

语言模型的训练目标是预测下一个 token。一句话因此变成了 4 道选择题，每道题只能看到前面的内容：

位置	模型看到的	正确答案
1	年假	需
2	年假 需	提前
3	年假 需 提前	三天
4	年假 需 提前 三天	申请

一次前向就能同时算出 4 个位置的预测，模型内部用遮罩保证每个位置看不到后面的词。这也是语言模型训练效率高的原因：一条 2048 token 的样本，一次前向就提供 2047 道题。

交叉熵：给正确答案的概率打分

每道选择题的选项是整个词表。Llama-3.1-8B 的词表有 128,256 个 token，模型在每个位置为每个 token 输出一个分数，叫 logit。分数经过 softmax 变成概率：每个分数取指数，再除以全部指数之和，保证概率都为正且加起来等于 1。

看第 3 个位置（“年假 需 提前”之后）。为了能手算，只保留 4 个候选：

候选 token	logit	e^logit	概率
三天	2.0	7.389	0.6095
一周	1.0	2.718	0.2242
申请	0.5	1.649	0.1360
下班	-1.0	0.368	0.0303
合计		12.124	1.0000

正确答案“三天”拿到 0.6095 的概率。交叉熵损失只看正确答案的概率，取负对数：

单个位置的损失 = −ln(正确 token 的概率)
              = −ln(0.6095) ≈ 0.495

概率为 1 时损失为 0；概率越小，损失增长越快，概率 0.01 时损失为 4.6。负对数的意义可以这样理解：模型给正确答案的概率每减半，损失就加 0.693（即 ln 2），无论起点在哪。

这个数字可以和一个基准对照。刚初始化、什么都没学的模型，近似在 128,256 个 token 里平均分配概率，损失约为 ln(128256) ≈ 11.76。训练日志里 loss 从十一点多开始往下掉，就是这个原因。

一条样本的损失是各位置的平均。假设 4 个位置的损失分别是 2.30、1.20、0.495、0.15，平均为 (2.30+1.20+0.495+0.15) ÷ 4 ≈ 1.04。一个 micro-batch 再在所有样本的所有位置上平均。SFT 时通常只对回答部分的 token 算损失，问题部分只作为条件，第 09 篇细讲。损失函数的一般概念见怎样衡量预测好坏。

反向：梯度从概率差开始

交叉熵配 softmax 有一个干净的性质：损失对每个 logit 的梯度，等于“预测概率减去正确答案的 one-hot 向量”（正确 token 处为 1，其余为 0）：

候选 token	概率 p	one-hot	梯度 p − one-hot
三天	0.6095	1	−0.3905
一周	0.2242	0	+0.2242
申请	0.1360	0	+0.1360
下班	0.0303	0	+0.0303

按“参数减去梯度”的方向走，“三天”的分数会被抬高，其他候选按各自概率被压低，概率越大的错误答案压得越多。四个梯度加起来为 0，抬高和压低正好平衡。

这只是梯度的起点。反向传播用链式法则把它一层层传回去，经过输出层、32 个 Transformer 层，直到词嵌入，给 80.3 亿个参数每个都算出一个梯度。单参数的链式法则在梯度与学习率里手算过，多层网络的传递见神经网络与反向传播。这里只需记住两个结论，下一篇算显存时要用：每个参数都有一个梯度，梯度和权重一样多；反向要用到前向的中间结果（激活），这些结果必须在前向时保存下来。

AdamW：给每个参数单独定步长

运费例子用的是最朴素的规则：新参数 = 原参数 − 学习率 × 梯度。在大模型上它有两个麻烦。一是梯度很嘈杂，每个 micro-batch 只有几条样本，这一步说往东、下一步说往西；二是不同参数的梯度大小可能差几个数量级，同一个学习率对一些参数太大、对另一些太小。

Adam 为每个参数额外记两个数来解决这两个问题。一阶动量 m 是梯度的滑动平均，记住“最近大体往哪个方向走”，抵消噪声；二阶动量 v 是梯度平方的滑动平均，记住“这个参数的梯度通常多大”，用来把步长归一化。AdamW 在 Adam 之上，把权重衰减单独作为一步执行：

θ ← θ × (1 − lr × λ)              权重衰减：每步把参数往 0 收一点
m ← β₁ × m + (1 − β₁) × g         一阶动量，β₁ 常取 0.9
v ← β₂ × v + (1 − β₂) × g²        二阶动量，β₂ 取 0.999 或 0.95
m̂ = m ÷ (1 − β₁ᵗ)，v̂ = v ÷ (1 − β₂ᵗ)   偏差修正，t 是第几步
θ ← θ − lr × m̂ ÷ (√v̂ + ε)         ε 很小（1e-8），防止除以 0

θ 是一个参数，g 是它这一步的梯度，lr 是学习率，λ 是权重衰减系数。偏差修正是因为 m、v 从 0 起步，前几步会被低估。

拿一个参数代入两步：θ = 0.5，lr = 1e-4，λ = 0.1，第 1 步梯度 g = −0.39，第 2 步 g = −0.20。

第 1 步：m = 0.1 × (−0.39) = −0.039       v = 0.001 × 0.1521 = 0.0001521
         m̂ = −0.039 ÷ 0.1 = −0.39         v̂ = 0.0001521 ÷ 0.001 = 0.1521
         m̂ ÷ √v̂ = −0.39 ÷ 0.39 = −1.00
         θ = 0.5 × (1 − 1e-5) + 1e-4 × 1.00 ≈ 0.500095

第 2 步：m = 0.9 × (−0.039) + 0.1 × (−0.20) = −0.0551
         v = 0.999 × 0.0001521 + 0.001 × 0.04 = 0.00019195
         m̂ = −0.0551 ÷ 0.19 = −0.290       v̂ = 0.00019195 ÷ 0.001999 = 0.0960
         m̂ ÷ √v̂ = −0.290 ÷ 0.310 = −0.936
         θ ≈ 0.500184

两个现象值得看清。第 1 步的 m̂ ÷ √v̂ 恰好是 −1，参数移动了整整一个学习率，与梯度是 0.39 还是 39 无关：Adam 的步长由学习率决定，梯度的绝对大小被约掉了。第 2 步梯度降到 −0.20，比值只降到 −0.936，因为动量记着上一步的方向。这让学习率有了直观含义：每步每个参数大约移动 lr 那么多。

权重衰减那一项每步把参数乘以 (1 − lr × λ)，长期看会压住那些没有梯度支持、只是越长越大的参数。PyTorch 的 torch.optim.AdamW 默认 betas=(0.9, 0.999)、eps=1e-8、weight_decay=0.01；大模型预训练常用 β₂=0.95、λ=0.1，Llama 2 论文报告的就是这组。

代价也在这里：每个参数要多存 m 和 v 两个数，而且通常用 FP32 存。这是下一篇显存账本里最大的一笔。

学习率先升后降

训练刚开始时，m、v 还没积累起可靠的统计，模型参数是随机的或刚换了任务，梯度方向很不稳定。这时用满学习率容易一步迈出合理范围，loss 突然冲高甚至变成 NaN。所以学习率先从很小的值线性升到峰值，叫预热（warmup）。之后再逐渐降低，让模型在接近终点时小步微调，这一段最常用余弦衰减：

预热阶段（s < W）：lr = 峰值 × (s + 1) ÷ W
衰减阶段（s ≥ W）：lr = 最低值 + 0.5 × (峰值 − 最低值) × (1 + cos(π × (s − W) ÷ (T − W)))

s 是当前 step（从 0 数），W 是预热步数，T 是总步数。代入资料助手的设置：峰值 2e-4，最低值取峰值的 10% 即 2e-5，预热 20 步，总共 564 步（下一节算出）。到衰减阶段正中间第 292 步，cos 项为 0，lr = 2e-5 + 0.5 × 1.8e-4 = 1.1e-4，正好在峰值和最低值的中点：

曲线在顶部和底部都比较平，中间下降最快。规模越大的训练预热越长：Llama 3 405B 预训练预热 8000 步，峰值 8×10⁻⁵，在 120 万步里余弦衰减到 8×10⁻⁷。资料助手的 2e-4 是 LoRA 微调的常见量级；全参数微调通常小一个数量级，因为改动的是全部参数。学习率是最影响结果的超参数，第 05 篇会讲怎样从 loss 曲线判断它是否合适。

batch、step、epoch 怎样换算

这几个词经常混用，按资料助手的数字把它们一次理清。训练集有 6000 条问答。

• micro-batch：一次前向反向同时处理的样本数，受显存限制。这里是 4。
• 梯度累积：连续做几个 micro-batch，把梯度加在一起再更新。这里是 8。
• 全局 batch：一次参数更新实际“看过”的样本数，等于 micro-batch × 累积次数 × GPU 数。
• step：一次参数更新。学习率调度、日志、检查点都按 step 计数。
• epoch：把整个训练集完整过一遍。

全局 batch = 4 × 8 × 1 张卡 = 32 条
每个 epoch 的 step 数 = 6000 ÷ 32 = 187.5 → 188（最后一个 step 只有 16 条）
3 个 epoch = 188 × 3 = 564 step
换成 4 张卡、其他不变：全局 batch = 128，每个 epoch 47 step，共 141 step

梯度累积存在的理由是显存。32 条 2048 token 的样本一起前向，激活放不下；拆成 8 次、每次 4 条，激活只需要 1/8，而参数更新看到的仍是 32 条的平均梯度。代价是时间：8 次 micro-batch 是串行的。

加卡之后 step 数变少，每步看的样本变多。这时学习率通常也要相应调整，否则同样 3 个 epoch 只更新 141 次，效果可能不同。预训练常用 token 数衡量 batch：Llama 3 405B 从每步 400 万 token 起步，后来加到 1600 万。

梯度累积的一个坑

梯度累积看起来和大 batch 完全等价，但实现时有一个容易出错的地方：语言模型的损失按 token 平均，而不同 micro-batch 的 token 数不一样。

假设一个 step 累积两个 micro-batch。A 里都是短回答，共 100 个回答 token，损失总和 50；B 里有长回答，共 900 个 token，损失总和 270。

正确做法：总损失 ÷ 总 token 数 = (50 + 270) ÷ (100 + 900) = 0.32
常见错误：先各自平均再平均 = (50÷100 + 270÷900) ÷ 2 = (0.5 + 0.3) ÷ 2 = 0.40

错误做法让 A 里的每个 token 权重是 B 里的 9 倍，短样本被过度重视，训练结果和真正的大 batch 不一致，而且 loss 曲线会随累积次数变化。2024 年 10 月，Hugging Face 确认 Transformers 的 Trainer 在因果语言模型上有这个问题并做了修正。判断方法很简单：同一份数据，用“micro-batch 8、累积 1”和“micro-batch 1、累积 8”各跑几十步，loss 曲线应该基本重合；明显分开就要检查损失的归一化方式。

混合精度：低精度算，高精度存

前向和反向的矩阵乘法用 BF16（16 位浮点）计算，比 FP32 快得多，数据量也减半。但参数更新这一步不能用 BF16，原因可以用上面的数字直接看出来。

BF16 只有 7 位显式尾数，在 0.5 到 1 之间，相邻两个可表示的数相差 2⁻⁸ ≈ 0.0039。AdamW 第 1 步要把 0.5 改成 0.500095，改动量 0.000095 不到这个间隔的 1/40，舍入后还是 0.5。每一步都这样，参数就永远不动。FP32 在同一区间的间隔约 6×10⁻⁸，能稳稳记住这次改动。

所以混合精度训练维护两份权重：前向反向用 BF16 副本，优化器在一份 FP32 主权重上更新，更新完再复制成 BF16 给下一步用（图中第 ⑦、⑧ 步）。m 和 v 也存 FP32。数值格式的细节和 autocast 的用法见混合精度怎样计算。BF16 的指数范围和 FP32 一样宽，所以不像 FP16 那样需要损失缩放。

梯度裁剪：限住异常的一步

偶尔有一个 micro-batch 特别“怪”，比如一条超长的乱码样本，算出的梯度比平时大几十倍。Adam 的归一化来不及适应，这一步可能把参数推得很远，表现为 loss 尖峰。梯度裁剪在更新前检查所有梯度合起来的大小：

总范数 = √(所有参数梯度的平方和)
若总范数 > 阈值：所有梯度 × (阈值 ÷ 总范数)

阈值常取 1.0，Llama 2 预训练和 Hugging Face Trainer 的默认值（max_grad_norm）都是这个数。某一步总范数是 4.0，所有梯度就乘以 0.25：方向不变，只是步子变小。这个范数值每步都会记进日志，它突然变大往往是数据或学习率出问题的先兆，第 05 篇会用到这条曲线。

一个 step 的工作可以连起来复述：一句话被拆成多道“预测下一个 token”的选择题，交叉熵按正确答案的概率打分，梯度从“概率减 one-hot”出发反向传到每个参数；几个 micro-batch 的梯度按 token 正确累加，裁剪住异常值，AdamW 用动量和二阶动量给每个参数定步长，在 FP32 主权重上完成更新，学习率按预热加余弦的曲线变化。这一路上出现了好几份和参数一样大的数据：BF16 权重、梯度、FP32 主权重、m、v，再加上为反向保存的激活。它们各占多少显存，8B 模型一张卡能不能训，是下一篇的问题：03. 训练显存账本。