技术博客

梯度检查点：用 33% 的额外计算换显存平方级的节省

第 144 / 147 篇 · 2026-01-06 · ai · 显存优化反向传播训练工程

反向传播的数学要求每个算子的局部梯度，而局部梯度要依赖前向的输入。朴素实现把前向的所有中间激活都留在显存里直到反向用完：一个 $L$ 层、激活宽度 $h$、序列 $s$ 的 Transformer，激活显存量级是

$$ M_{\text{act}} \sim L \cdot s \cdot h \cdot c \quad (c \approx 10\text{--}20,\ \text{每层每元素保留的中间量个数}). $$

层数翻倍显存翻倍，序列翻倍更是翻四倍（$s$ 以平方出现在注意力里）。梯度检查点（gradient checkpointing，也叫 activation recomputation）用一个朴素的观察打破这个线性增长：前向是可以重算的。

全量重计算：分段存储边界

把网络切成 $\sqrt{L}$ 个段，每段 $\sqrt{L}$ 层。前向时只在段边界保存激活（$\sqrt{L}$ 份），段内部的中间激活算完即弃。反向到某段时，从该段边界的保存激活出发重跑一遍段内前向，恢复中间激活，再做正常反向。

显存：存 $\sqrt{L}$ 份边界激活，加上正在重算段内的 $\sqrt{L}$ 份，总量 $O(\sqrt{L})$ 份——从 $O(L)$ 到 $O(\sqrt{L})$，层数方向是平方级的节省。取段数 $k = \sqrt{L}$ 不是巧合：设总层数 $L$、段数 $k$，显存 $\propto L/k + k$（边界激活 + 段内激活），对 $k$ 求极小值，$k^* = \sqrt{L}$，此时两部分相等，$M \approx 2\sqrt{L}$。通用优化结论：重算成本固定时，让"存的"和"重算的"相等。

计算代价：多出来的是被丢弃段的前向。FLOPs 账上，Transformer 前向约占 $1/3$、反向占 $2/3$（反向要对每个算子算输入梯度和权重梯度）。重算一次前向，总 FLOPs 从 $1+2=3$ 变成 $1+1+2=4$，+33%。这笔交换在现代 GPU 上几乎总是划算的：训练是访存密集型的（下一篇推理加速会展开），33% 的额外计算常常只换来 10–20% 的墙钟时间损失，而显存省下的部分可以拿去换 4 倍的 batch size 或序列长度。

梯度检查点的时间线：前向只存段边界，反向到段时重算段内前向 段1 前向（存边界） 段2 前向（丢弃） 段3 前向（存边界） 段4 前向（丢弃） 段5 前向（存边界） 反向开始： 重算段5 前向 段5 反向 重算段4 前向 段4 反向 重算段3 前向 段3 反向 内存占用曲线：前向期只有 O(k) 份边界激活；反向到某段时临时抬升到 O(k + L/k) 总 FLOPs：+33%（多跑一次前向）；最优段数 k = sqrt(L)，使边界存储与段内重算相等

选择性重计算：Megatron 的改良

全量重算对所有层一视同仁，但不同算子的"存储成本/重算成本"比差异巨大：

  • Dropout、残差加法、激活函数：反向只依赖输入（甚至只依赖输出），存储一份的代价小，重算的代价也小——不值得重算，直接存。
  • Attention 的 softmax(QK^T)V：中间矩阵是 $s \times s$，序列长了存储爆炸，但 FlashAttention 这类 kernel 本身就以分块方式重算（不物化完整注意力矩阵），天然配合重计算策略。
  • 大 FFN 的中间激活（SwiGLU 有两份）：宽度方向占显存大头，重算只是两次矩阵乘，划算。

Megatron-LM 的选择性重计算（selective recompute）就按这个清单来：只重算 attention 块和昂贵的 FFN 部分，其余照常存储。实测在 GPT-3 量级模型上，选择性方案以约 2/3 的额外 FLOPs 拿到全量重计算 90% 的显存收益——性价比曲线明显更陡。DeepSpeed 和 PyTorch 的 checkpoint API 都支持按函数粒度选择。

实现细节与坑

PyTorch 的 torch.utils.checkpoint.checkpoint(fn, *args) 把 fn 包成"前向不存中间激活、反向时重跑"的形式。用起来有三个细节：

  1. RNG 状态。重算必须复现原前向的随机性：dropout 的掩码、随机增广、（如果用到了）随机舍入。checkpoint 会保存和恢复 CUDA RNG 状态，但这个机制只在函数内部消费随机数时可靠——函数外的全局随机数使用顺序在两次前向之间变化，会造成重算与初算不一致。表现通常是 loss 曲线出现小的系统偏移而非崩溃，很难查。
  2. 与归一化层统计量的交互。前向时 BatchNorm 的 running stats 在重算时会再次更新（累加两次），要用 functional 接口或关掉动量。LN/RMSNorm 无此问题，又一个 Transformer 友好的点。
  3. 编译器组合。torch.compile 和 checkpoint 组合时，重算段会被编译器作为独立图处理，可能引入 graph break；新版本 PyTorch 的 torch.utils.checkpoint 已和编译器打通（preserve_rng_state + compiled region），但自定义 autograd 函数嵌套 checkpoint 仍容易踩到。

和流水并行组合时还有一个结构性收益：PP 的每段本身就需要保存边界激活供段间传递，checkpoint 的分段边界和 PP 的段边界天然对齐，两者叠加不增加额外的重算轮次。

小结

梯度检查点是把"显存-计算交换"做到极致的工具：理论账上 $O(L)$ 换 $O(\sqrt{L})$，代价 33% FLOPs，最优分段数取存储与重算成本相等处。工程上，选择性重计算把注意力这类"存不起又重算便宜"的算子挑出来，用 2/3 的代价拿到大部分收益。判断要不要开它，就看显存的瓶颈是不是激活（而不是模型状态或 KV cache）——长序列、大 batch 的预训练几乎总是答案为是。