前言
本文章是笔者在学习NCCL源码时候的感悟和理解,由于源码实在抽象,所以笔者尽量用更直白通俗的语言表达自己的理解。若有错误,望各位大神多多批评指教!该文章会不断更新修改,一些细枝末节也将会不断补充,新的文章未来也会采用链接形式纳入本文章中。
一、NCCL任务流总框架
NCCL的整体工作流无疑是复杂的,但是通过分层拆解又显得清晰明了。我根据我的理解划分为七个阶段。为了方便理解,我把这个流水线比作成一个外卖流水线。(我说白了就是HPC赛博外卖一条龙服务)
阶段 1:API 收请求。
这一层做的事是直接面向用户的。用户调用一个集合通信接口比如Allreduce,然后NCCL接收请求。这一层的代表文件是 collectives.cc。像 ncclAllReduce(...) 这样的 API,会先构造一个 struct ncclInfo info,把操作类型、buffer、count、datatype、op、comm、stream 等参数装进去,然后调用 ncclEnqueueCheck(&info)。所以这一层的职责不是执行,而是把用户请求标准化。
综上,可以理解为这一层就是一个服务员记录一下用户的订单,但是不能乱写,得格式化记录。
阶段 2:group 暂存请求
如果你看到一些使用NCCL的代码,比如torch的PG处理的部分,你会发现大量的GroupStart和GroupEnd。这个其实也非常好理解,Group相关操作设计的目的就是为了批量收集和整理用户请求。通过这两个接口包裹住了一个收集请求的上下文。这么说还是有点抽象,可以理解为用户向服务员点单,但是服务员总不能记录一个跑就一趟后厨,所以拿了个本子记下来一起送到后厨。 这样做当然好处多多,比如防止死锁等等。具体的就看官方文档把,我懒得写了。
这一层的代表文件是 group.cc。group.cc 里用 thread_local int ncclGroupDepth 记录当前线程 group 嵌套深度,还维护了当前线程的 comm 链表、异步 job 队列等 group 状态。这说明 group 本质上是线程级的批量提交上下文 。ncclInfo → 进入 enqueue 暂存。ncclGroupStart() 进入这个上下文,ncclGroupEnd() 调 ncclGroupEndInternal(),真正触发后续准备和发射。
插一句,这里的嵌套深度就是,一个Group上下文中可以嵌套好几层Group上下文,但是只有在最外层才会调 ncclGroupEndInternal()。
阶段 3:planner 整理任务
这一部分的任务就是要整理和组织规划已经格式化 一组任务(Group)。这里可以理解为,一大堆订单送到了厨房后厨,如果来一个菜就单独做,那就效率很低。(一下来100份拼好饭后厨也一个一个做要累死厨子)所以这时候就要有个chef主厨,负责专门来整理和分配这些任务给厨房的打工人。这里的打工人可以理解为就是设备。 这个chef在NCCL里边就是planner。这个planner挂在COMM下。
主要的一个函数是ncclPrepareTasks,文件在 enqueue.cc,接口声明在 enqueue.h。到了 groupEnd 触发处理时,planner 开始接手。ncclPrepareTasks(...) 的职责就是把已经收集到的用户任务整理到 planner 体系里,为后面剥离成 plan 做准备。从这一步开始,任务进入 planner 的内部组织结构。也就是暂存请求 → planner 内部 task 体系。
阶段 4:task 翻译成 devWork
这一步要做的事就是要做一个语义下沉。也就是主厨要把整理好的任务(task)翻译成厨房打工人能听的懂的话(devWork)(打工人一看要做"薯了算了"人都傻了,主厨说说白了就是薯片夹烤鹅干,我插一句,肉问屋的苹果包肉还挺好吃的但是不如烤肉有李的韩式烤肉)。 ncclTasksRegAndEnqueue() 明确会"为每个 task 构造一个 ncclDevWorkColl... 结构"。说明这一步是把"逻辑任务"翻译成"设备工作描述"。
阶段 5:把 devWork 装进 Plan
这里就要把devWork转变成Kernnel Plan了。这里的KernnelPlan中已经有之后启动Kernel参数了。核心的函数如ncclAddWorkBatchToPlan、finishPlan,都在 enqueue.cc。这一步可能会把多个work放在一个plan中。这样有点类似于算子融合的概念,但是本质不尽相同。主要目的是为了合并任务,减少launch。也就是打工人要把自己要做的事组织一下,把相同的订单直接一锅出了。
阶段 6:launch 生命周期执行
现在就是Plan准备好了,终于可以开始发射任务了!也就是终于要炒菜了! 相关函数有ncclLaunchPrepare、ncclLaunchKernelBefore_NoUncapturedCuda ncclLaunchKernel ncclLaunchKernelAfter_NoCuda ncclLaunchFinish
这些接口在 enqueue.h 中是明确分开的,所以要把它们理解成:NCCL 的发射不是一行函数就完了,而是一个小型生命周期
阶段 7: 真正执行
到这里,NCCL 已经不再处理"原始集合通信请求"了。GPU 看到的是:kernel args 、channel mask、devComm、work batch、devWork 描述。
然后设备侧 kernel 才去真正推进规约、搬运、同步等动作。finishPlan(...) 把这些 launch 所需的关键信息写入 plan->kernelArgs,就是为这一刻服务的。
二、NCCL任务流详解
阶段 1:API 收请求。
由于各种NCCL API的参数不尽相同,所以为了方便处理,会统一通过ncclInfo来收集。ncclInfo 是 NCCL 任务流中的第一张标准任务卡。它把不同集合通信入口的参数差异,统一规约成后续阶段都能处理的内部描述。这一部分的核心函数就是ncclEnqueueCheck(&info),对外屏蔽接口参数差异,对内把任务单向下层提交。
在源码中,不同操作构造 info 时会带不同的宏如:ALLREDUCE_CHUNKSTEPS。这说明 API 收请求阶段除了基本语义字段外,还会把与该通信类型默认切分方式相关的一些参数一并带上。也就是说,这张"订单卡"不只是写菜名,还顺手写了一些后厨处理这类菜常用的制作参数。
阶段 2:group 暂存请求
group 这一层并不负责执行通信本身,它负责提供一个"批量收集、延后统一处理"的上下文,把零散请求先攒起来,等到合适时机再统一进入准备、预连接和发射流程。 这和 group.cc 的实现是一致的:文件开头就维护了一组**thread_local** 的 group 状态,包括 group 嵌套深度ncclGroupDepth、group 错误状态ncclGroupError、按任务类型分类的 comm 链表ncclGroupCommHead、预连接链表ncclGroupCommPreconnectHead、异步 job 队列ncclAsyncJobs,以及 blocking/nonblocking 模式记录。要注意这些变量的定义都是线程级的,在同一个线程中是可见的,一个线程里连续调用多个 NCCL API,只要处于同一对 GroupStart/GroupEnd 之间,这些请求就会被归到同一个 thread-local 的 group 状态中。上文已经说过,这个上下文的定义是支持嵌套的,但是只有最外层的end才会把这一批任务向下提交。
当最外层 ncclGroupEnd() 到来时,NCCL 不会只做一件事,而是开始对这一批在 group 中暂存的请求做一次统一结算 。从源码看,这个结算大致会触发 5 类工作:先收束 group 状态,然后准备任务与建链,再处理异步 job,接着执行真正的 launch,最后做清理。group.cc 中对应的关键入口和辅助函数包括 ncclGroupEndInternal()、ncclPrepareTasksAndCollPreconnectFunc()、ncclP2PPreconnectFunc()、asyncJobLaunch()、doLaunches() 和 groupCleanup()。
1)确认这次 group 是否真的结束。 group 是允许嵌套的,因此不是每次 ncclGroupEnd() 都会立刻触发后续流程。只有当 thread_local 的 ncclGroupDepth 退回到最外层边界时,这一批请求才算真正"收齐",后面的 prepare、launch、cleanup 才会启动。源码里 ncclGroupDepth 就被明确定义为 group 嵌套深度。
2)对本组涉及的 comm 做准备工作。 对 group 中收集到的 communicator,NCCL 会开始真正的任务准备流程。ncclPrepareTasksAndCollPreconnectFunc() 里先切换到对应设备 cudaSetDevice(comm->cudaDev),必要时设置 CPU affinity,然后调用 ncclPrepareTasks(comm, ...)。如果这批 collective 之后需要连接某些算法路径,还会继续做 ncclCollPreconnect(...)。这说明 GroupEnd 后的第一件大事不是 launch,而是"先把待发射任务整理好,并把需要的连接建起来"。其实这也就是把任务整理成Task并且准备的过程。
3)处理 P2P / collective 相关的预连接。 对于不同类型的任务,group end 后还可能触发预连接逻辑。比如 ncclP2PPreconnectFunc() 会做 ncclTransportP2pSetup(...);而 collective 的预连接会根据算法类型走 ring/tree/NVLS/CollNet/PAT 等不同连接建立路径。也就是说,GroupEnd 后不仅是"准备发任务",还会确保后面的传输路径已经 ready。
4)把 group 里挂起的 async job 统一执行: 如果这一批 group 中还挂着内部异步 job,GroupEnd 也会把它们统一结算。源码里的 ncclAsyncLaunch(...) 表明:不在 group 里时 job 可以直接执行,而在 group 里时 job 会先被挂进 ncclAsyncJobs 队列;等到 group end,再由异步线程或主线程统一处理、回滚或清理。++注意,这里说的async job指的并不是我们说的通信语义,而是一些与正式通信相关的辅助性工作,借用上面的例子你可以把它理解成后厨里的"杂务工单"++ 。
**5)正式进入 launch 流程。**这一部分将会在后续部分详细讲解。
6)统一清理本组状态 。这一批任务处理完成后,groupCleanup(...) 会负责把本组的 thread-local group 状态清掉,同时把 comm 从 group 链表中摘出,清理 preconnect 标记,回收被放弃的 plan 内存,并对 async job 做 undo / destructor 等收尾。也就是说,GroupEnd 既是一次批处理的提交边界,也是一次批处理事务的清理边界。
从数据结构角度看,NCCL 的 group 与 planner 分属两个不同层次。group 是线程级上下文,通过 thread_local 变量维护本组状态,包括按任务类型组织的 communicator 链表 ncclGroupCommHead[...]、预连接链表 ncclGroupCommPreconnectHead 以及异步 job 队列 ncclAsyncJobs。因此,group 层面首先组织的是"本组涉及哪些 comm"。而每个 comm 内部又挂有自己的 planner,后者负责维护该 communicator 下尚未发射的 task、device work 和 kernel plan,例如 collSorter、peers[...]、wipPlan、planQueue 和 unlaunchedPlansHead 等结构。换句话说,group 解决的是"这批请求以哪些 communicator 为单位统一提交",planner 解决的是"某个 communicator 内部的任务如何整理、翻译并最终形成可发射 plan"。
当前线程(thread_local)
├── ncclGroupDepth
├── ncclGroupCommHead[type] ----> commA -> commB -> commC (本组 comm 链表)
├── ncclGroupCommPreconnectHead --> commX -> commY (需要预连接的 comm 链表)
└── ncclAsyncJobs ----> job queue
每个 comm
├── planner
│ ├── collSorter / 各类 task 统计信息
│ ├── peers[...] / bcastQueue / 各类 task 队列
│ ├── wipPlan
│ ├── planQueue
│ └── unlaunchedPlansHead ----> plan1 -> plan2 -> ...
├── groupNext[type] (把自己挂进当前线程 group 链表)
├── preconnectNext (把自己挂进 preconnect 链表)
├── connectSend/connectRecv (连接需求位图)
└── 其他 comm 自身状态
根据上面的示意图可以看出,comm 确实保存了"当前 communicator 尚待处理的任务状态",但这些状态不是单一一个链表,而是分散组织在 comm->planner 及若干专项任务队列中。 而Groupend阶段通过保存的相应comm链表或队列,找到对应的comm,对每个comm调用自己的planner相关流程。
说到这里,可能很多人都会发现一个问题,以上说的这些都是一个进程之内的事,收集任务统一处理固然好,但是多进程之间怎么协调呢。要想理解这个问题,要注意到两个前提,首先与MPI的一些通信组组织形式不同的是,每个rank都保有comm这个对象或者结构体,其次对于NCCL的API调用顺序是需要要调用层去保持顺序一致的,这些在torch的Processgroup的部分都有所体现。在这两个前提下,不同 rank 之间的 collective 协同,并不是先把本地任务以任意顺序异步排队,再通过 barrier 事后"等"出一致顺序。更准确地说,每个 rank 都在自己本地的 comm 和 planner 中形成一致对应的 task/plan 序列。随后在 GroupEnd 之后,NCCL 再通过 group launch 阶段的 barrier 与轮次推进机制,协调这些本地已经排好序的 plan 有序发射。因此,barrier 主要承担的是"推进节奏协调"的作用,而不是"乱序修复"的作用。
阶段 3-5:planner 组织任务
把阶段 3-5 合起来看,源码里的主线可以压成下面这一条:comm->planner 接住本地 task,→ 从 planner->collTaskQueue 逐个取 task→ 为每个 task 构造 ncclDevWork...→ 按 channel 和兼容性规则把 work 组织进 batch(这里就是一种灵活的增量思想)->把各 channel 的 batch、proxy op 和 kernel args 收束成 ncclKernelPlan.
ncclInfo -> task 的转变,不是在 ncclPrepareTasks() 里才发生的,而是更早,在 ncclEnqueueCheck(info) 这条入口链里就开始了。collectives.cc 里的各个 API 先构造 ncclInfo,然后统一调用 ncclEnqueueCheck(&info);而 enqueue.h 也把 ncclEnqueueCheck明确列为 enqueue 阶段的总入口。ncclEnqueueCheck(info) 根据 info 的内容,把它落成具体的 task 结构,ncclPrepareTasks() 再去整理这些已经存在的 task。这一阶段最关键的结构思想:planner 是 comm 内部的"任务组织器"。NCCL 没把任务整理做成一个跨 comm 的全局调度器,而是把 planner 内嵌到每个 comm 里,让每个 communicator 各自维护自己的任务组织状态。 从 ncclTasksRegAndEnqueue(struct ncclComm* comm) 一开始就取 &comm->planner 可以直接看出来,后续所有整理逻辑都围绕当前 comm 的 planner 展开。
cpp
struct ncclKernelPlanner* planner = &comm->planner;
当然,planner 不是直接面对原始 API 调用,而是面对已经进入本地任务队列的 task。
cpp
task = ncclIntruQueueHead(&planner->collTaskQueue);
while (task != nullptr) {
这里很值得点一下 intrusive queue 的使用风格。NCCL 大量用 ncclIntruQueueHead、ncclIntruQueueEnqueue、ncclIntruQueueDequeue 这种 intrusive queue,而不是简单数组,这说明它在这层更强调:增量组织,低开销拼接,顺序消费。相比于侵入式的单独构造节点,这种组织方式更有利于高性能的程序设计,但是对声明周期管理也需要更小心。
ncclTasksRegAndEnqueue 是 task → devWork 的核心入口。这个函数是阶段 3-5 的一个关键枢纽,是任务语义向设备层下沉的一个关键点。为什么要在kernel plan和task中插入这样的一层关于设备任务的中间态描述?这种分层的好处是:后面无论要做 batch 合并、channel 分治、proxy op 插入,还是 work 放进 args 还是 args 外,都有一层清晰的中间表示可以操作。
当一个 task 已经被翻译成设备 work 后,并不会立刻形成最终 plan。NCCL 不是把 work 平铺到一个全局列表里,而是先把它放进"当前 WIP plan 的某个 channel 槽位"里。NCCL 会调用 ncclAddWorkBatchToPlan(...)这个函数一上来就取:
cpp
ncclKernelPlanner::WipPlan::Channel* chan = &comm->planner.wipPlan.channels[channelId];
也就是说,阶段 3-5 从很早开始就是按 channel 组织的,channel 不是最后 launch 前才出现的概念。WIP plan 可以理解为"施工中的 plan",而 channel 则是这片施工现场中的多个工位 。 ncclAddWorkBatchToPlan() 最重要的工作不是"简单追加",而是判断 当前这个 work 能不能接到已有 batch 后面。源码里它会检查很多条件,比如:当前 batch 的 workType 是否一致、当前 batch 的 funcId 是否一致、P2P 场景下 p2pEpoch 是否一致、当前批次里是否已经出现过同样的 p2pRound、当前 batch 是否已经达到 work 数量或字节上限,在设备表示能力和执行约束下,把兼容的 work 合并进同一个 batch。如果不兼容,就必须开新 batch;如果偏移过大,还可能开 extension batch。
除了主线 work,NCCL 还会在这条流水线里处理需要的 proxy op。对应函数是 ncclAddProxyOpIfNeeded(...)。它会先判断这个 op 是否真的需要;如果需要,就把它分配出来,然后直接挂到:
cpp
comm->planner.wipPlan.channels[op->channelId].proxyOpQueue
这说明 proxy op 不是散落在 plan 外面的旁路信息,而是和主线 work 一样,也被纳入当前 WIP plan 的 channel 局部结构里。这一步的主要作用可以概括成:把辅助通信动作也收束到当前 plan 的施工现场里,方便后面统一封装。
当一个 WIP plan 已经积累了足够的 batch、work 和 proxy op 之后,就进入 finishPlan()。这个函数的主要作用,就是把前面那些分散在 wipPlan.channels[...] 里的局部状态,统一封装成一个最终可发射的 ncclKernelPlan。 它在源码里主要做几件事。第一,统计和决定存储方式:它先统计 workBytes 和 batchBytes,然后判断这次 plan 的 work 是否能直接塞进 kernel args;如果能,就使用 ncclDevWorkStorageTypeArgs。这一步决定的是:这包 work 最终放在 kernel 参数区里,还是放到参数区外的别处。第二,分配并填充 kernelArgs。finishPlan() 会计算 plan->kernelArgsSize,分配 plan->kernelArgs,然后把几个核心字段写进去:plan->kernelArgs->comm = comm->devComm、plan->kernelArgs->channelMask = plan->channelMask、plan->kernelArgs->workStorageType = plan->workStorageType。到 finishPlan() 结束时,这个 plan 已经带上了设备侧 launch 所需的关键信息。第三,把各 channel 的 batch 按设备要求的布局排进参数区,host 侧的封箱方式,是严格围绕设备端 kernel 的读取模式来设计的,而不是"怎么好写怎么排"。第四,把各 channel 的 proxy op 合并进 plan->proxyOpQueue。这一步可以概括为:finishPlan() 负责把施工中的 WIP plan 真正封成一个 launch-ready 的 ncclKernelPlan。
需要注意的是,对于任务的合并主要发生在work拼接入batch以及batch转化为kernel plan的阶段。,一般来说group上下文中收集到的通信请求会被一对一转化为task,每个task转化为一个devwork工作描述,但是多个兼容的work会被拼入某个channel的batch中进行任务合并,多个。batch可能会被拼入一个kernel plan
以一个Allreduce调用为例子,总结一下上面说的:
cpp
ncclAllReduce(...)
→ ncclInfo
→ 当前线程 group 中属于某个 comm 的本地请求
→ comm->planner 里的一个 collective task
→ ncclTasksRegAndEnqueue() 为它生成 ncclDevWork...
→ ncclAddWorkBatchToPlan() 把它并入 wipPlan.channels[channelId] 下的某个 batch
→ 相关 proxy op 进入 wipPlan.channels[channelId].proxyOpQueue
→ finishPlan() 把所有 channel 的 batch 和 proxy op 收束为一个 ncclKernelPlan
线程的 group 上下文
└── 收集本批次涉及的 comm
某个 comm
└── planner
├── collTaskQueue
│ └── task1 / task2 / task3 ...
│
└── wipPlan
└── channels[channelId]
├── workBatchQueue
│ └── batch1 / batch2 ...
└── proxyOpQueue
└── proxyOp1 / proxyOp2 ...
阶段 6:launch 生命周期执行
plan 已经在前面准备好了,launch 阶段要做的就是:按 comm 取 plan、按轮次推进、在正确的设备和同步语义下,把 plan 真正发出去。这条主线集中在 group.cc 的 doLaunches(...) 中。
cpp
static ncclResult_t doLaunches(struct ncclComm* head)
这个函数的输入不是 task,也不是 work,而是 group 中已经准备好待发射 plan 的 comm 链表头 。每个 comm 下面已经有 unlaunchedPlansHead 这样的待发 plan 队列。launch 阶段只关心 ,哪些 comm 要发,每个 comm 还有几个 plan 没发 ,本轮要发哪个 plan。
如果考虑的是主流训练场景(一卡一进程一rank),下面灰字可以直接跳过
launch 阶段的"通信组协同"并不是通过一个全局共享队列来完成的,而是通过 clique + 轮次 barrier 的方式完成的。具体来说,NCCL 会先把当前进程内 intraComm0 相同的本地 comm 视作同一个 clique;随后,doLaunches() 对该 clique 先统一执行 ncclLaunchPrepare,再在 group launch mode 下通过 ncclCommIntraBarrierIn/Out 进入按轮次推进的发射过程。每一轮中,每个 comm 最多从自己的 unlaunchedPlansHead 中弹出一个 ncclKernelPlan 发射;发完之后,再通过 ncclCommIntraBarrierIn(comm, hasMore ? 1 : 0) 上报自己是否还有待发 plan。barrier 的归约结果决定 clique 是否进入下一轮,因此同一组 comm 会以"每轮最多消费一个 plan"的方式协同推进,直到所有成员的 plan 队列都清空,最后再统一调用 ncclLaunchFinish 结束 launch 生命周期。
需要注意的是,在主流深度学习训练框架中,通常采用"一进程一 rank"的运行模式;但从 NCCL 源码设计看,rank 并不等同于操作系统进程。源码中不仅显式支持同一进程内多个 communicator 的协同,还考虑了"一个进程管理多个 local rank"的更一般场景。因此,intraComm0 等字段更准确地反映的是"同一进程内多个本地 comm 的组织与协调",而不能简单理解为"每个进程天然只会有一个 NCCL 参与者"。
在一进程一 rank 一 GPU 的主流训练模式下,NCCL 的 launch 生命周期仍然完整存在,但 clique 往往退化为单成员,ncclCommIntraBarrierIn/Out 也随之退化为单成员 barrier。此时 launch 的实质,就是对当前这个本地 comm 的 unlaunchedPlansHead 按轮次逐个弹出 kernelPlan,执行 Prepare → Before → Launch → After → Finish 这套生命周期,直到 plan 队列清空为止。
阶段 7: 真正执行
涉及到具体下层网络层级,我还没研究明白,先问问C老师和D老师以及G老师吧,后续待补充....
附录:名词解释(不是很严谨仅做参考)
Proxy:这里的proxy代理指的并不是在计算机网络中我们熟知的那种负责转发流量的中间服务器或者线程。在NCCL通信的时候,有的任务并不适合Gpu kennel来工作或启动,比如和一些网络设备协同,管理一些host侧资源,协调跨机网络传输等等。这里的Proxy指的就是一些处理更偏向host端相关网络事务的相关线程或者代理服务逻辑。
**batch:**这里的bacth不是模型训练里边的batch,而是NCCL内部的一组统一处理的工作,比如一组操作,一次打包提交,若干个task的集合。
**channel:**逻辑信道,可以理解为在物理链路至上的逻辑流水线,
**loop、chunk、slice、step:**这些是数据块大小的命名,对应了不同层级处理数据块的单位大小。比如对一个Allreduce,由于buffer太大一次处理不完,所以就会通过循环(loop)来循环处理。每个loop要处理的数据要被切分成chunk,这个chunk也是算法层面上的(比如典中典之ring)的数据粒度,各个 rank 在环上传递一个个 chunk。当网络层级再往下,若干chunk可能会被分配给某一个channel,为了实现流水线,更好地overlap,chunk又会被切成slice,这就是流水线级别的数据粒度。当然一个流水线处理slice的过程中又会有很多step。