用于大语言模型系统的 RDMA 点对点通信

用于处理新兴大语言模型通信模式的优雅工具

作者Perplexity Research

传统上,大语言模型(LLM)系统依赖集体通信来实现训练和推理中的结构化并行模式。然而,诸如解耦推理、专家混合(MoE)路由和异步强化学习微调等新兴工作负载,对点对点通信的灵活性提出了更高要求。

我们最新的研究论文介绍了 TransferEngine,这是一个基于 RDMA 的便携式通信库,可跨不同的网络接口控制器(NIC)运行,包括 NVIDIA ConnectX-7 和 AWS Elastic Fabric Adapter (EFA)。我们展示了它在 Perplexity AI 三个生产系统中的有效性:

解耦预填充与解码:使用分页写入和 UVM 监视器进行逐层 KvCache 传输,使预填充和解码实例能够独立扩展。

在不到 2 秒内完成强化学习后训练的权重传输:实现训练节点和推理节点之间极速的权重同步。

在 AWS EFA 上实现万亿参数模型:使用散点(scatter)和屏障(barrier)原语实现低延迟的 MoE 分发与合并操作,在 ConnectX-7 上达到业界领先的延迟,并成为首个与 AWS EFA 兼容且切实可行的实现。

本文聚焦于驱动这些系统的通信库 TransferEngine 本身。请在 arXiv 上阅读完整研究论文,并在 GitHub 上浏览代码。

动机

集体通信擅长处理张量并行和数据并行等结构化并行模式,但它们施加的约束限制了其在以下新兴工作负载中的适用性:

固定的成员身份:集体通信需要在初始化时建立全局通信域,并要求所有参与者进行同步设置。对于解耦推理而言,这是有问题的,因为预填充和解码实例的数量必须根据流量模式动态扩展。

统一的张量形状:集体通信假设所有参与者交换具有相同形状和数据类型的张量。在 MoE 路由中,发送给每个专家的令牌数量取决于数据,且在运行时发现,这迫使稠密的全对全(all-to-all)集体通信必须填充至最坏情况下的缓冲区大小。

操作顺序:集体通信 API 强制执行同步语义,即使应用程序仅需知道一组操作已完成。例如,在 KvCache 传输中,从预填充实例到解码实例的所有页面可以按任何顺序到达。同样,在强化学习权重更新中,所有权重张量都可以独立传输,而无需排序约束。

基于 RDMA 的点对点通信自然地解决了这些局限性,但有三个障碍阻碍了其采用:

库的局限性:虽然 torch.distributed 和 NCCL 提供了发送/接收(Send/Recv)等点对点 API,但这些操作仍然继承了集体通信的语义和限制。

硬件多样性:带有可靠连接(RC)传输的 NVIDIA ConnectX NIC 是最常见的 RDMA 设置,可保证按序交付。像 AWS 这样的云提供商使用不同的硬件,带有可扩展可靠数据报(SRD)传输的 Elastic Fabric Adapter (EFA) 以乱序方式传递消息。EFA 还要求聚合四个 100 Gbps NIC(或两个 200 Gbps NIC)以实现 400 Gbps 带宽。

软件可移植性:现有的 RDMA 库深受供应商锁定之苦。DeepEP 需要仅在 ConnectX 上提供的 GPU 发起的 RDMA(IBGDA)。NVSHMEM 在 EFA 上表现出严重的性能下降。像 Mooncake 和 NIXL 这样的库缺乏 EFA 支持或仍处于早期开发阶段。因此,对于大语言模型推理,目前尚不存在可跨提供商使用的可行解决方案。

TransferEngine

TransferEngine 通过以下四种核心能力解决了这些障碍:

通用的传输假设:TransferEngine 不针对特定的 NIC,而仅假设底层传输是可靠但无序的,即 RC 和 SRD 传输所提供保证的交集。

ImmCounter 同步:一种新颖的完成通知原语,它无需依赖消息排序,使得同一代码可以在有序(RC)和无序(SRD)传输上运行。

多 NIC 支持:透明地管理每个 GPU 的多个 NIC,自动分片传输以最大化带宽。

极简的 API 界面:足以构建复杂通信模式的一小套操作。

TransferEngine 系统设计图,展示了在两个 GPU/NUMA 域中,请求与回调路径、CPU 工作线程、UVM 缓冲区、GPU 和 NIC。

TransferEngine 为每个 GPU 生成一个工作线程。每个工作线程管理一个 DomainGroup,根据硬件平台协调 1 到 4 个 RDMA NIC(单个 ConnectX-7 提供 400 Gbps,而 EFA 需要四个 100 Gbps 适配器或两个 200 Gbps 适配器)。在 DomainGroup 内,每个 Domain 都是特定于硬件的,并管理一个单独的 NIC,处理队列对管理、工作提交和完成轮询。

