3. 梯度检查点 (Gradient Checkpointing)
好，咱们来聊聊梯度检查点。说实话，这玩意儿是我在训练一个 7B 模型时被逼着深入研究的技术。当时显存爆得我怀疑人生，恨不得把显卡拆了手动扩容。后来发现，其实有个更优雅的解法——用时间换空间。
3.1 核心原理：为什么它能省显存？
先想想看，训练一个模型时，显存都花在哪了？
- 模型参数（权重）
- 优化器状态（Adam 的动量、方差）
- 激活值（中间层的输出）
其中，激活值占了大头。尤其是 Transformer 这种结构，序列越长，激活值越恐怖。我见过一个 13B 模型，光激活值就吃了 40GB 显存。
梯度检查点的思路很简单：不保存所有中间激活值，而是在反向传播时重新计算它们。
具体来说，我们把网络分成若干段（checkpoint segments）。前向传播时，只保存每个段的输入和输出，中间的激活值全部扔掉。反向传播时，从保存的输入开始，重新跑一遍前向，把中间激活值算出来，再算梯度。
关键点： 这本质上是用计算量换显存。前向多算一次，显存省一半以上。
3.2 实现方式：PyTorch 里怎么搞？
PyTorch 官方提供了 torch.utils.checkpoint 模块，用起来非常直接。我个人习惯把它包装成一个函数，方便复用。
import torch
import torch.utils.checkpoint as checkpoint
def checkpoint_wrapper(module, *args):
"""对任意模块应用梯度检查点"""
return checkpoint.checkpoint(module, *args, use_reentrant=False)
# 使用示例
class MyModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.block1 = torch.nn.Linear(1024, 1024)
self.block2 = torch.nn.Linear(1024, 1024)
self.block3 = torch.nn.Linear(1024, 1024)
def forward(self, x):
# 只对中间层使用检查点
x = self.block1(x)
x = checkpoint_wrapper(self.block2, x) # 这里省显存
x = self.block3(x)
return x注意 use_reentrant=False 这个参数。我刚开始用的时候没加，结果在某个自定义层上出了 bug，排查了半天。后来发现这是 PyTorch 1.10 之后推荐的做法，能避免一些诡异的梯度问题。
小技巧： 不是所有层都需要加检查点。通常只对计算量大、激活值大的层（比如 Transformer 的 attention 和 FFN）使用。像 embedding 层这种，加了反而得不偿失。
3.3 时间与空间的权衡：到底值不值？
说白了，这就是个 trade-off。你省了显存，但多花了时间。具体多少呢？我拿一个实际案例给你看。
假设我们有一个 12 层 Transformer，每层 hidden size 为 4096，序列长度 2048，batch size 为 8。
| 配置 | 显存占用 | 训练时间（每步） | 备注 |
|---|---|---|---|
| 无检查点 | ~48 GB | ~1.2 秒 | 全部激活值都保存 |
| 每层都加检查点 | ~22 GB | ~2.1 秒 | 显存省了 54%，时间多了 75% |
| 隔层加检查点 | ~32 GB | ~1.5 秒 | 折中方案，显存省 33%，时间多 25% |
看到没？不是加得越多越好。我曾经在一个项目里，为了省显存，给所有层都加了检查点。结果训练时间翻了一倍，显存只省了 60%。后来改成每隔一层加一个，效果反而更好。
避坑指南： 我曾经在 batch size 很小的时候强行加检查点，结果显存没省多少，时间却暴涨。原因很简单——检查点本身也有开销。当 batch size 小于 4 时，我建议你慎重考虑是否真的需要它。
3.4 知识体系：一张图看懂
下面这张 SVG 图，把梯度检查点的核心逻辑串起来了。我画的时候特意把「时间换空间」这个 trade-off 放在中心，因为这是理解它的关键。
3.5 实战建议：怎么选 checkpoint 点？
根据我的经验，有几点可以帮你快速决策：
- 优先加在计算密集的层：比如 Transformer 的 attention 和 FFN，它们产生的激活值最大。
- 不要加在输入输出层：embedding 和最后的 linear 层，激活值很小，加了反而增加开销。
- 隔层加比全加好：比如 12 层 Transformer，每 2 层加一个检查点，效果往往最优。
- 结合 batch size 调整：batch size 越大，检查点的收益越明显。batch size 小的时候，不如直接减 batch size。
一句话总结： 梯度检查点不是银弹，但在显存不够时，它是性价比最高的解法之一。用 30% 的时间开销，换 50% 的显存节省，这笔账怎么算都划算。
嗯，关于梯度检查点，核心就是这些。下次你训练大模型爆显存时，不妨先试试这个方案，比换显卡便宜多了。