一、显存分配基础：GPU显存架构、CUDA内存模型、显存分配与释放机制、显存碎片化问题

大家好，我是你们这节课程的讲师。今天咱们聊聊模型并行里最基础、也最绕不开的话题——显存分配。说白了，你做大模型训练，显存就是你的“战场”。怎么分配、怎么回收、怎么避免踩坑，这些搞不清楚，后面跑分布式训练会非常痛苦。

我刚开始接触大规模训练时，就吃过显存分配的亏。有一次跑一个百亿参数模型，明明算好了显存够用，结果跑着跑着就OOM了。查了半天，发现是显存碎片化在作祟。嗯，从那以后，我就把显存分配机制彻底研究了一遍。今天把这些经验分享给你。

核心观点：显存分配不是简单的malloc/free，它涉及GPU硬件架构、CUDA驱动层、以及运行时库的三层协作。理解这三层，你才能写出高效的分布式训练代码。

1.1 GPU显存架构：从硬件说起

先看看GPU显存长什么样。我习惯把GPU显存想象成一个多层级的存储系统，跟CPU的缓存体系有点像，但更暴力、更直接。

GPU显存架构层级图 全局显存 (Global Memory) — 16GB~80GB HBM2e / HBM3，带宽可达 2TB/s+ L2 缓存 — 40MB~80MB 共享显存 (Shared Memory) + L1 缓存 — 每个SM 48KB~228KB 寄存器 (Registers) — 每个SM 65536个，每个线程最多255个 注：越往上容量越大、延迟越高；越往下容量越小、延迟越低

这张图我画了很多次，每次给团队新人讲都先用它。你想想看，全局显存是最大的，但访问延迟也最高。共享显存在SM内部，快得多，但容量有限。做模型并行时，你要频繁在这几层之间搬数据，怎么搬、搬多少，直接决定了训练效率。

我的经验：在A100上做训练时，我习惯把模型参数尽量留在全局显存里，把中间激活值放到共享显存里计算。这样能减少全局显存的带宽压力。但要注意，共享显存只有几十KB，别塞太多东西进去。

1.2 CUDA内存模型：软件视角

硬件看完了，咱们看看软件层面。CUDA把显存分成了几个逻辑区域，每个区域有不同的生命周期和访问权限。我刚开始学的时候，总觉得这些概念很抽象，后来踩了几个坑才真正理解。

内存类型 分配方式 生命周期 访问范围 典型用途
全局显存 cudaMalloc 程序显式释放 所有线程 模型参数、数据
共享显存 __shared__ 线程块内 同一线程块 中间结果缓存
寄存器 局部变量 线程内 单个线程 临时计算
常量显存 __constant__ 程序显式释放 所有线程（只读） 超参数、权重
纹理显存 cudaBindTexture 程序显式释放 所有线程（只读） 图像数据

这里有个坑，我必须要提醒你。全局显存用cudaMalloc分配，但释放必须用cudaFree。如果你忘了释放，或者释放顺序不对，就会造成显存泄漏。我曾经在一个分布式训练脚本里，因为一个循环中反复cudaMalloc但没有及时cudaFree，结果跑了两个小时就OOM了。排查起来特别痛苦。

注意：CUDA的cudaMalloc是同步操作，会阻塞CPU直到分配完成。而cudaFree是异步的，它不会立即释放显存，而是等到所有相关操作完成后才释放。这个特性在分布式训练中尤其重要，搞不好就会导致显存占用居高不下。

1.3 显存分配与释放机制

好，咱们深入看看分配和释放的具体机制。CUDA驱动层维护了一个显存分配器，它管理着GPU上的物理显存。当你调用cudaMalloc时，分配器会从空闲列表中找一块足够大的连续空间给你。如果没有合适的，它会尝试合并相邻的空闲块，或者向操作系统申请更多显存。

我个人的习惯是，在训练开始前就预分配好所有需要的显存，而不是边训练边分配。为什么呢？因为频繁的cudaMalloc和cudaFree会导致显存碎片化，后面我会详细讲。

