第 21 章训练、保存和审问 Mini GPT

上一章的模型已经能分配概率，但参数仍是刚初始化时的随机值。把那次前向计算重复一千遍，也不会让它学到东西。我们必须用损失来改变参数，然后分开回答两个问题：程序有没有遵照学习规则运行？得到的模型有没有学会我们希望它掌握的能力？

MG-2 在同一个 Mini GPT 上完成有界的 CPU 训练、状态保存、断点恢复、另起进程复跑，以及生成记录。下面的数值，要么来自明确说明的计算，要么来自共享的本地实测。小语料是本书原创的教学数据，这些结果不冒充论文基准或工业语言模型的表现。

21.1 让一次损失真正带来参数更新

第 7 章把梯度解释为：某个参数稍作变化时，损失在当前位置变化得有多快。反向传播沿前向计算反着走，算出这些变化率。模型结构没有改变，现在梯度从词表损失出发，经过输出头、归一化、FFN、注意力和嵌入，传回各个参数。优化器再据此更新保存的数值。一次前向计算中的注意力权重是临时结果，并不是优化器保存的参数。

例 21.1：第一次实际训练更新。 保持第 20 章的 24 条训练记录及其顺序，编号为 0—23。每次独立抽取一个编号，抽后放回，共抽八次。固定的批次随机数生成器第一次给出 [15,19,1,22,7,15,5,3]。15 号记录出现两次，是因为抽中后并未把它从集合里移走。按上一章的规则，为抽中的每条记录构造输入、有效位置标记和错开一位的目标。更新前，实际测得这批数据的平均损失为每个有效目标 5.672425134737334 奈特。

前向计算的细节已经实现，因此训练循环的关键步骤很短：

indices = torch.randint(len(corpus["train"]), (cfg["batch_size"],), generator=rng)
batch = tuple(array[indices] for array in train_batch)
opt.zero_grad(set_to_none=True)
loss = model(batch[0], batch[1], batch[2])[1]
loss.backward()
opt.step()

先清除梯度，可以避免无意中叠加上一步的贡献。backward() 填入梯度，step() 改变参数。第一次批次损失是在这次改变以前测得的。下面的曲线则在相应更新完成以后，评测两个固定划分中的全部记录。记录集合、平均分母和参数所处的时刻不同，就解释了这些数值为什么彼此不同，也不同于 MG-1 那两条记录的损失。

这里使用 Adam，即第 7 章非线性实验中已经用过的优化器。它持续保存梯度及梯度平方的统计量，所以更新既依赖当前梯度，也依赖此前的历史。本章把这条规则作为明确选定的实现条件，不把 Adam 的结果归到手算的普通 SGD 减法上。恢复训练时，这些历史状态会派上用场。

表 21.1：本次运行的固定约定。 设置保存在 data/mini-gpt/config.json 中，实测记录另附环境信息和完整依赖清单。

项目	数值与含义
模型与数据	同一个单块、13875 参数模型；24 条训练记录、8 条验证记录
上下文／批次	每条记录 64 个槽位；每次更新有放回地抽取 8 条记录
更新／测量	更新 600 次；在第 0、50、…、600 步评测完整划分
Adam	学习率 0.01；betas 为 0.9、0.999；epsilon 为 0.00000001；weight decay 为 0；`foreach=False`
初始化种子	1729，用于初始矩阵等初始化数值
批次种子	101，由独立的 CPU 随机数生成器使用
采样种子	31415，与训练分开
数值运行条件	CPU、float64、单线程，开启确定性算法；dropout 为 0
已测软件与系统	Python 3.12.10、NumPy 2.0.2、PyTorch 2.7.0；macOS 26.5.2 arm64
选择规则	在 13 个测量点中，取首次达到最低验证损失的那一步

种子为某种随机数生成过程规定初始状态，并不是一份脱离软件也通用的复现配方。分开使用生成器，可以防止“为了展示而抽一个结果”消耗掉下一批训练数据本来要用的随机数。源代码、数据及其顺序、tokenizer、配置和依赖，也都属于这次运行的身份信息。

21.2 先读完整条曲线，再选参数快照

