训练显存里最先爆的往往不是参数，而是每层都要留到反向的激活。论文把 n 层切成若干段，只保存段边界，反向时从边界重放前向——峰值从 O(n) 降到 O(n/k + k)，代价是每步多一次前向。
PyTorch 的 use_reentrant 又决定这次重放是整段执行还是算够即停。
反向要用的输入张量必须活到反向。矩阵乘 y = Wx 算 ∂L/∂W 要用 x，所以每层至少留一份激活；层数、batch 线性增长，seq 的 attention 项按 s² 增长（5bs²a），而参数和优化器状态不随 batch 变。
Full checkpointing 只保存每个段的入口，段内激活全部丢掉；反向从入口重算。段数 s、段长 k，峰值约为 (s + k)·A；k ≈ √L 时最小，约 2√L·A，代价是一次额外前向。
Selective checkpointing 只重算计算便宜的算子。Transformer 里典型做法是保留 MLP 激活、重算 attention 的 s² 矩阵（本页类比，推断）：显存省得比 full 少一点，算力几乎不涨。PyTorch 的 SAC（selective activation checkpointing）只支持 use_reentrant=False。
先算清楚反向到底要留什么，才知道 checkpoint 能省哪一块。
参数、梯度、优化器状态的大小只跟模型参数量有关；激活不一样。以最普通的矩阵乘为例：
所以每层的输入不能算完就释放。序列越长、batch 越大，这份“留给反向的张量”就越大。按 Transformer 的逐项估算，每层激活可以写成：
取一个参考配置：bs=1、seq=2048、h=4096、heads=32。代入 Table 1（本页推导）后每层约 956 MB（≈912 MiB），32 层约 30.6 GB（28.5 GiB）——还没算参数、梯度和优化器状态。长上下文训练里，激活经常先于参数把显存吃满。
如果只是把输出丢掉，反向就算不出梯度；如果把每层输入都留着，显存就是 O(L)。checkpoint 的做法是：只留少数几个入口，其他输入在反向时重新算出来。它不是减少计算，而是把“存”换成“再算一遍”。
论文给出的数量级是：1000 层 ResNet 从 48G 降到 7G，运行时间多 30%（Abstract、§5.4）。这两个数字是本页后面所有讨论的锚点，但它们有明确条件——ImageNet、batch 32、224×224、Titan X（§5.1、§5.2、§5.4）。
Full 把整段丢掉，selective 只丢重算便宜的那部分；两者的显存下限不一样。
每层输入都留到反向。显存 M = L·A，没有重算，训练最快，但 O(L) 随层数线性涨。
卡在：长上下文、大 batch 时直接 OOM；加模型并行或 offload 又会引入通信/PCIe 开销。
每 k 层为一段，只存段入口。反向从入口重放整段，峰值约 (s + k)·A，s = ceil(L/k)。
卡在：重算粒度粗。段内每个算子都被重放，哪怕其中很多算子的激活又小又便宜。
保留重算昂贵的算子（MLP），只丢重算便宜的算子（attention 的 s² 矩阵）。显存有下限，但算力开销小（本页对 Transformer 的类比，推断）。
卡在：要先知道哪些算子值得丢；策略错了会既没省多少显存、又增加重算。
粒度还可以再分一层：checkpoint(fn, ...) 由调用者手写边界；checkpoint_sequential 按层数均分成 segments 段；checkpoint_wrapper 以 module 为单位；apply_activation_checkpointing 用谓词选子模块；SAC 则细到 ATen 算子。粒度越细，越能贴近“只丢便宜货”，但实现和调参也越复杂。
论文的 §4.2 给的是同一原则：保留计算昂贵的算子，丢掉计算便宜的算子。原文的例子是 Conv-BatchNorm-Activation：保留 convolution 的结果，丢掉 batch norm、activation、pooling 的结果。把它套到 Transformer 上，就是保留 MLP、重算 attention——这是本页的类比，论文没有给 Transformer 的这张表（推断，依据 §4.2 的原则）。
论文 Eq.1 的全部内容；k ≈ √L 时最优，两端都退化到 L + 1 个单位。
论文把 n 层网络切成 k 段，每段只保留段输出，段内激活丢掉。用本页的记号重写（L 层、段长 k、段数 s = ceil(L/k)）：
把常数项也写出来：峰值约为 (s + k)·A，其中 A 是每层一份激活的字节数。对 s + k 求最小，s = L/k 时得到 k = √L，峰值 2√L·A。这是论文 Eq.1 的计法：段输出和段内激活分开算；精确到单个张量时，段输出会在两边各算一次，差 1 个单位，不影响 O(√L)。两端的退化也值得记住：
k = 1：每层一段，s = L，按论文计法峰值 L + 1 个单位；精确计法减掉重叠后就是 L。两种计法都说明：逐层 checkpoint 不比全存省。k = L：只有一段，s = 1，论文计法 1 + L；精确计法同样是 L。没有边界可省，只剩重算。k ≈ √L：段数和段长平衡，峰值最小。L=1000 时是 2√1000 ≈ 63 个单位，而不是 1000。计算侧：每个段在反向被重放一次，段长合计 L 层，所以额外计算约等于一次完整前向。论文说 backward 本身接近 forward 的两倍，因此总训练算力的增量约为 1/(1+2) = 33.3%；实测运行时间多 30%（§5.4）。这两个口径不同：33.3% 是按 §4.3 的时间比例折算（本页估算），30% 是论文的实测 runtime。
Table 1（本页推导，不是论文表格）：参考配置 b=1、seq=2048、h=4096、heads=32。激活公式按笔记《显存占用计算》§训练时；FLOPs 按每个矩阵乘 2×乘加数估算。Attention 的 s² 矩阵占激活 70.18%，但其 core FLOPs 只占前向 7.69%；投影 8sh² 另占 30.77%，MLP 占 61.54%——这正是 selective 只丢 s² 矩阵、只重算 core 的原因。
| 组成 | 公式 | 字节 / 层 | 占每层激活 | 占每层前向 FLOPs |
|---|---|---|---|---|
| Attention s² 矩阵 | 5bs²a | 671,088,640 | 70.18% | — |
| Attention 其余激活（QKV/输出等） | 11bsh | 92,274,688 | 9.65% | — |
| MLP + LayerNorm 激活 | 19bsh + 4bsh | 192,937,984 | 20.18% | — |
| 每层激活合计 | 34bsh + 5bs²a | 956,301,312 | 100% | — |
| Attention core FLOPs | 4s²h | — | — | 7.69% |
| Attention 投影 FLOPs | 8sh² | — | — | 30.77% |
| MLP FLOPs | 16sh² | — | — | 61.54% |
| 每层前向 FLOPs 合计 | 24sh² + 4s²h | — | — | 100% |
956,301,312 B ≈ 956 MB（≈912 MiB，1 MiB=1,048,576 B）；表里的 70.18% 和 7.69% 就是下面 selective 模式用的 m 和 f：m=5bs²a/(34bsh+5bs²a)，f=4s²h/(24sh²+4s²h)。分项百分比四舍五入，合计按未舍入值计；换 seq/h 这两个比例会变。本页 GB 一律按 1 GB=1000 MB 换算。
模型只算激活的峰值，不含参数、梯度、优化器状态和通信。A 是每层激活总量（MB），假设各层相同；k 是每段层数，s = ceil(L/k)。selective 只丢 attention 的 s² 矩阵，按 Table 1 的 m=5bs²a/(34bsh+5bs²a)=70.18%、f=4s²h/(24sh²+4s²h)=7.69% 计算。默认 A=956 MB 对应 Table 1 的 956,301,312 B（≈956.3 MB；912 MiB≈956 MB）；所有 GB 按 1 GB=1000 MB 换算。
L 个单位归一
k=L 时两种模式都没有显存收益：按本页段模型，full 是 L+1 个单位，selective 是 L+m≈L+0.70 个单位（m≈0.7018）。
Table 1 的参考 A 是 956 MB（912 MiB）；下面两组改用 A=912/512 MB 的整数示例，m、f 与 A 的取值无关。
组 1：L=32，A=912 MB，k=4。 s=ceil(32/4)=8。
M=32×912=29,184 MB=29.2 GB；无重算。M=(8+4)×912=10,944 MB=10.9 GB；1−12/32=62.5% 的激活显存被省掉；额外前向 32/32=100%；总算力 (32)/(32+64)=33.3%。Munits=32×(1−40/57)+(8+4)×(40/57)=544/57+480/57=1024/57=17.964912；M=17.964912×912=16,384.0 MB=16.4 GB；显存节省 1−17.964912/32=43.9%；额外前向 (1/13)×32=2.461538 层，即 7.7% 次前向；总算力 2.461538/(3×32)=2.6%。组 2：L=48，A=512 MB，k=6。 s=ceil(48/6)=8。
M=48×512=24,576 MB=24.6 GB；无重算。M=(8+6)×512=7,168 MB=7.2 GB；显存节省 1−14/48=70.8%；额外前向 100%；总算力 33.3%。Munits=48×(1−40/57)+(8+6)×(40/57)=816/57+560/57=1376/57=24.140351；M=24.140351×512=12,359.9 MB=12.4 GB；显存节省 1−24.140351/48=49.7%；额外前向 (1/13)×48=3.692308 层，即 7.7% 次前向；总算力 3.692308/(3×48)=2.6%。selective 的常数怎么来的：m=5·2048²·32/(34·2048·4096+5·2048²·32)=40/57=0.701754，f=4·2048²·4096/(24·2048·4096²+4·2048²·4096)=2/26=1/13=0.076923。这组数字对应 Table 1 的参考配置；换 seq/h 要重算。
这个模型简化了什么：它按论文 Eq.1 的段模型计账，假设各层激活相同、每个段都被完整重放；段输出和段内激活分开计数，精确到张量时会重叠 1 个单位（k=1 或 k=L 时正好等于全存）。没有计入参数/梯度/优化器状态、显存碎片、通信和 allocator 开销。PyTorch non-reentrant 的 early_stop 可能让重算提前停止，checkpoint_sequential 又会让最后一段不做 checkpoint——真实值会和这里略有出入。selective 的“只丢 s² 矩阵、只重算 attention core”是按 §4.2 原则做的类比，不是论文测量。
两者不是性能开关，而是 autograd 语义不同；选错会直接报错，或者得不到想要的梯度。
checkpoint() 时，省略 use_reentrant 仍会走 True在 pytorch@2ee7a63 里，use_reentrant=None 会先发 warning，然后赋值 True（torch/utils/checkpoint.py:650-660）。所以“推荐 False”不等于“默认 False”：要么显式传 False，要么用 checkpoint_wrapper（它默认 NO_REENTRANT，见 checkpoint_wrapper.py:120-135）。源码还说明未来会对省略参数直接报错。
Table 2（本页归纳）：checkpoint 的几种粒度，行号对应 pytorch@2ee7a63。
| API | 粒度 | 关键行为 |
|---|---|---|
checkpoint(fn, *args) | 一次调用包住的 region | 边界手写；推荐显式 use_reentrant=False（checkpoint.py:397-631） |
checkpoint_sequential(model, segments, input) | 按层数均分 segments 段 | use_reentrant=None 同样 warning 后取 True（:764-774）；除最后一段外都 checkpoint，最后一段正常前向（:716-805） |
checkpoint_wrapper(module) | 单个 module | 默认 NO_REENTRANT；包住 module 的 forward（checkpoint_wrapper.py:196-236） |
apply_activation_checkpointing(model, check_fn=...) | 谓词选中的子模块 | 递归替换选中的 module（checkpoint_wrapper.py:239-313） |
create_selective_checkpoint_contexts(policy) | ATen 算子 | 按 policy 返回 MUST_SAVE / PREFER_RECOMPUTE；仅 non-reentrant（checkpoint.py:1755-1794） |
Table 3（本页归纳）：use_reentrant=True/False 的语义差异。内容来自 torch/utils/checkpoint.py:480-520 的 note，不是凭经验补的。
| 行为 | use_reentrant=True | use_reentrant=False |
|---|---|---|
| 前向 autograd 图 | 不记录，forward 在 torch.no_grad() 下跑 | 记录 |
| 反向重算范围 | 整个 function 全量重算 | 需要的张量算完即停（early_stop=True 默认） |
| 反向 API | 只支持不带 inputs 的 torch.autograd.backward | 支持全部，包括 torch.autograd.grad |
| requires_grad 要求 | 至少一个输入/输出需要梯度 | 无此要求 |
| 嵌套结构 | 不把 list/dict 里的张量当作参与 autograd | 支持 |
| detach 张量 | 段内含 detach 会报错 | 支持 |
kwargs / context_fn | kwargs 不支持；context_fn 只支持 False | 支持 |
determinism_check / early_stop | determinism_check 在 True 下始终关闭；early_stop 在 True 下被忽略 | 均支持 |
non-reentrant 的重算钩子在“反向需要的张量都已重建”时抛出 _StopRecomputationError（checkpoint.py:1328-1331）。所以论文的“一次额外前向”是段级模型的上界；PyTorch 的实际重算量取决于这段反向到底要多少中间张量。reentrant 没有这个早停，early_stop 参数在 True 下被忽略（checkpoint.py:563-567）。
preserve_rng_state 默认 True：重算前恢复 forward 时的 CPU/device RNG 状态（checkpoint.py:1975-2018）。否则 dropout 的 mask 会变，重算出的激活和原来不一致，梯度就错了。源码还注明：在 torch.compile 下这个 flag 不生效，始终保留 RNG 状态（checkpoint.py:533-535）。
FSDP 在 HandleTrainingState.BACKWARD_PRE 状态下直接跳过 _pre_forward 的 unshard/注册逻辑（fsdp/_runtime_utils.py:373-380），_post_forward 也跳过 reshard/注册（:468-471）。反向重算发生在这个状态里，参数已经 unshard，跳过才能避免把同一套 forward 逻辑再走一遍。
DeepSpeed 的 CheckpointFunction 是从 torch.utils.checkpoint 改的，docstring 列了额外能力：RNG tracker、model parallel tracker、activation partitioning、contiguous memory、CPU checkpointing（deepspeed/runtime/activation_checkpointing/checkpointing.py:568-575）。配置键包括 partition_activations、number_checkpoints、contiguous_memory_optimization、cpu_checkpointing（config.py:27-43）；partition_activations 把激活按模型并行 rank 切开（checkpointing.py:1171-1175）。它在 torch.no_grad() 下跑前向（:626-627），不要把它和 PyTorch 的 non-reentrant 语义混为一谈。
还有一个实现细节：checkpoint_sequential 的 docstring 写明“除最后一段外都不保存中间激活”（checkpoint.py:720-723），代码里最后一段确实直接跑普通前向（:795-805）。最后一段不 checkpoint 不会增加峰值——反正反向第一步就要用它的激活——但会少一次重算。这也是为什么真实框架的“额外一次前向”通常略小于论文的段级上界。
12 层、每 4 层一段的例子：3 个检查点 + 4 层段内激活 = 7 个单位，而不是 12 个。
k=4 的段划分。前向只保留 c1–c3 三个段输出；反向从 c2 重放 L9–L12。峰值按论文 Eq.1 的段模型计为 s+k=3+4=7 个单位（段输出与段内激活各算一次；精确计法减 1 个重叠单位），不 checkpoint 是 12 个单位（少 41.7%）。图中省略了参数 unshard、通信和 allocator 细节。这张图也解释了为什么 k 不能随便取：段太长，段内激活 k 变大；段太短，检查点数量 s 变大。两者相加在 k≈√L 处最小。图里 L=12，√12≈3.46，所以 k=4 已经接近最优；k=3 或 k=4 都是 7 个单位。
论文的实测条件是 ImageNet、batch 32、224×224、Titan X；这些比例不能直接搬到别的模型。
Table 4：论文报告的结果。48G、7G、30% 来自 Abstract / §5.2 / §5.4；LSTM 的 4× 是相对“已做系统优化但不做重算”的方案，不是相对未优化的基线（§5.3）。
| 设置 | 结果 | 出处 |
|---|---|---|
| 1000 层 ResNet，ImageNet | 训练显存 48G → 7G | Abstract |
| 同设置，batch 32，输入 224×224 | 不到 7GB 即可训练 1000 层 | §5.2 |
| 运行时间（sublinear 方案整体） | 约多 30% | §5.4 |
| 4 层 LSTM，1024 hidden，unroll 64，batch 64 | 相对“优化但不重算”的方案，显存降到 1/4 以下（>4× reduction） | §5.3 |
| 只做系统优化、不重算 | 200 层 ResNet 已是当时 GPU 的上限 | §5.2 |
反直觉的地方在于“多一次前向”并没有让运行时间翻倍：论文的解释是 backward 本身接近 forward 的两倍，所以多一次 forward 只多约 1/3 的算力，实测是 30%。这也是为什么在显存受限的 GPU 上，重算常常比“省显存但增加 memory traffic”的方案更划算——但论文只报告了 ResNet/LSTM 的这组设置，不能把它当成所有模型的固定比例。
论文的另一半贡献是计算图层面的 in-place 与内存共享（§3）。Table 4 里的 48G→7G 是这些系统优化和 checkpoint 计划共同作用的结果；不要把它全部归给“每 k 层切一段”。
什么时候别用：激活不是瓶颈、前向不可重复、或者框架已经把 attention 重算掉了。
长上下文、大 batch、深层网络；参数/优化器状态已经通过 ZeRO、FSDP 或量化处理过，剩下的显存大头是激活；前向没有不可重放的副作用，或者 RNG 状态能被正确恢复。
短序列、小 batch、浅层模型：激活本来就不是大头，重算只会白花时间。训练已经 compute-bound、GPU 利用率打满时，再增加一次前向会直接拖慢吞吐。
checkpoint.py:466-471）。dropout 的 RNG 由 preserve_rng_state 处理，但自定义随机源、数据依赖的全局状态不在保护范围内。use_reentrant=True 不记录 autograd 图、只支持不带 inputs 的 torch.autograd.backward、要求至少一个输入/输出有梯度、不处理嵌套结构、段内不能有 detach 张量。需要 autograd.grad、kwargs、嵌套结构或 selective 策略时，必须用 False。create_selective_checkpoint_contexts 要给出每个算子的 MUST_SAVE / PREFER_RECOMPUTE 决策；策略错了可能既没省多少显存，又把重算成本加上去。而且 SAC 通过 context_fn 实现，context_fn 仅在 use_reentrant=False 下支持（checkpoint.py:550、674-675、1829）。fsdp/_runtime_utils.py:373-380、468-471）；DeepSpeed 的 partition_activations 依赖模型并行，cpu_checkpointing 则把激活挪到 CPU，换来 PCIe 传输。两者都不是“打开就变快”。checkpoint_sequential 不 checkpoint 最后一段（checkpoint.py:720-723、795-805）；non-reentrant 的 early_stop 可能提前停止重算（:1328-1331）。本页的 (s+k)·A 是论文 Eq.1 的段模型，不是每个框架的精确记账。offload_wrapper（checkpoint_wrapper.py:106-109、174-193），DeepSpeed 有 cpu_checkpointing（config.py:42-43）。它把激活放到 CPU，省 HBM 但吃 PCIe 带宽；和重算是正交的取舍，不是谁替代谁。当激活是主瓶颈、前向可重复、算力还有余量时，它是把 O(L) 压到 O(√L) 的最直接手段。先算清楚 (s+k)·A 和 1/(1+2) 这两个数，再决定 k 取多少、要不要上 selective——不要先打开开关再找原因。