首页 / 资讯中心 / 文章详情

NCCL Ring LL 源码解析与 PCIe IPC 实现 AllReduce 性能优化实战

NCCL Ring LL 源码解析与 PCIe IPC 实现 AllReduce 性能优化实战 ★ FEATURED ARTICLE
之前做性能优化的时候我基本都是把 NCCL 当黑盒用脚本里一个init_process_group完事底层跑的是 Ring 还是 Tree、用的 LL 协议还是 Simple 协议我完全不关心。直到这次要在一批只有 PCIe 互连的异构机器上压多卡 AllReduce 带宽我才被迫把 NCCL 源码从头翻了一遍从 Ring 算法的流水线看到 LL 协议的设计初衷最后甚至忍不住自己用进程间通信IPC在 PCIe 上撸了一个 AllReduce 实现来对比验证。这篇文章把这段从读源码到写实现、再回头查硬件的完整过程记录下来内容包括 Ring LL 的原理拆解、PCIe IPC 通信的工程细节以及我在实际环境中踩过的坑。1. 项目背景与核心诉求1.1 从黑盒调用到源码级理解先说清楚这个项目为什么会存在。当时我手里的机器是几台普通的 PCIe 双卡工作站没有 NVLinkGPU 之间全部走 PCIe Switch。跑一个 7B 参数模型的张量并行训练AllReduce 占了整个通信量的绝大多数实测吞吐比预期低了一大截。我做了几个常规尝试调大 batch size、换通信后端、开torch.compile效果都有限。最后只能走最笨的路——打开 NCCL 源码一行一行看它到底在 PCIe 环境下做了什么。这里我建议所有做分布式训练的人都至少完整读一遍nccl/src/collectives/device/all_reduce.h和nccl/src/transport/p2p.cc。之前我也觉得 NCCL 封装得太深源码又用了大量宏和模板读起来很劝退。但真正耐下心看进去就会发现它的设计其实非常朴素上层是算法选择Ring、Tree、CollNet中间是协议选择LL、LL128、Simple下层是传输通道P2P、SHM、NET。而我们这次的性能瓶颈恰好就集中在算法 协议 传输三层叠加在一起的那个路径上。读源码最大的收获不是记住了某个函数的名字而是搞懂了 NCCL 的决策逻辑它在什么时候选 Ring为什么 PCIe 环境下通常比 Tree 更适合小中型集群以及 LL 协议到底是牺牲了什么换来了低延迟。这些信息在文档和博客里也有但远没有直接从代码里来得准确、来得实在。1.2 Ring LL 在通信栈里的位置与价值多卡 AllReduce 的数学含义很简单把 N 张卡上的张量逐元素求和再广播回每张卡。但工程实现上如何在有限的拓扑带宽里把这个操作做得又快又稳才是真正复杂的地方。NCCL 提供的 Ring 算法在 PCIe 环境下几乎是默认答案因为它对拓扑结构不敏感只是把数据按环的方式转圈每张卡只跟相邻的卡通信不需要全局同步。而 LLLow Latency协议又是 Ring 算法之上的一个精妙设计它的核心目标是隐藏同步延迟。常规的做法是数据传完以后接收方需要等一个明确的通知再开始做归约。这个通知就是一次额外的通信虽然只有几个字节但在多卡高频小张量的场景下它带来的延迟占比非常可观。LL 协议把通知直接嵌入到数据流里让接收方只要看到某个标志位变化就知道数据已经完整到达省掉了整个同步回合。可以说Ring 负责把通信量摊平LL 负责把延迟压下去两者组合起来就是这次项目的理论起点。不过理论归理论等我真的在 PCIe 环境里跑起来才发现 Ring LL 也不是万能的它在某些 chunk 大小和包大小组合下有效带宽会被 flag 开销拖累。这个时候我才产生了能不能绕开 NCCL自己用 IPC 实现一版 AllReduce的念头。2. Ring AllReduce 与 LL 协议的原理拆解2.1 Ring 算法的数据流与时序Ring AllReduce 的过程可以拆成两个阶段reduce-scatter 和 all-gather。假设有 4 张卡每张卡持有完整的张量 K那我们先把 K 平均切成 4 份 chunk编号为 0 到 3。在 reduce-scatter 阶段每张卡将第 i 份 chunk 传给环上的下一个节点同时从上一个节点接收 chunk收到后与本地对应 chunk 做逐元素相加再把结果继续往下传。经过 N-1这里就是 3步之后每张卡上都有了一份完整的归约结果 chunk。第二阶段是 all-gather。每张卡把自己手里的那份归约结果 chunk 沿着环广播出去同样经过 N-1 步最终每张卡都拿到了所有 chunk 的归约结果拼起来就是完整的全局 AllReduce 结果。这里有个很关键的工程点Ring 算法之所以高效是因为它能把通信和计算重叠起来。在 N 卡环境下每张卡的通信总量大概是单卡发送总量(N-1) × (K/N)单卡接收总量(N-1) × (K/N)总通信量2 × (N-1) × K / N当 N 越大每卡额外承担的数据量越接近总数据量的 2 倍。这也是为什么 Ring 在卡数增多时带宽优势会减弱但对于中小规模 PCIe 集群它依然是最稳健的选择因为每一步都只依赖相邻节点环形拓扑和 PCIe Switch 的物理链路几乎完美对应。2.2 LL 协议的核心机制与设计意图NCCL 的协议层里LL 是最有代表性的一个。它把数据切成非常小的 slot每个 slot 只有 8 字节其中 4 字节放真实数据另外 4 字节放 flag。发送方每写完一个 slot就更新对应的 flag 值接收方通过轮询 flag 的变化来判断数据是否已经可读。这听起来很简单但它解决了两个问题。第一它消除了独立的同步消息。传统 Simple 协议里数据传送和我已经发完了的通知是分开的而 LL 协议把两者放在同一个内存区域flag 就是数据的影子。接收方不需要额外等待一个 doorbell 消息只需要不断地读取 flag。第二它把同步的粒度压缩到了极小。4 字节的 flag 意味着接收方可以逐 slot 地消费数据不需要等一个大的 chunk 全部到了才开始计算这在小消息场景下特别有利。不过我实际看代码的时候发现LL 协议对内存访问模式的要求很苛刻。它要求 GPU 之间必须支持直接的 peer-to-peer 访问否则 flag 轮询就要经过主机内存中转那延迟不仅没有降低反而会比 Simple 协议还高。NCCL 的ncclTransportP2pCanConnect函数里就是根据cudaDeviceCanAccessPeer的结果来做决策的如果不支持 P2PLL 根本不会被选上。2.3 Ring LL 的瓶颈到底在哪吹了这么多 LL 的好处也该说说它的短板。在 PCIe 环境下LL 最大的问题是有效带宽折损。因为每个 slot 只有一半的字节在传数据另一半是 flag理论上它把链路带宽利用率压到了 50% 左右。小消息场景下这完全值得因为延迟的收益远大于带宽的损失但消息一旦变大同样的物理链路上要传双倍的数据量整体吞吐反而会下降。NCCL 自己也意识到这一点所以它引入了 LL128 协议用 128 字节的 slot 格式来摊薄 flag 的占比同时在协议选择上做了启发式判断小的消息走 LL大的消息走 Simple。我在这次项目里最典型的场景是张量并行里的 AllReduce单次消息大小在几 MB 到几十 MB 之间正好落在 Simple 和 LL 的交界区。我试着强制开启 LL带宽直接掉了将近一半强制走 Simple小步长场景延迟又涨上去。这个现象让我意识到光靠 NCCL 的自动协议选择还不够得真正理解自己的数据特征之后才能定出合理的配置。还有一个容易被忽略的点是 chunk 大小。Ring 里数据的切分粒度直接影响流水线的深度。chunk 太小flag 的轮询开销会吃掉计算收益chunk 太大流水线的气泡又会变多。NCCL 里ncclNthreads和nchunks的取值组合在 PCIe 这类高延迟链路上需要仔细测试不能直接照搬 NVLink 环境下的参数。3. PCIe IPC AllReduce为什么会有这个方向3.1 什么是 PCIe IPC 通信聊完 NCCL 的 Ring LL现在回到标题的后半段PCIe IPC AllReduce。所谓 IPC在这里指的是多进程之间通过 CUDA 的 IPC 机制共享 GPU 设备内存。和普通的共享内存不一样GPU 内存的 IPC 是把显存的物理地址映射到另一个进程的地址空间里让进程 B 可以直接读写进程 A 的显存。具体到 PCIe 环境这就牵扯到了 GPU 之间怎么交换数据的问题。早期版本的 CUDA 不允许跨进程直接访问对方 GPU 的显存只能把数据先拷到主机内存再拷到目标 GPU。后来 CUDA 引入了cudaIpcGetMemHandle和cudaIpcGetMemHandle这一组 API只要硬件支持 peer access两个进程就能各自拿到对方的显存映射然后像读写本地内存一样读写远端显存。这个过程本质上是在利用 PCIe 的 DMA 能力。PCIe 协议本身定义了 Endpoint 之间的数据传输机制GPU 的 BAR 空间可以被映射到 CPU 或其它设备IPC 相当于把这条链路用软件方式直接打通了。我在项目里做的事情就是不再依赖 NCCL 的集成了整个上层逻辑的黑盒通道而是用cudaIpcOpenMemHandle打开对端进程的显存句柄拿到远端地址后自己控制数据怎么切、怎么发、怎么归约。3.2 把 Ring 的逻辑搬到 IPC 上既然 Ring 算法本身不依赖任何特殊网络硬件它只是定义了数据流转的时序那我完全可以把 NCCL 里那一套分块归约逻辑抽出来把传输层换成 IPC 映射后的 GPU 内存读写。实现思路分三部分。第一部分是建立连接。每个 rank 在自己的 GPU 上分配一块连续的缓冲区切成 N 份 chunk然后把这块缓冲区的 IPC 句柄广播给所有其它 rank。这里我用的是共享内存来交换句柄所有进程启动时先映射一段公共的 host 共享内存把cudaIpcMemHandle_t结构体塞进去再栅栏同步一次所有人就能拿到完整的对端句柄列表。第二部分是打开远端内存。每个 rank 用cudaIpcOpenMemHandle依次打开其它 rank 的句柄得到一组远端指针。这一步如果硬件支持 peer access返回的指针可以被 CUDA kernel 直接解引用如果不支持就只能退回到cudaMemcpyPeerAsync把数据从本地显存拷到远端显存。我在机器上测试时发现同为 PCIe 总线上的 GPU只要cudaDeviceCanAccessPeer返回 true直接读写和cudaMemcpyPeerAsync的路径都可用但性能差异很大。第三部分是流水线归约。每个 rank 维护两个 CUDA stream一个负责从环上游接收数据并做归约另一个负责把归约结果传给环下游这样收发和计算可以错开。这一步本质上就是我在源码里看到的ncclReduceScatter的 device loop只不过传输从 NCCL 的内部 pipe 换成了我自己的 IPC 缓冲区。3.3 为什么不用现成的 NCCL 通道有人可能会问折腾这么多为什么不直接用 NCCL我的回答是这个项目的目的本来就不是替代 NCCL而是验证在 PCIe IPC 原语之上能不能用最少的软件开销实现一个可用的 AllReduce。NCCL 作为一个通用库它要兼顾的硬件和场景太多了所以在某些特定的 PCIe IPC 直连场景下它的通用性反而成了开销。举个例子NCCL 的 P2P 传输层里为了保证多进程并发访问安全会引入额外的连接管理、内存注册和同步原语。这些逻辑在稳定性和扩展性上非常必要但如果我只想在一个固定拓扑的小集群里压榨最后一丁点带宽这些通用逻辑就会变成多余的开销。我自己实现的 IPC 版本里省略了动态拓扑发现和错误恢复只保留连接建立、环形传输、归约计算三条主线代码量大概只有 NCCL 对应路径的十分之一但控制权完全在自己手里。当然这种精简也意味着功能上的牺牲。NCCL 能自动 fallback 到主机内存传输能处理跨节点网络通信我的 IPC 版本一旦碰到底层 peer access 失败就只能干瞪眼。所以这个项目更准确地说是一个用于理解通信原理的最小复现而不是一个生产级库的竞争者。下面我直接给出我实际搭建和测试这套实现的过程包括环境和代码路径大家如果想在自己机器上复现可以直接参考。4. 实操过程与关键实现4.1 环境准备从 PCIe 枚举到拓扑确认动手写代码之前第一件事是确认硬件拓扑到底长什么样。Linux 下可以用lspci看到 GPU 的 BDFBus:Device.Function编号例如01:00.0表示 bus 1、device 0、function 0。PCIe 枚举的过程就是操作系统在启动时扫描总线上的设备为每个 Endpoint 分配 BDF 和 BAR 地址空间的过程。如果两个 GPU 挂在同一个 PCIe Switch 下面它们的物理距离比较近P2P 延迟更低如果挂在不同 Root Complex 下数据可能要走 QPI/UPI 或额外的 Switch 跳数延迟和带宽都会吃亏。我建议在开始前用nvidia-smi topo -m和lspci -tv各看一遍两个工具的输出对照起来能大致判断 GPU 之间的 PCIe 拓扑关系。特别注意是否出现了 CPU 侧的 NUMA 节点跨接。在我那台机器上GPU0 和 GPU1 在同一个 NUMA 节点下但 GPU2 在另一个 NUMA 节点实测 GPU0 到 GPU2 的 P2P 带宽只有 GPU0 到 GPU1 的 70% 左右。这种差距不提前测出来后面跑 AllReduce 的时候会非常疑惑为什么数据分布不同性能就会差一大截。拓扑确认完之后再验证 peer access 是否可用。直接在代码里调用int can_access 0; cudaDeviceCanAccessPeer(can_access, device_id_a, device_id_b); cudaDeviceEnablePeerAccess(device_id_b, 0);注意cudaDeviceEnablePeerAccess返回cudaErrorPeerAccessAlreadyEnabled在某些 CUDA 版本里是个坑你以为失败了其实已经启用了。正确的姿势是忽略这个错误码继续执行。另外如果进程里用cudaSetDevice切过设备peer access 的启用是针对当前 device不会自动切换多卡环境下要小心。4.2 基于 cudaIpc 的 AllReduce 核心路径代码结构上我把整个实现拆成了三步。首先是分配显存并导出句柄// rank 0 分配显存并导出句柄 float* d_buffer; cudaMalloc(d_buffer, size); cudaIpcMemHandle_t handle; cudaIpcGetMemHandle(handle, d_buffer); // 通过共享内存把 handle 发给其它 rank第二步是每个 rank 打开所有远端句柄// rank i 打开 rank j 的显存 void* remote_ptr nullptr; cudaIpcOpenMemHandle(remote_ptr, remote_handle[j], cudaIpcMemLazyEnablePeerAccess); // 之后可以在 kernel 里直接读写 remote_ptr注意cudaIpcMemLazyEnablePeerAccess这个 flag 很关键它允许打开句柄时不立刻做 peer access 的完整建立而是推迟到第一次访问。对于需要批量打开很多句柄的场景这个 flag 能明显缩短初始化时间。第三步是执行环形归约。我用最简单的实现做示例每个 rank 把本地数据按 rank 数切成 chunk然后循环N-1次。第 k 次循环里先从上游读取 chunk和本地对应 chunk 相加再写入下游。这里的读写如果走 kernel就是类似下面这样的形式// 伪代码省略了 chunk 索引和 stream 同步 kernel_add_and_copygrid, block, 0, stream_a( local_chunk, remote_ptr_upstream, remote_ptr_downstream); kernel_reducegrid, block, 0, stream_b( local_chunk, remote_ptr_downstream);实际测试中我发现直接解引用remote_ptr在 kernel 里做读写对 PCIe 链路的压力比cudaMemcpyPeerAsync更大因为每次内存访问都可能触发一次 PCIe 事务。如果 chunk 太小事务开销占比会很大如果remote_ptr恰好命中 CPU 不可缓存的映射区还会额外增加延迟。所以这里有个很反直觉的经验在 PCIe IPC 路径上与其在 kernel 里逐元素读写远端内存不如直接用cudaMemcpyPeerAsync把整块数据拷到本地再做归约计算。虽然多了一次拷贝但因为 PCIe 的事务是整块发起的实际吞吐反而更高。4.3 调参与性能对比记录我自己搭了一套双卡 PCIe Gen3 x16 的测试环境卡是两块消费级 GPUCUDA 版本 12.4。固定数据量为 128 MB分别测了三组实现的带宽和延迟实现方式平均带宽延迟128MB说明NCCL Ring Simple约 21.3 GB/s约 6.1 ms常规配置协议自动选择NCCL Ring LL强制约 10.8 GB/s约 6.8 ms带宽被 flag 开销拖累自写 PCIe IPC MemcpyPeer约 23.1 GB/s约 5.7 ms通过 cudaMemcpyPeerAsync 传 chunk这个结果在意料之外也符合我对 PCIe IPC 的理解。自写实现之所以能比 NCCL 的 Simple 略快一点主要原因是省掉了 NCCL 里大量通用逻辑比如更细粒度的内存切分和额外的 flag 同步。但在 8 卡环境下我复测了一次情况就反过来了NCCL 在 8 卡 Ring 上表现稳定而我的自写实现因为环上每步都要做一次全量句柄寻址在 chunk 多、跳数多时延迟线性增长反而比 NCCL 慢 15% 左右。所以我的最终建议是在 2-4 卡的 PCIe 直连场景下自写 IPC 版本确实有竞争力超过 4 卡或者拓扑变复杂直接用 NCCL让它的启发式算法去选协议和 chunk 大小。调参方面我重点调了两个参数chunk 数量和 CUDA stream 数量。chunk 数量从 4 调到 16 再到 32观察到 16 个 chunk 左右是最佳点再往上归约 kernel 的启动开销就开始拖累流水线。stream 数量方面两个 stream 刚好够用三个 stream 带来的收益已经小于多一个 stream 的调度成本。这些数值未必适用所有机器但调参的方向是一致的先固定数据量扫 chunk再固定 chunk 扫 stream最后看 NCCL_PROFILE 输出的时间线确认气泡位置。5. 常见问题与排查心得5.1 PCIe 链路稳定性掉卡、降速与 AER这类基于 PCIe 的通信项目最怕的不是代码 bug而是链路本身不稳定。我在测试过程中碰到过几次cudaErrorNoDevice重启设备后又能跑用dmesg查日志发现里面有几行PCIe Bus Error: severityCorrected正是典型的 PCIe AER 错误。AERAdvanced Error Reporting是 PCIe 协议里的错误上报机制它会把链路上的可纠正/不可纠正错误记录到内核日志里。对于 GPU 这类高功耗设备最容易出现的问题是供电波动导致链路降速。lspci -vvv里能看到当前链路状态LnkSta: Speed 8GT/s (downgraded), Width x16 (downgraded)如果看到括号里出现 downgraded说明链路没有跑在最高速率上比如从 Gen3 x16 掉到了 Gen2 x16带宽直接缩水一半性能测试结果会非常离谱。我建议在任何性能测试之前先固定记录一次lspci -vvv的链路状态如果发现降速优先检查供电和 PCIe 插槽接触再考虑 BIOS 里是不是开了什么省电策略。还有一个容易忽略的点是某些主板的 PCIe 插槽是分频的插了多张卡之后会自动把 x16 降成 x8这不是故障是主板的物理限制需要提前查清楚。在程序层面如果出现 AER 可纠正错误往往不会直接导致 crash但会带来额外的重传开销导致带宽出现周期性抖动。排查时可以用perf stat观察中断频率或者打开 NCCL 的NCCL_DEBUGINFO看是否有大量重试日志。5.2 IPC 句柄生命周期与同步陷阱自写 IPC AllReduce 有一个非常隐蔽的坑句柄的创建和打开时机。cudaIpcGetMemHandle必须在显存分配之后调用而且句柄是和创建它的进程绑定的如果创建进程退出了其它进程持有的句柄并不会自动失效但再次访问时可能触发段错误或者cudaErrorIpcDeleted。另一个更常踩的坑是 multi-process 的初始化顺序。我把句柄交换的共享内存定义成了固定大小一旦某个进程因为调度延迟没能及时把句柄写进去其它进程就会读到无效数据。解决方法是加一个两阶段栅栏第一阶段写入句柄第二阶段等所有进程都确认写完后再开始cudaIpcOpenMemHandle。同步方面IPC 打开的远端内存没有统一的跨进程事件机制两个进程对同一块显存的读写顺序必须完全依赖业务逻辑来控制。我有一次把 ring 的下游发送顺序写反了结果数据错乱但程序完全不会报错因为 GPU 端的读写都是异步的。排查这种问题唯一可靠的方法是在每轮传输结束后加一个cudaDeviceSynchronize打印 checksum确认数据前后一致再继续下一轮。5.3 调试工具与现场实战技巧最后说几个我实际用下来效率最高的排查技巧都是没有写在 CUDA 文档里的经验。第一用nvtop实时看 GPU 的 D2D 带宽和 SM 占用。IPC AllReduce 跑起来之后如果发现显存带宽很高但 GPU 利用率很低说明任务很可能被传输 stall 住了要优先优化数据流的重叠反过来如果 SM 利用率高但显存带宽一般那瓶颈就在 kernel 的归约计算本身。第二用cudaCtxResetPersistingL2Cache控制 L2 cache。LL 协议和 IPC 直读远端内存都极度依赖 L2 命中率如果发现测试结果受数据访问顺序影响巨大可以尝试在不同 chunk 之间插入 L2 刷新看看变化规律。第三用NCCL_P2P_DISABLE1强制 NCCL 走主机内存路径和自写 IPC 版本做对比。这个开关非常有用它能帮你快速定位问题是出在协议层还是传输层。如果禁用 P2P 后 NCCL 反而更快那说明你的 IPC 直接映射路径可能因为缺少缓存优化或事务合并不如走主机内存中转来得划算。最后再分享一个代码层面容易忽略的习惯所有cudaIpcOpenMemHandle打开的远端指针程序退出前必须调用cudaIpcCloseMemHandle关闭否则进程退出时 CUDA context 清理会变慢而且下次启动时可能因为资源未释放而随机失败。这个错误不会在第一次运行时就出现而是稳定复现于连续跑多次实验的过程中非常迷惑人。我在实际项目里最终没有把自写 IPC AllReduce 部署到所有机器上只在两卡和四卡直连的环境里保留了它作为默认路径其余场景还是回退到 NCCL。但这段把 Ring 算法从 NCCL 源码里拆出来、又用 PCIe IPC 重新搭一遍的经历让我对通信库的认识彻底改变了。以后再碰到性能问题我不会再盲目调参而是会先看一眼拓扑和协议栈判断瓶颈到底在哪一层。这种往下钻一层的视角可能就是这次折腾最大的收获。
阅读完成 · 觉得有帮助?
咨询建站