一、显存基础与瓶颈分析
大家好，我是你们的老朋友。今天咱们聊聊显存——这个在分布式训练里最让人头疼的东西。说实话，我入行那会儿，显存还只是个“够用就行”的参数，现在呢？它直接决定了你能不能跑起一个模型。
先看一张图，帮你快速建立本章的知识框架：
1.1 显存架构：HBM 与 GDDR
显存，说白了就是 GPU 的“工作台”。你模型的所有参数、梯度、中间结果，都得摆在这个台子上才能算。目前主流的有两种：HBM 和 GDDR。
HBM（High Bandwidth Memory），高端显卡和 AI 加速器在用。它通过 3D 堆叠和硅中介层，把多个 DRAM 芯片叠在一起，带宽极高。我手头这块 A100 用的就是 HBM2e，带宽超过 2 TB/s。嗯，听着就爽。
GDDR（Graphics Double Data Rate），消费级显卡的主力。像 RTX 3090 用的 GDDR6X，带宽也能到 1 TB/s 左右。但跟 HBM 比，功耗和延迟上还是差一截。
为什么会这样？因为 HBM 的位宽特别宽——1024 位甚至更高，而 GDDR 通常只有 384 位。位宽越宽，一次能搬的数据就越多。你想想看，同样是跑一个 175B 的模型，HBM 的加载速度能快出一倍不止。
关键区别：
- HBM：高带宽、低功耗、贵。适合 AI 训练。
- GDDR：性价比高、容量大。适合推理和消费级场景。
1.2 显存带宽与容量
带宽和容量，这两个指标经常被混为一谈。我见过不少新人，只看容量不看带宽，结果模型是装下了，但训练慢得像蜗牛。
这里有个公式你得记住：
计算强度 = 计算量（FLOPs） / 访存量（Bytes）如果计算强度低，说明你的模型是“访存密集型”——说白了就是带宽不够用，GPU 计算单元在等数据。我在项目中遇到过一个大模型，明明算力够，但训练速度就是上不去。一查，带宽利用率只有 30%。
| 显存类型 | 典型带宽 | 典型容量 | 适用场景 |
|---|---|---|---|
| HBM2e | 2 TB/s | 40-80 GB | A100, MI250 |
| HBM3 | 3.2 TB/s | 80-144 GB | H100, B200 |
| GDDR6X | 1 TB/s | 24-48 GB | RTX 4090 |
我的经验：选卡时别只看容量。带宽不够，大模型训练时你会哭的。我曾经用 GDDR6 跑一个 13B 模型，batch size 调到 1 都跑不动——不是容量不够，是带宽卡脖子。
1.3 训练中的显存占用分析
训练时显存都去哪了？我把它分成四块：
- 参数：模型权重。一个 1B 参数的模型，FP16 下占 2 GB。
- 梯度：跟参数一样大。反向传播时用。
- 优化器状态：Adam 要存一阶动量和二阶动量，每个参数多占 8 字节（FP32）。
- 激活值：前向传播的中间结果。这个最吃显存，尤其是大 batch size 时。
举个例子，一个 7B 模型用 Adam 优化器训练：
参数（FP16）：7B × 2B = 14 GB
梯度（FP16）：7B × 2B = 14 GB
优化器状态（FP32）：7B × 4B × 2 = 56 GB
激活值：取决于 batch size 和序列长度，通常 20-50 GB
合计：约 100-120 GB你看，光优化器状态就占了快一半。这也是为什么很多人用混合精度训练——能省下不少显存。
注意：激活值的显存占用是动态的，跟 batch size 和模型深度强相关。我曾经因为激活值没算好，直接 OOM 了。建议用 torch.cuda.max_memory_allocated() 实时监控。
1.4 显存瓶颈定位方法
怎么找到显存瓶颈？我一般用三板斧：
- Profiling：用 PyTorch Profiler 或 Nsight Systems 看显存分配曲线。哪个阶段显存飙升，一目了然。
- 内存快照：torch.cuda.memory_summary() 能列出每个张量的占用。我习惯在训练循环里定期打印。
- 计算强度分析：用 roofline 模型判断是计算瓶颈还是访存瓶颈。如果带宽利用率低，优先优化数据加载和通信。
举个例子，我调过一个分布式训练任务，发现通信阶段显存暴涨。一查，是 all-reduce 时梯度缓冲区没复用，每次重新分配。改成预分配后，显存占用降了 30%。
避坑指南：我曾经在 8 卡训练时，每张卡都单独分配了优化器状态，结果显存直接爆了。后来用了 ZeRO 优化器，把状态分片到各卡，才解决问题。记住：分布式不是简单的“卡多显存多”，通信和显存要协同优化。
好了，这一章就到这里。显存这东西，说白了就是“空间换时间”的游戏。你理解了它的架构、占用和瓶颈，后面讲多卡通信时才能游刃有余。