2、PyTorch显存分配器探秘：深入理解PyTorch的缓存分配器工作原理，学会使用torch.cuda.memory_summary()诊断工具。

说实话，很多同学跑深度学习模型时遇到OOM（显存溢出），第一反应就是“显存不够，换卡”。但我在实际项目中排查过上百次OOM问题，发现至少有一半的情况，并不是真的显存不够，而是不会用显存。

PyTorch的显存分配器，就是那个在幕后帮你管理显存的小管家。你每次调用.cuda()或者把张量放到GPU上，都是它在干活。但这个小管家有个“小毛病”——它喜欢囤货。

2.1 缓存分配器：为什么显存只增不减？

先看一个现象。你写个简单的代码：

import torch
a = torch.randn(1000, 1000).cuda()
print(torch.cuda.memory_allocated())  # 当前实际使用的显存
print(torch.cuda.memory_reserved())   # 分配器预留的显存

你可能会发现，memory_reserved比memory_allocated大得多。为什么会这样？

PyTorch的分配器，说白了就是一个“缓存池”。它从CUDA驱动那里申请一大块显存，然后自己切成小块分给你用。当你释放张量时，它并不把显存还给CUDA，而是留在自己的池子里，等下次再用。

我个人习惯把这个机制叫做“显存囤积症”。它这么做是有道理的——每次向CUDA申请显存都要经过内核态切换，开销很大。缓存起来可以加速后续分配。但代价就是，你看到的nvidia-smi显存占用，往往比实际需要的多。

2.2 分配器的工作流程：一张图看懂

我画了一张流程图，帮你理清分配器的内部逻辑。你看完就明白为什么显存会“只增不减”了。

你看，当用户释放张量时，分配器并不会把显存还回去，而是放回缓存池。这就是为什么你明明释放了张量，nvidia-smi显示的显存占用却没变。

2.3 诊断工具：torch.cuda.memory_summary()

好了，理论讲完了。现在说说怎么诊断。我最常用的工具就是torch.cuda.memory_summary()。这玩意儿比nvidia-smi详细得多。

用法很简单：

import torch

# 在训练循环中插入
model.train()
for batch in dataloader:
    outputs = model(batch)
    loss = criterion(outputs, targets)
    loss.backward()
    optimizer.step()
    
    # 打印显存快照
    if batch_idx % 100 == 0:
        print(torch.cuda.memory_summary())

输出会包含这些关键信息：

字段	含义	我的经验
Allocated	当前实际使用的显存	这个值才是模型真正占用的
Reserved	分配器预留的缓存	如果这个值远大于Allocated，说明缓存浪费严重
Free	设备上剩余的空闲显存	注意，这个值可能为0，但实际还有缓存可释放
Segment	分配器向CUDA申请的显存段	每个段默认2MB，大张量会申请大段

2.4 实战：用memory_summary()定位OOM

我曾经接手过一个项目，模型在训练到第3个epoch时总是OOM。代码看起来没问题，没有显式泄漏。我加了memory_summary()后，发现了一个有趣的现象：

# 第1个epoch结束时的显存状态
Allocated: 4.2GB
Reserved:  6.8GB
Free:      1.2GB

# 第2个epoch结束时的显存状态
Allocated: 4.3GB
Reserved:  8.5GB
Free:      0.0GB

看到没？Allocated几乎没变，但Reserved一直在涨。这说明分配器在不断申请新缓存，但旧的缓存没被回收。最终Free变成0，新请求触发OOM。

问题出在哪里？我检查了数据加载部分，发现每个epoch都会创建新的DataLoader实例，里面有一些临时张量没有被正确释放。修复后，显存就稳定了。

2.5 分配器的高级配置

PyTorch的分配器其实提供了几个环境变量，可以调整它的行为。我整理了一下：

环境变量	作用	建议值
PYTORCH_CUDA_ALLOC_CONF	配置分配器策略	max_split_size_mb:128
CUDA_LAUNCH_BLOCKING	同步CUDA操作（调试用）	1（仅调试时开启）

举个例子，如果你发现分配器频繁进行小块显存的分配和释放，可以限制最大分割块大小：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

这样分配器就不会把大块显存切得太碎，减少了碎片化。嗯，这里要注意，这个配置要在导入torch之前设置才有效。

2.6 总结一下

PyTorch的显存分配器，说白了就是一个“缓存池”机制。它帮你省去了频繁向CUDA申请显存的开销，但也带来了显存“虚高”的问题。

• 记住两个指标：memory_allocated是真实占用，memory_reserved是缓存占用。
• 用好诊断工具：memory_summary()能帮你看到分配器内部的细节。
• 遇到OOM别慌：先看看是不是缓存没释放，再考虑换卡。

你想想看，如果你能看懂分配器在干什么，很多显存问题其实都能自己解决。下一章我会讲如何用torch.cuda.set_per_process_memory_fraction限制显存使用上限，这个在生产环境中特别有用。