19 KiB

Raw Blame History

第 14 课：学习率、梯度累积与混合精度

这一课只解决一个问题：MiniMind 的训练循环如何把多个 micro-batch 的 loss 变成一次稳定的参数更新。

0. 本节主线

MiniMind 的一个训练 micro-step 是：

取一个 batch
-> 把 input_ids / labels 放到 device
-> 按当前全局 step 设置 learning rate
-> autocast forward 得到 lm loss 和 aux loss
-> loss 除以 accumulation_steps
-> backward，把梯度累积到参数上
-> 每 accumulation_steps 个 micro-step 做一次 optimizer step
-> step 前 unscale + clip grad
-> step 后 scaler.update + zero_grad
-> epoch 结束时处理不足 accumulation_steps 的尾巴

这一课的核心不是“调用 loss.backward()”，而是看懂这四件事如何配合：

learning rate schedule
gradient accumulation
mixed precision / GradScaler
gradient clipping + optimizer.step

1. 原理讲解

1.1 一个训练 step 到底更新了什么

模型训练的目标是让参数 \theta 沿着 loss 下降的方向移动。

如果一个 batch 的训练损失是 L(\theta)，反向传播计算的是：


g = \nabla_{\theta} L(\theta)

优化器做的事可以粗略理解成：


\theta_{\text{new}} = \theta_{\text{old}} - \eta \cdot \mathrm{Update}(g)

这里：

符号	形状	含义
`\theta`	每个参数自己的形状	模型参数
`L`	标量	当前 batch 的训练损失
`g`	和参数同形状	反向传播得到的梯度
`\eta`	标量	learning rate

对 MiniMind 来说，训练 loss 通常是：


L_{\text{train}} = L_{\text{lm}} + L_{\text{aux}}

其中 L_{\text{lm}} 是 next-token loss，L_{\text{aux}} 是 MoE router 的辅助损失。非 MoE 模型里，aux_loss 通常就是 0。

1.2 学习率不是固定不变的

Learning rate 控制每次参数更新走多大一步。太大容易震荡，太小训练太慢。

MiniMind 的 get_lr 用的是余弦衰减，但不是衰减到 0，而是从初始学习率衰减到 10% 初始学习率。

公式是：


\eta_s
= \eta_0 \left(0.1 + 0.45\left(1 + \cos\left(\frac{\pi s}{T}\right)\right)\right)

其中：

符号	形状	含义
`s`	标量	当前全局 micro-step
`T`	标量	总 micro-step 数
`\eta_0`	标量	命令行传入的初始学习率
`\eta_s`	标量	当前 step 实际使用的学习率

代入两个边界：


\begin{aligned}
\eta_0^{\text{actual}}
&= \eta_0 \left(0.1 + 0.45(1 + 1)\right)
= \eta_0 \\
\eta_T
&= \eta_0 \left(0.1 + 0.45(1 - 1)\right)
= 0.1\eta_0
\end{aligned}

也就是说，它从 learning_rate 平滑降到 0.1 * learning_rate。

1.3 梯度累积解决什么问题

显存放不下很大的 batch 时，可以用多个小 batch 模拟一个大 batch。

假设：

符号	形状	含义
`B`	标量	单个 micro-batch 的 batch size
`K`	标量	`accumulation_steps`
`B_{\text{eff}}`	标量	等效 batch size

单卡非 DDP 时：


B_{\text{eff}} = B \times K

如果是 DDP，多卡训练还要乘以进程数 W：


B_{\text{eff}} = B \times K \times W

为什么要把 loss 除以 K？

如果第 i 个 micro-batch 的原始 loss 是 L_i，MiniMind 用的是：


\tilde{L}_i = \frac{L_i}{K}

连续反向传播 K 次后，参数上累积的梯度是：


\begin{aligned}
g
&= \sum_{i=1}^{K}\nabla_{\theta}\tilde{L}_i \\
&= \sum_{i=1}^{K}\nabla_{\theta}\left(\frac{L_i}{K}\right) \\
&= \frac{1}{K}\sum_{i=1}^{K}\nabla_{\theta}L_i
\end{aligned}

