如何在 PyTorch 中可视化和理解 GPU 显存占用
Visualize and understand GPU memory in PyTorch
Hugging Face 博客给出用 PyTorch 的 torch.cuda.memory 记录并可视化 GPU 显存占用的完整方法，把 profile.pkl 拖入 pytorch.org/memory_viz 即可看到显存随时间变化的图形。
用 PyTorch 自带的内存快照工具逐步拆解训练显存构成，并给出可套用的估算公式。
你一定很熟悉这条消息 🤬：
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)
运行此代码会生成一个 profile.pkl 文件，其中包含执行期间 GPU 内存使用的历史记录。你可以在以下网址查看此历史记录：https://pytorch.org/memory_viz。
通过拖放你的 profile.pkl 文件，你将看到类似这样的图表：
让我们将这张图分解为几个关键部分：
模型创建：内存增加 2 GB，与模型大小相对应：
10,000×50,000 权重+50,000 偏置 以 float32 （4 字节）表示 ⟹ (5×108)×4 字节=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}。
这部分内存（蓝色）在整个执行过程中持续存在。
输入张量创建（第1次循环）：内存增加 200 MB，与输入张量大小相匹配：
5,000×10,000 元素 以 float32 （4 字节）表示 ⟹ (5×107)×4 字节=0.2 GB。 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}。
前向传播（第1次循环）：输出张量使内存增加 1 GB：
5,000×50,000 元素 以 float32 （4 字节）表示 ⟹ (25×107)×4 字节=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}。
输入张量创建（第2次循环）：新的输入张量使内存增加 200 MB。此时，你可能期望第2步的输入张量被释放。然而并没有：模型保留了它的激活值，因此即使该张量不再赋值给变量
inputs，它仍被模型的前向传播计算所引用。模型保留激活值，是因为这些张量是神经网络反向传播过程所必需的。尝试使用torch.no_grad()来查看差异。前向传播（第2次循环）：新的输出张量使内存增加 1 GB，计算方式同第3步。
释放第1次循环的激活值：在第二次循环的前向传播之后，第一次循环（第2步）的输入张量可以被释放。持有第一个输入张量的模型激活值被第二次循环的输入覆盖。一旦第二次循环完成，第一个张量不再被引用，其内存可以被释放。
更新
output：第3步的输出张量被重新赋值给变量output。之前的张量不再被引用并被删除，释放其内存。输入张量创建（第3次循环）：同第4步。
前向传播（第3次循环）：同第5步。
释放第2次循环的激活值：第4步的输入张量被释放。
再次更新
output：第5步的输出张量被重新赋值给变量output，释放之前的张量。代码执行结束：所有内存均已释放。
📊 训练过程中的内存可视化
之前的例子经过了简化。在实际场景中，我们通常训练的是复杂模型，而非单一的线性层。此外，之前的例子并未包含训练过程。这里，我们将考察在真实的大型语言模型（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 的文件。
以下是此示例的内存概况图：
这个图表比之前的例子更复杂，但我们仍可以一步步地拆解它。注意那三个尖峰，每个尖峰对应训练循环的一次迭代。让我们简化图表，以便更容易解读：
模型初始化（
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B").to("cuda")）：
第一步涉及将模型加载到 GPU 上。模型参数（蓝色部分）占用内存，并一直保留到训练结束。前向传播（
model(inputs)）：
在前向传播过程中，激活值（每层的中间输出）被计算并存储在内存中，以供反向传播使用。这些以橙色表示的激活值逐层增长，直至最后一层。损失值在橙色区域的峰值处计算得出。反向传播（
loss.backward()）：
在此阶段计算并存储梯度（黄色部分）。同时，不再需要的激活值被丢弃，导致橙色区域收缩。黄色区域代表了梯度计算的内存使用量。优化器步骤（
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
现在，最高峰值出现在优化器步骤期间，而不是前向传播期间。在这种情况下，内存需求变为（蓝色 + 绿色 + 黄色 + 红色）： Model Parameters+Optimizer State+Gradients+Optimizer Intermediates \text{Model Parameters} + \text{Optimizer State} + \text{Gradients} + \text{Optimizer Intermediates}
为了推广内存估算，我们需要考虑所有可能的峰值，无论它们出现在前向传播还是优化器步骤期间。 Model Parameters+Optimizer State+max(Gradients+Optimizer Intermediates,Activations) \text{Model Parameters} + \text{Optimizer State} + \max(\text{Gradients} + {\text{Optimizer Intermediates}, \text{Activations}})
现在我们有了公式，让我们看看如何估算每个组成部分。
模型参数
模型参数是最容易估算的。 Model Memory=N×P \text{Model Memory} = N \times P
其中：
- N N 是参数的数量。
- P P 是精度（以字节为单位，例如，对于
float32为 4）。
例如，一个具有 15 亿参数且精度为 4 字节的模型需要：
在上面的示例中，模型大小为： Model Memory=1.5×109×4 bytes=6 GB \text{Model Memory} = 1.5 \times 10^9 \times 4 \, \text{bytes} = 6 \, \text{GB}
优化器状态
优化器状态所需的内存取决于优化器类型和模型参数。例如，AdamW 优化器为每个参数存储两个矩（一阶和二阶）。这使得优化器状态大小为：
Optimizer State Size=2×N×P \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 个激活值。要估算输入张量的总激活内存，请使用： Activation Memory=A×B×L×P \text{Activation Memory} = A \times B \times L \times P
其中：
- A A 是每个 token 的激活值数量。
- B B 是批次大小。
- L L 是序列长度。
然而，直接使用这种方法并不总是实际的。理想情况下，我们希望有一个启发式方法来估算激活内存，而无需运行模型。此外，我们可以直观地看到更大的模型有更多的激活值。这就引出了一个问题：模型参数的数量和激活值的数量之间是否存在联系？
并非直接相关，因为每个 token 的激活值数量取决于模型架构。然而，LLMs 往往具有相似的结构。通过分析不同的模型，我们观察到参数数量与激活值数量之间存在粗略的线性关系：
这种线性关系允许我们使用启发式方法估算激活值： A=4.6894×10−4×N+1.8494×106 A = 4.6894 \times 10^{-4} \times N + 1.8494 \times 10^{6}
尽管这是一种近似方法，但它提供了一种实用的方式，无需为每个模型执行复杂计算即可估算激活内存。
梯度
梯度更容易估算。梯度所需的内存与模型参数相同： Gradients Memory=N×P \text{Gradients Memory} = N \times P
优化器中间值
更新模型参数时，优化器会存储中间值。这些值所需的内存与模型参数相同： Optimizer Intermediates Memory=N×P \text{Optimizer Intermediates Memory} = N \times P
总内存
总结一下，训练模型所需的总内存为： Total Memory=Model Memory+Optimizer State+max(Gradients,Optimizer Intermediates,Activations) \text{Total Memory} = \text{Model Memory} + \text{Optimizer State} + \max(\text{Gradients}, \text{Optimizer Intermediates}, \text{Activations})
包含以下组成部分：
- 模型内存：N×P N \times P
- 优化器状态：2×N×P 2 \times N \times P
- 梯度：N×P N \times P
- 优化器中间值：N×P N \times P
- 激活值：A×B×L×P A \times B \times L \times P ，使用启发式方法估算：A=4.6894×10−4×N+1.8494×106 A = 4.6894 \times 10^{-4} \times N + 1.8494 \times 10^{6}
为了让这个计算更简单，我为你创建了一个小工具：
🚀 后续步骤
你最初想要了解内存使用情况，很可能是因为某天你的内存耗尽了。这篇博客是否直接给出了解决这个问题的方案？可能没有。不过，既然你现在已经更好地理解了内存使用的工作原理以及如何对其进行性能分析，你就更有能力找到减少内存使用的方法。
关于优化 TRL 内存使用的具体技巧列表，你可以查看文档中的 减少内存使用 部分。不过，这些技巧并不局限于 TRL，可以应用于任何基于 PyTorch 的训练过程。
🤝 致谢
感谢 Kashif Rasul 对这篇博客文章提出的宝贵反馈和建议。
来源：Hugging Face Blog · hf-mirror.com