附加讲 A1 · 内存优化：重计算与卸载 · Rematerialization & Offloading

对应 Maps to：原课没有单独排这一讲。原课第 1 讲的”课程地图”在”内存优化”一层列出了重计算（rematerialization）和零冗余（zero redundancy）；第 2 讲提到用 √N 的检查点换显存；第 12–13 讲的必读 / 选读包括 ZeRO-Infinity 和 ZeRO-Offload；第 21 讲的 MegaBlocks 也说显存限制了 micro-batch、可以用选择性重计算缓解。本讲把这些线索整理成一讲。 精读论文：T. Chen, B. Xu, C. Zhang, C. Guestrin, Training Deep Nets with Sublinear Memory Cost（arXiv 2016）；V. Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models（MLSys 2023）。补充：P. Jain et al., Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization（MLSys 2020）。 依据说明：本讲没有对应的 slides，内容依据上面三篇论文和 ZeRO-Offload（ATC 2021）、ZeRO-Infinity（SC 2021）的公开内容。模型状态的卸载（ZeRO-Offload 的最优划分、一步延迟更新、ZeRO-Infinity 的带宽账）第 13 讲 §7 已经讲过，本讲只在 §4.1 回顾，重点放在激活上。论文数字只做转述。

remat.py（激活公式、按 Megatron 结构写的 GPT 层与实测、链式检查点的模拟与穷举、层内张量图的前沿、激活卸载的时间线）和全部实验都是本讲义自己设计的；GPU 上的时间是 roofline 估算，激活字节是在 CPU 上实测后按 16 位训练折算的。 配套代码 Code：code/appendixA1/（4 个脚本 + 一个共享模块；需要 numpy 和 torch；每个都在 5 秒以内） 预计阅读时间：2.5 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标

学完这一讲，你应该能够：

1. 说出训练时为反向传播保存了哪些张量，用 sbh·(34 + 5as/h) 估算一个 Transformer 层的激活，并在 PyTorch 里实测验证。
2. 推导 Chen et al. 的 √n 检查点（显存 O(√n)、多一遍前向）和递归方案（显存 O(log n)），并解释为什么均匀划分不是最优的。
3. 解释 Megatron 的选择性重计算为什么只重算注意力核心，以及”存哪些”的最优解怎样依赖代价模型（FLOPs vs 实际时间）。
4. 判断一个张量该留着、重算，还是卸载到 CPU，估算卸载的拷贝能不能藏在计算下面。
5. 给定模型、并行方式和显存预算，规划一个”先不增加计算、再便宜地重算、最后才完全重计算或卸载”的方案。

0.2 先修与本讲范围

0.3 本讲地图

§1  训练显存的另一半：激活                         ← 代码 01 演示 1–2
§2  链式网络上的检查点（Chen et al. 2016）          ← 代码 02
§3  一个 Transformer 层里存什么：选择性重计算        ← 代码 01 演示 3–4、03
§4  卸载：模型状态与激活                           ← 代码 04
§5  把它们放在一起：显存规划                        ← 代码 03 演示 4、04 演示 4
§6  论文精读：Chen et al. 2016；Korthikanti et al. 2022；（补充）Checkmate
§7 小结  §8 常见误区  §9 练习  §10 延伸阅读  §11 下一讲预告

0.4 记号

1. 训练显存的另一半：激活 · Activations

1.1 为什么要存激活

反向传播算梯度时，每个算子都要用到前向时的某些值：线性层 Y = XW 求 ∂L/∂W = Xᵀ·∂L/∂Y，要用输入 X；GeLU 要用它的输入；softmax 要用它的输出；dropout 要用掩码。前向时这些值必须存下来，等反向用完再释放。这就是激活显存——和模型状态（权重、梯度、优化器状态，16 字节 / 参数）不同，它和 batch、序列长度成正比。

1.2 一个 Transformer 层存多少（Korthikanti §4.1；代码 01 演示 2）

按 16 位激活、1 字节的 dropout 掩码逐个数：

我们按同样的结构写了一个 PyTorch 层，用 saved_tensors_hooks 统计它为反向保存的张量（同一块存储只算一次、不算参数，浮点按 2 字节、布尔按 1 字节）：

多出来的 0.03 是 LayerNorm 存的均值和标准差，论文把它们忽略了。