这相当于取了 K 个 micro-batch 梯度的平均值，而不是把它们直接相加。否则 accumulation_steps 越大，梯度尺度越大，学习率也会被隐式放大。

1.4 混合精度做了什么

混合精度的目标是减少显存和加速计算。

MiniMind 的 forward 放在 autocast_ctx 里：

with autocast_ctx:
    res = model(...)
    loss = ...

这表示某些矩阵乘法可以用 float16 或 bfloat16 执行，而不是全部用 float32。

但 float16 动态范围小，梯度可能 underflow。GradScaler 的思想是先把 loss 放大：


L_{\text{scaled}} = c \cdot L

反向传播得到放大后的梯度：


g_{\text{scaled}} = c \cdot g

真正更新前再除回来：


g = \frac{g_{\text{scaled}}}{c}

MiniMind 只在 dtype == float16 时启用 GradScaler。如果是 bfloat16，代码仍然调用 scaler.scale(...)，但 scaler 是 disabled 状态，本质上是 no-op。

1.5 为什么要先 unscale 再 clip grad

梯度裁剪限制梯度范数，避免一次异常梯度把参数推太远。

如果最大范数是 C，裁剪可以理解成：


g_{\text{clipped}}
= g \cdot \min\left(1, \frac{C}{\lVert g \rVert_2}\right)

如果用了 GradScaler，参数上的梯度可能还是放大后的 g_{\text{scaled}}。这时直接 clip，裁剪的就不是原始梯度。

所以 MiniMind 的顺序是：

scaler.unscale_(optimizer)
-> clip_grad_norm_
-> scaler.step(optimizer)
-> scaler.update()

先 unscale，再 clip，最后 step。

1.6 变量形状

训练循环里最重要的张量形状是：

变量	形状	含义
`input_ids`	`[B, S]`	token id
`labels`	`[B, S]`	训练目标，忽略位置是 `-100`
`logits`	`[B, S, V]`	每个位置对词表的预测
`res.loss`	标量	causal LM loss
`res.aux_loss`	标量	MoE aux loss，非 MoE 时通常是 0
`loss`	标量	`(res.loss + res.aux_loss) / accumulation_steps`
`param.grad`	和参数同形状	当前累计梯度

注意：loss 是标量，不是 [B, S]。前面第 10 课已经讲过，cross entropy 会把有效 token 的 loss 聚合成一个标量。

2. 源码阅读顺序图

这节源码按这个顺序读：

trainer_utils.get_lr
-> train_pretrain.py 参数
-> autocast / GradScaler 初始化
-> train_epoch 前半段：lr + forward + loss / accumulation_steps
-> train_epoch 中段：backward + unscale + clip + step + zero_grad
-> train_epoch 末尾：不足 accumulation_steps 的尾批处理

对应文件：

trainer/trainer_utils.py
trainer/train_pretrain.py

先看 get_lr，是为了知道每个 step 用什么学习率。
再看 train_epoch，是为了知道 loss、梯度、optimizer step 如何串起来。
最后看初始化部分，是为了知道 autocast_ctx 和 GradScaler 是怎么决定的。

3. MiniMind 源码走读

第 1 步：学习率公式

File: trainer/trainer_utils.py:40-41

Read this to understand: MiniMind 的 learning rate schedule 不是 PyTorch scheduler，而是每个 step 手动计算。

Code/config/template excerpt:

def get_lr(current_step, total_steps, lr):
    return lr*(0.1 + 0.45*(1 + math.cos(math.pi * current_step / total_steps)))

This code shows:

current_step 是当前训练进度。
total_steps 是总训练步数。
返回值直接写入 optimizer 的 param group。
它是 cosine decay，终点是 0.1 * lr，不是 0。

第 2 步：训练参数决定 batch、lr、累积步数和裁剪阈值

File: trainer/train_pretrain.py:87-95

Read this to understand: 训练机制的关键超参都来自命令行参数。

Code/config/template excerpt:

parser.add_argument("--epochs", type=int, default=2, help="训练轮数")
parser.add_argument("--batch_size", type=int, default=32, help="batch size")
parser.add_argument("--learning_rate", type=float, default=5e-4, help="初始学习率")
parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
parser.add_argument("--accumulation_steps", type=int, default=8, help="梯度累积步数")
parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")

