显存池回收与碎片整理：内存碎片产生原因，延迟回收与合并策略
显存碎片，说白了就是显存被切得七零八落。你明明看到剩余总量还够，但就是分配不出一块连续的大显存。我最早遇到这个问题是在做多卡流水线并行时，模型卡在中间层，报错说显存不足，可一看剩余还有4GB——嗯，全是碎片。
碎片是怎么产生的？
显存碎片主要分两种：内部碎片和外部碎片。内部碎片是分配粒度导致的，比如你申请3GB，但分配器按4GB块给你，多出来的1GB就浪费了。外部碎片更头疼——频繁申请释放后，空闲块被切得支离破碎。
我举个例子。你训练一个模型，前向传播时申请了A、B、C三块显存，反向传播时释放了B。这时候空闲块就不是一整块了，而是被A和C夹在中间。下次你想申请一块比B大的显存，对不起，放不下。
碎片产生的三大元凶：
- 频繁分配释放：每次迭代都申请临时张量，用完就丢
- 大小不一的请求：激活值、梯度、优化器状态，尺寸差异巨大
- 生命周期交错：不同张量的存活时间互相穿插，形成空洞
为什么会这样？你想想看，PyTorch默认的缓存分配器是贪心的——它尽量复用已有块，但不会主动整理。久而久之，显存就像一盘散沙。
延迟回收：给碎片一点时间
我个人的习惯是，不要急着把释放的显存还给CUDA。你释放了，CUDA driver不一定立刻回收。但如果你频繁调用torch.cuda.empty_cache()，反而会打断分配器的缓存策略。
延迟回收的核心思路是：把释放的块留在池子里，等凑够了再合并。就像你收拾房间，不会每扔一张纸就跑一趟垃圾桶，而是攒一堆再倒。
// 伪代码：延迟回收逻辑
class DelayedRecycler {
queue: Vec<Block>; // 待回收块队列
threshold: usize; // 触发回收的阈值
fn release(block: Block) {
if block.size < threshold {
// 小块直接进队列，等合并
this.queue.push(block);
} else {
// 大块立即回收
cudaFree(block.ptr);
}
}
fn try_merge() {
// 按地址排序，合并相邻块
this.queue.sort_by_key(|b| b.addr);
let merged = merge_adjacent(this.queue);
this.queue = merged;
if this.queue.len() > MAX_PENDING {
flush_to_cuda();
}
}
}我在项目中遇到过，延迟回收配合水位线机制效果最好。设置两个阈值：低水位和高水位。低于低水位时，尽量不回收；高于高水位时，强制回收一批。这样既避免了频繁的CUDA操作，又不会让待回收队列无限膨胀。
避坑指南：我曾经把延迟队列设得太大，结果显存占用居高不下，OOM了。后来加了超时机制——队列里的块超过一定时间（比如10秒）还没被复用，就强制回收。这个时间值需要根据你的训练步长来调。
合并策略：把碎片拼回去
合并说白了就是找相邻的空闲块，把它们粘成一块大的。但这里有个坑——你怎么知道哪些块是相邻的？
我常用的方法是地址排序+双向合并。每次释放时，检查前后邻居是否也是空闲的。如果是，直接合并。这需要维护一个红黑树或跳表来快速查找邻居。
// 合并逻辑示意
struct Block {
addr: u64,
size: u64,
is_free: bool,
}
fn coalesce(pool: &mut Vec<Block>, idx: usize) {
let block = &pool[idx];
// 检查前一个块
if idx > 0 && pool[idx-1].is_free {
let prev = &mut pool[idx-1];
if prev.addr + prev.size == block.addr {
prev.size += block.size;
pool.remove(idx);
// 递归检查前一个块是否还能再合并
coalesce(pool, idx-1);
}
}
// 检查后一个块
if idx < pool.len()-1 && pool[idx+1].is_free {
let next = &mut pool[idx+1];
if block.addr + block.size == next.addr {
block.size += next.size;
pool.remove(idx+1);
coalesce(pool, idx);
}
}
}这里有个细节：合并后要更新空闲块列表。我习惯用伙伴系统来管理——把显存按2的幂次分成不同大小的块。释放时，如果两个伙伴块都空闲，就合并成更大的块。这样碎片率能控制在5%以内。
注意：合并操作本身有开销。如果你每释放一个块都做一次完整合并，性能会下降。我建议批量合并——攒够一定数量的释放操作，或者显存碎片率达到阈值（比如30%），再触发一次全局合并。
实战中的碎片整理策略
我整理了一份常用策略对比表，方便你根据场景选择：
| 策略 | 适用场景 | 碎片率控制 | 性能开销 |
|---|---|---|---|
| 延迟回收+水位线 | 训练场景，迭代稳定 | 10-15% | 低 |
| 伙伴系统 | 大小请求均匀 | <5% | 中 |
| 全局压缩 | 碎片严重，紧急整理 | <1% | 高（需暂停训练） |
| 分代回收 | 生命周期差异大 | 8-12% | 中低 |
我个人最常用的是延迟回收+伙伴系统的组合。延迟回收处理临时张量，伙伴系统管理长期驻留的模型参数和优化器状态。这样既保证了低延迟，又控制了碎片率。
你想想看，如果每次迭代都做全局压缩，训练速度会慢得没法看。但完全不整理，显存利用率可能只有60%。这个平衡点需要根据你的模型和硬件来调。
核心要点总结：
- 碎片产生是必然的，但可以控制
- 延迟回收能减少CUDA操作次数，但要注意超时
- 合并策略要选对数据结构，红黑树或伙伴系统都行
- 不要追求零碎片，控制在10%以内就够用
这张图展示了完整的显存分配与释放流程。注意看释放路径——不是直接还给CUDA，而是先进延迟队列，等条件满足再合并。这样能有效减少碎片。
最后说一句，碎片整理不是银弹。如果你的模型显存需求波动特别大，比如训练GAN或者做强化学习，那碎片率天然就高。这时候可以考虑预分配——在训练开始前就把可能用到的显存都申请好，避免运行时频繁分配。当然，这需要你对模型的内存使用有精确的预估。
嗯，碎片整理这块就聊到这儿。记住一个原则：能不整理就不整理，要整理就批量整理。性能开销和碎片率之间，找到你的平衡点。