评价此页
★ ★ ★ ★ ★
intermediate/optimizer_step_in_backward_tutorial

注意

前往页面底部下载完整示例代码。

通过将优化器步骤融合到反向传播中来节省内存#

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

你好！本教程旨在展示一种通过减少梯度所占内存来降低训练循环内存占用量的方法。假设你有一个模型，并且正在寻找优化内存的方法，以避免内存不足 (OOM) 错误，或者仅仅是想从 GPU 中压榨出更多空间。好吧，你可能很幸运（如果梯度占用了你一部分内存，且你不需要进行梯度累积的话）。我们将探讨以下内容：

  1. 训练或微调循环中哪些部分占用了内存，

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

  3. 新的 Tensor.register_post_accumulate_grad_hook(hook) API，最后是，

  4. 如何用 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())

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%|          | 11.4M/1.13G [00:00<00:10, 119MB/s]
  2%|▏         | 23.8M/1.13G [00:00<00:09, 124MB/s]
  4%|▍         | 50.6M/1.13G [00:00<00:05, 195MB/s]
  6%|▋         | 73.4M/1.13G [00:00<00:05, 212MB/s]
  9%|▉         | 105M/1.13G [00:00<00:04, 253MB/s]
 12%|█▏        | 144M/1.13G [00:00<00:03, 308MB/s]
 16%|█▌        | 186M/1.13G [00:00<00:02, 351MB/s]
 20%|█▉        | 229M/1.13G [00:00<00:02, 383MB/s]
 23%|██▎       | 268M/1.13G [00:00<00:02, 390MB/s]
 27%|██▋       | 310M/1.13G [00:01<00:02, 405MB/s]
 30%|███       | 354M/1.13G [00:01<00:02, 420MB/s]
 34%|███▍      | 397M/1.13G [00:01<00:01, 430MB/s]
 38%|███▊      | 440M/1.13G [00:01<00:01, 437MB/s]
 42%|████▏     | 483M/1.13G [00:01<00:01, 441MB/s]
 45%|████▌     | 526M/1.13G [00:01<00:01, 430MB/s]
 49%|████▉     | 569M/1.13G [00:01<00:01, 436MB/s]
 53%|█████▎    | 610M/1.13G [00:01<00:01, 415MB/s]
 56%|█████▌    | 650M/1.13G [00:01<00:01, 410MB/s]
 60%|█████▉    | 694M/1.13G [00:01<00:01, 423MB/s]
 63%|██████▎   | 737M/1.13G [00:02<00:01, 432MB/s]
 67%|██████▋   | 780M/1.13G [00:02<00:00, 439MB/s]
 71%|███████   | 824M/1.13G [00:02<00:00, 444MB/s]
 75%|███████▍  | 867M/1.13G [00:02<00:00, 447MB/s]
 78%|███████▊  | 911M/1.13G [00:02<00:00, 449MB/s]
 82%|████████▏ | 954M/1.13G [00:02<00:00, 451MB/s]
 86%|████████▌ | 998M/1.13G [00:02<00:00, 452MB/s]
 90%|████████▉ | 1.02G/1.13G [00:02<00:00, 453MB/s]
 93%|█████████▎| 1.06G/1.13G [00:02<00:00, 453MB/s]
 97%|█████████▋| 1.10G/1.13G [00:02<00:00, 453MB/s]
100%|██████████| 1.13G/1.13G [00:03<00:00, 403MB/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)

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

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），那么梯度在你的训练循环中不会占用太多空间，带来的收益也就没那么令人兴奋了。在这种情况下，你应该首先尝试其他技术，如激活检查点 (activations checkpointing)、分布式训练、量化或减少批次大小。然后，当梯度再次成为瓶颈的一部分时，再回到这个教程！

还在吗？太棒了，让我们介绍 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()，它允许我们在张量的 .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``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？嗯，峰值已经转移了！峰值现在出现在反向传播步骤的开始附近，当时我们还在内存中保留了激活值；而之前，峰值出现在优化器步骤期间，此时激活值已经被释放了。因此，约 4.0GB - 约 3.6GB 之间的约 0.4GB 差异是由激活值内存造成的。人们可以想象，这项技术可以与激活检查点结合使用，以获得更多的内存收益。

结论#

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

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

在此页面上
PyTorch 库