2. NCCL基础：通信原语与Ring算法
聊到分布式训练，NCCL是绕不开的坎。我最早接触NCCL是在做多机多卡训练的时候，当时被各种通信原语搞得晕头转向。后来踩了不少坑，才慢慢摸清楚它的脾气。今天咱们就把NCCL的核心通信原语和Ring算法掰开揉碎讲清楚。
2.1 四大通信原语
NCCL提供了四种基础通信操作。说白了，分布式训练里所有的通信需求，都能用这四种原语拼出来。
2.1.1 AllReduce
这是最常用的原语。每个GPU上的数据，经过AllReduce后，所有GPU都会得到相同的结果。我习惯把它理解成「大家把数据交出来，算个总和，再每人发一份」。
核心特点：
- 输入：每个GPU有一份数据
- 输出：每个GPU得到所有数据的规约结果（通常是求和）
- 典型场景：梯度同步
举个例子，假设有4张卡，每张卡上有个数字[1,2,3,4]。AllReduce求和后，每张卡都会得到[10]。嗯，就是这么简单粗暴。
2.1.2 Broadcast
Broadcast是一对多的操作。一张卡把数据发给其他所有卡。我在项目中遇到过这样的场景：初始化模型参数时，用Broadcast把rank 0上的参数同步给所有卡。
我的经验：Broadcast常用于参数初始化或加载checkpoint后的同步。但要注意，如果数据量很大，Broadcast会成为瓶颈。我曾经在加载大模型时，光Broadcast就花了十几秒。
2.1.3 ReduceScatter
这个原语有点意思。它把AllReduce拆成了两步：先Reduce，再Scatter。每个GPU只负责计算一部分数据的规约结果。
你想想看，如果所有GPU都去算完整的结果，那通信量是O(N)的。但用ReduceScatter，每个GPU只算1/N，通信量就降下来了。这是Ring算法高效的关键。
2.1.4 AllGather
AllGather是Broadcast的逆操作。每个GPU把自己的数据发给别人，同时接收别人的数据。最终所有GPU都拥有完整的数据集。
注意：AllGather和AllReduce的区别在于，AllGather不做规约运算，只是单纯地收集数据。我刚开始学的时候经常搞混这两个。
2.2 Ring算法详解
为什么NCCL要用Ring算法？说白了，就是为了高效利用带宽。传统的树形算法，随着GPU数量增加，根节点的带宽会成为瓶颈。Ring算法把GPU连成一个环，每个GPU只和相邻的两个通信。
2.2.1 算法步骤
Ring AllReduce分两步走：ReduceScatter和AllGather。
- ReduceScatter阶段：数据被分成N份（N=GPU数）。每个GPU只负责计算一份数据的规约结果。数据在环上转一圈，每转一次，当前GPU就把收到的数据和自己本地的对应分片做规约。
- AllGather阶段：每个GPU已经有了完整的一份规约结果。现在把这些结果广播给所有GPU。数据再转一圈，每转一次，GPU就把收到的数据复制到自己的对应位置。
关键点：整个过程中，每个GPU只发送和接收N-1次数据。每次传输的数据量是总数据量的1/N。所以通信量是O(N)的，但带宽利用率很高。
2.2.2 为什么Ring算法快？
我刚开始也不理解，为什么绕一圈反而比树形快？后来在项目里实测才发现，Ring算法有两个优势：
- 带宽利用率高：每个GPU的收发带宽都被充分利用，没有空闲等待
- 可扩展性好：GPU数量增加时，通信时间线性增长，而不是指数增长
举个例子，8张卡用Ring做AllReduce，比用树形快了将近一倍。我在实际项目中测试过，32张卡时差距更明显。
2.3 代码示例
来看看NCCL的AllReduce怎么用。这段代码我改了好几次才跑通，分享给你：
#include <nccl.h>
#include <cuda_runtime.h>
int main() {
int nDevices = 4;
int64_t count = 1024 * 1024; // 1M个元素
// 初始化NCCL通信器
ncclComm_t comms[nDevices];
ncclCommInitAll(comms, nDevices, NULL);
// 分配GPU内存
float *sendbuf[nDevices], *recvbuf[nDevices];
cudaStream_t streams[nDevices];
for (int i = 0; i < nDevices; i++) {
cudaSetDevice(i);
cudaMalloc(&sendbuf[i], count * sizeof(float));
cudaMalloc(&recvbuf[i], count * sizeof(float));
cudaStreamCreate(&streams[i]);
}
// 执行AllReduce
ncclGroupStart();
for (int i = 0; i < nDevices; i++) {
ncclAllReduce(
sendbuf[i], recvbuf[i], count,
ncclFloat, ncclSum,
comms[i], streams[i]
);
}
ncclGroupEnd();
// 同步等待
for (int i = 0; i < nDevices; i++) {
cudaStreamSynchronize(streams[i]);
}
// 清理资源
for (int i = 0; i < nDevices; i++) {
ncclCommDestroy(comms[i]);
cudaFree(sendbuf[i]);
cudaFree(recvbuf[i]);
cudaStreamDestroy(streams[i]);
}
return 0;
}避坑指南：我曾经犯过一个错误，忘了调用ncclGroupStart/ncclGroupEnd。结果多个AllReduce调用互相干扰，数据全乱了。记住，多个NCCL操作一定要用Group包裹起来。
2.4 性能对比
不同通信原语的性能差异很大。我整理了一张表，方便你参考：
| 原语 | 通信量 | 延迟 | 适用场景 |
|---|---|---|---|
| AllReduce | O(N) | 中等 | 梯度同步 |
| Broadcast | O(1) | 低 | 参数初始化 |
| ReduceScatter | O(N) | 中等 | Ring算法第一步 |
| AllGather | O(N) | 中等 | Ring算法第二步 |
看到没？AllReduce的通信量是O(N)，但Ring算法通过分片把单次传输量降到了O(1)。这就是它高效的原因。
2.5 小结
NCCL的四个原语，每个都有自己的脾气。AllReduce最常用，Broadcast适合小数据同步，ReduceScatter和AllGather是Ring算法的左右手。Ring算法本身也不复杂，就是数据分片在环上转两圈。第一圈做规约，第二圈做广播。
嗯，掌握了这些基础，后面讲NVSHMEM和NCCL协同的时候，你就知道它们各自擅长什么了。