NCCL源码图解之allreduce

date
Jul 2, 2025
slug
nccl_allreduce
tags
NCCL
Source Code Diagram
CUDA
Distribute
summary
allreduce是最常用的集合通信算子, 同时也是primitives类涵盖最全的算子, 搞懂了allreduce的源码, 其他算子的源码就手到擒来了
type
Post
标签
状态
完成
描述
重要性
🌟🌟🌟
关键字
参考链接
status
Published
代码解读基于版本2.11.4；后续会基于最新版本重构，敬请期待！
 
NCCL里数据通信是通过CUDA kernel执行的, 其中每一个channel会绑定一个block（也就是一个SM）. kernel的代码在src/collectives/device 目录的各个.h中，
包含了all_reduce，all_gather等各种通信原语, 我们就以all_reduce入手, 着重讲讲其原理, 其他原语比较类似, 有区别的地方会讲到. 另外all_reduce原语还分别有ring及tree算法, 每种算法实现还有simple/ll(low latency)/ll_128三种协议
我们接下来以all_reduce原语的ring算法中的simple协议的实现入手, 逐行解析, 后面会对其中不一样的地方进行补充

Simple协议Ring算法的all_reduce原语

我们首先需要了解ring all_reduce的实现原理( ), 假设当前有4张卡, 有4个channel进行数据传输, 从上帝视角看, 其传输过程如下图:
ring all-reduce示意图
ring all-reduce示意图
这时候我们从时序的角度, 看下每张卡做的事情:
notion image
也就是说在任意时刻, 任意一张卡都是在做send/recv, 而且其要发送的rank及接收的rank是确定且固定的.
接下来我们看下代码

代码片段0

这里分别解释下含义:
  • nthreads: 一个block使用的线程数量
  • bid: block序号, NCCL中一个channel会绑定一个block, 那么也就代表了channel的序号
  • nChannels: 初始化时搜索出来的channel数量
  • nranks: 本次通信的GPU总数
  • size: 数据传输总量
还有一些参数比较复杂, 需要随着代码的深入逐渐解锁.

ncclRing

首先看下ncclRing这个数据结构
kernel中用的变量含义已经标出, 所以ringIx就是指当前rank在ring环内的index, 在我们的示例中就是等于实际的rank编号

代码片段1

首先进入一个循环, 每个线程都会迭代size/loopSize 次, 先忽略realChunkSize, 后面讲; 然后看下calcOffset函数和modRanks函数.其中calcOffset就是给定chunk的index 每个block要计算的数据的偏移.gridOffset就是每次循环的基地址.modRanks 实际就是在ring内部取模, 比如当前例子下, 1->1 ,4->0
我们来看看一个loopSize里面是如何操作的

loopSize

这里的逻辑可以用下图解释:
 
notion image
在一个loop内部, 数据先拆分成bid总数个块, 然后每个块会拆分成nranks个chunk, 最后完成一个小规模的ring all-reduce.
这里我们看其中的bid0, 一共nranks个chunk, 联系最上面原语的示意图, nranks个GPU的bid0做一次ring all_reduce. 每个rank内不同的block会并行完成不同的数据块. 然后再看loopSize = nChannels*nranks*chunkSize; 就可以解释loopSize的计算公式了, 接下来我们再看chunkSize又是怎么算的.

chunkSize

首先看下计算公式:
 
看下calcBytePerStep函数
这里的NCCL_STEPS是一个预设值:8, buffSizes就是simple协议的buffer大小, 这里的buffer就是通信过程中的数据缓冲区. 我们看看buffer的大小:
可以看到用户可以手动设置buffer的大小, buffer越大占用的显存越大(大模型训练中, 会有很多个context, 那么就会有很多buffer, 占用的显存通常在GB级别), 相应的传输效率会有一定提升(实际传输时buffer是一个FIFO队列, buffer越大越能掩盖”数据拷贝到缓冲区”的耗时). 所以这里有个traceoff.可以看到buffer的大小和CPU有关, 通常大小为4M, 只有ARM的CPU为1M
那么我们就可以算出默认的chunkSize大小:
就是下面的示意图:
notion image
其实看后面的代码, 实际的chunk大小用的是变量:realChunkSize.

realChunkSize

我们看下计算公式:
size-gridOffset 就是当前还剩余的数据量, 只有最后一次迭代realChunkSize ≤ chunkSize , 然后第二步就是向上取整, realChunkSize 必须是一个block内每个线程处理的数据的整数倍, 至于这里减去一个WARP_SIZE是因为第0个warp是控制线程, 不参与数据传输

代码片段2

这里是在执行一次reduce scatter(最后一步融合了all-gather的第一步, 暂时忽略), 我们看每一个传输步骤都先计算处理的chunk index, 然后计算offset决定处理的数据的位置, 最后执行具体的操作. 这里的prims就是实际执行通信操作的过程, 里面具体的逻辑先不讲, 看名字我们就能知道在做什么事情.

ReduceScatter(RS)

接下来我们把目光放在所有rank的bid0上, 看看他是如何完成一次RS的,
我们用一个表格看一下所有rank在第一个send过程这几个变量的值:
ringIx
send chunk
offset
0+3
3
offset+3
1+3
0
offset+0
2+3
1
offset+1
3+3
2
offset+2
可以看到每个rank处理不同index的chunk数据, 接下来扩大范围一下, 看下在整个RS过程中, offset的值(也就是原语示意图上数据的处理过程):
ringIx
send chunk
recv&send chunk(j=2)
recv&send chunk(j=2)
recv chunk
0
3
2
1
0
1
0
3
2
1
2
1
0
3
2
3
2
1
0
3
我们再结合代码看下示意图:
notion image
对于rank0:
  • 第0步(send): 发送a3给rank1
  • 第一步: 接收rank3发送的d2做一次reduce得到d2+a2 ,发送给rank1
  • 第二步(recvReduceSend, 图上未画出): 接收rank3发送的c1+d1做一次reduce得到c1+d1+a1 ,发送给rank1
  • 第三步(recvReduce), 接收rank3发送的b0+c0+d0做一次reduce得到a0+b0+c0+d0 ,得到最终的RS结果

代码片段3

这里就是在做一次all-gather(加上RS的最后一步的copySend), 把每个rank上reduce完成的部分数据发送给其他rank, 这个流程与上述RS类似, 就不再赘述.

Tree算法的All-reduce

首先需要了解DoubleBinaryTree的原理:
TODO

© 木白 2024 - 2025