显存分配与生命周期管理：PyTorch显存分配器机制

显存管理这事儿，说白了就是跟GPU要内存、用完了再还回去。但问题在于——怎么要、怎么还、还了之后能不能再快速借出来？我刚开始做大模型训练时，经常遇到OOM（显存溢出），后来才发现，很多时候不是显存真的不够，而是分配器没用好。

PyTorch显存分配器机制

PyTorch默认用的是自己的缓存分配器（caching allocator），而不是直接调用cudaMalloc。为什么要这么设计？因为cudaMalloc这玩意儿开销太大了，每次分配都要跟驱动打交道，频繁调用会严重拖慢训练速度。

PyTorch的做法是：先一次性从GPU申请一大块显存，然后自己内部做二次分配。你调用torch.cuda.empty_cache()时，其实只是把缓存标记为可用，并没有真正还给操作系统。嗯，这里要注意——这个函数只清空缓存，不释放已分配的内存块。

核心机制：PyTorch分配器维护了一个缓存池（cache pool），里面放着之前分配过但已经释放的显存块。下次需要同样大小的显存时，直接从池子里拿，省去了cudaMalloc的开销。

# 查看当前显存使用情况
import torch

# 分配一个张量
x = torch.randn(1024, 1024).cuda()
print(f"当前分配: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"缓存大小: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")

# 释放张量
del x
print(f"释放后分配: {torch.cuda.memory_allocated() / 1024**2:.2f} MB")
print(f"释放后缓存: {torch.cuda.memory_reserved() / 1024**2:.2f} MB")
# 注意：分配量归零，但缓存量不变

显存碎片化问题

碎片化是我在实际项目中踩过最大的坑。有一次训练一个BERT-large模型，明明总显存还有4GB空闲，但就是分配不出一个2GB的tensor。为什么会这样？因为显存被切成了无数个小碎片，每个碎片都不够大。

碎片化主要来自两个地方：

大小不一的张量反复分配释放——比如你循环里创建不同shape的中间变量
梯度检查点（checkpoint）的不当使用——保存的中间激活值大小差异太大

我的经验：训练前先跑一次warmup，让分配器把缓存池填满。这样后续训练时，分配器已经知道要分配哪些大小的块，碎片化会明显减少。

# 查看碎片化程度
def check_fragmentation():
    allocated = torch.cuda.memory_allocated()
    reserved = torch.cuda.memory_reserved()
    # 碎片率 = 1 - (已分配 / 已缓存)
    frag_ratio = 1 - allocated / max(reserved, 1)
    print(f"碎片率: {frag_ratio:.2%}")
    return frag_ratio

# 模拟碎片化场景
tensors = []
for i in range(100):
    size = (128, 128 * (i % 10 + 1))
    t = torch.randn(*size).cuda()
    if i % 3 == 0:
        del t  # 随机释放，制造碎片
    else:
        tensors.append(t)

check_fragmentation()

手动释放与缓存策略

我个人习惯在关键节点手动干预显存管理。虽然PyTorch的自动管理已经很智能了，但有些场景下，手动释放能带来质的提升。

1. 及时删除不再使用的张量

# 不好的做法：中间变量一直占用显存
def bad_forward(x):
    h1 = torch.relu(x @ w1 + b1)
    h2 = torch.relu(h1 @ w2 + b2)
    return h2 @ w3 + b3

# 好的做法：及时释放中间变量
def good_forward(x):
    h1 = torch.relu(x @ w1 + b1)
    h2 = torch.relu(h1 @ w2 + b2)
    del h1  # 显式释放
    result = h2 @ w3 + b3
    del h2
    return result

避坑指南：我曾经在循环里忘记释放中间变量，结果一个batch跑完显存涨了2GB。后来加了del和torch.cuda.empty_cache()才稳住。但注意——empty_cache不能频繁调用，否则会拖慢速度。

2. 使用上下文管理器自动释放

# 使用torch.no_grad()减少中间变量
with torch.no_grad():
    # 推理时不需要梯度，减少显存占用
    output = model(input_data)

# 自定义上下文管理器
class AutoRelease:
    def __enter__(self):
        return self
    
    def __exit__(self, *args):
        torch.cuda.empty_cache()

# 使用示例
with AutoRelease():
    # 做一些临时计算
    temp = torch.randn(10000, 10000).cuda()
    result = temp.sum()
    # 退出时自动清理缓存

3. 缓存策略对比

策略	适用场景	优点	缺点
默认缓存	常规训练	速度快，无需干预	碎片化风险
手动释放	显存紧张时	精确控制	代码侵入性强
定期清空	长训练任务	防止累积	性能损失
预分配	固定大小张量	零碎片	灵活性差

知识体系总览

下面这张图是我自己总结的显存管理核心逻辑，你一看就明白：

你看这个流程，其实核心就两件事：怎么快速拿到显存和怎么避免碎片化。缓存池是PyTorch给我们的礼物，但用不好反而会变成负担。

实用技巧：训练前用torch.cuda.empty_cache()清空一次缓存，然后跑一个batch做warmup。这样后续的分配都会走缓存命中路径，速度最快。

最后说一句——显存管理不是玄学，是工程。你只要理解了分配器的缓存机制，再配合手动释放和预分配策略，OOM基本就跟你无缘了。我在多个百亿参数模型上验证过这套方法，效果很稳定。

想要更多资料，联系可可

📚 课程目录

显存分配与生命周期管理：PyTorch显存分配器机制

PyTorch显存分配器机制

显存碎片化问题

手动释放与缓存策略

1. 及时删除不再使用的张量

2. 使用上下文管理器自动释放

3. 缓存策略对比

知识体系总览