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

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

- 作者：David（道雾轩）
- 专栏：nanoGPT 源码解读（https://daiw.org/manual/nanogpt.md）
- 最后更新：2026-07-13
- 原文：https://daiw.org/manual/nanogpt/training-loop
- 转载与引用：请注明出处并附原文链接（https://daiw.org/about/copyright）

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

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

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

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

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

主循环里用它：

```python
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 行）：

```python
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 行）是这样算的：

```python
@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 抖动大，多批平均才靠谱。这就是[数据那站](https://daiw.org/manual/nanogpt/data-prepare-char)留验证集的用处：`train loss` 和 `val loss` **并排打印**，一旦 val 不降反升，就是过拟合的信号。
- **checkpoint 存什么**：模型权重、优化器状态、`model_args`、`iter_num`、`best_val_loss`、`config`——**全套家当**。所以[断点续训(resume)](https://daiw.org/manual/nanogpt/train-setup)才能精确接上：不光模型，连优化器的动量、练到第几步都恢复。

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

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

```python
# 梯度累积:跑 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 训练更稳，但一批太大显存装不下。**梯度累积**的主意是：把一个「大 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](https://daiw.org/manual/nanogpt/get-batch)），让数据搬运和计算重叠，GPU 不空等。

### 讲究二：混合精度与 GradScaler

`with ctx:` 就是[配置那章](https://daiw.org/manual/nanogpt/train-setup)准备的 **autocast**——让前向计算用 bf16/fp16 等低精度跑，更快更省显存。但 **fp16 有个坑**：动态范围窄，小梯度会**下溢成 0**，训练废掉。

`GradScaler` 专治这个：

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

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

### 讲究三：梯度裁剪

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

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

### 讲究四：清空梯度

```python
optimizer.zero_grad(set_to_none=True)
```

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

## 结尾：计数与终止

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

外加每 `log_interval` 步打印一次 loss、耗时、**MFU**（模型算力利用率，衡量「GPU 跑得有多满」，细节留到[进阶那站](https://daiw.org/manual/nanogpt/scaling-and-speed)）。`iter_num` 超过 `max_iters` 就收工。

<Callout type="info">
  把这一步连起来看，它就是[第一站流水线图](https://daiw.org/manual/nanogpt/what-is-nanogpt)的代码化：**get_batch 取 (x，y) → forward 出 logits、算 loss → backward 得梯度 → 裁剪 → optimizer 更新参数**，循环几千上万次，loss 一路下探，模型就从吐乱码，慢慢学会写莎士比亚。**整个「学习」，就浓缩在这个 while 循环里。**
</Callout>

## 模型会学了，该让它说话

训练循环通了——模型能学了。练好之后，怎么让它**开口**生成新文本？这是第四站。我们逐行读那个自回归的 `generate`，以及把它串起来跑的 `sample.py`。👉 [采样：让模型开口](https://daiw.org/manual/nanogpt/generate-and-sample)

## 参考文献

- [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