第一版实测比公式多 6·sbh：PyTorch 自带的 F.dropout 在 CPU 上把掩码存成了浮点数，masked_fill 也会多存一个布尔掩码。改成自己写的布尔掩码 dropout 和加法的因果偏置之后，实测就和公式对上了——激活显存取决于实现，公式只是一个记账约定。

1.3 谁更大：模型状态还是激活（代码 01 演示 1）

GPT-3 175B 按 Megatron 的切分（张量并行 8、流水线并行 8、每块 GPU 一个微批）：

• 每块 GPU 的模型状态 = 16 × 174B / 64 = 43.5 GB；
• 流水线的第一段要存 p 个微批的激活，相当于存全部 96 层。只用张量并行时激活要 55.6 GB，比模型状态还大，一块 80 GB 的卡放不下。

所以对这一级的模型，处理激活和处理模型状态一样重要。

2. 链式网络上的检查点 · Checkpointing on a chain

2.1 Chen et al. 的算法 1（代码 02 演示 1）

把 n 层的链分成若干段，只存每段的输入（检查点），段内的输出算完就丢。反向时从最后一段开始，每段先从段首的检查点重新前向一遍、把段内所有输出存下，再逐层反向。

显存 = 检查点的个数 + 一段之内的输出。均匀地分成 k 段时是 k + n/k，在 k = √n 时最小，约 2√n；代价是多一遍前向。我们逐步模拟（每层输出大小 1）：

反向本来就是前向的 2 倍，所以多一遍前向约多 1/3 的时间——Chen et al. 在 Titan X 上实测约 30%。

2.2 递归：O(log n)

把每一段看成一个”大算子”，在段内递归地用同样的办法：g(n) = k + g(n/(k+1))，解出 g(n) = k·log_{k+1} n。k = 1 时显存 log₂ n，代价是前向次数 O(n log n)。论文的评价是”可以用但不常用”。从 O(n) 到 O(√n) 到 O(log n)，是一条连续的”显存换计算”曲线。

Chen et al. 还有两个常被忽略的贡献：一是原地运算和内存共享（按生命周期复用缓冲区，像编译器的寄存器分配）——光这一项就能省 2–3 倍；二是把”哪些结果可以丢”做成计算图上的一个属性（mirror count），由框架自动生成重计算的图，而不是让用户手写训练循环。

2.3 均匀划分不是最优的（代码 02 演示 2）

在 16 层的链上穷举所有 2¹⁵ 种检查点组合，算出精确的帕累托前沿：

• 均匀的链：√n（每 4 层一个，显存 9、重算 12）不在前沿上——同样的显存，最优只要重算 9。最优方案的段长是前长后短的（显存最小的点：段长 [1, 5, 4, 3, 2, 1]）。原因是反向从最后一段开始，处理后面的段时前面所有检查点都还在显存里，所以后面的段要短；轮到前面的段时，后面的检查点已经释放了。
• 不均匀的链（输出大小交替 1 / 4、重算代价交替 4 / 1）：√n 显存 27、重算 36；同样的显存下最优只要重算 5，同样的重算下最优只要显存 15。最优方案把检查点都放在小的输出上，重算的都是便宜的层。
• Chen 的算法 3（按预算贪心地切）只看累计的大小，在不均匀的链上可能正好切在大的输出上（B = 2 时显存 38，比 √n 还差）。

一般的计算图（有残差、有分支、每个算子的大小和代价都不同）上，“存哪些、什么时候重算”是一个组合优化问题——这就是 Checkmate（§6.3）。

2.4 PyTorch 里的检查点（代码 02 演示 3）

torch.utils.checkpoint.checkpoint(fn, x)：前向时在 no_grad 下运行 fn、只保存输入；反向时重新运行 fn 再求梯度。默认会保存并恢复随机数状态，所以有 dropout 时重算出来的掩码也一样。checkpoint_sequential(model, segments, x) 把一个 nn.Sequential 分成几段（最后一段不做检查点）。

32 层的 MLP（宽 1024、批 512）上：

（时间是多次运行的浮动范围。）最少的是分 8 段而不是 √32 ≈ 6 段：每层要存两个张量（Linear 的输入和 GELU 的输入），而每段只存一个输入，最优段数约 √(2n)。三种设置的梯度和全存时逐位相同。

