一、显存基础：GPU 显存架构、分配回收与溢出原因

各位同学，咱们今天聊分布式训练的第一个大坎儿——显存。

说句实话，做分布式训练，调通信是一方面，但真正卡脖子的，往往是显存。你模型再大，显存放不下，一切白搭。我最早接触分布式训练的时候，就吃过这个亏。一个 8 卡的任务跑了大半天，最后 OOM 崩了。查日志才发现，有一块的显存被一个没释放的中间变量吃掉了 2 个 G。

所以，这一章咱们把显存的基础打牢。只有把 GPU 里这点「地盘」搞清楚，后面讲的各种优化技巧才能落地。

---

1.1 GPU 显存架构：三层存储体系

很多人觉得 GPU 显存就是一个大池子。其实不是。

你把它看作一个分级存储体系，理解起来就顺了。我们可以分成三块：全局内存（Global Memory）、共享内存（Shared Memory）、寄存器（Registers）。

• 全局内存：最大，但最慢。模型权重、优化器状态、训练数据，都存在这里。我建议你把它想象成 CPU 端的「硬盘加内存」——容量大但是慢。
• 共享内存：片上，快，但小。一个线程块（Block）内的线程共用一个共享内存。用来做数据缓存和规约操作特别香。
• 寄存器：最快，最小。每个线程独享。变量多了，寄存器装不下，就会 spill 到全局内存，性能就崩了。

我个人习惯，在写 kernel 或者调框架的时候，心里会默念这个优先级：寄存器 > 共享内存 > 全局内存。训练时，很多显存优化思路，说白了就是想办法把数据往「小又快」的区域赶。

---

1.2 显存分配与回收：别踩 CUDAMalloc 的坑

说完架构，我们来看看显存是怎么分配和回收的。

TensorFlow、PyTorch 这些框架，底层都是走的 CUDA API。分配时调 cudaMalloc，回收时调 cudaFree。但实际跑起来，远没这么简单。

框架的缓存机制

你想想看，每次训练迭代都调 cudaMalloc 申请几 MB 的临时 buffer 是什么感觉？太慢了。所以框架会搞一个显存缓存池（Memory Pool）。申请显存时，框架先看看池子里有没有现成的。有，就拿过来复用。没有，才去调底层 API 申请新的。

我遇到过一个实际问题：某次跑大模型，显存占用盯着看一直在缓慢上涨，训练一天明天就要崩了。查了半天，发现是缓存池里积累了太多大小不一的碎片，新的请求匹配不上，只好再去要新显存。这种现象叫显存碎片化。

PyTorch 里，你可以用 torch.cuda.empty_cache() 清空缓存池，把空间还给其他进程。注意，它不会释放已经分配给张量的显存，只是把缓存池清一清。

# 查看当前显存使用情况
import torch
print(torch.cuda.memory_summary())

# 清空缓存池（释放已缓存但未使用的块）
torch.cuda.empty_cache()

# 查看当前已分配显存（不包括缓存池）
print(torch.cuda.memory_allocated())

还有一个重要概念——显存分配粒度。CUDA 底层的 cudaMalloc 是按页（通常 512 KB 一段）分配的。你申请 1 字节，它给你 512 KB。所以，小张量的浪费很严重。我建议，尽量把小的张量合并一下再申请。

---

1.3 显存溢出常见原因分析

OOM（Out of Memory）大概是分布式训练里最让人头疼的报错了。嗯，这里要注意，OOM 不一定是你模型太大了，很多时候是「管理不当」。

我把常见的原因整理成了一张表，你可以对照看看自己踩过哪个坑：

原因类型	具体表现	常见场景
模型过大	参数规模超过单卡显存上限	千亿参数模型直接往单卡里塞
batch size 过大	前向计算产生过大的激活值缓存	batch size 翻倍，显存翻倍不止
中间变量泄漏	Python 变量未释放，引用计数不为 0	循环内生成张量但赋值给全局变量
优化器状态开销	Adam 需要 2 倍模型大小	使用 AdamW 但没算足状态显存
数据加载缓存	Dataloader 的 pin_memory 用于 GPU 拷贝	num_workers 多，又不开启 `pin_memory`
分布式通信 buffer	all-reduce 需要额外 buffer	使用 NCCL 时，通信 buffer 未预估

