第 5 章 训练循环（train.py）· 下：逐行讲解

接 上篇。本篇逐函数剖析：种子、设备、余弦退火、AdamW、断点续训、训练循环 9 个细节。

5.5 可复现性：set_seed

def set_seed(seed: int) -> None:
    random.seed(seed)              # Python random
    np.random.seed(seed)           # NumPy
    torch.manual_seed(seed)        # CPU torch
    torch.cuda.manual_seed_all(seed)   # 所有 GPU torch
​

固定四个随机源。为什么是四个？因为 PyTorch 项目的随机性来自：

• 数据增强（如随机裁剪）—— Python random
• NumPy 操作（如 shuffle）—— NumPy
• 模型初始化、dropout —— CPU torch / GPU torch

少固定任何一个，结果就有偏差。

残留非确定性

同种子 + 同硬件 + 同代码 → 近似可复现，但不完全。原因：

• cuDNN 的某些卷积/attention 算子默认非确定。
• 多卡训练有 all-reduce 的浮点累加顺序问题。

完全复现要额外加：

torch.use_deterministic_algorithms(True)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
​

但会牺牲性能，通常只在调试/学术复现时开启。

5.6 设备选择：select_device

def select_device(name: str) -> torch.device:
    if name == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(name)
​

简洁的「auto + 显式」二选一。生产中还会区分多卡（cuda:0、cuda:1），本项目单卡够用。

5.7 学习率调度器：余弦退火（重难点）

def get_cosine_schedule_with_warmup(
    optimizer, num_warmup_steps, num_training_steps, min_lr_ratio=0.1,
):
    min_lr_ratio = max(0.0, min(1.0, min_lr_ratio))   # 钳制到 [0, 1]

    def lr_lambda(current_step: int) -> float:
        # 1) 预热阶段
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        # 2) 余弦退火阶段
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps)
        )
        cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
        return min_lr_ratio + (1.0 - min_lr_ratio) * cosine

    return LambdaLR(optimizer, lr_lambda)
​

工作原理

LambdaLR 接收一个函数 lr_lambda(step)，每步把基础学习率乘以这个函数的返回值。所以这个函数返回的是「当前学习率相对峰值的比例」，范围 [0, 1]：

实际 lr = optimizer.lr × lr_lambda(step) = 3e-4 × lr_lambda(step)
​

数学拆解

设 w = warmup_iters，T = max_iters，r = min_lr_ratio，当前步 s：

阶段 1 预热（s < w）：return s / w，从 0 线性升到 1。

max(1, num_warmup_steps) 是兜底：万一 warmup_iters=0，除以 0 会崩，用 1 兜底变成 0/1=0。

阶段 2 退火（s ≥ w）：

• progress = (s - w) / (T - w)：退火进度，从 0 到 1
• cos(π * progress)：从 cos(0)=1 到 cos(π)=-1
• cosine = 0.5 * (1 + cos(π * progress))：从 1 到 0
• r + (1-r) * cosine：起点 = r + (1-r)*1 = 1，终点 = r + (1-r)*0 = r

代入 r=0.1：终点 lr_ratio = 0.1 ✓，与配置一致。

曲线形状

lr_ratio
  ↑
1 │      ╭───╮  ← 峰值 (s=w 时达到 1)
  │     ╱     ╲
  │    ╱       ╲
  │   ╱         ╲
0.1 │ ╱           ╰───╯  ← 终点 (s=T 时降到 r)
  │╱
  └─────────────────────→ step
   0   w            T
​

为什么不直接用 HF 的同名函数

HuggingFace transformers 也提供了 get_cosine_schedule_with_warmup，但不支持 min_lr_ratio（它的终点固定是 0）。本项目自己实现一份，让退火终点可配——这是工业训练常用的小改进（终点留一点 lr，避免末期梯度完全消失、模型还能微调）。

5.8 优化器：AdamW

optimizer = AdamW(
    model.parameters(),
    lr=train_config.learning_rate,    # 3e-4
    betas=train_config.betas,         # (0.9, 0.95)
    weight_decay=train_config.weight_decay,   # 0.1
)
​

为什么是 AdamW 而不是 Adam

• Adam 把 weight decay 混在梯度里（L2 正则），与自适应学习率耦合，效果打折。
• AdamW 把 weight decay 解耦：先按 Adam 更新参数，再独立地按 w ← w - lr·wd·w 衰减权重。

这是 GPT-2/3 训练的标准做法。

betas=(0.9, 0.95) 的含义

Adam 维护梯度的一阶矩（动量）和二阶矩（梯度平方的指数移动平均）：

m_t = beta1 * m_{t-1} + (1 - beta1) * g_t        # 一阶矩
v_t = beta2 * v_{t-1} + (1 - beta2) * g_t²       # 二阶矩
update = lr * m_t / (sqrt(v_t) + eps)
​

• beta1=0.9：一阶矩衰减系数（标准值）。
• beta2=0.95：二阶矩衰减系数。比标准 Adam 的 0.999 小很多——这是 GPT 论文的经验值，让二阶矩更快跟随梯度变化，适合语言模型稀疏梯度的特点。

weight_decay=0.1 怎么作用

AdamW 的解耦权重衰减每步额外执行 w = w - lr · wd · w，把权重往 0 拉，防过拟合。

注意 bias 和 LayerNorm 的 gamma/beta 通常不衰减（HF 实现已处理）。

5.9 断点续训

start_step = 0
if args.resume:
    ckpt = torch.load(args.resume, map_location=device)
    model.load_state_dict(ckpt["model_state_dict"])
    optimizer.load_state_dict(ckpt["optimizer_state_dict"])
    scheduler.load_state_dict(ckpt["scheduler_state_dict"])
    start_step = ckpt["step"] + 1
