训练主循环 · 一步到底怎么走 · nanoGPT 源码解读

逐行读 train.py 的主循环——学习率的 warmup + 余弦退火、梯度累积怎么用小显存模拟大 batch、混合精度的 GradScaler 在防什么、梯度裁剪、以及评估与 checkpoint。一步「前向—反向—更新」的全部细节。

训练主循环 · 一步到底怎么走

第三站的高潮。前面配置、优化器、数据都备齐了，这一章走进 train.py 的 while True 主循环，把「训练一步」的每个动作拆开。它比前面几章密，但每一块都是你在别处也会反复见到的通用套路。

开场：学习率怎么随步数起落

循环第一件事，是给这一步定学习率。get_lr（约 231–242 行）实现了经典的 warmup + 余弦退火：

def get_lr(it):
    # 1) 前 warmup_iters 步,线性升温
    if it < warmup_iters:
        return learning_rate * (it + 1) / (warmup_iters + 1)
    # 2) 超过 lr_decay_iters,固定在最小值
    if it > lr_decay_iters:
        return min_lr
    # 3) 中间,余弦曲线从 learning_rate 平滑降到 min_lr
    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # 1 → 0
    return min_lr + coeff * (learning_rate - min_lr)

三段式：

升温(warmup)：头 2000 步，学习率从 0 线性爬到峰值。为什么？训练刚开始参数是随机的，梯度乱、方差大，一上来就用大学习率容易崩；先小步试探，稳住了再加速。
余弦退火：之后学习率沿余弦曲线从峰值平滑降到 min_lr。前期大步快跑、后期小步精调，是久经考验的调度。
收尾：超过 lr_decay_iters 就固定在 min_lr。

主循环里用它：

lr = get_lr(iter_num) if decay_lr else learning_rate
for param_group in optimizer.param_groups:
    param_group['lr'] = lr

定期评估 + 存档

每隔 eval_interval 步，评估一次并考虑存档（约 263–286 行）：

if iter_num % eval_interval == 0 and master_process:
    losses = estimate_loss()
    print(f"step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
    if losses['val'] < best_val_loss or always_save_checkpoint:
        best_val_loss = losses['val']
        if iter_num > 0:
            checkpoint = { 'model': raw_model.state_dict(), 'optimizer': optimizer.state_dict(),
                           'model_args': model_args, 'iter_num': iter_num,
                           'best_val_loss': best_val_loss, 'config': config }
            torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))

estimate_loss（约 215–228 行）是这样算的：

@torch.no_grad()
def estimate_loss():
    out = {}
    model.eval()
    for split in ['train', 'val']:
        losses = torch.zeros(eval_iters)
        for k in range(eval_iters):
            X, Y = get_batch(split)
            with ctx:
                logits, loss = model(X, Y)
            losses[k] = loss.item()
        out[split] = losses.mean()
    model.train()
    return out

几个要点：

@torch.no_grad() + model.eval()：评估时不算梯度、关掉 dropout，只看模型当前水平，省算力也更准。算完 model.train() 切回训练态。
训练集和验证集各测 eval_iters 批取平均——单批 loss 抖动大，多批平均才靠谱。这就是数据那站留验证集的用处：train loss 和 val loss 并排打印，一旦 val 不降反升，就是过拟合的信号。
checkpoint 存什么：模型权重、优化器状态、model_args、iter_num、best_val_loss、config——全套家当。所以断点续训(resume)才能精确接上：不光模型，连优化器的动量、练到第几步都恢复。

核心：前向—反向—更新（带四个讲究）

真正「学」的那几行（约 292–314 行），浓缩了四个重要概念：

# 梯度累积:跑 gradient_accumulation_steps 个 micro-batch,累加梯度,才更新一次
for micro_step in range(gradient_accumulation_steps):
    with ctx:                                   # 混合精度上下文
        logits, loss = model(X, Y)
        loss = loss / gradient_accumulation_steps  # 缩放,凑成平均
    X, Y = get_batch('train')                   # 预取下一批(与反向重叠)
    scaler.scale(loss).backward()               # 反向(带梯度缩放)
