在 PyTorch 中可视化并理解 GPU 内存

你一定很熟悉这条消息 🤬

RuntimeError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 7.93 GiB total capacity; 6.00 GiB already allocated; 14.88 MiB free; 6.00 GiB reserved in total by PyTorch)

虽然很容易看到 GPU 内存已满，但理解原因以及如何解决它可能更具挑战性。在本教程中，我们将逐步介绍如何在训练期间可视化并理解 PyTorch 中的 GPU 内存使用情况。我们还将了解如何估算内存需求并优化 GPU 内存使用情况。

🔎 PyTorch 可视化工具

PyTorch 提供了一个方便的工具用于可视化 GPU 内存使用情况

import torch
from torch import nn

# Start recording memory snapshot history
torch.cuda.memory._record_memory_history(max_entries=100000)

model = nn.Linear(10_000, 50_000, device ="cuda")
for _ in range(3):
    inputs = torch.randn(5_000, 10_000, device="cuda")
    outputs = model(inputs)

# Dump memory snapshot history to a file and stop recording
torch.cuda.memory._dump_snapshot("profile.pkl")
torch.cuda.memory._record_memory_history(enabled=None)

运行此代码会生成一个包含 GPU 内存使用历史记录的 profile.pkl 文件。你可以在以下网址可视化此历史记录：https://pytorch.ac.cn/memory_viz。

通过拖放 profile.pkl 文件，你将看到一个类似这样的图表

让我们将此图表分解为关键部分

1. 模型创建：内存增加 2 GB，对应于模型的大小

   $10{,}000 \times 50{,}000 \text{ weights} + 50{,}000 \text{ biases in } \texttt{float32 }\text{(4 bytes)} \implies (5 \times 10^8) \times 4 \, \text{bytes} = 2 \, \text{GB}.$

   此内存（蓝色部分）在整个执行过程中持续存在。

2. 输入张量创建（第一次循环）：内存增加 200 MB，与输入张量大小匹配

   $5{,}000 \times 10{,}000 \text{ elements in } \texttt{float32 }\text{(4 bytes)} \implies (5 \times 10^7) \times 4 \, \text{bytes} = 0.2 \, \text{GB}.$

3. 正向传播（第一次循环）：输出张量内存增加 1 GB

   $5{,}000 \times 50{,}000 \text{ elements in } \texttt{float32 }\text{(4 bytes)} \implies (25 \times 10^7) \times 4 \, \text{bytes} = 1 \, \text{GB}.$

4. 输入张量创建（第二次循环）：内存再次增加 200 MB，用于新的输入张量。此时，你可能期望步骤 2 中的输入张量被释放。然而它并没有被释放：模型保留了其激活值，因此即使该张量不再赋值给变量 inputs，它仍然被模型的正向传播计算所引用。模型保留其激活值是因为在神经网络中反向传播过程需要这些张量。尝试使用 torch.no_grad() 查看差异。

5. 正向传播（第二次循环）：内存增加 1 GB，用于新的输出张量，计算方式与步骤 3 相同。

6. 释放第一次循环激活值：在第二次循环的正向传播之后，第一次循环（步骤 2）的输入张量可以被释放。模型保留了第一个输入张量的激活值，这些激活值被第二次循环的输入覆盖。一旦第二次循环完成，第一个张量不再被引用，其内存可以被释放。

7. 更新 output：步骤 3 的输出张量重新赋值给变量 output。之前的张量不再被引用并被删除，释放了其内存。

8. 输入张量创建（第三次循环）：与步骤 4 相同。

9. 正向传播（第三次循环）：与步骤 5 相同。

10. 释放第二次循环激活值：步骤 4 的输入张量被释放。

11. 再次更新 output：步骤 5 的输出张量重新赋值给变量 output，释放了之前的张量。

12. 代码执行结束：所有内存被释放。

📊 训练期间的内存可视化

