PyTorch GPU 内存分析（memory profiling）的实践指南：梯度检查点、混...

来源：市场资讯

（来源：DeepHub IMBA）

模型有 2 亿个参数，fp32 精度下理论上只需 800 MB。为什么 24 GB 的 GPU 却满了？原因在于模型参数只是训练期间占用 GPU 内存的七种因素之一。了解这七个原因才能从靠猜转变为凭工程的判断。

GPU 内存的七大消耗项

调用 loss.backward() 和 optimizer.step() 时，GPU 中存储了：

1. 模型参数（Model parameters）——权重本身

2. 梯度（Gradients）——与参数等大，每个参数对应一个

3. 优化器状态（Optimizer state）——Adam 为每个参数额外存储 2 个张量（m 和 v）

4. 激活值（Activations）——每一层的输出，在反向传播时需要保留

5. 输入批次（Input batches）——加载到 GPU 上的数据

6. CUDA 工作区（CUDA workspace）——kernel 临时空间和 cuDNN 选择缓存

7. 内存碎片（Memory fragmentation）——已分配但因块间间隔而无法使用的内存

对于用 Adam 训练的 2 亿参数 fp32 模型：

• 参数：800 MB

• 梯度：800 MB（与参数等大）

• Adam 状态（m 和 v）：1600 MB（参数量的 2 倍）

• 激活值：差异很大，通常为参数量的 2–10 倍

• 输入批次：取决于批量大小

• CUDA 工作区：500 MB–1 GB

• 内存碎片：占总量的 5%–20%

所以保守估计总计，一个"理论上"只需 800 MB 的模型，实际要占用 5–8 GB。

如何测量实际情况

PyTorch 内置了相当精确的内存可见性机制，关键在于知道从哪里看。

 import torch# PyTorch 为张量实际分配的 GPU 内存量allocated = torch.cuda.memory_allocated() / 1024**3  # GB# PyTorch 从 CUDA 预留的内存量（含未使用部分）reserved = torch.cuda.memory_reserved() / 1024**3  # GB# 上次重置以来的峰值分配量peak = torch.cuda.max_memory_allocated() / 1024**3  # GB# 重置峰值计数器 torch.cuda.reset_peak_memory_stats()

allocated 与 reserved 之间的差值就是碎片量。如果 allocated 为 5 GB、reserved 为 8 GB，则有 3 GB 内存是 PyTorch 申请了但无法高效使用的。

 print(torch.cuda.memory_summary())

按分配器内存池（allocator pool）打印完整的内存分类统计，含大小分配对比、当前值与峰值，以及各类别明细。在一步训练后调用，能清楚看出内存去向。

大多数人不知道的杀手级功能

PyTorch 能记录每次内存分配，并以时间线形式将其可视化：

 torch.cuda.memory._record_memory_history(max_entries=100_000)# 执行一步训练output = model(x)  loss = criterion(output, y)  loss.backward()  optimizer.step()# 保存快照torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")   torch.cuda.memory._record_memory_history(enabled=None)

将生成的 pickle 文件上传到 https://pytorch.org/memory_viz，可以看到一个交互式可视化界面，展示每次分配、每次释放，以及触发它们的完整调用栈。

靠这个工具，几分钟内定位到用 print 语句排查要花几天的 OOM bug。

三种管用的优化方法

能测量了，才能优化。以下按影响大小排序：

1、梯度检查点（Gradient Checkpointing）以计算换内存

激活值通常是最大的内存消耗项。梯度检查点在反向传播时重新计算激活值，而不是将其存储下来。

 from torch.utils.checkpoint import checkpointclass MyBlock(nn.Module):      def forward(self, x):          return checkpoint(self._forward, x, use_reentrant=False)    def _forward(self, x):          # 此处为耗时操作         return x

典型节省幅度：激活值内存减少 40%–60%。代价：反向传播速度降低 20%–30%。

2、混合精度训练（Mixed Precision Training）内存减半，精度相近

 from torch.amp import autocast, GradScalerscaler = GradScaler('cuda')with autocast('cuda', dtype=torch.float16):      output = model(x)      loss = criterion(output, y)scaler.scale(loss).backward()  scaler.step(optimizer)   scaler.update()

激活值、梯度和大多数运算使用 fp16（每个值 2 字节，而非 4 字节），参数和优化器状态保持 fp32 以保证数值稳定性。

典型节省幅度：总内存减少 30%–50%。fp16 运算在现代 GPU 上速度更快，训练通常也会随之加速。

3. 优化器选择

Adam 为每个参数额外存储 2 个张量。对于 fp32 精度的 10 亿参数模型，仅优化器状态就占 8 GB。

有几个替代选择：

• SGD with momentum：每个参数额外 1 个张量（Adam 开销的一半）

• AdamW with bnb.optim.AdamW8bit：以 8 位存储优化器状态，内存减少 4 倍，精度损失极小

• Lion：内存占用与 SGD 相当，收敛效果通常接近 Adam

对于超过 10 亿参数的模型，优化器的选择可能直接决定训练是否能在现有硬件上运行。

总结

分布式系统领域有句话：无法测量的东西就无法优化。

大多数 PyTorch 团队完全跳过测量步骤：遇到 OOM 错误缩小批量大小然后继续。但是GPU 内存很贵，如果你分析过实际的内存使用情况，就可以将内存占用减半批量大小就能翻倍，这通常意味着更快的训练和更好的梯度估计。

作者：Aditya