评测通过 eval() 和 no_grad() 进行模型运行（inference），不执行优化器更新。每次完整训练集评测包含 521 个有效的下一个 token 预测事件，验证集包含 178 个。对任一划分，都把实际目标的负对数概率相加，再除以该划分的有效事件数。每个事件的权重相同；这里不先求各条记录的平均值，再让长短记录各占同样的总权重。

例如，一条记录有两个事件，损失分别为 1、3；另一条只有一个事件，损失为 5。按事件平均，结果为 $(1+3+5)/3=3$ 奈特。如果先求记录均值再平均，就会得到 $((1+3)/2+5)/2=3.5$ 。两种算法回答的不是同一个问题。训练和评测都沿用第 20 章按有效事件计数的分母。

实线、圆点：训练。虚线、方点：验证。点状竖线：所选第 50 步。线段只连接实测点。

图 21.1 — 第 50 步后，训练拟合继续改善，验证表现反而变差本书原创 CPU 实测，CC BY-SA 4.0。数据见 data/mini-gpt/training-run.json；所有绘图点也列于表 21.2。

表 21.2：全部实测点，单位为每个有效目标的奈特数。 保留六位小数，完整精度见 training-run.json。图 21.1 的线只连接这些点，并不表示中间各步也做过评测。

已完成更新次数	训练损失	验证损失
0	5.674427	5.669803
50	1.269421	2.405571
100	0.368759	2.703053
150	0.216352	2.993822
200	0.176652	3.324710
250	0.175422	3.522885
300	0.168024	3.595803
350	0.166762	3.748807
400	0.170217	3.741431
450	0.159973	3.896360
500	0.164760	3.901104
550	0.155814	4.009695
600	0.159599	4.068374

起初，两种损失都接近第 20 章的均匀概率基准。到第 50 步，两者都下降了。此后，训练损失大体继续下降，验证表现却越来越差。第 600 步给训练目标分配的概率，比第 50 步好得多，但在留出的记录上反而更差。在这个划分和目标下，这种分离是 过拟合（overfitting） 的证据。如果只看训练损失下降，就会把它藏起来。

规定的规则选择第 50 步，其验证损失为 2.405571231866604。训练结束，并不意味着该选择第 600 步。我们把两个参数快照都保留下来，以便比较行为。由于第 0 步和第 50 步之间没有测量，这份记录无法指出此前最好的停止时刻。这个被反复查看的验证集，也不能再充当最终独立测试。模板相似、数据极少，即使是最低验证损失，能够支持的结论也有限。

21.3 要保存什么，取决于下一次准备做什么

都叫“检查点（checkpoint）”的文件，内容可能很不一样。如果只是用固定模型运行，参数配上对应结构和 tokenizer 就可能足够。若要沿同一条优化轨迹继续训练，还得保存 Adam 的累积状态、步数和随机数生成器状态。项目里程碑又是另一回事：它还包括整条流程确实能工作的证据。

表 21.3：三种不能混为一谈的保存对象。

保存对象	包含或关联的内容	支持什么
权重快照	参数状态字典；另行版本化的结构和 tokenizer	用这组模型参数重新计算预测
训练状态检查点	参数、Adam 状态、步数、批次及 Torch 全局随机状态、配置和数据／tokenizer 身份	从相同的下一批数据和更新状态继续训练
MG-2 里程碑材料	MG-1 过程记录、语料／tokenizer／配置、代码和依赖身份、全部曲线点、检查点文件、恢复／复跑校验及生成过程	审计并复现这次小型端到端实验

本地的 weights-selected.pt 保存第 50 步的参数，weights-final.pt 保存第 600 步的参数。训练状态文件是 resume-step300.pt 和 resume-final.pt。四个文件的实测总大小为 966540 字节，约 0.967 MB（十进制），其中包含序列化开销。这是文件大小，不是训练时的内存峰值。所有权重都由原创的小型虚构语料训练得到，没有下载预训练权重。

只保存 model.state_dict()，会漏掉 Adam 的历史统计量。重新用种子 101 创建批次生成器，会从头开始它的序列，并不会自动接上第 300 步。我们的加载程序先核对语料、tokenizer 和配置身份，再按指定结构恢复模型、优化器与生成器。读取这些由本地程序生成、仅含张量和基础容器的文件时，使用 torch.load(..., map_location="cpu", weights_only=True)。这个设置不能补出缺失的配置，也不能修复被换过的词表。