// 一个典型的显存分配示例
float* d_params;
float* d_grads;
float* d_optimizer_state;

// 预分配所有显存
cudaMalloc(&d_params, params_size * sizeof(float));
cudaMalloc(&d_grads, params_size * sizeof(float));
cudaMalloc(&d_optimizer_state, params_size * 2 * sizeof(float)); // Adam需要两个状态

// 训练循环中不再分配显存
for (int epoch = 0; epoch < num_epochs; epoch++) {
    // 前向、反向、更新...
}

// 训练结束后统一释放
cudaFree(d_params);
cudaFree(d_grads);
cudaFree(d_optimizer_state);

你看，这样写的好处是显存分配只发生一次，后续的训练过程不会因为分配显存而产生额外开销。而且，避免了在循环中反复分配释放导致的碎片问题。

最佳实践：在分布式训练中，我建议使用显存池（Memory Pool）技术。比如NVIDIA的cudaMemPool，或者PyTorch的torch.cuda.memory.CUDAPluggableAllocator。它们能复用显存块，减少分配次数，降低碎片化。

1.4 显存碎片化问题

说到碎片化，这可是个大问题。你想想看，显存就像一块大蛋糕，你切一块、我切一块，切得多了，就会留下很多小缝隙。这些小缝隙单独看都不够大，但加起来却占了不少空间。结果就是，你明明还有空闲显存，但就是分配不出一块连续的大空间。

为什么会这样？主要有两个原因：

  • 外部碎片：频繁分配和释放不同大小的显存块，导致空闲空间不连续。比如你先分配了100MB，又分配了200MB，然后释放了100MB的那块。这时候空闲空间是100MB，但如果你要分配150MB的连续空间，就分配不了。
  • 内部碎片：分配器为了对齐，会多分配一些空间。比如你申请了100字节，但分配器可能给你128字节，多出来的28字节就是内部碎片。单个看起来不多，但累积起来也很可观。

我记得有一次，在训练一个GPT-3规模的模型时，显存占用率显示只有85%，但就是跑不起来，一直报OOM。我查了半天，发现是碎片化导致的。后来我用cudaMemGetInfo查看了显存使用情况，发现空闲空间被分割成了几十个小块，最大的连续块只有不到2GB。

// 查看显存使用情况的代码
size_t free_bytes, total_bytes;
cudaMemGetInfo(&free_bytes, &total_bytes);
printf("总显存: %.2f GB\n", total_bytes / 1024.0 / 1024.0 / 1024.0);
printf("空闲显存: %.2f GB\n", free_bytes / 1024.0 / 1024.0 / 1024.0);
printf("已用显存: %.2f GB\n", (total_bytes - free_bytes) / 1024.0 / 1024.0 / 1024.0);

避坑指南：我曾经在训练一个混合精度模型时，因为频繁创建和销毁临时张量，导致显存碎片化严重。后来我改用torch.cuda.empty_cache()来清理缓存，但效果有限。最终解决方案是使用显存池，把所有临时张量都从池里分配，训练结束后统一回收。碎片率从30%降到了5%以下。

总结一下，显存碎片化是分布式训练中一个非常隐蔽但破坏力巨大的问题。它不会直接报错，但会让你莫名其妙地OOM。解决思路有三个：

  1. 预分配策略：训练开始前就把所有需要的显存分配好，避免训练过程中频繁分配释放。
  2. 显存池技术：使用CUDA的显存池或PyTorch的自定义分配器，复用显存块。
  3. 碎片整理：定期调用cudaMemGetInfo监控碎片情况，必要时手动整理。但这个方法比较暴力，不推荐在生产环境使用。

好了，这一节的内容就到这里。显存分配看起来是个基础问题，但里面门道很多。你把这些搞清楚了，后面讲模型并行、流水线并行时，才能理解为什么有些策略能省显存、有些策略会浪费显存。记住，显存是你最宝贵的资源之一，用好它，你的训练效率能翻倍。

想要更多资料，联系可可

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