3. 一个 Transformer 层里存什么：选择性重计算 · Selective recomputation

3.1 张量并行、序列并行与激活（Korthikanti 表 2）

张量并行把注意力和前馈的矩阵切到 t 块卡上，但 LayerNorm、dropout、残差这些部分在每块卡上是重复的，于是每卡每层是 sbh·(10 + 24/t + 5as/(ht))——那个 10 不随 t 缩小。序列并行把这些部分沿序列维度切开，配合把 AllReduce 拆成 reduce-scatter + all-gather（通信量不变），变成 sbh·(34/t + 5as/(ht))。

张量 + 序列 + 选择性比只用张量并行少 5.4 倍（论文图 7：约 5 倍）。GPT-3 的整个第一段：激活 55.6 GB → 10.3 GB，加上模型状态 53.8 GB，放得下，而且不用完全重计算。

3.2 为什么只重算注意力核心（代码 01 演示 3）

5as/h 这一项来自注意力矩阵：softmax 的输出、dropout 掩码和 dropout 的输出。GPT-3（a = 96、s = 2048、h = 12288）上 5as/h = 80，MT-NLG 上 64，都大于其余部分的 34。而重算它们只要 QKᵀ 和 softmax——大而便宜。选择性重计算只把这一段包进检查点：

• 省下的激活：GPT-3 70%、MT-NLG 65%；
• 额外的 FLOPs：按论文附录 A 的式 (8)–(9) 是 2.7% / 1.6%。但如果只把”重算 QKᵀ 和 ×V 的前向”（每层 4Bs²h）加到 model FLOPs 上，是 0.9% / 0.5%——论文的式 (8) 相当于每层多加了 12Bs²h（口径不同或笔误）。两种算法都说明开销很小；
• 论文表 4（22B 模型的一层，A100）：选择性重计算让一层慢 7%，完全重计算 39%；再加上序列并行只慢 4%。

⏱ 我们在 CPU 上测一个”像 GPT 比例”（s/h = 1/8）的层：多次运行时选择性 +8% 到 +21%，完全重计算 +37% 到 +49%，梯度和不重算时逐位相同。长序列（s/h = 1）时选择性的开销升到 +33% 以上：softmax、exp、dropout 的元素个数是 a·s²·b，在单线程 CPU 上很贵。

3.3 存哪些：穷举最优与代价模型（代码 03）

把一个 GPT 层要存的 13 个张量写成一张小图（每个张量有大小、重算的 FLOPs、要搬的字节、依赖哪些张量），穷举 2¹² 种”存 / 不存”的组合。没存的张量要从存着的张量出发重算，每个最多重算一次。用两种代价模型：

（重算代价占一层前向 + 反向的比例。）

• 按时间算，Megatron 的选择正好在前沿上：34·sbh 时最少的重算是 5.20%，就是它的方案。按 FLOPs 算只有 0.46%，比论文实测的 7% 小得多——softmax 和 dropout 是访存受限的，按 FLOPs 算严重低估了它们的代价。
• 前沿在 34 以下还很平：再重算两个 LayerNorm、GeLU 和 dropout 掩码，降到 20·sbh 只要 5.89%；再往下就要重算矩阵乘法（q,k,v、gelu_in），代价跳到完全重计算的 24.5%。
• 我们的”完全重计算”少于一遍前向（约 25% 而不是 33%）：反向用不到层的最终输出，按需重算就省掉了最后一个矩阵乘法；PyTorch 的 checkpoint 会把整层重跑一遍。

最优解取决于代价模型。这是 Checkmate 的核心主张：用在目标硬件上实测的代价，而不是假设每个算子代价相同。

3.4 长序列与 FlashAttention（代码 03 演示 3）

s = 32K 时不重算要 70 GB / 层——长上下文训练离不开选择性重计算。但它的时间开销也随 s 增长：注意力本身是 O(s²) 的。

第 5 讲的 FlashAttention 从根上去掉了这一项：注意力矩阵从不写进显存（只多存每行的 softmax 统计量），反向时在内核里分块重算。它就是”融合进内核的选择性重计算”，显存和选择性一样，而且重算时不用来回读写 a·s²·b 个元素。今天的训练栈里两者通常二选一。

