上承〔06 · 一张图、一百行:MiniMind 的骨架〕。上一篇那一百行伪代码,从头到尾只有
forward。模型能算,能出分数,但不会学。参数是随机初始化的,输出是乱码。这一篇补上另一半:backward。看看那 6400 万个旋钮,是怎么从随机数,一步一步被拧到有用的位置的。
一、训练和推理,差在哪儿
推理的时候,模型是只读的。输入一段话,前向传播一次,输出下一个 token,参数一动不动。
训练的时候,模型是可写的。同样是前向传播,但拿到输出之后,不是直接拿去用,而是先和正确答案比一比,算出一个“错得有多离谱”的分数,然后根据这个分数,把每个旋钮往“错得更少”的方向拧一点点。
流程比推理多了三步:
推理: 前向传播 → 输出。
训练: 前向传播 → 算损失 → 反向传播 → 更新参数。
多出来的三步,就是这一篇要拆的东西。
二、损失函数:怎么衡量“猜得准不准”
2.1 模型输出的是一组概率
第五篇讲过,模型最后输出 6400 个 logits,经过 Softmax 变成 6400 个概率。每个概率对应词表里的一个 token,表示“模型认为下一个词是它的可能性有多大”。
比如输入“今天天气”,模型给出:
好 → 0.45
啊 → 0.20
, → 0.12
冷 → 0.08
...
2.2 训练时有“标准答案”
训练数据里,这句话的下一词是“好”。所以正确答案是“好”。
模型给“好”分配了 0.45 的概率。这个概率高不高?损失函数来打分。
2.3 交叉熵损失
MiniMind 用的损失函数是交叉熵损失(Cross-Entropy Loss)。
交叉熵衡量的是:模型给正确答案分配的概率,和理想情况(100% 给正确答案)之间的差距。
是模型给正确答案分配的概率。
- 如果模型给“好”的概率是 1.0,损失是 。完美。
- 如果模型给“好”的概率是 0.45,损失是 。
- 如果模型给“好”的概率是 0.01,损失是 。错得离谱。
给正确答案的概率越高,损失越小;概率越低,损失越大。 这就是交叉熵做的事。
整个序列的损失,是所有位置损失的平均值。每个位置都要猜下一个词,每个位置都有损失,最后取平均。
2.4 为什么用交叉熵,不用别的
交叉熵有一个重要性质:它的梯度很“温柔”。
如果模型给正确答案的概率很低,损失很大,梯度也很大,参数会被大力调整。如果概率已经很高,损失很小,梯度也很小,参数只是微调。
这种“错得越多、改得越猛”的性质,正好适合训练。用别的损失函数,比如均方误差,梯度行为可能就不这么理想。
源码里,MiniMind 的 forward 函数接收 labels,内部调用 PyTorch 的 cross_entropy:
loss = F.cross_entropy(logits.view(-1, vocab_size), labels.view(-1))
三、反向传播:误差怎么传回每一个旋钮
3.1 一句话说清
损失算出来是一个数,比如 0.80。这个数要告诉 6400 万个旋钮:你该往哪个方向拧。
方向怎么算?链式法则,一层一层往回传。
3.2 梯度的含义
每个旋钮收到一个梯度,表示“你往这个方向拧一点点,损失会变小”。
比如某个旋钮当前值是 0.5,梯度是 -0.02,意思是:把 0.5 调大一点点,损失会变小。另一个旋钮梯度是 +0.03,意思是:把它的值调小一点点,损失会变小。
梯度不告诉步长,只告诉方向。
3.3 链式法则的直觉
模型是一层一层算出来的。损失在最外层,参数在最内层。要算“最内层的参数变化,会让最外层的损失变化多少”,就得把中间每一层的“变化传递率”连乘起来。
这就是链式法则。前向传播时,数据从内往外流。反向传播时,梯度从外往内流。
3.4 PyTorch 的自动求导
MiniMind 不手写反向传播,全靠 PyTorch 的 autograd。每次前向传播,PyTorch 会记录所有操作,构建一张计算图。调用 loss.backward() 时,自动沿计算图反向遍历,算出每个参数的梯度。
loss.backward() # 一行代码,算完所有梯度
四、优化器:怎么拧旋钮
4.1 梯度只给方向,不给步长
有了梯度,知道往哪拧。但拧多少?每步拧 0.001,还是 0.1?
步长太大,一下子冲过头,损失反弹。步长太小,拧一万次也到不了位。
这个“步长”就是学习率(learning rate)。学习率的调度,是训练中最重要的超参数之一。
4.2 为什么用 AdamW
MiniMind 用的优化器是 AdamW。
它的前身是 Adam。Adam 给每个旋钮维护两个状态:
- 一阶动量:梯度的指数移动平均,表示“这个旋钮最近一直往哪个方向拧”。
- 二阶动量:梯度平方的指数移动平均,表示“这个旋钮最近的梯度波动有多大”。
更新时,用一阶动量除以二阶动量的平方根,得到一个自适应的步长。梯度波动大的旋钮,步子小一点;波动小的旋钮,步子大一点。
AdamW 在 Adam 的基础上,把权重衰减(weight decay)从梯度里解耦出来,单独施加在参数上。这让权重衰减的行为更可预测。
4.3 MiniMind 的优化器配置
从 train_pretrain.py 里能看到:
optimizer = torch.optim.AdamW(
model.parameters(),
lr=args.learning_rate,
betas=(0.9, 0.95),
weight_decay=0.1
)
lr:初始学习率,预训练默认 5e-4。betas:两个动量衰减率。0.9 控制一阶动量,0.95 控制二阶动量。weight_decay:权重衰减系数,0.1。防止参数变得过大。
4.4 学习率调度
学习率不是固定不变的。MiniMind 用的是余弦退火(Cosine Annealing)。
公式大概是这样:
其中 是当前步数, 是总步数。
效果是:训练开始时,学习率从最大值开始,缓慢下降,到训练结束时降到最小值。
为什么开头大、结尾小? 开头参数还很随机,需要大步伐探索;结尾参数已经接近最优,需要小步伐微调。
MiniMind 的学习率从 5e-4 开始,通过余弦调度逐渐降到 5e-5 左右。
学习率调度在训练循环里的位置,是每次参数更新之前,先更新学习率:
lr = get_lr(current_step, total_steps, args.learning_rate)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
先调学习率,再用新学习率去更新参数。
4.5 梯度裁剪
还有一个防爆炸的保险:梯度裁剪(gradient clipping)。
有时候个别参数的梯度会变得极大,一步更新就把模型带飞到不归路。梯度裁剪的做法是:把所有参数的梯度拼成一个向量,算它的总长度(范数)。如果超过了阈值 grad_clip,就按比例缩放,让总长度刚好等于阈值。
MiniMind 的默认阈值是 1.0。
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
它和 AdamW、学习率调度一起,构成训练稳定性的三道保险。
五、训练循环:一遍一遍过数据
现在把上面的零件串起来。
5.1 一个完整的训练步骤
从 train_pretrain.py 里,一个 step 的流程是这样的:
1. 取一个 batch 的数据 (input_ids, labels)
2. 前向传播:model(input_ids) → logits
3. 算损失:cross_entropy(logits, labels) → loss
4. 损失缩放:loss = loss / accumulation_steps
5. 反向传播:loss.backward()
6. 梯度累积:等攒够 accumulation_steps 个 batch
7. 梯度裁剪:clip_grad_norm_(model.parameters(), grad_clip)
8. 参数更新:optimizer.step()
9. 清空梯度:optimizer.zero_grad()
第 4 步和第 6 步是梯度累积(gradient accumulation)。
5.2 梯度累积:小显存模拟大 batch
如果显存不够,放不下大 batch,怎么办?
用梯度累积。每次只放一个小 batch,算梯度,但先不更新参数。攒够 N 个小 batch 的梯度之后,再一起更新。
效果上,和用一个 N 倍大的 batch 差不多。代价是训练变慢,但显存占用大大降低。
MiniMind 预训练的默认配置是:
batch_size:32(每 GPU)accumulation_steps:8
所以有效 batch size 是 32 × 8 = 256。
5.3 混合精度训练
MiniMind 用 bfloat16 做训练精度。
正常的浮点数是 fp32(4 字节)。bfloat16 是 2 字节,省一半显存。它和 fp16 不同——fp16 的数值范围小,容易溢出;bfloat16 的范围和 fp32 一样大,只是精度低一些。对训练来说,范围比精度更重要。
代码里用 torch.amp.autocast 自动切换:
with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
logits = model(input_ids)
loss = F.cross_entropy(...)
前向传播用 bfloat16 算,参数更新还是用 fp32。这样既省显存,又保持训练稳定。
5.4 一个 epoch 是什么意思
整个数据集过一遍,叫一个 epoch。
MiniMind 预训练默认跑 2 个 epoch。
数据量不大的时候,多个 epoch 能让模型多看几遍数据。数据量极大的时候(比如万亿 token),一个 epoch 可能就跑不完了,会用“步数”而不是“epoch”来控制训练长度。
六、MiniMind 的训练脚本长什么样
把上面的零件拼起来,train_pretrain.py 的核心结构大概是这样:
# 1. 解析参数
args = parse_args()
# 2. 构造模型
model = MiniMindForCausalLM(config)
# 3. 构造数据集
dataset = PretrainDataset(data_path, tokenizer, max_length)
# 4. 构造优化器
optimizer = AdamW(model.parameters(), lr=args.learning_rate, ...)
# 5. 构造学习率调度
scheduler = CosineAnnealingLR(optimizer, T_max=total_steps)
# 6. 训练循环
for epoch in range(args.epochs):
for step, (input_ids, labels) in enumerate(dataloader):
# 更新学习率
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# 前向传播
with autocast(dtype=bfloat16):
logits = model(input_ids)
loss = cross_entropy(logits, labels)
# 损失缩放
loss = loss / accumulation_steps
# 反向传播
loss.backward()
# 梯度累积
if (step + 1) % accumulation_steps == 0:
# 梯度裁剪
clip_grad_norm_(model.parameters(), grad_clip)
# 参数更新
optimizer.step()
optimizer.zero_grad()
# 每个 epoch 结束,保存 checkpoint
save_checkpoint(model, epoch)
和上一篇那一百行伪代码对比一下:模型部分一模一样,多出来的全是训练的部分。
forward 没变,加上了 loss、backward、optimizer、scheduler、梯度累积、梯度裁剪、混合精度。这些就是“训练”的全部内容。
七、回到开头的问题
上一篇那一百行,只有 forward。模型能算,但不会学。
这一篇补上了 backward 和 optimizer。旋钮从随机数变成有用的值,靠的就是这套流程:
前向传播 → 交叉熵损失 → 反向传播 → 梯度裁剪 → AdamW 更新 → 学习率余弦退火。
每一步都在拧旋钮。拧几百万次,模型就学会了。
MiniMind 的预训练跑完,模型学会了“话赶话”——给它一段文本,它能接着往下续写。但它还不会回答问题。下一篇讲的是:怎么把“接话机”教成“助手”。
八、本篇概念清单
| 概念 | 本篇交代到什么程度 |
|---|---|
| 交叉熵损失 | 讲透:公式、含义、为什么用它 |
| 反向传播 | 讲清:链式法则的直觉、PyTorch autograd |
| 梯度 | 讲透:方向和步长的区别 |
| AdamW | 讲透:一阶动量、二阶动量、权重衰减 |
| 学习率调度 | 讲透:余弦退火、为什么开头大结尾小 |
| 梯度裁剪 | 讲清:防爆炸、阈值 1.0 |
| 梯度累积 | 讲清:小显存模拟大 batch |
| 混合精度训练 | 讲清:bfloat16 和 fp16 的区别 |
| 训练循环 | 讲透:一个 step 的完整流程 |
本篇要牢记的只有三个词:交叉熵、AdamW、梯度。
九、思考题
- 交叉熵损失用 。如果正确答案的概率是 0.5,损失是多少?如果改成均方误差 ,损失又是多少?两者在“模型已经很自信但猜错”时的行为有什么不同?
- 学习率余弦退火是“从大到小”。如果反过来,从小到大,会怎样?
- 梯度累积默认 8 步。如果改成 1,batch size 不变,训练会有什么变化?
- bfloat16 的范围和 fp32 一样大,但精度低。为什么训练时“范围比精度更重要”?
备注:本篇的源码分析基于 MiniMind 仓库的 trainer/train_pretrain.py。训练参数(batch size、学习率、accumulation steps 等)来自 MiniMind 的默认配置,实际使用时可能根据硬件调整。