尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

一文读懂分布式训练梯度同步:从显存到显存的 RDMA 完整旅程

一文读懂分布式训练梯度同步:从显存到显存的 RDMA 完整旅程 一块梯度从 A 机的 GPU 显存出发经过 RDMA 网卡和网络链路最终落在 B 机的 GPU 显存里——中间没有 CPU 参与数据拷贝也没有 TCP/IP 协议栈逐层打包。这是 AI Infra 领域最值得搞清楚的链路之一。今天把这块拆开数据并行训练下梯度为什么要“出门”、RDMA 到底解决了什么以及一块梯度从显存到对方显存的完整旅程适合正在做分布式训练、AI Infra 调优或者准备相关面试的朋友。这期是第 4 课的 Loop2我直接把视角放在一个训练 step 里最烧钱的环节梯度同步。跑 7B 以上模型时通信耗时经常能占到整个迭代的 30% 到 50%搞不懂这一步里的原理和坑后面做性能优化基本就是瞎调参。1. 先搞清楚梯度为什么必须“出门”1.1 数据并行训练的分工与矛盾现在训练大模型基本躲不开数据并行。做法是把一个 batch 的数据切成 N 份分给 N 个设备每个设备上放一份完整的模型副本各自做前向和反向。这样算力可以线性扩展但问题也随之而来每个设备只看到自己那一份数据算出来的梯度只反映局部样本的统计信息不能让模型参数各走各的否则训练就乱套了。所以每次反向传播结束所有设备必须把各自算出来的梯度拿来做一次全局归约得到一份所有设备都认可的“平均梯度”再各自用这份平均梯度更新本地模型。这套机制在今天已经是分布式训练的标准动作了一线工程师必须理解在数据并行的语义下梯度是所有设备共享的中间产物天然需要“出门”交换。1.2 反向传播结束梯度躺在显存里很多人会把“反向传播”和“梯度同步”混在一起其实它们是完全不同的两件事。反向传播是每个设备内部的事它负责从 loss 出发沿着计算图逐层求出 loss 对每个参数的偏导也就是梯度。反向传播本身并不解决梯度下降是否陷入局部极小的问题它只负责把梯度算出来那是优化器、学习率调度和初始化策略该去管的。拿一个 7B 参数的模型举例如果梯度用 FP32 存一份完整梯度就是 7B 乘 4 字节约等于 28GB。这 28GB 数据在反向传播结束时静静地躺在 GPU 显存里但此刻它只是“本地版本”。如果 64 张卡各算各的64 份梯度互不相通模型更新就变成 64 个分叉了。因此接下来必须把这些梯度“汇合”成一份而汇合这件事就是梯度同步也就是让梯度跨设备甚至跨机器挪动。1.3 梯度同步到底有多耗时直观感受一下梯度同步的代价。8 台机器、每台 8 卡总共 64 卡做数据并行模型还是 7B 规模。使用 ring allreduce 做梯度归约时每台机器需要发送的数据量大约接近 2 倍的单卡梯度大小也就是接近 56GB。如果节点间只有 200Gbps 网络理论极限带宽不过 25GB/s一次同步就要吃掉 2 秒以上。而一个迭代的前向加反向可能也就几秒钟通信开销占比一下子就上来了。很多团队刚开始做大模型训练时发现多机扩展效率上不去第一反应是调学习率实际上用nsys拉一下时间线就会发现GPU 一直闲着等梯度同步。这也是 RDMA 必须在 AI Infra 里登场的根本原因梯度太大了以 TCP 那套数据路径来搬根本搬不动。2. RDMA 技术拆解内核旁路与 Queue Pair2.1 TCP 为什么会卡在分布式训练里先想想如果梯度走 TCP 会是什么场景。每台机器要把 28GB 数据发出去传统 TCP socket 的数据路径是应用层把数据从显存拷到 CPU 内存然后交给内核协议栈经过系统调用、协议栈封包、内存拷贝、网卡中断到达对端后再一层层解包最后拷贝到应用缓冲。问题不仅仅是慢而是 CPU 成了瓶颈。几十 GB 的数据全要过 CPU 走一遍CPU 会被中断和内存拷贝打满留给计算的时间就没了。而且小包场景下 TCP 更加吃力七层协议处理、ACK 确认、拥塞控制全在 CPU 里跑100Gbps 网络下想单靠 CPU 跑满基本是做梦。分布式训练的梯度同步是大流量、低延迟、高并发TCP 这条路径从根上就不适合。RDMA 的核心思路是内核旁路把数据搬运和协议处理全部卸载到网卡硬件上应用直接和网卡交换数据数据从用户态内存甚至 GPU 显存通过 DMA 直达网卡对端网卡通过 DMA 直接把数据写进目标内存整个过程 CPU 只负责发起和确认。2.2 RDMA 的三种实现IB、RoCE 与 iWARPRDMA 不是一个单一协议而是三种主要实现先分清楚再谈链路。InfiniBandIB专用网络体系网卡、交换机、线缆整套都是为 RDMA 设计的硬件成本高但性能、可靠性、生态成熟度都是最好的。大模型训练集群里IB 是首选。RoCE基于以太网的 RDMA 实现v2 版本把 IB 报文封装进 UDP/IP 包可以跑在普通以太网交换机上成本低很多。但以太网天生有丢包问题RoCE 对网络质量极其敏感通常要开启 PFC、ECN 这类机制把网络改造成“无损网络”否则性能会断崖式下跌。iWARP基于 TCP 实现 RDMA兼容性最好但 TCP 协议栈的处理逻辑拖累了性能实际高性能场景用得很少。做 AI Infra 时最常见的组合就是 IB 或者 RoCE。关键词里有人搜“rdma qp 是什么”这说明很多人已经遇到了实际配置问题QP 恰恰是 RDMA 里最核心的概念下一节单独展开。2.3 QP 是什么怎么理解它QP 的全称是 Queue Pair中文叫队列对由一对队列组成一个发送队列SQ和一个接收队列RQ。RDMA 通信前通信双方需要各自创建一个或多个 QP然后交换 QP 编号、LID、GID 这些连接信息建立起一条逻辑上的“可靠连接”。打个比方QP 就像一条快递专线。你给这条专线起了名字、编了号对方也知道这条专线的编号两边约好在哪收件。发送方不断往这条专线投递包裹发送请求接收方提前准备好空货架接收缓冲区快递员也就是网卡硬件负责把包裹送过去。每次投递和收货都会产生一张回执单这在 RDMA 里叫 CQE。软件只需要盯住回执单确认哪个包裹送达了收尾。QP 的传输模式里最常用的是 RC可靠连接提供有序、不丢包、有确认和重传的端到端通信UD不可靠数据报支持一对多但单包长度受限用得少。在分布式训练场景几乎全是 RC因为梯度数据量大不容许任何丢失。单张网卡支持的 QP 数量有上限而训练框架通常需要同时开多个 QP 并行传输所以 QP 的规划和配置直接决定了能否把网络带宽打满。2.4 与 QP 配套的 MR、WQE、CQE 等抽象光有 QP 还不够RDMA 还有几个绕不开的概念贴一下我的理解。MRMemory Region叫内存区域注册。因为 RDMA 网卡要用 DMA 直接读写内存必须提前把这块内存的物理地址锁定并告诉网卡网卡才能绕过 CPU 去访问。注册时会生成一个 rkey相当于这把内存的钥匙。在分布式训练里每轮迭代都要发梯度所以显存缓冲区的 MR 都是提前注册好、反复复用的绝对不会每轮重新注册那开销太大。WQEWork Queue Element是投递给 QP 的工作请求。发送方调用post_send把一个 WQE 放进发送队列里面描述了本地内存地址、长度、目标 QP 编号等接收方调用post_recv把接收缓冲区放进接收队列。CQECompletion Queue Element是完成事件网卡干完活之后往完成队列写一条记录应用去轮询 CQ 来确认。把这些概念串起来一次 RDMA 发送就是注册 MR → 创建 QP → 交换连接信息 → 投递 WQE → 轮询 CQE。理解了这条主线再看直方图、报错日志就不会一头雾水。3. 一块梯度的完整旅程从显存到另一块显存3.1 出发前梯度切片与 QP 连接准备在训练框架里梯度搬运并不是把 28GB 当成一个整体丢到网络上。NCCL 这类通信库会把梯度 buffer 拆成很多 chunk每个 chunk 几十 KB 到几 MB分散到多个 QP 上并行发送。为什么要拆因为一个 QP 是一条串行流水线单 QP 很难打满高速网卡多发几个 QP 并行跑可以更接近硬件极限另外小 chunk 还能让归约操作更早开始避免接收端一直空等。在发送之前通信库还会确认连接状态。多机训练时每台机器上的每个 GPU 都有对应的通信 RankNCCL 启动时会建立一个全局通信域Rank 之间通过 TCP 或已有的 RDMA 连接交换 QP 信息相当于先修好路再让梯度出门。3.2 在途网卡如何把数据“搬”过去一旦 WQE 被投递到 QP网卡硬件就会接管。它通过 DMA 引擎直接读取显存或内存里的数据组装成 RDMA 报文通过物理链路发出去。如果是 RoCEv2报文结构大致是以太网头 IP 头 UDP 头目标端口 4791 IB BTH 数据如果是 IB则走 IB 自己的报文帧。接收端网卡收到报文后解析头、校验完整性再把数据通过 DMA 写到预先注册好的接收缓冲区。整个过程中CPU 只负责最初投递 WQE 和最终读取 CQE搬运本身完全不经过 CPU。这也是“零拷贝”的意义所在数据不落在用户态缓冲、不经过内核协议栈、不经过 CPU 寄存器搬运相当于快递员直接带着包裹从发货仓库开到收货仓库中间没有中转分拣。3.3 到达后接收、归约与参数更新数据到达对端的显存后真正的“归约”才开始。以 ring allreduce 为例所有设备被组织成一个逻辑环。第一步 reduce-scatter每个设备把自己的梯度切成 N 块和相邻设备交换其中一块累加后继续传给下一个设备直到每个设备持有某一小块的全量累加结果。第二步 allgather再把累加好的小块依次传递出去最终每个设备都拿到完整的全局梯度。这个过程中梯度在机器之间“搬”了不止一次而是每个人都在收、都在发链路始终是满的。这也是 ring 算法设计上的巧妙之处虽然总通信量变多了但没有任何一条链路闲置整体带宽利用率很高。当所有设备拿到全局梯度optimizer 才能拿着它去更新参数一个训练循环的 Loop2 到这里才算真正闭合。3.4 中间少走的一道弯GPUDirect RDMA前面说的“网卡通过 DMA 直接读显存”有条件网卡必须支持 GPUDirect RDMA。没有这个能力时梯度得先从显存拷贝到 CPU 内存网卡再从 CPU 内存发走接收方向反过来。多了两次 PCIe 上的内存拷贝延迟和 CPU 占用都会显著上升。开启 GPUDirect RDMA 之后网卡可以直接把 GPU 显存地址当作 DMA 地址来访问数据路径变成“显存 → 网卡 → 网络 → 对端网卡 → 对端显存”全程没有 CPU 内存参与。实测下来开启前后小消息延迟能差几倍大消息聚合带宽也更高。想确认当前环境是否生效可以用nvidia-smi topo -m看 GPU 和网卡的拓扑检查它们是否挂在同一个 PCIe switch 下也可以看 NCCL 日志里是否启用了 GPU Direct 相关路径。如果 GPU 和网卡跨 NUMA、跨 PCIe 根节点中间要绕远路性能会打折扣。4. 训练框架视角NCCL、梯度累积与梯度裁剪4.1 NCCL 在训练循环里扮演的角色NCCLNVIDIA Collective Communications Library是现在多卡多机训练的事实标准通信库PyTorch 的torch.distributed底层调用的就是它。它做的事情远不止把数据发出去还包括拓扑探测、通信算法选择、缓冲区管理、与 NVLink/RDMA 的协同。对开发者来说梯度同步对应一次allreduce调用但 NCCL 内部会根据拓扑自动选择 ring 或 tree 算法。如果机内有 NVLink、机间是 IBNCCL 会先在节点内部做一次归约把跨机流量缩到最小再走 RDMA 做节点间归约最后在节点内广播回来。这个优化的收益非常大单机 8 卡有 NVLink 的情况下跨机通信量可以降到原来的 1/8省下的全是宝贵的 IB 带宽。4.2 梯度累积如何影响通信梯度累积是很多训练脚本里的常见参数它和通信的关系紧密到必须单独说。思路很简单连续跑几个 micro-batch 的反向传播梯度先本地累加攒够一定步数后再做一次 allreduce通信频率降下来了网络压力也小了。但代价也很实际参数更新变得滞后模型看到的等效 batch size 变大了训练收敛特性和小 batch 不一样需要配合学习率调整。另外累积时梯度是 FP32 累加还是 FP16 累加、loss scale 怎么处理都容易踩坑。如果你的集群网络带宽有限又不想改模型并行策略梯度累积是一个简单的“减通信”手段但它不是免费的午餐调参时要盯着收敛曲线看不能只看吞吐。4.3 梯度裁剪、梯度同步与全局范数有人会搜“梯度裁剪有哪些方法”这里也说清楚它和通信的关系。常见的梯度裁剪分为按值裁剪clip_by_value和按全局范数裁剪clip_grad_norm。按值裁剪好实现但会破坏梯度各维度的相对比例按全局范数裁剪更常用做法是算所有梯度向量的 L2 范数如果超过阈值就整体等比缩放。问题在于数据并行下“全局范数”必须是所有 Rank 的梯度合起来算的。也就是说裁剪前需要先对“梯度的范数”做一次跨设备归约或者更高效的做法是在做梯度 allreduce 的同时把范数归约也合并进去。我之前在实践里做过一次优化把这两个通信合并后整个迭代的通信量直接少了近一半。梯度消失和梯度爆炸虽然本质上是训练稳定性的问题但如果你用 FP16 训练梯度数值范围控制不好通信过程里就可能出现溢出把没用对归约算法的问题进一步放大。5. 实操观察与问题排查实录5.1 自测网络与定位通信瓶颈拿到一个新的多机训练环境我一般不看理论参数先跑一遍实测脚本确认底子什么样。第一步用ibv_devinfo看 IB 网卡状态端口是否 Active、速率是否是预期值第二步用 perftest 工具集里的ib_write_bw和ib_send_bw在机器之间直接打带宽。ib_write_bw -a -d mlx5_0这类命令很简单但能立刻暴露问题光模块没插好、换了几台交换机之后路由不对、RoCE 流控没开都会在这里现出原形。第三步把 NCCL_DEBUGINFO 打开跑一个小规模 allreduce 测试看日志里是否走NET/IB而不是NET/Socket。如果看到 Socket说明 RDMA 没通框架悄悄回落到 TCP性能一定上不去。5.2 常见问题速查表现象常见原因排查思路NCCL 报net poll failed、训练卡死IB/RoCE 链路不稳定QP 超时用 perftest 跑持续带宽和延迟测试检查光模块和交换机端口日志显示走了NET/SocketRDMA 设备不可见或驱动问题ibv_devinfo检查网卡状态确认驱动和固件版本多机吞吐远低于预期RoCE 网络有丢包或 PFC 配置不当在交换机上查端口丢包计数确认无损网络参数开启GPU 和网卡通信延迟高拓扑跨 NUMAGPUDirect RDMA 未生效nvidia-smi topo -m看拓扑调整设备亲和性allreduce 小消息速度快、大消息带宽低QP 数量不足调大NCCL_IB_QPS_PER_CONNECTION观察变化训练时好时坏、偶发超时未开启 ECN 或交换机缓冲区不足开启 RoCE 的 ECN/PFC检查交换机的 buffer 模型排查通信问题最忌一上来就堆环境变量先把物理链路和驱动基础打牢再谈调参。5.3 几个 QP 和调优相关的经验最后分享几个我实际跑训练攒下来的调优经验。QP 数量不是越大越好虽然更多 QP 能提升并发但每个 QP 也有内存和硬件资源开销开太多反而增加管理负担。我习惯从默认值开始用 allreduce 基准测试逐步往上调找到带宽不再增长的拐点就停。NCCL_BUFFSIZE这个参数控制通信缓冲区大小和模型尺寸、消息大小要匹配设太大会多吃显存设太小又会在内存注册和 DMA 分片上浪费性能。RoCE 环境下确保交换机开了无损网络相关机制否则一切参数都是白调IB 环境下反而简单专网专设备基本不用管拥塞控制。另外强烈建议养成长跑之前先做压力测试的习惯至少跑十几分钟的 allreduce 基准观察带宽曲线是否有周期性坍塌。很多“训练卡死”其实是网络在半小时后才开始丢包前几分钟正常掩盖了问题。数据要连续记录不要只跑一次就下结论。我个人在排查分布式训练问题时最常用的一招是把数据路径画出来从显存到网卡、从网卡到交换机、从交换机到对端网卡、再到对端显存每一个环节标清楚数据是“从哪里来、经没经过 CPU”。很多时候问题不是出在某一个复杂参数上而是某个环节偷偷多了一次拷贝或者本来该走 GPU Direct 却走了 CPU 中转。把这层想通了梳理 Loop2 里梯度搬运的细节会比记住一堆环境变量有用得多。
返回列表