这里的恢复比较针对参数更新。MG-2 还单独保留此前的曲线和已选权重快照，因为只凭第 300 步的训练状态，无法知道此前哪一步的验证损失最低。保留这些记录，才能延续整次运行的选择规则，而不是只在恢复以后的步骤中重新挑选。

例 21.2：恢复训练，先看下一次更新能否接上。 在不中断的运行中，第 301 步抽到 [12,16,9,15,22,6,3,1]，更新前的批次损失为 0.17280904760175986。另一次运行加载第 300 步的检查点，从这里继续。批次编号和损失都一致，随后直到第 600 步的全部记录也一致。

表 21.4：实际恢复和复跑的证据。 这里的相等，是指定软件环境下的精确相等，不只是舍入后看起来一样。

比较对象	实际结果
恢复运行：第 301—600 步的批次编号与更新前损失	全部相等
恢复运行：第 300—600 步的完整划分曲线点	全部相等
恢复运行：最终参数、优化器及两种随机状态	按固定算法计算的张量／状态摘要全部相等
恢复运行：第 600 步后的下一批	都是 `[15,21,7,2,23,16,7,16]`
新进程复跑：全部 600 步及全部 13 个曲线点	全部相等
新进程复跑：所选步数、最终状态与下一批	全部相等

新进程从初始化重新开始，不加载已经训练好的权重。因此，它检查的是另一条路径，不等同于断点恢复。记录中的最终参数摘要为：

6de30d8677afce8e39691f446344d4bd67ae0eda5136743e1894152b78c54c01

代码按固定顺序，把各张量的类型、shape 和字节内容纳入摘要；检查点文件另有文件哈希，用来标识序列化后的文件。摘要相等可以检查所覆盖的状态是否改变，却不能证明实现符合预期数学规则。前面各章独立完成的机制、因果可见性和损失校验仍然必要。

21.4 一句生成结果，来自一连串重新计算的选择

生成时，把所选参数固定下来。先把提示文本编码一次，加上 BOS，运行模型，取最后一个真实输入位置的 logits。选中 token 后，追加它的 ID，再运行一次。这个简单实现每次重算整个前缀，不使用键值缓存，也不会在每次选择后重新给不断增长的文本分词。

例 21.3：完整跟踪三次选择。 这次明确加载最后第 600 步的快照，目的是检查它记住的句式。提示是 The key is red. Answer: ，冒号后有一个空格。包含 BOS 的输入 ID 为 [256,84,265,268,32,277,270,261,65,276,119,262,259]。采用贪心选择（greedy selection）：从允许输出的候选中选 logit 最高者，完全并列时取较小 ID。禁止输出 BOS 和 PAD，允许 EOS，最多生成 24 个新 token。

选择 1 · 第 600 步权重

行编号13 个 ID；末项 259 为冒号加空格
处理步骤前向计算 → 最后位置 logits → 贪心选择
行编号追加 270 = red；p = 0.986202

模型此时还没有收到生成的 red；BOS、PAD 不能被选中。

选择 2 · 权重不变

行编号14 个 ID；末项 270 = red
处理步骤用新前缀重新计算
行编号追加 46 = 句点；p = 0.999693

句点概率来自一次新的词表概率分布。

选择 3 · 权重不变

行编号15 个 ID；末项 46 = 句点
处理步骤重新计算 → 选择 257 = EOS
文本片段停止；普通文本为 red.

p(EOS) = 0.999583。图中均为排除 BOS、PAD 以前的原始模型概率。

图 21.2 — 参数不变，前缀改变三次本书原创教学图，CC BY-SA 4.0。数值与规则见正文及共享数据。

表 21.5：每次追加 ID，都会改变下一次输入。 概率列是模型在全部 291 个候选上的原始概率，尚未排除 BOS 和 PAD。

第几次选择	选择前的输入长度	所选 ID／片段	所选 logit	模型概率	结果
1	13	270／`red`	12.756594	0.986202	追加 ID 270
2	14	46／`.`	17.023494	0.999693	追加 ID 46
3	15	257／EOS	16.481496	0.999583	停止；显示 `red.`

