实践 P08：训练迷你 Transformer

本实践把第 20 章的完整数据流缩小到一台普通电脑能够检查的规模。你将训练一个字符级、仅解码器的 Transformer，让它根据前文预测下一个汉字或标点。

这不是对 GPT 等大模型规模的复现。它复现的是核心机制：Token 编号、Embedding、位置、因果多头注意力、前馈网络、残差、归一化、交叉熵和自回归生成。模型很小，语料也很短，所以生成内容会重复或断裂；这些失败恰好能帮助我们区分“架构能够运行”和“系统获得通用语言能力”。

实践概览

项目 本实践配置
学习目标 亲手验证因果注意力、张量形状、损失下降与逐 Token 生成
先修知识 Python 基础；第 13、15、19、20 章主线
默认 Tokenizer 字符级，一个不同字符对应一个 Token
数据 随项目提供的原创中文科普短文，不依赖外部下载
快速验证 A 级，CPU，约数十秒内完成安装后的代码检查
标准训练 B 级，消费级 CUDA GPU；CPU 也能运行但更慢
软件 Python 3.11 或 3.12，PyTorch 2.8.0
成功标准 测试通过；损失低于随机基线；保存检查点；生成可重复执行

先明确我们复现什么

大模型训练需要大量数据、计算集群和复杂工程。把参数量缩小之后，我们不能声称复现它们的知识、语言水平或训练结论。但只要缩小版本仍保留相同的信息流，就能验证基础机制。

本实践保留：

  • 输入与目标错开一个 Token；
  • Q、K、V 和缩放点积注意力；
  • 未来位置被因果掩码屏蔽；
  • 多个注意力头并行后重新合并；
  • Pre-LN、残差连接与逐位置前馈网络；
  • 多层输出到词表 logits；
  • 交叉熵训练与自回归采样。

本实践有意省略：

  • 大规模子词 Tokenizer 与清洗后的海量语料；
  • 分布式训练、混合精度和学习率预热；
  • 旋转位置编码、分组查询注意力、SwiGLU 等现代变体；
  • 指令微调、偏好优化、工具调用和安全系统；
  • 面向生产推理的 KV Cache 与批处理服务。

把边界说清，是可复现实验的一部分。

项目文件

实践代码位于 examples/p08-mini-transformer/：

