显存管理：从碎片到回收，把每一兆都榨干

显存这东西，说白了就是大模型训练的命根子。我见过太多团队，模型架构设计得漂漂亮亮，一跑起来就OOM（Out of Memory），然后开始怀疑人生。其实很多时候不是显存不够，是你没管好。

今天咱们就聊聊显存管理的四个核心问题：碎片化、梯度检查点、混合精度训练、以及监控回收。每个坑我都踩过，每个方案我都亲手调过。

一、显存碎片化：看不见的内存杀手

你想想看，训练一个70B的模型，显存占用曲线是不是经常像锯齿一样？一会儿涨上去，一会儿掉下来。这就是碎片化在作祟。

为什么会这样？PyTorch的CUDA缓存分配器默认会缓存已释放的显存块，下次分配时复用。但问题是，频繁的tensor创建和销毁会导致显存块大小不一，就像硬盘碎片一样。

我在项目中遇到过最夸张的一次：一张A100 80G，实际只用了60G，但分配一个32G的tensor时直接报错。查了半天，发现是碎片化导致没有连续32G的空闲块。

解决方案：

• 设置PYTORCH_CUDA_ALLOC_CONF：用 max_split_size_mb 限制最大碎片块大小
• 使用expand_as替代view：减少tensor形状变化带来的重新分配
• 预分配策略：在训练开始前预先分配好所有需要的显存

# 我常用的配置
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

# 或者更激进一点
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

二、梯度检查点：用时间换空间的艺术

梯度检查点（Gradient Checkpointing），说白了就是牺牲一点计算时间，换来显存的大幅降低。原理很简单：前向传播时不保存中间激活值，反向传播时重新计算。

我刚开始用的时候觉得这玩意儿太慢了，后来发现——嗯，真香。尤其是训练大模型时，显存瓶颈远比计算瓶颈更致命。

怎么用？

from torch.utils.checkpoint import checkpoint

# 原始前向
def forward(self, x):
    return self.layer2(self.layer1(x))

# 使用梯度检查点
def forward(self, x):
    return checkpoint(self.layer2, checkpoint(self.layer1, x))

但要注意，不是所有层都适合用检查点。我个人习惯只对计算量大的层（比如Transformer的Attention和FFN）使用，对embedding层和layernorm层不用。

经验数据：

模型规模	不使用检查点	使用检查点	显存节省
7B	48G	28G	~42%
13B	76G	44G	~42%
70B	需要多卡	单卡可跑	质变

三、混合精度训练：FP16/BF16的实战心法

混合精度训练现在已经是标配了。但你真的用对了吗？

我记得第一次用AMP（Automatic Mixed Precision）时，模型直接不收敛。后来才发现是loss scale设置的问题。AMP默认是动态loss scale，但有些场景下需要手动调整。

实战配置：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler(init_scale=2**16)  # 初始scale值

for data, target in dataloader:
    optimizer.zero_grad()
    
    with autocast(dtype=torch.bfloat16):  # 推荐用BF16
        output = model(data)
        loss = loss_fn(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

你可能会问：什么时候用FP16？我建议只有在显存极度紧张且模型对精度不敏感时才用FP16。比如推理场景，或者小模型的微调。

混合精度的三个层次：

1. Level 1 - 自动混合精度：用PyTorch的AMP，适合大多数场景
2. Level 2 - 手动控制：对特定层强制使用FP32（比如BatchNorm、损失计算）
3. Level 3 - 自定义策略：根据梯度分布动态调整精度，适合极致优化

四、显存监控与回收：别等爆了再处理

显存监控这事儿，很多人都是等到OOM了才想起来看。其实应该像看仪表盘一样，时刻关注。

我习惯在训练脚本里加一个显存监控的装饰器，每N步打印一次显存状态。这样能及时发现显存泄漏或者异常增长。

监控工具：

import torch.cuda as cuda

def print_memory_stats():
    allocated = cuda.memory_allocated() / 1024**3
    reserved = cuda.memory_reserved() / 1024**3
    print(f"已分配: {allocated:.2f}G, 已预留: {reserved:.2f}G")
    
# 更详细的
print(cuda.memory_summary())

显存回收策略：

• 手动回收：torch.cuda.empty_cache() 释放未使用的缓存
• 自动回收：设置 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
• 对象回收：及时删除不再使用的tensor，用 del 配合 gc.collect()

实战中的显存泄漏排查：

如果你发现显存随着训练步数线性增长，大概率是显存泄漏。常见原因：

1. 在循环中创建了tensor但没有释放
2. DataLoader的num_workers设置不当导致内存泄漏
3. 自定义的hook或者callback没有正确清理

排查方法很简单：在每个epoch结束时打印显存，如果持续增长，就二分法注释代码找问题。

知识体系总览

下面这张图是我自己总结的显存管理全景图，你可以把它当作一个检查清单：

这四个模块不是孤立的。比如你用混合精度训练，显存占用降低了，碎片化问题可能就缓解了。反过来，如果你用了梯度检查点，显存占用曲线更平滑，碎片化也会减少。它们之间是相互影响的。

最后说一句：显存管理没有银弹。每个方案都有trade-off，关键是你得清楚自己的瓶颈在哪。是显存不够？还是碎片化严重？还是精度损失？对症下药才是王道。

---

公众号:蓝海资料掘金营，微信deep3321