2、显存池初探：什么是显存池？它和普通显存分配有什么不同？

好，咱们正式开始聊显存池。

在讲具体实现之前，我想先问你一个问题：你平时写深度学习模型的时候，有没有遇到过显存突然爆掉的情况？明明算了一下觉得够用，结果一跑就 OOM（Out of Memory）。

我遇到过太多次了。有一次调一个 7B 参数的模型，前向传播好好的，反向传播刚走到一半，啪，显存炸了。当时我盯着 nvidia-smi 看了十分钟，发现峰值显存比平均值高出了将近 40%。

嗯，这就是我们今天要解决的问题。

2.1 普通显存分配：随用随取，随取随还

先看看常规做法。PyTorch 或者 TensorFlow 里，你每创建一个张量，框架就会向 CUDA 驱动申请一块显存。

import torch

# 普通分配：每次创建张量都去申请显存
a = torch.randn(1024, 1024).cuda()  # 申请 4MB
b = torch.randn(1024, 1024).cuda()  # 再申请 4MB
c = a + b                           # 临时张量，再申请 4MB
del a                               # 释放 4MB
del b                               # 释放 4MB

这个过程看起来挺自然的对吧？但问题出在哪呢？

说白了，每次申请和释放显存，都是一次系统调用。CUDA 驱动需要去维护一个空闲列表，找到合适的内存块，更新页表……这些操作都是有开销的。

更麻烦的是，显存碎片化。你申请 4MB，释放 4MB，再申请 8MB，系统可能找不到连续的空间。我见过一个极端案例：显存总量 24GB，空闲 18GB，但就是申请不下一块 12GB 的张量。为什么？因为空闲显存被切成了几百个小碎片。

2.2 显存池：先囤一批，随用随取

显存池的思路其实很朴素：与其每次去系统申请，不如我一次性先拿一大块，自己管自己用。

你可以把它想象成一个仓库。普通分配是每次需要工具了，跑去五金店买一个，用完再退回去。显存池呢？我直接包下一个仓库，把常用的工具都备齐了，用的时候直接拿，用完放回架子上。

# 显存池的思路（伪代码）
memory_pool = cuda_alloc(10 * 1024 * 1024 * 1024)  # 一次性申请 10GB

def pool_alloc(size):
    # 从池子里找一块够大的空闲块
    block = find_free_block(memory_pool, size)
    if block is None:
        # 池子不够了，才去系统申请
        block = cuda_alloc(size)
    return block

def pool_free(block):
    # 不真正释放，标记为空闲
    mark_as_free(block)

你看，核心区别就两点：

• 申请时：优先从池子里拿，池子不够才找系统
• 释放时：不还给系统，放回池子里复用

这样做的好处很明显：

1. 减少系统调用：大部分分配释放都在用户态完成，快得多
2. 降低碎片化：池子内部可以自己整理碎片，合并小空闲块
3. 平滑峰值：因为池子提前占了一块显存，临时申请的次数变少了

2.3 两者对比：一张表说清楚

对比维度	普通显存分配	显存池
申请方式	每次调用 cudaMalloc	优先从池中分配
释放方式	调用 cudaFree 归还系统	标记为空闲，留在池中
系统调用次数	高（每次分配释放都涉及）	低（仅池扩容时涉及）
显存碎片	容易产生外部碎片	内部可整理，碎片可控
峰值控制	被动，取决于瞬时使用量	主动，池大小可预设
适用场景	小模型、一次性推理	大模型训练、服务部署

我个人习惯在训练大模型时，一定会用显存池。尤其是做多卡并行的时候，每张卡的显存分配如果不统一，很容易出现某张卡先爆掉的情况。池化之后，每张卡的显存使用曲线平滑很多。

2.4 一个直观的例子：峰值是怎么降下来的

咱们看一个具体场景。假设你的模型在反向传播时需要创建一些临时张量：

# 普通分配下的显存变化
# 时间点1: 前向结束，占用 8GB
# 时间点2: 开始反向，创建梯度张量，占用 8GB + 2GB = 10GB
# 时间点3: 创建临时计算张量，占用 10GB + 4GB = 14GB  ← 峰值
# 时间点4: 释放临时张量，回到 10GB
# 时间点5: 释放梯度张量，回到 8GB

峰值是 14GB。如果用显存池呢？

# 显存池分配下的显存变化
# 时间点0: 池子预分配 12GB
# 时间点1: 前向结束，池内使用 8GB，空闲 4GB
# 时间点2: 反向开始，从池内拿 2GB，池内使用 10GB，空闲 2GB
# 时间点3: 创建临时张量，池内空闲 2GB 不够，从系统再申请 2GB
#          池外申请 2GB，总占用 12GB + 2GB = 14GB  ← 峰值一样？

等等，峰值好像没降？

嗯，这里要注意：显存池本身不能凭空减少显存使用量。它降低的是「系统层面的峰值波动」。你看，普通分配下，系统看到的是 8→10→14→10→8 的剧烈波动。显存池下，系统看到的是 12→12→12→14→12 的平稳曲线。

更重要的是，如果池子预分配得够大（比如 14GB），那整个训练过程中，系统层面根本看不到任何波动。这就是为什么我说显存池能「平滑峰值」。

2.5 显存池的核心逻辑图

下面这张图展示了显存池的工作流程。我画得比较简洁，但核心逻辑都在里面了。

这张图你看懂了吗？核心就两个分支：池里有就用池里的，池里没有才去找系统。释放的时候，一律放回池子，不还给系统。这样，大部分分配释放操作都在用户态完成，速度飞快。

2.6 避坑指南

显存池也不是银弹。我踩过几个坑，跟你分享一下：

• 池子设太大：我曾经为了省事，直接预分配了全部显存。结果模型加载时发现显存不够用，因为 CUDA 驱动自己也要留一点。建议留 5%-10% 给系统。
• 碎片整理不及时：池子用久了，内部也会碎片化。需要定期做一次合并操作，把相邻的空闲块合并成大块。我一般每 100 次迭代做一次。
• 多线程安全：如果你的模型用到了多线程数据加载，池子的分配和释放需要加锁。否则两个线程同时拿同一块显存，后果你懂的。

好了，这一章就到这里。下一章我们深入代码，看看怎么手写一个简单的显存池。

---

想要更多资料，联系可可