注意
转到末尾 下载完整的示例代码。
通过将优化器步骤融合到反向传播中来节省内存#
创建于：2023 年 10 月 02 日 | 最后更新于：2024 年 01 月 16 日 | 最后验证于：2024 年 11 月 05 日
你好！本教程旨在展示一种通过减少梯度占用的内存来降低训练循环内存占用的方法。假设你有一个模型，并且你对优化内存以避免 内存不足 (OOM) 错误或仅仅是为了从 GPU 中榨取更多性能感兴趣。那么，你可能走运了（如果梯度占用了部分内存并且你不需要进行梯度累积）。我们将探讨以下内容：
在训练或微调循环中什么会占用内存，
如何捕获和可视化内存快照以确定瓶颈，
新的
Tensor.register_post_accumulate_grad_hook(hook)API，最后，在 10 行代码中如何将所有内容整合在一起以实现内存节省。
要运行本教程，您需要：
PyTorch 2.1.0 或更新版本及
torchvision1 个 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]
3%|▎ | 29.5M/1.13G [00:00<00:03, 309MB/s]
6%|▌ | 70.9M/1.13G [00:00<00:02, 382MB/s]
10%|▉ | 114M/1.13G [00:00<00:02, 416MB/s]
14%|█▎ | 158M/1.13G [00:00<00:02, 432MB/s]
17%|█▋ | 201M/1.13G [00:00<00:02, 440MB/s]
21%|██ | 245M/1.13G [00:00<00:02, 446MB/s]
25%|██▍ | 289M/1.13G [00:00<00:02, 449MB/s]
29%|██▊ | 332M/1.13G [00:00<00:01, 451MB/s]
32%|███▏ | 375M/1.13G [00:00<00:01, 452MB/s]
36%|███▌ | 419M/1.13G [00:01<00:01, 453MB/s]
40%|███▉ | 462M/1.13G [00:01<00:01, 450MB/s]
44%|████▎ | 505M/1.13G [00:01<00:01, 451MB/s]
47%|████▋ | 549M/1.13G [00:01<00:01, 452MB/s]
51%|█████ | 592M/1.13G [00:01<00:01, 453MB/s]
55%|█████▍ | 636M/1.13G [00:01<00:01, 454MB/s]
58%|█████▊ | 679M/1.13G [00:01<00:01, 453MB/s]
62%|██████▏ | 722M/1.13G [00:01<00:01, 453MB/s]
66%|██████▌ | 766M/1.13G [00:01<00:00, 454MB/s]
70%|██████▉ | 809M/1.13G [00:01<00:00, 453MB/s]
73%|███████▎ | 852M/1.13G [00:02<00:00, 453MB/s]
77%|███████▋ | 896M/1.13G [00:02<00:00, 453MB/s]
81%|████████ | 939M/1.13G [00:02<00:00, 453MB/s]
85%|████████▍ | 982M/1.13G [00:02<00:00, 453MB/s]
88%|████████▊ | 1.00G/1.13G [00:02<00:00, 454MB/s]
92%|█████████▏| 1.04G/1.13G [00:02<00:00, 453MB/s]
96%|█████████▌| 1.09G/1.13G [00:02<00:00, 453MB/s]
100%|█████████▉| 1.13G/1.13G [00:02<00:00, 453MB/s]
100%|██████████| 1.13G/1.13G [00:02<00:00, 448MB/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 内存可视化工具中的快照。内存时间线是否符合您的预期？
模型参数已在训练步骤之前加载到内存中，因此我们一开始就能看到一部分内存专门用于权重。当我们开始前向传播时，内存会逐渐分配给激活，或者说我们保存以便能够在反向传播中计算梯度的张量。一旦我们开始反向传播，激活会逐渐释放，而梯度的内存则开始累积。
最后，当优化器启动时，它的状态将被惰性初始化，因此我们应该只在第一次训练循环的优化器步骤中看到优化器状态内存逐渐增加。在未来的循环中，优化器内存将保持不变并在原地更新。当调用 zero_grad 时，梯度内存会在每个训练循环结束时相应地被释放。
此训练循环中的内存瓶颈在哪里？换句话说，峰值内存是多少？
峰值内存使用量在优化器步骤期间！请注意，内存此时包括约 1.2GB 的参数、约 1.2GB 的梯度以及约 2.4GB=2*1.2GB 的优化器状态，正如预期。最后的约 1.2GB 来自 Adam 优化器为中间值所需的内存，总计约 6GB 的峰值内存。严格来说，如果您将 Adam(model.parameters(), foreach=False) 设置为 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 设置为 False 就可以节省优化器中间值内存，而无需此技术。
您可能正在正确地思考：如果我们节省了 2.4GB 的内存，为什么峰值内存不是 6GB - 2.4GB = 3.6GB？嗯，峰值已经移动了！峰值现在位于反向传播步骤的开始附近，此时我们仍然在内存中有激活，而在之前，峰值是在优化器步骤期间，此时激活已被释放。因此，约 0.4GB 的差异（占 4.0GB - 3.6GB）是由于激活内存。然后，您可以想象这项技术可以与激活检查点结合使用以获得更多内存收益。
结论#
在本教程中，我们通过新的 Tensor.register_post_accumulate_grad_hook() API 了解了将优化器融合到反向传播步骤中以节省内存的技术，以及何时应用此技术（当梯度内存很重要时）。在此过程中，我们还学习了内存快照，它在内存优化方面非常有用。
脚本总运行时间： (0 分钟 8.892 秒)
