我们同学在压测时遇到过压测进行到一段时间之后会几率性阻塞, 于是让我帮忙看一下. 于是我拉起了我的开发调试镜像, 里面存放着各种预编译好的, 带有调试信息的 pytorch, nccl 等一大坨我们会用到的组件, 然后开始启动压测复现, 复现过程不再多提, 最终我拿到了一个现场, 赶紧挂起 cuda-gdb 到一个 worker 进程, 可以看到:
0x00007f6489d08c9b in sched_yield () from /usr/lib/x86_64-linux-gnu/libc.so.6
(cuda-gdb) bt
#0 0x00007f6489d08c9b in sched_yield () from /usr/lib/x86_64-linux-gnu/libc.so.6
#1 0x00007f634666dd86 in waitWorkFifoAvailable (comm=0x5651ee3c3a30, desiredSent=4324720) at enqueue.cc:1030
#2 0x00007f634666ded3 in uploadWork (comm=0x5651ee3c3a30, plan=0x565211c5cc80) at enqueue.cc:1063
#3 0x00007f634666f776 in ncclLaunchKernelBefore_NoUncapturedCuda (comm=0x5651ee3c3a30, plan=0x565211c5cc80) at enqueue.cc:1342
对应着代码:
comm->workFifoAckdMin = ackdAll;
// See if that was enough.
if (!rollingLess32(comm->workFifoAckdMin + comm->workFifoDepth, desiredSent)) break;
sched_yield(); // 1030 行
嗯…! workFifo 满了?! 首先简单介绍一下 workFifo 的作用. 关于 nccl 具体各个模块源码介绍可以参考我的 NCCL 源码解码系列.
在 nccl 中, 首先会使用结构 ncclInfo 来存放着一次集合通信相关参数信息, 在一次 ncclGroupStart, ncclGroupEnd 期间用户可能会发起多个集合通信操作. 在 ncclGroupEnd 中, nccl 会为当前 group 中发起的每一个集合通信操作创建对应的 ncclTaskColl 结构/ncclDevWorkColl 结构, ncclDevWorkColl 用于给 kernel 层提供相关信息. ncclKernelPlanner.collWorkQueue/.collTaskQueue 分别等价于 List<ncclDevWorkColl>, List<ncclTaskColl> 存放着本 group 中发起的所有集合通信操作相关结构.
之后 planner 会根据任务情况创建多个 ncclKernelPlan, 一个 plan 对应着一次 launch kernel. ncclKernelPlan 中存放着从 planner.collWorkQueue/.collTaskQueue 抽取出来的若干个集合通信操作相关信息.
+-------+-------+-------+-------+-------+-------+-------+-------+
| task1 | task2 | task3 | task4 | task5 | task6 | task7 | task8 |
+-------+-------+-------+-------+-------+-------+-------+-------+
| plan1 | plan2 | plan3 |
countOneBits(ncclKernelPlan.channelMask) 确定了 plan 包含多少个 channel, 一个 channel 对应着一个 thread block. blockIdx.x 即等同于 channel id. scheduleCollTasksToPlan 会确定 plan 使用多少个 channel, 以及 plan 内每个 task 具体使用哪些 channel, 比如上图 plan2 假设会有 4 个 channel, 其中 plan2 channelId=1 的 channel 需要处理 task3 后半部分数据, task4 所有数据, task5 前半部分数据
task 分配到的 channel id 集合
task3: [0, 1]
task4: [1, 1]
task5: [1, 2]
task6: [2, 3]
ncclKernelPlan.kernelArgs.workBuf, ncclDevKernelArgs.workBuf; 其指向着一个内存块, 这个内存块存放着归属于 plan 所有 taskcoll 对应 ncclDevWorkColl 结构. ncclKernelPlan.workBytes 表明所有 taskcoll 对应 ncclDevWorkColl 即 workBuf 指向的内存块至少要具有 workBytes 字节才能容纳所有 task. ncclKernelPlan.workStorageType 指定了 workBuf 指向内存块的位置, 若其为 ncclDevWorkStorageTypeArgs, 则如 finishPlan 所示, 此时 ncclKernelPlan.kernelArgs 指向的内存块布局:
batchZero
|
+-------------------+--------------------------+----------------------+
| ncclDevKernelArgs | channel ncclDevWorkBatch | task ncclDevWorkColl |
+-------------------+--------------------------+----------------------+
若其为 ncclDevWorkStorageTypeFifo, 则如 uploadWork 所示, 意味着 workBuf 指向着 ncclComm.workFifoBuf. 若其为 ncclDevWorkStorageTypePersistent, 则如 uploadWork 所示, 会使用 aligned_alloc 为其分配空间.
ncclDevWorkBatch, 如上所示已知一个特定的 channel 可能会处理多个 taskcoll, channel 会将具有相同 workType/devFuncId 配置且连续的 taskcoll 划分为一个组, 这里 “连续” 是指 taskcoll 对应 dev work coll 结构在 workBuf 中连续存放. channel 使用 ncclDevWorkBatch 来存放同一组 taskcoll 相关信息. ncclDevWorkBatch.offsetBase 指定了这组 taskcoll 第一个 task 在 workBuf 的偏移. offsetBitset ncclDevWorkBatch 包含了哪些 task, 比如 offsetBitset = 0b111, 则意味着 ncclDevWorkBatch 包含了 3 个 task, task0 在 workBuf 中偏移是 offsetBase + 0 * sizeof(WorkStructType), task1 是 offsetBase + 1 * sizeof(WorkStructType), task2 是 offsetBase + 2 * sizeof(WorkStructType). 这里 WorkStructType 是指 ncclDevWorkColl/ncclDevWorkCollReg/ncclDevWorkP2p. channel 会根据这些信息在运行时从 workBuf 中取出 ncclDevWorkColl 结构并执行. 在了解这些设定之后, 这个 BUG 基本上一眼就能看出来了:
static ncclResult_t uploadWork(struct ncclComm* comm, struct ncclKernelPlan* plan) {
bool persistent = plan->persistent;
int channelUbound = plan->channelUbound;
struct ncclWork* workHeap;
if (!persistent) {
workHeap = comm->workFifoHeap;
} else {
workHeap = ncclMemoryStackAlloc<struct ncclWork>(&comm->memScoped, nWork);
}
while (q != nullptr) {
q->work.header.inFifo = !persistent ? 1 : 0;
q->work.header.doneAcks = ix+1;
// BUG! persistent=true 时 q 并没有放入 fifo 中, 没有必要更新 sent.
comm->channels[c].workFifoSent = ix+1;
}
}
如上所示: 在 persistent=true 即表明当前正在进行 graph capturing 时, workBuf 指向着 ncclDevWorkStorageTypePersistent. 在 persistent=false 时 workBuf 指向着 ncclDevWorkStorageTypeFifo. 但在后面将 q, 对应着 ncclDevWorkColl 结构保存在 workBuf 中却总是无脑更新了 fifo 相关字段信息. 这可不就导致 fifo 被误认为已满了么! 问题的修复也很简单: 升级到 nccl 最新版本即可, 如上所示在 nccl 引入了 ncclDevWorkStorageTypeArgs 类型 workBuf 之后, 针对 LLM 推理过程中对集合通信操作的使用情况, 其基本上是用不到 fifo 了.
B.T.W 关于 persistent 之前看代码时也是能发现有点问题的:
ncclKernelPlanner.persistent, ncclKernelPlan.persistent; 若为 true 则表明当前正在进行 capture. 其赋值位于 groupLaunch -> doLaunches -> ncclLaunchPrepare 链路, 但却在 ncclPrepareTasks -> registerCollBuffers 使用, 存在 use-before-assign 情况. 目前 master 分支(20250214) 已经修复了这一情况, 会在 ncclPrepareTasks 就提前设置为 persistent.
