显存池基础：什么是显存池、为什么需要显存池、显存池的核心设计目标

大家好，我是你们这门课的主讲。今天咱们来聊聊显存池。

说实话，我刚开始做深度学习训练的时候，根本没想过显存还要“池”一下。那时候模型小，显存够用，直接 cudaMalloc 一把梭。直到后来做大模型，显存动不动就爆，我才意识到——嗯，这里头有大学问。

什么是显存池？

显存池，说白了就是一个预先分配好的显存“仓库”。

你想想看，平时我们写 CUDA 程序，需要显存了，就调用 cudaMalloc 去申请；用完了，再 cudaFree 释放。这就像每次做饭都要去超市买一棵葱，用完就扔掉。是不是很浪费？

显存池的做法是：一次性从 GPU 那里申请一大块显存，放在池子里。谁要用，就从池子里“借”一块；用完了，再“还”回来。这样就不用频繁跟 GPU 驱动打交道了。

核心思想： 把显存分配从“按需申请”变成“池化管理”。

我在项目中遇到过这样一个场景：训练一个 70 亿参数的模型，每次迭代都要创建和销毁几十个中间张量。如果用原始的 cudaMalloc/Free，光分配开销就能吃掉 30% 的训练时间。换成显存池后，速度直接翻倍。

为什么需要显存池？

这个问题，我分三点来讲。

1. 减少分配开销

cudaMalloc 这个操作，可不是简单的“从内存里划一块出来”。它涉及到 GPU 驱动、内存管理、页表更新等一系列操作。我做过测试，一次 cudaMalloc 大概需要 10-50 微秒。听起来不多对吧？

但你要知道，在大模型训练中，一次前向+反向传播可能要申请释放上千次显存。算下来，光分配开销就能达到几十毫秒。而显存池的分配操作，只需要修改几个指针，纳秒级别就能完成。

2. 避免显存碎片

这是个很头疼的问题。我记得有一次，明明显存还剩 4GB，但申请一个 3GB 的张量却失败了。为什么？因为显存被切成了无数个小碎片，没有一块连续空间能放下 3GB。

显存池通过预分配大块连续内存，然后自己管理内部碎片，能有效缓解这个问题。当然，不能完全消除，但比裸用 cudaMalloc 好太多了。

3. 支持内存复用

这是显存池最大的价值。训练过程中，很多张量的生命周期是错开的。比如 A 张量用完后，B 张量才开始用。显存池可以检测到这种模式，把 A 的显存直接给 B 用，省去重复分配的开销。

我的经验： 在混合精度训练中，显存池的复用率能达到 60%-80%。这意味着你只需要 20%-40% 的额外显存，就能跑原本需要双倍显存的模型。

显存池的核心设计目标

设计一个显存池，不是随便搞个链表就行的。我总结了四个核心目标：

设计目标 说明 我的踩坑经历
低延迟 分配和释放操作要在纳秒级完成 曾经用红黑树管理空闲块，结果分配延迟飙到微秒级，被老板骂了一顿
高利用率 尽量减少内部碎片和外部碎片 有个版本为了追求速度，每个块都按 2MB 对齐，结果碎片率高达 40%
线程安全 多线程同时分配时不能出问题 第一次写显存池没加锁，结果两个线程拿到了同一块显存，数据全乱了
可扩展性 支持不同大小的显存块，适配各种算子 早期只支持固定大小块，遇到变长张量就傻眼了

这四个目标之间是有 trade-off 的。比如，想提高利用率，就得做更精细的内存管理，但延迟就会上升。想降低延迟，用最简单的 free list，但碎片率就会增加。

我个人习惯的做法是：先满足延迟要求，再优化利用率。因为在大模型训练中，延迟是硬指标，利用率可以慢慢调。

显存池的典型架构

下面这张图是我自己画的，展示了显存池的核心结构：

显存池核心架构图 分配接口：allocate(size) / deallocate(ptr) 内存管理器 空闲块列表（Free List） | 哈希索引 | 缓存策略 线程安全锁 | 碎片整理 | 大小分级 预分配的显存块（GPU Memory Chunks） 64KB 256KB 1MB 4MB 16MB+ 图1：显存池三层架构——接口层、管理层、存储层

从图上你能看到，显存池分三层：

  • 接口层： 对外提供 allocate 和 deallocate 两个函数，用户不用关心底层实现
  • 管理层： 核心逻辑都在这里，包括空闲块管理、大小分级、线程安全等
  • 存储层： 真正从 GPU 申请的大块显存，按大小分成不同等级

注意： 显存池不是万能的。如果你的模型显存需求变化特别剧烈（比如一会儿要 1GB，一会儿要 10GB），显存池的利用率会下降。我曾经在一个动态图模型上吃过这个亏，后来加了自适应扩容机制才解决。

一个简单的显存池实现思路

这里给个伪代码，让你感受一下显存池的基本操作：

class MemoryPool {
    // 预分配一大块显存
    void init(size_t total_size) {
        cudaMalloc(&base_ptr, total_size);
        free_list.push_back({base_ptr, total_size});
    }
    
    // 从池中分配
    void* allocate(size_t size) {
        lock();
        // 找到第一个足够大的空闲块
        block = find_first_fit(size);
        if (block) {
            // 切分：用掉一部分，剩下的放回空闲列表
            split_block(block, size);
            unlock();
            return block->ptr;
        }
        // 没有足够大的块，触发扩容或报错
        unlock();
        return nullptr;
    }
    
    // 归还到池中
    void deallocate(void* ptr) {
        lock();
        // 把块放回空闲列表，尝试合并相邻块
        merge_adjacent(ptr);
        unlock();
    }
};

这段代码虽然简单，但已经包含了显存池的核心逻辑：预分配、查找、切分、合并。实际工程中，你还需要考虑对齐、缓存、线程安全等细节。

避坑指南： 我曾经在 deallocate 里忘记做相邻块合并，结果跑了一天训练后，显存池里全是小碎片，再也分配不出大块了。所以合并操作一定要做，而且要做对。

好了，关于显存池的基础知识就讲到这里。记住三个关键词：预分配、复用、减少碎片。后面的章节，我们会深入具体的实现细节和优化技巧。

想要更多资料，联系可可

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