# 梯度裁剪
if grad_clip != 0.0:
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
# 更新
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)

大 batch 训练更稳，但一批太大显存装不下。梯度累积的主意是：把一个「大 batch」拆成 gradient_accumulation_steps 个小批(micro-batch)，一个个前向+反向，梯度在参数上自动累加，攒够了才更新一次。等效于一次用了 batch_size × gradient_accumulation_steps 的大 batch，显存却只需装下一个 micro-batch。

loss = loss / gradient_accumulation_steps 是关键一步：因为梯度会累加 k 次，先把每次的 loss 除以 k，累加起来才等于「平均」而非「求和」，梯度尺度才对。

那句 X, Y = get_batch('train') 放在前向之后、反向之前，是故意的：趁 GPU 做反向传播时，异步预取下一批数据（配合上一章的 pin_memory），让数据搬运和计算重叠，GPU 不空等。

讲究二：混合精度与 GradScaler

with ctx: 就是配置那章准备的 autocast——让前向计算用 bf16/fp16 等低精度跑，更快更省显存。但 fp16 有个坑：动态范围窄，小梯度会下溢成 0，训练废掉。

GradScaler 专治这个：

scaler.scale(loss).backward()：反向前把 loss 乘一个大因子放大，梯度跟着放大、逃离下溢区；
scaler.unscale_(optimizer)：裁剪前先还原回真实尺度（否则裁的是放大后的假梯度）；
scaler.step(optimizer)：更新参数（内部会检查有没有 inf/nan，有就跳过这步）；
scaler.update()：动态调整那个放大因子。

（用 bf16 时 GradScaler 是空操作——bf16 动态范围够，不需要这套。）

讲究三：梯度裁剪

torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)

把所有梯度的总范数限制在 grad_clip（默认 1.0）以内：一旦某步梯度「爆」了（范数过大），就整体等比缩小。这是防止训练被个别异常大梯度一步带崩的安全带。

讲究四：清空梯度

optimizer.zero_grad(set_to_none=True)

更新完，清空梯度，准备下一轮累积。set_to_none=True 是把梯度置为 None 而非填 0——省一次内存写、略快。

结尾：计数与终止

iter_num += 1
local_iter_num += 1
if iter_num > max_iters:
    break

外加每 log_interval 步打印一次 loss、耗时、MFU（模型算力利用率，衡量「GPU 跑得有多满」，细节留到进阶那站）。iter_num 超过 max_iters 就收工。

把这一步连起来看，它就是第一站流水线图的代码化：get_batch 取 (x，y) → forward 出 logits、算 loss → backward 得梯度 → 裁剪 → optimizer 更新参数，循环几千上万次，loss 一路下探，模型就从吐乱码，慢慢学会写莎士比亚。整个「学习」，就浓缩在这个 while 循环里。

模型会学了，该让它说话

训练循环通了——模型能学了。练好之后，怎么让它开口生成新文本？这是第四站。我们逐行读那个自回归的 generate，以及把它串起来跑的 sample.py。👉 采样：让模型开口

参考文献

[1] Loshchilov & Hutter. “SGDR: Stochastic Gradient Descent with Warm Restarts.” ICLR 2017. arXiv:1608.03983. —— 余弦学习率退火。
[2] Micikevicius et al. “Mixed Precision Training.” ICLR 2018. arXiv:1710.03740. —— 混合精度与损失缩放(GradScaler)。
[3] Pascanu et al. “On the difficulty of training Recurrent Neural Networks.” ICML 2013. arXiv:1211.5063. —— 梯度裁剪。
[4] Hoffmann et al. “Training Compute-Optimal Large Language Models (Chinchilla).” 2022. arXiv:2203.15556. —— 学习率与迭代步数的设定思路。
源码：Andrej Karpathy, nanoGPT, train.py. https://github.com/karpathy/nanoGPT

训练主循环 · 一步到底怎么走