前面的例子被简化了。在实际场景中，我们通常训练复杂的模型，而不是单个线性层。此外，前面的例子没有包括训练过程。在这里，我们将研究在一个真实的大型语言模型（LLM）的完整训练循环中 GPU 内存的行为。

import torch
from transformers import AutoModelForCausalLM

# Start recording memory snapshot history
torch.cuda.memory._record_memory_history(max_entries=100000)

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

for _ in range(3):
    inputs = torch.randint(0, 100, (16, 256), device="cuda")  # Dummy input
    loss = torch.mean(model(inputs).logits)  # Dummy loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# Dump memory snapshot history to a file and stop recording
torch.cuda.memory._dump_snapshot("profile.pkl")
torch.cuda.memory._record_memory_history(enabled=None)

💡 提示： 在进行性能分析时，限制步数。每个 GPU 内存事件都会被记录，文件可能会变得非常大。例如，上述代码生成了一个 8 MB 的文件。

这是此示例的内存配置文件

此图表比前一个示例更复杂，但我们仍然可以逐步分解它。注意三个峰值，每个峰值对应于训练循环的一次迭代。让我们简化图表，使其更容易解释

1. 模型初始化 (model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda"))
   第一步涉及将模型加载到 GPU。模型参数（蓝色）占用内存并保持在那里直到训练结束。

2. 正向传播 (model(inputs))
   在正向传播过程中，激活值（每层的中间输出）被计算并存储在内存中用于反向传播。这些激活值（橙色表示）逐层增长，直到最后一层。损失在橙色区域的峰值处计算。

3. 反向传播 (loss.backward())
   梯度（黄色）在此阶段计算并存储。同时，激活值被丢弃，因为它们不再需要，导致橙色区域缩小。黄色区域表示梯度计算的内存使用情况。

4. 优化器步骤 (optimizer.step())
   梯度用于更新模型的参数。最初，优化器本身被初始化（绿色区域）。此初始化只进行一次。之后，优化器使用梯度更新模型的参数。为了更新参数，优化器会临时存储中间值（红色区域）。更新后，梯度（黄色）和优化器中间值（红色）都被丢弃，释放内存。

至此，一次训练迭代完成。此过程对剩余的迭代重复，从而在图表中产生可见的三个内存峰值。

像这样的训练配置文件通常遵循一致的模式，这使得它们对于估算给定模型和训练循环的 GPU 内存需求非常有用。

📐 估算内存需求

从上面一节来看，估算 GPU 内存需求似乎很简单。所需的总内存应对应于内存配置文件中的最高峰值，这发生在正向传播期间。在这种情况下，内存需求是（蓝色 + 绿色 + 橙色）：$\text{Model Parameters} + \text{Optimizer State} + \text{Activations}$

就这么简单吗？实际上，这里有一个陷阱。配置文件可能因训练设置而异。例如，将批量大小从 16 减少到 2 会改变情况

- inputs = torch.randint(0, 100, (16, 256), device="cuda")  # Dummy input
+ inputs = torch.randint(0, 100, (2, 256), device="cuda")  # Dummy input

现在，最高峰值出现在优化器步骤期间，而不是正向传播期间。在这种情况下，内存需求变为（蓝色 + 绿色 + 黄色 + 红色）：$\text{Model Parameters} + \text{Optimizer State} + \text{Gradients} + \text{Optimizer Intermediates}$

为了概括内存估算，我们需要考虑所有可能的峰值，无论它们发生在正向传播还是优化器步骤期间。$\text{Model Parameters} + \text{Optimizer State} + \max(\text{Gradients} + {\text{Optimizer Intermediates}, \text{Activations}})$

现在我们有了公式，让我们看看如何估算每个组件。

模型参数

模型参数最容易估算。$\text{Model Memory} = N \times P$

其中：

• $N$ 是参数数量。
• $P$ 是精度（以字节为单位，例如 float32 为 4）。

例如，一个具有 15 亿参数且精度为 4 字节的模型需要

在上述示例中，模型大小为：$\text{Model Memory} = 1.5 \times 10^9 \times 4 \, \text{bytes} = 6 \, \text{GB}$

优化器状态

优化器状态所需的内存取决于优化器类型和模型参数。例如，AdamW 优化器为每个参数存储两个动量（一阶和二阶）。这使得优化器状态大小为：$\text{Optimizer State Size} = 2 \times N \times P$

激活值

激活值所需的内存更难估算，因为它包括在正向传播过程中计算的所有中间值。要计算激活内存，我们可以使用正向钩子来测量输出的大小

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda")

activation_sizes = []

def forward_hook(module, input, output):
    """
    Hook to calculate activation size for each module.
    """
    if isinstance(output, torch.Tensor):
        activation_sizes.append(output.numel() * output.element_size())
    elif isinstance(output, (tuple, list)):
        for tensor in output:
            if isinstance(tensor, torch.Tensor):
                activation_sizes.append(tensor.numel() * tensor.element_size())

# Register hooks for each submodule
hooks = []
for submodule in model.modules():
    hooks.append(submodule.register_forward_hook(forward_hook))

# Perform a forward pass with a dummy input
dummy_input = torch.zeros((1, 1), dtype=torch.int64, device="cuda")
model.eval()  # No gradients needed for memory measurement
with torch.no_grad():
    model(dummy_input)

# Clean up hooks
for hook in hooks:
    hook.remove()

print(sum(activation_sizes))  # Output: 5065216

对于 Qwen2.5-1.5B 模型，这会给出每个输入 token 5,065,216 个激活值。要估算输入张量的总激活内存，请使用：$\text{Activation Memory} = A \times B \times L \times P$

其中：

• $A$ 是每个 token 的激活值数量。
• $B$ 是批量大小。
• $L$ 是序列长度。

然而，直接使用这种方法并不总是实用的。理想情况下，我们希望有一种启发式方法来估算激活内存，而无需运行模型。此外，我们可以直观地看到更大的模型具有更多的激活。这引出了一个问题：模型参数数量与激活数量之间是否存在关联？

并非直接相关，因为每个 token 的激活数量取决于模型架构。然而，LLM 倾向于具有相似的结构。通过分析不同的模型，我们观察到参数数量与激活数量之间存在大致的线性关系

这种线性关系使我们能够使用启发式方法估算激活值：$A = 4.6894 \times 10^{-4} \times N + 1.8494 \times 10^{6}$

尽管这只是一个近似值，但它提供了一种无需为每个模型执行复杂计算即可估算激活内存的实用方法。

梯度

梯度更容易估算。梯度所需的内存与模型参数相同：$\text{Gradients Memory} = N \times P$

优化器中间值

在更新模型参数时，优化器会存储中间值。这些值所需的内存与模型参数相同：$\text{Optimizer Intermediates Memory} = N \times P$

总内存

总而言之，训练模型所需的总内存为：$\text{Total Memory} = \text{Model Memory} + \text{Optimizer State} + \max(\text{Gradients}, \text{Optimizer Intermediates}, \text{Activations})$

组件如下

• 模型内存：$N \times P$
• 优化器状态：$2 \times N \times P$
• 梯度：$N \times P$
• 优化器中间值：$N \times P$
• 激活值：$A \times B \times L \times P$，使用启发式方法估算：$A = 4.6894 \times 10^{-4} \times N + 1.8494 \times 10^{6}$

为了让这个计算更容易，我为你创建了一个小工具

🚀 后续步骤

你最初理解内存使用情况的动机很可能是因为有一天你内存不足了。这篇博客为你提供了直接的解决方案来解决这个问题吗？可能没有。然而，现在你对内存使用情况有了更好的理解，并且知道了如何分析它，你就能更好地找到减少内存使用的方法了。

有关优化 TRL 内存使用的具体技巧列表，你可以查看文档的减少内存使用部分。不过，这些技巧不限于 TRL，可以应用于任何基于 PyTorch 的训练过程。

🤝 致谢

感谢 Kashif Rasul 对本博客文章提出的宝贵反馈和建议。