4. 卸载 · Offloading

4.1 回顾：模型状态的卸载（第 13 讲）

• ZeRO-Offload：把训练画成数据流图，前向反向（O(M·B) 的计算）必须留在 GPU；在通信最少（4M 字节）的划分里，把 fp16 梯度和”更新超级节点”（fp32 参数、动量、方差和 Adam 的计算）都放到 CPU，GPU 上只剩 fp16 参数，显存省 8 倍。一步延迟更新（DPU）把 CPU 的更新和下一步的 GPU 计算重叠，代价是一步陈旧。
• ZeRO-Infinity：效率 = ait·bw / (ait·bw + peak)。参数和梯度的算术强度是 seq·bsz，优化器状态 seq·bsz/4，激活检查点 24·hd·ci——激活检查点最便宜，卸载到 CPU 几乎没有代价。

4.2 CPU 上的 Adam 有多快（代码 04 演示 1）

⏱ 本机实测 PyTorch 的 Adam（2000 万个 fp32 参数）：每十亿参数单线程约 0.72 秒、12 线程约 0.35 秒（读写约 16 字节 / 参数，相当于 22 / 46 GB/s 的内存带宽）。ZeRO-Offload 论文表 4 里，它的 CPU-Adam（AVX-512、多线程、分块拷贝）每十亿参数 0.22–0.26 秒，比 PyTorch 的 CPU 实现快 5–6 倍。Adam 是访存受限的，CPU 的内存带宽决定了它的速度。10B 模型一次更新约 2.6 秒，所以每步的 GPU 计算要足够长，或者用 DPU 把它藏起来。

4.3 激活：留着、重算，还是卸载（代码 04 演示 2）

GPT-3 一层每卡（张量并行 8，前向约 1.18 ms）：

（链路带宽取每方向 25 / 50 / 450 GB/s 的量级；“拷出 + 拷回”是不和计算重叠时的时间。）

• 注意力矩阵：重算比经 PCIe 拷贝便宜几十倍——大而便宜的张量应该重算，不该卸载。
• 矩阵乘法的输出：按不重叠的时间算，连它们也是重算更快；但”重算 / 拷贝”的比值是 31%，而注意力矩阵只有 2%，要卸载的话应该先卸载这些。
• 卸载的价值在于和计算重叠：藏住了就不花时间，藏不住就比重算还慢。

4.4 卸载的时间线（代码 04 演示 3）

模拟两条流（计算流、拷贝流）：每层前向后把要存的激活拷到 CPU，反向前提前拷回来，GPU 上只留 2 个缓冲；缓冲满了计算就得等。“计算等拷贝”的时间占纯计算的比例：

• 序列越长、模型越宽，每字节激活对应的计算越多，卸载越容易藏住——这就是 ZeRO-Infinity 的 AIT = 24·hd·ci。
• 只卸载检查点（每层 2·sbh）几乎总能藏住：GPU 上只留两层的缓冲。
• GH200 的 NVLink-C2C 让几乎所有情况都藏得住：CPU 内存变成了 GPU 显存的”二级缓存”。硬件的链路带宽直接改变了”重算还是卸载”的答案。

5. 把它们放在一起：显存规划 · Putting it together

GPT-3（张量并行 8、流水线并行 8），每卡模型状态 43.5 GB，80 GB 的卡留 10 GB 给临时缓冲：

（时间开销是 H100 roofline 估算。）

规划的顺序：

1. 先用不增加计算的办法：序列并行、ZeRO（把模型状态切开）、更大的并行度；
2. 再用最便宜的重计算：只重算大而便宜的张量（注意力核心，或者直接用 FlashAttention），再考虑 LayerNorm、GeLU 这类便宜的；
3. 链路够快时，把小而贵的张量或检查点卸载到 CPU，前提是拷贝能藏在计算下面；
4. 实在不够才完全重计算（约多 1/4 到 1/3 的时间）。

s = 8192 时，PCIe 5.0 或 NVLink-C2C 上的”选择性 + 卸载”比完全重计算快；PCIe 4.0 上反而更慢。没有放之四海皆准的答案，要看模型形状和硬件。

6. 论文精读 · Paper reading

6.1 Training Deep Nets with Sublinear Memory Cost（Chen, Xu, Zhang, Guestrin, arXiv 2016）

