2. NCCL通信原语：AllReduce、AllGather、ReduceScatter、Broadcast的原理与对比
聊到NCCL，绕不开的就是这几个通信原语。说白了，它们就是多卡通信的「基本动作」。我刚开始接触分布式训练时，总觉得这些概念差不多，都是卡之间传数据嘛。但实际调优时才发现，选错原语，性能可能差一个数量级。
今天咱们就把这四个原语掰开揉碎，看看它们到底怎么工作的，各自适合什么场景。
2.1 Broadcast：一卡发声，全员收听
Broadcast是最直观的原语。它的任务很简单：把一张卡上的数据，复制到所有其他卡上。
工作原理：
- 有一个根节点（root），通常是rank 0
- 根节点把数据发出去，其他卡接收
- 所有卡最终拥有相同的数据副本
我个人的经验：Broadcast常用于模型参数初始化。比如你用rank 0加载预训练权重，然后广播给所有卡。我曾经犯过一个错——在每张卡上都单独加载一次权重，结果不仅慢，还容易因为文件系统并发读而出错。后来改成Broadcast，问题就解决了。
代码示例很简单：
// 在rank 0上准备数据
float* sendbuf = ...; // 只有rank 0有数据
float* recvbuf = ...; // 所有卡都有接收缓冲区
ncclBroadcast(sendbuf, recvbuf, count, ncclFloat, root, comm, stream);嗯，这里要注意：sendbuf只在root卡上有意义，其他卡传什么值都行，NCCL会忽略。
2.2 Reduce：多卡数据，归约到一卡
Reduce和Broadcast正好相反。它是把多张卡上的数据，通过某种操作（通常是求和）归约到一张卡上。
为什么需要Reduce？ 你想想看，每张卡算完自己的梯度后，需要把梯度加起来才能更新模型。Reduce就是干这个活的。
避坑指南：我曾经在调试时发现梯度一直不对，排查了半天才发现——Reduce默认是求和操作，但我传了ncclProd（乘积）。结果梯度越乘越大，模型直接炸了。所以一定要确认你用的操作符是对的。
ncclReduce(sendbuf, recvbuf, count, ncclFloat, ncclSum, root, comm, stream);2.3 AllReduce：全员归约，全员拥有
AllReduce是分布式训练里最常用的原语，没有之一。它相当于「Reduce + Broadcast」的组合：先把所有卡的数据归约到一张卡，再广播回去。但实际实现中，NCCL用了更高效的算法，比如Ring AllReduce。
Ring AllReduce的原理：
- 所有卡连成一个环
- 数据分成N份（N等于卡数）
- 每张卡把一份数据传给下一张，同时从上一张接收一份数据
- 经过N-1次传输，每张卡都拥有部分数据的完整归约结果
- 再经过N-1次传输，把结果扩散到所有卡
说白了，就是「分而治之」的思路。每张卡只和邻居通信，避免了单点瓶颈。
我个人的习惯：在8卡机器上，Ring AllReduce的带宽利用率可以做到接近线性。但跨机时，我建议用Tree算法，因为跨机带宽低，Ring的延迟累积会很明显。
ncclAllReduce(sendbuf, recvbuf, count, ncclFloat, ncclSum, comm, stream);你看，API就一行。但背后的算法选择，NCCL会根据通信拓扑自动决定。
2.4 ReduceScatter：归约并分发
ReduceScatter是AllReduce的前半段。它把数据归约后，不是广播给所有人，而是分散到各张卡上。每张卡只拿到一部分结果。
举个例子：假设有4张卡，每张卡有4个元素。ReduceScatter之后，rank 0拿到前1/4的归约结果，rank 1拿到第二个1/4，以此类推。
这个原语在混合并行中特别有用。比如模型并行时，每张卡只需要一部分梯度，用ReduceScatter就比AllReduce省一半通信量。
注意：ReduceScatter要求数据量能被卡数整除。我遇到过因为batch size没对齐，导致最后一组数据不够分的情况。嗯，这个坑踩过一次就记住了。
ncclReduceScatter(sendbuf, recvbuf, recvcount, ncclFloat, ncclSum, comm, stream);2.5 AllGather：收集并分发
AllGather是ReduceScatter的逆操作，也是AllReduce的后半段。它把各张卡上的数据收集起来，然后广播给所有人。
工作流程：
- 每张卡有一份独特的数据
- 所有卡把自己的数据发给别人
- 最终每张卡都拥有完整的数据集合
你可能会问：这和AllReduce有什么区别？区别在于操作：AllReduce有归约（求和），AllGather只是拼接。说白了，AllGather不改变数据值，只是把数据拼在一起。
我的经验：在做数据并行时，如果你需要每张卡都看到完整的batch（比如Batch Normalization的全局统计量），可以用AllGather。但要注意显存——每张卡都要存所有卡的数据，显存占用会翻N倍。
ncclAllGather(sendbuf, recvbuf, sendcount, ncclFloat, comm, stream);2.6 四个原语的对比
为了方便你快速查阅，我整理了一张对比表：
| 原语 | 输入 | 输出 | 通信量 | 典型场景 |
|---|---|---|---|---|
| Broadcast | 1卡有数据 | 所有卡有相同数据 | O(N) | 参数初始化、加载权重 |
| Reduce | 所有卡有数据 | 1卡有归约结果 | O(N) | 梯度汇总到主卡 |
| AllReduce | 所有卡有数据 | 所有卡有归约结果 | O(2N) | 数据并行梯度同步 |
| ReduceScatter | 所有卡有数据 | 每卡有部分归约结果 | O(N) | 混合并行、模型并行 |
| AllGather | 每卡有部分数据 | 所有卡有完整数据 | O(N) | 全局统计量收集 |
注意：通信量里的N是数据总量。AllReduce的2N是因为它相当于ReduceScatter + AllGather。
2.7 知识体系总览
下面这张图展示了四个原语之间的关系和核心逻辑：
2.8 如何选择？
说了这么多，到底什么时候用哪个？我总结几条经验：
- 数据并行训练：无脑用AllReduce。NCCL内部会选最优算法，你不用操心。
- 模型并行：用ReduceScatter + AllGather的组合，比AllReduce省一半通信量。
- 参数初始化：用Broadcast，简单高效。
- 梯度裁剪或特殊操作：用Reduce把梯度汇总到主卡处理，再Broadcast回去。
最后提醒一句：别迷信「AllReduce万能」。我见过有人在小batch size场景下硬用AllReduce，结果通信开销比计算还大。这时候用Reduce + Broadcast反而更快。选原语，一定要结合你的实际场景。
好了，这四个原语就聊到这儿。理解了它们，NCCL的调度策略你就掌握了一半。