可视化并理解 PyTorch 中的 GPU 内存 – Hugging Face 博客摘要

TL;DR

该文章展示了如何在 PyTorch 中记录和可视化 GPU 内存快照，解释了内存配置文件的每个部分代表什么（模型参数、优化器状态、激活值、梯度、优化器中间量），并给出了估算训练期间峰值内存使用的通用公式。

🔎 PyTorch 可视化工具

PyTorch 提供了一个内置工具来记录和可视化 GPU 内存使用情况。通过在运行代码之前调用 torch.cuda.memory._record_memory_history，随后调用 torch.cuda.memory._dump_snapshot，您可以获得一个 profile.pkl 文件，可在 https://pytorch.org/memory_viz 查看。该可视化工具显示内存分配和释放的时间线。

在简单的 nn.Linear 层示例中，图表显示：

  • 模型创建为权重和偏置分配约 2 GB（float32）。
  • 每个输入张量增加约 200 MB。
  • 每次前向传播输出增加约 1 GB。
  • 之前步骤的激活值会被保留，直到不再需要用于反向传播，之后其内存被释放。
  • 重新分配变量会释放之前引用的张量。

📊 训练期间的内存可视化

对于使用大型语言模型（Qwen/Qwen2.5-1.5B）和 AdamW 优化器的实际训练循环，内存配置文件显示三个峰值，每个训练迭代对应一个峰值。

该配置文件可以按以下方式分解：

  1. 模型初始化 – 模型参数（蓝色）占用内存，并在训练结束前保持分配状态。
  2. 前向传播 – 激活值（橙色）逐层计算并存储；它们在损失计算时达到峰值。
  3. 反向传播 – 梯度（黄色）被计算；激活值被丢弃，导致橙色区域缩小。
  4. 优化器步骤 – 优化器状态（绿色）仅初始化一次；优化器使用梯度更新参数，临时存储优化器中间量（红色）。更新后，梯度和中间量被释放。

该模式在每次迭代中重复，从而产生观察到的峰值。

📐 估算内存需求

峰值内存使用是配置文件中的最高点，可能出现在前向传播或优化器步骤中，这取决于批次大小。

一个同时适用于两种情况的通用表达式是：

Total Memory = Model Memory + Optimizer State + max(Gradients + Optimizer Intermediates, Activations)

其中每个术语的定义如下。

模型参数

模型内存 = N × P

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

以 Qwen2.5-1.5B 为例（1.5 B 参数，float32）： 模型内存 = 1.5 × 10⁹ × 4 字节 = 6 GB。

优化器状态

对于 AdamW，它为每个参数存储两个矩： 优化器状态大小 = 2 × N × P

梯度

梯度内存 = N × P（与模型参数大小相同）。

优化器中间量

优化器中间量内存 = N × P（与模型参数大小相同）。

激活值

激活内存取决于批次大小 (B)、序列长度 (L) 和每个标记的激活数量 (A)。 激活内存 = A × B × L × P

A 可以通过前向钩子直接测量，但该文章提供了一个基于跨模型线性拟合的启发式方法：

A = 4.6894 × 10⁻⁴ × N + 1.8494 × 10⁶

使用此启发式方法，您可以在不运行完整前向传播的情况下估算激活内存。

总内存公式（合并）

将组件公式代入得到：

Total Memory = N×P + 2×N×P + max(N×P, N×P, A×B×L×P)
             = 3×N×P + max(N×P, A×B×L×P)

由于 max 项始终至少为 N×P，表达式可简化为：

Total Memory = 3×N×P + A×B×L×P

当激活值占主导时；否则优化器步骤占主导，总内存为 4×N×P。

该文章包含一个小工具，可输入 N、P、B、L 以获得估算值。

🚀 下一步

理解内存配置文件有助于您寻找降低使用量的方法，例如减小批次大小、使用梯度检查点或切换到混合精度训练。TRL 文档中“减少内存使用”部分的建议广泛适用于任何基于 PyTorch 的训练。

🤝 致谢

感谢 Kashif Rasul 对博客文章的反馈和建议。

Sources