Previewing UCCL-EP: Flexible and Efficient Expert Parallelism for Cloud and Beyond
Previewing UCCL-EP: Flexible and Efficient Expert Parallelism for Cloud and Beyond
作者/机构: Ziming Mao (加州大学伯克利分校), Yang Zhou (加州大学戴维斯分校), Yihan Zhang (加州大学戴维斯分校), Chihan Cui (威斯康星大学麦迪逊分校), Zhongjie Chen (清华大学), Zhiying Xu (AWS) 等 UCCL-EP 贡献者
A1 主要贡献
本文针对专家并行(Expert Parallelism, EP)中存在的通信效率和平台可移植性问题,提出了UCCL-EP,一个灵活且高效的专家并行解决方案,旨在推动EP在公有云和异构设备(包括不同厂商的GPU和NIC)上的普及。
核心问题:
1. 现有方案的局限性:当前高效的EP通信库,如DeepEP,严重依赖于NVIDIA特定的硬件(NVIDIA GPU、NVIDIA NIC)和软件栈(NVSHMEM、GPUDirect、IBGDA),导致其无法在采用不同硬件(如AWS EFA网卡)的公有云或其他异构环境中运行,限制了其可移植性和部署灵活性。
2. 细粒度通信的挑战:EP通信涉及大量小消息(例如,推理时7KB,训练时256KB)的频繁交换,这使得为大负载优化的通用集合通信库(如NCCL)效率低下,因为延迟和同步开销成为主要瓶瓶颈。
3. 控制与可观察性缺失:DeepEP通过IBGDA技术将NIC驱动逻辑移至GPU核心执行,虽然提升了小消息效率,但也牺牲了对通信过程的精细控制和可观察性。GPU线程“盲目”地发出RDMA操作,难以实现流量控制、拥塞避免和完成状态监控,可能导致网络拥塞或高尾延迟。
研究目标与创新点:
UCCL-EP旨在解决上述挑战,其核心设计思想是在不牺牲性能的前提下,实现GPU与CPU职责的清晰分离,从而获得高度的灵活性和可移植性。
1. 解耦GPU与NIC的控制:UCCL-EP的核心创新在于,虽然通信仍由GPU发起以实现计算与通信的重叠,但实际的RDMA网络操作由一个轻量级的多线程CPU代理来处理。GPU通过高效的共享内存通道向CPU发送轻量级控制命令,CPU代理再使用标准化的libibverbs接口代表GPU执行这些网络操作。
2. 跨厂商的硬件支持:通过依赖于Linux-RDMA社区维护的、厂商中立的libibverbs库,UCCL-EP能够支持各种符合该标准的RDMA网卡,如AWS EFA、Broadcom Thor、AMD Pensando等,并计划支持AMD和Intel的GPU,从而打破厂商生态系统的锁定。
3. 高效的GPU-CPU通信机制:UCCL-EP设计并实现了一个无锁的FIFO队列作为GPU(生产者)和CPU(消费者)之间的共享内存通道。通过每个GPU配置多个通道和专用的RDMA队列对(QP),该机制能够支持每秒数千万次的RDMA请求,确保GPU到CPU的命令转发不会成为性能瓶颈。
4. 软件层面的功能模拟与增强:针对不同NIC硬件能力的差异,UCCL-EP在软件层面实现了关键功能。例如,为解决AWS EFA网卡不保证包内有序性的问题,UCCL-EP通过在控制消息中嵌入序列号实现软件层面的消息重排;为解决EFA不支持硬件RDMA原子操作的问题,它通过常规的RDMA写操作和立即数(immediate data)来模拟原子操作,甚至通过将原子信号附加到数据写操作中,将两次操作合并为一次,提高了效率。
通过这些设计,UCCL-EP不仅在功能和接口上与DeepEP保持一致,支持解码的低延迟模式和预填充/训练的常规模式,还在AWS等异构云平台上展现出超越现有技术的卓越性能。
A3 背景知识与设计原则
专家并行(EP)的通信模式。专家并行(EP)被广泛应用于大规模混合专家(MoE)模型中,模型的不同“专家”子集被放置在跨多个节点的不同GPU上。在推理或训练过程中,每个输入token都根据一个学习到的门控函数被路由到一个或几个选定的专家。这种选择性路由需要在网络中进行频繁的“分发”(dispatch,将token嵌入发送到正确的专家GPU)和“合并”(combine,将专家输出收集回发送方rank,然后进行加权求和)操作。这些数据交换通常通过远程直接内存访问(RDMA)在高速互连(如InfiniBand或RoCE)上执行。
EP通信的挑战:细粒度消息。与传统的数据或张量并行(通信涉及兆字节或千兆字节级别的大型连续张量)不同,EP通信是高度细粒度的。每次分发或合并操作通常涉及小消息尺寸——例如,在像DeepSeek V3这样的模型中,从7 KB(推理)到256 KB(训练)。如此小的消息尺寸对像NCCL这样的通用集合通信库构成了挑战,因为这些库是为大负载的高吞吐量传输(例如,在all-reduce或all-gather操作中)而优化的。当消息如此之小时,单次传输的延迟和同步开销占据主导地位,导致网络带宽利用率低下。因此,EP系统通常需要定制的、低延迟的通信运行时,能够有效地重叠计算和通信,并处理大量并发的小消息操作。
现有方案DeepEP及其局限性。一个流行的EP库是DeepEP,它利用NVIDIA特有的NVSHMEM/IBGDA技术,让NVIDIA GPU能够直接向NVIDIA NIC发出RDMA操作,以实现小消息的高效处理。IBGDA本质上是在GPU的SM核心内运行NIC驱动函数,这样GPU就可以绕过CPU直接与NIC通信。因此,GPU可以直接将RDMA写、读或原子操作排队到NIC的门铃寄存器。然而,尽管DeepEP性能很高,但由于GPU和NIC之间的紧密耦合,它存在两个局限性。
DeepEP的第一个局限:厂商锁定。DeepEP与NVIDIA的软硬件生态系统紧密耦合。它依赖于NVIDIA GPU、NVIDIA NIC及其专有网络栈(例如,NVSHMEM、GPUDirect和IBGDA)。因此,DeepEP只能在这些组件被协同设计和支持的NVIDIA控制的平台上运行。这种设计极大地限制了可移植性。例如,DeepEP无法在AWS云实例上运行,因为AWS使用弹性结构适配器(EFA)RDMA NIC,而不是NVIDIA硬件。在部署非NVIDIA RDMA解决方案的其他公有云和数据中心环境中,例如Broadcom Thor NIC、Google Cloud Falcon NIC和AMD Pensando NIC,也会出现类似的不兼容性。同样的限制也适用于GPU厂商——DeepEP对NVIDIA特定API和设备驱动接口的依赖,使得它即便在存在可比的RDMA网络硬件时,也难以(如果不是不可能的话)在AMD或Intel GPU上运行。随着现代AI集群在GPU架构和网络结构方面变得越来越异构,这种缺乏跨厂商可移植性的问题日益限制了部署的灵活性。
DeepEP的第二个局限:缺乏精细控制与可观察性。通过将NIC驱动逻辑移入GPU线程,DeepEP牺牲了对通信过程的精细控制和可观察性。在传统的CPU驱动的RDMA系统中,主机管理流量控制、队列深度、完成通知以及跨多个网络队列的负载均衡。这些机制对于确保低尾延迟、避免拥塞以及在高网络压力下进行恢复至关重要。然而,在IBGDA模型中,GPU直接发出RDMA操作,无需CPU的协调。换句话说,GPU线程是“盲目地”发出一系列单边操作。这使得监控或调节流量变得困难。例如,GPU可能会在不了解NIC队列利用率的情况下发布大量未完成的RDMA写操作,导致拥塞或高尾延迟。在GPU或DeepEP中,检测传输完成、测量传输延迟或处理网络背压也是不可能的,因为IBGDA或NVSHMEM没有暴露相关接口。
A2 方法细节
UCCL-EP核心思想:解耦GPU与CPU的职责。UCCL-EP的核心见解是,高效的专家并行通信虽然受益于GPU发起,但并不需要GPU直接控制NIC。相反,UCCL-EP恢复了GPU和CPU之间明确的职责分离:GPU保留其在大数据密集型任务上的大规模并行能力——例如token打包、专家合并、NVL转发和本地RDMA缓冲,并与后台RDMA通信有效重叠。GPU仍然负责发起通信。CPU则通过一个轻量级的多线程CPU代理来处理网络中控制密集型的方面——包括队列管理、流控制、完成处理和负载均衡。本质上,每个GPU不是直接向NIC发布RDMA操作(如NVIDIA的IBGDA模型),而是通过一个高效的共享内存通道将轻量级控制命令——例如“将此token写入对等方X”——转发给CPU。一个多线程CPU代理池随后解释这些命令,并代表GPU向NIC发出实际的RDMA verbs。这样,GPU仍然发起通信,从而与GPU上的数据复制和计算重叠,而CPU则管理RDMA通信的复杂性和控制决策。我们注意到,UCCL-EP的方法与同样使用CPU代理的NVSHMEM的IBRC解决方案有相似之处,但通过利用多个CPU代理线程来提升性能,并支持广泛的供应商以实现可移植性,从而与之产生显著区别。
利用标准化接口libibverbs实现跨厂商兼容性。该设计利用了一个关键观察:每个RDMA NIC已经通过由Linux-RDMA社区维护的libibverbs库暴露了一个标准化的、厂商中立的接口。通过让GPU将RDMA请求通过PCIe转发给CPU线程,虽然通信仍由GPU发起,但UCCL-EP可以使用任何NIC驱动程序都支持的相同verbs API来代表GPU发出网络操作。
CPU-GPU通信延迟并非主要瓶颈。UCCL-EP设计的第二个观察是,CPU-GPU通信延迟不是主要的瓶颈。现代互连技术如PCIe Gen5、NVLink和C2C(芯片到芯片)链路提供了微秒级的延迟和数十至数百GB/s的CPU与GPU之间的带宽。这意味着将一个控制命令从GPU转发到CPU的速度极快——尤其是与穿越网络的端到端RDMA操作延迟相比。此外,专家并行中的每个控制命令通常代表一个涉及多个token的批量数据移动。因此,通过PCIe发送一个命令描述符的摊销成本相对于传输的数据来说是微不足道的。
构建高效的GPU-CPU转发通道。UCCL-EP中的一个核心挑战是构建一个高效的GPU和CPU之间的转发通道,该通道必须能够每秒维持数千万个RDMA请求而不会成为瓶颈。UCCL-EP将此通道实现为一个在GPU生产者和CPU消费者之间共享的无锁FIFO队列。每个GPU将轻量级的RDMA传输描述符入队,而多个CPU线程则出队并执行它们。
每GPU多通道设计以提升性能。UCCL-EP为每个GPU部署了多个通道;每个通道都有自己专用的RDMA队列对(QP)。这种设计使得每个GPU能够以适度的延迟开销实现每秒超过5000万次RDMA操作(如UCCL PR #454所示),此时NIC的固有延迟和网络延迟——而不是CPU-GPU通道——成为主要成本。每个TransferCmd很小,仅占用128位。通道的尾指针(tail)由CPU线程读取并分配在主机内存中,而头指针(head)由GPU线程读取和更新并分配在设备内存中。它进一步在GPU上缓存尾指针的值以加快访问速度。
应对特定NIC的挑战:乱序交付。不同的NIC供应商由于传输协议和硬件能力的变化,引入了额外的系统级挑战。例如,AWS EFA NIC使用可扩展可靠数据报(SRD)协议,该协议采用先进的多路径技术来缓解大规模拥塞。虽然这种设计提高了吞吐量和可靠性,但它不提供单个SRD队列对(QP)内的严格有序交付保证。这对于依赖有序RDMA写操作后跟原子操作来通知远程GPU写操作已传递到RDMA传输缓冲区中指定位置的DeepEP式通信来说是个问题。为了解决这个问题,UCCL-EP利用其CPU端的灵活性来实施软件级的消息排序。每个RDMA写操作都携带一个立即数(immediate data),该立即数编码了一个每个RDMA通道的序列号,接收方使用这个序列号在将乱序消息提交到GPU内存之前对其进行重新排序。重要的是,这只适用于控制消息(例如原子操作),而不适用于数据载荷。
软件模拟RDMA原子操作。此外,在DeepEP的NVIDIA特定IBGDA路径中,GPU依赖硬件RDMA原子操作来发出远程完成信号。然而,一些NIC提供商(例如EFA)本身不支持RDMA原子操作,这带来了正确性挑战:接收方仍然必须知道载荷何时被完全写入,然后才能继续读取或合并它。为了模拟这种行为,UCCL-EP使用常规的RDMA写操作和立即数来实现软件级原子操作。发送方首先写入载荷,然后发出一个携带立即数的小型RDMA写操作,该立即数充当原子消息(例如,新的计数器值或标志)。在接收方,CPU代理更新一个本地完成计数器——从而有效地再现了硬件原子操作的行为。这种方法一个令人惊讶的好处是,原子操作可以附加在现有的写操作上(通过一次RDMA操作实现),而在DeepEP中,这需要两次操作(写+原子)。根本原因在于UCCL-EP采用双边通信模型,而原始的DeepEP由于GPU需要直接与NIC对话,是单边通信模型。
移除NVSHMEM依赖以增强可移植性。为了使UCCL-EP能够与不同的GPU供应商合作,我们迈出了消除nvshmem依赖的第一步,这对于可移植性以及其他特性(如弹性伸缩)至关重要。我们还有趣地观察到,移除nvshmem依赖有时可以带来性能提升,我们怀疑这是由于nvshmem库的内部开销所致。
A4 实验环境与结果
实验环境
-
硬件配置:
- 测试平台1 (EFA性能测试): 2个节点,每个节点配备H200 GPU和400 Gbps的AWS EFA RDMA网卡。
- 测试平台2 (EFA扩展性测试): 多个节点(支持高达32个专家的并行度),每个节点包含8块H200 GPU,并通过一个400 Gb/s的EFA RDMA网卡互连。
- 测试平台3 (GH200性能测试): 一个小型的GH200测试平台,其特点是CPU和GPU之间通过高速的NVLink-C2C互连,并具备缓存一致性。
-
软件与模型配置:
- 测试核心: 使用未经修改的Perplexity MoE Kernels进行基准测试。
- 基线方案:
- Torch: 使用其AlltoAll API,并为其编写了高效的打包和解包核心。
- NVSHMEM: 同样使用其AlltoAll API和定制的打包/解包核心。
- PPLX。
- 模型参数 (扩展性测试): 遵循DeepSeek-V3的预训练配置,具体为:每个批次4096个token,隐藏层维度7168,top-4门控分组,top-8专家选择,分发(dispatch)操作使用FP8精度,合并(combine)操作使用BF16精度。
实验结果
实验一:在EFA平台上对比不同方案的Dispatch吞吐量
- 实验内容: 在2节点的H200 + EFA平台上,比较UCCL-EP与Torch、NVSHMEM、PPLX等基线方案在不同token数量下的FP8 Dispatch吞吐量。
- 实验结果: 如图4所示,UCCL-EP的性能显著优于所有基线方案。随着token数量的增加,UCCL-EP的吞吐量线性增长,在4096个token时接近50 GB/s。相比之下,Torch和NVSHMEM的吞吐量则停滞在10 GB/s以下,而PPLX在处理4096个token时直接出错。
- 分析结论: 该结果表明UCCL-EP能够在AWS EFA等非NVIDIA原生硬件上高效运行,其架构设计成功克服了异构平台的挑战,并展现出卓越的性能和可扩展性。
实验二:在EFA平台上测试UCCL-EP的扩展性
- 实验内容: 在多节点H200 + EFA平台上,使用DeepSeek-V3的配置,测试UCCL-EP在不同专家并行规模(EP=8, 16, 24, 32)下的Dispatch和Combine带宽。
- 实验结果:
- 节点内 (EP=8): Dispatch和Combine的瓶颈带宽均超过300 GB/s,受限于NVLink。
- 节点间 (EP≥16): Dispatch带宽随着并行规模的扩大而提升,超过50 GB/s。而Combine带宽稳定在40 GB/s左右,略低于Dispatch带宽。
- 分析结论: Dispatch操作表现出良好的扩展性。Combine带宽较低的原因可能有两个:一是合并操作本身包含额外的计算开销(如对专家输出进行加权求和);二是由于每个节点会进行本地聚合,导致Combine操作发出的RDMA消息数量少于Dispatch(例如在TopK=8时),使得RDMA流水线未能完全填满。文章指出,这个问题在原始DeepEP中也有被提及(Issue #82),可能与拥塞有关,团队仍在调查EP=16时Combine吞吐量相对较低的原因。
实验三:在GH200平台上与DeepEP的初步比较
- 实验内容: 在一个小型GH200测试平台上,将UCCL-EP的性能与原始的DeepEP进行比较。
- 实验结果: 令人惊讶的是,UCCL-EP的性能甚至超过了原始的DeepEP。
- 分析结论: 研究人员推测了两个可能的原因:
- GH200平台上高速的NVLink-C2C互连和CPU-GPU缓存一致性,使得CPU与GPU之间的通信开销几乎可以忽略不计,这非常有利于UCCL-EP的架构。
nvshmem库本身可能存在一定的内部开销,移除它带来了性能增益。
研究人员强调,这一发现需要在更大规模的测试平台上进行验证。
A5 结论与未来工作
UCCL-EP目前仍处于积极开发阶段,预计在未来几个月内会有新的成果发布。团队计划发布一篇正式的博客,介绍其在应用层面的性能以及在AMD GPU和其他NIC厂商硬件上的表现。
当前路线图包括:
* 进一步提升UCCL-EP在EFA平台上的性能。
* 完成向AMD GPU和Broadcom NIC的移植工作(PR #457)。
* 在CPU代理中实现更高级的流控制和拥塞管理机制。
* 将UCCL-EP集成到vLLM和SGLang等流行的推理框架中,并欢迎社区贡献。
💬 评论讨论
欢迎在这里分享您的想法和见解!