目录
  • 教程 >
  • 如何通过将优化器步骤融入反向传播来节省内存
快捷方式
intermediate/optimizer_step_in_backward_tutorial

注意

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

如何通过将优化器步骤融入反向传播来节省内存

创建于：2023 年 10 月 02 日 | 最后更新：2024 年 1 月 16 日 | 最后验证：2024 年 11 月 05 日

大家好！本教程旨在展示一种通过减少梯度占用的内存量来减少训练循环内存占用量的方法。假设您有一个模型，并且您对优化内存以避免 Out of Memory (OOM) 错误或仅仅是从 GPU 中榨取更多性能的方法感兴趣。那么，您_可能_很幸运（如果梯度占用了一部分内存，并且您不需要进行梯度累积）。我们将探讨以下内容

  1. 训练或微调循环期间哪些因素占用内存，

  2. 如何捕获和可视化内存快照以确定瓶颈，

  3. 新的 Tensor.register_post_accumulate_grad_hook(hook) API，以及最终，

  4. 所有内容如何在 10 行代码中组合在一起以实现内存节省。

要运行本教程，您将需要

  • PyTorch 2.1.0 或更新版本，带有 torchvision

  • 1 个 CUDA GPU（如果您想在本地运行内存可视化）。否则，这项技术在任何设备上都将获得类似的好处。

让我们从导入所需的模块和模型开始。我们将使用 torchvision 中的 vision 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())

Downloading: "https://download.pytorch.org/models/vit_l_16-852ce7e3.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/vit_l_16-852ce7e3.pth

  0%|          | 0.00/1.13G [00:00<?, ?B/s]
  1%|          | 7.88M/1.13G [00:00<00:14, 82.2MB/s]
  2%|1         | 22.2M/1.13G [00:00<00:09, 122MB/s]
  3%|3         | 39.8M/1.13G [00:00<00:07, 150MB/s]
  5%|4         | 57.4M/1.13G [00:00<00:07, 163MB/s]
  6%|6         | 75.0M/1.13G [00:00<00:06, 171MB/s]
  8%|7         | 92.6M/1.13G [00:00<00:06, 175MB/s]
  9%|9         | 110M/1.13G [00:00<00:06, 178MB/s]
 11%|#1        | 128M/1.13G [00:00<00:06, 180MB/s]
 13%|#2        | 145M/1.13G [00:00<00:05, 180MB/s]
 14%|#3        | 162M/1.13G [00:01<00:05, 179MB/s]
 15%|#5        | 180M/1.13G [00:01<00:05, 179MB/s]
 17%|#6        | 197M/1.13G [00:01<00:05, 179MB/s]
 18%|#8        | 214M/1.13G [00:01<00:06, 160MB/s]
 20%|#9        | 231M/1.13G [00:01<00:05, 165MB/s]
 21%|##1       | 248M/1.13G [00:01<00:05, 170MB/s]
 23%|##2       | 266M/1.13G [00:01<00:05, 173MB/s]
 24%|##4       | 283M/1.13G [00:01<00:05, 175MB/s]
 26%|##5       | 300M/1.13G [00:01<00:05, 177MB/s]
 27%|##7       | 318M/1.13G [00:01<00:04, 181MB/s]
 29%|##8       | 336M/1.13G [00:02<00:04, 183MB/s]
 31%|###       | 354M/1.13G [00:02<00:04, 184MB/s]
 32%|###2      | 372M/1.13G [00:02<00:04, 186MB/s]
 34%|###3      | 390M/1.13G [00:02<00:04, 186MB/s]
 35%|###5      | 408M/1.13G [00:02<00:04, 187MB/s]
 37%|###6      | 426M/1.13G [00:02<00:04, 187MB/s]
 38%|###8      | 444M/1.13G [00:02<00:04, 187MB/s]
 40%|###9      | 462M/1.13G [00:02<00:03, 188MB/s]
 41%|####1     | 480M/1.13G [00:02<00:03, 188MB/s]
 43%|####2     | 498M/1.13G [00:02<00:03, 188MB/s]
 44%|####4     | 516M/1.13G [00:03<00:03, 188MB/s]
 46%|####6     | 534M/1.13G [00:03<00:03, 188MB/s]
 48%|####7     | 552M/1.13G [00:03<00:03, 188MB/s]
 49%|####9     | 570M/1.13G [00:03<00:03, 188MB/s]
 51%|#####     | 588M/1.13G [00:03<00:03, 188MB/s]
 52%|#####2    | 607M/1.13G [00:03<00:03, 189MB/s]
 54%|#####3    | 625M/1.13G [00:03<00:02, 189MB/s]
 55%|#####5    | 643M/1.13G [00:03<00:02, 189MB/s]
 57%|#####6    | 661M/1.13G [00:03<00:02, 188MB/s]
 58%|#####8    | 679M/1.13G [00:03<00:02, 186MB/s]
 60%|#####9    | 696M/1.13G [00:04<00:02, 185MB/s]
 61%|######1   | 714M/1.13G [00:04<00:02, 184MB/s]
 63%|######3   | 732M/1.13G [00:04<00:02, 183MB/s]
 65%|######4   | 749M/1.13G [00:04<00:02, 182MB/s]
 66%|######6   | 767M/1.13G [00:04<00:02, 182MB/s]
 68%|######7   | 784M/1.13G [00:04<00:02, 182MB/s]
 69%|######9   | 801M/1.13G [00:04<00:02, 182MB/s]
 71%|#######   | 819M/1.13G [00:04<00:01, 181MB/s]
 72%|#######2  | 836M/1.13G [00:04<00:01, 181MB/s]
 74%|#######3  | 854M/1.13G [00:04<00:01, 182MB/s]
 75%|#######5  | 871M/1.13G [00:05<00:01, 181MB/s]
 77%|#######6  | 888M/1.13G [00:05<00:01, 182MB/s]
 78%|#######8  | 906M/1.13G [00:05<00:01, 182MB/s]
 79%|#######9  | 923M/1.13G [00:05<00:01, 181MB/s]
 81%|########  | 940M/1.13G [00:05<00:01, 181MB/s]
 82%|########2 | 958M/1.13G [00:05<00:01, 181MB/s]
 84%|########3 | 975M/1.13G [00:05<00:01, 181MB/s]
 85%|########5 | 992M/1.13G [00:05<00:00, 181MB/s]
 87%|########6 | 0.99G/1.13G [00:05<00:00, 181MB/s]
 88%|########8 | 1.00G/1.13G [00:05<00:00, 181MB/s]
 90%|########9 | 1.02G/1.13G [00:06<00:00, 181MB/s]
 92%|#########1| 1.04G/1.13G [00:06<00:00, 183MB/s]
 93%|#########3| 1.06G/1.13G [00:06<00:00, 185MB/s]
 95%|#########4| 1.07G/1.13G [00:06<00:00, 186MB/s]
 96%|#########6| 1.09G/1.13G [00:06<00:00, 186MB/s]
 98%|#########7| 1.11G/1.13G [00:06<00:00, 187MB/s]
 99%|#########9| 1.13G/1.13G [00:06<00:00, 187MB/s]
