在使用 MinerU、Stable Diffusion 或 LLM 本地部署时，我们经常会遇到一种“灵异现象”：模型明明跑完了，任务也结束了，但显存（VRAM）占用依然居高不下。如果有两个模型想在同一张显卡上轮流跑，经常会因为前一个模型“占着茅坑不拉屎”而导致后一个模型 OOM（Out Of Memory）。

为什么显存会“无休止”地扩张？为什么框架不愿意把显存还给操作系统？这背后其实是一套精心设计的“贪婪”机制。

今天我们就来扒一扒深度学习框架（以 PyTorch 为例，TensorFlow 同理）的显存管理底层逻辑，以及如何在单卡上优雅地运行多个模型。

一、 为什么显存“只吃不吐”？

很多开发者会发现，Python 脚本运行过程中，显存占用往往是一条只增不减的曲线。这其实是因为框架内部实现了一个 Caching Allocator（缓存分配器）。

1. 昂贵的系统调用

在 GPU 上分配内存（cudaMalloc）和释放内存（cudaFree）是系统级调用，这些操作非常“昂贵”，会强制 GPU 同步，打断计算流水线。

如果模型在推理过程中，每算一层网络就申请一次显存，算完就释放，那么 GPU 大部分时间都在等待内存分配，而不是在做计算。这会导致推理速度慢如蜗牛。

2. 框架的“贪婪”策略

为了极致的性能，PyTorch 采用了一种“池化管理”策略：

  • 申请：当你需要显存时，PyTorch 会先找 CUDA 要。
  • 释放：当你用完显存（比如变量被销毁），PyTorch 不会把显存还给 CUDA（操作系统），而是把这块显存标记为“空闲”，放入自己的内部缓存池中。
  • 复用：下次你需要显存时，PyTorch 直接从缓存池里拿，完全跳过了昂贵的 cudaMalloc。

这就是为什么你用 nvidia-smi 看到的显存占用一直很高，但程序内部其实并没有真的在用那么多。框架就像一个“守财奴”，把要来的显存都存着，以备不时之需。

二、 显存碎片的诅咒

你可能遇到过这种情况：nvidia-smi 显示还有 4GB 空闲显存，但你申请 2GB 显存时却报错 OOM。

这是“显存碎片”在作祟。

想象显存是一条长长的面包：

  1. 你切了 1GB 给变量 A。
  2. 你切了 2GB 给变量 B。
  3. 你切了 1GB 给变量 C。
  4. 现在变量 B 用完了，释放了。中间空出了 2GB。
  5. 现在你需要申请 3GB 给变量 D。

虽然总空闲显存有 2GB（中间）+ 剩余的，但中间那块 2GB 的洞塞不下 3GB 的变量 D。这就是外部碎片。

MinerU 或 Transformer 类模型在处理变长序列时，经常频繁申请不同大小的显存，极易导致碎片化。这也是为什么有时候重启一下代码就能跑通的原因——重启重置了显存布局。

三、 如何实现多模型共存？

既然知道了底层原因，我们在单卡上同时跑多个模型（比如一个 OCR，一个 LLM）时，就有章可循了。

1. 主动清理缓存：empty_cache()

这是最直接的手段。

1
2
3
4
5
6
7
8
9
10
11
import torch

# 运行模型 A
model_a.predict()
del model_a  # 删除引用

# 强制归还显存给 OS
torch.cuda.empty_cache()

# 现在运行模型 B，显存就像刚开机一样干净
model_b.predict()
  • 原理：告诉缓存分配器，“别贪了，把缓存池里没用的显存统统还给操作系统”。
  • 代价：下次再申请显存时，需要重新进行系统调用，会稍微慢一点点，但对于多模型切换场景，这是必要的。

2. 显存隔离与限制

如果你希望两个模型同时运行（并发），而不是轮流运行，你需要限制每个进程的显存上限，防止一个模型吃光所有资源。

1
2
# 限制当前进程只能使用 50% 的显存
torch.cuda.set_per_process_memory_fraction(0.5, 0)

这样，即使模型想“贪婪”地缓存更多显存，也会被框架强制制止，从而给另一个模型留出空间。

3. 终极方案：模型卸载 (CPU Offload)

如果显存实在太小，可以采用“时间换空间”的策略。模型权重平时存在 CPU 内存（RAM）里，计算时才加载到 GPU。

1
2
3
4
5
6
# 伪代码示例
device = "cuda"
model.to(device)  # 加载到 GPU
output = model(input)
model.to("cpu")   # 用完立刻踢回 CPU
torch.cuda.empty_cache() # 清理残留

像 accelerate 库或 LangChain 中的很多组件都内置了这种 device_map="auto" 或 offload_folder 机制，自动帮你做这件事。

四、 总结

深度学习框架的显存管理机制，本质上是在做 Speed（速度） vs Memory（显存） 的权衡。

  • 为什么不释放？ 为了快，避免系统调用。
  • 为什么会 OOM？ 除了真不够用，还可能是碎片化导致的有空地没法用。
  • 怎么解决？
    • 单任务大文件：拆分输入（如 MinerU 的单页处理）。
    • 多模型切换：用完及时 del 并 empty_cache()。
    • 多模型并发：设置显存上限 set_per_process_memory_fraction。

理解了这些，当你下次再看到显存占用居高不下时，就不会惊慌了——那只是你的框架在为你“未雨绸缪”。