p08-mini-transformer/
|-- corpus.txt       # 随书原创 UTF-8 语料
|-- requirements.txt
|-- train.py         # 模型、训练、评估、生成和检查点保存
`-- test_model.py    # 形状、因果性和参数更新测试

所有命令都从本书项目根目录运行。Windows PowerShell 示例使用 .venv；Linux 或 macOS 把激活脚本路径换成 .venv/bin/activate 即可。

第 1 步：建立隔离环境

cd examples/p08-mini-transformer
python -m venv .venv
\.venv\Scripts\Activate.ps1
python -m pip install -r requirements.txt

如果 PowerShell 的当前执行策略不允许运行激活脚本，不需要修改系统策略，可以直接调用虚拟环境中的 Python：

.venv\Scripts\python.exe test_model.py

安装完成后记录版本：

.venv\Scripts\python.exe --version
.venv\Scripts\python.exe -m pip show torch

requirements.txt 锁定 PyTorch 2.8.0。CUDA 是否可用取决于你安装的 PyTorch 构建和显卡驱动；代码默认自动选择 CUDA，否则回退到 CPU。

第 2 步：先运行结构测试

不要先看生成文字是否有趣。一个掩码错误的模型也可能在训练阶段得到很好看的损失。先执行：

.venv\Scripts\python.exe test_model.py

测试包含三项：

  1. 输入 [2,8] 能否得到 logits [2,8,11] 和有限损失；
  2. 修改序列后半段，前半段的 logits 是否完全不变；
  3. 一次反向传播和优化步骤是否真正改变参数。

第二项是因果性的直接检查。假设两条序列前四个 Token 相同、后四个不同。因果模型在前四个位置看不到未来，因此这四个位置的输出必须一致。测试使用 torch.testing.assert_close 比较结果。

预期输出：

P08 structural tests passed

第 3 步：检查训练样本怎样产生

corpus.txt 是一篇随书原创的短文。代码收集其中所有不同字符，排序后建立词表：

vocabulary = sorted(set(text))
stoi = {character: index for index, character in enumerate(vocabulary)}

这不是适用于真实产品的 Tokenizer，但每一步都可观察。设一段编号为：

[12, 7, 31, 5, 9]

上下文窗口为 4 时，输入和目标是：

输入：[12, 7, 31, 5]
目标：[ 7,31,  5, 9]

一次前向传播同时产生 4 个位置的预测。第一个位置根据 12 预测 7，第二个根据 12,7 预测 31，依次类推。因果掩码保证第一个位置不能直接读取后面的正确答案。

语料按连续位置切分：前 90% 用于训练，后 10% 用于验证。这种切分保留了时间顺序，但数据太小，不足以估计真实语言泛化。验证损失在这里主要用于演示接口和发现明显过拟合。

第 4 步：跟踪多头注意力形状

代码中，输入 x 的形状是 [B,T,C]，其中 \(C=d_{\text{model}}\)。一个线性层同时产生三组投影：

q, k, v = self.qkv(x).chunk(3, dim=-1)

每一组随后从 [B,T,C] 改成 [B,H,T,D_h]：

tensor.view(batch, time, n_heads, head_dim).transpose(1, 2)

注意力分数由下面一行产生：

scores = q @ k.transpose(-2, -1) / math.sqrt(head_dim)

结果形状是 [B,H,T,T]。掩码在 softmax 以前加入：

future = ~causal_mask[:, :, :time, :time]
minimum = torch.finfo(scores.dtype).min
scores = scores.masked_fill(future, minimum)
weights = softmax(scores, dim=-1)

如果先 softmax 再把未来位置乘 0，保留下来的权重之和通常不再是 1。这个实现细节会改变输出尺度。

各头与 Value 相乘后得到 [B,H,T,D_h]，再转置、合并回 [B,T,C]。阅读代码时建议在纸上写出每一次 view、transpose 前后的形状。

第 5 步：运行冒烟训练

先进入实践目录，再运行：

cd examples/p08-mini-transformer
.venv\Scripts\python.exe train.py --mode smoke --device cpu

冒烟模式只训练 3 步，使用 1 层、32 维模型。它的目标是检查：数据能够读取、前向传播和反向传播能够完成、损失是有限数、检查点能够写入、生成循环能够结束。

程序开头打印一行 JSON，其中 random_loss_baseline 约为 \(\log V\)。如果模型对 \(V\) 个字符给出均匀概率，交叉熵就是这个值。3 步训练不保证明显低于基线；它只验证链路。

随书版本的实测记录

2026 年 8 月 18 日，锁定环境使用 Python 3.12.13、PyTorch 2.8.0 和 NumPy 2.3.2，在 CPU 上通过了三项结构测试。冒烟模型包含 26,144 个参数，词表为 402 个字符，随机基线损失为 5.9965；3 步训练完成了前向、反向、评估、生成和检查点保存。

同一台机器还用标准模型配置运行了 100 步 CPU 验证。模型包含 852,992 个参数，结果如下：

步骤 训练损失 验证损失
0 6.0170 5.9891
100 2.8955 5.0275

训练损失下降得明显更快，验证损失也下降但仍高得多。这与短语料下迅速记忆局部模式的预期一致。它验证了学习链路，不应解释为获得通用语言能力。

随后会看到类似结构：

step    0 | train ... | validation ...
step    1 | train ... | validation ...
...
generated:
人工智能...
checkpoint: ...\checkpoint.pt

具体数字受设备和 PyTorch 实现影响。固定随机种子减少随机差异，但跨硬件不承诺逐位相同。

第 6 步：运行标准训练

CUDA 环境：

.venv\Scripts\python.exe train.py --mode train --device cuda

没有 CUDA 时也可明确使用 CPU，并先减少步数：

.venv\Scripts\python.exe train.py --mode train `
  --device cpu --steps 100

标准配置使用 4 层、128 维、4 个注意力头、64 Token 上下文和 1000 个训练步骤。每 100 步分别估计训练与验证损失。训练中应观察：

  • 训练损失是否从接近随机基线的位置下降；
  • 验证损失是否也先下降，而不是从一开始就持续上升；
  • 生成结果是否从近似随机字符变为包含语料中的局部结构；
  • 后期是否出现重复，暴露小数据记忆与过拟合。

不要用一条最好看的生成样本判断成功。至少保留完整损失记录，并使用同一提示词比较不同检查点。

第 7 步：读懂模型的生成

