一、显存基础与通信瓶颈
大家好，我是你们的老朋友。今天咱们聊聊多卡训练里最让人头疼的两个问题——显存墙和通信墙。
说实话，我刚开始做分布式训练那会儿，也踩过不少坑。有一次跑一个百亿参数模型，8张A100跑着跑着就OOM了。排查了半天，发现是显存分配策略出了问题。嗯，从那以后我就特别重视显存这块的基础知识。
1.1 显存墙：你的模型到底吃多少显存？
先问大家一个问题：一个70B参数的模型，用FP16存，需要多少显存？
简单算一下：70B × 2字节 = 140GB。这还没算优化器状态、梯度、激活值呢。你想想看，一张A100才80GB，单卡根本放不下。
这就是所谓的显存墙——模型太大，单卡装不下。
显存消耗的主要来源：
- 模型参数：权重矩阵，FP16下每参数2字节
- 优化器状态：Adam需要2倍参数量的额外显存（动量+方差）
- 梯度：与参数同大小，FP16下也是2字节/参数
- 激活值：这个最容易被忽略，训练时中间结果会占大量显存
我个人习惯用这个公式估算：
总显存 ≈ 参数显存 × (1 + 优化器倍数 + 梯度倍数) + 激活值
以70B模型为例（FP16）：
- 参数：140GB
- Adam优化器：280GB（2倍参数）
- 梯度：140GB
- 激活值：取决于batch size，通常几十GB
合计：轻松超过500GB所以，单卡搞不定，必须上多卡。但多卡又带来了新问题——通信。
1.2 通信墙：NVLink vs PCIe，差距有多大？
多卡训练时，卡之间要频繁交换梯度。这时候通信带宽就成了瓶颈。
我遇到过最夸张的一次，通信时间占到了总训练时间的60%以上。说白了，GPU在等数据，算力白白浪费了。
目前主流的卡间互联方式有两种：
| 特性 | NVLink（以A100为例） | PCIe 4.0 x16 |
|---|---|---|
| 单向带宽 | 600 GB/s（单卡总带宽） | 32 GB/s |
| 双向带宽 | 1200 GB/s | 64 GB/s |
| 延迟 | ~1μs | ~5μs |
| 拓扑 | 全互联或混合立方体 | 树形或环形 |
看到了吗？NVLink比PCIe快了将近20倍。这就是为什么高端服务器都用NVLink。
我的经验：如果只能用PCIe，尽量用梯度压缩或梯度累积来减少通信量。我曾经在一个PCIe 3.0的集群上训练，通信开销大到让人崩溃，后来用了TopK梯度压缩，通信量降了90%，训练速度直接翻倍。
1.3 通信与计算重叠：让GPU不再干等
核心思路很简单：别让GPU闲着。
传统做法是：先算梯度 → 等通信完成 → 再算下一轮。这就像流水线工人，干完活就站着等。
更好的做法是：算一部分梯度，就开始通信这部分，同时继续算剩下的梯度。这就是通信与计算重叠。
实现方式：
- 异步AllReduce：用CUDA stream将通信操作和计算操作放到不同流上
- 梯度分桶：把梯度分成多个桶，每个桶算完就发，不用等全部算完
- 流水线并行：微批次之间做通信和计算的重叠
我曾经在一个8卡训练任务中，通过梯度分桶+异步通信，把通信开销从30%降到了8%。效果立竿见影。
注意：重叠不是万能的。如果通信带宽已经饱和，再重叠也没用。这时候要先优化通信本身，比如用NVLink或梯度压缩。
1.4 知识体系总览
下面这张图总结了本章的核心逻辑：
说白了，显存墙逼着我们用多卡，通信墙逼着我们优化通信方式。而通信与计算重叠，就是在这两面墙之间找到一条路。
一句话总结：先算清楚显存够不够，再想通信怎么优化，最后用重叠技术把通信时间藏起来。这三步走完，多卡训练才算入门。