每个 TransferEngine 实例都公开一个用于对等识别和发现的主地址。远程对等方交换编码网络地址的 NetAddr 结构,这些结构可序列化以用于带外通信。系统中的所有对等方必须在每个 GPU 上使用相同数量的 NIC。这种统一的配置允许 TransferEngine 智能地分片传输:在发起传输时,源端和目的端都知道哪些 NIC 可用,从而实现跨所有适配器的自动负载均衡。

API 设计

TransferEngine 公开了一个极简 API,抽象了底层 RDMA 接口的复杂性。

rust
#[serde] struct NetAddr(Bytes); #[serde] struct MrDesc{ ptr: u64, rkeys: Vec<(NetAddr, u64)> } struct MrHandle(NonNull<c_void>); type Offset = u64; struct Pages{ indices: Vec<u32>, stride: u64, offset: Offset } struct PeerGroupHandle(u64); struct ScatterDst{ len: u64, src: Offset, dst: (MrDesc,Offset)} enum OnDone { Callback(fn () -> ()), Flag(Atomic<bool>) } trait TransferEngine { fn main_address() -> NetAddr; // Memory Region Management fn reg_mr(ptr, len, device) -> (MrHandle, MrDesc); // Two-sided Send/Recv fn submit_send(addr: NetAddr, msg: &[u8], cb: fn () -> ()); fn submit_recvs(len: u64, cnt: u64, cb: fn (&[u8]) -> ()); // One-sided Write fn expect_imm_count(imm: u32, count: u32, cb: fn () -> ()); fn submit_single_write(len: u64, imm: Option<u32> src: (MrHandle, Offset), dst: (MrDesc, Offset), OnDone); fn submit_paged_writes(page_len: u64, imm: Option<u32>, src: (MrHandle, Pages), dst: (MrDesc, Pages), OnDone); // One-sided Write to a group of peers fn add_peer_group(addrs: Vec<NetAddr>) -> PeerGroupHandle; fn submit_scatter(h: Option<PeerGroupHandle>, OnDone, imm: Option<u32>, src: MrHandle, dst: Vec<ScatterDst>); fn submit_barrier(h: Option<PeerGroupHandle>, OnDone, imm: u32, dst: Vec<MrDesc>); // Watcher for CPU-GPU synchronization fn alloc_uvm_watcher(cb: fn(u64,u64) -> ()) -> NonNull<u64>; }

该 API 由以下几个关键组件组成:

内存注册

RDMA 操作要求在数据传输前将内存区域注册到 NIC。 reg_mr 函数接收指针、长度和设备,并返回两个句柄:

MrHandle:用作传输源的本地句柄

MrDesc:可与远程对等方交换的序列化描述符,使其能够写入注册的内存

MrDesc 结构包含授予远程访问内存区域权限的远程网络地址和远程密钥列表(每个 NIC 一个)。这种设计自然地处理了多 NIC 设置。

双向发送/接收

对于小消息和 RPC 风格的控制平面消息,TransferEngine 提供了 submit_send 和 submit_recvs:

submit_send:将消息发送到由 NetAddr 标识的远程对等方。消息在内部进行复制,因此调用者可以立即重用或释放缓冲区。

submit_recvs:发布一个接收缓冲区池。当消息到达时,将调用包含接收数据的回调函数。回调完成后,缓冲区自动返回到池中。

单向写入

对于批量数据传输,TransferEngine 提供零拷贝单向写入操作:

submit_single_write:将连续区域从源传输到目的地

submit_paged_writes:传输由索引、步长和偏移量指定的非连续页面

这两个操作均接受一个可选的 32 位立即值,该值在完成时原子地交付给接收方。

引擎自动跨多个 NIC 对大文件传输进行分片,以最大化带宽。

对等组和批量操作

对于针对多个对等方的操作(常见于 MoE 路由),TransferEngine 提供了优化的原语:

add_peer_group:预注册一组对等地址以便快速查找

submit_scatter:向组中的每个对等方发送不同的数据切片

submit_barrier:向所有对等方发送仅含立即值的通知(零字节传输)

这些操作通过缓存内部数据结构和模板化 RDMA 工作请求来优化低延迟,避免了重复设置的开销。

UVM 监视器:GPU 到 CPU 的信号通知

大语言模型系统中的一个关键挑战是协调 RDMA 传输与 GPU 进度。理想情况下,数据传输应在 GPU 产生输出后立即开始,但从 CPU 轮询 GPU 是低效的。

TransferEngine 通过 UVM 监视器 解决了这个问题:这是统一的虚拟内存位置,可由 GPU 内核(甚至在 CUDA 图中)更新,并由 CPU 使用 GDRCopy 高效轮询。分配 UVM 监视器时,会注册一个回调函数,每当 GPU 更新该值时,该回调就会被调用。

这实现了逐层 KvCache 传输等模式:在每个 Transformer 层完成后,GPU 会增加监视器的值,触发该层 KvCache 的 RDMA 传输,同时后续层继续进行计算。

完成通知

TransferEngine 为发起方(submit_write 的调用者)和接收方(写入操作的目标)提供了完成通知机制。

发起方通知:当提交写入操作时,调用者可以选择通过回调函数或原子标志在完成后获得通知。此通知确认数据已成功从源缓冲区传输。

接收方通知:核心创新是 ImmCounter,这是一个跟踪由 32 位立即数索引的计数器的组件。当写入操作携带立即值时,接收方的 ImmCounter 在数据交付时原子地增加。expect_imm_count 函数允许接收方注册一个回调,当特定的立即值达到目标计数时触发。计数器可通过 GDRCopy 与 GPU 透明同步、直接轮询或通过回调监控。

这种接收方机制提供了可靠的同步,而无需消息排序。它在 ConnectX NIC(保证 RC 传输的按序交付)和 EFA NIC(通过 SRD 传输乱序交付消息)上同样有效。

实现

TransferEngine 使用 Rust 实现,仔细优化了分配、问题(线程)和同步,以最大限度地减少延迟并实现高吞吐量。该引擎为每个 DomainGroup 生成一个工作线程,每个线程绑定到设备所连接 NUMA 节点上的 CPU 核心,从而最小化调度和内存访问延迟。特定于域的数据结构在绑定后分配,以确保在正确的 NUMA 节点中保留内存。一个工作线程最多处理 4 个 Domain(每个管理一个 NIC),而另一个专用线程则轮询 GPU 以更新 UVM 监视器。跨线程通信通过无锁队列完成。

API 将请求转发给服务于相应设备的 DomainGroup,其中第一个也服务于主机。在紧密循环中,域工作问题(线程)轮询新工作,优先提交新请求。根据硬件和配置,请求会在可用的 Domain 之间进行分片和负载均衡。复合请求的第一次写入会立即发布到 NIC 的发送队列。一旦新请求耗尽,工作问题(线程)就会继续处理待定操作,发布写入以填满硬件管道。随后,轮询完成队列以查询已完成的传输和立即数计数器的增加。事件会被聚合以提供每个传输的通知,并将传输移交给所有组共享的专用回调工作问题(线程)。

DomainGroup 内部的分片是灵活的。传输可以通过索引针对特定的 NIC。单个写入可以拆分。Paged 传输、散点和屏障操作(所有这些都会转换为多次写入)可以在所有 NIC 之间分片。

硬件特定优化

TransferEngine 内的 Domain 针对其控制下的硬件进行了专门化和优化:

AWS EFA:EFA 实现使用 libfabric,在 DomainGroup 内为每个 NIC 管理一个 fabric 域。关键实现决策和优化包括:

零大小写入:EFA 不同于 RDMA 规范,它要求即使对于零大小的写入,也需要有效的目标描述符。因此,TransferEngine 强制要求所有传输必须具有有效的描述符。

工作请求模板化:在发布前预填充并保留 libfabric 结构通用字段,减少批量传输和对等组操作的每个操作开销

NVIDIA ConnectX-7:ConnectX-7 实现使用带有 RC 传输的 libibverbs。关键实现决策和优化包括:

用于握手的 UD 队列对:每个 Domain 使用不可靠数据报(UD)队列对来交换 RC 连接握手。

独立的 RC 队列对:为每个对等方创建 2 个 RC 队列对,一个用于双向发送/接收操作,另一个用于单向写入和 WriteImm 操作。这种分离是必要的,因为 Recv 和 WriteImm 完成均以发布顺序消耗工作请求,从而在允许高层级 Recv 语义的同时支持 WriteImm 而不会产生干扰。

工作请求模板化:为批量传输和对等组预填充工作请求结构的通用字段。

工作请求链接:通过 ibv_send_wr 的 next 指针链接最多 4 个工作请求,减少对 NIC 的门铃信号次数。

宽松排序:启用 IBV_ACCESS_RELAXED_ORDERING 以允许 NIC 和 GPU 内存之间出现乱序的 PCIe 事务。

结论

通过针对不同 RDMA 传输保证的交集,并引入通过 ImmCounter 实现的独立于顺序的同步,TransferEngine 为大语言模型系统中的点对点通信提供了便携式基础。该库是开源的,可在 GitHub 上获取,我们鼓励社区探索并在此基础上构建大规模机器学习系统中的新兴通信模式。

引用

BibTeX
@online{perplexity2025RdmaPointToPointCommunicationForLlmSystems,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{RDMA point-to-point communication for LLM systems}},
  date = {2025-11-05},
  url = {https://www.perplexity.ai/hub/blog/rdma-point-to-point-communication-for-llm-systems}
}