2、显存基础：GPU显存架构、显存分配机制、常见显存消耗来源
在排查显存溢出（OOM）之前，必须先建立对 GPU 显存“物理结构—分配机制—消耗来源”三层认知。很多工程师一遇到 OOM 就盲目调小 batch size，却不知道显存究竟被谁吃掉，结果要么治标不治本，要么牺牲了吞吐量。本节把显存拆开来看，为后续的定位与 ZeRO 调优打下基础。
一、GPU 显存架构：不止是一块“大内存”
从硬件视角看，一块 NVIDIA GPU 的显存（HBM/GDDR）并不是一个均质的池子。它由显存控制器管理，通过高带宽总线与 SM（流多处理器）相连。真正影响训练的是“可用显存”与“碎片化”两个概念：
- 总显存：如 A100 80GB、V100 32GB，是物理上限。
- 已分配显存：PyTorch 缓存分配器向驱动申请、当前被张量占用的部分。
- 已保留显存：分配器缓存起来、暂未归还给驱动的部分，用于加速后续分配。
- 碎片：由于反复申请/释放不同大小块，导致大块连续显存被切碎，即使总量够也可能 OOM。
此外，CUDA 上下文、cuDNN/cuBLAS 句柄、通信缓冲区（NCCL）也会占用一部分显存，这部分常被忽略，却可能达到数百 MB 到数 GB。
二、显存分配机制：allocated 与 reserved 的区别
PyTorch 默认使用缓存分配器（Caching Allocator）。当你创建一个张量时，分配器先看缓存里有没有合适大小的空闲块，有就直接复用，没有才向 CUDA 驱动申请。释放张量时，显存不会立刻还给驱动，而是回到缓存池。这带来两个关键指标：
torch.cuda.memory_allocated()：当前张量实际占用的显存。torch.cuda.memory_reserved()：分配器向驱动保留的总显存，通常大于 allocated。
OOM 报错时，PyTorch 会打印这两者以及碎片情况。如果 reserved 远大于 allocated，说明碎片或缓存未释放是主因；如果两者接近且都逼近上限，说明真实需求确实超了。
import torch
print(torch.cuda.memory_allocated() / 1024**3, "GB allocated")
print(torch.cuda.memory_reserved() / 1024**3, "GB reserved")
print(torch.cuda.memory_summary())
# 手动释放缓存（谨慎使用，会降低后续分配速度）
torch.cuda.empty_cache()
memory_summary() 中的 allocated 与 reserved 差值。差值大，优先怀疑碎片和未释放的中间变量；差值小，则要削减真实显存需求。三、常见显存消耗来源与量级估算
训练一个模型，显存主要被四类对象占据。以参数量为 P、精度为 fp16（2 字节）为例，可以做一个快速估算：
| 消耗来源 | 典型大小 | 是否随 batch 增长 | 是否随序列长度增长 | 优化手段 |
|---|---|---|---|---|
| 模型参数 | 2P（fp16）或 4P（fp32） | 否 | 否 | 混合精度、量化、ZeRO-3 分片 |
| 梯度 | 约 2P（fp16） | 否 | 否 | 梯度累积、ZeRO-2 分片 |
| 优化器状态 | Adam 约 8P（fp32 动量+方差） | 否 | 否 | ZeRO-1/2/3、8-bit Adam |
| 激活值 | 与层数、batch、序列长度成正比 | 是 | 是 | 梯度检查点、序列并行、ZeRO-3 |
| 临时/通信缓冲 | 数百 MB ~ 数 GB | 部分 | 部分 | 调整 bucket、减少 all-gather 频率 |
以一个 7B 参数模型、fp16 混合精度训练为例：参数约 14GB，梯度约 14GB，Adam 优化器状态（fp32）约 56GB，仅这三项就接近 84GB，还没算激活值。这就是为什么单卡 80GB 也常常 OOM，必须借助 ZeRO 把优化器状态和梯度分片到多卡。
四、案例与误区复盘
案例：某团队在 8×A100 40GB 上训练 13B 模型，batch size 设为 4 仍 OOM。查看 memory_summary() 发现 reserved 比 allocated 高出 6GB，且碎片严重。原因是他们在训练循环里反复创建临时张量并拼接，导致缓存池被切碎。改为预分配缓冲区并复用后，OOM 消失，batch size 还能提到 8。
empty_cache() 能解决一切。它只归还缓存，不减少真实需求，频繁调用反而拖慢训练。误区三：忽略通信缓冲区。NCCL 的 all-reduce 缓冲、梯度 bucket 在 ZeRO 下会显著增加显存占用。建立“架构—分配—来源”的完整认知后，下一步才能用 torch.cuda.memory_summary()、NVIDIA nvidia-smi、PyTorch Profiler 等工具精确定位是哪一类对象在膨胀，从而选择对应的 ZeRO Stage 与调优参数。