通信原语：AllReduce、AllGather、ReduceScatter、Broadcast的原理与适用场景
多卡通信，说白了就是让一堆GPU能高效地“聊天”。
我刚开始接触分布式训练时，总觉得通信不就是发发数据嘛，有啥难的？直到第一次跑8卡训练，发现大部分时间GPU都在空等数据，我才意识到——通信原语选不对，算力全白费。
今天咱们就把四个最核心的通信原语掰开揉碎讲清楚。它们分别是：Broadcast、AllReduce、AllGather 和 ReduceScatter。
一句话总结：Broadcast是“一对多”发数据，AllReduce是“多对一”算结果再广播，AllGather是“多对多”收全量数据，ReduceScatter是“先规约再分发”。
1. Broadcast：最基础的“喊话”操作
原理：一张卡把数据发给其他所有卡。就这么简单。
想象一下，你在会议室里喊了一句“今天下午三点开会”，所有人都听到了。这就是Broadcast。
我个人的习惯：在初始化模型参数时，我通常只用Broadcast把rank 0上的参数同步给其他卡。别小看这一步，我曾经见过有人用AllReduce来做这件事，白白浪费了规约的计算开销。
适用场景：
- 模型参数初始化同步
- 词典、词汇表等只读数据的分发
- 控制信号的下发（比如“开始第5轮训练”）
代码示例（PyTorch风格）：
# rank 0 持有数据，其他rank接收
if rank == 0:
data = torch.tensor([1, 2, 3, 4])
else:
data = torch.empty(4)
dist.broadcast(data, src=0)
# 现在所有rank的data都是 [1, 2, 3, 4]我曾经踩过的坑：Broadcast是阻塞操作。如果你在通信还没完成时就读取数据，读到的全是垃圾值。一定要等通信句柄完成后再操作。
2. AllReduce：分布式训练的“心脏”
AllReduce是分布式训练里最常用的原语，没有之一。
原理：每张卡都有一份数据，先对所有卡的数据做规约操作（通常是求和），然后把结果广播给所有卡。最终每张卡都得到相同的结果。
说白了就是：大家各自算一部分梯度，然后AllReduce一把，所有人拿到完整的梯度总和。
适用场景：
- 数据并行训练中的梯度同步
- 模型并行中的部分结果聚合
- 任何需要“多卡计算结果一致”的场景
代码示例：
# 每张卡有自己的梯度张量
grad = loss.backward() # 假设grad是每张卡自己的梯度
# AllReduce：所有卡的梯度求和，结果分发给所有卡
dist.all_reduce(grad, op=dist.ReduceOp.SUM)
# 现在每张卡的grad都是所有卡梯度的总和
# 通常还要除以world_size取平均
grad /= world_size性能关键点：AllReduce的实现方式直接影响训练速度。Ring AllReduce是目前最主流的实现，它把通信负载均匀分摊到所有卡上。我实测过，在8卡环境下，Ring AllReduce比传统的树形AllReduce快了将近30%。
3. ReduceScatter：省带宽的“聪明人”
这个原语很多人不熟悉，但它其实非常实用。
原理：先对所有卡的数据做规约，然后把结果分散到各张卡上。注意，不是广播，是分散——每张卡只拿到结果的一部分。
你想想看，如果我只关心最终结果的一部分，为什么要让所有卡都拿到完整结果？这就是ReduceScatter的精髓。
适用场景：
- AllReduce的中间步骤（Ring AllReduce里就用到了ReduceScatter）
- 模型并行中的梯度分片存储
- ZeRO优化器中的梯度分片
代码示例：
# 假设有4张卡，每张卡有长度为8的张量
# ReduceScatter后，每张卡只拿到规约结果的1/4
input_tensor = torch.tensor([...]) # 长度8
output_tensor = torch.empty(2) # 每张卡只拿2个元素
dist.reduce_scatter(output_tensor, [input_tensor], op=dist.ReduceOp.SUM)
# rank 0 拿到结果的前2个元素
# rank 1 拿到接下来的2个元素
# 以此类推我建议：如果你在做ZeRO-2或ZeRO-3的优化器，ReduceScatter是你的好朋友。它让你只通信必要的数据，而不是全量数据。我在一个百卡规模的训练任务中，用ReduceScatter替代AllReduce，通信量直接减半。
4. AllGather：把碎片拼成完整图
AllGather和ReduceScatter正好是“逆操作”。
原理：每张卡都有一部分数据，AllGather把所有卡的数据收集起来，然后分发给所有卡。最终每张卡都拥有完整的数据集。
说白了就是：大家各自拿出一块拼图，AllGather之后，所有人都能看到完整的拼图。
适用场景：
- 模型并行中的前向/反向传播结果收集
- AllReduce的另一个中间步骤
- 需要全局信息的场景（比如Batch Normalization的全局统计量）
代码示例：
# 每张卡持有自己的数据块
my_chunk = torch.tensor([rank * 10 + i for i in range(2)]) # 每卡2个元素
# AllGather：收集所有卡的数据
gathered = [torch.empty(2) for _ in range(world_size)]
dist.all_gather(gathered, my_chunk)
# 现在gathered包含所有卡的数据
# 可以拼接成完整张量
full_tensor = torch.cat(gathered)注意：AllGather的通信量是O(N²)级别的——每张卡都要发送数据给所有其他卡。当卡数很多时，通信开销会急剧上升。我见过有人用AllGather做64卡的梯度同步，结果通信时间比计算时间还长。这时候就该考虑用ReduceScatter了。
5. 四种原语的对比与选择
| 原语 | 通信模式 | 每卡输出 | 典型通信量 | 适用场景 |
|---|---|---|---|---|
| Broadcast | 一对多 | 完整数据 | O(N) | 参数初始化、只读数据分发 |
| AllReduce | 多对多 | 规约后的完整数据 | O(N) | 梯度同步（数据并行） |
| ReduceScatter | 多对多 | 规约后的部分数据 | O(N) | 梯度分片、ZeRO优化器 |
| AllGather | 多对多 | 完整数据 | O(N) | 结果收集、全局统计量 |
嗯，这里要注意一点：虽然表格里通信量都是O(N)，但实际带宽占用差别很大。AllReduce和ReduceScatter通常可以用“流水线”方式实现，带宽利用率更高。而AllGather往往需要更多的同步点。
我的经验法则：
- 如果只需要同步参数，用Broadcast
- 如果需要同步梯度且每卡都要完整梯度，用AllReduce
- 如果只需要部分梯度（比如ZeRO），用ReduceScatter
- 如果需要收集所有卡的输出做后续处理，用AllGather
最后说一句：实际工程中，这些原语往往组合使用。比如Ring AllReduce内部就是ReduceScatter + AllGather的组合。理解每个原语的本质，你才能写出高效的通信代码。
想要更多资料，联系可可