3. 显存池配置参数
好，咱们接着聊显存池的配置。说实话，这部分内容我当年刚接触时也觉得有点枯燥，不就是几个参数嘛。但后来在项目中吃过亏，才明白这几个参数调不好，轻则显存浪费，重则直接OOM崩掉。
显存池的配置，说白了就是四个核心参数：初始池大小、最大池限制、扩展步长、空闲块回收阈值。它们共同决定了显存池怎么分配、怎么增长、怎么回收。我一个个来讲。
3.1 初始池大小
初始池大小，就是显存池刚创建时预分配多少显存。这个值设多大，挺有讲究的。
设小了会怎样？ 模型加载时频繁触发扩展操作，每次扩展都要走cudaMalloc，开销不小。我在项目中遇到过，初始池设了512MB，但模型需要2GB，结果训练前几分钟一直在做扩展，GPU利用率上不去。
设大了呢？ 显存白白占着，别的进程用不了。尤其多卡训练时，每张卡都预分配一大块，其他任务直接没资源了。
我个人习惯的做法是：先跑一次小batch，观察峰值显存占用，然后取这个值的80%作为初始池大小。比如峰值是2.5GB，初始池就设2GB。留点余量，但别太多。
经验公式： 初始池大小 ≈ 模型参数显存 + 最大batch的激活显存 × 0.7
3.2 最大池限制
最大池限制，就是显存池最多能膨胀到多大。这个参数其实是个安全阀。
为什么要设这个？你想想看，如果没有上限，显存池会一直扩展，直到把GPU显存吃光。到时候系统OOM，整个训练进程被kill掉，哭都来不及。
我建议设成GPU总显存的90%～95%。留5%～10%给系统和其他进程用。比如24GB的卡，最大池限制设22GB左右。别贪心，留点余地。
注意： 最大池限制不要超过GPU显存总量。我曾经见过有人设了和显存一样大，结果系统显存被占满，连驱动都挂了，只能重启机器。
3.3 扩展步长
扩展步长，就是每次显存不够用时，一次性扩展多少显存。这个参数直接影响扩展频率和显存碎片。
步长太小： 频繁扩展，每次开销大，而且容易产生大量小碎片。我之前调一个NLP模型，步长设了128MB，结果训练过程中扩展了上百次，显存碎片率高达30%。
步长太大： 一次扩展太多，显存浪费严重。比如只需要256MB，结果一次扩展了2GB，后面用不完就闲置了。
我个人的经验是：扩展步长设为初始池大小的10%～20%。比如初始池2GB，步长就设256MB～512MB。这样既不会频繁扩展，也不会浪费太多。
| 初始池大小 | 推荐扩展步长 | 适用场景 |
|---|---|---|
| 512MB | 64MB～128MB | 小模型、单卡训练 |
| 2GB | 256MB～512MB | 中等模型、单卡/多卡 |
| 8GB | 1GB～2GB | 大模型、多卡训练 |
3.4 空闲块回收阈值
这个参数挺有意思的。显存池里会有一些空闲块，是之前分配过但已经释放的。回收阈值决定了这些空闲块什么时候被真正归还给系统。
阈值设得低： 空闲块很快被回收，显存利用率高，但下次再分配时又要重新申请，开销大。
阈值设得高： 空闲块留在池里，下次分配直接复用，速度快，但显存占用高。
我一般这样设：空闲块回收阈值 = 最大池限制 × 5%。也就是说，当空闲块总量超过最大池的5%时，开始回收。这样既保证了复用效率，又不会让空闲块堆积太多。
小技巧： 如果你发现训练过程中显存占用忽高忽低，波动很大，很可能是回收阈值设得太低了。调高一点，让空闲块多留一会儿，显存波动会平滑很多。
3.5 四个参数的关系
这四个参数不是孤立的，它们互相影响。我画个图帮你理清关系：
3.6 实战调参建议
说了这么多理论，来点实际的。我总结了一套调参流程，你可以照着试试：
- 先确定最大池限制：设为GPU显存的90%～95%。这是底线，不能动。
- 再设初始池大小：跑一次小batch，取峰值显存的80%。
- 扩展步长取中间值：初始池的15%左右。后续根据扩展次数调整。
- 回收阈值先设5%：观察显存波动，如果波动大就调高到8%～10%。
避坑指南： 我曾经在调一个BERT-large模型时，初始池设了4GB，步长设了1GB，结果训练到一半显存池扩展到了22GB，直接OOM。后来发现是最大池限制没设，默认无限制。从那以后，我每次配置显存池，第一件事就是设最大池限制。
嗯，这四个参数说完了。其实没那么复杂，核心就是平衡：既要够用，又不能浪费；既要快，又不能碎片太多。多试几次，找到适合你模型的组合就好。