4、显存池预分配策略:固定大小池、动态增长池、混合策略

显存池怎么分配,这事儿我琢磨了好几年。

说白了,预分配策略就是决定「什么时候向GPU要内存,要多少」。策略选对了,你的训练任务能快30%以上。选错了,要么显存炸了,要么性能拉胯。

我个人习惯把主流策略分成三类:固定大小池、动态增长池、混合策略。咱们一个一个聊。

4.1 固定大小池:简单粗暴,但够稳

固定大小池的思路很直接——程序启动时,一次性向GPU申请一大块显存。比如8GB,切分成若干固定大小的块。后面所有张量分配,都从池子里拿。

核心特点: 分配速度极快,O(1)时间复杂度。没有碎片问题。

我在项目中遇到过这样一个场景:一个推理服务,每次请求的输入尺寸固定。用固定大小池,分配延迟从微秒级降到了纳秒级。你想想看,这差别有多大。

// 固定大小池的简化实现
class FixedPool {
    char* base;          // 显存基地址
    size_t block_size;   // 每个块大小,比如 4MB
    int num_blocks;      // 块数量,比如 2048
    std::vector<int> free_list;  // 空闲块索引
    
public:
    void* allocate() {
        if (free_list.empty()) return nullptr;
        int idx = free_list.back();
        free_list.pop_back();
        return base + idx * block_size;
    }
    
    void deallocate(void* ptr) {
        int idx = (char*)ptr - base / block_size;
        free_list.push_back(idx);
    }
};
我的建议: 如果你的工作负载很稳定,比如固定batch size的训练,固定大小池是最优解。别折腾。

但缺点也很明显——灵活性差。一旦遇到超大张量,池子里的块不够大,就得走fallback路径。我曾经因为这个踩过坑,一个模型升级后参数变大了,固定池直接罢工。

4.2 动态增长池:灵活,但有代价

动态增长池的策略是「按需分配,用多少要多少」。池子初始很小,比如256MB。当分配请求超过当前容量时,自动向GPU申请更大的块。

嗯,这里要注意。动态增长不是无限增长。一般会设置一个上限,比如总显存的80%。超过上限就抛出OOM错误。

避坑指南: 我曾经在动态增长池里忘记设置上限,结果一个内存泄漏把整张卡搞崩了。一定要设硬上限!

动态增长池的实现,核心是「增长因子」。我见过两种主流做法:

增长策略 倍数 适用场景
线性增长 每次+512MB 内存需求缓慢增长
指数增长 每次×2 内存需求波动大
自适应增长 根据历史分配量动态调整 复杂工作负载

我个人习惯用指数增长,但增长因子设为1.5而不是2。为什么?2倍太激进,容易浪费。1.5倍是个不错的折中。

// 动态增长池的核心逻辑
class DynamicPool {
    size_t current_size = 256 * 1024 * 1024;  // 初始256MB
    size_t max_size = 8LL * 1024 * 1024 * 1024; // 上限8GB
    float growth_factor = 1.5f;
    
    void grow(size_t required_size) {
        size_t new_size = current_size;
        while (new_size < required_size) {
            new_size = static_cast<size_t>(new_size * growth_factor);
        }
        if (new_size > max_size) {
            throw std::runtime_error("OOM: 超过显存上限");
        }
        // 向GPU申请新显存
        cudaMalloc(&new_ptr, new_size);
        current_size = new_size;
    }
};
小技巧: 动态增长池配合「延迟释放」效果更好。释放的显存先不还给GPU,留在池子里给后续分配用。

4.3 混合策略:取长补短

实际工程中,纯固定或纯动态都不够完美。混合策略才是王道。

我常用的混合策略是这样的:

  1. 预分配一个固定大小的核心池,覆盖90%的常见分配请求
  2. 预留一个动态增长区,处理那些超大或超小的异常请求
  3. 设置水位线,当固定池使用率超过80%时,提前扩容动态区

说白了,就是用固定池保证性能,用动态池兜底。

核心思想: 80%的场景用固定池(O(1)分配),20%的异常场景用动态池(O(log n)分配)。整体性能接近固定池,灵活性接近动态池。

我在一个推荐系统项目中实践过混合策略。训练时batch size偶尔会波动,固定池处理大部分情况,动态池处理峰值。最终显存利用率从65%提升到了92%。

4.4 三种策略对比

做个总结,方便你选型:

策略 分配速度 内存利用率 灵活性 实现复杂度
固定大小池 ⭐⭐⭐⭐⭐ ⭐⭐ ⭐ ⭐
动态增长池 ⭐⭐⭐ ⭐⭐⭐⭐ ⭐⭐⭐⭐⭐ ⭐⭐⭐
混合策略 ⭐⭐⭐⭐ ⭐⭐⭐⭐⭐ ⭐⭐⭐⭐ ⭐⭐⭐⭐⭐

你想想看,如果你的场景是「模型固定、batch固定」,固定池就够了。如果是「多模型混跑、负载波动大」,混合策略更靠谱。

4.5 知识体系图

下面这张图,把三种策略的核心逻辑串起来了:

显存池预分配策略知识体系 显存池预分配策略 固定大小池 动态增长池 混合策略 O(1)分配 无碎片 按需分配 增长因子 延迟释放 固定池+动态区 水位线预警 性能+灵活 选型原则:稳定场景用固定,波动场景用混合

最后说一句,没有银弹。我见过有人把固定池玩出花来,也见过混合策略搞得太复杂反而性能下降。关键还是理解你的业务场景,选最合适的策略。

想要更多资料,联系可可