第一行中，概率靠后的候选包括 blu（ID 274，概率 0.006150）和 key（ID 268，0.001701）。模型给 red 的概率很高，但下一行已经把 red 放进输入，重新算了一次。那时句点的概率，并不是从第一行的分布中查出的条件值。EOS 结束记录，不作为普通文本显示。

采样改变的是选择规则。给定大于零的温度（temperature）tau，先把允许候选的 logits 除以 tau，再取指数。如果还使用 top-k 采样，就只保留其中 logit 最高的 k 个候选，然后在保留的集合内重新归一化。因此，设 R 是保留的候选 ID 集合，有：

p_i^{\text{sample}}=\frac{\exp(z_i/\tau)}{\sum_{j\in R}\exp(z_j/\tau)},\qquad i\in R.

这里的 z_i 是最后一个位置给候选 i 的分数，不是注意力分数。被排除的候选概率为零。较低温度会放大分数差异，较高温度会让保留集合内的概率更平缓。例如，两个候选的 logits 分别为 2、0，且都被保留：tau=1 时，第一个候选的概率为 $e^2/(e^2+1)\approx0.880797$ ；tau=2 时为 $e/(e+1)\approx0.731059$ 。选择规则本身不会补进模型缺少的信息。

实现中，在取指数之前，会先减去保留候选中最大的温度缩放后 logit，以避免指数过大；这不会改变概率。日志同时记录原始模型概率和实际选择分布的概率。即使是贪心生成，排除 BOS、PAD 后重新归一化的概率也略有不同：上面第一行中，前者是 0.9862019510381717，后者是 0.9862021226445513。

21.5 改变线索，看看成功还能不能成立

原始的红钥匙诊断提示没有末尾空格：The key is red. Answer:。最后一个 token 是冒号字节，ID 58。训练记录却把 : 合成了一个 token，ID 259。文本能精确往返，并不意味着较短文本的编码一定是较长文本编码的前缀。我们先观察到了这个边界差异，才添加末尾带空格的诊断；原始失败仍保存在 generation-baseline.json 和 generation.json 中。

表 21.6：在相同 token 预算下的贪心续写。 代码格式保留标点；“有空格”指冒号后恰好一个普通空格。表中是实测续写，没有替模型修改拼写。

提示条件	所选第 50 步	最后第 600 步
`The key is red. Answer:`，无空格	`of Anspitap the blue.`	`Result: keyred red.`
`The key is red. Answer:` ，有空格	`bluesuersther: blue key.`	`red.`
`The key is blue. Answer:` ，有空格	`blue.`	`blue.`
`The key was red. Now blue. Answer:` ，有空格	`blue key.`	`red.`
`Old limit: 600. New limit: 750. Answer:` ，有空格	`peng.`	`shipedwion: old.`

边界对齐后，最终模型能接出熟悉的红钥匙答案。但这并不证明它能稳妥地处理改变后的线索。第四条提示已经明确把钥匙改为蓝色，最终模型却仍回答 red。第五条给出了虚构限额的新值 750，两个快照都没有接出它。这些有意改变条件的例子，足以让我们否决“模型能可靠遵循更新”这个宽泛结论；它们不是面向总体的准确率估计。

原始语料还有一条诊断：Old limit: 750. New limit:，作者给定的期望后缀是 600.。但提示并没有告诉模型新值是多少。所以，没接出 600，本身不能证明模型违背了明确给出的更新。我们保留这条诊断并说明限制，不把作者想要的后缀偷偷当成逻辑上确定的答案。后来加入的明确数值诊断，问题更完整；它有单独版本，既没有替换不利的基线，也没有用于重新选择检查点。

表 21.7：在边界对齐的红钥匙提示上改变温度。 两次采样都使用 top-k=5、种子 31415，最多新增 24 个 token。

选择规则	所选第 50 步	最后第 600 步
贪心	`bluesuersther: blue key.`	`red.`
采样，温度 0.7	`blueshindiwer: ng.`	`red.`
采样，温度 1.2	`blue.`	`red.`

