3. NCCL通信原语：AllReduce原理与实现、Broadcast原理与实现、AllGather原理与实现、ReduceScatter原理与实现、Send/Recv点对点通信
好，咱们直接进入正题。这一章聊的是NCCL最核心的通信原语。说白了，就是多卡之间怎么“说话”的几种基本方式。
我刚开始接触分布式训练时，觉得这些原语名字挺唬人。什么AllReduce、AllGather，听着像魔法咒语。后来调了无数次bug才明白——它们本质上就是几种固定的数据搬运+计算模式。你掌握了这几种模式，多卡协同就是个拼图游戏。
3.1 AllReduce：最常用的“求和”操作
AllReduce是分布式训练里出现频率最高的原语。梯度同步全靠它。
原理一句话： 每张卡都有一个数据块，AllReduce让所有卡最终都拿到这些数据块的“归约结果”（通常是求和）。
举个例子：4张卡，每张卡有个数字[1, 2, 3, 4]。AllReduce求和后，每张卡都得到[10]。嗯，就是这么简单。
但实现上，NCCL用了Ring AllReduce算法。我画个图帮你理解：
Ring AllReduce分两步走：
- Reduce-Scatter：数据切成N份（N=卡数），每张卡只归约自己负责的那一份。数据在环上转一圈，每转一次做一次局部归约。
- AllGather：把归约好的结果广播给所有卡。同样在环上转一圈，每张卡收集完整结果。
核心要点： Ring AllReduce的通信量是2*(N-1)*K（K是数据量），而朴素做法是N*(N-1)*K。卡越多，优势越明显。
代码示例：
// 最简单的AllReduce调用
ncclAllReduce(sendbuf, recvbuf, count, ncclFloat, ncclSum, comm, stream);
// 实际项目中我习惯这样用
ncclAllReduce(gradients, gradients, param_count,
ncclFloat32, ncclSum, comm, stream);
// 注意：原地操作时sendbuf和recvbuf可以指向同一地址我的经验： 刚开始用AllReduce时，我总忘记设置正确的数据类型。ncclFloat和ncclFloat32不一样——前者是half精度，后者才是float。踩过这个坑之后，我每次都会double-check。
3.2 Broadcast：一对多的数据分发
Broadcast就是“一个人说话，所有人听”。一张卡把数据发给所有其他卡。
应用场景： 参数初始化、加载checkpoint后广播给所有卡。
实现上，NCCL用树形广播或环状广播。我个人更关注它的实际用法：
// root=0表示GPU0是发送方
ncclBroadcast(sendbuf, recvbuf, count, ncclFloat32, 0, comm, stream);
// 非root卡：sendbuf传NULL即可
ncclBroadcast(NULL, recvbuf, count, ncclFloat32, 0, comm, stream);注意： 所有卡必须传入相同的root参数。我曾经见过一个bug——某张卡root设成了1，其他卡设成了0，结果数据全乱了。排查了整整一下午。
3.3 AllGather：收集所有卡的数据
AllGather和Broadcast有点像，但区别在于：每张卡都贡献一份数据，最终每张卡拿到所有人的数据。
你想想看，这其实就是Ring AllReduce的第二阶段。单独拿出来用也很常见——比如做模型并行时，需要收集各卡上的激活值。
// 每张卡有chunk_size个元素
ncclAllGather(sendbuf, recvbuf, chunk_size, ncclFloat32, comm, stream);
// recvbuf大小 = chunk_size * nranks数据排布： 收集后的数据按rank顺序排列。rank0的数据在最前面，rank1紧随其后。这个顺序在后续处理时要注意。
3.4 ReduceScatter：先归约再分发
ReduceScatter是AllReduce的“前半段”。它把数据归约后，分散到各张卡上——每张卡只拿到一部分结果。
我刚开始觉得这个原语有点鸡肋。后来做超大模型训练时才发现它的妙用：
- 减少单卡显存占用（不用存完整结果）
- 为后续计算做准备（比如做局部更新）
// 每张卡输入N个元素，输出N/nranks个元素
ncclReduceScatter(sendbuf, recvbuf, recv_count, ncclFloat32, ncclSum, comm, stream);避坑指南： recv_count是每张卡接收的元素数，不是总数。我见过有人传了总数，结果NCCL直接报错——嗯，这个错误信息还挺友好的，会告诉你“recvcount * nranks != sendcount”。
3.5 Send/Recv：最灵活的点对点通信
前面几个都是集合通信——所有卡一起参与。Send/Recv则是一对一的通信，想怎么传就怎么传。
为什么需要这个？因为有些场景下，你不需要所有卡都通信。比如：
- Pipeline并行中，前一张卡把中间结果传给下一张卡
- 自定义的通信模式，集合原语搞不定
// GPU0发送
ncclSend(sendbuf, count, ncclFloat32, 1, comm, stream);
// GPU1接收
ncclRecv(recvbuf, count, ncclFloat32, 0, comm, stream);关键点： Send和Recv必须配对。发送方指定peer=1，接收方指定peer=0。不配对的话，程序会卡死——我调试时遇到过，当时还以为是NCCL的bug，后来发现是自己写错了peer编号。
| 原语 | 通信模式 | 典型场景 | 我的建议 |
|---|---|---|---|
| AllReduce | 多对多 | 梯度同步 | 最常用，优先用Ring算法 |
| Broadcast | 一对多 | 参数广播 | 注意root一致性 |
| AllGather | 多对多 | 收集激活值 | 注意输出缓冲区大小 |
| ReduceScatter | 多对多 | 显存优化 | recv_count别传错 |
| Send/Recv | 一对一 | 自定义通信 | 必须配对，否则死锁 |
最后说一句：这些原语是NCCL的基石。你不需要记住每个算法的细节，但一定要理解什么时候用哪个。我在项目中见过有人用AllReduce做Broadcast的事——虽然也能跑，但效率差了一截。选对原语，性能就赢了一半。