从零实现 nanoGPT（六）：完善训练流程，实现模型评估与断点续训

从零实现 nanoGPT（六）：完善训练流程，实现模型评估与断点续训

上一篇，我们完成了模型检查点的保存，并编写了独立的 sample.py，让训练好的 GPT 可以重新加载并生成文本。

到目前为止，我们已经实现了一个简单但完整的字符级语言模型：

文本数据
   ↓
字符编码
   ↓
GPT 模型
   ↓
计算 Loss
   ↓
反向传播
   ↓
更新参数
   ↓
保存模型
   ↓
生成文本

但是，当我们准备让模型训练更长时间时，就会遇到几个新的问题：

1. 训练 Loss 不断下降，是否意味着模型的预测能力一定在提高？
2. 最后一次训练得到的模型，是否一定是效果最好的模型？
3. 如果训练过程中程序意外退出，能否从上次的位置继续训练？
4. 当显存或内存不足以支持更大的 Batch Size 时，怎样让模型每次更新使用更多训练数据？

这些问题都与模型的训练管理有关。

因此，这一篇我们暂时不修改 Transformer 的内部结构，而是完善之前的 train.py，实现以下功能：

• 验证集评估：同时观察训练 Loss 和验证 Loss。
• 模型检查点管理：分别保存最新模型和最佳模型。
• 断点续训：恢复模型参数、优化器状态和训练进度。
• 学习率调度：使用 Warmup 和余弦退火调整学习率。
• 梯度累积：在有限设备内存下模拟更大的有效 Batch Size。
• 梯度裁剪：限制过大的梯度，降低训练不稳定的风险。

本文继续使用前几篇实现的 GPTLanguageModel，并在 train.py 中完善训练逻辑。

---

一、为什么还需要验证集？

1.1 训练 Loss 下降不代表模型一定变好

在第二篇中，我们实现了基本的模型训练循环：

logits, loss = model(x, y)

optimizer.zero_grad()
loss.backward()
optimizer.step()

其中，loss 衡量模型预测结果与真实目标之间的差异。

对于我们的字符级语言模型，使用的是交叉熵损失：
$$L=-1N\; \sum i=1N\; logP(yi\mid \; x\le i\; )\; L=-\backslash frac\{1\}\{N\}\backslash sum\_\{i=1\}^\{N\}\backslash log\; P(y\_i\backslash mid\; x\_\{\backslash le\; i\})$$ L=−N1i=1∑NlogP(yi∣x≤i)

这里：

• $NN$ N 表示参与计算的预测位置数量。
• $yi\; y\_i$ yi 表示第 $ii$ i 个位置真正需要预测的字符。
• $x\le i\; x\_\{\backslash le\; i\}$ x≤i 表示该位置能够看到的输入上下文。
• $P(yi\mid \; x\le i\; )\; P(y\_i\backslash mid\; x\_\{\backslash le\; i\})$ P(yi∣x≤i) 表示模型分配给正确字符的概率。

训练过程中，我们通过反向传播不断调整模型参数，目标就是降低这个损失。

但需要注意：

训练 Loss 下降，只能说明模型越来越擅长拟合训练数据，并不能直接证明模型在新数据上的表现也越来越好。

例如，假设我们记录了以下结果：

训练步数	Train Loss	Val Loss
100	3.20	3.30
300	2.80	2.95
500	2.40	2.70
700	2.10	2.85
900	1.80	3.10

以上数值仅用于说明，不代表实际训练结果。

可以看到，训练 Loss 一直下降，但验证 Loss 在第 500 步之后反而开始上升。

这说明模型虽然越来越擅长预测训练集中的字符，却可能逐渐失去对未参与训练文本的泛化能力。

这就是常见的**过拟合（Overfitting）**现象。

当然，真实训练中的 Loss 可能存在随机波动，不能仅凭一次验证 Loss 上升就断定模型已经过拟合。

1.2 实现验证集评估

还记得第一篇中，我们把数据按照顺序划分为了训练集和验证集：

原始 Shakespeare 文本
         ↓
     字符编码
         ↓
   ┌─────┴─────┐
   │           │
训练集        验证集
 train.bin     val.bin

之前，我们主要使用训练集更新模型参数。

现在，就可以使用验证集评估模型的泛化表现。

首先实现 estimate_loss()：

@torch.no_grad()
def estimate_loss(model, datasets, batch_args, eval_iters):
    was_training = model.training
    model.eval()

    result = {}

    try:
        for split in ("train", "val"):
            losses = []

            for _ in range(eval_iters):
                x, y = get_batch(
                    split,
                    datasets,
                    **batch_args,
                )

                _, loss = model(x, y)
                losses.append(loss.item())

            result[split] = sum(losses) / len(losses)

    finally:
        model.train(was_training)

    return result

这段代码有几个需要理解的地方。

（1）为什么使用 @torch.no_grad() ？

