备注

点击此处下载完整示例代码

如何通过融合优化器步骤到反向传播中节省内存 ¶

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

你好！本教程旨在展示一种减少训练循环内存占用的一种方法，通过减少梯度的内存占用。假设你有一个模型，并且你对优化内存以避免 Out of Memory （内存不足错误）或简单地从你的 GPU 中挤出更多内容感兴趣。你可能很幸运（如果梯度占用了你的一部分内存，而你不需要进行梯度累积）。我们将探讨以下内容：

在你的训练或微调循环中占用内存的是什么，
如何捕获和可视化内存快照以确定瓶颈，
新的 Tensor.register_post_accumulate_grad_hook(hook) API，以及最后，
10 行代码实现内存节省，一切如何融合在一起。

运行本教程，您需要：

PyTorch 2.1.0 或更高版本，并使用 torchvision
如果您想在本地上运行内存可视化，则需要 1 个 CUDA GPU。否则，这项技术对任何设备都有同样的好处。

让我们先导入所需的模块和模型。我们将使用 torchvision 中的视觉 Transformer 模型，但您也可以替换为自己的模型。我们还将使用 torch.optim.Adam 作为我们的优化器，但同样，您也可以替换为自己的优化器。

import torch
from torchvision import models
from pickle import dump

model = models.vit_l_16(weights='DEFAULT').cuda()
optimizer = torch.optim.Adam(model.parameters())

现在我们来定义典型的训练循环。在训练时应使用真实图像，但在这个教程中，我们传递的是假输入，并不关心加载任何实际数据。

IMAGE_SIZE = 224

def train(model, optimizer):
  # create our fake image input: tensor shape is batch_size, channels, height, width
  fake_image = torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE).cuda()

  # call our forward and backward
  loss = model.forward(fake_image)
  loss.sum().backward()

  # optimizer update
  optimizer.step()
  optimizer.zero_grad()

训练过程中的内存使用

我们将查看一些内存快照，因此我们应该准备好正确分析它们。通常，训练内存包括：

模型参数（大小 P）

为反向传播保存的激活（大小 A）

梯度，与模型参数大小相同，因此大小 G = P。

优化器状态，与参数大小成比例。在这种情况下，Adam 的状态需要 2 倍的模型参数，因此大小 O = 2P。

中间张量，在整个计算过程中分配。我们现在不必担心它们，因为它们通常很小且短暂。

捕获和可视化内存快照

让我们获取一个内存快照！当你的代码运行时，考虑一下你预期的 CUDA 内存时间线可能是什么样子。

# tell CUDA to start recording memory allocations
torch.cuda.memory._record_memory_history(enabled='all')

# train 3 steps
for _ in range(3):
  train(model, optimizer)

# save a snapshot of the memory allocations
s = torch.cuda.memory._snapshot()
with open(f"snapshot.pickle", "wb") as f:
    dump(s, f)

# tell CUDA to stop recording memory allocations now
torch.cuda.memory._record_memory_history(enabled=None)

现在，在 https://maskerprc.github.io/memory_viz 的 CUDA 内存可视化器中打开快照，通过拖放 snapshot.pickle 文件。内存时间线是否符合你的预期？

snapshot.png loaded into CUDA Memory Visualizer

模型参数在训练步骤之前已经被加载到内存中，所以我们一开始就看到一大块内存用于权重。当我们开始正向传递时，会逐渐为激活分配内存，或者说是我们保存的张量，以便在反向传递中计算梯度。一旦我们开始反向传递，激活会逐渐释放，而梯度内存开始积累。

最后，当优化器开始工作时，其状态将懒加载初始化，因此我们只会在第一个训练循环的优化器步骤中看到优化器状态内存逐渐增加。在未来的循环中，优化器内存将保持并就地更新。梯度内存将在每次训练循环结束时相应地释放，当调用 zero_grad 时。

这个训练循环中的内存瓶颈在哪里？换句话说，峰值内存在哪里？

峰值内存使用发生在优化器步骤！注意此时的内存由大约 1.2GB 的参数、大约 1.2GB 的梯度以及预期的 2.4GB=2*1.2GB 的优化器状态组成。最后的 1.2GB 来自 Adam 优化器，它需要存储中间值，总计达到大约 6GB 的峰值内存。从技术上讲，如果您设置 Adam(model.parameters(), foreach=False) ，可以消除对最后 1.2GB 优化器中间值的需要，这将牺牲运行时间以换取内存。如果您关闭 foreach 运行时优化就能在内存节省方面满足您的需求，那很好，但如果您对如何通过本教程做得更好感兴趣，请继续阅读！我们将很快介绍一种技术，通过消除对大约 1.2GB 梯度内存和优化器中间值内存的需求来降低峰值内存。那么，您预计新的峰值内存会是多少？答案将在下一个快照中揭晓。

声明：此技术并不适用于所有情况

在我们过于兴奋之前，我们必须考虑这种技术是否适用于您的用例。这绝对不是万能药！将优化器步骤融合到反向传播中仅针对减少梯度内存（以及作为副作用也减少了优化器中间变量内存）。因此，梯度占用的内存越大，内存减少的效果就越显著。在我们的例子中，梯度占用了内存的 20%，这是一个相当大的比例！

