3、PyTorch显存管理：torch.cuda机制、显存分配与释放、缓存分配器

显存管理，说白了就是PyTorch怎么跟GPU要内存、怎么还内存。很多同学模型跑着跑着就OOM了，其实不是显存真的不够，而是你没搞懂PyTorch的分配策略。我刚开始调大模型时也吃过这个亏，今天咱们就把这块彻底讲透。

3.1 torch.cuda机制：谁在管显存？

PyTorch的显存管理，核心就两个角色：CUDA运行时和PyTorch缓存分配器。

• CUDA运行时：真正的“地主”。它负责向GPU驱动申请物理显存。
• PyTorch缓存分配器：二房东。它从CUDA运行时那里批量化拿一大块显存，然后切成小块分给张量用。

为什么要搞这么复杂？直接每次要多少给多少不行吗？

不行。你想想看，每次调用cudaMalloc都是系统调用，开销极大。如果每个张量都单独申请释放，训练一个ResNet可能要几万次系统调用，速度直接崩掉。

3.2 显存分配：二房东的生意经

PyTorch的缓存分配器是怎么工作的？我画了一张图，你看完就明白了。

分配器内部维护了一个空闲块列表。当你创建一个张量时，它先去空闲列表里找大小合适的块。找到了就直接用，找不到才去CUDA运行时那里申请新的。

3.3 显存释放：不是你想的那么简单

很多人以为del tensor之后显存就立刻释放了。嗯，这里要注意——并不是。

PyTorch的释放分两步：

1. 张量引用计数归零：Python的垃圾回收器标记该内存可回收。
2. 缓存分配器回收：内存块回到空闲列表，但不归还给CUDA运行时。

也就是说，你del掉张量后，memory_allocated()会下降，但memory_reserved()纹丝不动。这些显存还在PyTorch手里，下次有张量需要时直接复用。

3.4 缓存分配器的配置与调优

PyTorch的缓存分配器支持一些环境变量配置，我整理了一张表：

环境变量	作用	默认值	我的建议
PYTORCH_CUDA_ALLOC_CONF	分配器配置字符串	空	OOM时必调
max_split_size_mb	最大分割块大小	无限制	设为512可减少碎片
roundup_power2_divisions	向上取整对齐策略	1	设为16减少碎片
expandable_segments	是否启用可扩展段	False	PyTorch 2.0+推荐True

举个例子，如果你遇到显存碎片化严重的问题，可以这样配置：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512,roundup_power2_divisions:16'

# 或者直接在命令行启动时设置
# PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python train.py

3.5 实战：监控与调试显存

最后，分享几个我常用的调试手段：

• torch.cuda.memory_allocated()：当前张量占用的显存，单位字节。
• torch.cuda.memory_reserved()：分配器持有的总显存。
• torch.cuda.max_memory_allocated()：历史峰值占用，排查OOM利器。
• torch.cuda.reset_peak_memory_stats()：重置峰值统计，方便分段测量。

我习惯在训练脚本里加个显存监控装饰器：

import torch
import functools

def monitor_memory(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        torch.cuda.reset_peak_memory_stats()
        result = func(*args, **kwargs)
        peak = torch.cuda.max_memory_allocated() / 1024**3
        current = torch.cuda.memory_allocated() / 1024**3
        print(f"[显存] 峰值: {peak:.2f}GB, 当前: {current:.2f}GB")
        return result
    return wrapper

@monitor_memory
def train_step(batch):
    # 你的训练逻辑
    pass

说白了，显存管理就是理解“二房东”的运作逻辑。你只要记住：分配器为了性能会囤积显存，OOM不一定是真不够，可能是缓存池太大。下次遇到OOM，先跑个memory_summary()看看，八成能找到问题。

想要更多资料，联系可可