1. 问题：训练 n 层的网络，中间结果（特征图）和它们的梯度要 O(n) 的显存。很多当时最好的模型已经碰到 GPU 显存的上限，更深的网络没法试。能不能用亚线性的显存训练？

2. 为什么重要又困难：

• 显存决定了能试多深的网络、多长的序列（RNN 展开多少步）、多大的 batch；
• 参数本身相对小，大头是中间结果；
• 框架当时要么在运行时用引用计数和垃圾回收，要么静态分配，都没有讨论”计算换显存”。

3. 以前的工作为什么不行：

• 自动微分里的检查点（Griewank）早就有，但没有用到一般的深度网络上；
• 框架只在构造好梯度图之后做优化，不考虑丢掉一些结果再重算；
• 在 CPU / GPU 之间换入换出（swap）要占 PCIe 带宽，和数据并行抢带宽。

4. 新意：

• 计算图上的原地运算和内存共享（按生命周期的线性时间启发式，类似寄存器分配）；
• 用计算换显存：只存每段的输入，反向时重算段内的结果（算法 1）；用 mirror count 把这个模式推广到一般的计算图上，由框架自动生成重计算的梯度图（算法 2）；
• k = √n 段时显存 O(√n)、只多一遍前向；按预算贪心地切（算法 3），在预算上搜索；
• 递归地应用得到 O(log n) 的显存，前向 O(n log n)；
• 给框架的建议：允许丢掉便宜算子（BN、ReLU）的结果，提供显存规划算法，让用户设置 mirror 属性。

5. 结果：

• 1000 层的 ResNet 在 ImageNet 上，显存从 48 GB 降到 7 GB；
• LSTM 长序列上比只做系统优化再省 4 倍以上；
• 速度：亚线性方案多约 30% 的运行时间。

本讲的验证：逐步模拟了算法 1 和递归方案（n = 1024：65 和 13 个单位的显存，1.97n 和 5.5n 次前向）；在 PyTorch 里实测了分段检查点的显存和时间，梯度逐位相同。也看到它的局限：均匀的 √n 划分不是最优的（最优的段长前长后短），贪心的算法 3 在大小不均匀的链上可能比 √n 还差。

局限：

• 主要针对链状结构，一般图上的切分要靠启发式；
• 假设每层代价相近，不考虑硬件上的实际代价；
• 只处理特征图，不处理参数和优化器状态（后来由 ZeRO 解决）。

6.2 Reducing Activation Recomputation in Large Transformer Models（Korthikanti et al., MLSys 2023）

1. 问题：训练千亿参数的 Transformer 时，即使用了张量并行和流水线并行，激活仍然放不下，大家只好对每一层做完全重计算——这要多 30%–40% 的时间。能不能少重算一些？

2. 为什么重要又困难：

• 这些模型训练一次要几个月、几千块 GPU，30% 的时间就是几百万美元；
• 完全重计算是”全有或全无”的：只检查点一部分层，粒度太粗（MT-NLG 每块卡只有 3 层）；
• 张量并行下有一部分激活（LayerNorm、dropout）在每块卡上重复，张量并行越大越浪费。

3. 以前的工作为什么不行：

• 完全的激活检查点：省显存但多一遍前向；
• 只检查点部分层：粒度粗，不同的激活重算代价也不同；
• 把激活切到张量并行的各卡上再 all-gather（Megatron 早期的做法）：多一次通信。

4. 新意：

• 逐项算出每层的激活：sbh·(34 + 5as/h)，以及各种并行方式下的公式（表 2）；
• 序列并行：把 LayerNorm、dropout 这些在张量并行下重复的部分沿序列维度切开，AllReduce 换成 reduce-scatter + all-gather，通信量不变；
• 选择性重计算：只重算注意力核心（QKᵀ、softmax、dropout、×V）——它占激活的大头（5as/h），计算却很少；
• 两者结合：激活约少 5 倍；流水线并行时第一段的显存分析（要存 p 个微批）。

5. 结果：