训练时，PyTorch 需要记录计算过程，构建计算图，以便后续执行反向传播。

但是评估模型时，我们只需要计算 Loss，不需要更新参数。

因此可以使用：

@torch.no_grad()

关闭梯度记录，减少不必要的计算图开销。

（2）为什么调用 model.eval() ？

PyTorch 中的模型有两种主要运行模式：

model.train()
model.eval()

它们不会决定是否计算梯度，而是影响某些模块的行为。

例如，我们在第四篇的 Transformer Block 中使用了 Dropout。

训练模式下，Dropout 会随机将部分特征置零。

评估模式下，Dropout 不再随机丢弃特征。

因此，在验证模型时，应该先调用：

model.eval()

（3）为什么最后还要恢复模型模式？

因为评估结束之后，我们还需要继续训练。

was_training = model.training

这行代码记录模型进入评估函数之前的状态。

评估结束后：

model.train(was_training)

将模型恢复为原来的模式。

这里使用 try...finally，可以保证即使评估过程中出现异常，也会尝试恢复模型的运行模式。

（4）为什么计算多个 Batch 的平均 Loss？

因为 get_batch() 是随机抽取训练样本的。

如果只计算一个 Batch，它可能恰好比较容易，也可能比较困难，导致 Loss 产生较大的波动。

因此，我们连续抽取多个 Batch：

for _ in range(eval_iters):
    ...

然后计算平均值：

result[split] = sum(losses) / len(losses)

这样能够减小单个随机 Batch 对结果的影响。

但要注意，这仍然是抽样估计，并不是对整个验证集的精确评估。

---

二、为什么要保存两份模型？

第五篇中，我们已经学会了使用：

torch.save(checkpoint, checkpoint_path)

保存模型检查点。

但之前主要考虑的是：

如何让训练好的模型可以被重新加载？

现在还需要进一步考虑：

如何保证我们既能继续训练，又不会丢失训练过程中表现更好的模型？

2.1 Last 和 Best 的区别

我们可以同时维护两份检查点：

检查点	保存条件	主要用途
last.pt	每次定期评估后	恢复最近的训练进度
best.pt	验证 Loss 创下新低时	加载表现最好的模型

例如，假设模型训练了 1000 步。

第 800 步时：

Val Loss = 2.15

第 1000 步时：

Val Loss = 2.30

虽然第 1000 步的模型训练时间更长，但它在验证集上的损失反而更大。

这时：

last.pt → 第 1000 步模型
best.pt → 第 800 步模型

如果我们希望继续训练，就加载 last.pt。

如果希望使用验证集表现更好的模型生成文本，就加载 best.pt。

2.2 代码实现

首先记录历史最低验证 Loss：

best_val_loss = float("inf")

这里使用正无穷，意味着第一次正常得到的有限验证 Loss 都会成为最佳结果。

每次评估结束后：

improved = val_loss < best_val_loss

if improved:
    best_val_loss = val_loss

随后保存：

torch.save(checkpoint, out_dir / "last.pt")

if improved:
    torch.save(checkpoint, out_dir / "best.pt")

这样便实现了两种不同的保存策略。

需要注意，本篇中的"最佳模型"指的是：

在定期抽样验证中，取得最低验证 Loss 的模型。

它不一定是整个训练过程中理论上最优的模型，因为我们并没有在每一步都进行验证。

---

三、断点续训：为什么只保存模型参数还不够？

3.1 回顾模型参数保存

第五篇中，我们使用：

model.state_dict()

获得模型参数。

之后可以通过：

model.load_state_dict(checkpoint["model"])

恢复模型权重。

对于文本生成而言，这通常已经足够。

但是，如果我们希望继续训练，仅仅恢复模型参数是不完整的。

这是因为我们还使用了优化器。

3.2 AdamW 也存在内部状态

之前的训练代码中，我们使用：

optimizer = torch.optim.AdamW(
    model.parameters(),
    lr=3e-4
)

AdamW 不仅依赖当前计算得到的梯度，还会维护梯度的一阶矩和二阶矩估计。

以 Adam 的核心更新思想为例：
$$mt=\beta 1\; mt-1\; +(1-\beta 1)gt\; m\_t=\backslash beta\_1m\_\{t-1\}+(1-\backslash beta\_1)g\_t$$ mt=β1mt−1+(1−β1)gt
$$vt=\beta 2\; vt-1\; +(1-\beta 2)gt2\; v\_t=\backslash beta\_2v\_\{t-1\}+(1-\backslash beta\_2)g\_t^2$$ vt=β2vt−1+(1−β2)gt2

其中：

• $gt\; g\_t$ gt 是当前梯度。
• $mt\; m\_t$ mt 是梯度的一阶矩估计。
• $vt\; v\_t$ vt 是梯度平方的二阶矩估计。
• $\beta 1,\beta 2\; \backslash beta\_1,\backslash beta\_2$ β1,β2 是控制历史信息衰减的系数。

