梯度检查点技术：用时间换显存的经典套路

说实话，我第一次接触梯度检查点时，心里是拒绝的。

你想啊，本来训练一个模型就够慢了，还要在 forward 的时候故意不存中间结果，backward 的时候再重新算一遍——这不是自己给自己找麻烦吗？

但后来我在一个 80GB 显存的 A100 上跑一个 7B 的 LLaMA 微调任务，batch size 只能设到 2。换了梯度检查点之后，batch size 直接翻到了 8。嗯，真香。

什么是梯度检查点？

梯度检查点（Gradient Checkpointing），也叫「激活重计算」。说白了就是：forward 的时候只存一小部分关键激活值，backward 的时候需要哪些中间结果，就现场重新算一遍。

我习惯把它理解成「显存版的虚拟内存」——把不常用的数据暂时扔掉，用的时候再从磁盘（这里是计算）换回来。

核心思想：用额外的计算开销，换取显存空间的释放。

原理：到底省了哪部分显存？

我们先看看一个普通的训练流程里，显存都花在哪了：

  • 模型参数：权重本身，这个省不了
  • 优化器状态：Adam 的 momentum 和 variance，也省不了
  • 激活值（Activations）：每一层的中间输出——这才是大头
  • 梯度：反向传播时产生的中间梯度

激活值占了多少？我举个例子。一个 24 层的 Transformer，序列长度 1024，隐藏维度 1024，光激活值就能吃掉 3-4GB 显存。层数越多、序列越长，这个数字越恐怖。

梯度检查点就是针对「激活值」下手的。

实现方式：怎么在代码里用？

如果你用 PyTorch，其实就一行代码的事：

from torch.utils.checkpoint import checkpoint

# 原始 forward
def forward(self, x):
    return self.ffn(self.attention(x))

# 加了检查点
def forward(self, x):
    return checkpoint(self.ffn, self.attention(x))

等等，别急着复制。这里有个坑——checkpoint 函数要求输入不能有 in-place 操作。我曾经在一个 LayerNorm 里用了 x += residual，结果 backward 的时候梯度死活不对。排查了半天，才发现是 in-place 惹的祸。

避坑指南：使用 checkpoint 时，确保模块内部没有 in-place 操作（如 +=、.add_()）。否则梯度会算错，而且不会报错——这才是最坑的。

如果你用 HuggingFace Transformers，更简单：

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./output",
    gradient_checkpointing=True,  # 就这一行
    per_device_train_batch_size=4,
)

嗯，HuggingFace 已经帮你封装好了。但我建议你还是理解一下底层原理，不然出了问题都不知道怎么调。

时间与空间的权衡：到底值不值？

这是最核心的问题。我们来看一组实测数据（以 GPT-2 1.5B 为例，batch size=8，序列长度=1024）：

配置 显存占用 训练时间（每步） 加速比
无检查点 24.3 GB 1.2s 1.0x
全层检查点 12.1 GB 1.9s 0.63x
选择性检查点（每2层） 16.8 GB 1.4s 0.86x

看到没？全层检查点省了 50% 的显存，但慢了 58%。而选择性检查点只省了 30% 的显存，但只慢了 16%。

我个人习惯的做法是：先看显存瓶颈在哪。如果就差那么 2-3GB，用选择性检查点就够了。如果差得很多，那就全层检查点，配合梯度累积一起用。

什么时候该用？什么时候不该用？

我总结了几条经验：

  • 该用：
    • 模型太大，batch size 上不去
    • 序列长度很长（比如 4096 以上）
    • 显存刚好差那么一点，换卡又太贵
  • 不该用：
    • 模型本身很小，显存绰绰有余
    • 训练速度是瓶颈，显存不是
    • 你的模型里有大量 in-place 操作（改起来太麻烦）

小技巧：如果你用 PyTorch 的 torch.utils.checkpoint，可以配合 torch.no_grad() 一起用。在不需要梯度的推理阶段，直接关掉检查点，省掉不必要的重计算开销。

核心逻辑流程图

下面这张图展示了梯度检查点的工作流程，我画了个 SVG 方便你理解：

梯度检查点工作流程 普通训练 Forward 存激活值 Backward 更新参数 显存占用：高（存了所有激活值） 梯度检查点训练 Forward 只存关键点 重计算 Backward 更新参数 显存占用：低（只存关键点） 权衡分析 • 时间开销：重计算阶段增加了约 30%-60% 的 backward 时间 • 空间收益：显存占用降低 40%-60%，可以支持更大的 batch size • 适用场景：显存受限但可以接受训练速度下降的情况

一些实战中的坑

最后分享几个我踩过的坑：

  1. 不要对输入做检查点——如果你把 embedding 层也包进 checkpoint，反而会浪费显存，因为 embedding 的激活值本来就不大。
  2. 注意梯度累积的配合——用了梯度检查点之后，batch size 上去了，但梯度累积步数要相应减少。我见过有人两个一起开，结果显存是省了，但训练效果变差了。
  3. 检查点粒度不是越细越好——每层都做检查点，重计算开销会爆炸。我一般建议每 2-4 层设一个检查点。
  4. 混合精度训练下要小心——fp16 的梯度检查点有时会引入数值误差。我习惯在关键层（比如 attention 的 softmax）保留 fp32 精度。

嗯，梯度检查点就是这么个东西。说白了就是「用时间换空间」——在显存不够用的时候，这是一个性价比很高的选择。但别盲目用，先算清楚你的显存瓶颈到底在哪，再决定要不要上。

想要更多资料，联系可可

← 上一章 📖 返回目录 下一章 →