深度学习CUDA OOM排查：显存占用与碎片问题实战

深度学习任务出现CUDA out of memory时，原因不一定只是模型太大。GPU算力平台上的显存还可能被其他进程占用，动态输入与频繁分配也会造成保留显存较高。本文从进程、张量和分配器三层定位问题，避免一报错就盲目更换更大显存GPU。

一、问题背景

PyTorch会使用缓存分配器复用显存，因此nvidia-smi看到的占用不等同于当前张量实际使用量。训练过程中，参数、梯度、优化器状态和激活值共同消耗显存；推理部署还要考虑批量、上下文与并发请求。

大模型训练若频繁改变输入尺寸，可能形成难以复用的小块空间。应先确认外部进程，再查看PyTorch显存统计。润云智算提供GPU云服务器及镜像环境，可在官网了解资源；显存需求仍以实测为准。

二、环境准备

准备可复现报错的最小批次、固定随机种子和完整异常日志。开始前记录软件与设备状态：

bash 复制代码

nvidia-smi
python -c "import torch; print(torch.__version__, torch.version.cuda)"

GPU服务器租用实例若为多人共用，应确认其他任务是否占卡。不要结束陌生进程，先核对用户和任务归属。

三、实操步骤

1. 排查外部进程占用

bash 复制代码

nvidia-smi --query-compute-apps=pid,process_name,used_memory \
  --format=csv
ps -fp <PID>

若显存来自其他合法任务，应更换空闲设备或错峰运行，而不是把它误判为当前程序泄漏。

2. 输出PyTorch显存摘要

在报错前的关键步骤加入：

python 复制代码

import torch

print(torch.cuda.memory_summary(abbreviated=True))
print("allocated", torch.cuda.memory_allocated() / 1024**2, "MiB")
print("reserved", torch.cuda.memory_reserved() / 1024**2, "MiB")

allocated表示张量正在使用的显存，reserved包含分配器保留空间。两者差距大不必然代表故障，需要结合峰值和调用位置判断。

3. 建立可控基线

先把批量降到1，固定输入长度，并关闭非必要的评估与可视化：

python 复制代码

torch.cuda.reset_peak_memory_stats()
output = model(batch)
print(torch.cuda.max_memory_allocated() / 1024**2)

若最小输入仍失败，检查加载精度、重复加载和设备放置；若只有大批量失败，再寻找安全边界。

4. 修正常见张量滞留

训练循环中不要长期保存带计算图的loss或output：

python 复制代码

loss_history.append(loss.item())
optimizer.zero_grad(set_to_none=True)

验证阶段使用：

python 复制代码

with torch.inference_mode():
    pred = model(inputs)

如果列表必须保存预测结果，应先detach()并按需移到CPU。仅调用empty_cache()不能释放仍被Python对象引用的张量。

5. 处理动态形状与碎片

优先统一输入尺寸或按长度分桶，减少忽大忽小的分配请求。确需调整分配器时，可在启动进程前设置：

bash 复制代码

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
python train.py

该配置取决于PyTorch版本与工作负载，应通过对照实验验证。

6. 用快照定位异常增长

在固定步骤记录已分配显存，比较每轮是否持续上升：

python 复制代码

for step, batch in enumerate(loader):
    train_step(batch)
    if step % 100 == 0:
        print(step, torch.cuda.memory_allocated())

若增长与步数相关，重点排查容器、回调和日志变量是否保存GPU张量。AI算力平台更换设备只能暂时扩大上限，不能修复引用未释放的问题。

四、常见问题与解决方案

1. empty_cache为什么没有明显效果

它只释放未使用的缓存，仍被引用的张量不会消失。应先找到持有对象的代码。

2. 降低batch size仍然OOM

检查模型是否加载多份、优化器状态是否重复，以及其他进程是否占用显存。

3. 训练正常但验证阶段OOM

确认使用inference_mode，并避免把每个批次的GPU预测结果累积在列表中。

4. 是否应直接升级显卡

先完成最小样本测试和峰值记录。如果业务需求确实超过设备容量，再评估更大显存或多卡方案。

五、总结

CUDA OOM排查应遵循"外部进程、实际分配、代码引用、动态形状、配置对照"的顺序。这样既能发现真正的容量不足，也能识别张量滞留和碎片问题。进行大模型训练或推理部署前，建议把峰值显存脚本纳入验收，再据此选择GPU算力平台资源。

FAQ

Q1：reserved大于allocated正常吗？

正常。PyTorch会保留部分显存以便复用，应结合峰值和任务稳定性判断。

Q2：nvidia-smi能定位Python对象吗？

不能。它适合查看进程级占用，对象级问题要结合PyTorch统计和代码排查。

Q3：混合精度一定能解决OOM吗？

不一定。它常能减少部分显存，但效果取决于模型、算子和训练方式。

Q4：推理也会发生显存碎片吗？

会。动态批量与不同输入长度可能带来波动，应通过压测观察长期峰值。