AdamW 在自适应梯度更新的基础上，采用解耦权重衰减。

这意味着，优化器的更新过程不仅与当前模型参数有关，还与之前积累的优化状态有关。

如果只恢复模型参数，却重新创建一个没有历史状态的优化器：

optimizer = torch.optim.AdamW(model.parameters())

优化器之前积累的一阶矩和二阶矩信息就丢失了。

所以，为了支持更完整的续训，我们还需要保存：

optimizer.state_dict()

恢复时执行：

optimizer.load_state_dict(checkpoint["optimizer"])

3.3 设计新的检查点结构

本篇将检查点组织为：

checkpoint = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "model_args": model_args,
    "train_args": train_args,
    "meta": meta,
    "completed_steps": completed_steps,
    "best_val_loss": best_val_loss,
    "val_loss": val_loss,
}

各字段作用如下：

字段	作用
model	模型参数及持久化 Buffer
optimizer	优化器状态
model_args	模型结构配置
train_args	训练超参数
meta	字符词表和编码信息
completed_steps	已完成的参数更新次数
best_val_loss	历史最低验证 Loss
val_loss	当前检查点的验证 Loss

这里特别需要理解：

completed_steps

在本篇中，我们统一约定：

一次训练 Step 指的是一次参数更新，即执行一次 optimizer.step() 。

它不是从数据集中抽取一次 Batch。

为什么需要这样约定？

因为接下来我们会加入梯度累积。

加入梯度累积后，模型可能需要读取多个 Batch，才执行一次参数更新。

---

四、学习率调度：Warmup + 余弦退火

4.1 为什么不一直使用固定学习率？

在之前的训练中，我们使用固定学习率：

learning_rate = 3e-4

整个训练过程都保持不变。

这种方式比较简单，但在较长的训练过程中，我们希望能够根据训练进度调整学习率。

常见策略是：

1. 训练刚开始时使用较小的学习率。
2. 经过一段 Warmup，逐渐增加到设定的最大学习率。
3. 随着训练继续，逐渐降低学习率。

这样既可以减轻初始阶段参数更新过大的问题，也可以让后期更新更加细致。

4.2 Warmup：学习率预热

假设：

learning_rate = 3e-4
warmup_steps = 100

我们希望前 100 次更新逐渐把学习率提高到：
$$3\times 10-43\backslash times10^\{-4\}$$ 3×10−4

可以使用线性增长：
$$\eta t=\eta max\; t+1Twarmup\; \backslash eta\_t=\backslash eta\_\{\backslash max\}\backslash frac\{t+1\}\{T\_\{\backslash text\{warmup\}\}\}$$ ηt=ηmaxTwarmupt+1

其中：

• $\eta t\; \backslash eta\_t$ ηt 是当前学习率。
• $\eta max\; \backslash eta\_\{\backslash max\}$ ηmax 是最大学习率。
• $tt$ t 是从 0 开始的更新索引。
• $Twarmup\; T\_\{\backslash text\{warmup\}\}$ Twarmup 是预热步数。

对应代码：

if step < warmup_steps:
    return peak_lr * (step + 1) / warmup_steps

这里使用 step + 1，是为了让第一次参数更新就获得非零学习率。

4.3 余弦退火

Warmup 结束后，我们使用余弦函数逐渐降低学习率：
$$\eta t=\eta min+12(\eta max-\eta min)$$1+cos⁡(πp) \eta_t= \eta_{\min} +\frac{1}{2}(\eta_{\max}-\eta_{\min}) \left1+\\cos(\\pi p)\\right ηt=ηmin+21(ηmax−ηmin)1+cos(πp)

其中：
$$p=\; t-Twarmup\; Tdecay-Twarmup\; p=\; \backslash frac\{t-T\_\{\backslash text\{warmup\}\}\}\; \{T\_\{\backslash text\{decay\}\}-T\_\{\backslash text\{warmup\}\}\}$$ p=Tdecay−Twarmupt−Twarmup

这里的 $pp$ p 表示学习率衰减阶段的进度。

当 $p=0p=0$ p=0 时：
$$\eta t=\eta max\; \backslash eta\_t=\backslash eta\_\{\backslash max\}$$ ηt=ηmax

当 $p=1p=1$ p=1 时：
$$\eta t=\eta min\; \backslash eta\_t=\backslash eta\_\{\backslash min\}$$ ηt=ηmin

这样，学习率就会沿着余弦曲线逐渐下降。

4.4 完整实现

def get_lr(step, config):
    peak_lr = config["learning_rate"]
    min_lr = config["min_lr"]
    warmup_steps = config["warmup_steps"]
    decay_steps = config["decay_steps"]

    if step < warmup_steps:
        return peak_lr * (step + 1) / warmup_steps

    if step >= decay_steps:
        return min_lr

    progress = (
        (step - warmup_steps)
        / (decay_steps - warmup_steps)
    )

    coefficient = 0.5 * (
        1.0 + math.cos(math.pi * progress)
    )

    return min_lr + coefficient * (peak_lr - min_lr)

