4、梯度与中间变量：自动求导机制、计算图与中间变量、梯度累积与显存、retain_graph的影响

好，咱们今天聊一个非常实在的话题——梯度与中间变量。

说实话，很多同学跑深度学习模型，loss 一设、backward 一调，就完事了。但模型一旦大起来，显存动不动就爆了，这时候你才意识到：梯度是怎么算的？中间变量到底占了多少显存？

我当年第一次调一个 7B 参数的大模型时，就栽在这个坑里。明明模型参数只占 14GB，但一跑训练，显存直接飙到 40GB。后来一查，好家伙，中间变量占了 20 多 GB。嗯，今天咱们就把这笔账算清楚。

4.1 自动求导机制：PyTorch 是怎么算梯度的？

说白了，自动求导就是 PyTorch 帮你把链式法则自动跑一遍。你只需要定义前向计算，框架自动帮你反向传播。

但这里有个关键点：PyTorch 不是算完就扔，它会保留中间结果。为什么呢？因为反向传播时需要用到它们。

核心机制：

  • 前向传播时，PyTorch 会构建一个动态计算图
  • 每个张量操作都会记录在图中，包括输入、输出、梯度函数
  • 调用 .backward() 时，从输出节点开始反向遍历计算图
  • 每经过一个节点，就计算局部梯度，并累加到叶子节点的 .grad 属性上

举个例子，你写了个 z = x * y，然后 z.backward()。PyTorch 内部会记录：

  • x 的梯度是 y
  • y 的梯度是 x
  • 这些计算需要用到前向时的 x 和 y 的值

所以，前向时的中间变量必须保留到反向结束。这就是显存占用的根源之一。

4.2 计算图与中间变量：显存到底被谁吃了？

我习惯把计算图想象成一个「有向无环图」。每个节点是一个张量操作，每条边是数据流动。

你想想看，一个简单的线性层 y = Wx + b，中间变量就有：

  • Wx 的乘积结果
  • Wx + b 的加法结果
  • 激活函数的输入和输出

这些中间变量，在反向传播完成之前，一个都不能丢。

我的经验：一个典型的 Transformer 层，中间变量占的显存大约是参数显存的 3-5 倍。层数越深，这个比例越大。

咱们画个图看看，计算图到底长什么样：

计算图与中间变量示意 输入 x 权重 W Wx 中间变量 偏置 b Wx+b 输出 y 前向计算 反向传播时需保留 中间变量区域（反向前必须保留）

看到没？中间变量就是那个粉色的椭圆。它在前向时产生，在反向时被消费。如果网络很深，这些中间变量层层叠加，显存就上去了。

4.3 梯度累积与显存：小 batch 也能跑大模型

我遇到过很多情况：显存只够跑 batch_size=2，但模型需要 batch_size=32 才能收敛。怎么办？

答案就是梯度累积。

梯度累积的原理：

  • 不更新参数，只累积梯度
  • 每跑一个 micro batch，调用 .backward() 计算梯度
  • 梯度自动累加到 .grad 中
  • 累积 N 个 micro batch 后，调用 optimizer.step() 更新参数
  • 然后手动清零梯度 optimizer.zero_grad()

代码实现很简单：

accumulation_steps = 4  # 等效 batch_size = micro_batch * 4

for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    
    # 除以累积步数，保持 loss 尺度一致
    loss = loss / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

我个人习惯：梯度累积时，记得把 loss 除以累积步数。否则梯度会放大 N 倍，学习率得重新调。我曾经因为这个疏忽，模型训了三天发现 loss 一直在震荡……

但注意，梯度累积不减少中间变量的显存。它只减少参数更新的频率。中间变量在每个 micro batch 前向时依然会创建，反向后释放。

4.4 retain_graph 的影响：什么时候需要它？

这个参数坑过不少人。默认情况下，.backward() 执行完后，计算图会被释放。这是为了省显存。

但有些场景需要多次反向传播，比如：

  • GAN 训练：生成器和判别器需要分别计算梯度
  • 多任务学习：多个 loss 需要各自反向
  • 梯度裁剪前需要先计算梯度范数

这时候就需要 retain_graph=True：

# 第一次反向，保留计算图
loss1.backward(retain_graph=True)

# 第二次反向，可以复用同一个计算图
loss2.backward()

# 注意：第二次不需要 retain_graph，因为这是最后一次使用

警告：retain_graph=True 会让计算图一直留在显存里。如果你有 N 个 loss 需要反向，显存占用会翻 N 倍。我曾经在训练一个三任务模型时，忘了及时释放，显存直接爆了 3 倍。

更好的做法是：能合并就合并，能少用就少用。比如两个 loss 可以加权求和后一次反向：

# 推荐做法：合并 loss
total_loss = loss1 + 0.5 * loss2
total_loss.backward()

这样只需要一次反向，计算图用完就释放，显存省一半。

4.5 实战建议：显存优化 checklist

场景 问题 解决方案
模型太大，batch_size 跑不动 中间变量太多 梯度累积 + 混合精度训练
多 loss 训练 retain_graph 导致显存翻倍 合并 loss 或使用梯度累积
调试时需要查看中间变量 显存占用过高 用 torch.no_grad() 或 detach() 隔离
长序列模型（如 Transformer） 中间变量随序列长度平方增长 使用 checkpointing（梯度检查点）

避坑指南：我曾经在调试一个 100 层的 ResNet 时，发现显存占用异常高。排查了半天，发现是某个中间变量被不小心引用了，导致计算图无法释放。记住：用完的变量及时 del，或者用 .detach() 切断梯度追踪。

最后总结一句话：中间变量是显存占用的隐形杀手，梯度累积是省钱利器，retain_graph 是双刃剑。理解了这三者的关系，你的显存优化就成功了一半。

核心要点回顾：

  • 自动求导依赖计算图，中间变量在反向前必须保留
  • 梯度累积通过多次小 batch 模拟大 batch，不减少中间变量显存
  • retain_graph 保留计算图用于多次反向，但显存翻倍
  • 优先合并 loss，减少 retain_graph 的使用

想要更多资料，联系可可

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