3. 梯度检查点技术：原理、实现方式、与计算时间的权衡

梯度检查点，说白了就是「拿计算换内存」。

我刚开始接触大模型训练时，最头疼的就是显存爆炸。一个 7B 的模型，光中间激活值就能吃掉几十 GB。那时候我就在想：能不能不存那么多中间结果？

嗯，梯度检查点就是干这个的。

3.1 原理：为什么能省显存？

先说说正常训练时显存都花在哪了。你想想看，前向传播时，每一层的输出都得留着。为什么？因为反向传播要用啊。比如一个 100 层的 ResNet，每层都存一份激活值，那显存压力可想而知。

梯度检查点的思路很简单：我不存所有中间结果，只存一部分。反向传播时，没存的那部分怎么办？重新算一遍。

这就好比你去图书馆查资料。正常做法是把所有书都借回家（存着），随时翻看。梯度检查点呢？我只借几本关键的书，其他书需要时再去图书馆翻。省了家里的书架空间，但多跑了几趟路。

核心公式：

显存节省 ≈ (1 - 1/N) × 原始激活值显存

计算开销增加 ≈ (N - 1) × 额外一次前向传播

其中 N 是检查点间隔的层数

3.2 实现方式：PyTorch 里怎么用？

PyTorch 官方提供了 torch.utils.checkpoint 模块。用起来其实很简单：

import torch
from torch.utils.checkpoint import 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(self.layer2, x)  # 这一层不存激活值
        x = self.layer3(x)
        return x

我个人习惯把检查点用在那些计算量不大但显存占用大的层上。比如 Transformer 里的 attention 层，激活值特别大，但计算量相对较小，用检查点就很划算。

小技巧：

我在项目中遇到过这种情况：把检查点用在所有层上反而更慢。后来发现，检查点间隔 4-8 层是最优的。间隔太小，重计算次数太多；间隔太大，省不了多少显存。

3.3 计算时间的权衡：到底值不值？

你可能会问：多算一遍前向传播，时间上划得来吗？

我们来算笔账。假设一个模型有 L 层，每层前向传播耗时 T_f，反向传播耗时 T_b。正常训练总时间：

正常训练： T_normal = L × (T_f + T_b)

使用检查点（间隔 N 层）： T_checkpoint = L × T_f + L × T_b + (L/N) × T_f

看到没？多出来的就是 (L/N) × T_f。这个值取决于 N 的选择。

检查点间隔 N	显存节省比例	时间增加比例	我的建议
1（每层都检查）	~90%	~30-40%	不推荐，太慢了
4	~75%	~10-15%	推荐，性价比高
8	~60%	~5-8%	最常用，平衡点
16	~40%	~2-3%	显存够用时可选

我曾经在一个 32 层的 Transformer 上做过实验。间隔 8 层时，显存从 48GB 降到了 20GB，训练时间只增加了 6%。这个 trade-off 我觉得很值。

避坑指南：

我曾经踩过一个坑：在 BatchNorm 层上用检查点，结果训练崩了。为什么？因为 BatchNorm 需要统计均值和方差，重计算时这些统计量会变。后来我改用 GroupNorm 或者 LayerNorm 就没事了。

另外，Dropout 层也要小心。重计算时 dropout mask 会重新生成，导致训练不一致。建议把 dropout 放在检查点外面。

3.4 高级技巧：自定义检查点策略

PyTorch 的 checkpoint 默认是「均匀间隔」策略。但实际项目中，不同层的显存占用差异很大。比如 Transformer 里，attention 层的激活值比 FFN 层大得多。

我建议的做法是：

先 profiling：跑一次小 batch，看看每层的显存占用
按需设置：显存大的层单独加检查点，小的层不用
动态调整：训练过程中监控显存，如果接近上限，自动增加检查点

# 自定义检查点策略示例
def custom_checkpoint_forward(model, x, layer_indices):
    """
    layer_indices: 需要检查点的层索引列表
    """
    for i, layer in enumerate(model.layers):
        if i in layer_indices:
            x = checkpoint(layer, x)
        else:
            x = layer(x)
    return x

嗯，这里要注意：自定义策略虽然灵活，但代码可读性会下降。我一般只在关键模块上用，不会全模型都搞这么复杂。

3.5 与其他技术的配合

梯度检查点不是万能的。它最适合的场景是：

模型太大：显存不够，但计算资源充足
batch size 太小：想增大 batch size 但显存不够
长序列：Transformer 处理长文本时，激活值爆炸

但如果你已经用了混合精度训练、ZeRO 优化、激活值压缩等技术，显存还是不够，那检查点就是最后的救命稻草。

一句话总结：

梯度检查点是用 5-15% 的训练时间，换 50-80% 的显存节省。具体值多少，取决于你的模型结构和检查点策略。

最后说一句：梯度检查点不是银弹。如果你的模型已经能塞进显存，就别用。但如果显存是瓶颈，它绝对是最简单有效的方案之一。

想要更多资料，联系可可