注意，这个函数要求配置满足：

0 < warmup_steps < decay_steps

本篇还会把这些调度参数放进 train_args，与检查点一起保存。

这样在断点续训时，学习率就能够按照已经完成的训练步数继续变化，而不是重新从 Warmup 开始。

---

五、梯度累积：如何模拟更大的 Batch Size？

5.1 为什么需要梯度累积？

在第一篇中，我们介绍了两个重要超参数：

batch_size
block_size

其中：

• batch_size 表示一次处理多少条序列。
• block_size 表示每条序列包含多少个输入 Token。

例如：

batch_size = 8
block_size = 64

得到的输入形状是：

x.shape = (8, 64)

即一次前向传播处理 8 条序列，每条序列包含 64 个 Token。

现在假设，我们希望把 Batch Size 提高到 32。

最直接的方法是：

batch_size = 32

但这样通常也会增加训练时的内存或显存占用。

如果设备无法一次处理这么大的 Batch，我们可以使用梯度累积。

5.2 梯度为什么可以累积？

先回顾 PyTorch 的反向传播：

loss.backward()

默认情况下，PyTorch 不会在每次反向传播时自动清空参数的梯度。

而是把计算出的梯度累加到：

parameter.grad

因此：

loss1.backward()
loss2.backward()

会累积两次反向传播产生的梯度。

利用这个特点，我们就可以分多个小批次计算梯度，最后统一更新参数。

5.3 一个简单的例子

假设：

batch_size = 8
gradient_accumulation_steps = 4

我们连续处理 4 个 Batch：

Batch 1：8 条序列 → 计算梯度
Batch 2：8 条序列 → 累积梯度
Batch 3：8 条序列 → 累积梯度
Batch 4：8 条序列 → 累积梯度
                         ↓
                  optimizer.step()

这样一次参数更新就使用了：
$$8\times 4=328\backslash times4=32$$ 8×4=32

条序列。

因此可以定义有效批次大小：
$$Bmicro\times Naccumulation\; B\_\{\backslash text\{micro\}\}\; \backslash times\; N\_\{\backslash text\{accumulation\}\}$$ Bmicro×Naccumulation

对于本例：
$$Beffective=8\times 4=32\; B\_\{\backslash text\{effective\}\}=8\backslash times4=32$$ Beffective=8×4=32

5.4 为什么 Loss 要除以累积次数？

实现代码如下：

optimizer.zero_grad(set_to_none=True)

for _ in range(4):
    x, y = get_batch(...)

    _, loss = model(x, y)

    loss = loss / 4
    loss.backward()

optimizer.step()

注意这里：

loss = loss / 4

如果直接执行 4 次反向传播，就会把 4 个 Batch 的梯度直接相加。

但我们希望计算的是平均梯度。

假设四个 Batch 对应的梯度为：
$$g1,g2,g3,g4\; g\_1,g\_2,g\_3,g\_4$$ g1,g2,g3,g4

期望得到：
$$g=\; g1+g2+g3+g4\; 4\; g=\backslash frac\{g\_1+g\_2+g\_3+g\_4\}\{4\}$$ g=4g1+g2+g3+g4

而不是：
$$g=g1+g2+g3+g4\; g=g\_1+g\_2+g\_3+g\_4$$ g=g1+g2+g3+g4

因此，我们在反向传播之前，将每个 Batch 的 Loss 除以 4。

在每个小批次包含相同数量的有效预测位置时，这样可以得到与大批次平均梯度相对应的结果。

但由于数值计算顺序及某些随机操作可能不同，梯度累积不一定与直接使用大 Batch 产生完全相同的浮点结果。

5.5 每次参数更新处理多少 Token？

除了有效批次大小，我们还可以计算每次更新处理的输入 Token 数量：
$$Bmicro\times Naccumulation\times T\; B\_\{\backslash text\{micro\}\}\; \backslash times\; N\_\{\backslash text\{accumulation\}\}\; \backslash times\; T$$ Bmicro×Naccumulation×T

其中， $TT$ T 是 block_size。

代入本篇的配置：

batch_size = 8
gradient_accumulation_steps = 4
block_size = 64

得到：
$$8\times 4\times 64=20488\backslash times4\backslash times64=2048$$ 8×4×64=2048

也就是说，每次参数更新处理 2048 个输入 Token。

需要注意：

梯度累积降低的是单次前向和反向传播需要同时处理的数据规模，并不会减少这些小批次本身需要执行的计算。

---

六、梯度裁剪：限制过大的梯度

训练神经网络时，可能会遇到梯度突然变得很大的情况。

假设某次更新中，一个参数的梯度为：
$$g=1000g=1000$$ g=1000

如果学习率为：
$$\eta =0.001\backslash eta=0.001$$ η=0.001