100%|##########| 1.13G/1.13G [00:06<00:00, 181MB/s]

现在让我们定义我们典型的训练循环。在训练时您应该使用真实图像，但为了本教程的目的，我们传入的是虚假输入，并且不担心加载任何实际数据。

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)

现在在 CUDA 内存可视化工具中打开快照，网址为 https://pytorch.ac.cn/memory_viz，方法是拖放 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 字段累积后在其上添加一个 hook。我们将把优化器步骤封装到这个 hook 中。如何封装？

所有内容如何在 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``s 或在整个训练周期中操作优化器配置的人来说。使用这些更改来解决此 API 将更加复杂，并且可能需要将更多配置移动到全局状态，但这并非不可能。也就是说，PyTorch 的下一步是使此 API 更容易与 LRScheduler 以及您已经习惯的其他功能一起采用。

但让我回到说服您这项技术值得尝试。我们将咨询我们的朋友，内存快照。

# 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 内存可视化工具中。

snapshot.png loaded into CUDA Memory Visualizer
几个主要观察结果

  1. 不再有优化器步骤了！是的……我们将其融入了反向传播中。

  2. 同样，反向传播拖得更长，并且有更多随机分配用于中间值。这是预期的，因为优化器步骤需要中间值。

  3. 最重要的是！峰值内存更低了！现在约为 4GB（我希望这与您之前的期望值接近）。

请注意，与之前相比，不再有任何为梯度分配的大块内存，节省了约 1.2GB 的内存。相反，我们通过尽可能提前移动优化器步骤，在计算出每个梯度后立即快速释放它们。哇！顺便说一句，另外约 1.2GB 的内存节省来自将优化器分解为每个参数优化器，因此中间值也相应缩小了。这个细节不如梯度内存节省重要，因为您只需关闭 foreach=False 即可获得优化器中间值节省，而无需使用此技术。

您可能会正确地想知道：如果我们节省了 2.4GB 的内存，为什么峰值内存不是 6GB - 2.4GB = 3.6GB？好吧，峰值已经移动了！峰值现在靠近反向传播步骤的开始，此时我们仍然在内存中保留激活值，而在之前，峰值是在优化器步骤期间，此时激活值已被释放。约 0.4GB 的差异（约 4.0GB - 约 3.6GB）因此是由于激活值内存造成的。然后可以想象，这项技术可以与激活值检查点结合使用，以获得更多的内存优势。

结论

在本教程中，我们学习了通过新的 Tensor.register_post_accumulate_grad_hook() API 将优化器融入反向传播步骤的内存节省技术，以及何时应用此技术（当梯度内存非常重要时）。在此过程中，我们还了解了内存快照，这在内存优化中通常很有用。

脚本的总运行时间： ( 0 分钟 16.185 秒)

由 Sphinx-Gallery 生成的图库

评价本教程

© Copyright 2024, PyTorch。

使用 Sphinx 构建，主题由 theme 提供，托管于 Read the Docs。

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源

为了分析流量并优化您的体验，我们在本网站上使用 Cookie。通过点击或导航，您同意允许我们使用 Cookie。作为本网站的当前维护者，Facebook 的 Cookie 政策适用。了解更多信息，包括关于可用控件：Cookie 政策。