---

一、一句话概括

Activation Checkpointing，中文通常叫梯度检查点，也有叫激活重计算（Activation Recomputation）。

它的核心思想是：前向传播时不保存全部中间激活值，只保存少量检查点；反向传播时，从最近的检查点重新计算需要的激活值。

这样做的直接效果是：（1）显存占用大幅下降；（2）计算量增加，训练变慢。

一句话总结：用计算换显存。

---

二、为什么需要 Activation Checkpointing

1、训练时的显存都花在哪里

训练一个神经网络，显存主要花在：

• 模型参数
• 梯度
• 优化器状态
• 中间激活值

其中，中间激活值往往是最容易被忽视、但占用极大的一部分。

2、为什么激活值占用大

反向传播计算梯度时，需要用到前向传播时每一层的输出，也就是激活值。

假设网络有 $L$ 层，每层的激活值大小为 $A$，那么保存全部激活值需要：

$$ \text{Activation Memory} \approx L \times A $$

对于大模型：

• 层数 $L$ 很大，比如几十到上百层；
• 每层激活值 $A$ 也很大，因为 batch size 和隐藏维度都很大。

所以激活值显存会随层数线性增长，成为训练显存的主要瓶颈之一。

3、和 batch size 的关系

激活值大小和 batch size 成正比。batch size 越大，激活值显存越大。

很多时候，训练 OOM 不是因为模型放不下，而是因为激活值放不下，导致 batch size 上不去。

Activation Checkpointing 就是为了解决这个问题。

---

三、Activation Checkpointing 的基本原理

1、普通反向传播

普通训练流程是：

• 1、前向传播，逐层计算激活值；
• 2、把每一层的激活值都保存下来；
• 3、反向传播时，直接用保存的激活值计算梯度。

优点是反向传播快，因为激活值现成的。

缺点是显存占用高，因为要保存所有层的激活值。

2、使用 Checkpointing 后

使用 Activation Checkpointing 后：

• 1、把网络分成若干段；
• 2、每段只保存段边界处的激活值，也就是检查点；
• 3、段内部的中间激活值不保存，前向算完就丢掉；
• 4、反向传播到某一段时，从该段最近的检查点重新做一次前向，重建段内需要的激活值；
• 5、用重建的激活值计算梯度。

这样，显存里只需要保存少量检查点，而不是全部激活值。

3、一个直观例子

假设有 12 层，每 4 层设一个检查点：

• 只保存第 0、4、8、12 层的激活值；
• 第 1 到 3 层的激活值不保存；
• 反向传播到第 3 层时，从第 0 层检查点重新前向计算第 1 到 3 层。

这样显存从保存 12 层激活值，降到保存 4 层检查点，再加上临时重建的一小段。

---

四、为什么能省显存

1、显存节省的来源

普通训练保存全部 $L$ 层激活值，显存约为：

$$ O(L \times A) $$

使用 Checkpointing 后，假设每 $k$ 层设一个检查点，则只需要保存：

$$ O\left(\frac{L}{k} \times A\right) $$

再加上反向时临时重建的 $O(k \times A)$。

当 $k$ 较大时，显存占用显著下降。

2、极端情况下

如果整个网络只设一个检查点，也就是只保存输入：

• 显存里几乎不存中间激活值；
• 反向传播时几乎要从头重算一遍前向。

显存最省，但计算代价最大。

3、和平方根法则

如果网络有 $L$ 层，每层激活值相同，理论最优的检查点间隔约为：

$$ \sqrt{L} $$

此时显存占用和计算代价之间达到较好的平衡：

• 显存从 $O(L)$ 降到 $O(\sqrt{L})$；
• 计算量从 $O(L)$ 增加到约 $O(2L)$，也就是多一次前向。

这就是经典的平方根法则。

---

五、会带来什么代价

1、计算量增加

反向传播时，需要重新计算被丢弃的激活值，相当于多做一次前向。