那么按照最简单的 SGD 更新公式：
$$\theta t+1\; =\theta t-\eta g\; \backslash theta\_\{t+1\}=\backslash theta\_t-\backslash eta\; g$$ θt+1=θt−ηg

这次更新对应的参数变化量就是：
$$\Delta \theta =-1\backslash Delta\backslash theta=-1$$ Δθ=−1

当模型参数规模较小，或者某些方向上的梯度异常大时，这种更新可能导致训练不稳定。

为此，我们可以使用梯度裁剪。

6.1 什么是梯度范数裁剪？

PyTorch 提供了：

torch.nn.utils.clip_grad_norm_(
    model.parameters(),
    max_norm=1.0
)

它会计算所有参数梯度的整体范数。

假设把所有参数梯度展平后得到向量：
$$g=(g1,g2,...,gn)\; g=(g\_1,g\_2,\backslash ldots,g\_n)$$ g=(g1,g2,...,gn)

其 L2 范数为：
$$\mid g\mid \ast 2=\; \sum \ast i=1ngi2\; |g|*2=\; \backslash sqrt\{\backslash sum*\{i=1\}^\{n\}g\_i^2\}$$ ∣g∣∗2=∑∗i=1ngi2

如果：
$$\mid g\mid 2>c\; |g|\_2>c$$ ∣g∣2>c

就按照比例缩小梯度：
$$g\; c\; \mid g\mid 2\; g\backslash frac\{c\}\{|g|\_2\}$$ g∣g∣2c

其中， $cc$ c 是允许的最大梯度范数。

在实际实现中，还会考虑数值稳定性。

梯度裁剪不是把每个梯度元素直接限制到 [-1, 1]，而是限制整体梯度范数。

6.2 梯度裁剪应该放在哪里？

本篇还使用了梯度累积，所以训练顺序非常重要：

optimizer.zero_grad()
        ↓
第一个 Batch 前向传播
        ↓
第一个 Batch 反向传播
        ↓
第二个 Batch 前向传播
        ↓
第二个 Batch 反向传播
        ↓
       ...
        ↓
梯度累积完成
        ↓
  clip_grad_norm_()
        ↓
  optimizer.step()

对应代码：

optimizer.zero_grad(set_to_none=True)

for _ in range(accumulation):
    x, y = get_batch(...)

    _, loss = model(x, y)

    (loss / accumulation).backward()

torch.nn.utils.clip_grad_norm_(
    model.parameters(),
    max_norm=1.0
)

optimizer.step()

需要特别注意：

梯度裁剪应该放在所有小批次的梯度累积完成之后、参数更新之前。

如果在每个小批次之后分别裁剪，就会改变我们最终希望累积得到的梯度。

---

七、完整实现 train.py

前面已经分别介绍了验证集评估、检查点保存、断点续训、学习率调度、梯度累积和梯度裁剪。

现在，我们把这些功能整合到同一个训练脚本中。

继续沿用之前的项目结构：

my-nanogpt/
├── data/
│   └── shakespeare_char/
│       ├── train.bin
│       ├── val.bin
│       └── meta.pkl
├── dataset.py
├── model.py
├── train.py
└── sample.py

其中：

• dataset.py 负责加载数据和生成 Batch。
• model.py 保存第四篇实现的 GPTLanguageModel。
• train.py 负责训练、评估、保存和断点续训。
• sample.py 负责加载模型并生成文本。

这里继续使用前面定义的数据接口：

load_data(data_dir)

以及：

get_batch(
    split,
    datasets,
    *,
    batch_size,
    block_size,
    device,
)

下面给出完整的 train.py。

from pathlib import Path
import argparse
import math
import pickle

import torch

from dataset import load_data, get_batch
from model import GPTLanguageModel


BASE_DIR = Path(__file__).resolve().parent


# -------------------------
# 学习率调度
# -------------------------

def get_lr(step, config):
    peak_lr = config["learning_rate"]
    min_lr = config["min_lr"]
    warmup_steps = config["warmup_steps"]
    decay_steps = config["decay_steps"]

    if step < warmup_steps:
        return peak_lr * (step + 1) / warmup_steps

    if step >= decay_steps:
        return min_lr

    progress = (
        (step - warmup_steps)
        / (decay_steps - warmup_steps)
    )

    coefficient = 0.5 * (
        1.0 + math.cos(math.pi * progress)
    )

    return min_lr + coefficient * (peak_lr - min_lr)


# -------------------------
# 评估模型
# -------------------------

@torch.no_grad()
def estimate_loss(model, datasets, batch_args, eval_iters):
    was_training = model.training
    model.eval()

    result = {}

    try:
        for split in ("train", "val"):
            losses = []

            for _ in range(eval_iters):
                x, y = get_batch(
                    split,
                    datasets,
                    **batch_args,
                )

                _, loss = model(x, y)
                losses.append(loss.item())

            result[split] = sum(losses) / len(losses)

    finally:
        model.train(was_training)

    return result


