NCCL源码图解之基本通信算子

date
Aug 2, 2025
slug
nccl_primitives_ops
tags
NCCL
Source Code Diagram
CUDA
Distribute
summary
NCCL会将通信原语拆分成更细粒度的primitives op的组合,
type
Post
标签
状态
进行中
描述
重要性
🌟🌟🌟
关键字
参考链接
status
Published

🥰总结

notion image
notion image

⁉️问题

 

🧐内容

我们先看下NCCL中allreduce的原理图:
notion image

genericOp

上面每一个操作都是Primitives类中一个基本的通信算子:
上节我们讲了
NCCL源码图解之Primitives类的初始化, 接下来我们看下这些算子的含义.
首先我们会发现所有Op都是调用genericOp这个函数, 区别只在于模板参数和函数参数. 所以我们讲清楚这一个函数就可以了.
讲函数之前我们先了解下这些Op都做了哪些事情, 毕竟上面太多内容了, 很容易看晕, 看图是最直观的

算子图示

Send

notion image
  • 最基本的send就是把用户输入数据发送给send peer的buffer里;
  • sendFromOutput 就是把用户输出数据发送给对端buffer
  • 上节初始化中, 大篇幅讲的Direct就体现在这里, directSend 就是把input直接发送到对端的Output, 不经过buffer
  • directSendFromOutput 也好理解, output发送到对端output

Recv

notion image
  • 最基本的recv 就是从当前rank的buffer中把数据拷贝到output
  • directRecv 就是啥都不干(因为directSend就已经把数据传输到output了)

CopySend

notion image
  • copySend 相比send 多了一步, 就是将数据从input拷贝到output里
  • directCopySend 就是copy +directSend

recvCopySend

notion image
  • recvCopySend 这里实际上是将数据从buffer中读到寄存器(kernel)然后分别写到本地DDR和发送给远端buffer
  • directRecvCopySend , 前面讲directRecv 就是啥都不干, 由于数据已经在output里了, 所以copy 也啥都不干, 最后就是一个directSendFromOutput
  • RecvCopyDirectSend 就是recv + directSendFromOutput

recvReduce

notion image
前面都是直接的数据拷贝, 现在开始我们加上了reduce操作, 也就是有了计算的参与;
  • recvReduceCopy 是从buffer中接收数据, 和本地的input做一次reduce, 然后写入output
  • recvReduceSend 就是recvReduceCopy的基础上不写入output, reduce之后直接send 到对端buffer
  • recvReduceCopySend 就是上面二者的组合, reduce完后写入本地output一份, 再发送给对端一份

directRecvReduceSend

notion image
  • directRecvReduceSend 这里的direct修饰recv , 所以是将output中的数据和input一起做reduce, 然后将结果send 到对端
  • directRecvReduceCopySend 这个比较特殊, 名字不是很准确, 这里的direct修饰send , 所以是在recvReduceCopySend 的基础上将数据send 到对端的output(正确的名字应该是RecvReduceCopyDirectSend , 这样写应该是读起来顺一点~)

总结

  • send 是将数据发送给远端, 没加修饰词就是发给buffer
  • recv 是将数据从本地buffer拷贝到用户输出output
  • FromOutput 修饰send , 表示发送端不是input而是output
  • direct 修饰send/recv, 表示数据不经过buffer, 直接操作input和output
  • copy 指的是将数据从buffer或input拷贝到output
  • reduce 是会做计算(SUM, MIN 等), 而且需要两个输入
以上就是基于genericOp的所有基本通信算子. 对于想了解原理的同学可以到此为止了. 接下来我们就讲解一下源码

genericOp源码剖析

API

我们先看下genericOp的接口:
首先是模板参数:
  • DirectRecv1: 控制是否需要从buffer中recv
  • DirectSend1: 控制发送到对端的buffer还是output中
  • Recv/Send: 控制当前操作是发送还是接收(可能同时存在)
  • SrcBuf/DstBuf: 控制是否要从用户输入/输出中获取数据
然后是函数参数:
  • srcIx/dstIx:上面模板参数SrcBuf/DstBuf存在时才会有值传入
  • remoteIx: 只有开启Direct Send时才会有值, 表示此时不往对端buffer中写输入, 而是往给定位置写
  • nelem: 发送的数据量
  • postOp: 控制当前操作是否需要计算, 也就是reduce操作