• GPT-3 175B：5as/h = 80，选择性重计算省 70% 的激活；MT-NLG 530B：省 65%；FLOPs 开销 2.7% / 1.6%（论文口径）；
• 22B 模型的一层：完全重计算 +39%，选择性 +7%，选择性 + 序列并行 +4%；530B 和 1T 的模型上开销只有 2%，完全重计算是 36%；
• 序列并行本身让一层快 6%（LayerNorm、dropout 只算 1/t 的数据）；
• 整体上消除了完全重计算带来的 90% 以上的时间开销：530B 的模型在 2240 块 A100 上 MFU 54.2%，用完全重计算时是 42.1%（快 29%）。

本讲的验证：在 PyTorch 里按同样的结构实测了每层保存的激活（54.031 / 34.031 / 2.000 个 sbh），和公式只差 LayerNorm 的统计量；复现了 5as/h、70% / 65%、“约 5 倍”；穷举 2¹² 种组合后发现，按 H100 的 roofline 时间算，选择性重计算正好在帕累托前沿上。和论文不同的地方：FLOPs 开销的式 (8) 看起来多加了 12Bs²h；按 FLOPs 估计会严重低估访存受限的 softmax 的代价。

局限：

• 只针对标准的 Transformer 层；
• 选择性重计算的开销随序列长度增长，长序列时不再便宜——FlashAttention 后来在内核层面解决了这个问题；
• 序列并行和张量并行绑定在一起（同一组卡）。

6.3 （补充）Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization（Jain et al., MLSys 2020）

• 问题：一般的计算图（残差、分支、代价各不相同）上，最优的重计算计划是什么？Chen 和 Griewank 的方法只适用于链状、代价相同的图。
• 方法：把问题写成混合整数线性规划。变量 R_{t,i}（第 t 阶段是否计算算子 i）、S_{t,i}（是否把 i 的结果从 t − 1 阶段保留到 t）、FREE（何时释放）、U（每一步之后的显存）；目标是最小化总计算代价，约束是依赖关系和显存上限。用”前沿推进的阶段”（第 t 阶段算出第 t 个算子）减少变量；代价来自在目标硬件上的实测。也给了 LP 松弛 + 确定性取整的近似算法。
• 结果：现成的求解器在一小时内解出最优计划，能训练最多 5.1 倍大的输入。
• 和本讲的关系：我们在一个 Transformer 层的 12 个张量上做的穷举，是 Checkmate 思路的一个极小版本（只决定”存哪些”，不涉及跨层的调度和峰值显存）。它也说明了 Checkmate 坚持”按硬件实测代价”的原因：同一张图，按 FLOPs 和按时间得到的前沿不一样。

6.4 论文之间的联系

• 同一个问题的三种粒度：Chen 在”层”的粒度上切（存每段的输入），Korthikanti 在”层内张量”的粒度上挑（只重算注意力核心），FlashAttention 在”内核”的粒度上做（注意力矩阵从不离开片上存储）。粒度越细，越能挑出大而便宜的部分，开销越小。
• 代价模型决定答案：Chen 假设每层代价相同，所以均匀划分；Korthikanti 用”大而便宜”的直觉挑出注意力矩阵；Checkmate 把代价交给实测、把选择交给求解器。我们的穷举说明这三者是一致的：按实际时间算，“大而便宜”的就是注意力矩阵。
• 重算和卸载是同一个账本的两边：一个张量的”代价”可以是重算的时间，也可以是拷贝的时间（如果藏不住）；ZeRO-Infinity 卸载检查点、Korthikanti 重算注意力，在我们的模型里都是各自链路和形状下的最优选择。
• 一个可以写成评审的新想法：把 Checkmate 的规划推广到”存 / 重算 / 卸载”三选一，代价按目标硬件（PCIe 还是 NVLink-C2C）建模，并把拷贝和计算的重叠写进约束——GH200 这样的硬件上，答案可能和 A100 时代完全不同。

7. 小结 · Summary

