激活值检查点：用时间换空间的经典策略

说到显存优化，激活值检查点（Activation Checkpointing）是我最常用的手段之一。说白了，它的核心思想就是——不存那么多中间结果，用的时候再算一遍。

你想想看，训练一个大模型时，前向传播会生成大量激活值（activation）。这些值本来是要留着给反向传播用的。但如果我们全存下来，显存很快就爆了。那怎么办？扔掉一部分，反向传播时重新计算。这就是检查点技术的本质。

  核心公式：显存占用 ≈ O(√N) 而不是 O(N)，其中 N 是网络层数。代价是大约 20-30% 的额外计算开销。

原理：为什么能省显存？

正常的训练流程是这样的：前向传播时，每一层都保存激活值 → 反向传播时直接用这些值算梯度。这就像你写作业时，把所有草稿纸都留着，方便后面检查。

检查点技术则不同。它只保存少数几个关键位置的激活值（比如每隔几层存一次）。反向传播时，如果发现需要的激活值没存，就从最近的一个检查点开始重新计算。这就像你只保留每道题的最终答案，中间步骤忘了就重新推导一遍。

我在项目中遇到过最极端的案例：一个 100 层的 ResNet，如果不做检查点，单卡 32GB 显存根本跑不了 batch size 32。用了检查点后，batch size 直接翻倍到 64，训练速度只慢了 15%。这个 trade-off 非常划算。

Checkpoint 策略选择

检查点不是非黑即白的选择。我个人习惯把策略分成三类：

策略类型	做法	省显存效果	额外计算开销
全部检查点	每一层都不存激活值	最大（约 70-80%）	最大（约 30-50%）
选择性检查点	只对特定模块做检查点	中等（约 40-60%）	中等（约 15-25%）
无检查点	全部保存激活值	无	无

嗯，这里要注意：不是所有层都适合做检查点。比如 embedding 层和 loss 计算层，它们的激活值很小，重新计算反而浪费。我一般只对 transformer 的 attention 和 FFN 模块做检查点。

我的经验：对于 12 层以上的 transformer，建议每 3-4 层设置一个检查点。这样显存和计算开销的平衡最好。

PyTorch 实现

PyTorch 提供了非常方便的接口——torch.utils.checkpoint。用起来很简单：

import torch
import torch.utils.checkpoint as checkpoint

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = torch.nn.Linear(1024, 1024)
        self.layer2 = torch.nn.Linear(1024, 1024)
        self.layer3 = torch.nn.Linear(1024, 1024)
    
    def forward(self, x):
        # 对 layer2 使用检查点
        x = self.layer1(x)
        x = checkpoint.checkpoint(self.layer2, x)
        x = self.layer3(x)
        return x

你看，只需要把想检查点的层包在 checkpoint.checkpoint() 里就行。PyTorch 会自动处理前向传播时不保存激活值，反向传播时重新计算。

但有个坑——检查点函数不能有副作用。比如你不能在检查点函数里修改全局变量或者 in-place 操作。我曾经因为这个 bug 排查了整整两天，梯度死活不对，最后发现是某个 in-place 的 ReLU 搞的鬼。

避坑指南：使用检查点时，确保被包裹的函数是纯函数（无副作用）。避免使用 in-place 操作，比如 x.relu_() 要改成 torch.relu(x)。

与梯度累积的配合

检查点 + 梯度累积，这两个技术配合起来效果非常好。梯度累积本身就是为了解决显存不足的问题——把一个大 batch 拆成多个小 batch，分别计算梯度再累加。

但梯度累积有个问题：每个 micro-batch 的前向传播都会产生激活值。如果 micro-batch 数量很多，激活值占用的显存会成倍增加。这时候检查点就派上用场了。

我建议的做法是：

先确定你能承受的最大 micro-batch size
在这个 size 下，用检查点把显存压到安全线以下
然后增加梯度累积步数来达到目标 batch size

举个例子，假设你想跑 batch size 256，但显存只够跑 batch size 32：

# 梯度累积步数 = 256 / 32 = 8
accumulation_steps = 8
optimizer.zero_grad()

for micro_batch in range(accumulation_steps):
    # 每个 micro-batch 都使用检查点
    output = checkpoint.checkpoint(model, input_data)
    loss = criterion(output, target)
    loss.backward()

optimizer.step()

这样组合使用，显存占用只跟 micro-batch size 有关，跟总 batch size 无关。我在训练 175B 参数的大模型时，就是靠这个组合拳把 8 卡 80GB 的显存利用率拉满的。

  关键点：检查点解决的是单次前向传播的显存问题，梯度累积解决的是总 batch size 的显存问题。两者互补，缺一不可。

知识体系图

下面这张图展示了激活值检查点的核心逻辑和决策流程：

这张图帮你理清思路：先判断显存是否够用，不够就选策略，然后用 PyTorch 实现，最后配合梯度累积把效果拉满。我在实际项目中就是按这个流程走的，屡试不爽。

小技巧：如果你用的是 Hugging Face 的 Transformers 库，可以直接传 gradient_checkpointing=True 参数，它会自动帮你做检查点。但底层原理还是我们上面讲的这套。

好了，激活值检查点就讲到这里。记住核心思想：用计算换显存。在显存紧张的时候，这招比换显卡便宜多了。

📚 课程目录

激活值检查点：用时间换空间的经典策略

原理：为什么能省显存？

Checkpoint 策略选择

PyTorch 实现

与梯度累积的配合

知识体系图