3、PyTorch显存管理机制：缓存分配器、显存池与碎片化问题

聊到显存溢出，很多人第一反应就是“模型太大，显存放不下”。但我在实际项目中踩过不少坑，发现很多时候并不是模型真的装不下，而是PyTorch的显存管理机制在“捣鬼”。说白了，你看到的显存占用，和模型实际需要的显存，中间差了好几个“隐形消耗”。

这一节，我们就来扒一扒PyTorch的显存管理到底是怎么运作的。理解了它，你才能精准定位显存溢出的根因。

3.1 PyTorch的缓存分配器（Caching Allocator）

PyTorch默认使用一个叫做 缓存分配器 的东西来管理显存。它不像Python的垃圾回收那样用完就扔，而是把释放的显存块缓存起来，留给后续的Tensor分配使用。

为什么要这么做？因为每次向GPU申请显存（cudaMalloc）和释放显存（cudaFree）都是很昂贵的操作。频繁调用会严重拖慢训练速度。所以PyTorch的策略是：先缓存，再复用。

我个人习惯用 torch.cuda.memory_summary() 来查看当前分配器的状态。这个命令会告诉你：当前分配了多少显存、缓存了多少、峰值用了多少。嗯，这里要注意，你看到的 allocated 是实际Tensor占用的，而 cached 是分配器手里攥着没还给GPU的。

import torch

# 查看当前显存分配器状态
print(torch.cuda.memory_summary())

# 更简洁的方式
allocated = torch.cuda.memory_allocated()  # 实际占用的
cached = torch.cuda.memory_reserved()      # 缓存的总量
print(f"Allocated: {allocated / 1024**3:.2f} GB")
print(f"Cached: {cached / 1024**3:.2f} GB")

我在项目中遇到过这样一个情况：明明只用了6GB显存，但 nvidia-smi 显示占用了10GB。多出来的4GB就是缓存分配器“囤”的货。这时候如果你去排查显存溢出，只看 nvidia-smi 会被误导。

3.2 显存池（Memory Pool）的工作原理

显存池是缓存分配器的底层实现。它把显存划分成不同大小的块，按大小分桶管理。PyTorch内部维护了多个“桶”（bins），每个桶里放着大小相近的空闲块。

分配流程大致是这样的：

1. 请求分配一块 N 字节的显存
2. 分配器在对应大小的桶里找空闲块
3. 找到大小 >= N 的块，直接返回（可能拆分，剩余部分放回小桶）
4. 找不到，就向GPU申请一大块（通常是2MB的倍数），切分后分配

释放流程：

1. Tensor被销毁，显存块标记为空闲
2. 分配器检查相邻块是否也是空闲，如果是，合并成更大的块
3. 合并后的块放回对应大小的桶里

你想想看，这种设计在大多数场景下是高效的。但有一个致命问题——显存碎片化。

3.3 显存碎片化问题

碎片化，说白了就是显存被切成了很多小块，虽然总空闲量够，但找不到一块连续的大块来满足分配请求。这就像你硬盘里有很多小文件碎片，虽然剩余空间还有50GB，但存一个40GB的电影却提示“磁盘空间不足”。

为什么会这样？我举个例子：

• 你分配了一个4GB的Tensor，然后又分配了8个512MB的Tensor
• 你把4GB的Tensor释放了，显存池里多了一个4GB的空闲块
• 接着你分配一个3.5GB的Tensor，分配器从4GB块里切出3.5GB，剩下0.5GB的小块
• 然后你又释放了那8个512MB中的4个，池子里多了4个512MB的块
• 现在你有一个0.5GB的块，和4个512MB的块，它们可能不连续
• 这时候你想分配一个2GB的Tensor，虽然总空闲量有2.5GB，但没有一个连续块 >= 2GB

嗯，这就是典型的显存碎片化。在训练大模型时，这种问题尤其常见。我曾经在训练一个BERT-large模型时，明明显存还有3GB空闲，但就是报OOM。查了半天，发现是碎片化导致的。

3.4 如何诊断碎片化

PyTorch提供了一些工具来诊断显存碎片化程度：

# 查看显存块的详细信息
print(torch.cuda.memory_summary(abbreviated=False))

# 查看每个桶的空闲块数量
stats = torch.cuda.memory_stats()
for key, value in stats.items():
    if 'segment' in key or 'block' in key:
        print(f"{key}: {value}")

关键指标：

指标	含义	健康值
num_alloc_retries	分配重试次数（因碎片导致）	接近0
max_split_size	最大空闲块大小	接近总显存
num_segments	显存段数量	越小越好
segment_split_ratio	碎片化比例	< 0.3

3.5 实战建议：减少碎片化

结合我的经验，给你几个实用的建议：

1. 使用固定大小的Tensor：尽量让所有中间Tensor的大小保持一致，避免大小频繁变化。
2. 复用Tensor池：对于循环中反复使用的Tensor，创建一个池子来复用，而不是每次新建。
3. 调整缓存分配器参数：通过环境变量 PYTORCH_CUDA_ALLOC_CONF 调整分配器行为。比如设置 max_split_size_mb:128 可以限制最大拆分块大小，减少小碎片的产生。

# 设置环境变量，限制最大拆分块为128MB
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

# 或者使用更激进的策略：启用expandable_segments（PyTorch 2.0+）
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

我记得有一次，一个同事的模型在训练到第100个batch时突然OOM。排查后发现是某个中间Tensor的大小随着输入变化而变化，导致显存碎片化越来越严重。我们改成固定大小后，问题就解决了。

最后说一句：显存管理不是玄学，是工程。理解了PyTorch的缓存分配器和显存池机制，你就能在遇到OOM时，快速判断是“真不够”还是“碎片化”，从而对症下药。