显存池核心机制：预分配策略、缓存复用机制、碎片整理原理
聊到显存池，我得先跟你交个底——这东西看着简单，但坑是真不少。我最早接触显存池是在做大规模推荐系统训练的时候，那时候batch size一上去，显存直接炸了。后来我才意识到，问题不在于显存不够，而在于你不会管。
显存池说白了就三件事：预分配、缓存复用、碎片整理。这三板斧砍下去，显存利用率能提升30%以上。咱们一个一个拆开讲。
预分配策略：别等用的时候再要
你想想看，每次训练迭代都要申请显存，再释放，再申请……这就像你每次做饭都要去超市买菜，做完一顿再买下一顿的。效率能高吗？
预分配的核心思路就是：一次性把可能用到的显存都申请好，放到池子里。训练的时候直接从池子里拿，用完还回去。
关键点：预分配不是越多越好。我见过有人直接预分配了80%的显存，结果模型本身只用了60%，剩下的20%全浪费了。预分配的大小要根据模型的最大显存需求来算。
具体怎么做？我一般用两种策略：
- 固定大小池：提前算好最大batch size下的显存需求，一次性分配。适合batch size固定的场景。
- 动态增长池：先分配一个基础大小，不够了再按倍数增长。适合batch size动态调整的场景。
我个人习惯用动态增长池。为什么呢？因为实际训练中，你很难保证batch size一直不变。比如你调参的时候，可能从32调到64，再调到128。固定池就得重新分配，动态池直接扩容就行。
# 伪代码：动态增长池
class DynamicMemoryPool:
def __init__(self, initial_size=1GB, growth_factor=2):
self.pool = allocate(initial_size)
self.size = initial_size
self.growth_factor = growth_factor
def request(self, need_size):
if need_size > self.size:
# 扩容
new_size = max(need_size, self.size * self.growth_factor)
new_pool = allocate(new_size)
copy(self.pool, new_pool)
free(self.pool)
self.pool = new_pool
self.size = new_size
return self.pool[:need_size]小技巧：预分配的时候，可以留10%-20%的余量。别问我为什么，问就是踩过坑——有一次模型稍微改了个结构，显存需求涨了5%，结果池子不够用，直接OOM了。
缓存复用机制：用完别扔，洗洗还能用
预分配解决了「申请慢」的问题，但还有一个问题：显存碎片。每次申请和释放，都会在显存上留下大小不一的空洞。时间一长，显存利用率直线下降。
缓存复用的思路很简单：用完的显存块别急着释放，先缓存起来。下次有相同大小的请求，直接复用。
我做过一个实验：在ResNet-50训练中，不加缓存复用，显存碎片率在10轮训练后达到15%。加了缓存复用后，碎片率稳定在3%以内。
| 策略 | 显存碎片率（10轮后） | 训练速度 |
|---|---|---|
| 无缓存复用 | 15% | 基准 |
| 简单缓存复用 | 8% | +5% |
| 智能缓存复用 | 3% | +12% |
智能缓存复用是怎么做的？说白了就是按大小分类。把显存块按大小分成几个桶：小块（<1MB）、中块（1MB-100MB）、大块（>100MB）。每个桶用LRU（最近最少使用）策略管理。
注意：缓存复用不是万能的。我曾经在一个NLP模型上试过，结果发现缓存命中率只有20%。为什么？因为Transformer的显存需求变化太大，每次请求的大小都不一样。这种情况下，缓存复用反而增加了管理开销。
所以我的建议是：先做profiling。跑几轮训练，统计一下显存请求的大小分布。如果分布集中（比如80%的请求都在某个范围内），那缓存复用效果会很好。如果分布太散，不如直接用预分配。
碎片整理原理：给显存做个「大扫除」
即使有预分配和缓存复用，显存碎片还是会出现。为什么？因为有些显存块被长期占用（比如模型参数），有些是临时申请（比如中间激活值）。长期占用的块会把显存空间「切」成很多小段。
碎片整理，说白了就是把分散的小块合并成大块。但这里有个问题：你不能随便移动显存里的数据，因为可能有指针在引用它。
我常用的碎片整理策略有两种：
- 压缩式整理：把所有空闲块移动到一端，已用块移动到另一端。需要暂停训练，开销较大。
- 分代式整理：把显存块按「年龄」分组。年轻块（刚分配的）经常整理，年老块（长期存在的）少整理。
我个人更推荐分代式整理。为什么呢？因为模型参数这种「老块」基本不会变，整理它们纯属浪费时间。而中间激活值这种「年轻块」频繁申请释放，才是碎片的罪魁祸首。
# 伪代码：分代碎片整理
class GenerationalDefrag:
def __init__(self):
self.young_gen = [] # 年轻代，频繁整理
self.old_gen = [] # 老年代，少整理
def defrag(self):
# 年轻代：每次训练迭代后整理
self.young_gen = compact(self.young_gen)
# 老年代：每100轮整理一次
if step % 100 == 0:
self.old_gen = compact(self.old_gen)核心原则：碎片整理的频率要和训练开销做平衡。我一般控制在每轮训练增加不超过5%的时间开销。如果整理太频繁，反而得不偿失。
嗯，这里要注意一个细节：碎片整理的时候，要保证所有正在使用的显存块不被移动。我见过一个坑——有人在整理的时候没检查引用计数，结果把正在计算的张量给移动了，训练直接崩了。
所以我的做法是：在每次迭代的间隙做整理。具体来说，就是在反向传播结束、优化器更新之前，这个时间段没有张量在计算，可以安全地移动数据。
三者的配合：一个完整的显存池
预分配、缓存复用、碎片整理，这三者不是孤立的。它们需要配合使用：
- 预分配：先申请一大块显存，作为池子的基础。
- 缓存复用：训练过程中，用完的显存块先缓存起来，下次复用。
- 碎片整理：当缓存复用找不到合适大小的块时，触发碎片整理，合并小空闲块。
我画了一张流程图，帮你理解这三者的关系：
你看这个流程：预分配是基础，缓存复用是核心，碎片整理是兜底。三者缺一不可。
实战建议：刚开始做显存池的时候，别想着一步到位。先实现预分配，跑通。再加缓存复用，观察命中率。最后加碎片整理，看能不能进一步降低碎片率。我当年就是这么一步步来的，稳得很。
最后说一句：显存池不是银弹。如果你的模型本身就很小，显存绰绰有余，那搞这些反而增加复杂度。但如果你像我一样，经常跟大模型、大batch size打交道，那显存池就是你的救命稻草。
想要更多资料，联系可可