# -------------------------
# 主训练函数
# -------------------------

def main(argv=None):

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--steps",
        type=int,
        default=1000,
    )

    parser.add_argument(
        "--resume",
        action="store_true",
    )

    parser.add_argument(
        "--eval-every",
        type=int,
        default=100,
    )

    parser.add_argument(
        "--eval-iters",
        type=int,
        default=20,
    )

    parser.add_argument(
        "--data-dir",
        type=Path,
        default=BASE_DIR / "data" / "shakespeare_char",
    )

    parser.add_argument(
        "--out-dir",
        type=Path,
        default=BASE_DIR / "out-v2",
    )

    parser.add_argument(
        "--device",
        choices=["auto", "cpu", "mps", "cuda"],
        default="auto",
    )

    args = parser.parse_args(argv)

    if min(args.steps, args.eval_every, args.eval_iters) <= 0:
        parser.error(
            "steps、eval-every、eval-iters 必须大于 0"
        )

    # -------------------------
    # 选择运行设备
    # -------------------------

    device = args.device

    if device == "auto":
        if torch.cuda.is_available():
            device = "cuda"
        elif torch.backends.mps.is_available():
            device = "mps"
        else:
            device = "cpu"

    # -------------------------
    # 检查点路径
    # -------------------------

    last_path = args.out_dir / "last.pt"
    best_path = args.out_dir / "best.pt"

    if not args.resume and (
        last_path.exists() or best_path.exists()
    ):
        parser.error(
            "输出目录已有检查点，请使用 --resume 或更换 --out-dir"
        )

    checkpoint = None

    # -------------------------
    # 加载续训检查点
    # -------------------------

    if args.resume:
        if not last_path.is_file():
            parser.error(
                f"找不到续训检查点：{last_path}"
            )

        checkpoint = torch.load(
            last_path,
            map_location="cpu",
            weights_only=True,
        )

    # -------------------------
    # 加载数据
    # -------------------------

    with open(args.data_dir / "meta.pkl", "rb") as f:
        meta = pickle.load(f)

    datasets = load_data(args.data_dir)

    # -------------------------
    # 初始化配置或恢复配置
    # -------------------------

    if checkpoint is None:

        model_args = {
            "vocab_size": meta["vocab_size"],
            "n_embd": 64,
            "n_head": 4,
            "n_layer": 2,
            "block_size": 64,
            "dropout": 0.1,
        }

        train_args = {
            "batch_size": 8,
            "gradient_accumulation_steps": 4,
            "learning_rate": 3e-4,
            "min_lr": 3e-5,
            "warmup_steps": 100,
            "decay_steps": 5000,
            "grad_clip": 1.0,
            "weight_decay": 0.01,
        }

        completed_steps = 0
        best_val_loss = float("inf")

    else:

        if checkpoint["meta"] != meta:
            raise ValueError(
                "当前数据词表与检查点不一致"
            )

        model_args = checkpoint["model_args"]
        train_args = checkpoint["train_args"]

        completed_steps = checkpoint["completed_steps"]
        best_val_loss = checkpoint["best_val_loss"]

    # -------------------------
    # 检查训练配置
    # -------------------------

    if not (
        0 < train_args["warmup_steps"]
        < train_args["decay_steps"]
    ):
        raise ValueError(
            "需要满足 0 < warmup_steps < decay_steps"
        )

    torch.manual_seed(1337)

    # -------------------------
    # 创建模型与优化器
    # -------------------------

    model = GPTLanguageModel(**model_args).to(device)

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=train_args["learning_rate"],
        weight_decay=train_args["weight_decay"],
    )

    if checkpoint is not None:
        model.load_state_dict(checkpoint["model"])
        optimizer.load_state_dict(checkpoint["optimizer"])

    # -------------------------
    # Batch 配置
    # -------------------------

    batch_args = {
        "batch_size": train_args["batch_size"],
        "block_size": model_args["block_size"],
        "device": device,
    }

    for split in ("train", "val"):
        if len(datasets[split]) <= model_args["block_size"]:
            raise ValueError(
                f"{split} 数据长度必须大于 block_size"
            )

    args.out_dir.mkdir(
        parents=True,
        exist_ok=True,
    )

    accumulation = train_args[
        "gradient_accumulation_steps"
    ]

    target_steps = completed_steps + args.steps

    tokens_per_update = (
        train_args["batch_size"]
        * accumulation
        * model_args["block_size"]
    )

    print(f"设备：{device}")
    print(
        f"从 {completed_steps} 步训练到 "
        f"{target_steps} 步"
    )
    print(
        f"每次更新处理 "
        f"{tokens_per_update} 个输入 token"
    )

    # -------------------------
    # 开始训练
    # -------------------------

    model.train()

    for step in range(completed_steps, target_steps):

        # 更新学习率
        lr = get_lr(step, train_args)

        for group in optimizer.param_groups:
            group["lr"] = lr

        # 清空上一轮梯度
        optimizer.zero_grad(set_to_none=True)

        update_loss = 0.0

        # 梯度累积
        for _ in range(accumulation):

            x, y = get_batch(
                "train",
                datasets,
                **batch_args,
            )

            _, loss = model(x, y)

            update_loss += (
                loss.item() / accumulation
            )

            (loss / accumulation).backward()

        # 梯度裁剪
        torch.nn.utils.clip_grad_norm_(
            model.parameters(),
            train_args["grad_clip"],
        )

        # 更新参数
        optimizer.step()

        completed_steps = step + 1

        # 打印训练日志
        if completed_steps % 20 == 0:
            print(
                f"step {completed_steps} | "
                f"batch loss {update_loss:.4f} | "
                f"lr {lr:.2e}"
            )

        # -------------------------
        # 定期评估
        # -------------------------

        should_eval = (
            completed_steps % args.eval_every == 0
            or completed_steps == target_steps
        )

        if should_eval:

            losses = estimate_loss(
                model,
                datasets,
                batch_args,
                args.eval_iters,
            )

            improved = (
                losses["val"] < best_val_loss
            )

            if improved:
                best_val_loss = losses["val"]

            # 构造检查点
            state = {
                "model": model.state_dict(),
                "optimizer": optimizer.state_dict(),
                "model_args": model_args,
                "train_args": train_args,
                "meta": meta,
                "completed_steps": completed_steps,
                "best_val_loss": best_val_loss,
                "val_loss": losses["val"],
            }

            # 保存最新检查点
            torch.save(
                state,
                last_path,
            )

            # 保存最佳检查点
            if improved:
                torch.save(
                    state,
                    best_path,
                )

            print(
                f"评估 step {completed_steps} | "
                f"train {losses['train']:.4f} | "
                f"val {losses['val']:.4f} | "
                f"best {best_val_loss:.4f}"
            )