两种温度得到相同的最终续写，并不矛盾：分布可以不同，某一串具体抽取仍可能选中同样的 ID。所选快照的验证损失最低，也不保证它在每条提示上都成功。如果只挑最终模型的红钥匙成功案例展示，就会同时藏起验证曲线和改变线索后的失败。

表中展示的运行都在 EOS 处停止，字节也都能正确解码为 UTF-8。生成器还明确处理两种停止条件：新增 token 预算耗尽，以及输入缓冲区达到 64 个 ID 的上限（包含 BOS）。它不会滑动上下文窗口，也不会凭空补出新位置的参数。任意生成的字节未必能组成合法 UTF-8；这条失败路径有单独测试：保留原始字节，标记 valid_utf8 为 false，并把含替换字符的显示注明为显示结果，不能称为精确往返。

21.6 复跑这次有界实验

按 code/part-iv/README.md 配好 CPU 环境后，在书籍仓库根目录使用该环境的 Python 运行：

python code/mini-gpt/forward.py
python code/mini-gpt/run.py --output /tmp/mini-gpt-reproduction
python -m unittest discover -s code/mini-gpt -p 'test_*.py' -v

训练命令会执行完整训练、断点恢复校验、新进程复跑和生成记录。把结果写入独立目录，可以保留已接受的参考证据。code/mini-gpt/README.md 说明各文件和命令的关系；notebooks/part-iv/03-mini-gpt.ipynb 在新的 CPU 内核中执行同一套核心函数，并保留输出。Notebook 是补充实践，读懂关键证据已经可以依靠上面的静态表格。

参考记录按 Asia/Shanghai 时区，生成于 2026 年 9 月 13 日。主运行的 600 步训练耗时 1.7165 秒，从第 300 步继续耗时 0.7023 秒，新进程中的训练部分耗时 1.7270 秒，整条流程耗时 4.7675 秒。这些是在同一台机器上测得的墙钟时间，各项包含的周边工作不同，进程启动和缓存条件也会变化。它们不是硬件性能对比，也不是向读者承诺的运行时间。依赖安装另算，首次安装需要网络；实验计算本身只使用本地数据。

在记录的软件环境中，精确复跑成功。跨操作系统、设备或库版本时，即使种子相同，也不能保证浮点结果和采样文本完全一致。做这种复跑，应先核对数学不变量、划分与 tokenizer 身份、损失是否有限、是否没有未来泄漏，以及过拟合证据的方向，再明确报告数值差异。本次没有实测跨平台或 GPU 复现。

21.7 Part IV 概念检查点：让证据接住结论

现在，我们能从输入坐标一直跟到一个头的加权结果，准确指出查询能看到哪些键，把独立的头组合起来，再把注意力放进残差、归一化和 FFN 构成的块。位置信息和 token 边界都有明确接口。架构与训练目标是不同的选择，单凭任何一种选择，都不能把网络变成可靠助手。

同一张结构图如今还对应着真实的 CPU 数组和代码。MG-1 验证从文本到 logits／损失的路径；MG-2 验证参数更新、状态保存、断点恢复、指定条件下的可复现性，以及可逐步检查的生成。曲线和改变线索的诊断既展示进展，也标出限制。生成文本没有证明模型执行过真实订单操作、查过事实数据库，或掌握了通用解题推理（reasoning）能力。

当数据、模型规模和周围的学习过程扩大时，又会发生什么变化？第 22 章将通过现代 LLM 的生命周期来回答。这个小模型让我们有了可以具体分析的机制，也留下了一种检验习惯：说清希望得到什么能力，选择可能揭穿失败的证据，并把失败保留下来。

21.8 习题