上面这些参数不用记住, 有个大概的了解即可, 后面使用到的时候结合源码就会很容易理解.
接下来只讲send, recvReduceCopySend; 这两个比较有代表性, 剩余的基本就是在其基础上做了一些改动, 可以自己结合源码看看

send

我们先贴上send的调用:
接下来我们看看genericOp的代码:
首先注意到DirectRecv/DirectSend的两个都是0, 所以我们会把数据传输给对端的buffer; 然后就是Input给了值((Src=1), ouput没有((Dst=0), 因为send只有数据的输入,没有输出;

sliceSize

我们先了解下NCCL中很重要的三个概念: chunk, slice, step, 上图:
notion image
首先我们的buffer在初始化的时候就已经申请好了, 默认是4M, 然后在simple协议下会等分成8份, 其中一份的大小是512K, 也就是一个step, 是NCCL数据传输的基本单位; 然后一个slice的大小是两个step, 一个chunk的大小是4个step(2个slice); 在simple协议下一次传输是一个slice.
再回到源码中, nelem在
NCCL源码图解之allreduce 讲过, 就是实际的chunk_size, 也就是2个slice大小, 那么这行代码:
就是首先chunk_size要16B对齐, 然后除以2(就是上图中slice大小, 1M); 但是我们还要考虑一种情况: 最后的一个chunk可能就剩几个数据了,或者数据量本身就很小, 假设只有2K, 那么max前面的值只有1K, 在NCCL的实现中不好传输, 所以这里就是说最小传输32K

UserBuff

接下来就是进入数据传输的循环中了, 看下这几行代码:
前面我们讲过src=1, dst=0; 所以只有上面的if会进入, 下面的不会; 而且对于上面的if只有RoleInput 这一个线程会进入, 功能是把userBuff 这个地址拷贝到共享内存中, 这里大家可能会疑惑: 什么是userBuff? 这个和我们前面讲到的buffer是一个东西吗?
首先说结论: 这个userBuff 实际就是用户输入数据的地址, 这个命名有误导性, 看下图:
notion image
所以和我们上面讲到的buffer完全没有关系, buffer是NCCL为了数据传输申请的临时缓冲区.

waitPeer

看下这两行代码:
waitPeer就是让我的对端准备好进行数据接收, subBarrier 是做一个同步操作, 我们先讲前一个函数:
先贴一下这几个变量的值:
  • isSendNotRecv = 1 : 表示这是一个send操作
  • noRecvWait = 0 send不会用到
  • noSendWait = 0 这里的判断条件是Direct, 因为Direct模式是直接写用户的输出, 所以内存肯定是申请好的, 所以放心写;
但是非Direct模式, 我们要往buffer里面写, buffer满了的情况下就不能继续发了, 需要对端完成recv, 腾出buffer的空间我们才能继续send; 因为我们的buffer是一个循环队列.
接下来看下面的代码:
满足if的第二个条件, 只有RoleWaitSend 这个线程才会进入, 而且里面是一个while循环(所以外面会有一个subBarrier函数同步线程;这里的spins 是一个计数器, 当长时间无法退出循环时, 就需要触发abort, 防止对端无法接收数据(对端进程hang)而导致无法退出的情况.
接下来我们回忆下connStepPtr , 在
NCCL源码图解之Primitives类的初始化 我们讲过loadSendConn函数:
send端connStepPtr 就是buffer上的head指针, 这是一个本地的值;那么为什么是volatile呢? 因为对端做完一次recv会更新这个值; connStepCache是本地的一个值, 不会被远端修改, 所以这就是一个读本地写远端的场景.
我们回到上面的while循环里, 看下循环的条件, send场景下:
一共有三种情况:
notion image
 
  • 当buffer为空, 这时候待发送的数据量小于buffer大小, 所以满足发送条件, 退出循环, 开始发送
  • 当buffer有数据, 这时候待发送的数据量小于buffer的剩余大小, 满足发送条件
  • 当buffer满了, 这时候待发送的数据量大于buffer剩余大小, 不满足发送条件(因为buffer是循环队列, 此时写入会覆盖掉还未recv的数据,导致出错), 此时就需要一直循环下去, 同时读取*connStepPtr中的值, 直到远端recv完成数据后更新这个值(也就是会变成第二种情况)
接下来就是退出循环, 这时候我们需要计算接收地址(也就是buffer的地址+偏移):
最后根据slice的大小更新step的值.

subBarrier

waitPeer结束后会进行一次同步, 这里也比较有意思, 我们看一下
因为waitPeer只有RoleWaitSend线程在检查是否可以发送, 那么剩余的工作线程就需要等待该线程, 决定何时进行下一步的数据发送. 另外我们看这里只会对工作线程进行同步(辅助线程不进行同步), 然后根据不同的情况进行了优化, 以保证在包含所有工作线程的最小的范围内进行同步

数据发送

这里就是调用ReduceOrCopyMulti这个函数在进行真正的数据发送了, 这个函数需要单独讲一讲.

同步

我们看这里做了三个不同粒度的同步:
  • barrier(): 当前block内所有线程同步, 这个后面详细讲
  • __threadfence_system(): 该同步的作用是确保线程的内存操作在整个系统中的可见性和顺序一致性, 也就是告诉对端我写完了, 你可以去内存中取数据了(而不是缓存)
  • __syncwarp(): 上一个同步只有一个线程操作, 所以需要warp内进行同步

barrier():

这个同步非常有意思, 也是能做下面”两个循环”优化的核心所在, 我们先看下源码:
对于多个warp的情况会调用PTX的bar.sync ,这里的第一个参数指的是使用block内第几个bar资源(0-15), 第二个参数指的是block内要同步的线程数, barrier 的功能是在一个block内部, 当前bar资源内等到了nthreads个线程才会继续往下走, 否则就在这里等待, 看下图:
notion image
上图中圆圈代表warp, 工作线程和非工作线程会走到不同的分支, 做不同的事, 然后两个不同的分支都会有一个barrier , 这两个barrier 会使用同一个bar资源, 然后在红线处进行同步. 那么我们就实现了在不同的分支同步所有线程的目的.

PostPeer

数据发送完毕后, 我们还需要更新buffer的tail位置, 让recv端知道哪些数据是需要接收的, 我们看下代码:
Send场景只有RolePostSend线程才会进入, 更新step值, 然后把新的step值写入到*connStepPtr中, 这个值就是tail:

两个循环

接下来我们从整个函数的角度看一下, 下面是精简后的代码:
可以看到工作线程会进入上面的do-while循环, 非工作线程进入下面的while循环, 其中do-while循环负责真正的数据传输, 进入条件:
而while循环则是处理空的slice的情况(因为slice的长度要要对齐的, 那么就可能存在空的情况), 进入条件:
至于为什么不在工作线程中做这件事, 代码中的注释也解释了, 这样有助于减少分支, 提升性能.
这里我们看一下两个循环的相同位置处都有个barrier() , 这两个是一对, 也就是说block内工作线程和非工作线程都需要在这里同步到. 这里有非常多的同步, 直接上图:
notion image
上面图中一串圆圈代表一个线程, 每一串上的一个圆圈代表当前线程要参与的工作, 可以直观的看到每个同步的范围.

ReduceOrCopyMulti

该函数是NCCL里真正做数据搬运及reduce的核心函数, 逻辑也不复杂. 首先看下调用:
对于send 操作, 进入了else里, 我们把形参的值补充上去:
我们需要注意以下几个参数:
  • UNROLL: cuda常用的优化手段: 循环展开; 需要注意Ampere及之后的架构这个值是8, 我们以8为例
    • PreOpN: 这个值和reduce有关, send 用不到, 先不考虑

    代码片段一

    结合上面形参的值, 上面代码补充了参数的含义, warp和thread相关的参数后面再讲, 先看下对齐的部分, 上一节我们提到slice_size 会进行16对齐, 这里还会判断一次输入/输出是否是128(16 * 8 bit)对齐的, 原因在这里:
    NCCL数据加载操作一次是128bit, 这是一个权衡下的最优解, 至于原因我们后面会讲.

    代码片段二

    下面是核心处理逻辑:
    我们先看图(一个warp的视角下, 处理的数据), 再理解代码:
    notion image
    根据数据是否对齐, 会分别走不同的两个分支:
    • 对齐分支, 调用ReduceCopy128bMulti:
      • 第一步数据会对PACKELEMS*UNROLL*WARP_SIZE 进行对齐, 一个warp也处理这么多数据, 其中一个线程一次循环处理PACKELEMS个数据, 循环UNROLL次, 如果数据处理完毕就在这里退出;
      • 如果还有数据, 就会进入第二步, 这时不进行循环展开, 处理完剩下的数据(只有部分线程工作)
    • 非对齐分支, 调用ReduceCopyMulti:
      • 这时每个线程一次循环只处理一个数据, 会循环展开UNROLL * PACKELEMS/2 次, 如果数据处理完毕就在这里退出;
      • 如果还有数据, 就会进入第二步, 这时不进行循环展开, 处理完剩下的数据(只有部分线程工作)
    这两个函数都是数据搬运操作, 只是数据的粒度不一样, 逻辑基本一致, 这里我们只讲解第一个. 第一个ReduceCopy128bMulti 是主要的数据搬运过程(main loop), 因为前面我们将数据切分为slice并且进行了对齐, 我们看下其中的逻辑

    ReduceCopy128bMulti

    依然先把形参补充上:

    代码片段一

    首先计算了inc , 这个变量表示整个block一次循环处理的数据量(单位是Pack128, 这次考虑了整体的warp数量), 那么每一个线程要处理的数据就是上面offset 的位置, 关系如下图:
    notion image
    然后就是从srcs/dsts中加载数据的起始指针位置, 对于send 场景, 只有一个输入和一个输出

    代码片段二

    上面是精简后的代码, 逻辑比较清晰, 我们先看几个重要的点

    Fetch128

    这段代码的目的是从GPU的全局内存(global memory)中加载128位数据到当前线程的寄存器中, 这是一个内联汇编(PTX)代码, 为什么使用汇编? 因为C++加载(如 v=*p )可能会被编译器优化并且也不好保证内存一致性. 讲解一些关键点:
    • asm volatile : 告诉编译器不要优化这个指令, 避免缓存不一致, 因为我们是跨GPU, 跨进程搬运数据
    • ld.volatile : 依然是避免优化, 确保数据从内存中加载, 而不是缓存(跨GPU的数据搬运, 数据可能被其他线程修改)
    • v2.u64 : 向量加载2个无符号64位值, 也就是Pack128
    • memory : 表示这个指令修改了内存, 编译器需要在前后插入内存barrier, 防止指令重排, 确保数据加载的原子性和一致性

    为什么使用128位取数据?

    1. 内存访问效率
        2. NVIDIA GPU想要最高效的访问全局内存, 需要线程访问是”合并”的(coalesced): 一个warp(32线程)内的线程应该访问连续的内存地址, 形成一整块数据的访问.
    1. 指令数量
        2. GPU是根据指令操作数据的, 一般而言, 取同样的大小的数据使用的指令越少就越高效. 那么就会采用向量化加载, 减少内存事务
    1. 缓存大小及寄存器数量
        2. 而最高效的访问单元是32字节(256位), 但实际上, 128位才是常见的向量大小, 因为单个线程加载128位, 一个warp加载32*128bit=4096bit=512bytes 接近缓存大小(也正好是一个step的大小), 能最大化利用带宽; 另外线程的寄存器也是有限的, 256位在部分场景(多个输入, 多个输出)下可能会溢出, 导致性能下降
    1. 兼容性
        2. NCCL需要支持多个架构的高效传输, 128位在不同的架构下都可以高效加载
    所以128位是权衡了以上因素的综合考量

    while循环

    上面的循环就很简单了:
    1. 将本地数据加载到寄存器, 粒度是Pack128
    1. 将寄存器中的数据存储到对端的buffer
    1. 更新src/dst的位置, 更新offset的位置, 继续迭代
     
    到这里send场景才算结束, 完完整整走了一遍流程, 我们接下来讲RecvReduceCopySend , 看看复杂的场景是如何操作的, 这里我们的讲解就不会很细致了, 只会针对与send场景不同的地方去梳理

    RecvReduceCopySend

    TODO

    © 木白 2024 - 2025