主题
第 18 章 训练性能及 CPU、MPS 数值差异
优化一个实现时，应该同时回答两个问题：它少花了哪些时间或内存，仍然在计算什么。把模型转为 bfloat16、替换注意力内核、增加批量或扩大词表，都可能提高吞吐，却也可能改变数值、目标函数或训练预算。
本章覆盖上游 00_orig.py、01_opt_single_gpu.py、02_opt_multi_gpu_ddp.py，以及 CPU/MPS 随机采样 notebook。只运行了 CPU 数学和随机数检查，没有执行 CUDA、MPS、torch.compile 性能或真实 DDP。性能脚本目录
18.1 先建立可比较的基线
上游 baseline 基于主训练代码，但增大上下文和批量，并换成更长文本，新增区间 token/s 和内存统计。它与第 13 章短文训练的资源占用不同，不能把两者的差异全部归因于某个优化。
优化版集中做了合并 QKV、SDPA、原生 LayerNorm/GELU、bfloat16、fused AdamW、pinned memory、编译、词表对齐和更大批量。要知道单项贡献，最好每次只改变一项，并保留数值与训练轨迹检查。累积加入多项的表格描述的是那条实验路线，不保证各项独立加速比可以相乘。
上游 README 的 A100 吞吐数字是作者在其硬件与软件组合上的结果，本书没有复测，不把它们列为本地实验成绩。读者应优先复现固定输入、固定步数和计时范围，而不是追逐一个没有环境说明的目标数字。
18.2 减少中间张量和调用开销
原生 LayerNorm 和 GELU 可以调用优化过的实现，避免教学版的若干逐步运算。替换时仍要一致：LayerNorm epsilon、总体方差定义、GELU 是否采用 tanh 近似都影响结果。本书第 6 章已经展示即使公式相同，常数精度也会造成微小差异。
SDPA 把注意力核心交给库，可能使用适合设备的融合实现，减少显式 [B,H,T,T] 中间存储。合并 QKV 减少三次独立线性投影调用。普通因果模式通过 is_causal=True 表达，不再让每层永久保存完整 float mask。这些机制的正确性可用第 4—5 章的前向与梯度对照检查。
18.3 TF32、bfloat16 与 autocast 不是同一件事
TF32 相关设置影响某些支持设备上的 float32 矩阵乘法路径；bfloat16 改变张量数据类型；autocast 则在运算级选择类型，通常不等于把整个模型参数永久改成 bfloat16。上游优化脚本明确执行 model.to(device).to(torch.bfloat16)，并非一个保留全部 float32 参数、只在前向里 autocast 的示例。
bfloat16 指数范围较宽，但有效尾数位少，较小更新可能因舍入难以反映在低精度参数里。优化器状态的 dtype 与实现同样要核对。低精度方案应检查 loss 是否有限、梯度是否合理，以及与高精度参考的误差，而非仅看速度。
上游按 CUDA capability 条件设置 float32 matmul precision，并注明这不影响已使用 bfloat16 的矩阵乘法。不能把“开启 TF32”与“bfloat16 自动使用合适内核”重复当成同一运算的两份收益。
18.4 pinned memory 在主机，不是预留显存
DataLoader 的 pin_memory=True 使用页锁定的主机内存，可能改善 CPU 到 GPU 的传输。它不是预分配 GPU 权重或重复使用显存。与合适的非阻塞传输、流水线配合时，才可能更好地重叠读取和计算。
上游数据传输仍主要写 .to(device)，因此不能仅凭 pin_memory 开关就声称所有数据搬运都异步重叠。增加 num_workers 也可能增加进程、内存和序列化成本。短数据、内存中已构造好的 tensor 数据集，未必因更多 worker 更快。
fused AdamW 把部分更新操作融合，减少调用与中间数据开销。它受设备、dtype 和版本支持限制。CPU fallback 的存在不意味着脚本里所有优化选项都在 CPU 可用，应分别检查 compile、fused optimizer 和 bfloat16 算子支持。
18.5 编译收益需要排除首次成本
torch.compile 可能合并操作、优化图和生成更合适的代码。第一次调用含编译成本，输入形状或控制流变化可能触发新图或回退。因此要同时报告冷启动和稳定态，而不是丢掉慢的几步后只说“总共加速多少”。
编译后的模块保存名可能带包装前缀，保存时应理解是否需要访问原始模型；否则另一个未编译模型加载会出现键名不匹配。上游文件末尾提供了相关保存提示，但没有构成完整跨版本检查点协议。
编译工具链失败可以先跳过编译继续验证数学正确性，不必为了运行教学实验而安装不明来源编译组件。本书共享 CPU 环境没有执行编译性能路径。
18.6 词表对齐增加的行不是合法文本 token
上游把词表从 50,257 扩为 50,304，使矩阵维度更适合某些硬件路径。额外 47 行不自动对应 tokenizer 的有效 ID。训练标签仍只含原词表 ID，但输出 softmax 若包含新增行，分母和优化目标会改变；生成时若选到新增 ID，解码器可能无法处理。
本书小反例把一个无效补齐类的 logit 设为 100，贪心立刻选到它。明确屏蔽 logits[..., valid_vocab:] = -inf 后才回到合法候选。是否在训练也屏蔽，应作为目标定义的一部分记录。矩阵对齐是工程选择，不应悄悄改变 tokenizer 协议。
18.7 计时边界能让同一个数字有不同含义
上游 train_model_simple_with_timing 在评价间隔计算新增 token 数与经过时间，CUDA 路径使用 event 和 synchronize，CPU 用墙钟。第一次区间很短且可能含初始化开销，单 GPU 平均值跳过第一次；DDP 变体的累计规则不同，比较时必须看实际实现。
定时器在评价前重启，因此下一段时间可能包含前一次评价、打印和 epoch 末尾生成的开销；分子却只增加训练 token。这可以作为某种包含附带工作的有效吞吐，但不能称作纯优化器步骤吞吐。若想只测训练，应在评价完成后重新开始计时；若想端到端，应清楚包含数据加载、评价和保存。
allocated memory 是当前被张量使用的分配量，reserved memory 是分配器保留量，二者都不等于峰值。某些脚本把 device 局部变量重新赋成当前 CUDA 整数编号，后续传递时应保持类型约定清楚。测量代码本身也需要审查。
18.8 DDP 同步参数梯度，不会自动汇总所有指标
DDP 版由 torchrun 提供进程环境，设置设备和进程组，使用 DistributedSampler，编译后包装模型。数据下载由指定进程协调，并通过 barrier 等待其他进程看到文件；这隐含文件路径可访问的条件。训练每个 epoch 调 set_epoch，各进程执行相同数量的同步步骤。
梯度平均在等量局部样本时对应整体平均。配套 CPU 数学模拟把四个样本分成两半，局部梯度平均与整体梯度最大差约
上游吞吐把各 rank 的 token 数 all-reduce 求和，再除以本地 elapsed；更严格的全局端到端测量应考虑最慢 rank 和同步边界。验证 loss 没有同样自动加权 all-reduce，rank 0 打印自己的值不能被称为全局验证 loss。DistributedSampler 默认补齐样本也可能造成评估重复，需要记录有效样本权重。
只有 rank 0 生成样例时，还要考虑 DDP forward 内部可能发生的同步或 buffer 广播。其他 rank 如果进入下一轮训练，控制流不一致可能等待或死锁。独立生成通常应使用未包装模块，并在各 rank 明确同步阶段；不能因为只想“少打印几行”就忽略分布式调用顺序。上游脚本是单机多 GPU 教学路径，不能直接等同已验证的多机部署。
18.9 相同 seed 不保证跨设备相同随机数
上游 CPU/MPS notebook 先创建概率，再分别在 CPU、MPS、CUDA 调用 multinomial，每次设置相同 seed；第二个函数把概率搬回 CPU 抽样，再把结果搬到原设备。它隔离的是随机采样位置，不是在证明所有设备上的模型 logits 一样。随机采样 notebook
不同设备可能使用不同随机数生成器与实现，相同整数种子没有跨算法逐位一致保证。本书只验证同一 CPU、同一 seed、同一调用顺序会重现六个随机数；在前面额外抽一个数，后续序列就改变。把采样统一放 CPU 有助于控制一个变量，但引入设备传输，也不能消除前向数值差异。
18.10 浮点加法不满足实数里的结合律
float32 小例子中，
排查跨设备差异应分层：先固定输入和参数，在 eval/no-grad 下比较 logits；再比较 softmax；最后把同一概率交给统一采样器。还应记录 dtype、库版本、后端和误差容差。若直接比较两段长随机文本，一次微小分叉会让后续全部不同，几乎无法定位原因。
上游其他 notebook 对部分 MPS 版本采用 CPU 回退，这属于固定代码的兼容策略。本书运行环境没有 MPS，没有验证某个新版是否已经消除特定问题，也没有把 CPU 检查写成跨设备测试。
18.11 练习与答案
补充题一：提高 batch_size 后 tokens/s 增加，是否说明同样的训练步骤数具有相同预算？答案：不是，每步 token 数也增加了，应同时记录总 token 和更新次数。
补充题二：DDP 已 all-reduce 梯度，验证指标为何还需聚合？答案：梯度同步与用户计算的统计量是两条路径，准确率、损失分子分母不会自动变成全局值。
补充题三：把所有采样放 CPU 后，跨设备生成一定一致吗？答案：不一定，模型前向分布可能不同，调用顺序和传输类型也可能不同。补充题四：pinned memory 是哪一侧的内存？答案：页锁定主机内存，不是 GPU 显存。
执行 python examples/core_training/optimizer_numerics.py qa/core_training/optimizer_numerics.json 可复查 CPU 随机数、结合律、补齐词表和等量梯度平均反例。GPU、MPS、编译和通信性能明确未运行。
参考与覆盖
- baseline、单 GPU、DDP 三脚本及 README：复用数据/模型/损失见第 3—6、13 章，新增计时、精度、编译、词表与同步路径在本章。
- CPU/MPS 随机采样：同种子设备分支、CPU 采样桥接与边界。