This code shows:

batch_size 是单个 micro-batch 的大小。
accumulation_steps 决定多少个 micro-batch 做一次 optimizer update。
learning_rate 是 schedule 的初始值。
dtype 决定 autocast 用 float16 还是 bfloat16。
grad_clip 决定梯度裁剪上限。

第 3 步：初始化 autocast 和 GradScaler

File: trainer/train_pretrain.py:119-138

Read this to understand: 混合精度由 autocast_ctx 和 GradScaler 两部分组成。

Code/config/template excerpt:

device_type = "cuda" if "cuda" in args.device else "cpu"
dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
...
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)

This code shows:

CPU 上不用 autocast，直接是 nullcontext()。
CUDA 上会进入 torch.cuda.amp.autocast(dtype=dtype)。
bfloat16 只用 autocast，不启用 GradScaler。
float16 同时使用 autocast 和 GradScaler。
optimizer 是 AdamW。

第 4 步：每个 micro-step 设置学习率

File: trainer/train_pretrain.py:24-34

Read this to understand: 学习率不是 epoch 级别更新，而是每个 micro-step 更新。

Code/config/template excerpt:

for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
    input_ids = input_ids.to(args.device)
    labels = labels.to(args.device)
    last_step = step
    lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

This code shows:

step 从 start_step + 1 开始，支持断点续训。
epoch * iters + step 是跨 epoch 的全局 step。
每个 param group 都会被写入当前 lr。
input_ids 和 labels 先搬到训练 device。

第 5 步：forward、loss 和梯度累积缩放

File: trainer/train_pretrain.py:35-40

Read this to understand: MiniMind 是先求完整 loss，再除以 accumulation_steps。

Code/config/template excerpt:

with autocast_ctx:
    res = model(input_ids, labels=labels)
    loss = res.loss + res.aux_loss
    loss = loss / args.accumulation_steps

scaler.scale(loss).backward()

This code shows:

res.loss 是语言模型主 loss。
res.aux_loss 是 MoE 辅助 loss。
loss / accumulation_steps 是为了让累积梯度保持平均尺度。
scaler.scale(loss).backward() 兼容 float16 AMP；如果 scaler disabled，就等价于普通 backward。

第 6 步：到达累积步数后才更新参数

File: trainer/train_pretrain.py:42-50

Read this to understand: 不是每个 micro-batch 都调用 optimizer step。

Code/config/template excerpt:

if step % args.accumulation_steps == 0:
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)

    scaler.step(optimizer)
    scaler.update()

    optimizer.zero_grad(set_to_none=True)

This code shows:

只有 step % accumulation_steps == 0 时才更新参数。
unscale_ 必须在 clip_grad_norm_ 前面。
scaler.step(optimizer) 内部会在 scaler enabled 时检查溢出。
optimizer.zero_grad(set_to_none=True) 清掉已使用的梯度，为下一组累积做准备。

第 7 步：日志里的 loss 要乘回去

File: trainer/train_pretrain.py:51-59

Read this to understand: 训练内部用的是缩放后的 loss，但日志要展示原始尺度。

Code/config/template excerpt:

current_loss = loss.item() * args.accumulation_steps
current_aux_loss = res.aux_loss.item() if res.aux_loss is not None else 0.0
current_logits_loss = current_loss - current_aux_loss
current_lr = optimizer.param_groups[-1]['lr']

This code shows:

loss 已经除过 accumulation_steps。
打日志时乘回 accumulation_steps，得到原始训练 loss 尺度。
current_logits_loss 是总 loss 减掉 aux loss。
current_lr 从 optimizer 里读当前实际学习率。

第 8 步：epoch 结束时处理尾巴

File: trainer/train_pretrain.py:75-80

Read this to understand: 如果最后剩下不足 accumulation_steps 的梯度，也要做一次 optimizer step。

Code/config/template excerpt:

if last_step > start_step and last_step % args.accumulation_steps != 0:
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)

This code shows:

如果 epoch 里真的跑过 batch，last_step > start_step。
如果最后一个 step 不是累积边界，说明还有没 step 的梯度。
这时需要 flush tail，否则最后几个 micro-batch 的梯度会被丢掉。

