Activation Checkpointing：用一次前向换激活显存

训练基础 的四份显存里，激活是唯一随 batch 和序列长度增长的一份，也是大 batch、长序列训练的第一约束。Activation checkpointing（又叫 gradient checkpointing、重计算）是省这份显存的标准手段：不存中间激活，反向需要时就地重算。

1. 问题：激活显存为什么先爆

静态三份（Θ、∇W、optimizer states）是常数，激活是变量：

$$M_{\text{act}} \propto B \times S \times H \times L$$

batch size 翻倍、序列翻倍、层数翻倍，激活都跟着线性翻倍。7B 模型 bf16 混合精度静态约 112 GB 尚可分片，但 $B{=}32$、$S{=}8\text{k}$ 时仅激活就能到数百 GB——不处理它，batch 根本开不大。

2. 机制：存边界，重算段内

把一个连续算子段（通常一个 Transformer block）声明为 checkpoint 段：

• 前向：段内激活算出即弃，只把段的输入（边界）留档；
• 反向：逆序到达该段时，从边界激活就地重跑一遍该段前向，补出段内全部中间激活，立刻交给该段的反向算子消费，用完即弃；
• 同一时刻显存里只有「所有段的边界 + 当前正在重算的那一段」。

语义上完全等价：重算的输入是同一批边界激活，算子是同一批确定性算子，产出的梯度逐位一致（不含 dropout 时；含 dropout 要靠 RNG state 重放保证 mask 一致，框架已处理）。

3. 两本账

显存账：默认全存是 $M \approx c \cdot B \cdot S \cdot H \cdot L$（$c$ 是每层保存量的系数，标准 Transformer 约 34 bytes/元素量级）。full checkpointing 后只剩边界：

$$M_{\text{ckpt}} \approx \underbrace{2 \cdot B \cdot S \cdot H \cdot n_{\text{seg}}}_{\text{边界}} + \underbrace{c' \cdot B \cdot S \cdot H}_{\text{单段峰值}}$$

段数 $n_{\text{seg}}$ 可取 $\sqrt{L}$ 使总量最优，工程上更常见的是每层一个边界、只重算当前层。

算力账：反向依赖的段内激活多算了一次前向，step 算力从 6P 涨到约 8P FLOPs/token，即 step 时间 +33% 左右。

4. Selective Checkpointing：更划算的切法

Megatron 的 selective（选择性）重算观察到一个不对称：attention 内部的中间量（$S \times S$ 的 score 矩阵、softmax 统计量）占显存随 $S^2$ 涨，但产出它们的算力相对便宜；而 GEMM 输出占显存少、算力贵。于是只重算 attention 内部，GEMM 输出照常保存——以极小的额外算力换掉增长最快的那部分显存。长序列训练里 selective 几乎总是优于 full；很多框架把它包进 FlashAttention 的反向实现（本来就不存 $S^2$ 矩阵），不用单独开。

5. 实践口径

• PyTorch：torch.utils.checkpoint.checkpoint(block, x)，或 HF gradient_checkpointing_enable()；
• 什么时候开：激活显存是 OOM 元凶、长序列（S 大）、想加大 micro-batch 改善流水线气泡占比；
• 什么时候不开：MFU 已经很高（没有算力余量换）、显存富余（白送 33% 时间）；
• 与并行的交互：TP 下每层激活按 TP 度切分，单卡账本除以 TP；PP 下每个 stage 只保存自己那几层的激活与边界。

相关页面

• 训练基础：一个 Step 的执行图四份显存数据与激活的生命周期
• 混合精度训练激活以 bf16 保存的精度前提
• Megatron 并行selective 重算与 TP 的出处语境
• GPU Trace 分析重算在 trace 里表现为反向段中多出的前向 kernel