4. 梯度检查点（Gradient Checkpointing）：时间换空间的哲学
说到显存优化，梯度检查点是我个人最常用的技巧之一。它的核心思想很简单——用计算量换显存。说白了，就是让模型在前向传播时只保留一部分中间激活值，剩下的在反向传播时重新算一遍。
你可能会问：「这听起来有点亏啊，多算一遍不是更慢吗？」嗯，确实会慢一些，但换来的是显存的大幅下降。我曾在一次大模型训练任务中，靠这个技巧把 batch size 从 8 提到了 32，训练速度反而更快了——因为减少了梯度同步次数。
4.1 为什么需要梯度检查点？
先看一个典型场景。你训练一个 10 层 Transformer，每层都会产生大量中间激活值。这些值在反向传播时要用到，所以必须存着。显存就这么被吃掉了。
我遇到过最夸张的一次，一个 7B 参数的模型，光中间激活值就占了 40GB 显存。参数本身才 14GB 左右。你想想看，这多浪费。
梯度检查点的做法是：只保存少数关键层的激活值，其他层的激活值在反向传播时重新计算。这样显存占用可以降低 50%-80%。
核心公式：
显存节省 ≈ (1 - 1/N) × 中间激活值大小
其中 N 是检查点间隔的层数。N 越大，省得越多，但重计算开销也越大。
4.2 checkpoint 函数的用法
PyTorch 提供了现成的接口——torch.utils.checkpoint.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):
# 普通前向
x = self.layer1(x)
# 使用梯度检查点
x = checkpoint(self.layer2, x)
x = self.layer3(x)
return x这里 self.layer2 的中间激活值就不会被保存。反向传播时，PyTorch 会重新运行 self.layer2 的前向计算来获取梯度。
个人习惯：我一般把检查点放在计算量最大的层上，比如 Attention 或 FFN。这样重计算的开销相对较小，但显存节省很明显。
4.3 适用场景
不是所有地方都适合用梯度检查点。我总结了几条经验：
- 深层网络：层数越多，效果越好。20 层以上的模型收益明显。
- 大 batch 训练：当显存成为瓶颈时，用检查点换 batch size 很划算。
- 长序列任务：比如 NLP 里的长文本、CV 里的高分辨率图像，中间激活值特别大。
- 单卡训练：多卡场景下，检查点可以和模型并行配合使用。
我曾经踩过的坑：在非常浅的网络（比如 4 层）上用检查点，结果训练时间翻了一倍，显存只省了 10%。得不偿失。建议至少 10 层以上再考虑。
4.4 开销分析
梯度检查点不是免费的午餐。我们来看具体数据：
| 模型层数 | 检查点间隔 | 显存节省 | 训练时间增加 |
|---|---|---|---|
| 12 | 4 | ~40% | ~15% |
| 24 | 6 | ~60% | ~25% |
| 48 | 8 | ~75% | ~35% |
从表格可以看出，层数越多，性价比越高。48 层的模型用检查点，显存省了 75%，时间只多了 35%。这在很多场景下是可以接受的。
为什么会这样？因为重计算只发生在反向传播阶段，而反向传播本身只占整个训练流程的一小部分。前向传播、参数更新、数据加载这些环节不受影响。
4.5 进阶用法
除了基本的 checkpoint 函数，还有一些高级技巧：
- 自定义检查点策略：不是每层都检查，而是每隔几层设一个。我习惯在每 4-6 层设一个检查点。
- 结合混合精度：检查点 + FP16 训练，显存能省到极致。但要注意数值稳定性。
- 动态调整：训练初期用检查点，后期显存压力小了可以关掉。不过这个实现起来有点麻烦。
一句话总结：梯度检查点是用 20%-30% 的训练时间，换 50%-80% 的显存。在显存受限的场景下，这笔交易非常划算。
嗯，这就是梯度检查点的核心内容。记住一点：不要盲目使用。先分析你的模型结构，找到显存瓶颈，再有针对性地设置检查点。这样效果最好。