4. 本节必须会写 / 暂时不要求

必须会写：

MiniMind 的 cosine lr：


\eta_s
= \eta_0 \left(0.1 + 0.45\left(1 + \cos\left(\frac{\pi s}{T}\right)\right)\right)

梯度累积的 loss 缩放：


\tilde{L} = \frac{L}{K}

optimizer step 判断：


\text{should\_step}(s, K) =
\begin{cases}
\text{True}, & s \bmod K = 0 \\
\text{False}, & s \bmod K \ne 0
\end{cases}

尾批 flush 判断：


\text{flush\_tail} =
(s_{\text{last}} > s_{\text{start}})
\land
(s_{\text{last}} \bmod K \ne 0)

AMP update 顺序：

scale loss
-> backward
-> unscale grad
-> clip grad
-> optimizer step
-> scaler update
-> zero grad

暂时不要求：

1. DDP 梯度同步细节
2. checkpoint / resume 的完整状态恢复
3. torch.compile 对训练循环的影响
4. wandb / swanlab 日志工程
5. PPO / GRPO 里的 scheduler 版本

这一课只要求你能手写 pretrain/SFT 这类监督训练最小循环。

5. 手写模块

本节你要补的是：

course/impl/core/train_loop.py

5.1 补 `cosine_lr`

接口：

def cosine_lr(current_step: int, total_steps: int, base_lr: float) -> float:
    ...

对齐源码：

trainer/trainer_utils.py:40-41

你要实现的行为：

输入 current_step、total_steps、base_lr。
输出和 MiniMind get_lr 完全一致的 float。
不要调用原项目的 get_lr，要自己写公式。

5.2 补 `scale_loss_for_accumulation`

接口：

def scale_loss_for_accumulation(loss: torch.Tensor, accumulation_steps: int) -> torch.Tensor:
    ...

对齐源码：

trainer/train_pretrain.py:35-40

你要实现的行为：


\tilde{L} = \frac{L}{K}

其中 K 就是 accumulation_steps。

5.3 补 `should_step_optimizer`

接口：

def should_step_optimizer(step: int, accumulation_steps: int) -> bool:
    ...

对齐源码：

trainer/train_pretrain.py:42

你要实现的行为：

step=1, K=4 -> False
step=2, K=4 -> False
step=3, K=4 -> False
step=4, K=4 -> True

5.4 补 `should_flush_tail`

接口：

def should_flush_tail(last_step: int, start_step: int, accumulation_steps: int) -> bool:
    ...

对齐源码：

trainer/train_pretrain.py:75

你要实现的行为：

last_step=10, start_step=0, K=4 -> True
last_step=8,  start_step=0, K=4 -> False
last_step=0,  start_step=0, K=4 -> False

5.5 补 `train_one_epoch`

接口已经在骨架里。第一版只需要对齐监督训练主循环，不要写 DDP、wandb、checkpoint。

对齐源码：

trainer/train_pretrain.py:24-80

建议实现顺序：

1. 遍历 loader，得到 input_ids / labels。
2. 搬到 device。
3. 用 cosine_lr 设置 optimizer lr。
4. 进入 autocast_ctx。
5. 调 model(input_ids, labels=labels)。
6. loss = outputs.loss + outputs.aux_loss。
7. loss = scale_loss_for_accumulation(loss, accumulation_steps)。
8. scaler.scale(loss).backward()。
9. 如果 should_step_optimizer 为 True：
   unscale -> clip -> step -> update -> zero_grad。
10. epoch 结束后，如果 should_flush_tail 为 True：
    再执行一次 unscale -> clip -> step -> update -> zero_grad。

注意：

如果 scaler is None，你可以创建 torch.cuda.amp.GradScaler(enabled=False)，让 CPU 路径也能跑。
如果 autocast_ctx is None，可以用 contextlib.nullcontext()。
outputs.aux_loss 可能是 0 标量，不要特殊跳过。

6. 对齐测试

本节新增对齐测试：

course/impl/tests/test_train_loop_mechanics.py

运行命令：

cd /home/sun/minimind
python course/impl/tests/test_train_loop_mechanics.py

