激活值检查点：用空间换时间的艺术

大家好，今天我们来聊聊激活值检查点（Activation Checkpointing）。

说实话，这玩意儿刚出来的时候，我第一反应是：这不就是拿计算换显存吗？后来在项目中真正用起来才发现，里面的门道比想象中多得多。

Checkpointing 原理：到底在干什么？

先说说基本原理。训练一个深度网络，前向传播时我们会保存每一层的激活值。反向传播时，这些激活值用来计算梯度。你想想看，一个 7B 的模型，光激活值就能吃掉几十 GB 显存。

Checkpointing 的思路很简单：我不保存所有激活值了，只保存一部分关键节点。反向传播时，如果发现某个激活值没保存，我就从最近的一个检查点重新算一遍前向传播。

核心思想：用额外的计算开销，换取显存空间的释放。

举个例子。假设你的网络有 100 层。正常训练时，你需要保存 100 层的激活值。用了 Checkpointing 后，你可能只保存第 10、20、30...90 层的激活值。反向传播时，比如你要算第 25 层的梯度，发现没保存？没关系，从第 20 层的检查点重新算一遍 20→25 的前向传播就行。

我刚开始接触这个技术时，总觉得这太浪费计算了。后来在训练一个 13B 的 GPT 模型时，显存实在撑不住了，才硬着头皮上了 Checkpointing。结果发现，实际训练时间只增加了 20-30%，但显存占用直接砍半。嗯，真香。

选择 Checkpointing 策略：不是所有层都值得

那么问题来了：哪些层该设成检查点？哪些不该？

我个人习惯把网络分成三类：

  • 计算密集层（比如大矩阵乘法）：重计算代价高，尽量少设检查点
  • 激活值巨大层（比如注意力层）：显存占用大头，必须设检查点
  • 计算轻量层（比如 LayerNorm、激活函数）：重计算几乎没成本，大胆设检查点

我曾经在一个项目中犯过错误：把所有的 Transformer 层都设成了检查点。结果呢？训练速度慢了 3 倍，显存倒是省了不少，但完全没法接受。后来我调整策略，只在每 4 层中选 1 层作为检查点，效果就好多了。

我的经验：对于 Transformer 模型，建议每 2-4 层设置一个检查点。具体多少层，取决于你的显存压力和可接受的速度损失。

Checkpointing 与重计算开销：这笔账怎么算？

我们来算一笔账。假设你的网络有 N 层，你设置了 K 个检查点。

  • 正常训练：需要保存 N 层的激活值
  • Checkpointing：只需要保存 K 层的激活值
  • 额外计算：反向传播时，需要重计算 (N/K - 1) 次前向传播

举个例子，N=32，K=8。那么你需要重计算 3 次完整的前向传播。每次前向传播的 FLOPs 是固定的，所以额外计算开销大约是 3 倍的前向传播时间。

但注意，前向传播通常比反向传播快很多。所以实际训练时间增加的比例，远小于重计算次数增加的比例。我实测过，32 层 Transformer，8 个检查点，训练时间只增加了约 25%。

注意：重计算会占用额外的 GPU 算力。如果你的 GPU 已经算力吃紧（比如 batch size 已经很大了），Checkpointing 可能会导致训练速度大幅下降。这时候你需要权衡：是降 batch size 还是接受更慢的训练速度。

Checkpointing 在 Transformer 中的应用

Transformer 是 Checkpointing 的典型应用场景。为什么？因为 Transformer 的激活值太大了。

一个标准的 Transformer 层包含：

  • 多头注意力（MHA）：需要保存 Q、K、V、注意力分数、输出等
  • 前馈网络（FFN）：需要保存中间激活值
  • LayerNorm：需要保存均值和方差

我算过一笔账：一个 12 层、隐藏维度 768、序列长度 2048 的 Transformer，单层激活值大约 50 MB。12 层就是 600 MB。如果 batch size 是 32，那就是 19.2 GB。这还没算上模型参数和优化器状态。

所以，在 Transformer 中，Checkpointing 几乎是必选项。

下面是我常用的 Checkpointing 配置：

# PyTorch 中的 Checkpointing 示例
import torch
from torch.utils.checkpoint import checkpoint

class TransformerLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.attention = MultiHeadAttention()
        self.ffn = FeedForward()
        self.norm1 = LayerNorm()
        self.norm2 = LayerNorm()
    
    def forward(self, x):
        # 对注意力层使用 Checkpointing
        def attention_fn(x):
            return self.attention(self.norm1(x))
        
        # 对 FFN 层使用 Checkpointing
        def ffn_fn(x):
            return self.ffn(self.norm2(x))
        
        # 实际使用时，可以只对其中一个使用 Checkpointing
        x = x + checkpoint(attention_fn, x)
        x = x + checkpoint(ffn_fn, x)
        return x

注意看，我在这里对 attention 和 FFN 都用了 Checkpointing。但实际项目中，我通常只对 FFN 使用 Checkpointing。为什么？因为 attention 的重计算代价太高了——它需要重新计算 Q、K、V 和注意力分数，这些操作在长序列下非常耗时。

我的建议：在 Transformer 中，优先对 FFN 层使用 Checkpointing。如果显存还不够，再考虑对 attention 层使用。不要一股脑全用上。

另外，还有一个细节：Checkpointing 的粒度。你可以对整个 Transformer 层使用 Checkpointing，也可以对层内的子模块使用。粒度越细，显存节省越多，但重计算开销也越大。我一般选择对整层使用，这样代码更简洁，性能也够用。

最后，给大家一个避坑指南：不要在第一个和最后一个 Transformer 层使用 Checkpointing。第一个层的输入是 embedding，激活值很小；最后一个层的输出直接连 loss，重计算会引入额外的梯度传播路径，容易出问题。我曾经在这上面踩过坑，训练 loss 死活不降，排查了半天才发现是最后一个层的 Checkpointing 搞的鬼。

好了，关于激活值检查点，今天就聊到这里。记住一句话：Checkpointing 不是银弹，但用好了，能帮你省下一半的显存。

激活值检查点核心逻辑 正常训练（无Checkpointing） 层1 层2 层3 层4 激活值 激活值 激活值 激活值 显存占用：4层激活值（高） Checkpointing训练 层1 层2 层3 层4 ✓ ✓ 保存 丢弃 丢弃 保存 显存占用：2层激活值（低） 反向传播时的重计算 反向传播需要层2的激活值 → 从层1检查点重新计算 层1 → 层2 → 层3 → 层4 重计算路径：从最近的检查点（层1）重新计算到目标层（层2） 核心权衡：显存节省 vs 计算开销 检查点越多 → 显存越省 → 重计算越多 → 训练越慢

总结一下：

  • Checkpointing 的核心是用计算换显存
  • 不是所有层都适合设检查点，要区分计算密集型和激活值密集型
  • Transformer 中优先对 FFN 使用 Checkpointing
  • 避免在首尾层使用 Checkpointing
  • 实际项目中，训练时间增加 20-30%，显存节省 40-60%

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →