Activation Checkpointing#

本页内容4 个部分

1. 为什么反向传播需要保存激活#

普通训练在前向中保留反向求梯度所需的中间张量。Activation checkpointing（激活重算，也称梯度检查点）只保存选定边界的输入，反向需要内部张量时重新执行前向片段，用额外计算换取较低的激活峰值。它不直接减少参数、梯度和优化器状态。

2. 分段保存与重算的计算链#

2.1 等大线性链的理想模型#

普通保存#

设有 $L$ 层，每层需保存的激活大小近似相同。普通方式保存 $O(L)$ 份激活；这里只计激活，不是训练总显存。

分段重算#

每 $k$ 层分一段，保留约 $L/k$ 个边界。反向处理某段时，重算并暂存约 $k$ 层，因此理想存储阶数为：

M_{\text{act}}\propto\frac{L}{k}+k.

对 $k$ 求导， $-L/k^2+1=0$ 给出 $k\approx\sqrt L$ ，得到 $O(\sqrt L)$ 的激活存储。该推导依赖链式结构、近似等大的层与具体分段；逐层包装 checkpoint 并不自动得到同样的阶数。Sublinear Memory，2016。

2.2 48 层的手算示意#

每 8 层一段，保存约 6 个边界；重算当前段时最多再暂存约 8 层，共约 14 份等大激活，对比原来的 48 份，理想减少 $1-14/48\approx70.8\%$ 。忽略边界计数常数和其他张量，此数不能用作真实模型的显存承诺。

若完整重算约增加一次前向，设一次前向成本为 $F$ 、反向约为 $2F$ ，则总量从 $3F$ 变成 $4F$ ，计算量增加约 $33.3\%$ 。实际时间受重算范围、kernel、通信重叠与内存带宽影响，不能把这个 FLOP 估算直接等同于耗时增幅。

3. 实现边界与随机状态#

from torch.utils.checkpoint import checkpoint

# 模型 forward 中的示意片段
for block in self.blocks:
    if self.training:
        x = checkpoint(block, x, use_reentrant=False)
    else:
        x = block(x)

PyTorch 2.9 推荐显式使用 use_reentrant=False。默认会处理所支持设备的 RNG 状态，因此含 Dropout 不代表不能重算；自定义随机源、跨新设备迁移或改变外部可变状态仍可能造成前向与重算不一致。checkpoint 文档。

Hugging Face 模型若支持，可调用 gradient_checkpointing_enable()；需检查实际重算粒度，而不是只确认开关已打开。

4. 正确性与性能怎样检查#

固定参数、输入与随机性，比较开启前后的 loss 和梯度，再测峰值显存与端到端步时。检查重算函数是否依赖已变化的全局状态。按同一有效 batch 和精度比较，并区分 allocated、reserved 与设备总占用。

全量重算不一定最优；Transformer 可只重算高存储、低计算成本的部分，见选择性激活重算。显存其他部分由分布式训练与显存核算维护。

累积、AMP 与重算分别改变 microbatch 激活、算子 dtype 和保存张量的方式。组合后的收益要实测，不能把各自百分比直接相乘，更不能据此保证指定显卡容纳某个模型。