对于您来说可能并非如此，例如，如果您的权重已经很小（可能是由于应用了 LoRa），那么梯度在您的训练循环中不会占用太多空间，因此收益也就不那么令人兴奋了。在这种情况下，您应该首先尝试其他技术，如激活检查点、分布式训练、量化或减少批量大小。然后，当梯度再次成为瓶颈时，再回到这篇教程！

您还在吗？太好了，让我们介绍 Tensor 的新 register_post_accumulate_grad_hook(hook) API。

`Tensor.register_post_accumulate_grad_hook(hook)` API 和我们的技术¶

我们的技术依赖于不需要在 backward() 期间保存梯度。相反，一旦累积了一个梯度，我们就会立即将优化器应用于相应的参数，并完全丢弃该梯度！这消除了在优化器步骤之前保持大量梯度缓冲区所需的需求。

那么，我们如何解锁更积极地应用优化器的行为呢？在我们的 2.1 版本中，我们增加了一个新的 API torch.Tensor.register_post_accumulate_grad_hook() ，允许我们在 Tensor 的 .grad 字段累积后添加一个钩子。我们将优化器步骤封装在这个钩子中。怎么做？

10 行代码如何实现一切

记得我们从一开始的模型和优化器设置吗？我将在下面将它们注释出来，以免我们浪费资源重新运行代码。

model = models.vit_l_16(weights='DEFAULT').cuda()
optimizer = torch.optim.Adam(model.parameters())

# Instead of having just *one* optimizer, we will have a ``dict`` of optimizers
# for every parameter so we could reference them in our hook.
optimizer_dict = {p: torch.optim.Adam([p], foreach=False) for p in model.parameters()}

# Define our hook, which will call the optimizer ``step()`` and ``zero_grad()``
def optimizer_hook(parameter) -> None:
  optimizer_dict[parameter].step()
  optimizer_dict[parameter].zero_grad()

# Register the hook onto every parameter
for p in model.parameters():
   p.register_post_accumulate_grad_hook(optimizer_hook)

# Now remember our previous ``train()`` function? Since the optimizer has been
# fused into the backward, we can remove the optimizer step and zero_grad calls.
def train(model):
  # create our fake image input: tensor shape is batch_size, channels, height, width
  fake_image = torch.rand(1, 3, IMAGE_SIZE, IMAGE_SIZE).cuda()

  # call our forward and backward
  loss = model.forward(fake_image)
  loss.sum().backward()

  # optimizer update --> no longer needed!
  # optimizer.step()
  # optimizer.zero_grad()

这在我们的样本模型中大约需要 10 行更改，这很棒。然而，对于真实模型来说，将优化器替换为优化器字典可能是一个相当侵入性的更改，尤其是对于那些在整个训练周期中使用 `LRScheduler` 或操作优化器配置的人来说。处理这些更改的 API 将更加复杂，可能需要将更多配置移动到全局状态，但这并不应该是不可能的。话虽如此，PyTorch 的下一步是使这个 API 更容易与 LRSchedulers 和其他您已经习惯的功能兼容。

但让我回到说服您这个技术值得做的观点。我们将咨询我们的朋友，内存快照。

# delete optimizer memory from before to get a clean slate for the next
# memory snapshot
del optimizer

# tell CUDA to start recording memory allocations
torch.cuda.memory._record_memory_history(enabled='all')

# train 3 steps. note that we no longer pass the optimizer into train()
for _ in range(3):
  train(model)

# save a snapshot of the memory allocations
s = torch.cuda.memory._snapshot()
with open(f"snapshot-opt-in-bwd.pickle", "wb") as f:
    dump(s, f)

# tell CUDA to stop recording memory allocations now
torch.cuda.memory._record_memory_history(enabled=None)

是的，花点时间将您的快照拖入 CUDA 内存可视化器。

几个主要观察结果：

没有更多的优化步骤了！对吧…我们把它融合到了反向过程中。
同样，反向过程耗时更长，中间变量分配也更随机。这是预期的，因为优化步骤需要中间变量。
最重要的是！峰值内存降低了！现在大约是 4GB（我希望这接近你之前的预期）。

注意，与之前相比，不再有大量内存分配给梯度，这节省了大约 1.2GB 的内存。相反，我们在计算完每个梯度后非常快地释放了它们，将优化步骤尽可能提前。太棒了！顺便说一下，其他大约 1.2GB 的内存节省来自于将优化器拆分为每个参数的优化器，因此中间变量成比例缩小。这个细节不如梯度内存节省重要，因为即使不使用这种技术，也可以通过将 foreach=False 关闭来获得优化器中间变量的节省。

你可能正确地想知道：如果我们节省了 2.4GB 的内存，为什么峰值内存不是 6GB - 2.4GB = 3.6GB 呢？嗯，峰值移动了！现在峰值出现在反向步骤的开始阶段，那时我们仍然有激活在内存中，而之前，峰值出现在优化器步骤中，那时激活已经被释放。因此，解释了大约 0.4GB 的差异，即大约 4.0GB - 大约 3.6GB 的差异，是由于激活内存。然后可以想象，这种技术可以与激活检查点结合使用，以获得更多的内存优势。

结论 ¶

在本教程中，我们学习了将优化器融合到反向步骤中的内存节省技术，以及何时应用这种技术（当梯度内存很大时）。在这个过程中，我们还学习了内存快照，这在内存优化中通常很有用。

脚本总运行时间：（0 分钟 0.000 秒）

由 Sphinx-Gallery 生成的画廊