梯度检查点：用时间换空间的经典套路

说到显存优化，梯度检查点（Gradient Checkpointing）是我最常用的手段之一。说白了，它的核心思想就一句话：不存中间激活值，用的时候重新算。

你想想看，训练一个深度网络，最吃显存的是什么？不是模型参数，而是中间层的激活值。一个 ResNet-152，光中间激活就能吃掉好几个 GB。我当年做视频理解模型时，一个 16 帧的输入，显存直接飙到 24GB 以上，根本跑不动。后来用了梯度检查点，显存降到 12GB，虽然训练慢了 20%，但至少能跑了。

原理与计算图

先看一个简单的计算图。假设我们有三个操作：A → B → C → loss。

正常前向传播时，每个节点的输出都会被保存下来：

# 普通前向
x = A(input)      # 保存 x
y = B(x)          # 保存 y
z = C(y)          # 保存 z
loss = loss_fn(z) # 保存 loss

反向传播时，这些保存的激活值直接拿来用：

# 反向传播
dz = grad(loss, z)  # 直接用 z
dy = grad(z, y)     # 直接用 y
dx = grad(y, x)     # 直接用 x

梯度检查点是怎么做的？它只保存输入和输出，中间的激活值全部丢弃：

# 检查点模式
x = A(input)        # 保存 x（检查点）
# y 和 z 不保存！
z = C(B(x))         # 前向计算，但不存中间结果

# 反向传播时
# 需要 y 怎么办？重新算！
x_restored = A(input)  # 重新前向计算 x（其实已经有了）
y_restored = B(x_restored)  # 重新计算 y
dy = grad(z, y_restored)

嗯，这里要注意：检查点之间的区域，前向计算会执行两次。一次是正常前向（不存中间值），一次是反向传播时的重计算。

核心公式：显存节省 ≈ O(N) 的中间激活 → O(1) 的检查点存储，但计算量增加约 1 倍前向时间。

我画了一张图，帮你理解这个过程：

torch.utils.checkpoint 实现

PyTorch 提供了现成的接口，用起来非常简单。我个人习惯把它封装成一个装饰器或者上下文管理器。

先看最基础的用法：

import torch
import torch.utils.checkpoint as cp

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 = cp.checkpoint(self.layer2, x)  # 这里不保存中间激活
        x = self.layer3(x)
        return x

嗯，就这么简单。但实际项目中，我经常遇到几个坑：

避坑指南：我曾经在一个 Transformer 模型上直接用 checkpoint 包装整个 forward，结果发现反向传播时重复计算了 embedding 层，导致训练时间翻倍。正确的做法是只包装计算密集的部分，比如 attention 和 FFN。

更实用的做法是包装一个函数块：

def transformer_block(x, attn, ffn):
    # 注意力层
    x = x + attn(x)
    # FFN 层
    x = x + ffn(x)
    return x

# 在 forward 中使用
x = cp.checkpoint(transformer_block, x, self.attn, self.ffn)

这里有个细节：checkpoint 函数的输入不能有 in-place 操作。为什么？因为重计算时输入可能已经被修改了。我建议在包装的函数内部避免使用 x += ... 这种写法，改用 x = x + ...。

小技巧：如果你需要保存某些中间结果（比如 attention 的权重用于可视化），可以用 cp.checkpoint 的 preserve_rng_state=False 参数，能省一点显存。不过要注意，这会丢失随机数状态，如果函数里有 dropout 就别用。

时间与空间权衡分析

梯度检查点不是银弹。它用时间换空间，这个「汇率」是多少？我实测过一些典型场景：

模型类型	显存节省	时间增加	适用场景
ResNet-50	~40%	~15%	小 batch 训练
BERT-base	~60%	~25%	长序列训练
GPT-2 1.5B	~70%	~30%	大模型微调
ViT-Large	~50%	~20%	高分辨率输入

从表格能看出来，模型越大、中间激活越多，检查点的收益越明显。但时间开销也不是线性的——因为重计算只发生在反向传播时，前向传播只慢一点点。

我个人的经验是：当显存成为瓶颈时，先用检查点把 batch size 翻倍。比如原来 batch size=8 显存爆了，用了检查点后 batch size=16 还能跑，那训练速度反而可能更快。因为更大的 batch 意味着更少的参数更新次数，整体收敛更快。

但要注意一个边界：检查点粒度过细反而会降低效率。我曾经试过对每个线性层都用 checkpoint，结果显存没省多少，时间翻了一倍。因为每个小操作的重计算开销太大了。一般建议对 3-5 个连续层或一个完整的 block 使用一次 checkpoint。

最佳实践：

对计算密集的模块使用（Transformer block、CNN stage）
避免对轻量操作使用（激活函数、LayerNorm）
检查点之间的跨度建议 4-8 层
结合混合精度训练，效果更佳

最后说一个我踩过的坑：检查点不能和 DataParallel 混用。因为 DataParallel 会在每个 GPU 上复制模型，而检查点的重计算逻辑会破坏梯度同步。用 DistributedDataParallel 就没问题。嗯，这个坑我当年调了两天才发现。

梯度检查点是个很实用的工具，但别滥用。我一般只在显存差一点点就能跑的时候才用。如果显存差太多，我会先考虑梯度累积、混合精度、模型并行这些手段。毕竟，时间也是成本嘛。

想要更多资料，联系可可

📚 课程目录

梯度检查点：用时间换空间的经典套路

原理与计算图

torch.utils.checkpoint 实现

时间与空间权衡分析