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

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

## 目录

- [0. 本节主线](#l14-mainline)
- [1. 原理讲解](#l14-principle)
- [2. 源码阅读顺序图](#l14-reading-order)
- [3. MiniMind 源码走读](#l14-source-walkthrough)
- [4. 本节必须会写 / 暂时不要求](#l14-must-write)
- [5. 手写模块](#l14-handwrite)
- [6. 对齐测试](#l14-alignment-test)
- [7. 阶段组装](#l14-stage-assembly)
- [8. 本节检查](#l14-check)
- [9. 下一课](#l14-next)

<a id="l14-mainline"></a>
## 0. 本节主线

MiniMind 的一个训练 micro-step 是：

```text
取一个 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()`”，而是看懂这四件事如何配合：

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

<a id="l14-principle"></a>
## 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` 里：

```text
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 的顺序是：

```text
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 聚合成一个标量。

<a id="l14-reading-order"></a>
## 2. 源码阅读顺序图

这节源码按这个顺序读：

```text
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 的尾批处理
```

对应文件：

```text
trainer/trainer_utils.py
trainer/train_pretrain.py
```

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

<a id="l14-source-walkthrough"></a>
## 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:

```python
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:

```python
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:

```python
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:

```python
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:

```python
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:

```python
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:

```python
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:

```python
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 的梯度会被丢掉。

<a id="l14-must-write"></a>
## 4. 本节必须会写 / 暂时不要求

必须会写：

1. MiniMind 的 cosine lr：

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

2. 梯度累积的 loss 缩放：

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

3. optimizer step 判断：

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

4. 尾批 flush 判断：

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

5. AMP update 顺序：

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

暂时不要求：

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

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

<a id="l14-handwrite"></a>
## 5. 手写模块

本节你要补的是：

```text
course/impl/core/train_loop.py
```

### 5.1 补 `cosine_lr`

接口：

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

对齐源码：

```text
trainer/trainer_utils.py:40-41
```

你要实现的行为：

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

### 5.2 补 `scale_loss_for_accumulation`

接口：

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

对齐源码：

```text
trainer/train_pretrain.py:35-40
```

你要实现的行为：

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

其中 $K$ 就是 `accumulation_steps`。

### 5.3 补 `should_step_optimizer`

接口：

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

对齐源码：

```text
trainer/train_pretrain.py:42
```

你要实现的行为：

```text
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`

接口：

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

对齐源码：

```text
trainer/train_pretrain.py:75
```

你要实现的行为：

```text
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。

对齐源码：

```text
trainer/train_pretrain.py:24-80
```

建议实现顺序：

```text
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 标量，不要特殊跳过。

<a id="l14-alignment-test"></a>
## 6. 对齐测试

本节新增对齐测试：

```text
course/impl/tests/test_train_loop_mechanics.py
```

运行命令：

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

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

```text
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：

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

这个实验会打印：

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

<a id="l14-stage-assembly"></a>
## 7. 阶段组装

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

```text
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
```

它们以后会接到：

```text
course/impl/train_pretrain_impl.py
```

阶段目标是：

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

当前 Pretrain 阶段还缺：

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

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

<a id="l14-check"></a>
## 8. 本节检查

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

<a id="l14-next"></a>
## 9. 下一课

第 15 课讲 checkpoint 与断点续训。

下一课要解决：

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