什么是梯度检查点（激活重计算）？它如何用时间换显存？

简化版

梯度检查点（gradient checkpointing，又叫激活重计算） 是一种“用计算换激活显存”的技巧：前向只保存选定边界，其余中间激活丢弃，反向到该段时再重跑前向。均匀分段的经典分析可把层方向峰值从 O(L) 降到约 O(√L)，但实际节省与减速取决于切分、算子和硬件，不能固定写成 30%。

详细版

为什么激活这么占显存

反向传播算梯度用到链式法则，需要前向时每一层的中间激活。训练时的显存主要由四块组成：参数、梯度、优化器状态、激活。其中激活显存 ∝ batch × 序列长度 × 隐藏维 × 层数，在长序列/大 batch 时会非常庞大，常常是显存瓶颈。

梯度检查点的做法

普通前向：保存每一层的激活 → 反向直接用      → 激活显存 O(L)
检查点：  只保存 √L 个检查点的激活 → 反向时从检查点重算中间激活 → 激活显存 O(√L)

把网络分成若干段，只在段边界保存激活（检查点）。
反向传播进入某段时，以该段起点的检查点为输入，重新前向一遍算出段内各层激活，再算梯度。
用完即丢，不长期占显存。

代价与收益

指标	普通	梯度检查点
激活显存	O(L)	约 O(√L)
计算量	1 次前向 + 1 次反向	额外重算被检查点覆盖的前向算子

适用场景

长上下文训练、大 batch、单卡显存吃紧时。
和混合精度、ZeRO、并行等叠加使用。

完整版教学

一、先看清训练显存都花在哪

训练一个模型，显存主要被四部分吃掉：

参数：模型权重。
梯度：与参数等大。
优化器状态：Adam 的动量 m、方差 v、FP32 副本等（往往最大头之一）。
激活（activations）：前向过程中每一层产生的中间结果，反向要用它们算梯度。

前三块由参数量决定，相对固定；而激活显存随 batch、序列长度、层数线性增长，在长序列训练里会爆炸式增大，常常成为压垮单卡的最后一根稻草。梯度检查点专治这一块。

二、为什么反向传播非要激活不可

反向传播用链式法则逐层往回算梯度。算某一层参数的梯度时，公式里会用到该层前向时的输入/中间激活。举个直觉例子：一个乘法 y = w·x，∂y/∂w = x——要算对 w 的梯度，就必须知道当时的输入 x。所以普通实现会在前向时把所有层的激活缓存下来，供反向使用。层越多、序列越长，这个缓存越大。

三、核心思想：不存就重算

梯度检查点的洞察是：激活是可以”重新算出来”的——只要保留了某个位置的输入，就能通过前向计算再得到后面的激活。既然如此，何必全都存着占显存？不如丢掉大部分，需要时重算。

具体做法（以按段划分为例）：

把 L 层网络切成若干段，只在段的边界保存激活（这些就是”检查点”）。
前向时，段内的中间激活算完就丢弃，不缓存。
反向传播到某段时，取该段起点的检查点作为输入，重新前向计算一遍这一段，临时得到段内激活，用来算梯度，算完再丢。

这样常驻的激活显存只剩下「检查点」的量，大幅下降。

四、为什么是 O(√L)：一个漂亮的权衡

如果把 L 层均匀分成 √L 段、每段 √L 层：

常驻激活 = 检查点数量 ≈ √L 段的边界激活 → O(√L)。
重算开销 = 若几乎完整覆盖网络，合计可能接近额外一次前向；在“反向约为前向两倍 FLOPs”的简化模型中，总 FLOPs 增幅约三分之一。真实 step time 还受通信重叠、内存带宽和选择性覆盖影响，必须实测。

均匀分段时，段数越多意味着边界检查点越多、单段临时激活越少；段数越少则边界少、单次重算段更长。让“边界激活数”和“最大段内临时激活数”量级接近，可得到约 √L 的峰值量级。真实 Transformer 还受 Attention、MLP、序列并行和框架保存策略影响，因此常用 selective checkpoint 做实测取舍。

记忆钩子：梯度检查点就是“少存多算”：节省的是被丢弃的前向激活，付出的是反向前重算；比例由检查点边界决定，不是固定的 30%。

五、在整套显存优化里的位置

梯度检查点只省激活这一块，要和其他手段配合才能训动大模型：

梯度检查点：省激活显存（用计算换），并要求正确处理随机状态。
混合精度（BF16）：参数/激活/梯度降到 16 位，整体减半。
ZeRO/FSDP：分片参数/梯度/优化器状态，省前三块的冗余。
张量/流水线并行：把模型拆到多卡。

它是专门针对激活的常用手段，在长上下文训练中价值尤其大；副作用包括额外计算、step time 增长，以及随机状态处理不当时的梯度偏差。是否启用和覆盖哪些层应由显存目标与训练吞吐共同决定。

六、用 16 层分段理解重算

若 16 层网络原本保存每层激活，需要保留 16 份；每 4 层设置一个检查点时，前向只长期保存第 0、4、8、12、16 层附近的边界。反向处理第 13—16 层前，从第 12 层激活重算该段中间值，用完即释放，再处理前一段。

策略	保存激活	额外计算	特点
不重算	约 16 层	0	显存最高
每 4 层检查点	约 4 个边界 + 一段临时值	多次段内前向	折中
更激进重算	更少	更多	速度代价更大

forward: save L0 -> recompute L1..L4 -> save L4 -> ...
backward: reload L12 -> recompute L13..L16 -> gradients -> release

重算要求随机操作可复现，否则 dropout 掩码不同会让反向不再对应原前向。验收时应同时比较峰值显存、step time、loss 曲线和梯度一致性，而不是只确认“不 OOM”。

七、常见误区与追问

误区：梯度检查点能省参数/优化器显存。 不能，它只省激活显存；参数/梯度/优化器状态要靠 ZeRO、混合精度等。
误区：重算天然保证与原前向完全一致。 确定性算子且正确保存/恢复 RNG 时应数值等价；Dropout 掩码、带状态算子或非确定性 kernel 处理不当会产生差异。
追问：为什么反向需要激活？ 链式法则里参数梯度依赖前向的输入/中间值（如 ∂(wx)/∂w = x）。
追问：省了多少、代价多少？ 均匀分段的经典分析可降到约 O(√L)，代价是重算覆盖段；实际比例必须 Profile。
追问：为什么长序列训练一定要它？ 激活显存 ∝ 序列长度，长上下文时它是最大瓶颈。
追问：和混合精度冲突吗？ 不冲突，正交叠加，都是省显存的不同维度。

八、加强记忆

梯度检查点记“少存多算、时间换激活显存”：前向只保留边界，反向从边界重算段内激活。均匀分段的经典分析可达到约 O(√L) 的层方向峰值，但具体节省和减速取决于切分、序列长度、算子与硬件，不能背固定 30%。它不减少参数、梯度或优化器状态，还必须正确保存 RNG 才能让 Dropout 重算一致；最终应和混合精度、ZeRO/FSDP 并用，并同时验收峰值显存、step time、loss 与梯度。

简化版

详细版