现在还没有实现时，这个测试会因为 NotImplementedError 失败。等你补完前四个小函数后，它应该打印类似：

lr_step_0_diff=0.000000000000
lr_step_1_diff=0.000000000000
...
scaled_loss=2.000000
optimizer_step_micro_steps=[4, 8]
train_loop_mechanics=passed

这个测试先不检查完整 train_one_epoch，因为完整训练循环需要模型、数据和 optimizer 共同参与。先把机制函数对齐，再进入阶段组装更稳。

另外可以观察原项目 tiny 训练 step：

cd /home/sun/minimind
python course/labs/trace_pretrain_step.py --batch_size 2 --max_length 48 --hidden_size 64 --num_hidden_layers 2

这个实验会打印：

input_ids.shape
labels.shape
logits.shape
loss
aux_loss
lr
grad_norm_before_clip
max_first_param_delta_after_step

7. 阶段组装

本节完成后，Pretrain 阶段会多出训练循环核心能力：

course/impl/core/train_loop.py::cosine_lr
course/impl/core/train_loop.py::scale_loss_for_accumulation
course/impl/core/train_loop.py::should_step_optimizer
course/impl/core/train_loop.py::should_flush_tail
course/impl/core/train_loop.py::train_one_epoch

它们以后会接到：

course/impl/train_pretrain_impl.py

阶段目标是：

dataset
-> DataLoader
-> CourseMiniMindForCausalLM
-> train_one_epoch
-> save course_pretrain checkpoint

当前 Pretrain 阶段还缺：

1. 教学版 CausalLM 完整 forward/loss
2. 教学版 pretrain dataset 或 dataset 适配
3. checkpoint 保存和加载
4. tiny pretrain 脚本组装

第 14 课只先把训练循环的机械部分写清楚。

8. 本节检查

MiniMind 的 get_lr 起点和终点分别是多少？
为什么梯度累积时要把 loss 除以 accumulation_steps？
如果 batch_size=4，accumulation_steps=8，单卡训练的等效 batch size 是多少？
autocast 和 GradScaler 分别解决什么问题？
为什么 clip_grad_norm_ 前要先 scaler.unscale_(optimizer)？
step % accumulation_steps == 0 控制的是什么？
为什么 epoch 结束时还要处理不足 accumulation_steps 的尾巴？
写出 input_ids、labels、logits、loss、param.grad 的形状。

9. 下一课

第 15 课讲 checkpoint 与断点续训。

下一课要解决：

保存哪些状态；
为什么只保存 model 不够；
optimizer / scaler / epoch / step 如何恢复；
SkipBatchSampler 如何跳过已经训练过的 batch；
普通权重文件和 resume checkpoint 有什么区别。

19 KiB Raw Blame History Unescape Escape

第 14 课：学习率、梯度累积与混合精度

目录

0. 本节主线

1. 原理讲解

1.1 一个训练 step 到底更新了什么

1.2 学习率不是固定不变的

1.3 梯度累积解决什么问题

1.4 混合精度做了什么

1.5 为什么要先 unscale 再 clip grad

1.6 变量形状

2. 源码阅读顺序图

3. MiniMind 源码走读

第 1 步：学习率公式

第 2 步：训练参数决定 batch、lr、累积步数和裁剪阈值

第 3 步：初始化 autocast 和 GradScaler

第 4 步：每个 micro-step 设置学习率

第 5 步：forward、loss 和梯度累积缩放

第 6 步：到达累积步数后才更新参数

第 7 步：日志里的 loss 要乘回去

第 8 步：epoch 结束时处理尾巴

4. 本节必须会写 / 暂时不要求

5. 手写模块

5.1 补 cosine_lr

5.2 补 scale_loss_for_accumulation

5.3 补 should_step_optimizer

5.4 补 should_flush_tail

5.5 补 train_one_epoch

6. 对齐测试

7. 阶段组装

8. 本节检查

9. 下一课

19 KiB

Raw Blame History

5.1 补 `cosine_lr`

5.2 补 `scale_loss_for_accumulation`

5.3 补 `should_step_optimizer`

5.4 补 `should_flush_tail`

5.5 补 `train_one_epoch`