2. PyTorch显存管理机制

好，咱们今天聊聊PyTorch的显存管理。说实话，很多同学跑模型遇到OOM（Out Of Memory），第一反应就是“显存不够，换卡”。但我在实际项目中踩过不少坑，发现很多时候不是显存真的不够，而是你没用好PyTorch的显存管理机制。

你想想看，PyTorch底层是怎么跟GPU打交道的？它不会每次需要显存都去跟驱动申请一次，那样太慢了。它有一套自己的分配器和缓存策略。搞懂这套机制，你就能在有限的显存里塞下更大的模型。

2.1 torch.cuda的核心机制

torch.cuda 是PyTorch与GPU交互的桥梁。它负责管理设备、流（Stream）、事件（Event）以及最关键的——显存。

我个人习惯，在代码开头先检查一下GPU可用性：

import torch
print(torch.cuda.is_available())      # True or False
print(torch.cuda.device_count())      # 有几张卡
print(torch.cuda.get_device_name(0))  # 卡的名字

嗯，这里要注意：torch.cuda 默认使用当前设备（device 0）。如果你有多卡，记得用 torch.cuda.set_device(1) 切换，或者用 with torch.cuda.device(1): 上下文管理器。

2.2 显存分配器：缓存池与块管理

PyTorch的显存分配器，说白了就是一个“内存池”。它从GPU驱动那里一次性申请一大块显存（比如2GB），然后自己切成小块，按需分配给你。

为什么要这么做？因为每次调用 cudaMalloc 和 cudaFree 都很慢，尤其是频繁分配小张量时，性能会崩。我在项目中遇到过，一个简单的数据加载循环，因为频繁创建销毁小Tensor，导致训练速度慢了3倍。

分配器的核心逻辑是这样的：

• 大块缓存：从驱动申请2MB以上的连续显存块
• 小块缓存：小于2MB的块，用专门的缓存池管理
• 回收机制：Tensor释放后，显存不还给驱动，而是留在缓存池里复用

你可以用 torch.cuda.memory_summary() 查看当前分配器的状态：

print(torch.cuda.memory_summary())

输出会告诉你：当前用了多少显存、缓存了多少、峰值是多少。这个命令我几乎每天都会用，尤其是在调试OOM的时候。

2.3 缓存策略：如何复用显存

PyTorch的缓存策略，核心就一句话：尽量复用，减少分配。

具体来说，分配器维护了一个“空闲块列表”。当你释放一个Tensor时，它的显存块不会立刻销毁，而是被标记为“空闲”。下次有相同大小的Tensor需要分配时，直接从这个列表里取。

为什么会这样？因为GPU显存分配是昂贵的操作。你想想看，如果每次前向传播都重新分配一遍所有中间激活值，那训练速度得慢成什么样。

我建议你养成一个习惯：尽量复用Tensor，而不是反复创建。比如：

# 不好的做法
for i in range(1000):
    x = torch.randn(1024, 1024, device='cuda')
    y = x * 2
    # ... 计算完后x和y被销毁

# 好的做法
x = torch.empty(1024, 1024, device='cuda')
y = torch.empty_like(x)
for i in range(1000):
    x.random_()  # 原地更新
    torch.mul(x, 2, out=y)  # 写入已有Tensor

嗯，这里要注意：torch.empty 不会初始化数据，但会分配显存。配合 out= 参数，可以避免重复分配。

2.4 内存碎片化：隐形杀手

内存碎片化，是显存优化的头号敌人。它不像OOM那么明显，但会让你莫名其妙地跑不了大模型。

碎片化是怎么产生的？简单说：你分配了不同大小的Tensor，释放后，显存被切成了一块块“空洞”。下次要分配一个大Tensor时，虽然总空闲显存够，但没有连续的大块空间，于是分配失败。

我曾经在一个NLP项目里，训练一个BERT-large模型。刚开始好好的，跑了200步后突然OOM。我一看显存占用，才用了70%，怎么就OOM了？后来用 torch.cuda.memory_summary() 一查，发现空闲块有300多个，但最大的连续块只有128MB。这就是典型的碎片化。

如何检测碎片化？你可以用这个公式：

def fragmentation_rate():
    total = torch.cuda.memory_stats()['allocated_bytes.all.current']
    max_block = torch.cuda.memory_stats()['reserved_bytes.large_pool.free']
    if total == 0:
        return 0.0
    return 1.0 - max_block / total

碎片率超过30%就要警惕了。我一般控制在10%以内。

2.5 实战：用SVG图理解显存管理流程

说了这么多，咱们用一张图把整个流程串起来。下面这张SVG图展示了PyTorch显存管理的核心路径：

2.6 总结与建议

好了，咱们把PyTorch显存管理的核心捋一遍：

• torch.cuda 是入口，负责设备管理和显存分配
• 分配器 用缓存池减少cudaMalloc/cudaFree调用
• 缓存策略 复用空闲块，提升性能
• 碎片化 是隐形杀手，需要主动监控和预防

我个人建议，每次训练前先跑一下 torch.cuda.memory_summary()，看看基线。训练过程中每隔100步再查一次，对比碎片率的变化。如果发现碎片率飙升，优先检查是不是动态shape或者频繁创建小Tensor导致的。

嗯，最后说一句：显存优化不是玄学，是工程。你掌握了这套机制，就能在有限的硬件上跑更大的模型。下次遇到OOM，别急着换卡，先看看是不是分配器在“捣乱”。