​

关键认知

恢复训练不是只加载模型权重！优化器内部有动量、二阶矩估计，调度器有当前步数。三者必须一起恢复：

续训的典型场景

1. 训练中断：停电、OOM、Ctrl-C。--resume checkpoints/gpt_step2000.pt 从 2001 步继续。
2. 继续训更长：原本 5000 步不够，想加 5000 步。改 --max-iters 10000 --resume .../gpt_step5000.pt。
3. 换硬件：在 A100 上训到一半，转到另一台 A100 上继续（同架构）。

5.10 训练循环 9 个细节

model.train()                          # ① 训练模式
while step < train_config.max_iters:
    try:
        x, y = next(data_iter)         # ② 无限数据
    except StopIteration:
        data_iter = iter(dataloader)
        x, y = next(data_iter)

    x = x.to(device, non_blocking=True)  # ③ 异步拷贝
    y = y.to(device, non_blocking=True)

    outputs = model(input_ids=x, labels=y)  # ④ 自动 loss
    loss = outputs.loss
    optimizer.zero_grad(set_to_none=True)   # ⑤ None 替代 0
    loss.backward()

    torch.nn.utils.clip_grad_norm_(model.parameters(), train_config.grad_clip)  # ⑥ 梯度裁剪

    optimizer.step()                    # ⑦ 先 step
    scheduler.step()                    # ⑦ 后 schedule

    running_loss += loss.item()         # ⑧ 同步开销
    step += 1
    ...

    if step % save_iter == 0 or step == max_iters:   # ⑨ 双触发保存
        save_checkpoint(...)
​

① model.train() vs model.eval()

• train()：dropout 生效、BatchNorm 用 batch 统计。
• eval()：dropout 关闭、用全局统计。

训练循环开头必须切 train()。本项目无验证集，只在开头切一次。

② try/except StopIteration 实现「无限数据」

DataLoader 迭代完一个 epoch 会抛 StopIteration。捕获后重建迭代器，实现「按步数而非 epoch 控制训练」——LM 训练的标准模式。

③ non_blocking=True

配合 pin_memory=True，让 CPU→GPU 数据拷贝异步进行，可与上一步反向传播重叠。

④ model(input_ids=x, labels=y) 自动算 loss

GPT2LMHeadModel 的 forward 接受 labels 时会内部算交叉熵 loss，返回 outputs.loss。比手写 F.cross_entropy(...) 方便。

⑤ zero_grad(set_to_none=True)

PyTorch 2.0+ 推荐：把梯度置 None 而非填 0。省内存、优化器可跳过冻结参数。

⑥ clip_grad_norm_：防梯度爆炸

把所有参数梯度当作一个大向量，若 L2 范数超过 max_norm（默认 1.0），按比例缩放。Transformer 训练的「保险丝」。

注意是 clip_grad_norm_（按总范数），不是 clip_grad_value_（按每元素绝对值）。

⑦ optimizer.step() 在前、scheduler.step() 在后

顺序不能反。调度器要记录「这一步用过的 lr」。

⑧ loss.item() 会触发同步

.item() 把 GPU 标量拷到 CPU，强制 CPU/GPU 同步。每步都打 tqdm 会拖慢训练，严格优化时可改成每 N 步取一次。

⑨ 双触发条件的 checkpoint 保存

if step % save_iter == 0 or step == max_iters:
​

or 后半句保证训练正常结束的最后一步一定保存，即使 max_iters 不是 save_iter 的整数倍。

5.11 最终保存：final 模型

final_path = Path(ckpt_dir) / "gpt_final.pt"
torch.save({
    "step": step,
    "model_state_dict": model.state_dict(),
    "gpt_config": gpt_config.__dict__,
    "train_config": train_config.__dict__,
    "loss": loss.item(),
}, final_path)
​

final 与中途 checkpoint 的区别

final 不存优化器/调度器——它只用来推理，不再训练，存了浪费空间（优化器状态与模型权重等大）。

5.12 动手实验

1. 画学习率曲线：训练循环里加每 100 步打印 scheduler.get_last_lr()[0]，对照 5.7 节曲线是否一致。

2. 断点续训验证：训练 100 步，记下 loss；中断，用 --resume 继续 100 步，对比与连续训 200 步的 loss 是否吻合。

3. 梯度裁剪对比：把 grad_clip 改成 0.01（过小），观察 loss 是否停滞或 NaN。

4. 优化器状态重要性：训 200 步存 checkpoint；写脚本只加载 model_state_dict 不加载 optimizer，继续训练，观察前几十步 loss 是否抖动更厉害。

5. 思考题：为什么 optimizer.step() 必须在 scheduler.step() 之前？

6. 进阶：在日志点加 TensorBoard：

   from torch.utils.tensorboard import SummaryWriter
   writer = SummaryWriter("runs/exp1")
   writer.add_scalar("loss", avg_loss, step)
   writer.add_scalar("lr", current_lr, step)
   ​

5.13 小结

• 训练 8 阶段：配置 → 种子/设备 → 数据 → 模型 → 优化器/调度器 → 续训 → 循环 → 保存。
• set_seed 固定 4 个随机源；select_device 简单的 auto/显式选择。
• get_cosine_schedule_with_warmup 自定义实现「预热 + 余弦退火 + 可配终点比例」。
• AdamW 用 betas=(0.9, 0.95)（GPT 经验值）+ weight_decay=0.1（解耦）。
• 断点续训必须同时恢复模型、优化器、调度器、step。
• 训练循环 9 个细节：train 模式、StopIteration 重启、non_blocking、自动 loss、set_to_none、clip_grad_norm、step 顺序、loss.item 同步、双触发保存。

5.14 下一章

模型训完了，去 第 6 章 推理与采样 看怎么让它生成文本。