说几个我个人经验里踩过的典型坑：

1.3.1 中间变量泄漏：一个常见的低级错误

我曾经在调试一个分布式训练脚本时发现，第二张卡的显存总是比第一张卡多占 200 MB。查了好几遍代码，才发现是在一个循环里，我无意中把 loss 赋给了全局列表，导致 loss.item() 对应的张量一直没法回收。

代码犯傻长这样：

# 错误做法：全局列表持有引用，导致无法释放
global_losses = []
for batch in dataloader:
    loss = model(batch)
    global_losses.append(loss.item())  # .item() 返回标量没问题
    # 但这里如果我直接写 global_losses.append(loss) 就凉了
    loss.backward()
    optimizer.step()

你看看，loss 如果没引用，反向传播完就被释放了。但一旦被全局列表持有，训练迭代到第 100 步时，你还存着 100 个 loss 张量。能不崩吗？

1.3.2 优化器状态：被低估的显存消耗

这点真的太容易被忽视了。你想想看，Adam 优化器需要给每个参数存两样东西：一阶梯度和二阶梯度。假设你模型参数是 10 GB，优化器状态就得占 20 GB。加上模型本身 10 GB，前向激活值再占 5~10 GB，一张 80 GB 的 A100 就快吃不消了。

我建议在模型初始化之后，先打印一下显存占用，做个心理预期：

model = MyLargeModel().cuda()
optimizer = torch.optim.Adam(model.parameters())

# 打印模型 + 优化器状态占用的显存
model_mem = sum(p.numel() * p.element_size() for p in model.parameters())
optim_mem = model_mem * 3  # param + exp_avg + exp_avg_sq

print(f"模型参数占用: {model_mem / 1024**3:.2f} GB")
print(f"Adam 状态占用: {optim_mem / 1024**3:.2f} GB")
print(f"合计: {(model_mem + optim_mem) / 1024**3:.2f} GB")

后面我会讲到用 ZeRO 优化器 怎么分散存储这些状态，但这里你先记住：计算显存预算时，别忘了优化器状态。

1.3.3 分布式通信 buffer：隐蔽的显存黑洞

分布式训练里，NCCL 在做 all-reduce 通信时，会在每张卡上申请一块额外的 buffer 用于通信。这块缓冲区的大小跟模型参数的总大小基本一致。

假设你用的是数据并行，每张卡都有模型参数的完整副本。all-reduce 梯度时，NCCL 需要一个跟梯度同样大的 buffer 来通信。这等于显存消耗又翻了一倍。

我见过不止一次，单机 8 卡跑大模型，模型 + 优化器 + 激活值刚好占满显存，all-reduce 起步后直接 OOM。你查代码怎么都查不到问题，逼得你去看 NCCL 的文档，看完才明白是通信 buffer 不够了。

# 可以通过环境变量设置 NCCL 通信 buffer 的大小（单位：B）
# 默认一般是 20 MB，如果你的模型非常大，可以适当调低
import os
os.environ["NCCL_IB_TIMEOUT"] = "22"  # 通信超时
os.environ["NCCL_BUFFSIZE"] = "20971520"  # 20 MB 够小模型，大模型需要 64 MB 以上

---

本章小结

这一章，咱们聊了三件事：

• GPU 显存三层架构：全局内存（大而慢）、共享内存（中而快）、寄存器（小而极速）。优化思路全在「把数据往快的地儿搬」。
• 显存分配与回收机制：框架有缓存池，cudaMalloc 有分配粒度，小张量浪费显存。
• 显存溢出常见原因：模型大、batch 大、变量泄漏、优化器状态、通信 buffer。我把它们记在了一张表里。

我刚开始做分布式训练的时候，显存这块是踩坑最多的。有时候明明卡着 79.8 GB 在跑，多存一个日志就 OOM。后来我养成了一个习惯——每加一个中间变量，就问自己一句：它什么时候释放？它的生命周期跟谁的引用有关？

嗯，从小处做起，后面的章节里我们再深入聊那些重量级的显存技巧，比如重计算、ZeRO、模型并行。但基础不牢，后面用什么都容易翻车。希望这一章帮你把地基打好了。