训练循环与评估

不会测量的 loop，一定会撒谎。这节课构建真正驱动 GPT 模型训练的那条 loop：带 weight decay 分组的 AdamW、warmup + cosine 学习率计划、calc_loss_batch helper、在 held-out 数据上跑的 evaluate_model、每隔 K 步就跑一次的 generate_and_print_sample 定性探针，以及一份可画图的 JSONL loss 日志。以后你搭的每个 decoder LLM，骨架都离不开这套东西。

类型： Build 语言： Python 前置要求： 第 19 阶段第 30-35 课 预计时间： ~90 分钟

学习目标

• 构建一条训练 loop，用正确的 input/target 对齐方式为 next-token prediction 计算 cross-entropy loss。
• 配置 AdamW，让 weight decay 只落在权重张量上，而不落在 LayerNorm 和 bias 上。
• 实现带线性 warmup 与 cosine decay 的学习率计划，并能读出每一步实际 LR。
• 在 held-out split 上用 evaluate_model 跑评估，让 eval loss 可以跨 run 对比。
• 每隔 K 步用 generate_and_print_sample 生成一段定性样本，在 loss curve 先崩之前抓到异常。
• 把每一步的 loss 持久化到 JSONL，这样训练日志能被重载、绘图和交付。

问题背景

一个训练脚本如果只会打印 loss，但除此之外什么都不做，会在三处同时失明。它不知道 loss 降下来到底是不是对的（模型也许只是过拟合训练集）；它不知道发散是不是已经开始（有的 run 会先 spike 一步再恢复，有的则一步直接崩掉）；它更不知道模型到底学会了什么（loss 只是标量，样本生成出来才像人话）。这三类失败，只有肯测才看得到。

本课的 loop 从三个方向测：

• 每一步都记录训练 batch 上的 loss
• 每隔 K 步在 held-out batch 上记录一次 loss
• 每隔 K 步对固定 prompt 生成一段 continuation

最终日志落到 JSONL，上面写的就是 loop 自己留下的证词。

核心概念

flowchart TB
  D[(Token 张量<br/>训练 + 验证)] --> B[构造 batch<br/>输入 + 目标左移一位]
  B --> F[前向<br/>logits]
  F --> L[交叉熵<br/>展平 batch 和时间维]
  L --> Bw[反向传播]
  Bw --> Cg[梯度范数裁剪]
  Cg --> Step[AdamW 步进]
  Step --> Sched[Cosine 学习率调度]
  Sched --> JL[追加步记录<br/>到 losses.jsonl]
  JL --> Probe{是否为探测步?}
  Probe -- 是 --> Eval[在验证集上 evaluate_model]
  Probe -- 是 --> Sample[generate_and_print_sample]
  Probe -- 否 --> Next[下一步]
  Eval --> Next
  Sample --> Next

真正不那么显然的两块，是 loss 对齐和 AdamW 的 decay 分组。

Loss 对齐

模型要在每个位置预测“下一个 token”。若输入 batch 是 [t0, t1, t2, t3]，那目标 batch 必须是 [t1, t2, t3, t4]。cross-entropy 是在展平后的 (batch * seq, vocab) 上对 (batch * seq,) 的目标做计算。若你忘了做左移，模型学的就是“预测自己”，loss 可以轻松收敛到 0，但模型实际上什么都没学到。

AdamW 的 Decay Split

weight decay 要正则化的是权重张量，不是归一化层的 scale，也不是 bias。把 decay 落到 LayerNorm scale 上，会慢慢把 scale 压到 0，直接毁掉归一化。落到 bias 上通常数学上问题不大，但纯属浪费。标准做法是：所有矩阵形状的张量（线性层权重、embedding 表）放 decay 组；看起来像 scale 或 shift 的参数全部放 no-decay 组。

Warmup + Cosine Schedule

warmup 会在最初几百步把学习率从 0 线性爬到目标值，让 optimizer state 有时间长出来。之后 cosine decay 再把学习率慢慢衰回 0，给训练末期一个更细的步长。这一套是开源 LLM 训练里最常见的 schedule，因为它能显著减少“前一千步”和“最后一千步”最脆的阶段。

Held-Out 评估

evaluate_model 会在 validation split 上跑固定批次数，累积 loss，再除以 batch 数返回均值。不开梯度，不开 dropout。只要 seed 和 split 固定，这个数就是可复现的。训练 loss 和 held-out loss 并排看，才看得出过拟合。

定性采样是更早的预警

一个训练 loss 看起来一路往下的模型，也可能生成出来全是同一个 token，那它就是坏了。反过来，一个 loss 曲线看起来没那么漂亮的模型，也可能已经能吐出明显更成形的词。定性探针常常比单个标量更早暴露问题。