if __name__ == "__main__":
    main()

这里需要特别说明：

本文通过 --resume 恢复模型参数、优化器状态、训练配置和已经完成的更新次数。

它能够支持正常的断点续训，但还不是严格意义上的逐位一致恢复。

例如，本篇尚未保存 PyTorch 随机数生成器状态，也没有保存 CUDA 或 MPS 对应的全部随机状态。

因此，中断后恢复训练得到的参数，不保证与完全不中断的训练结果逐位一致。

另外，本篇保存的检查点结构比第五篇更加完整，但仍保留了第五篇 sample.py 所需的 model、model_args 和 meta 字段。

因此，之前的生成程序仍然可以按照相同方式构建并加载模型。

---

八、先进行小规模测试

修改完 train.py 后，不建议立即开始长时间训练。

我们应该先确认三个功能：

1. 训练能够正常进行。
2. 检查点能够正常保存。
3. 程序能够正确恢复训练进度。

8.1 先训练两步

在项目根目录执行：

python train.py \
  --steps 2 \
  --eval-every 1 \
  --eval-iters 2 \
  --out-dir out-smoke

这里：

• --steps 2 表示训练 2 次参数更新。
• --eval-every 1 表示每次更新后都评估。
• --eval-iters 2 表示每个数据集抽取 2 个 Batch 估计 Loss。
• --out-dir out-smoke 表示把检查点保存到 out-smoke 目录。

这种测试通常被称为 Smoke Test，主要用于检查程序能否完成基本流程。

由于训练步数很少，它不能用于判断模型最终效果。

如果程序正常运行，目录中应该出现：

out-smoke/
├── last.pt
└── best.pt

8.2 测试断点续训

接下来执行：

python train.py \
  --resume \
  --steps 2 \
  --eval-every 1 \
  --eval-iters 2 \
  --out-dir out-smoke

注意这次增加了：

--resume

程序就会从：

out-smoke/last.pt

加载检查点。

启动信息应该包含：

从 2 步训练到 4 步

这里需要理解一个重要约定：

本篇的 --steps 表示本次运行还要额外训练多少步，而不是训练结束时的总步数。

例如：

python train.py --resume --steps 500

如果当前检查点已经完成了 1000 步，程序就会继续训练至 1500 步，而不是重新训练到 500 步。

另外，last.pt 是定期保存的，而不是每次参数更新后都保存。

如果程序在两次保存之间意外中断，就会丢失这段尚未写入检查点的训练进度。

---

九、正式训练与文本生成

9.1 启动正式训练

确认小规模测试通过之后，我们就可以开始一个新的训练实验。

python train.py \
  --steps 1000 \
  --out-dir out-v2

程序会自动选择可用设备。

如果在支持 MPS 的 Mac 上运行，且没有可用的 CUDA 设备，则会选择：

设备：mps

本篇的默认配置为：

batch_size = 8
gradient_accumulation_steps = 4
block_size = 64

所以每次参数更新处理的输入 Token 数为：

