1. 显存池基础概念

大家好，我是你们的老朋友。今天咱们聊聊显存池——这个在分布式训练里经常被忽视，但实际非常关键的东西。

说实话，我刚开始做大规模训练的时候，根本没想过显存分配还能搞出什么花样。那时候觉得，需要显存就 cudaMalloc，用完就 cudaFree，多简单。直到有一次，我在一个 8 卡 A100 的集群上跑一个 175B 的模型，训练到一半突然卡住了。排查了半天，发现是频繁的显存分配和释放把驱动搞崩了。嗯，从那以后，我就开始认真研究显存池了。

为什么需要显存池？

你想想看，深度学习训练过程中，显存的分配和释放有多频繁？

• 每个 batch 的前向传播，要分配中间激活值
• 反向传播时，又要释放这些激活值
• 优化器状态、梯度缓冲区，来回切换
• 通信操作（如 all-reduce）需要临时缓冲区

这些操作加在一起，一个典型的训练循环里，每秒可能发生几百次显存分配和释放。如果每次都要走 cudaMalloc/cudaFree，那性能损失就大了。

显存分配与释放的开销分析

咱们来拆解一下，一次 cudaMalloc 到底干了什么：

1. 驱动层调用：从用户态切换到内核态，这个切换本身就有几百纳秒的开销
2. 内存管理：驱动需要查找空闲页表，分配物理内存，更新页表映射
3. 清零操作：默认情况下，cudaMalloc 返回的内存是清零的，这个操作在显存带宽有限的情况下很耗时
4. 同步屏障：cudaMalloc 会隐式同步所有正在执行的 CUDA 流，确保之前的内存操作完成

我做过一个测试，在 A100 上连续分配和释放 10000 次 1MB 的显存块，总耗时接近 2 秒。而如果用显存池，同样的操作只需要 0.1 秒。差距就是这么大。

操作	单次耗时（微秒）	10000 次耗时（秒）
cudaMalloc (1MB)	~180	1.8
cudaFree (1MB)	~20	0.2
显存池分配	~0.5	0.005
显存池释放	~0.3	0.003

显存池的核心设计思想

说白了，显存池的核心思想就四个字：预分配 + 复用。

具体来说：

• 预分配：在训练开始前，一次性向 GPU 申请一大块显存（比如总显存的 90%）
• 分块管理：把这块显存切成不同大小的块，用空闲链表或红黑树管理
• 按需分配：训练过程中需要显存时，直接从池子里拿一块，不需要走驱动
• 延迟释放：释放时并不真正还给驱动，而是标记为空闲，供后续复用

这里有一个关键点：碎片管理。如果分配和释放的块大小不一，时间长了就会产生碎片。我见过一个项目，因为碎片太多，明明总显存还有 20GB，但就是分配不出一个 1GB 的连续块。

下面我用一个简单的示意图，展示显存池的核心逻辑：

从这张图可以看得很清楚：左边是每次分配释放都要走驱动，右边是预分配好一块大显存，然后在用户态管理。性能差距一目了然。

显存池的几种常见实现策略

在实际工程中，显存池的实现有很多种。我根据自己的经验，总结了三种主流方案：

• 简单空闲链表：把所有空闲块串成一个链表，分配时遍历查找合适大小的块。实现简单，但容易产生碎片。
• 伙伴系统（Buddy System）：把显存按 2 的幂次分成不同大小的块。分配时向上取整到最近的 2 的幂，释放时合并相邻的空闲块。碎片少，但内部碎片（浪费的空间）较多。
• 大小分级（Size Class）：预先定义好一组常用的大小（如 1MB、4MB、16MB...），每个大小维护一个空闲链表。分配时直接取对应大小的块，释放时放回。内部碎片可控，性能最好。

最后，我想强调一点：显存池不是银弹。如果你的模型显存需求变化非常大（比如动态 shape 的输入），显存池的碎片问题可能会很严重。这时候，你可能需要结合显存压缩、梯度检查点等技术一起使用。

好了，这一章的内容就到这里。显存池是分布式训练优化的基础，理解了它，后面讲多卡通信优化时你就能更清楚地知道为什么需要那些技巧。