Activation recomputation（activation checkpointing）怎么做？
full 和 selective 有什么区别？代价多少？

英文问法How does activation checkpointing work? · You run out of GPU memory during training. What do you do? · Full vs selective recomputation? · How much extra compute does it cost?

考过字节（"oom怎么办，activation checkpointing原理"）· 牛客（"解决显存不够的方法有哪些"）

先弄清楚：每个概念是什么

Activation: forward 时存给 backward 用的中间结果（构成见 training memory）。例：GPT-3 一层，\(s = 2048\)、\(b = 1\)，不切分时约 2.9 GB。
Activation recomputation: 像打游戏只在每关开头存档（checkpoint），backward 到这一关时从存档重跑 forward。
Full recomputation: 每层只存输入（\(2sbh\) bytes），整层 forward 重跑，省得最多。
Selective recomputation: Megatron 只重算 attention 里的 \(QK^\top\)、softmax、softmax dropout、乘 \(V\)。这些 \(s\times s\) 的 tensor 占 GPU memory 大头（\(5as/h\) 项），FLOPs 却很少。
\(\sqrt n\) checkpointing: \(n\) 层切成 \(k\) 段，只存段边界，backward 时一次只重算一段。例：\(n = 96\)、\(k = 10\) 时只存约 \(10 + 96/10 \approx 20\) 层的量。

看图：存什么、多算什么

图 1：GPT-3 一层在 TP = 8 + SP 下的 activation，按 sbh/t 的倍数画。橙色是 attention 里 s×s 的部分，selective 只重算这一块。full 每层只留输入 2sbh。

图 2：3 层的时间线，按 FLOPs 画长度（forward 2、backward 4）。full recompute 时 backward 到哪一层就先重跑这层的 forward（虚线框），多出一次 forward，6 变 8，多 33%。

30 秒回答

Activations saved for backward often dominate training memory, so we keep only checkpoints, usually each layer's input, and rerun that layer's forward during backward, trading compute for memory. Full recomputation keeps only 2sbh bytes per layer but costs one extra forward pass, so training goes from about 6ND to 8ND FLOPs, roughly 30–40% slower. Megatron's selective recomputation redoes only the attention core, QKᵀ, softmax, dropout and attention over V, which holds most of the memory but few FLOPs: for GPT-3 it saves 70% of activation memory for about 2.7% more FLOPs.

关键公式

存几个 checkpoint

\(l\) 层、\(c\) 个 checkpoint，\(A_\text{in}\) = 每层输入，\(A_\text{int}\) = 每层中间结果：

\[ \text{memory} = c\,A_\text{in} + \frac{l}{c}\,A_\text{int},\qquad c^* = \sqrt{l\,A_\text{int}/A_\text{in}} \]

\(A_\text{in} = A_\text{int}\) 时就是 \(\sqrt n\) 方案，memory \(O(\sqrt n)\)。Transformer 里 \(A_\text{int}/A_\text{in} = (34 + 5as/h)/2\)，GPT-3 是 57，\(c^* \approx 74\)（共 96 层），所以实践中每 1–2 层存一个。

代价

\[ \text{full：} \underbrace{2}_{\text{fwd}} + \underbrace{4}_{\text{bwd}} + \underbrace{2}_{\text{重算}} = 8\Psi D\quad(+33\%) \] \[ \text{selective：} \frac{\text{hardware FLOPs}}{\text{model FLOPs}} \approx 1 + \frac{s}{6h} \]

MFU 不算重算的 FLOPs，HFU 算（见 FLOPs 与 MFU）。

手算例子：GPT-3 175B

\(s = 2048\)、\(h = 12288\)、\(a = 96\)、\(b = 1\)、TP \(t = 8\) + SP，96 层：

方案	每层	96 层	额外计算
全存：\(sbh(34 + 5as/h)/t\)	359 MB	34.4 GB	0
selective：\(34sbh/t\)	107 MB	10.3 GB	≈ 2.7% FLOPs
full：\(2sbh\)	50 MB	4.8 GB	+1 次 forward

\(5as/h = 80\) 比 34 还大，selective 一下去掉 \(80/114 \approx 70\%\)。实测（22B，一层）：full 慢 39%，selective 慢 7%，selective + SP 只慢 4%。

必答点

原理：只存 checkpoint，backward 时重跑 forward，用 compute 换 memory
full：每层只剩 \(2sbh\)，多一次 forward：\(6\Psi D \to 8\Psi D\)，实测慢 30%–40%
selective：只重算 attention 的 \(s^2\) 部分；GPT-3 省 70% activation，多约 2.7% FLOPs
\(k\) 段 → memory \(O(n/k + k)\)，\(k = \sqrt n\) 时 \(O(\sqrt n)\)

追问

PyTorch 里怎么用？

torch.utils.checkpoint.checkpoint(block, x, use_reentrant=False) 包住一个 Transformer block。默认保存 RNG 状态，重算时 dropout 的 mask 一致。

和 FlashAttention 什么关系？

FlashAttention 不存 attention matrix，backward 时分块重算，本身就是一种 selective checkpointing。用了它，selective recompute 要省的那部分基本已经省掉。

显存不够还有哪些办法？

减 micro-batch + gradient accumulation、mixed precision、ZeRO/FSDP、加 TP/PP/CP、FlashAttention、CPU offload。先用不掉速的，full recompute 和 offload 放后面。

为什么 full recompute 实测慢 39%，不是 33%？

33% 只算 FLOPs。论文的解释是，baseline 在 backward 里把 all-reduce 和 weight gradient 的计算 overlap 了，backward 变短，多出来的 forward 占比就超过 33%。

来源Chen et al. 2016 §4.3（p.5）、§5.4（p.10）· Narayanan et al. 2021 §3.5（p.6）· Korthikanti et al. 2022 §5（pp.8–10）、Table 2（p.9）、§6.2 Table 4（pp.10–11）、Appendix A（p.15）· Dao et al. 2022, FlashAttention §3.1（p.5）· PyTorch torch.utils.checkpoint

Activation recomputation（activation checkpointing）怎么做？full 和 selective 有什么区别？代价多少？