普通训练一次迭代的计算量约为：

$$ \text{前向} + \text{反向} \approx 1 + 2 = 3 $$

使用 Checkpointing 后，计算量约为：

$$ \text{前向} + \text{重算前向} + \text{反向} \approx 1 + 1 + 2 = 4 $$

所以训练时间通常会增加约 20% 到 30%，具体取决于检查点粒度和实现。

2、训练速度变慢

因为多了重算，GPU 计算时间增加，吞吐下降。

如果训练本身是计算受限，而不是显存受限，那么开启 Checkpointing 会得不偿失。

3、实现复杂度增加

需要：

• 合理划分检查点段；
• 处理反向传播时的重计算逻辑；
• 和流水线并行、张量并行等配合；
• 避免重计算时显存峰值过高。

4、不是所有层都适合

有些层计算便宜但激活值大，适合设检查点；有些层计算昂贵但激活值小，设检查点反而亏。

所以实际中往往只对部分层开启，比如 Transformer 的某些块。

---

六、在什么场景下值得开启

1、值得开启的场景

• 显存是硬瓶颈：模型能算，但激活值放不下，batch size 上不去；
• 需要更大 batch size：大 batch 对训练稳定性、吞吐都有好处；
• 模型层数很深：激活值随层数线性增长，Checkpointing 收益明显；
• 训练大模型：LLM 训练中几乎标配；
• 显存换计算划算：计算资源相对充足，显存资源紧张。

2、不太值得开启的场景

• 显存充足：没有 OOM 压力，开启只会变慢；
• 计算是瓶颈：GPU 已经跑满，再加重算会更慢；
• 浅层小模型：激活值本来就不大，收益有限；
• 对训练速度极度敏感：比如需要快速迭代的实验。

3、实际中的折中做法

• 只对部分层开启，比如每隔几层设一个检查点；
• 结合 ZeRO、混合精度、梯度累积一起使用；
• 用工具自动决定哪些层设检查点；
• 在显存和速度之间做实验调优。

---

七、和其他显存优化技术的关系

Activation Checkpointing 通常不是单独使用，而是和以下技术组合：

• 混合精度训练：降低参数和激活值的存储精度；
• ZeRO：切分优化器状态、梯度、参数；
• 梯度累积：用小 batch 模拟大 batch；
• 张量并行、流水线并行：切分模型和层。

其中，Checkpointing 主要解决激活值显存，而 ZeRO 主要解决模型状态显存，两者互补。

---

八、简单评价这道面试题

1、这道题的价值

• 这是大模型训练中非常实用的一道题；
• 一道题能考察激活值显存、重计算原理、计算与显存权衡；
• 区分度好，初级可能只知道“能省显存”，中级能说清原理，高级能讲清平方根法则和实际调优；
• 贴近实战，PyTorch 的 torch.utils.checkpoint、DeepSpeed、Megatron 都支持。

2、这道题的不足

• 题目偏概念，如果没有实际训练经验，容易只答表面；
• 没有涉及具体显存计算，比如给定模型算激活值；
• 没有涉及和流水线并行、张量并行的交互；
• 对“什么时候不该开”的考察可以更深入。

3、建议的追问方向

• 给定一个 $L$ 层网络，Checkpointing 的显存和计算量怎么估算？
• 为什么说最优检查点间隔是 $\sqrt{L}$？
• Checkpointing 和流水线并行一起用时要注意什么？
• PyTorch 的 torch.utils.checkpoint 是怎么实现的？
• 如果重计算时显存峰值反而变高，可能是什么原因？

4、总体评价

这是一道很好的大模型训练显存优化面试题，适合考察候选人对激活值显存和计算换显存思想的理解。它既能考察基础原理，也能延伸到工程调优。如果候选人能把 Checkpointing 的原理、显存节省来源、计算代价、平方根法则以及适用场景都讲清楚，基本可以判断其具备大模型训练显存优化的扎实基础。

---

---