简介:PD 分离推理的加速大招,百度智能云网络基础设施和通信组件的优化实践
为了适应 PD 分离式推理部署架构,百度智能云从物理网络层面的「4us 端到端低时延」HPN 集群建设,到网络流量层面的设备配置和管理,再到通信组件和算子层面的优化,显著提升了上层推理服务的整体性能。
百度智能云在大规模 PD 分离式推理基础设施优化的实践中,充分展现了网络基础设施、通信组件与上层业务特征深度融合的重要性。
传统的推理服务均是集中式,大多是单机部署。即使是多机部署,机器规模也非常小,对网络的带宽和时延需求都不大。当前大规模 PD 分离式推理系统来说,对网络通信的需求则发生了变化:
为了提升大规模 PD 分离式推理系统的效率,百度智能云针对性地优化了网络基础设施和通信组件:
下面我们逐一介绍百度智能云在以上几个层面的优化实践。
百度智能云在训练场景下的 HPN 网络架构设计已经有着丰富的经验,AIPod 使用多导轨网络架构,GPU 服务器配有 8 张网卡,然后每张网卡分别连到一个汇聚组的不同 LEAF 上。在 LEAF 和 SPINE 层面,通过 Full Mesh 的方式进行互联。
以下图为例,考虑一个训练场景下的 3 层架构 HPN 网络:

在传统非 MoE 的训练场景下,跨机通信产生的流量大多数都是同号卡流量。例如在梯度同步时候产生的 AllReduce 或者 ReduceScatter 或者 AllGather,PP 之间的 SendRecv 等。同号卡通信最佳情况可以只经过一跳,以上图为例,每个 LEAF 交换机有 64 个下联口,因此 64 台服务器规模同号卡通信理论上可以做到一跳可达。
规模再大,就只能经过 SPINE 或者最差经过 SUPER SPINE 来进行通信。为了减少流量上送 SPINE,百度百舸在任务调度的时候会自动进行服务器的亲和性调度。在创建任务的时候,尽量把同一通信组下的 Rank 排布在同一 LEAF 交换机下的服务器内,那么理论上大部分流量都可以收敛在 LEAF 下。
对于推理服务来说,MoE EP 之间的 Alltoall 通信流量模式与 AllReduce 等不同,会产生大量的跨导轨流量。虽然对于 Prefill 阶段来说,可以通过软件实现层面规避掉跨导轨的流量,但是 Decode 阶段仍然无法避免跨导轨,这会导致多机之间的通信不只是同号卡通信,跨机流量大部分并不能一跳可达,会有大量的流量上到 SPINE 或者 SUPER SPINE,从而导致时延增加。
对于 MoE 训练的流量会更加复杂,综合了训练和推理的流量特征,既存在传统的梯度同步产生的 AllReduce 或者 ReduceScatter 或者 AllGather,PP 之间的 SendRecv,也存在 EP 之间的 Alltoall 流量。这些流量不但会出现跨导轨传输的问题,他们之间可能会存在 overlap 导致互相干扰。
鉴于 Alltoall 通信的特点,我们在设计 HPN 网络的时候,会考虑优先保证跨导轨流量至多 2 跳可达,让 Alltoall 流量收敛到 SPINE 层,以这种方式尽量减少跨导轨的通信时延。如下图所示:

LEAF 层所有设备全部有一根线接入同一台 SPINE 交换机,这样可以让集群内 Alltoall 跨导轨流量全部收敛到 SPINE 层,跨导轨通信时延可以进一步从 5us+ 缩小为 4us。
这种经过优化后的 HPN 网络架构,能接入的卡数主要取决于交换机芯片支持的最大的下联口有多少。虽然对于超大模型的训练任务来说,这个集群规模可能不够,但是对于推理来说,通常不需要那么大规模的机器,是可以满足需求的。
同时,由于 Alltoall 通信流量的特征,LEAF 到 SPINE 之间的通信流量会成为常态。当流量需要通过 SPINE 传输的时候,会由 hash 选择 SPINE 出口的过程,这时候有可能会产生 hash 冲突,导致网络抖动。因此为了避免 hash 冲突,百度智能云基于自研交换机实现自适应路由。如下图所示:
假设 A 和 C 进行 Alltoall 跨导轨通信,A 发出的流量必然要经过 SPINE,那么流量在到达 LEAF 的时候,会基于 packet 做 hash,并结合链路的负载情况动态选择最优的出口,将报文发送到多个 SPINE 上。
基于报文 hash 到不同的物理路径,百度智能云实现了链路负载均衡,消除因 hash 冲突时延增加导致的性能抖动,实现稳定的低时延网络。
详情可参考:彻底解决网络哈希冲突,百度百舸的高性能网络 HPN 落地实践
在推理服务中,EP 间的 Alltoall 通信流量特性与传统训练中的 AllReduce 完全不同,网络上多打一造成的 incast 流量非常常见。这种 incast 的严重程度会随着规模的增大而增大。incast 流量的突发,可能会造成接收侧网卡上联的交换机端口向发送侧反压 PFC,导致网络降速。
传统 Alltoall 流量多打一的示意图如下,假设机器 A 和机器 C 的 GPU0、GPU2、GPU4、GPU6 都需要给机器 B 的 GPU0 发送数据,那么在网络上就会出现 8 打 1 的情况。

传统的 Alltoall 实现,例如 PyTorch 内部调用的 Alltoall,是使用 send recv 去实现的,如果使用 PXN 可以缩小网络上的发生多打一的规模,但是多打一依然存在,如下图所示:

因此无论 Alltoall 的算子实现方式如何,网络上的多打一都无法避免。此时如果网络侧的拥塞控制算法的配置不合理,对拥塞过于敏感,就会产生降速,进而对整体吞吐造成影响。
除此之外,如果集群内还存在其他流量,例如训练任务 DP(数据并行)之间的 AllReduce 或者 ReduceScatter 或者 AllGather,或者 PD(Prefill-Decode)之间的 KV Cache 传输,也会对 Alltoall 的流量造成影响,从而进一步降低推理引擎的整体吞吐。
因此无论是端侧网卡的配置,或者是交换机的配置,都需要针对 Alltoall 这种多打一 incast 流量做针对性优化,同时尽量避免集群内其他流量对 Alltoall 流量造成影响。
针对这种情况,我们给出的解决方案如下:
在经过端侧网卡和网侧交换机配合调整后,可以保障 Alltoall 通信流量的通信带宽和传输时延,实现训推任务的互不干扰,并有效的缓解 incast 流量带来的非预期的降速而造成的性能抖动。
经过测试,在我们部署的推理服务中,Alltoall 过程的整体通信时延有 5% 的降低。
在 PD 分离式推理系统中,还存在 PD 之间 KV Cache 传输的流量。相比 Alltoall 虽然他的带宽需求不算大,但为了避免二者的流量互相干扰,通常我们会让 KV Cache 的传输流量单独走 DCN 网络,使其与 Alltoall 的流量完全隔离开。

在 DCN 网络的设计上,为了保证 KV Cache 流量的传输带宽,其网络架构收敛比采用 1:1。端侧网卡支持弹性 RDMA,使用 RDMA 协议保证 KV Cache 的高性能传输。
在传输库层面,百度智能云使用自研的高性能 KV Cache RDMA 传输库,其接口设计与框架层深度定制,支持上层框架分层传输,也支持多层 KV Cache 的批量传输,便于在框架层做计算与传输的 overlap。
通过以上设计优化,KV Cache 传输在主网卡可以用满带宽,传输时间可以完全被计算 overlap 住,不成为推理系统的瓶颈。
在有了高带宽低时延的网络基础设施的基础上,如何用好网络基础设施,是推理服务软件层面需要重点考虑的事情。
在我们对 PD 分离推理服务的 profile 分析当中,发现了一些影响网络通信效率的关键因素。
目前社区开源的 DeepEP 已经给出了推理系统中 dispatch 和 combine 过程的 Alltoall 高性能的算子的实现,且性能表现优异。
对于 Prefill 来说,由于输入的 batch size 较大,Alltoall 通信算子的同号卡传输阶段为了平衡显存资源和性能,采用分 chunk 传输的方式,发送和接收会循环使用一小块显存,并对每次 RDMA 发送以及机内 NVLink 传输的 token 数做了限制。
通过实际观测网卡的传输带宽,发现其并没有被打满。在此基础上,我们对网络传输的显存的大小,以及每一轮发送接收的最大 token 数等配置,针对不同的 GPU 芯片,做了一些精细化的调整,使之在性能上有进一步的提升。通过优化,DeepEP 的传输性能有大概 20% 的性能提升,网络带宽已经基本被打满。
对于 Decode 来说,DeepEP 的实现是多机之间的 EP 通信,不区分机内和机间,一律采用网络发送。这样做的考虑是为了机内传输也不消耗 GPU 的 SM 资源,完成网络发送后算子即可退出。在网络传输的时间内做计算,完成后再调用 Alltoall 的接收算子,以此来实现计算和通信的 overlap。但这样做的缺点是机内的 NVLink 的带宽并没有被高效的利用起来,网络传输的数据量会变大。
因此,百度智能云通过在 GPU 算子内使用 CE 引擎做异步拷贝,在不占用 GPU SM 资源的情况下,也能实现机内 NVLink 带宽的高效利用,同时不影响计算和通信的 overlap。
EP 专家之间如果出现处理 token 不均衡的情况,将会导致 Alltoall 通信算子的不同 SM 之间,以及不同 GPU 的通信算子之间,出现负载不均的情况,导致的结果就是整体通信时间会被拉长。
由于 EP 专家之间的负载均衡是推理服务引擎提升吞吐的非常重要的一环,经过百度智能云的大规模的线上实践的经验来看,静态冗余专家并不能很好的保证专家均衡。因此我们专门适配了针对 batch size 级别的动态冗余专家,把专家均衡度(max token/avg token)基本控制在了 1.2 以下,不会出现明显的「快慢卡」的情况
通信和计算 overlap,隐藏通信的开销,一直都是推理服务,或者大模型训练中,非常重要的课题。
在百度智能云的实践中,我们在线上大规模的推理服务中开启了双流。为了尽量隐藏掉通信的开销,达到最好的 overlap 的效果,除了做 EP 之间的专家均衡以外,对计算算子也做了针对性的优化,例如对计算算子和通信算子 kernel launch 的顺序做合理排布,对二者所需的 SM 资源做合理的分配,避免出现计算算子占满 SM 导致通信算子 launch 不进去的情况,尽可能的消灭掉 GPU 间隙的资源浪费。通过这些优化,整体的吞吐可以提升 20% 以上。
百度智能云在大规模 PD 分离式推理基础设施优化的实践中,充分展现了网络基础设施、通信组件与上层业务特征深度融合的重要性。这种融合不仅是技术层面的创新,更是对实际业务需求的深刻理解和响应。