1. 训练显存的另一半是激活：一个 Transformer 层存 sbh·(34 + 5as/h) 字节，大模型、长序列时可以比模型状态还大。具体存多少取决于实现（dropout 掩码存成什么类型）。
2. 检查点：只存一部分，反向时重算。链上 √n 段时显存 O(√n)、多一遍前向；递归时 O(log n)。均匀划分不是最优的：后面的段应该更短，检查点应该放在小的输出上。
3. 选择性重计算：Transformer 里最大最便宜的是注意力矩阵（5as/h 项），只重算它就省 65%–70% 的激活、只多几个百分点的时间；配合序列并行，激活少约 5 倍。FlashAttention 把这件事做进了内核。
4. “存哪些”的最优解依赖代价模型：按 FLOPs 算会低估访存受限的算子；按目标硬件的实际时间算，Megatron 的选择正好在前沿上。
5. 卸载：模型状态放到 CPU（ZeRO-Offload / Infinity），CPU 的 Adam 受内存带宽限制；激活能不能卸载取决于拷贝能否藏在计算下面——序列越长、链路越快越容易。
6. 规划的顺序：先不增加计算（序列并行、ZeRO），再便宜地重算（注意力核心），链路够快时卸载，最后才完全重计算。

8. 常见误区 · Common misconceptions

1. “激活检查点就是多一遍前向，开销固定 33%。” 完全重计算大约多一遍前向；选择性重计算只多几个百分点；按需重算（不重算反向用不到的东西）还能少一点。开销取决于重算什么。
2. “√n 检查点是最优的。” 只在”每层大小、代价都一样”且只允许重算一次的模型里近似最优；即使这样，均匀划分也不是精确最优（§2.3）。
3. “重算代价看 FLOPs 就够了。” softmax、dropout、LayerNorm 是访存受限的，按 FLOPs 算会把它们的重算代价低估一个数量级（§3.3）。
4. “PyTorch 测出来的激活就等于论文公式。” 取决于实现：CPU 上的 F.dropout 把掩码存成浮点数，masked_fill 会多存一个掩码；融合内核（FlashAttention）根本不存注意力矩阵。
5. “卸载总比重算好，反正 CPU 内存大。” 不重叠时拷贝往往比重算还慢；能不能藏住取决于每字节激活对应多少计算和链路带宽（PCIe 4.0 上 s = 2048 时，卸载选择性的激活要多 131% 的时间）。
6. “有了 FlashAttention 就不需要重计算了。” FlashAttention 去掉的只是 5as/h 那一项，剩下的 34·sbh / 层在大模型上仍然很大，长序列、深模型时还是要序列并行、检查点或卸载。
7. “重计算会改变训练结果。” 只要保存并恢复随机数状态，重算出来的值（包括 dropout 掩码）和原来逐位相同，梯度也逐位相同（代码 01、02 都验证了）。

9. 练习 · Exercises（含参考答案）

练习 1（激活公式）：一个模型 h = 4096、a = 32、s = 4096、b = 1，不并行。每层激活多少？5as/h 占多少？选择性重计算后剩多少？

答案：sbh = 4096 × 1 × 4096 = 16.8M；5as/h = 5 × 32 × 4096 / 4096 = 160；每层 sbh·(34 + 160) = 16.8M × 194 ≈ 3.25 GB，注意力矩阵占 160/194 = 82%；选择性后 34·sbh ≈ 0.57 GB。

练习 2（√n）：一条 400 层的链，每层输出 10 MB。全存要多少？均匀地每 20 层一个检查点呢？大约多几遍前向？

答案：全存约 400 × 10 MB = 4 GB。每 20 层一个：20 个检查点（200 MB）+ 一段 20 层（200 MB）≈ 400 MB，约少 10 倍；多约一遍前向，总时间约多 1/3。

练习 3（递归）：用 g(n) = k + g(n/(k+1)) 估算 n = 1024、k = 3 的显存。和代码 02 演示 1 的结果（18）比，差在哪里？

答案：g(n) ≈ k·log_{k+1} n = 3 × log₄ 1024 = 3 × 5 = 15。模拟的 18 多了递归最底层存下的几段输出和梯度缓冲，以及每一层递归时上一层的检查点还在显存里。

练习 4（选择性的开销）：按”只加 4Bs²h”的口径，GPT-3（s = 2048、h = 12288）选择性重计算的 FLOPs 开销是多少？如果 s 变成 32768 呢？

答案：开销 = 4Bs²h / (72Bsh²(1 + s/6h)) = s / (18h(1 + s/6h))。s = 2048：2048 / (221184 × 1.028) ≈ 0.9%；s = 32768：32768 / (221184 × 1.444) ≈ 10.3%。长序列时注意力的计算比例变大，选择性重计算不再便宜。