cap-training-loop

动手构建

code/main.py 会实现：

• make_batches(token_ids, batch_size, context_length)：把长 token tensor 切成 input/target 对
• calc_loss_batch(model, inputs, targets)：前向、展平并返回标量 cross-entropy
• evaluate_model(model, val_loader, max_batches)：无梯度跑固定数量 validation batch，并返回平均 loss
• generate_and_print_sample(model, prompt, max_new_tokens)：调用第 35 课的 generation 函数，在固定 prompt 上打印样本
• build_param_groups(model, weight_decay)：生成 AdamW 需要的两组参数
• cosine_with_warmup(step, warmup_steps, total_steps, max_lr, min_lr)：给定 step 返回当前 LR
• train(...)：主训练 loop，负责写 outputs/losses.jsonl，并每隔 eval_every 步打印 eval loss 与样本
• 一个 demo：在合成数据上训练一个 tiny model，跑少量 step，输出 JSONL 日志，并在 probe 点打印 eval loss 与生成样本；CPU 下不到一分钟

运行方式：

python3 code/main.py

输出包括：逐 step 的 loss、每个 probe step 的 eval loss、每个 probe step 的生成样本，以及最终的 outputs/losses.jsonl。这份 JSONL 可以逐行 json.loads 再画图。

Stack

• torch：负责 autograd、optimizer 和模块系统
• main.py 本地重实现了第 35 课的 GPTModel 及相关模块

生产里常见的三个模式

Gradient norm clipping 不是可选项。 某个坏 batch（异常数据、学习率尖峰、数值边角）完全可能给你一记超大梯度，直接抹掉几个小时训练成果。backward 之后、step 之前加一发 torch.nn.utils.clip_grad_norm_(params, max_norm=1.0)，能把 optimizer 控在安全区间。1.0 是多数配置都扛得住的默认值。

可恢复的 JSONL 日志，而不是 pickle。 每一步一条 \{"step": int, "train_loss": float, "lr": float\} 的 JSONL 记录非常耐摔：中途崩了，文件还是可读；可以 grep，可以几十行 Python 画图，也可以直接从最后一条恢复。反过来，pickle 状态和模块布局绑得太死，一重构就容易废。

Eval batch 要来自固定切片。 validation token 在脚本启动时就要切好，而不是边跑边抽。可复现的前提，就是每次 run 看到的 eval batch 完全一致。否则你比较的不是模型差异，而是 batch 顺序噪声。

实际使用

• 本课的 loop 与真实 124M 模型训练用的是同一骨架。把合成 token tensor 换成 datasets 风格的 loader，主体逻辑根本不用改。
• JSONL 日志本身就是交付物，它把“训练过程”变成证据。下一课会直接拿它来对比 freshly trained checkpoint 与 pretrained one。
• 定性 sample probe 没有任何单一标量能完全替代。

练习

1. 给 weight_decay_groups() 写单测，确认 scale / bias 落进 no-decay 组，而线性层和 embedding 权重落进 decay 组。
2. 把合成随机 token 换成一个小文本文件的 bytes，让 demo 在真正可读文本上训练。确认生成样本使用的是文件里出现过的字符。
3. 给 cosine schedule 加一个 min_lr = 0.1 * max_lr 的下限，再画一次图。
4. 除了 JSONL 日志，每逢 eval_every 再存一份 checkpoint。再加一个 resume_from 开关，能读回 model state 和 optimizer state。
5. 每步记录 throughput（tokens per second），确认它稳定落在一个窄带内。

关键术语

英文	大家嘴上怎么说	它实际指什么
Loss alignment	“Shift by one”	输入位置 0..T-1，目标位置 1..T；cross-entropy 在展平后的形状上计算
Decay split	“Two groups”	AdamW 接收两组参数：带 decay 的矩阵权重，和不带 decay 的 scale/bias
Warmup	“Ramp”	学习率先从 0 爬到目标值，让 optimizer state 有时间稳定
Eval batches	“Held-out batches”	validation token 的一段固定切片，在脚本启动时就确定好，每次 probe 都完全一致
Qualitative probe	“Sample print”	每隔 K 步用固定 prompt 生成一小段文本，用来抓 loss 看不出的坏模式

延伸阅读

• 第 19 阶段第 35 课：本训练 loop 驱动的模型
• 第 19 阶段第 37 课：如何把预训练权重装进同一套模型
• 第 10 阶段第 04 课：在真实数据上的 pretraining 流程
• 第 10 阶段第 10 课：超出 cross-entropy loss 的更完整 eval 面