3. 梯度检查点（Checkpointing）：原理、实现、与重计算开销权衡

梯度检查点，说白了就是「用时间换空间」。

我刚开始接触大模型训练时，最头疼的就是显存爆炸。明明一个 7B 的模型，光参数就占了 14GB，加上优化器状态、激活值，一张 A100 根本塞不下。后来我发现了梯度检查点这个神器——虽然训练慢了 20%，但显存直接省了一半。

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

先想想看，正常训练一个网络，显存里要存什么？

• 模型参数：这个省不了
• 优化器状态：Adam 的话要存一阶矩和二阶矩
• 激活值：前向传播时每一层的输出
• 梯度：反向传播时计算出来的

其中激活值占大头。比如一个 1024x1024 的卷积层，batch size 64，光这一层的激活值就要 256MB。模型深了，几十层堆下来，显存直接爆炸。

梯度检查点的思路很简单：前向传播时只保存少数关键层的激活值（检查点），其他层的激活值全部扔掉。反向传播时，如果发现某层的激活值丢了，就从最近的检查点重新计算一遍。

我习惯把检查点比喻成「书签」。你读一本厚书，不可能每页都记住。你只记几个关键页码（检查点），需要某页内容时，翻回最近的书签重新读一遍。嗯，就是这个道理。

3.2 实现方式：PyTorch 中的 checkpoint

PyTorch 提供了现成的 torch.utils.checkpoint 模块。用法非常简单：

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

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

你看，只需要把你想节省显存的层包在 checkpoint() 里就行。PyTorch 会自动处理：

• 前向传播时，只保存该层的输入
• 反向传播时，如果发现激活值丢了，就用保存的输入重新计算一遍

3.3 重计算开销：到底慢了多少？

这是大家最关心的问题。我直接给结论：通常增加 20%-30% 的训练时间，但显存节省 50% 以上。

为什么会这样？因为反向传播时，你不仅要计算梯度，还要重新计算激活值。相当于每个检查点之间的层，前向传播做了两次。

策略	显存占用	训练时间	适用场景
无检查点	100%	100%	小模型、大显存
每层检查点	~30%	~130%	大模型、显存紧张
每 N 层检查点	~50%	~115%	平衡方案

我曾经在一个 13B 的 LLaMA 模型上做过测试。不加检查点，batch size 只能开到 4。加了检查点后，batch size 直接翻倍到 8，训练时间只增加了 18%。你说值不值？

3.4 检查点策略：怎么选？

不是所有层都适合加检查点。我总结了几个原则：

1. 计算密集的层优先：比如卷积、Transformer 的 attention。这些层重计算开销相对较小。
2. 激活值大的层优先：比如输出通道多的卷积层。省显存效果最明显。
3. 避免对简单层使用：比如 ReLU、LayerNorm。重计算的开销可能超过省显存的好处。

3.5 高级技巧：自定义检查点

PyTorch 的 checkpoint 虽然方便，但有时候不够灵活。比如你想对某个模块内部的部分层使用检查点，或者想控制检查点的保存频率。

我习惯自己写一个包装器：

class CheckpointWrapper(nn.Module):
    def __init__(self, module, interval=2):
        super().__init__()
        self.module = module
        self.interval = interval  # 每 interval 层保存一个检查点
    
    def forward(self, x):
        # 手动实现检查点逻辑
        for i, layer in enumerate(self.module):
            if i % self.interval == 0:
                # 保存检查点
                x = checkpoint(layer, x)
            else:
                # 正常前向
                x = layer(x)
        return x

这样你就可以控制检查点的密度了。interval 越大，显存占用越大，但重计算开销越小。我一般从 interval=4 开始调，然后根据显存占用情况逐步减小。

3.6 权衡的艺术：到底用不用？

说实话，没有银弹。我建议你按这个流程决策：

1. 先试试不加检查点，看 batch size 能不能满足需求
2. 如果显存不够，从计算量最大的层开始加检查点
3. 观察训练时间变化，如果超过 30%，考虑减少检查点数量
4. 如果还是不行，考虑其他显存优化技术（比如混合精度、梯度累积）

你想想看，显存优化说白了就是一场「空间换时间」的游戏。梯度检查点只是其中一张牌，但确实是最常用、最有效的一张。

嗯，关于梯度检查点，核心就是这些。记住一句话：检查点不是越多越好，找到那个「显存省得够多，时间慢得能接受」的平衡点，才是关键。