训练循环与评估

训练循环与评估

本节摘要:不测量的训练循环是会撒谎的循环。本节为第 33 节组装好的 GPT 模型配上训练循环:AdamW(带权重衰减分组)、线性预热 + 余弦退火的学习率调度、calc_loss_batch 损失计算助手、在留出集上跑 evaluate_model 的评估遍、每隔 K 步用 generate_and_print_sample 做定性探测、以及一份可事后绘图的 losses.jsonl 训练日志。损失对齐(输入与目标错位一位)、AdamW 衰减分组(矩阵张量衰减、缩放/偏置不衰减)是两处不显然的陷阱,本节一并讲清。这套骨架能训练你今后构建的每一款解码器 LLM。

对应原课程:Phase 19 · Lesson 36 · training-loop-eval(原英文 phases/19-capstone-projects/36-training-loop-eval/docs/en.md)。本节属「从零构建 GPT」赛道第七节。

学习目标

阅读完本节,你应当能够:

1. 构建训练循环,正确对齐输入与目标(下一 token 预测的位移),用交叉熵计算损失。
2. 配置 AdamW,把权重衰减只施加在矩阵张量上,LayerNorm 缩放与偏置不衰减。
3. 实现线性预热 + 余弦退火的学习率调度,并在任意步数读出精确学习率。
4. 在留出集上用 evaluate_model 评估,使评估损失跨运行可比。
5. 每 K 步用 generate_and_print_sample 生成定性样本,赶在损失曲线之前抓到发散。
6. 把每步损失持久化到 JSONL,使训练日志可重载、可绘图、可交付。

一、问题与直觉

一个只打印损失、其他什么都不做的训练脚本会以三种方式失败。它无法告诉你损失下降的原因是否正确(模型可能只过拟合训练集,根本没学到泛化);它无法告诉你发散是否正在开始(损失可能某一步尖峰后恢复,也可能某一步尖峰后崩溃);它无法告诉你模型学到了什么(损失是标量,而生成的样本是一段话)。这三种失败只有在循环去测量时才暴露。

本节的循环用三种方式测量:每步在训练批上算损失;每 K 步在留出批上算损失;每 K 步从固定提示生成一段续写。训练日志落到 JSONL,这份产物就是循环的证词。

两处不显然的件是损失对齐与AdamW 衰减分组。

损失对齐

模型在每个位置预测下一个 token。若输入批是 [t0, t1, t2, t3],目标批必须是 [t1, t2, t3, t4]。交叉熵在展平的 (batch*seq, vocab) logits 与展平的 (batch*seq,) 目标上计算。忘记位移,你就在训练模型预测自己——损失收敛到零,却什么有用的都没学到。

AdamW 衰减分组

权重衰减正则权重张量,但不该施加在归一化缩放或偏置上。把衰减放在 LayerNorm 缩放上会把缩放慢慢推向零、破坏归一化;把衰减放在偏置上数学无害但浪费算力。标准分组是:矩阵形张量(线性权重、嵌入表)衰减,任何像缩放或位移的张量不衰减。

预热加余弦退火

预热在几百步内把学习率从零爬到目标,让优化器状态有时间填满;余弦退火在剩余步数内把学习率降回接近零,使最后阶段用小步长精调权重。两者结合是开源 LLM 训练中最常见的调度,因为它消除了前一千步与最后一千步里大部分脆弱时刻。

留出评估与定性探测

evaluate_model 跑固定几批验证集、累计损失、除以批数返回,无梯度、无 dropout——给定相同种子与划分,数字跨运行可复现。把留出损失与训练损失并报,就是抓过拟合的方式。

定性探测更快:一个训练损失漂亮下降、但生成样本全是同一个 token 的模型是坏的;一个损失曲线看着平、但生成样本锐化成连贯词的模型在学。定性探测比读整条曲线更快,还能抓到标量漏掉的模式。

二、从零实现

code/main.py 实现:

• make_batches(token_ids, batch_size, context_length):把长 token 张量切成输入/目标对。
• calc_loss_batch(model, inputs, targets):前向、展平、返回标量交叉熵。
• evaluate_model(model, val_loader, max_batches):在无梯度下迭代固定几批验证集,返回平均损失。
• generate_and_print_sample(model, prompt, max_new_tokens):在固定提示上跑第 33 节的生成函数并打印。
• build_param_groups(model, weight_decay):产出两组 AdamW 参数列表。
• cosine_with_warmup(step, warmup_steps, total_steps, max_lr, min_lr):返回给定步数的学习率。
• train(...):跑循环、持久化 outputs/losses.jsonl、每 eval_every 步打印评估损失与样本。

核心骨架:

def train(model, train_loader, val_loader, steps, warmup, total,
          lr_max, lr_min, eval_every):
    opt = AdamW(build_param_groups(model, weight_decay=0.1),
                lr=lr_max, betas=(0.9, 0.95))
    losses = []
    for step in range(steps):
        x, y = next(train_loader)                       # input 与 target 已错位一位
        lr = cosine_with_warmup(step, warmup, total, lr_max, lr_min)
        for g in opt.param_groups: g["lr"] = lr
        loss = calc_loss_batch(model, x, y)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step(); opt.zero_grad()
        losses.append({"step": step, "train_loss": loss.item(), "lr": lr})
        if step % eval_every == 0:
            print("eval", evaluate_model(model, val_loader, max_batches=5))
            generate_and_print_sample(model, "Once upon")
    jsonl_dump(losses, "outputs/losses.jsonl")
​

损失计算的关键是展平:

def calc_loss_batch(model, inputs, targets):
    logits = model(inputs)                              # (B, T, vocab)
    B, T, V = logits.shape
    return F.cross_entropy(logits.view(B*T, V), targets.view(B*T))
​

参数分组的判定逻辑:

def build_param_groups(model, weight_decay):
    decay, no_decay = [], []
    for name, p in model.named_parameters():
        if not p.requires_grad: continue
        if p.ndim == 2 or "embed" in name: decay.append(p)      # 矩阵/嵌入 -> 衰减
        else: no_decay.append(p)                                # 缩放/偏置 -> 不衰减
    return [{"params": decay,    "weight_decay": weight_decay},
            {"params": no_decay, "weight_decay": 0.0}]
​

设计要点:用 p.ndim == 2 判定矩阵张量是通用启发式——二维权重衰减,一维缩放/偏置不衰减。嵌入表虽是二维也命中,符合预期。

三、框架对比

本节手写的「预热 + 余弦 + 梯度裁剪 + AdamW 衰减分组」是开源 LLM 训练的事实标准。换成 HuggingFace Trainer,只需配置 optim="adamw_torch"、lr_scheduler_type="cosine"、warmup_steps、max_grad_norm,底层逻辑与手写一致——但黑盒后你不再清楚每步到底发生了什么。PyTorch Lightning 把这些包进 Trainer.fit 的回调里,CosineAnnealingWarmRestarts 提供类似调度。手写一遍的价值在于:当训练在凌晨三点崩溃时,你知道每一步该看哪里。

四、可复用产物

code/main.py 是独立训练脚本,demo 用合成 token 在 CPU 上几十秒跑完一节小训练,写出 outputs/losses.jsonl,每条形如 {"step":int,"train_loss":float,"lr":float},可用 json.loads 逐行加载绘图。把合成 token 张量换成 datasets 式加载器,循环原样训练 124M 模型;JSONL 日志把一次训练跑变成证据,下一节用它比较从零训练的检查点与预训练检查点。

五、生产模式

三模式把教科书循环变成能过夜的循环。梯度范数裁剪不可省:坏批(异常数据、学习率尖峰、数值边界)产生巨大梯度,抹掉几小时训练;clip_grad_norm_(params, max_norm=1.0) 在 backward 后 step 前把优化器锁在安全范围,1.0 是多数配置能存活的默认。可恢复的 JSONL 日志而非 pickle 状态:每步损失记成 JSONL 行是持久的——任何崩溃留下可读产物、可 grep、可用三十行 Python 绘图、可读最后一步恢复训练;pickle 状态把你绑死在产出文件的模块布局上,重一次构就崩。评估批来自固定切片:验证 token 在脚本启动时一次性切批,而非即时切——可复现性依赖评估批逐运行一致,否则两次运行的评估损失比较里混进了批洗牌的噪声。

六、练习

1. 分组单测:给 build_param_groups 加单测,确认缩放与偏置落在不衰减组、线性与嵌入权重落在衰减组。
2. 真文本:把合成随机 token 换成小文本文件的字节,让 demo 训练在可读内容上,验证生成样本用了文件中出现的字符。
3. 学习率地板:给余弦调度加 min_lr = 0.1 * max_lr 的地板,重新绘曲线。
4. 检查点恢复:每 eval_every 步除 JSONL 外再存检查点,加 resume_from 标志重载模型与优化器状态。
5. 吞吐日志:在损失旁记每步吞吐(token/秒),确认它稳在一个带内。

本节要点回顾

1. 不测量就撒谎:训练批损失、留出批损失、生成样本三管齐下,缺一都会漏掉某类失败。
2. 损失对齐:输入错位一位当目标,忘位移会把训练变成「预测自己」,损失归零却没学到。
3. AdamW 衰减分组:矩阵张量(线性权重、嵌入)衰减,缩放/偏置不衰减,否则破坏 LayerNorm。
4. 预热 + 余弦:前几百步线性爬、剩余步数余弦降,开源 LLM 训练的事实标准。
5. 留出评估可复现:固定切片、无梯度、无 dropout,逐运行可比是抓过拟合的前提。
6. 定性探测早于标量:从固定提示生成续写,抓到损失曲线漏掉的模式崩溃。
7. JSONL 日志即证据:可 grep、可绘图、可恢复,优于 pickle 的脆弱状态。

下一节,我们做「加载预训练权重」——从 safetensors 文件把 GPT-2 风格权重灌进这套架构,逐张量走参数名映射,并生成续写证明加载成功。