一个新验证集的两条记录，分别包含 2 个和 3 个有效目标；逐目标损失为 [1,3]、[2,2,7]。求按事件平均的损失，以及让两条记录等权的平均损失。本项目用哪一种？
一位同事认为应选择第 600 步，因为其训练损失为 0.159599，远低于第 50 步的 1.269421。用表 21.2 和既定规则评价这个选择。这张表能否找出第 1—49 步中的最佳一步？
两个程序恢复相同的第 300 步权重。一个重新创建 Adam，并把批次种子重设为 101；另一个恢复 Adam 和批次生成器状态。它们下一次预测必然不同吗？下一次参数更新必然相同吗？解释两者的区别。
两个允许候选的 logits 为 2、0。温度为 2 时，第一个候选的概率是多少？top-k=1 会怎样改变这个概率？此时再改变温度，一定能改变续写吗？
表 21.6 中，最终模型对熟悉的红钥匙提示和改为蓝钥匙的提示都回答 red。分别说出证据支持的狭窄成功，以及不能支持的宽泛结论。为什么原始“新限额未给定”的诊断证据更弱？
新 tokenizer 仍有 291 项，但交换了 ID 270 和 274 的含义。Shape 和权重加载都能通过。什么会出错？哪项检查点校验应该拒绝它？
一条生成过程在达到 24 个新增 token 的预算后结束，没有出现 EOS。显示界面呈现了一句整齐的话。记录必须保留什么，别人才有依据解释这个显示结果？
准备把表示宽度扩展到 32，仍保留两个头。为验证这个改动列一份证据清单，包括一个 shape 后果、一项信息流测试、一项损失检查，以及一种避免只挑好看输出的评测选择。

21.9 参考解答

总损失为 15，共五个事件，本项目的平均值为 3。两条记录的均值分别为 2 和 11/3，再等权平均得到 17/6，约 2.833333。后一种算法让短记录中每个事件的权重更大，不符合规定的损失平均规则。
按规则应选第 50 步，因为其验证损失 2.405571 是全部测量值中的最低值。第 600 步的验证损失更差，为 4.068374。训练拟合程度有用，但不是这里的选择标准。表中没有测量第 1—49 步；把曲线点连起来，不会凭空产生这些评测。
在下一次更新以前，只要模型参数、输入和评测设置相同，就可能得到相同预测，不需要优化器状态。后续更新却未必相同：Adam 的累积状态不同，重新开始的生成器也未必抽到同一批。要按约定继续训练，必须恢复这两类状态，同时保持身份信息及其余运行条件一致。
除以 2 后，logits 为 1、0，所以概率为 $e/(e+1)\approx0.731059$ 。top-k=1 只保留较大的 logit，任何正温度下它的选择概率都是 1。k 更大时，即使分布不同，具体抽取也可能生成同一段文本。
模型在 token 边界对齐时，能复现熟悉训练句式的后缀。它没有证明自己能可靠使用修订后的线索；明确改成蓝色后仍答红色，正好反驳了这个结论。未给出新限额的提示并没有提供答案，因此没接出作者期望的 600，本身不构成对已给信息的违背。
嵌入查表和输出列的含义将不再对应保存的模型，即使所有 shape 都合法。Tokenizer 身份校验应拒绝这次加载。如果要有意迁移，就必须更新所有受影响的参数和数据接口，并记录新版本。
保留提示及其 token ID、所选检查点和 tokenizer、采样设置及种子、全部生成 ID／字节、停止原因，以及 UTF-8 是否有效。预算耗尽不同于模型选择 EOS；不能把补过标点或用替换字符解码的显示结果，悄悄当成原始输出。
每头宽度变为 16，缩放分母变为 16 的平方根；Q/K/V 和残差宽度都要同步。替换未来 ID 后，较早真实位置的 logits 应保持不变。全零 logits 的平均损失仍应只对有效目标计算，等于 $\ln(291)$ 。保留固定划分和步数选择规则，保存整条验证曲线与全部原始诊断，再为希望更宽模型改善的能力添加有独立版本的新测试。

21.10 参考文献

PyTorch，Adam，2.7 版。说明优化器保存的梯度统计量、配置和状态；本章曲线来自本书自己的实验。
PyTorch，Saving and Loading Models 与 torch.load，2.7 版。支持模型状态与继续训练所需状态的区别，以及本项目明确在 CPU 上读取张量／容器的做法。配套记录固定了实现和已测版本。
PyTorch，Reproducibility，2.7 版。说明如何控制随机来源、使用确定性算法，以及为什么仅有种子不能保证不同软件版本和平台间完全一致。

第 21 章 训练、保存和审问 Mini GPT