主题
第 26 章 让同一份计算跑得更快：精度、编译和注意力
模型写对之后，训练时间仍可能差很多。瓶颈有时是矩阵乘法，有时是显存读写，有时是 Python 发起太多小操作。把所有问题都归结为“GPU 不够快”，就很难知道下一步该改哪里。
GPT-2 复现课第二部分依次讨论 TF32、BF16、torch.compile、FlashAttention 和词表补齐。本章解释每一步减少了什么成本，以及怎样检查它没有把问题改掉。视频章节标题中的毫秒数属于讲者当时的硬件、软件和输入配置；本书没有在 GPU 上重新测量，也不把这些数字当作今天的性能承诺。官方章节描述。
26.1 先定义你到底在计时什么
一次训练迭代包括取数据、前向、反向、梯度处理和参数更新。GPU 调用通常是异步的：CPU 发出任务后就能继续往下走，所以用两个 CPU 时间点夹住一段代码，可能只测到提交任务的时间。要测完成时间，需要在合适的位置同步，或者使用 GPU 事件。本书只在 CPU 上做数值测试，不用 CPU 时间推断 GPU 吞吐。
第一次调用还可能包含编译、内存分配和后端初始化。测稳定吞吐时，应先预热，再多次测量；测服务首次响应时，初始化反而可能是用户真实承受的成本。两种指标都合理，不能把它们混为一个数字。
固定课程代码每次优化处理的 token 数为 B×T×K×R，其中 K 是梯度累积次数，R 是进程数。吞吐定义为这个数量除以迭代耗时。但课程计时从循环顶部开始，某些步骤还包含验证、HellaSwag 和生成。因此这些特殊步骤的 tok/sec 不能与纯训练步骤直接比较。看日志时先识别采样和评估周期，再判断吞吐是否真的退化。训练循环与计时位置。
26.2 少用一些有效位，换来什么
浮点格式有三部分：符号、指数和尾数。指数决定能表示多大、多小的数；尾数决定相邻可表示数之间的间隔。把一种格式称为“16 位”，并不能说明它适合哪类训练。
| 格式或计算模式 | 指数位 | 尾数位（不含隐含首位） | 本章关注的作用与限制 |
|---|---|---|---|
| FP32 | 8 | 23 | 范围和精度较充足，常用于参数与敏感计算 |
| FP16 | 5 | 10 | 存储较小，数值范围较窄，小梯度可能下溢 |
| BF16 | 8 | 7 | 范围接近 FP32，但尾数较粗；不保证任意训练都稳定 |
| TF32 矩阵乘计算模式 | 8 | 10 | 支持的 NVIDIA 硬件可在 FP32 接口下使用较短乘法有效位；不是独立的普通张量存储 dtype |
上述格式的位数与混合精度背景可核对 NVIDIA 混合精度说明。
把 FP32 权重直接全部改成 BF16，与开启 autocast 不是同一件事。autocast 根据算子选择计算类型；通常仍可保留 FP32 参数，并让部分矩阵乘法采用更低精度。某些求和、归一化或损失计算需要更谨慎地处理。课程使用 BF16 autocast 包围前向和损失，然后在外面调用 backward。PyTorch 2.6 自动混合精度示例。
FP16 常配合 loss scaling：先把损失乘以较大的正数 S，反向得到 Sg，再在更新前除以 S。这样能让原本过小的梯度在中间计算中不那么容易归零。但缩放不能修复所有溢出，也不能补回已经丢掉的有效信息。若要裁剪梯度，必须先解除缩放；否则阈值 1 裁剪的是 Sg，实际允许的原始梯度范数会变成 1/S。BF16 的指数范围较大，课程没有为它使用 GradScaler；这不等于 BF16 的数值误差可以忽略。
固定代码中的 set_float32_matmul_precision('high') 是矩阵乘法计算精度策略，不会把所有张量改成 TF32。它的效果依赖后端和硬件。复现时记录策略，比只记录“dtype=float32”更完整。相关 API 可以在本书固定环境中读取 help(torch.set_float32_matmul_precision)，避免把未来版本默认值倒写进 2024 年课程。
26.3 编译减少的是调度和中间开销
一次 LayerNorm 或 GELU 看起来只有一行 Python，背后却可能包含多次读写、启动多个 GPU kernel。kernel 是在设备上执行的一段计算程序。若输入很小，启动和调度成本可能比计算本身还显著。
torch.compile 可以捕获一部分程序结构、生成优化后的执行方案，并在条件允许时融合操作。比如原本依次计算中间张量 a、b、c，需要多次写回显存；融合后可能在更近的存储中直接传递中间结果。它不会自动证明你的损失写对，也不会保证任意输入模式都更快。变化的形状、数据依赖的 Python 控制流和不支持的操作都可能影响编译效果。
课程中的历史演示展示编译后的加速，但固定仓库最终把 use_compile 设为 False，并注释说明与 HellaSwag 评估和生成有交互问题，待修复。解释这门课时不能同时说“按最终文件直接运行”和“默认使用编译后的性能”。这两个描述对应不同状态。固定版本编译开关。
工程上应该先保留不编译的正确基线，再对固定批次比较前向损失、梯度和训练短跑。编译后的微小浮点差异不一定是 bug，但差异容限应根据 dtype 和任务设定，而不是看到文本还算通顺就忽略。
26.4 FlashAttention 没有把注意力公式换掉
普通实现常把整个 T×T 注意力矩阵写到显存。T 翻倍，这一矩阵的元素数变成四倍。FlashAttention 的关键思路是分块组织精确注意力计算，减少慢速显存与片上存储之间的数据搬运；它不是用稀疏连接近似原来的注意力。FlashAttention 原论文。
为什么 softmax 可以分块？对一行分数 s，常见的数值稳定写法先取最大值 m，再计算 Z=Σexp(s_j−m)，输出是 Σexp(s_j−m)u_j/Z。处理了一部分位置后，只需保留当前 m、Z 和加权和。设新块最大值为 m_block，更新参照点 m′=max(m,m_block)。把旧的分子与分母都乘以 exp(m−m′)，再加上新块按 m′ 缩放后的贡献。保留迄今见过的最大值，才能继续避免指数计算溢出。这样不用永久保存整行概率，也可以得到同一个数学结果。
举个缩小的例子，旧块最大分数为 2，新块出现最大分数 5。旧分母原来统计的是 Σexp(s_j−2)，换到新的参照点后要乘 exp(2−5)=exp(−3)。分子做同样缩放，分子分母的共同参照点保持一致。反向时可以重算部分中间量，用额外计算换取更少的存储。具体 GPU kernel 还有更多线程与内存安排，本书不把这一段解释冒充完整实现。
PyTorch 的 SDPA 接口负责根据输入选择可用实现。CPU 上调用同一个接口也能得到正确注意力结果，但这不构成 FlashAttention GPU 性能验证。还要注意 dropout_p：这个函数根据传入概率执行 dropout，调用方应在评估时传 0，而不能仅依赖外围 model.eval()。本书测试显式使用 dropout_p=0。SDPA 文档中的后端与 dropout 说明。
26.5 用前向和反向一起验证替换
本书创建形状为 2×2×5×4 的 Q、K、U，以 float64 分别计算手写的因果注意力与 SDPA。两者输入相同，遮罩相同，缩放均为 1/√4，且没有 dropout。为比较反向，定义标量损失 L=ΣY²，然后分别对 Q、K、U 求导。
实际结果如下。这些是 CPU 数值检查，不是速度测试。
| 检查 | 最大绝对误差 |
|---|---|
| 注意力前向输出 | 4.44×10⁻¹⁶ |
| Q、K、U 的反向梯度 | 4.88×10⁻¹⁵ |
| 修改未来 token 后，模型早先位置的 logits | 0 |
为什么还需要第三项？前两个实现可能同时犯了同一个遮罩错误。独立的因果性检查改变从零起编号的下标 4 及其之后的 token，要求下标 0—3 的输出完全不变。它验证的是模型应满足的性质，而不是两个函数是否碰巧一致。完整执行命令、输入和断言在 lab.py，结果保存在 results.json。
这类测试也有边界：只检查一组尺寸，不能覆盖非连续输入、混合精度、长序列和所有后端。迁移到 GPU 后，要重新检查目标配置。float64 CPU 的容限不能原样用于 BF16。
26.6 为什么把 50257 改成 50304
硬件通常喜欢对齐良好的矩阵尺寸。课程把随机训练模型的词表维度由 50257 补到 50304，使最后的矩阵乘法更适合当时的执行路径。差值是 47，新增共享嵌入参数为 47×768=36096。这个改变不会自动给 tokenizer 添加 47 个有意义的新 token。
这里有两个容易漏掉的后果。第一，原版预训练权重仍对应 50257 的词表，不能直接把新的配置交给旧权重导入函数，忽略形状不匹配。第二，如果输出 softmax 包含这 47 个额外位置，它们会占据概率质量；随机初始化时总损失基线也跟着词表大小变化。即使这些位置从未成为正确标签，也不代表它们永远没有概率。
因此，用补齐模型生成文本时应限制可采样编号到 tokenizer 真正支持的范围，例如在 softmax 前把 logits[...,50257:] 设为负无穷。固定课程脚本的采样没有显式做这一限制；本书把它当作需要留心的实现边界，不能假定任意训练阶段都不会抽到无效编号。不要把“计算维度补齐”和“扩充 tokenizer 词表”当作一个操作。
练习与答案
以下为本书补充练习。
- T 从 1024 增到 2048，普通注意力分数矩阵的存储量变化多少？投影矩阵参数量是否随之变化？
- FP16 使用缩放 S=1024，直接对缩放后的梯度裁剪到范数 1，相当于对原始梯度设什么阈值？
- 为什么第一次 torch.compile 调用的时间不适合直接作为稳定训练吞吐？
- 若 CPU 上 SDPA 与手写注意力完全一致，可否据此报告 FlashAttention 加速倍数？
答案：分数矩阵元素变为四倍，QKV 投影参数并不随 T 变化；原始梯度阈值变成 1/1024；第一次调用可能包括编译和初始化，应区分冷启动与预热后性能；数值等价不提供 GPU 性能证据，不能报告加速倍数。
本章来源与验证边界
- GPT-2 复现课及官方章节表：历史演示顺序。
- build-nanogpt 固定代码：开关、dtype、词表补齐与计时位置。
- PyTorch 2.6 AMP 与 SDPA：接口语义。
- FlashAttention 论文：精确注意力的 IO 优化思想。
本书验证了 CPU 注意力数值与因果性质，没有验证 GPU kernel 选择、编译加速、显存节省或 BF16 收敛。