练习 5（代价模型）：一个张量 100 MB，重算它要 40 GFLOPs 并读写 300 MB。在 H100（989 TFLOP/s、3.35 TB/s）上按 FLOPs 和按 roofline 时间各要多久？哪个更接近真实？

答案：按 FLOPs：40e9 / 989e12 ≈ 40 µs；按访存：300e6 / 3.35e12 ≈ 90 µs；roofline 取较大的 90 µs。这个算子是访存受限的，按 FLOPs 会低估一半以上——正是 softmax 一类算子的情况。

练习 6（卸载能否藏住）：一层每卡前向 5 ms、反向 10 ms，每层要卸载 150 MB 的激活。PCIe 每方向 25 GB/s，能藏住吗？最少要多大的带宽？

答案：拷出 150 MB 要 6 ms > 前向 5 ms，前向时藏不住（反向 10 ms 时拷回 6 ms 可以藏住）。前向要藏住至少需要 150 MB / 5 ms = 30 GB/s（PCIe 5.0 可以）。

练习 7（CPU Adam）：一个 13B 的模型在一块 GPU 上用 ZeRO-Offload，CPU-Adam 每十亿参数 0.25 秒。每步的 GPU 计算要多长，才不至于（不用 DPU 时）让一半时间花在 CPU 上？每步要多少 token（按 8 FLOPs / 参数 / token、GPU 有效 300 TFLOP/s）？如果模型换成 3B 呢？

答案：CPU 更新 13 × 0.25 = 3.25 秒，GPU 计算也要约 3.25 秒；token 数 = 3.25 × 300e12 / (8 × 13e9) ≈ 9,400 个 token。换成 3B：CPU 更新 0.75 秒，token 数 = 0.75 × 300e12 / (8 × 3e9) ≈ 9,400——和模型大小无关：token 数 = (CPU 每参数的秒数) × GPU 算力 / 8。GPU 越快，每步需要的 token 越多，要么梯度累积，要么用 DPU 把更新藏起来。

练习 8（规划）：你要在 80 GB 的卡上训练一个模型，每卡模型状态 40 GB；不重算时激活 60 GB，选择性重计算后 25 GB，完全重计算 8 GB；选择性开销 6%，完全重计算 30%。你的卡和 CPU 之间是 PCIe 4.0。你会怎么选？如果换成 GH200 呢？

答案：留 10 GB 缓冲，激活预算 30 GB。选择性（25 GB）放得下，开销 6%，是 PCIe 4.0 上的首选；不需要卸载（藏不住的话反而更慢）。GH200 上可以考虑”选择性 + 卸载一部分小而贵的张量”，把显存腾出来换更大的微批（提高 GPU 利用率），只要拷贝能藏在计算下面。

10. 延伸阅读 · Further reading

• A. Griewank, A. Walther, Algorithm 799: revolve（ACM TOMS 2000）——自动微分里的二项式检查点，多级重计算的最优调度。
• A. Gruslys et al., Memory-Efficient Backpropagation Through Time（NeurIPS 2016）——RNN 上的动态规划检查点。
• M. Kirisame et al., Dynamic Tensor Rematerialization（ICLR 2021）——运行时贪心地丢掉和重算张量，不需要提前规划。
• J. Ren et al., ZeRO-Offload（ATC 2021）；S. Rajbhandari et al., ZeRO-Infinity（SC 2021）——见第 13 讲。
• T. Dao et al., FlashAttention（NeurIPS 2022）——内核层面的”选择性重计算”，见第 5 讲。
• A. Gomez et al., The Reversible Residual Network（NeurIPS 2017）——用可逆结构从输出算回输入，不存激活。
• 本仓库 CMU11-868 第 18 讲（ZeRO 与卸载的实现）。

11. 下一讲预告 · Next lecture

附加讲 A2 · 模型压缩：量化、稀疏与低秩——这一讲是”少存一点”，A2 是”少算一点、少读一点”：

• 训练后量化：逐通道、分组的 RTN，GPTQ 用二阶信息补偿误差，AWQ / SmoothQuant 按激活缩放；
• 2:4 结构化稀疏与硬件支持；低秩分解；
• 一条贯穿全课的教训：FLOPs 少了不等于时延低了——压缩要和内核、硬件配合才能真的变快。