通信原语详解：AllReduce、AllGather、ReduceScatter、Broadcast
大家好，我是你们的老朋友。今天咱们来聊聊分布式训练里最核心的几个通信原语。说实话，我刚入行那会儿，觉得这些概念就是一堆花里胡哨的术语。直到我在一个 64 卡的项目里，因为选错了通信方式，训练速度直接掉了 30%。嗯，从那以后，我再也不敢小看它们了。
你想想看，多卡训练的本质是什么？说白了，就是让一堆 GPU 协同干活。但干活过程中，数据怎么分、结果怎么合，全靠这些通信原语来搞定。今天我就带你一个一个捋清楚。
1. Broadcast：一卡广播，全员接收
Broadcast 是最简单的通信模式。它的语义是：一张卡把自己的数据，复制给所有其他卡。
什么时候用？最常见的就是参数初始化。比如你用一张卡加载了预训练模型，然后通过 Broadcast 把权重广播到所有卡上。我习惯在数据并行训练的开头用这个，省得每张卡都去读一遍磁盘。
// 伪代码示例：Broadcast 语义
// root_rank = 0 的卡将 buffer 数据发送给所有其他卡
ncclBroadcast(buffer, size, ncclFloat32, root_rank, comm, stream);2. Reduce：多卡归约，结果合一
Reduce 是 Broadcast 的反向操作。多张卡各自有数据，通过某种操作（比如求和、求最大值）合并成一个结果，最终落到某一张卡上。
我在项目中遇到过这样一个场景：多卡各自计算了梯度，需要把梯度求和后更新到主卡。这时候用 Reduce 就对了。不过要注意，Reduce 的结果只存在 root 卡上，其他卡拿不到。
3. AllReduce：全员归约，人人有份
AllReduce 是分布式训练里最常用的原语，没有之一。它的语义是：所有卡的数据先归约（比如求和），再把结果分发给每一张卡。说白了，就是 Reduce + Broadcast 的合体。
数据并行训练中，梯度同步就是靠 AllReduce 完成的。每张卡算完梯度后，AllReduce 把梯度求和，然后每张卡都拿到完整的梯度，各自更新参数。
| 通信原语 | 输入 | 输出 | 典型场景 |
|---|---|---|---|
| AllReduce | 每张卡有数据 X_i | 每张卡都得到 sum(X_i) | 梯度同步、参数平均 |
| Reduce | 每张卡有数据 X_i | root 卡得到 sum(X_i) | 单卡收集结果 |
| Broadcast | root 卡有数据 X | 所有卡都得到 X | 参数广播、初始化 |
// NCCL 中 AllReduce 的调用方式
ncclAllReduce(sendbuff, recvbuff, count, ncclFloat32, ncclSum, comm, stream);4. ReduceScatter：先归约，再分发
ReduceScatter 是 AllReduce 的「前一半」。它的语义是：所有卡的数据先做归约，但结果不是全量发给每张卡，而是把归约后的数据切分成块，每张卡只拿属于自己的那一块。
你想想看，如果数据量很大，AllReduce 的通信量是 O(N) 的（N 是卡数）。但 ReduceScatter 的通信量只有 O(N/K)，K 是分块数。说白了，它省了一半的通信量。
我在做超大模型训练时，经常用 ReduceScatter 配合 AllGather 来替代 AllReduce。比如在 DeepSpeed 的 ZeRO 优化器里，梯度同步就用到了这个技巧。
// ReduceScatter 示例：每张卡只拿归约后的 1/N 数据
ncclReduceScatter(sendbuff, recvbuff, recvcount, ncclFloat32, ncclSum, comm, stream);5. AllGather：全员收集，拼成完整
AllGather 是 ReduceScatter 的逆操作。它的语义是：每张卡把自己的一块数据发给所有其他卡，最终每张卡都拼出完整的数据。
什么时候用？比如模型并行中，每张卡只存了模型的一部分参数，但前向计算时需要完整的参数。这时候就用 AllGather 把参数碎片拼回来。
我记得有一次调试一个 MoE 模型，专家参数分散在 8 张卡上。每次前向都要 AllGather 一次，通信开销大得吓人。后来我改成只 AllGather 当前层需要的专家，才把性能拉回来。
// AllGather 调用：每张卡贡献自己的数据块，最终每张卡都有完整数据
ncclAllGather(sendbuff, recvbuff, sendcount, ncclFloat32, comm, stream);6. 一张图看懂它们的关系
下面这张 SVG 图，我画了这四个原语的核心逻辑。你看一眼就能明白它们之间的区别。
7. 如何选择？我的实战建议
说了这么多，到底怎么选？我总结几条经验：
- 数据并行训练：直接用 AllReduce。NCCL 内部会帮你选最优实现，省心。
- 模型并行训练：用 AllGather 收集参数，用 ReduceScatter 同步梯度。
- 参数服务器架构：用 Reduce 把梯度推到 server，用 Broadcast 把参数拉回来。
- 超大模型（ZeRO、FSDP）：手动拆成 ReduceScatter + AllGather，中间可以插梯度压缩。
好了，通信原语就聊到这儿。这些概念是分布式训练的基石，搞懂了它们，后面讲算子优化和带宽压榨的时候，你才能理解为什么某些技巧能生效。下一章，我们聊聊 NCCL 的 Ring 算法到底是怎么工作的。