训练结束后，默认提示词是“人工智能”。每次循环只取最后一个位置的 logits：

logits, _ = model(context)
next_logits = logits[:, -1, :] / temperature
next_token = torch.multinomial(softmax(next_logits, dim=-1), 1)

新 Token 拼到上下文末尾，然后再次调用模型。为了不超过位置嵌入范围，代码只保留最近的 block_size 个 Token。

这个教学实现每一步都会重新计算上下文，没有加入 KV Cache。这样写速度较慢，却能让生成逻辑保持清晰。第 20 章已经解释，生产推理通常缓存旧 Token 的 Key 和 Value。

成功标准

完成实践时，用下面的证据验收，而不是凭感觉判断：

检查项 通过条件 它能证明什么
结构测试 三项测试全部通过 形状、因果性和更新链路基本正确
随机基线 初始损失与 \(\log V\) 数量级接近 词表与损失计算没有明显异常
训练损失 标准训练后显著低于初始值 参数学会了语料中的可预测模式
验证损失 至少在早期下降或保持有限 模型没有只靠直接读取未来答案
检查点 checkpoint.pt 可生成且可读取 训练状态能够保存
生成 能从提示词逐 Token 继续 自回归路径完整

这些结果不能证明模型理解文本、掌握事实或能迁移到开放任务。

常见错误与排查

找不到 PyTorch

如果看到 ModuleNotFoundError: No module named 'torch'，确认运行脚本的 Python 与安装依赖的 Python 是同一个。使用完整路径 .venv\Scripts\python.exe 最容易排除环境混淆。

CUDA 不可用

如果看到 CUDA was requested but is not available，先运行 --device cpu。若必须使用 GPU，再核对显卡驱动、PyTorch 构建和 torch.cuda.is_available()，不要把模型代码问题与 CUDA 环境问题混在一起。

损失不再是有限数

先恢复默认学习率和配置，检查输入编号是否落在词表范围内。代码已经包含梯度裁剪；若自行扩大模型，可以进一步降低学习率并检查混合精度设置。

训练损失异常低，生成却完全失败

首先重新运行因果性测试。然后检查训练目标是否错开一位、生成时是否只读取最后位置、是否误把真实后续文本送入生成循环。

生成内容重复

对这个小语料模型而言，重复是预期现象。可以调高温度或扩大 top_k，但采样参数只能改变选择方式，不能补足语料和模型容量。

三个值得亲手做的对照实验

  1. 把 block_size 从 64 改为 16，保持其他配置相同，比较验证损失和生成中的长距离连贯性。
  2. 暂时破坏因果掩码并运行测试，观察自动检查如何在训练前发现信息泄漏；实验后恢复代码。
  3. 把语料复制多次再训练，观察训练损失下降与生成质量之间并非线性关系：更多重复字符不是更多知识。

每次只改变一个因素，并记录配置、随机种子、运行时间和损失。否则无法知道差异来自哪里。

进阶复现路线

完成最小版本后，可以按以下顺序增加复杂度：

  1. 使用成熟的 BPE 或字节级 Tokenizer，比较词表大小与序列长度；
  2. 换成更大、许可清楚的中文语料，并建立真正独立的验证集；
  3. 加入学习率预热与余弦衰减，记录相同计算预算下的差异；
  4. 使用自动混合精度，测量显存和吞吐变化；
  5. 实现 KV Cache，验证生成结果一致并比较单 Token 延迟；
  6. 阅读并对照《Attention Is All You Need》的编码器-解码器实现。(Vaswani 等 2017年)

从第 1 步到第 6 步，实验逐渐接近真实训练工程，但仍不等于复现某个商业大模型。复现报告必须始终注明数据、参数量、计算量和训练目标的差异。

实践小结

P08 把 Transformer 从结构图还原成了可检查的数据流：字符变成编号，编号变成向量，因果注意力只读取过去，Block 反复加工表示，输出层预测词表概率，交叉熵把错误送回所有参数，生成循环再逐 Token 使用模型。

最重要的实验习惯也同样明确：先测试机制，再观察指标；先定义成功标准，再看漂亮样本；缩小模型时说明保留了什么，也说明没有复现什么。

参考文献

Vaswani, Ashish, Noam Shazeer, Niki Parmar, 等. 2017年. 《Attention Is All You Need》. Advances in Neural Information Processing Systems 30, 5998～6008. https://papers.nips.cc/paper/7181-attention-is-all-you-need.