ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

分布式训练梯度跨机同步:RDMA、QP与Ring AllReduce原理解析

分布式训练梯度跨机同步:RDMA、QP与Ring AllReduce原理解析 一次200B参数模型的训练迭代跑完反向传播每一张GPU上都会落下一份梯度张量。如果训练是中单机单卡直接更新参数就行但只要上了多机多卡问题就立刻变得棘手——这几份梯度要怎么合并成一份并广播回去中间如果还跨了机器那本质上就是在回答一个问题一块梯度到底是怎么通过 RDMA 从这台机器的显存里搬到另一台机器的显存里的这个系列走到第4课前面已经聊过分布式训练的基本并行策略、通信原语和单机内的数据传输方式。今天这篇内容就是把镜头推到最底层专门把“梯度跨机搬运”这条链路上的每一站都拆开看一遍包括梯度同步在训练循环里的位置、RDMA 为什么是标配、QP/MR/CQ 这些抽象到底在干什么、Ring AllReduce 为什么能把带宽用满以及实际操作中那些跑一次就踩一次的性能坑。适合正准备深入 AI Infra、分布式训练框架或者被 NCCL 黑盒性能问题折腾过的人读。1. 先搞懂Loop2在训练循环里的位置1.1 一次训练迭代里发生了什么分布式训练一个 step 如果拆开看其实不是一个“前向-反向-更新”的简单循环。尤其在用梯度累积gradient accumulation做大等效 batch 的时候真实的流程会更接近下面这段伪代码for step in range(total_steps): # 计算循环Loop1 for micro_batch in range(gradient_accumulation_steps): loss forward(micro_batch) loss.backward() # 通信循环Loop2 allreduce(gradients) optimizer.step()Loop1 是每张卡自己闷头算前向、反向把梯度算出来放在显存 buffer 里。Loop2 则是所有设备开始把各自手里的梯度做一次全局同步同步完之后每个 rank 拿到的是所有 rank 梯度的平均值然后才能安全地调用 optimizer.step() 更新参数。在分布式训练的语境下我习惯把 Loop2 理解成“梯度通信与同步循环”它和 Loop1 是天然紧耦合的。Loop1 决定的是算得有多快Loop2 决定的是同步有多快。当模型规模上去之后同步效率往往比计算效率更容易成为瓶颈。1.2 梯度同步需要解决的核心问题梯度同步的本质是把 N 份分布在不同机器上的梯度张量先归约reduction比如求和或求平均再广播broadcast回所有节点让每个 rank 拿到的参数更新量完全一致。这里有几个很要命的约束条件第一梯度数据量极大。一个 70B 参数的模型光梯度就是 70B 个浮点数FP32 下约 280GB。一次反向传播结束集群内部要搬运的量级就是这么恐怖。第二同步等待不可接受。Loop2 必须等到 Loop1 结束下一个 step 又必须等到 Loop2 结束。通信拖得越久GPU 空转时间越长训练吞吐就越难看。第三任何一台机器都不能掉队。同步是全局性的一台卡住整个训练 job 就 hang 住。这也是为什么很多分布式训练故障最终都指向网络层。所以梯度跨机搬运的方案选型核心就一个目标用尽量短的时间把尽量大的数据量可靠地从每台机器搬到所有机器。这几乎是为 RDMA 量身定做的问题。1.3 为什么不是把梯度打包发过去那么简单很多人第一次接触分布式训练时容易产生一个错觉梯度同步嘛不就是每台机器把梯度通过 TCP 发到一个中心节点求个平均再发回去8 台机器各发 1GB加起来 8GB听起来也不离谱。但真实情况是TCP 的收发路径上有太多环节GPU 显存里的数据要先复制到 CPU 内存CPU 再通过内核协议栈把数据拷贝到 socket 缓冲区网卡发出去接收端反过来还要再经过内核、再复制到用户态、再拷贝到显存。数据在内存里被反复搬来搬去CPU 每一轮都得全程参与。这还没算 TCP 本身的重传、拥塞控制带来的延迟抖动。这套路径在小规模、小数据量的场景里没问题放到 400Gbps 网卡、TB 级梯度同步的集群里CPU 早就被打满网络却喂不饱。这也是为什么现代 AI Infra 里梯度跨机几乎默认走 RDMA——它走的是完全另一套数据传输路径。2. RDMA凭什么成为梯度搬运的“高速公路”2.1 RDMA 的三大技术支柱RDMA 全称 Remote Direct Memory Access直译过来是“远程直接内存访问”。它之所以能在 AI 训练集群里站稳脚跟靠的是三条硬功夫内核旁路、零拷贝、CPU 卸载。内核旁路Kernel Bypass传统网络通信要多次经过内核协议栈数据包从应用态到内核态再出去回来还要再走一遍。RDMA 让应用在用户态直接和网卡交换数据协议栈的大部分处理都在硬件或专用固件里完成省掉了内核参与带来的上下文切换和拷贝开销。零拷贝Zero Copy传统 TCP 路径里数据至少要经历“应用缓冲 → 内核缓冲 → 网卡”的复制链条。RDMA 网卡可以直接通过 DMA 读取用户态注册好的内存区域数据从源端内存到目的端内存之间全程没有一个 CPU 参与的数据拷贝动作。CPU 卸载Offload数据分片、传输确认、重传、乱序重组、完成通知这些脏活累活全部由网卡硬件RNIC完成。CPU 只需要在最初把工作请求Work RequestWR扔给网卡然后该干嘛干嘛去。这三件事叠加起来的效果非常像修了一条从 A 机房直接通到 B 机房的专用地下管道。传统 TCP 是快递员带着包裹走公共道路每个站点都要验货签收重新分拣RDMA 是管道两端直接对接货物进了管道口另一头直接入库中间不需要任何人搬动。2.2 InfiniBand 与 RoCE两种主流实现RDMA 并不是特指某一种网络硬件而是一类技术的总称。现在工业界用得最多的主要是两种类别网络载体优势劣势InfiniBandIB专用 IB 网络有独立交换机原生 RDMA 设计延迟最低隔离性好生态成熟硬件成本高需要单独组网RoCEv2以太网 RDMA 协议栈复用现有以太网设备成本低运维体系通用依赖无损网络配置PFC/ECN调优要求高iWARPTCP 网络能用普通网卡跑 RDMA性能受 TCP 栈限制实际部署占比小NVIDIA 的 DGX 系列机器通常标配 InfiniBand很多自建 AI 集群为了平衡成本和性能会选 RoCEv2。无论走哪种底层用的都是同一套 RDMA 编程模型所以“梯度怎么通过 RDMA 搬过去”的原理是通用的。2.3 单边操作与双边操作的区别RDMA 有两种经典的通信模式理解这个对后续看集合通信实现非常有帮助。双边操作Two-sidedSEND/RECV发送方调用 SEND接收方必须提前调用 RECV 并准备好接收缓冲区两端 CPU 都要参与投递工作请求。类比是“打电话”必须双方都在线才能通话。这种模式适合控制信息、握手消息等小数据量交互。单边操作One-sidedREAD/WRITE发送方直接用 RDMA WRITE 把数据写到远端已经注册好的内存地址里远端 CPU 完全无感不需要参与RDMA READ 则是直接去远端内存读取数据。类比是“往别人家的仓库里直接送货”只要你有钥匙远端内存地址和 RKey货放进去就可以了。梯度同步这种大流量传输用的正是单边 WRITE 为主。这里有个关键点RDMA WRITE 是“哑”的操作它只负责把数据从一个内存区域搬到另一个内存区域并不通知对端“数据已到达”。所以实际工程里大块梯度用 RDMA WRITE 传输传完之后往往再补一个很小的显式通知或者借助集合通信框架里自带的同步逻辑来收尾。2.4 为什么跨机必须靠网卡而不是只靠 NVLink单机内 GPU 之间传数据可以用 NVLink带宽高、延迟低但 NVLink 是板级或机箱级互联跨机柜就无能为力了。多机训练必须通过网络把数据搬过去。在机间这个尺度上IB 或 RoCE 提供的 400Gbps 带宽就是最主流的选项。所以整个数据路径其实有个明显的分界单机内尽量用 NVLink/PCIe跨机一律走上行网卡走 RDMA。明白了这条主线再去看各种集合通信框架的拓扑优化就知道它们本质上是在解决“哪段路用什么交通工具”的问题。3. 一块梯度数据包从显存到对端显存的全链路拆解3.1 数据流动路线图把视角收回到单块梯度张量。假设它此刻已经在 GPU 显存里随着反向传播的结束“就绪”了。从它开始跨机到最终落进对端显存完整路径大致是这样的GPU显存里的梯度张量 ↓ DMAGPUDirect RDMA 或经host staging buffer 用户态内存/GPU显存已注册为MR锁页 ↓ 网卡硬件DMA读取 发送端RNIC ↓ 网络传输IB/RoCE 接收端RNIC ↓ 硬件直接写入远端MR地址 对端用户态内存/GPU显存 ↓ 集合通信框架做归约 对端最终梯度buffer每一步拆开来看最有技术含量的其实在两端的内存注册和网卡处理逻辑上。3.2 必须理解的关键实体QP、MR、CQRDMA 编程模型里有几个绕不开的概念把它们搞明白看任何 RDMA 代码都不会再犯怵。QPQueue Pair队列对每个 RDMA 连接由一对队列组成一个是发送队列SQ一个是接收队列RQ。发数据时把工作请求 WR 扔进 SQ网卡按顺序处理。QP 可以理解成两台机器之间的一条逻辑连接。梯度同步这种大流量通常一个连接配多个 QP 来提升并发度。这也就是热搜词里“rdma qp是什么”的答案。MRMemory Region内存区域RDMA 要操作的内存必须先注册成 MR注册时要把物理页锁定pin 住并记录虚拟地址到物理地址的映射。注册完成后会得到两个关键值lkey本地访问 key和 rkey远端访问 key。远端拿到 rkey才允许对它发起 RDMA WRITE/READ。你可以把 rkey 想象成一把只针对这块内存仓库的钥匙。CQCompletion Queue完成队列网卡每处理完一批 WR会往 CQ 里扔一个完成事件。应用程序通过轮询 CQ 知道数据是否已经发送完成或接收完成。CQ 是异步程序的“收件箱”。3.3 建立一条 RDMA 连接要做什么假设 rank0 要把梯度写到 rank1 的内存里用极简的流程描述需要走这几步两端都打开 RDMA 设备申请保护域Protection DomainPD。PD 是资源的归属容器。两端各自创建 QP并分配好 send queue 和 receive queue 的深度。两端各自注册一块接收缓冲区为 MR。交换“连接信息”包括 QP 编号、端口 LID/GID、MR 的远端地址和 rkey。这一步通常需要一个旁路通道比如 TCP 或共享文件来传递。QP 状态移到 RTSReady To Send之后连接就真正可用。发送端构建一个 RDMA WRITE 的 WR指定远端地址和远端 rkeypost 到 SQ。网卡自动从本地 MR 地址读数据并写入远端 MR 地址。完成之后 CQ 收到完成通知。整个过程里第 6 步之后发送端 CPU 就不再参与数据搬运了接收端 CPU 从头到尾可能根本不知道数据已经到了。3.4 发送时机梯度就绪后怎么触发通信梯度什么时候“就绪”这本身是个需要精确控制的问题。在 GPU 编程里反向传播结束意味着显存里的梯度 buffer 已经写完但这个写完动作对 CPU 来说可能是异步的。所以 NCCL 或者自研框架在 AllReduce 之前通常要确保相应的 CUDA stream 已经完成写操作或者说在同一个 stream 上排队保证通信操作不会读到半新不旧的梯度。我在早期调试自研通信模块时就栽过这个跟头AllReduce 总是快那么一点点导致同步结果偶尔对偶尔错而且出错概率极低极难复现。后来查到底就是“计算完成”和“通信开始”之间缺了一个 stream 同步。在 Pytorch DDP 里这一切都被封装好了但如果你自己写集合通信这一步千万不能省。4. 从单块梯度到集群梯度AllReduce与Ring算法4.1 为什么单块搬运不够必须要集合通信如果训练过程里只有两块卡互相传梯度直接写一个点对点 RDMA WRITE 就能解决。但真实的分布式训练面对的是 N 个 rank 同时手握一份梯度而且最终每一份都要变成全局平均梯度这个操作在集合通信里叫 AllReduce。最笨的 AllReduce 实现是“中心汇聚”所有 rank 把梯度发给 rank0rank0 做归约后再广播给所有人。这个方案逻辑简单但有两个致命问题一是 rank0 要收发 N-1 份全量数据带宽瓶颈死死卡在单点二是只要 rank0 网络慢一点整个集群都得等它。所以工程上不会用它来做大规模梯度同步。4.2 Ring AllReduce 的设计思路Ring AllReduce 的聪明之处在于它把每个 rank 的梯度先切成 N 份N 是所有参与通信的 rank 数然后让所有 rank 组成一个逻辑环。整个过程分成两个阶段第一阶段Reduce-Scatter归约-分散。环上每个节点把本地梯度切成 N 份然后按照环的方向每轮把自己手上的第 i 份发给下一个节点同时从上一个节点接收对应份的数据并做归约。经过 N-1 轮之后每个节点手上有 1 份“全局归约后的部分结果”但这 N 份结果分散在 N 个不同节点上。第二阶段All-Gather全收集。每个节点再把手里那份全局归约结果沿着环继续转圈每轮发给下一个节点同时接收上一个节点的数据。经过 N-1 轮之后每个节点都收到了完整 N 份全局归约结果拼接起来就是全局梯度。这样做的好处是每个节点的收发带宽始终被数据填满没有任何一个节点会成为单点瓶颈。而且从总通信量来看它非常接近理论下界。4.3 通信量计算Ring 为什么带宽最优把数学拆开看。假设参与通信的节点数是 N每个节点梯度的总数据量是 D。Ring AllReduce 中每个节点在 reduce-scatter 阶段要发出 N-1 份大小为 D/N 的数据块所以发送量是 (N-1) × D/Nall-gather 阶段同样还要发送 (N-1) × D/N。合计每个节点的发送量就是总发送量 2 × (N-1) / N × D以 8 个节点、每个节点梯度 1GB 为例每个节点实际需要发送的数据量是 2 × 7/8 × 1GB 1.75GB。而“中心汇聚”方案呢rank0 要接收 7GB再广播 7GB单点总流量高达 14GB。对比一下就知道 Ring 的优势有多明显。这也是为什么在带宽受限的跨机场景里Ring AllReduce 会成为 NCCL 的默认选择之一。它的精髓在于把“单点瓶颈”变成了“流水线式的全员协作”让每张网卡都能以接近满速的状态持续工作。4.4 Tree AllReduce 与多算法取舍Ring 不是唯一选择。在节点数量较少、单节点带宽很高的场景里Tree AllReduce 也能有不错的表现。它的思路是构建一个树形结构叶子节点向上归约根节点拿到完全结果后向下广播。树形结构的好处是归约过程可以并行坏处是根节点的出向带宽仍然可能是瓶颈而且树深度增加时同步等待变长。NCCL 内部其实会根据拓扑自动选择算法。单机内多用 NVLink 直连跨机时在 ring 和 tree 之间做启发式选择还会根据带宽、延迟、报文大小做 tuning。对我们做 AI Infra 的人来说知道这些算法原理的价值在于当一个问题出现时你能判断它到底出在“算法选型”还是“底层网络传输”上。5. 实操环节写一份“梯度搬运”最小示例5.1 用 RDMA verbs 直接发送梯度虽然生产环境很少直接对着 verbs API 写代码但用一段极简伪代码理解发送流程比看任何文档都直观。下面这段只展示了发送端最关键的操作省略了大量错误处理和建连细节// 伪代码用 RDMA WRITE 把本地缓冲区里的梯度写到远端 MR 地址 struct ibv_context *ctx ibv_open_device(ibv_get_device_list(NULL)); // 打开设备 struct ibv_pd *pd ibv_alloc_pd(ctx); // 分配保护域 struct ibv_mr *mr ibv_reg_mr(pd, local_grad, size, IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); // 关键注册内存区域锁定物理页得到 lkey 和 rkey // 假设 remote_mr_addr 和 remote_rkey 已经通过配置/握手拿到 struct ibv_qp *qp create_qp_with_req(pd); // 创建QP exchange_and_setup(qp, remote_info); // 交换QP状态进入RTS post_recv(qp, recv_buf, size); // 如果是SEND模式需要提前post接收 struct ibv_sge sge { .addr (uint64_t)local_grad, .length size, .lkey mr-lkey, }; struct ibv_send_wr wr { .opcode IBV_WR_RDMA_WRITE, // 单边写 .send_flags IBV_SEND_SIGNALED, .sg_list sge, .num_sge 1, .wr.rdma.remote_addr remote_mr_addr, // 远端MR地址 .wr.rdma.rkey remote_rkey, // 远端钥匙 }; ibv_post_send(qp, wr, bad_wr); // 投递之后CPU可以立刻走人 poll_cq(cq, timeout); // 等待完成通知这段代码的核心动作就是把“我要把一个内存区域的数据写到另一个机器的指定地址”这件事告诉网卡然后网卡自己干完所有活。看起来只有几行但背后涉及的内存注册、QP 状态机、连接信息交换、错误重试每一项都能写一篇长贴。5.2 生产环境真正用到的NCCL 一行搞定直接写 verbs 代码解决梯度同步不现实因为在多机多卡场景下要处理的问题太多了拓扑感知、通道划分、并行度、pipeline 重叠、错误恢复、多 QP 负载均衡……这些都是框架干的事。生产环境里梯度 AllReduce 通常就是一行 NCCL 调用ncclAllReduce(sendbuff, recvbuff, count, ncclFloat, ncclSum, comm, stream);这一行的背后NCCL 会负责把每个 rank 的梯度切块根据拓扑在 ring/tree 之间选算法在底层创建多个 RDMA QP 或多个网络 channel用 GPUDirect RDMA 直接把显存里的梯度交给网卡算完后再把结果广播回所有 rank。对于绝大多数用 PyTorch 的同学这一行也被封装进了 DDP 的all_reduce调用里所以日常开发时你可能感受不到它的存在。但一旦遇到 NCCL timeout、带宽跑不满、训练诡异的 hang 住就得回到前面这些底层概念里找答案。5.3 梯度累积一个降低同步频率的有效手段梯度累积和 Loop2 的关系非常直接你累积多少个 micro-batch 的梯度就决定你多久做一次 AllReduce。如果每个 micro-batch 都做一次同步通信次数会非常密集通信延迟可能比通信量本身更让人头疼如果累积 K 个 micro-batch 的梯度再做一次 AllReduce通信次数立刻降为原来的 1/K。注意梯度累积并不减少总通信数据量它减少的是“同步次数”和“同步次数带来的等待开销”。尤其在网络 RTT 较高的跨机训练场景把很多小步通信合并成一大步通信吞吐提升是肉眼可见的。但梯度累积要小心两个点一是累积时要用一个独立梯度 buffer不要覆盖当前 micro-batch 的梯度二是更新前要除以累积步数 K或者调整学习率否则等效 batch size 变大后收敛行为会漂移。5.4 容易忽略的隐藏坑内存注册与锁页RDMA 的数据搬运并不是从任意内存地址都可以直接开始。前面提到RDMA 操作的对象必须注册为 MR注册过程会锁定物理页面防止操作系统把这些页 swap 到磁盘上。如果内存没有锁页网卡的 DMA 操作就可能读到过期的数据甚至导致系统异常。另一个坑是注册内存有开销。频繁注册/注销 MR 会吃掉不少 CPU 周期。所以框架级的实现通常会做 MR 池化提前注册好一批大 buffer反复复用。NCCL 在初始化时会为每个 rank 预留足够的内存池而不是每次 AllReduce 都去注册新内存。这个细节在我们自己写通信库时非常关键。6. 实际部署中的常见故障与排查技巧6.1 带宽跑不满先查 MTU 和拥塞控制RDMA 网络的性能问题里带宽跑不满是最常见的。很多新集群上来一测吞吐和纸面带宽差一大截排查第一个动作就是看 MTU。InfiniBand 通常用 4096 的 MTURoCEv2 一般也建议设成 4096。如果协议栈把大包拆成小包报头开销占比会膨胀有效吞吐自然下降。另外RoCE 因为跑在以太网上对丢包极其敏感需要交换机侧的 PFC优先流控制和网卡侧/交换机侧的 ECN显式拥塞通知配合才能保证无损网络环境。没有 PFCRoCE 一遇到拥塞就丢包RDMA 的重传机制会让性能直接雪崩。注意很多人的第一反应是“调大 QP 深度、调大 buffer”但先确认网络丢包率和 MTU 往往见效更快。6.2 训练 hang 住多半是 QP/连接状态异常分布式训练最常见的故障之一就是 NCCL timeout。从网络层看最常见的元凶包括QP 状态异常比如对端没准备好就发数据接收端 MR 的 rkey 不匹配交换机或网卡固件问题导致静默丢包连接信息交换时用了不稳定的旁路通道遇到 hang 住先不要急着怀疑算法。打开 NCCL 调试日志NCCL_DEBUGINFO再看dmesg里有没有 RDMA 相关的报错然后用ibstat或ibstatus确认端口状态是否 ACTIVE。如果底层链路是通的再把怀疑范围缩小到通信模式或者 buffer 管理上。6.3 网卡和 GPU 没对齐NUMA 亲和性在高性能服务器上GPU、PCIe 交换机、网卡都挂在特定的 NUMA 节点上。如果 CPU 分配线程和内存分配没有绑到对应 NUMA 节点跨 NUMA 访问带来的额外延迟和带宽损耗在 RDMA 场景里会被放大好几倍。排查时用nvidia-smi topo -m看 GPU 和网卡的拓扑关系用numactl --hardware看节点分布尽量保证通信线程和接收 buffer 的内存分配都落在网卡所在 NUMA 节点上。这个优化成本最低收益却非常可观。6.4 常用排障工具清单场景命令/工具作用验证 RDMA 读写带宽ib_write_bw/ib_read_bw直连带宽测试排除上层干扰验证延迟ib_write_lat/ib_read_lat查看端到端延迟是否正常查看端口状态ibstatus/ibstat确认链路是否 ACTIVE查看网卡丢包/错误ethtool -S看 RLROE/ICRC 错误之类计数器NCCL 调试NCCL_DEBUGINFO / NCCL_DEBUGTRACE看初始化信息、选用的通道和协议CUDA 通信分析nsys profile看通信时间线定位是否计算通信重叠不足有一次我在调一个 64 卡集群吞吐始终上不去最后扒了半天发现是某台机器的 RoCE 网卡丢包计数一直在涨而交换机端口上 PFC 没有开齐。用ethtool -S一眼就看出来了但之前不看计数器的话可能调 QP 深度、调 buffer 大小调几天都找不到原因。6.5 调试中的三条经验第一永远先验证物理链路再验证上层逻辑。带宽不对劲先用ib_write_bw把网卡到网卡的裸性能跑出来链路不通就谈不上调框架。第二日志能救命但别默认开 VERBOSE。NCCL_DEBUGINFO 在平时建议关闭只在排障时打开。TRACE 日志量巨大开一次能产生几十 GB 文本别在生产环境常开。第三升级网卡固件和 OFED 驱动要慎重。RDMA 的 bug 常在驱动/firmware 层面升级前一定要在测试集群上跑一轮长稳测试否则可能引入比旧版本更奇怪的网络异常。最后分享一点个人体会。刚接触 RDMA 的时候我去读那些协议文档总觉得抽象难啃后来有一次在集群上做性能对比随手记录了一次 RDMA WRITE 的延迟只有 1.2 微秒左右而同样尺寸的数据走 TCP 却要 30 多微秒。那一刻我突然理解了为什么整个 AI Infra 的梯度同步栈最终都会一层层沉淀到 RDMA 上。这不是某个团队拍脑袋的选择而是性能指标摆在那里算力再强梯度搬不过去一切都白搭。如果你也在做分布式训练或者自研通信库建议花点时间把 QP、MR、CQ、Ring AllReduce 这些基本功吃透遇到问题的时候你会感谢当年多看的这几段底层原理。
返回列表