8 × 4 × 64 = 2048

默认每 100 步进行一次训练集和验证集评估。

最终会在输出目录中保存：

out-v2/
├── last.pt
└── best.pt

9.2 继续训练

假设我们已经完成了 1000 步，现在希望再训练 1000 步。

执行：

python train.py \
  --resume \
  --steps 1000 \
  --out-dir out-v2

程序将从最新的检查点继续训练。

如果此前恰好在第 1000 步保存了检查点，恢复后将继续训练至第 2000 步。

需要注意，本篇的学习率调度默认在第 5000 步达到最小学习率。

所以即使我们分多次启动程序，学习率仍然根据累计更新步数计算，不会在每次运行时重新开始 Warmup。

9.3 使用最佳模型生成文本

第五篇已经实现了独立的 sample.py。

现在，我们可以指定最佳检查点生成文本：

python sample.py \
  --checkpoint out-v2/best.pt \
  --prompt "ROMEO:" \
  --tokens 300 \
  --temperature 0.8 \
  --top-k 20

这里继续沿用上一篇的参数：

• --checkpoint：需要加载的检查点。
• --prompt：文本生成的起始提示。
• --tokens：需要生成的 Token 数。
• --temperature：控制采样概率分布的平滑程度。
• --top-k：限制每次采样考虑的候选 Token 数量。

由于我们训练的是字符级语言模型，这里的一个 Token 通常对应一个字符。

同样，我们也可以加载最新检查点：

python sample.py \
  --checkpoint out-v2/last.pt \
  --prompt "ROMEO:" \
  --tokens 300

这样就可以比较最新检查点与最佳检查点的生成效果。

---

十、训练过程中应该观察什么？

现在，我们已经拥有一个比之前更加完整的训练系统。

但训练程序能够运行，并不等于模型已经学到了有效的语言规律。

在训练过程中，我建议重点观察三个方面。

10.1 训练 Loss 是否整体下降

训练 Loss 下降说明模型逐渐能够更好地拟合训练样本。

但它只能作为一个训练指标。

10.2 验证 Loss 是否同步改善

验证 Loss 反映模型对未参与训练的数据的预测表现。

如果训练 Loss 继续下降，而验证 Loss 持续上升，就应该考虑是否已经出现过拟合。

10.3 生成文本的结构是否改善

对于我们的 Shakespeare 字符级模型，可以使用固定的提示词：

ROMEO:

定期生成文本进行对比。

随着训练进行，我们希望模型逐渐学到：

• 常见英文单词的拼写模式。
• 合理的单词间隔。
• 标点和换行的使用。
• Shakespeare 文本中常见的人物名字与对话形式。

不过，生成文本具有随机性。即使模型参数完全相同，不同的采样结果也可能存在明显差异。

因此，生成样例适合用于直观检查模型表现，但不应该完全代替验证 Loss 等定量指标。

---

十一、本篇总结

从第一篇到现在，我们已经逐渐完成了一个小型字符级 GPT 项目的主要部分。

第一篇，我们把 Shakespeare 文本转换为模型可以处理的整数，并构造训练样本。

第二篇，我们实现 Bigram 语言模型，理解语言模型的训练与生成过程。

第三篇，我们引入因果自注意力，让每个位置能够利用之前的上下文。

第四篇，我们加入多头注意力、前馈网络、残差连接和 LayerNorm，搭建了完整的小型 GPT。

第五篇，我们实现了模型检查点保存，并通过独立的生成脚本加载模型。

而这一篇，我们进一步完善了训练流程：

功能	解决的问题
验证集评估	判断模型在未参与训练数据上的表现
Last / Best 检查点	同时保存最新训练进度和最佳验证结果
优化器状态恢复	支持更加完整的断点续训
Warmup + 余弦退火	根据训练进度调整学习率
梯度累积	在有限内存下增大有效批次大小
梯度裁剪	限制过大的梯度范数

到这里，我们的项目不仅可以完成文本生成，也已经具备了持续训练、定期评估、保存检查点和恢复训练进度的基本能力。

当然，它仍然是一个以理解原理为目标的教学实现，与真正的 nanoGPT 在模型计算效率、训练工程和性能优化方面还有一定差距。

下一篇预告

《从零实现 nanoGPT（七）：优化注意力计算，让 GPT 训练更高效》

接下来，我们将重新回到模型内部，逐步优化之前实现的 Transformer。

主要包括：

1. 将多个独立的注意力头改为融合 QKV 投影，减少分散的线性计算。
2. 使用 PyTorch 的 scaled_dot_product_attention 接口，了解它如何根据设备和条件选择注意力实现。
3. 实现输入 Embedding 与输出层的权重共享。
4. 了解 GPT 中常见的参数初始化方法。

通过这些优化，我们将进一步理解 nanoGPT 是怎样从一个能够运行的教学模型，逐渐发展为更加高效的训练实现的。