显存池核心设计：内存池架构、预分配策略、块大小管理

好，咱们直接进入正题。显存池这东西，说白了就是给GPU搞一个「内存超市」。你想想看，每次训练模型都去cudaMalloc申请显存，就像每次做饭都去菜市场买菜——来回跑路的时间比做饭还长。我刚开始做分布式训练时，就吃过这个亏。

核心思想：提前申请一大块显存，自己管理分配和回收。避免频繁调用cudaMalloc/cudaFree带来的开销和碎片。

内存池架构：两级结构

我个人习惯把显存池设计成两级：全局池和线程池。为什么这么搞？

  • 全局池：负责大块显存的预分配。比如一次申请2GB，切成若干块。
  • 线程池：每个CUDA流（stream）有自己的小缓存。小请求直接在线程池里拿，不用加锁。

我曾在项目中遇到过一个问题：多个线程同时申请显存，全局锁成了瓶颈。后来改成两级架构，性能直接翻倍。嗯，这里要注意：线程池的大小要有限制，不然显存会被吃光。

// 伪代码：两级池架构
class MemoryPool {
    GlobalPool global;   // 全局池，线程安全
    ThreadLocalPool local[MAX_STREAMS]; // 每个流一个本地池
};

预分配策略：什么时候申请？申请多少？

预分配策略是个技术活。你想想看，申请少了不够用，申请多了浪费。我一般用懒分配 + 按需增长的策略。

策略优点缺点适用场景
全量预分配零延迟分配显存浪费严重显存充足，延迟敏感
懒分配按需使用，节省显存首次分配有延迟显存紧张，批次大小变化大
混合策略平衡延迟和利用率实现复杂大多数生产环境

我曾经踩过一个坑：全量预分配了4GB，结果模型只用了1.2GB，剩下的全浪费了。后来改成懒分配，每次翻倍增长——第一次申请256MB，不够就512MB，再不够1GB。这样既保证了性能，又不会浪费。

我的经验：初始块大小设为模型最大tensor的1.5倍。太小了频繁扩容，太大了浪费。这个值我调了不下20次才找到平衡点。

块大小管理：如何切分显存？

块大小管理是显存池的核心。说白了就是：你手里的显存怎么切？切多大一块？

我常用的策略是伙伴系统（Buddy System）。把显存按2的幂次切分：2MB、4MB、8MB、16MB……这样分配和合并都很方便。

// 伙伴系统块大小示例
const size_t BLOCK_SIZES[] = {
    2 * 1024 * 1024,    // 2MB
    4 * 1024 * 1024,    // 4MB
    8 * 1024 * 1024,    // 8MB
    16 * 1024 * 1024,   // 16MB
    32 * 1024 * 1024,   // 32MB
    64 * 1024 * 1024,   // 64MB
    128 * 1024 * 1024,  // 128MB
    256 * 1024 * 1024,  // 256MB
};

为什么用2的幂次？因为合并方便。比如两个相邻的16MB块，如果都没在用，可以合并成一个32MB块。这个操作在伙伴系统里是O(1)的。

注意：伙伴系统有个问题——内部碎片。比如你申请10MB，系统会给你16MB的块，浪费了6MB。我一般会加一个「小对象缓存」，专门处理小于2MB的请求，减少浪费。

分配与回收：核心流程

分配流程其实不复杂，我画个图你就明白了。

显存池分配与回收流程 申请显存请求 线程池有空闲块？ 是 从线程池分配 否 从全局池分配 返回显存指针

回收流程正好反过来。线程用完了，先还到线程池。如果线程池满了，再还给全局池。全局池会尝试合并相邻的空闲块。

小技巧：回收时别急着合并。我一般设置一个阈值——空闲块数量超过10个才触发合并。这样可以避免频繁合并带来的开销。

实际项目中的坑

讲几个我踩过的坑，你们引以为戒。

  1. 显存泄漏：有一次我忘了在析构函数里释放全局池，结果程序跑了一周，显存爆了。排查了两天才找到。后来我加了引用计数，每个块被引用时才分配，引用归零就回收。
  2. 碎片化：频繁分配释放不同大小的块，会导致显存碎片化。我试过用「最佳适配」算法，效果不错，但性能有损耗。最后用了伙伴系统 + 小对象缓存，算是平衡了。
  3. 多线程竞争：全局池的锁竞争是个大问题。我后来改成读写锁——读操作（分配）可以并发，写操作（合并）才加独占锁。性能提升明显。
// 读写锁示例
class GlobalPool {
    std::shared_mutex mutex;
    
    void* allocate(size_t size) {
        std::shared_lock lock(mutex); // 读锁，可并发
        // 查找空闲块
    }
    
    void deallocate(void* ptr) {
        std::unique_lock lock(mutex); // 写锁，独占
        // 回收并尝试合并
    }
};

性能调优建议

最后给几个调优建议，都是实战经验。

  • 块大小对齐：所有块大小对齐到256字节。为什么？因为GPU的缓存行是128字节，对齐后可以减少缓存未命中。
  • 预分配时机：我一般在模型加载完成后、训练开始前做预分配。这时候显存使用情况最清晰。
  • 监控指标：一定要监控池的使用率、分配延迟、碎片率。我习惯用NVIDIA的Nsight Systems做性能分析。

一句话总结：显存池的核心就是「提前准备、按需分配、及时回收」。架构上两级分离，策略上懒分配+伙伴系统，调优上注意对齐和监控。做到这几点，你的显存池基本就稳了。

好了，这一章就到这里。记住，显存池不是万能的，但用好了能让你的训练效率提升30%以上。我当年优化完显存池后，训练时间从3天缩短到2天，那种成就感，啧，真不错。

← 上一章 📖 返回目录 下一章 →