Lequn Chen || abcdabcd987

大语言模型系统中RDMA通信的一些探索

2025-11-09T00:00:00-08:00

上周我们公司把最近对大语言模型中点对点通信的一些成果总结了一下，写了一篇论文挂在了 arXiv 上面，同时也在 GitHub 上面开源了代码。

我们构建了一套基于“无序可靠数据报传输（Unordered Reliable Datagram）”语义的 RDMA 通信库，既能跑在 AWS EFA，又能跑在 NVIDIA ConnectX。我们把这套 RDMA 通信库应用在了三个场景下面：分离式推理的 KvCache 传输、强化学习后训练中的模型参数更新、以及 MoE 通信。这个 MoE kernel 在 ConnectX-7 上面跑 decode 甚至比 DeepEP 还快一点点，在 EFA 上也是首次达到了可用的性能。

这篇文章我跟大家讲讲其中的来龙去脉，更多的是想分享一下做这些工作的动机以及背后的故事。对具体的技术细节感兴趣的读者可以在文章末尾的链接里面找到我们的论文、代码、以及相关的博客。

集合通信的痛点

首先，我们觉得集合通信虽然说在结构化并行的场景下久经考验（例如数据并行和张量并行），但是在一些新的应用场景用起来会很变扭，不如直接用 RDMA 点对点来得方便和高效。为什么说变扭呢？

第一，集合通信需要构建一个通信“世界”，而且这是一个成员固定的世界，没法增加或者删除成员，这就会给生产部署带来非常大的难题。举个例子，在分离式推理架构里面，Prefiller 和 Decoder 之间需要传输 KvCache。而我们知道，生产环境的线上流量总是会有涨落的，所以总是需要根据负载弹性地增加或者减少副本的数量。另一方面，无论是硬件故障还是软件有瑕疵，总是会有需要把某个坏掉的节点去掉的时候。而如果通信的成员只能是固定的，那既没有办法弹性伸缩，也没办法替换坏节点。

第二，通信世界的建立是阻塞的，而且需要所有成员的参与。这就意味着每次增减节点都需要让全世界停下来。虽然能用一些比较复杂的工程代码来异步地建立新的通信世界，但是这就很麻烦，而且也很容易出错。

第三，集合通信的接口保证了一个顺序执行的语义。其实这个本意是很好的，可以简化开发者的负担。然而，网络可以是乱序到达的，所以有可能就需要做一些额外的缓冲或者同步来保证顺序执行的语义。可惜的是，集合通信库辛辛苦苦为大家做了一个这么有序的抽象，有的应用场景却不领情。比方说还是 KvCache 的传输，我们只是想把所有的页面都传出去，至于哪个页面先到哪个后到则是不重要的。另外一个例子就是强化学习后训练当中训练节点往推理节点更新模型参数，我们在乎的是所有的模型参数都更新好，而不在乎哪个先哪个后。

第四，集合通信的接口要求所有参与者都有相同的张量形状和类型，这样就会带来一些编程上的不便，有些时候还会带来很严重的性能损失。比方说如果用集合通信来做 RPC，那么每次就都得传最大的消息大小。再比方说 MoE，往每个 Rank 要发多少个 token 是没法提前知道的，所以如果用集合通信来做 all-to-all 的话，那么就得按照最坏情况来算 token 数量（也就是所有 rank 都发往同一个 rank）。这种最坏情况跟平均情况可以差非常多，如果用 DP64 EP64 那就是每个 rank 要多传64倍数据，而且还是无用的数据。

RDMA 的痛点

那为什么大家还是执着于集合通信呢？我感觉一个原因是缺少简单易用的 RDMA 通信库，另一个原因是现有的 RDMA 库基本上都绑定英伟达网卡 ConnectX。众所周知，“N卡网速快”，这句话本来只是个2010年代的梗，结果在英伟达收购了 Mellanox 之后，这句话就变成了现实。我们在论文里面对比了，是真的快。而且 ConnectX 网卡就是直接可以在市场上面购买到的商品。又快，又好买，又有大量的软件支持，也不难怪大家都在 ConnectX 的生态上继续开发。

但是对于技术实力强劲的云服务商来说，他们投入资源开发自己的网卡也是对的。最开始的目的应该还是为了解决云服务器的存储和网络的虚拟化、安全性、减少网络通信的 CPU 占用。后面高性能计算以及机器学习的应用需求来了，再缝缝补补接着往上尽量补全功能，尽量能兼容已有的生态。比如 AWS 就有专门的 NCCL 插件，能让 NCCL 跑在 AWS EFA 上面。

能跑 NCCL 并不代表就能跑其他的通信库。比方说 Mooncake、DeepEP 这样非常优秀的通信库，都没有办法在 AWS EFA 上面跑。我们3月份的时候也尝试过依托 NVSHMEM 来写一个 MoE 通信库，但是效果非常不好。在 ConnectX-7 上面，哪怕是用了 NVSHMEM 的 IBGDA 实现，速度也远比直接用 mlx5 驱动的 DeepEP 要慢。虽然 NVSHMEM 也支持 EFA，但是我们当时测下来发现要比 ConnectX-7 慢40倍，比 NCCL 还慢很多，完全不可用。

为什么 RDMA 程序难以移植到不同的网卡呢？我觉得有下面几个原因：

第一，程序的逻辑会依赖于底层通信协议的特性。大家用得最多的 RDMA 通信协议是可靠连接（Reliable Connection, RC）。RC 保证了消息是按顺序到达的。比方说如果我有一堆的消息要发，要怎么让远端知道这些消息都送到了呢？如果依赖于消息按序送达的假设，那就可以在最后一个消息上面打上一个记号，这样远端一旦收到了这个记号，它就知道了前面所有的消息都送达了。这里“打记号”的方式可以有很多种，比方说可以在最后一个消息上面带上一个32位的立即数（Immediate），再比方说对一个计数器进行原子操作。然而 EFA 用的是自研的可扩展可靠数据报协议（Scalable Reliable Datagram, SRD），虽然它保证消息能送达，但是送达的顺序是乱序的。

这个时候我们想象一下像 NVSHMEM 这样的通信库要怎么实现可移植性，让一个基于按序送达假设的程序能在一个无序送达的硬件上跑起来。显然什么都不做是会破坏正确性的。所以为了兼容这个语义，通信库就不得不多做很多额外的工作，比如可以使用缓冲区来进行消息重排，比如可以插入同步点来等待消息队列清空。显然这些权宜之计都会严重地损失性能。

第二，网卡本身的带宽也不同。ConnectX-7 网卡，一张网卡就能达到 400 Gbps 的带宽。然而 AWS p5 实例上，一张 EFA 网卡就只有 100 Gbps 带宽。p5en 实例好一点，一张 EFA 网卡有 200 Gbps 带宽。总之，需要把多张网卡聚合起来才能达到跟 ConnectX-7 一样的带宽。

第三，网卡支持的功能也不同。之前提到的 IBGDA 技术能让显卡直接向网卡发起通信，然而到目前为止这个功能只有 ConnectX 网卡支持。在缺少这个功能的情况下，又要让显卡能往网卡发起消息，那就只能让 CPU 来充当中间商了。如果这个软件实现做得比较差，那就会带来很多额外的开销。后文我还会展开讨论这个点。

第四，程序为了追求极致的性能，用的接口可能本身就是不可移植的。比方说 DeepEP 虽然说表面上用了 NVSHMEM，但其实只是把 NVSHMEM 当作一个启动器，真正跟网卡打交道的是直接用 mlx5dv，这就直接绑定了 ConnectX 网卡。

解决痛点的思路

基于上面这些背景，我解决上面这些问题的思路是取不同通信协议和硬件特性的交集，在此之上构建一套尽可能通用的 RDMA 通信库：

可靠传输：网卡和网络层需要保证消息能送达。一来是省得在通信库里面实现重传等麻烦的逻辑，二来是充分利用硬件卸载的优势，尽量减少 CPU 的参与，避免 CPU 拖慢性能。
无序送达：前面提到，如果硬要在无序送达的硬件上实现按序送达的语义，那就会损失性能。另一方面，应用场景也不一定需要按序送达的语义。所以不如直接放弃按序送达的语义。
支持双边 RDMA 操作：RECV / SEND 非常适合用来实现 RPC。
支持单边 WRITE_IMM：单边 WRITE 可以提供极致的性能，而带上一个32位的立即数又可以用来通知目标方操作已经完成。
基于 ImmCounter 的同步机制：虽然应用场景不一定需要按序送达，但是还是需要一些同步机制来保证所有的消息都送达了。要在无序送达的硬件上实现这个语义，我们构建了一个 ImmCounter 抽象，即每一个 WRITE 操作都带上一个用户指定的立即数，然后目标方的通信库会对这个立即数进行计数，当计数器达到预期的值时，就认为所有的消息都送达了。
针对常见的通信模式提供高性能的实现优化：比方说同样是 WRITE，除了传输一个连续的内存区域，还可以是多个不连续的内存区域（例如 KvCache 的传输），还可以是多个接收方并且每个接收方传输的数据量不同（例如 MoE 中的 scatter）。虽然理论上来说这些操作都可以在应用层通过单个 WRITE 操作来实现，但是在网络库提供这些额外的接口还是有很多好处。一是可以简化应用层的实现；二是可以减少通信库和应用程序之间跨语言及跨线程的通信开销；三是可以在内部缓存一些数据结构来优化性能。
支持多网卡聚合：网络库可以提供一些默认的聚合策略，例如对于单个大消息，可以对消息字节数进行切分；对于多个页面的传输，可以对页面的列表进行切分；对于多个接收方的传输，可以对接收方的列表进行切分。当然，如果应用程序想要特殊的控制，也可以手动指定网卡聚合策略。
借助 CPU 来支持显卡发起 RDMA 通信（Host-Proxy 模式）：KvCache 传输和 MoE 都需要在 Cuda Graph 里面发起 RDMA 操作，那我们就往内存和显存里面放一些特殊标记。因为 CPU 和显卡各自都可以访问自己和对方的内存，所以两者都可以通过写入和轮询这些标记来传递消息。

注意到，这里我舍弃掉了对已有 RDMA 程序的兼容性。其一是因为对网络底层假设发生了变化，为了保证性能，我们要求应用程序必须按照我们这套更松弛的假设来编写。其二是因为我和我同事们都喜欢造轮子，所以上层的应用迟早都会被我们造一遍，也没有去兼容已有程序的需求。

我造的这套通信库一开始只是为了 KvCache 传输一个用处，后面发现用在强化学习后训练的模型参数更新上面也非常方便。让我比较吃惊的是，这套通信库用在 MoE 上面效果也出奇的好。

MoE Kernel

MoE 的故事还得回到今年5月份去 MLSys 开会的时候，在跟字节的朋友们交流的时候得到了一个非常重要的洞见：CPU 和显卡之间的 PCIe 延迟只有2微秒。我跑了一下 GDRCopy 自带的性能测试，确实如此。所以开完会之后我就想开一个坑，用这套通信库实现类似 NVSHMEM 的功能，来弥补没有 IBGDA 的遗憾，毕竟一次 MoE 通信也得好几百微秒，就算这套 host-proxy 带来15微秒的额外开销，也完全能够接受。然后在考虑 EFA 的乱序送达特性的时候，我意识到了直接移植现有的 DeepEP 代码或者我们之前的 NVSHMEM 实现都可能带来性能损失，要达到最好的性能最好还是按照乱序送达的语义来重新设计一套 MoE kernel。正好我同事也对他上次用 NVSHMEM 实现的 MoE kernel 不太满意，所以这个艰巨的任务就交给他了。我的同事也是非常给力，单枪匹马很快就实现好了一个新的 MoE kernel。

看到新的 MoE kernel 能在 EFA 上面跑起来，我们都很兴奋。但是我们不知道它跑得慢到底是我们优化得不够，还是 EFA 本身的性能不好。所以我们就想着要在 ConnectX-7 上跑起来，然后跟 DeepEP 比一比。算上 CPU 和显卡之间的 PCIe 延迟，以及 CPU 上面的各种开销，我们觉得优化的终点应该就是比 DeepEP 慢个15微秒就差不多可以收工了。

要让我们的 MoE kernel 能在 ConnectX-7 上面跑起来，首先我们得有 ConnectX-7 机器，于是就到处化缘。在这里我要非常感谢英伟达，他们非常慷慨地给我们提供了一个64卡的开发集群。

有了开发机，我就开始往我们的 RDMA 通信库里面加 ConnectX-7 的支持。我本来以为 libfabric 本身就能支持 ConnectX-7，结果发现怎么都跑不起来，问了一下 libfabric 社区和英伟达，他们都说 ConnectX 没有作为开发的重点，所以我猜就算能跑性能也不好。于是我就用 libibverbs 自己写了一套。好在我之前造的这套 RDMA 通信库的抽象做得还行，所以只需要最底下硬件相关的代码换一下就好了；中间的抽象不需要做大的改动；暴露给应用的接口更是不需要改。

在加 verbs 支持的时候，也算是对比了一下 SRD 和 RC 两种协议的编程接口。虽然说 EFA 跑得慢，但从编程接口的角度来说，我更喜欢 SRD 一点：

SRD 是基于数据报的协议，知道地址就能直接发消息，很方便。RC 是基于连接的协议，所以跟每一个远端节点都需要建立一个连接。然而要建立 RC 连接，需要用另一个信道来交换彼此的 QPN 和 PSN，这个握手的过程就变成一个需要双方参与的异步操作。为了不引入额外的通信信道，所以我就用 UD 实现了这个握手过程。
我有考虑过维护连接这个事情到底是在网络库里面做比较好还是让应用程序来做比较好。最后我觉得，还是让网络库把这个复杂度吃掉比较好。毕竟我在写网络库的时候都已经觉得维护连接是一个很麻烦的事情了，更不用说应用的开发者了。（结合上前面的乱序送达特性，我觉得可以把我们这个网络库对网络的假设称作 Relaxed Reliable Datagram，RRD。）
libfabric 接收 WRITE_IMM 的立即数是不需要消耗 RECV 操作的，然而 verbs 会消耗，而且 verbs 的 RECV 操作是按照提交的顺序消耗的。这就意味着没法把 WRITE_IMM 和正常的双边 RECV 操作混合在一起跑。好在我后面想了一个挺简单的解决方法，开两个 RC QP，把单边操作和双边操作分开跑。

在 ConnectX-7 上面跑通了我们 RDMA 通信库的点对点测试之后，我就开始胆战心惊地尝试跑 MoE kernel。本来以为肯定还要在折腾各种奇怪的问题的，结果没想到竟然第一次就丝滑地跑通了！这也是对我们这套 RDMA 通信库的可移植性的一个很好的证明。

跑通了之后同事和我就开始做优化。优化的维度有很多：调整算法增加并行度；调整 RDMA 和 NVLink 的同步；调整 Cuda kernel 里面的 PTX 指令、内存屏障；优化 load/store；减少 CPU 的内存分配次数；在 RDMA 通信库里面引入多个接收方的优化接口；ConnectX 和 EFA 通用的优化以及各自单独的优化；等等。具体的细节就不展开了，有兴趣的读者可以看我们的论文、其他几篇博客、以及具体的代码。

最终优化完了之后，没想到在 ConnectX-7 上面 Decode 竟然跑的比 DeepEP 还要快。为了确保结果的可比性，我还特地让 DeepEP 用跟我们相同的性能测试代码来跑，测出来的结果也跟 DeepEP 官方的性能测试结果差不多。另外一方面，虽然我们一直是在 ConnectX-7 上面做优化，但这些优化大多也同时作用于 EFA 上，所以最终我们 EFA 上面的性能也提升了不少。

在这里我也要非常感谢 AWS 尤其是 EFA 团队。他们回复我们问题的速度非常快，给了我们很多关于 EFA 的建议。我们之后也会接着和 AWS 合作，继续优化我们的通信库以及 MoE kernel 在 EFA 上的性能。

在对 Decode 结果满意之后，我们又试着跑了跑 Prefill。跟 DeepEP 不同，我们 Prefill 和 Decode 用的是同一个实现，不过我们还没专门针对 Prefill 做优化。测出来我们能打满 RDMA 带宽，不过还是比 DeepEP 慢了不少。最让我们觉得奇怪的是，我们用 RDMA 带宽以及传输的字节总数算了一下理论上的最低延迟，结果发现 DeepEP 的延迟比理论的更低。在仔细检查了性能测试程序，确认一些无误之后，唯一的解释就是 DeepEP 传的字节总数比我们的少。后面又请教了一下 DeepEP 的朋友，证实了这个猜测。在 dispatch 的时候，如果同一个 token 发给同一台机器的不同 rank，DeepEP 可以只用 RDMA 发一次，然后用 NVLink 在机内进行转发；在 combine 的时候，如果同一台机器的不同 rank 会发同一个 token 给同一个远程的 rank，DeepEP 可以先在 NVLink 上做一个部分和，然后在 RDMA 上就只需要发送一个 token 就行了。不得不说 DeepEP 真的是太精妙了！这也解释了为什么 DeepSeek-V3/R1 里面选专家的时候会限定最多分到4个节点上。

感慨万千

到这里，我们的优化工作也暂告一个段落，于是我们整理了一下思路，写了一篇论文，也把开源代码放出来了。感兴趣的读者可以点击下面的链接查看：

arXiv 2510.27656: RDMA Point-to-Point Communication for LLM Systems
GitHub: pplx-garden
公司博客：
我的博客：
- 跨机秒传RL模型参数更新的一些探索及后续
- 驾驭3200Gbps网络

这一路上我从同行朋友们的交流以及开源项目中学到了太多东西。我自己写博客以及在公司里面推动技术博客和代码开源，也算是向社区做出的一点微小的贡献吧。

文章的最后我想感慨一下，在不到一年的时间里面，在 RDMA 这方面我真的学习到了很多。去年12月初的时候我还不懂怎么在 AWS 上跑 RDMA，趁着圣诞节放假摸索了怎么用 EFA，到现在有了一套比较完整的解决方案。我们小小的4人 Inference Team （这两个月加到了7个人）在一年半里面也做了很多事。去年4月份我加入的时候，我们还完全依赖 TensorRT-LLM，后面我们搞了自己的推理引擎、搞了上层的智能路由、搞了下层的算子和网络。一开始我只是把我已经知道的推理优化认真的做工程落地，做着做着发现我们有了一些工程上的创新，再到现在感觉有了学术研究贡献的创新，真的是收获的远远比我预想的多。

（如果大家对我为什么选择了我司以及对我司的工作环境感兴趣，我后面也可以专门讲讲。）

如果在美国的读者有对我司的系统优化感兴趣的，欢迎直接联系我。对我司后训练、AI 产品、或者其他职位感兴趣的，我也可以帮忙内推。

跨机秒传RL模型参数更新的一些后续

2025-09-17T00:00:00-08:00

上一篇博客介绍了我们如何实现跨机秒传RL模型参数更新。这一篇博客简单补充一些后续：

Kimi-K2 (1T params)，256卡BF16训练，128卡FP8推理，参数更新只需要不到 1.3 秒。
参数更新的流水线稍微再优化了一些，增加了两个可以并行的项目：H2D Memcpy 和全局通讯屏障。
跑了一遍 PyTorch Profiler，方便直观地分析参数更新流水线，看看时间都花在哪里了。
加了一些图方便理解。

流水线优化

之前的版本只有GPU和网络这两个流水线阶段，现在增加到了四个：

H2D Memcpy
GPU 操作：full_tensor()、投影融合、量化
RDMA 传输
全局通讯屏障：在每个 Mesh Group 的最后一次 full_tensor() 之后就可以加一个 async_op=True 的全局通讯屏障。不用等到 RDMA 传输结束。因为 Gloo 走 Ethernet，所以这又是可以并行的。

PyTorch Profiler

首先看一次完整的参数更新。从上图可以看到，总时间大概1.2秒出头。

参数更新分成了两个 Mesh Group，中间有全局通讯屏障。前面一次是在传非 MoE 参数，FSDP 在机内 NVLink 维度上做切分。后面一次是传 MoE 参数，FSDP 在跨机 RDMA 维度上做切分。

对于非 MoE 参数，full_tensor()（也就是 NCCL all-gather）跑得很快，因为是通过 NVLink 传输的。所以总体上是在等 RDMA 传输。中间的全局通讯屏障也完美地被藏在 RDMA 传输里面了。

对于 MoE 参数，full_tensor() 肉眼可见地拖慢了非常多的速度。全局通讯屏障也没能够被藏起来，也就是 RDMA 传输比全局通讯屏障完成得更早，应该是在等某个跑得特别慢的 Rank。

上面是放大看非 MoE 参数的传输。

上面是放大看 MoE 参数的传输。

一个明显的区别是，对于非 MoE 参数，H2D Memcpy 是一口气完成的，而对于 MoE 参数，H2D Memcpy 始终是掺杂在整个时间线里面的。这说明 MoE 参数的传输任务并不是在一开始就全部开始的，而是随着前面一些传输任务的完成而逐步开始的。这说明 MoE 参数的传输可能撞到了我之前设置的最大临时内存限制。

补充一些图

Rank0-based 和 P2P 参数更新的对比

Torch.Distributed + RPC 和 RDMA 参数更新的对比

跨机秒传RL模型参数更新的一些探索

2025-09-07T00:00:00-08:00

我最近花了两周时间把 Qwen3-235B （BF16 训练，FP8 推理）的跨机（128卡训练，32卡推理）参数更新跑通了，只需要2秒。这篇博客我打算不单单是给读者呈现一个解决方案，而是记录一下我的探索过程，以及我的一些思考。过几天也会在公司博客上发一篇精简版。

起因

之前公司博客发的文章里面有提到，我们自己手搓了一个 LLM 推理引擎。另外大家也知道我司虽然不做预训练，但在后训练上花了非常多力气。所以说把我们自己的推理引擎接到我们的后训练框架上是一个很自然的事情。要做这个事情，第一步就是先搞定从训练节点到推理节点的参数更新。

因为我们用的是异步 RL 训练，所以训练和推理跑在不同的机器上。这样一来我在实现参数更新的时候就更简单了，连用 cuMem 把推理引擎的参数卸载掉都不用，直接把参数从训练节点传到推理节点就可以了。

当同事告诉我开源的 RL 框架需要好几分钟来更新参数的时候，我还是有点震惊的。比方说我们拿 DeepSeek 671B 来举例，就假设把参数切分到8张卡上好了，每张卡有 400 Gbps 也就是 50 GB/s 的网络，那把所有参数传出去也只需要 671 GB / (8 * 50 GB/s) = 1.68 秒。当然网络不可能跑到这么快，接收方也不止有一台机器，但反过来说模型也可能被切得更碎，并且训练节点也可能比推理节点更多。所以这些正负一抵消，我感觉参数更新的时间大概就是1秒这个数量级，再怎么样偏差也不可能到分钟级别。

不过现有方案跑得咋样跟我也没关系，反正我要做的是把我们自己的推理引擎接上去。我之前对训练这块也完全没了解过（读者读到后文就会惊呼，不是哥们，这你都不懂？），所以也懒得再去学习现有的方案是啥了，直接就开干。

预想的解决方案

读过我之前的教程或者读过公司博客的另两篇文章的读者应该有印象，我手搓了一个 RDMA 的通信库。所以在我的脑海里，传参数应该是一个和把大象放进冰箱里面一样简单的事情：

主控节点先把每张训练 GPU 和每张推理 GPU 上面的参数的元数据拿出来。
主控节点跑一个参数名的匹配
主控节点算一个固定的参数更新路由表，也就是每一个训练 GPU 需要把哪些参数发给哪个推理 GPU
主控节点把路由表发给每张训练 GPU
需要做参数更新的时候，主控节点只需要告诉所有训练 GPU 现在开始传参数啦。
每张训练 GPU 只需要按照路由表用 RDMA 把参数传给对应的推理 GPU 就好了。
推理 GPU 甚至都不知道自己的参数被更新了。

看起来这甚至都不需要修改我们的推理引擎，只需要在主控节点写一些很细致的参数匹配和路由表计算逻辑，以及在训练节点写一点调用 RDMA 通信库的代码就行了。

写成伪代码的话大概是这样子：

@ray.remote
class TrainingWorker:
    def get_param_metadata(self) -> dict[str, ParamMeta]: ...
    def set_routing_table(self, routing_table: RoutingTable) -> None: ...
    def transfer_weights(self) -> None:
        for entry in self.routing_table:
            src_mr = get_mr(entry.src_ptr)
            submit_rdma_write(src_mr, *entry)
        wait_for_rdma_writes_to_complete()

@ray.remote
class RolloutWorker:
    def get_param_metadata(self) -> dict[str, ParamMeta]: ...
    def get_memory_regions(self) -> list[MemoryRegion]: ...

@dataclass
class WeightTransferEntry:
    src_ptr: int
    src_size: int
    dst_ptr: int
    dst_size: int
    dst_mr: MemoryRegion

RoutingTable: TypeAlias = list[WeightTransferEntry]

def controller_main() -> None:
    trainers: list[TrainingWorker] = ...
    rollouts: list[RolloutWorker] = ...
    trainer_params = ray.get([x.get_param_metadata.remote() for x in trainers])
    rollout_params = ray.get([x.get_param_metadata.remote() for x in rollouts])
    rollout_mrs = ray.get([x.get_memory_regions.remote() for x in rollouts])

    schedule: list[RoutingTable] = compute_weight_transfer_schedule(trainer_params, rollout_params, rollout_mrs)
    ray.get([trainer.set_routing_table.remote(routing_table)
             for trainer, routing_table in zip(trainers, schedule)])

    while training_not_done:
        train()
        ray.get([trainer.transfer_weights.remote() for trainer in trainers])
        rollout()

RDMA 传输参数的概念验证

因为之前没了解过我们的训练框架，也没接触过 Ray，也没用过 slurm script，所以我打算一切从简，先验证一下我们的推理引擎能够接受 RDMA 传输的模型参数。

我的打算是用两个 GPU，一个跑我们的推理引擎，但是不加载模型参数，另一个按照我们推理引擎的格式加载模型参数，再把参数传给推理引擎。然后我们让推理引擎随便生成一点文字，如果参数传输是成功的，那么生成的文字应该是正常的。

对于这个概念验证，我也不需要用到我们训练框架的代码，所以我也不用着急去了解里面有什么。Ray 看起来也很好入门的样子，就一个 @ray.remote 一个 ray.get() 就能用了。然后 slurm script 也可以省了，反正我就用两块 GPU，在开发机上起一个 Ray 头节点就行了。

这样一来开发效率非常高，两天就搞定了。其中唯一遇到的困难是，有4个 tensor 总是会传输失败。

传输失败

首先我先尝试了一下遇到失败的时候重新传一次，但是无论怎么重试都不行，先睡几秒再传也没用。libfabric 给的错误信息也很模糊：

Unexpected status returned by responder

我在 libfabric 和 rdma-core 的代码搜了一下，不是特别能看出来是啥问题，这个错误是更底层的 EFA 驱动抛出来的。我隐约觉得可能是我的代码里面内存区域算错了，但是一般来说如果内存区域算错了，libfabric 的错误信息会说是 rkey 不对。

抱着路是问出来的心态，我先去问了问 AWS 的工程师。AWS EFA 团队的工程师也很给力，消息回的很快，他们也觉得这个错误消息很奇怪，他们也怀疑是内存区域算错了。

于是我就着重朝着这个方向去排查。首先我注意到的一个特点是，这4个传输失败的 tensor 都非常小，都不到 1KB（就是 FP8 Block Quant 的 inv_scale）。我怀疑就是我的 RDMA 通信库的分片算错了。

因为一张 EFA 网卡的速度不到 400 Gbps，所以 AWS 会给每张 GPU 多插几张 EFA 网卡（比如 p5 是4张100 Gbps，p5en 是2张200 Gbps）。所以我在 RDMA 通信库里面会把单个 RDMA WRITE 均匀地拆到所有 EFA 网卡上。然后我有一个小小的策略是，如果传输的数据太小（小于 MTU），那后面几张网卡就只传0个字节（0-length WRITE）。

然后我在算内存区域的偏移量的时候，没有特殊考虑 0-length，所以有可能出现最后一个 0-length WRITE 的指针正好在内存区域的末尾。众所周知，末尾都是开区间，所以其实这并不是一个合法的内存区域。

于是我在计算内存区域偏移量的时候，特判了一下 0-length WRITE，发现问题就解决了。

AWS EFA 的工程师告诉我，在 RDMA 规范里面，对于 0-length WRITE，是不需要验证目标方的内存区域的，但是 EFA 却还是做了检验。他们会把这个跟规范有偏差的地方改过来，也会改进错误信息。这里给 AWS EFA 团队点个赞！

获取 PyTorch 内存区域

熟悉 RDMA 或者读过我之前的 RDMA 教程的读者都知道，所谓内存区域就是4个东西：ptr, size, rkey, lkey。拿着 ptr 和 size 就可以在 RDMA 通信库里面注册内存区域了。但是我们要怎么拿到 PyTorch 分配的内存区域的元数据呢？

最朴素的办法当然是对于每一个 tensor，用 data_ptr() 得到指针，用 numel() * element_size() 得到大小。但这样需要注册非常多非常小的内存区域，感觉不是很优雅。我感觉既然 PyTorch 有内存池，那就把内存池里面的每个内存块都注册成内存区域不就好啦。

我本来以为得写一点 C++ 代码来把 PyTorch 内存池的信息导出来，好在在动手之前先搜索了一遍文档。我在 PyTorch 的文档里面找到了这么个函数：torch.cuda.memory.memory_snapshot()。这个函数可以返回 PyTorch CUDACachingAllocator 分配的所有内存块的指针和大小，然后我们就可以愉快地拿着这两个信息去注册 RDMA 内存区域了。

传输速度

试了一下传输 DeepSeek-V2-Lite 的模型参数，大概 16 GB，花了0.44秒，大概 36 GB/s，这还是在没有预热的情况下跑出来的，我感觉很可以了。然后推理引擎在拿到参数后也能够输出正常的文字，说明参数传输是成功的。

真实的参数更新（惨败）

有了概念验证给的信心，我接下来就开始实现真实的参数更新了。我想着说先来个简单点的吧，就先考虑 DeepSeek-V2-Lite，也不考虑量化了，就先训练和推理都用 BF16。那在这样简化的情况下，唯一需要额外考虑的就是投影的融合了（比如把 {q,k,v}_proj 融合成一个 qkv_proj）。考虑到所有的融合都是在 dim0 上拼接的，而 PyTorch tensor 又是 row-major 的，所以其实在传输的时候也只需要加一个偏移量就好了。

于是我花了几天把参数名映射和路由表计算写了，训练节点只用1个 GPU 的时候也能跑通。我觉得这样拓展到8卡甚至多机训练也应该是没问题的，下一步就可以考虑量化了。然而现实却狠狠地打了脸。只要训练节点用了不止1个 GPU，我这套方案就跑不通，连 shape 都对不上。

仔细观察 shape 的差异，我发现差了8倍。我突然意识到，我一直以为 FSDP 就是普通的 DP，没想到它还把模型参数给做了切分。看了一眼 FSDP 的文档，一边在赞叹这个 API 做得好，一边又默默悲伤我的方案要推倒重来了。

真实的参数更新（推倒重来）

在看了 FSDP、DTensor 以及 DeviceMesh 的文档后，我意识到在计算路由表的时候，得按照 DTensor 的 Placement 来对参数进行切分。Placement 分成两种（不会出现 Partial），一种是 Replicate，另一种是沿着 tensor 的某一维 Shard。观察了一下我们训练框架产生的 Placement，最多切分一个维度。所以相对来说这个复杂性还比较低，要写的话很快可以写出来。总体的框架不用变，还是算好路由表直接 RDMA WRITE 就完事了。

DTensor

但是我马上想到了，如果我自己做这个切分的话，我还得自己算投影融合的各种偏移量，这个就更麻烦了。更加棘手的是量化，如果每个训练 GPU 的模型分片大小不能被 block quant 的大小整除，那还要二次计算那些处在边界上的元素。另外，FSDP 还可以把参数卸载到 CPU 上，这样我还得手动把数据从 CPU 传输到 GPU。

想了想，不如还是跟着 DTensor 的逻辑走吧。这样虽然做不到最快，但是起码很好写。只要做完 full_tensor()，DeviceMesh 里面的每一个 GPU 都拥有了完整的参数，可以很方便地计算投影融合以及进行量化。然后 DeviceMesh 里面的任意一个 GPU 都可以作为源，向有需要的推理节点传输参数。

DeviceMesh

注意到 full_tensor() 是一个集合通讯操作（对于 Shard Placement 而言，就是一个 all-gather），所以 DeviceMesh 里面的所有 GPU 都需要按照同样的顺序调用 full_tensor()。不同的 DeviceMesh 之间倒是不需要进行同样的集合通讯操作。

很自然地，我的下一个想法就是，看看到底有多少不同的 DeviceMesh，以及他们之间能不能并行呢？

好在这一题的答案也很简单。考虑 FSDP, PP, EP 这三个维度，假设一共有32卡， FSDP=2, PP=2, EP=8。首先无视掉 PP，因为它不对参数进行切分。对于非 MoE 参数（比如 {q,k,v}_proj），我们的训练框架会在 EP 维度上 Shard，在 FSDP 维度上 Replicate，所以一个 DeviceMesh 里面会有 2 * 8 = 16 卡，整个集群被分成两个不相交的 DeviceMesh，我们可以把它称为一个 DeviceMesh Group：

Non-MoE DeviceMesh Group:
DeviceMesh 0: 0~7 , 16~23
DeviceMesh 1: 8~15, 24~31

对于 MoE 参数，我们的框架在 FSDP 维度上 Shard，所以一个 DeviceMesh 里面会有2张卡，整个集群被分成16个不相交的 DeviceMesh：

MoE DeviceMesh Group:
DeviceMesh  0:  0, 16
DeviceMesh  1:  1, 17
...
DeviceMesh 15: 15, 31

因为一个 DeviceMesh Group 里面的 DeviceMesh 之间是没有交集的，所以他们可以并行地进行通信。我们只需要把整个参数传输的过程按照 DeviceMesh Group 进行划分就行了，在每个 DeviceMesh Group 传输结束后，需要放入一个全局的通讯屏障。

这个寻找 DeviceMesh Group 的过程很简单，随便写个贪心啥的就行了。

def find_disjoint_mesh_groups(mesh_set: set[Mesh]) -> list[set[Mesh]]:
    """Example:
    Input: {
        Mesh([range(0, 8), range(16, 24)]),
        Mesh([range(8, 16), range(24, 32)]),
        Mesh([0, 16]),
        Mesh([1, 17]),
        ...
        Mesh([15, 31]),
    }
    Output: [
        {
            Mesh([range(0, 8), range(16, 24)]),
            Mesh([range(8, 16), range(24, 32)]),
        },
        {
            Mesh([0, 16]),
            Mesh([1, 17]),
            ...
            Mesh([15, 31]),
        },
    ]
    """

训推参数匹配

参数匹配是一个比较细节的活，当然我们可以根据肉眼观察规律来写一个匹配器，但是我还是想把这个匹配流程写的更加系统化一点。想象这是一个编译器，经过一个一个的 pass，我想要能够优雅地把路由表算出来。

首先为了区分有没有 quant scale，以及把 quant scale 和参数打包在一起，我先引入两个类：

@dataclass(slots=True, frozen=True)
class Identity:
    base_name: str

    @property
    def weight_name(self) -> str:
        return self.base_name + ".weight"


@dataclass(slots=True, frozen=True)
class Quantization:
    base_name: str
    scale_suffix: str

    @property
    def weight_name(self) -> str:
        return self.base_name + ".weight"

    @property
    def scale_name(self) -> str:
        return self.base_name + self.scale_suffix

然后我们可以写一个函数去把所有的参数名变成变成这两者之一：

def group_quant_weight_names(
    scale_suffix: str,
    weight_names: set[str],
) -> list[Identity | Quantization]:
    """
    For example,
    - ["foo.weight", "foo.weight_inv_scale"] -> [Quantization("foo", ".weight_inv_scale")]
    - ["bar.weight"] -> [Identity("bar")]
    """

对于训练节点，我们需要考虑重命名以及投影融合。因为每个模型的规则是不一样的，所以每个模型需要写一个函数，输出一个从推理节点参数名到训练节点参数名的一个映射。在训练节点参数名这边，除了上面引入到两种类型，我们额外增加一种 ProjectionFusion，用来表示投影融合。

@dataclass(slots=True, frozen=True)
class ProjectionFusion:
    weight_names: tuple[Identity, ...] | tuple[Quantization, ...]

class ModelWeightMatcher(Protocol):
    def map_trainer_weight_name(
        self,
        trainer_quant_weight_names: list[Identity | Quantization],
    ) -> dict[str, Identity | Quantization | ProjectionFusion]:
        """
        For example:
        Input: [
            Quantization("layers.0.attn.q_proj", ".weight_inv_scale"),
            Quantization("layers.0.attn.k_proj", ".weight_inv_scale"),
            Quantization("layers.0.attn.v_proj", ".weight_inv_scale"),
            Identity("layers.0.mlp.down_proj"),
        ]
        Output: {
            "layers.0.attn.qkv_proj": ProjectionFusion(
                weight_names=(
                    Quantization("layers.0.attn.q_proj", ".weight_inv_scale"),
                    Quantization("layers.0.attn.k_proj", ".weight_inv_scale"),
                    Quantization("layers.0.attn.v_proj", ".weight_inv_scale"),
                )
            ),
            "layers.0.mlp.down_proj": Identity("layers.0.mlp.down_proj"),
        }
        """

有了这些铺垫之后，参数匹配就只剩下一些细致的检查了：

@dataclass(slots=True, frozen=True)
class WeightNameMapping:
    trainer: Identity | Quantization | ProjectionFusion
    rollout: Identity | Quantization
    do_quant: bool

def match_weight_names(
    trainer_named_parameters: dict[str, ParamMeta],
    rollout_named_parameters: dict[str, ParamMeta],
    matcher: ModelWeightMatcher,
) -> list[WeightNameMapping]:
    trainer_quant_weight_names = group_quant_weight_names(trainer_named_parameters)
    rollout_naming_map = {x.weight_name: x for x in group_quant_weight_names(rollout_named_parameters)}

    trainer_weight_name_map = matcher.map_trainer_weight_name(trainer_quant_weight_names)
    for expected_rollout_weight_name, trainer_naming in trainer_weight_name_map.items():
        rollout_naming = rollout_naming_map[expected_rollout_weight_name]
        # Expand shape by mesh placement
        # Check shape, dtype, duplicate, etc...

生成路由表

生成路由表的逻辑也很简单，我们要生成的就是给每一个训练 GPU 算好要按照什么顺序把哪些参数传给哪些推理 GPU。

@dataclass(slots=True)
class WeightTransferEntry:
    name_mapping: WeightNameMapping
    rollout_workers: tuple[int, ...]

@dataclass(slots=True)
class WeightTransferGroup:
    mesh_group: set[Mesh]
    transfer_entries: list[WeightTransferEntry]

@dataclass(slots=True)
class WeightTransferRoutingTable:
    groups: list[WeightTransferGroup]

@dataclass(slots=True)
class WeightTransferSchedule:
    trainers: list[WeightTransferRoutingTable]

def generate_table_for_mesh_group(
    trainers: list[ParametersMetadata],
    rollouts: list[ParametersMetadata],
    name_mappings_with_mesh: list[tuple[WeightNameMapping, Mesh]],
    mesh_group: set[Mesh],
) -> list[WeightTransferGroup]:
    transfer_entries_list: list[list[WeightTransferEntry]] = [[] for _ in trainers]
    send_bytes = [0] * len(trainers)
    for name_mapping, mesh in name_mappings_with_mesh:
        owners = mesh.members()
        for i_rollout, rollout in enumerate(rollouts):
            winner = ... # min owner by send_bytes
            # Assign winner to i_rollout
        for owner in owners:
            transfer_entries_list[owner].append(...)

def compute_weight_transfer_schedule(
    trainers: list[ParametersMetadata],
    rollouts: list[ParametersMetadata],
    trainer_named_parameters: dict[str, ParameterMetadata],
    name_mappings: list[WeightNameMapping],
) -> WeightTransferSchedule:
    mesh_set: set[Mesh] = ...
    mesh_groups = find_disjoint_mesh_groups(mesh_set)
    trainer_tables = [WeightTransferRoutingTable([]) for _ in trainers]
    for mesh_group in mesh_groups:
        trainer_groups = generate_table_for_mesh_group(trainers, rollouts, name_mappings_with_mesh, mesh_group)
        for table, group in zip(trainer_tables, trainer_groups):
            table.groups.append(group)
    return WeightTransferSchedule(trainer_tables)

唯一要注意的就是，前面提到了 full_tensor() 是一个集合通讯操作，所以这意味着一下两件事：

DeviceMesh 里面的所有 GPU 都要按照同样的顺序执行 full_tensor()。哪怕某张训练 GPU 不需要把当前的参数传给任何推理 GPU，它也一样得调用 full_tensor()。
DeviceMesh Group 之间需要插入一个全局同步屏障 torch.distributed.barrier()。所以我们在生成路由表的时候，要把每个 DeviceMesh Group 的路由表区分开来。

在执行 full_tensor() 之后，DeviceMesh 里面所有 GPU 都拥有了完整的参数，都可以作为源，那么我们要怎么进行负载均衡呢？

因为 GPU 两两之间的 RDMA 带宽都是相同的，所以我这里用了一个非常简单的办法：统计每个训练 GPU 需要传输的字节数。对于每个参数，每个推理 GPU 都选择当前累计传输字节数最小的训练 GPU 作为源。

这个当然不是最优的解法，但是我感觉这个简单粗暴的方法应该已经够用了。

执行参数更新

实现完生成路由表，下一步就是实现在训练 GPU 上运行的传输逻辑了。这里有 GPU 操作（把参数从 CPU 读到 GPU、full_tensor()、投影融合、量化），还有 RDMA 网络传输操作。我想要把这两者并行起来。另外，为了防止爆显存，我还想能够在传输的过程中（大致地）限制最大的显存使用量。

我的做法大概是这样的。先给每一个传输任务搞一个类来存状态：

@dataclass(slots=True)
class _WeightTransferTask:
    spec: WeightTransferEntry
    num_transfers: int
    total_bytes: int

    # Step 1: Gather full_tensor() (async GPU operations)
    weight_full_tensors: list[torch.Tensor] | None = None
    scale_full_tensors: list[torch.Tensor] | None = None
    # Step 2: Transform tensors (async GPU operations)
    # Step 2a: Fuse projection
    # Step 2b: Quantize on-the-fly
    weight_tensor: torch.Tensor | None = None
    scale_tensor: torch.Tensor | None = None
    # Step 3: Wait for async GPU operations to finish
    gpu_op_done: torch.cuda.Event | None = None
    # Step 4: Submit RDMA transfers
    submitted_transfer: bool = False
    # Step 5: Wait for transfers to finish
    finished_transfers: int = 0

    def is_done(self) -> bool:
        return self.submitted_transfer and self.finished_transfers == self.num_transfers

然后我们把所有的任务分成三类：

还没开始跑的
在等 GPU 操作跑完的
在等 RDMA 传输跑完的

class transfer_weights:
    def __init__(
        self,
        fabric: TransferEngine,
        model: nn.Module,
        transfer_entries: list[WeightTransferEntry],
        max_tmp_bytes: int = 1<<30,
    ) -> None:
        self.tmp_bytes = 0
        self.rdma_done_queue: Queue[_WeightTransferTask] = Queue()

        tasks: deque[_WeightTransferTask] = deque()
        for entry in transfer_entries:
            tasks.append(_WeightTransferTask(...))

        self.tasks_not_started = tasks
        self.tasks_waiting_gpu_op: deque[_WeightTransferTask] = deque()
        self.tasks_waiting_transfer: deque[_WeightTransferTask] = deque()

        while self.tasks_not_started or self.tasks_waiting_gpu_op or self.tasks_waiting_transfer:
            self._poll_progress()

然后我们每一次执行的时候，就根据任务的状态，决定跑什么，以及是否转变到下一个状态：

class transfer_weights:
    ...

    def _poll_progress(self) -> None:
        # Clear finished tasks
        while self.tasks_waiting_transfer:
            task = self.tasks_waiting_transfer[0]
            if not task.is_done():
                break
            task = self.tasks_waiting_transfer.popleft()
            self.tmp_bytes -= task.total_bytes

        # Kick off async GPU operations
        while self.tasks_not_started:
            task = self.tasks_not_started[0]
            if self.tmp_bytes + task.total_bytes > self.max_tmp_bytes:
                break

            task = self.tasks_not_started.popleft()
            self.tasks_waiting_gpu_op.append(task)
            self.tmp_bytes += task.total_bytes

            self._to_device(task)            # roughly: .to(device, non_blocking=True)
            self._gather_full_tensors(task)  # roughly: .full_tensor()
            self._fuse_projection(task)      # roughly: torch.cat()
            self._quantize(task)

            task.gpu_op_done = torch.cuda.Event()
            task.gpu_op_done.record()

        # Wait for async GPU operations to finish and submit transfers
        while self.tasks_waiting_gpu_op:
            task = self.tasks_waiting_gpu_op[0]
            assert task.gpu_op_done is not None
            if not task.gpu_op_done.query():
                break

            task = self.tasks_waiting_gpu_op.popleft()
            self.tasks_waiting_transfer.append(task)

            self.fabric.submit_write(
                ...,
                callback=(lambda task: lambda: self.rdma_done_queue.put(task))(task),
            )
            task.submitted_transfer = True

        # Handle completed transfer
        while True:
            try:
                task = self.queue.get_nowait()
            except Empty:
                break
            task.finished_transfers += 1

细心的读者可能注意到了，在执行 GPU 操作的时候，我们看起来调用的是阻塞式的函数，但是其实这里并不会真的阻塞。我在这里走的一个弯路是我试图在 .full_tensor() 的时候使用 async_op=True 改成异步的，然后用 AsyncCollectiveTensor.completed 来检测是否完成。然而这样做始终会传出错误的参数，除非我加一个 time.sleep(0.01)。在仔细看了 Torch Distributed 关于异步的语义的文档之后，我才明白，这里 completed 只是说提交到了 Cuda Stream，并不是代表这个操作完成了。

而 PyTorch 的 kernel 调用其实也都只是把操作提交到 Cuda Stream，所以除非我们显式地要求同步，不然这些 GPU 操作也不会阻塞我们的 Python CPU 线程。所以我之前在这里用 async_op=True 完全是画蛇添足。

那么怎么知道 GPU 操作完成了呢？一般来说我们写 PyTorch 代码的时候不需要关心这个，因为所有的操作都是提交到 default stream 上面的。但是在这里，我们需要一个非阻塞的方法来知道 GPU 操作完成了，然后才能让 RDMA 传输开始。要实现这个，我们可以在提交完所有的 GPU 操作之后，插入一个 Cuda Event，然后通过 event.query() 来知道 GPU 操作是否完成了。

另外，以上代码是用来传单个 DeviceMesh Group 的，别忘了要加入一个全局通讯屏障。

@ray.remote
class TrainingWorker:
    def set_routing_table(self, routing_table: WeightTransferRoutingTable) -> None: ...

    def transfer_weights(self) -> None:
        for group in self._routing_table.groups:
            transfer_weights(..., group.transfer_entries)
            torch.distributed.barrier()

至此，我们的实现就完成了。

秒传 Qwen3-235B FP8

经过一些简单的调试，这一次的代码很轻松地就跑通了 DeepSeek-V2-Lite（BF16 训练，BF16 推理），然后我又试了一下加上 FP8 block quant，也顺利地秒传了。对于 Qwen3-235B（BF16 训练，FP8 推理），我稍微改了一下投影融合的规则，然后也非常轻松地跑起来了（128卡训练，32卡推理）。

上图可以看到根据路由表算出来每张训练 GPU 需要发送的字节数，以及每张推理 GPU 需要接收的字节数。虽然有些不平衡，但也马马虎虎。

上图可以看到，一次完整的传输啪地一下就完成了，只花了不到2秒。虽然传输的带宽只有大概 5 GB/s，跟 50 GB/s 的上限以及之前概念验证跑出来的 36 GB/s 有很大的差距，但是考虑到这里我们做了很多的 GPU 操作，或许也是可以理解的。

为什么跑不满带宽

虽然我没有打算着急做更多的优化，不过我还是想知道为什么离跑满带宽还有那么大的差距，我想知道这些时间都花在哪里了。所以我就多插了一些 Cuda Event，算了算时间。

目测了一眼数据，发现主要还是慢在提交 GPU 操作以及等待 GPU 操作完成上面。过两天我再仔细汇总一下数据，然后回来补个分析。

为什么现有框架跑不快

（这里没有拉踩同行的意思，我只是想从门外汉的角度猜测一下为什么现有框架跑不快，如果说得不对欢迎指正。另外，我们自己的推理引擎和后训练框架支持的功能都非常有限，只关注几个我们特定的场景，而不是各种复杂功能的排列组合，所以工程实现以及优化都不需要考虑太多兼容性问题，比较容易落地。用我们这种精简功能的框架去跟成熟的框架比，也对成熟的框架不公平。更不用说在开源社区能够协调这么多贡献者和用户也是很不容易的事情。我这里只有对同行深深的敬意。）

对现有框架的快速调研

因为我对 RL Infra 这个领域不熟悉，所以不太清楚现有的框架更新参数具体需要多少时间。看到一个比较新的数据是来自两周前的博客：

目前 slime 可以做到 7s 完成训推一体下 Qwen3 30B-A3B 模型 bf16 权重的参数同步。100s 完成 GLM4.5 355B-A32B 的 fp8 blockwise 量化 + 参数更新。

（UPDATE: 2025-09-08 @朱小霖提供了一个更新的数据）

在训推分离的配置下（64 卡训练，64 卡推理），slime 中 Qwen3 235B-A22B 从训练这边的 bf16 参数更新至推理测的 fp8 参数大约需要 8s

NeMo-RL 上个月也有一篇博客介绍他们对权重更新的优化。

另外Chayenne的这篇博客也对权重更新的几种主流的方案做了一个很好的总结：

update_weights_from_disk：写磁盘读磁盘，简单透明，性能取决于存储系统的快慢。
update_weights_from_distributed：把参数聚合到训练节点的 rank 0，然后发到推理节点的 rank 0，然后再从推理节点的 rank 0 发到推理节点的其他 rank。
update_weights_from_tensor：类似 update_weights_from_distributed，对训推一体有特殊优化（比如传递 Cuda IPC Handle），但是对推理引擎需要有侵入式修改。

一些猜测

没跑过现有的框架，没做过性能测试，没仔细读过代码，所以这里我就开始胡说八道了：

只依赖 rank 0 进行训练节点和推理节点之间的通信，那么 rank 0 GPU 就是整个通信的瓶颈。GPU 和网卡通过 PCIe 连接，所以再快也不可能快过 PCIe 的带宽。因此整机的通信速度就被卡死在了单卡的 400 Gbps (50 GB/s)，而不是 3200 Gbps (400 GB/s)。
RPC 额外开销太大。如果每传一个 tensor 都需要 RPC 调用一次推理引擎的方法，那么这些控制平面的消息传递本身就会花很多时间。我猜 RPC 是跑在 TCP 上的。另外，序列化反序列化也是需要时间的。
没有充分利用并行和流水线。每个 tensor 的更新是可以按照硬件资源的不同划分成不同阶段的。不同的硬件资源可以并行地进行不同的阶段。比如 GPU 操作和 RDMA 传输可以并行。但是因为担心爆显存，很多时候不得不一个 tensor 一个 tensor 地传（或者一层一层地传），这就导致不能充分利用并行和流水线。
重复计算及重复传递传输方案。
工程实现上将太多步骤耦合在了一起（例如参数匹配、传输方案计算、投影融合、量化、机内传输、跨机传输、显存用量控制），在一堆乱麻中可能很难理清思绪。

而我上面提到的方案就像把冰箱门打开，把大象装进去，然后关上冰箱门，因为：

训练节点的每张卡都可以向外传参数，推理节点的每张卡都可以接收参数。因此所有的 RDMA 带宽都能用上。
推理框架不需要修改，不需要任何控制平面的 RPC 消息。推理节点甚至都不知道自己的参数被改了。
每个 tensor 的传输任务划分成了3个阶段：提交 GPU 操作，等待 GPU 操作完成并提交 RDMA 传输，等待 RDMA 传输完成。每个任务只要完成了上一阶段，马上可以进入下一个阶段。只要硬件腾出资源，就可以开启下一个 tensor 的传输。
启动训练的时候计算一次传输方案并下发，之后每次需要更新权重的时候只需要执行传输方案就行了。
传输方案的计算以及传输方案的执行各自都分成了几个阶段，方便实现，也方便写单元测试。

那既然用 RDMA 点对点传输这么简单方便，为什么现有框架不这么做呢？我猜测是因为没有好用的 API。

大家手头上能用的 API 就是 torch.distributed，而 torch.distributed 底下跑的是 NCCL。虽然 NCCL 和 torch.distributed 也支持点对点通信，但是非常不灵活：

发送方 send(tensor) 和接收方 recv(tensor) 需要给定同样类型、同样大小的 tensor。
传输是阻塞的，也就是发送方 send(t2) 需要等待 send(t1) 完成，接收方同理。

然而 RDMA 的点对点传输非常灵活且高效：

接收方无需任何操作。接收方甚至都不知道自己的内存被改了。
传输的大小及目的地内存地址都是可以灵活指定的。
传输是异步的，不会阻碍发送方或接收方。
所有的数据传输都不会经过 Linux 内核。用户态程序直接向网卡提交 RDMA 操作。网卡可以直接从内存或者显存中读取数据。

你需要的或许是一个 RDMA 通信库？

经过一段时间的探索，我们发现有了 RDMA 通信库之后可以很方便地做很多事情，比如 KvCache 传输，训练参数更新，甚至一些你想不到的事情也能做。然后我们的通信库在 AWS EFA 上也越来越稳定，性能也越来越好。最近我正在加 ConnectX-7 网卡的支持。我们很快会开源这个通信库，敬请期待。

另外，对我司感兴趣的读者，也欢迎给我们投简历。

更新

2025-09-17 更新：跨机秒传RL模型参数更新的一些后续

驾驭3200Gbps网络(15): 惰性提交操作

2024-12-31T00:00:00-08:00

在上一章中，批量提交新的 WRITE 操作，使得传输速度达到了 2589.488 Gbps，达到了总带宽 3200 Gbps 的 80.9%。现在满速 3200 Gbps 已经近在咫尺了。

注意到我们每次提交 WRITE 操作的时候，PostWrite() 函数都会调用 ProgressPendingOps()。可以想象的是，如果在同一批次的操作中，如果第一次的 ProgressPendingOps() 无法进步，那么后续的操作必然也无法进步。然而每一次的 ProgressPendingOps() 都会将操作从队头取出，调用 fi_writemsg()，再放回队头。这样的操作是非常低效的。我们可以尝试让 PostWrite() 懒一点，不要每次都调用 ProgressPendingOps()。我们把本章的程序命名为 15_lazy.cpp。

惰性提交操作

为了体现操作的惰性，我们把函数名从 PostWrite() 改为 LazyPostWrite()，并且去掉了 ProgressPendingOps() 的调用。

void Network::LazyPostWrite(
    RdmaWriteOp &&write, std::function<void(Network &, RdmaOp &)> &&callback) {
  auto *op = new RdmaOp{
      .type = RdmaOpType::kWrite,
      .write = std::move(write),
      .callback = std::move(callback),
  };
  pending_ops.push_back(op);
  // Caller needs to poll completion to progress pending ops
}

在后续的主循环中，我们会通过 PollCompletion() 调用 ProgressPendingOps()，因此我们无需做更多的修改。

运行效果

从上面的视频中可以看到，我们达到了 3108.283 Gbps，即总带宽 3200 Gbps 的 97.1%。对比第七章中单网卡的 97.433 Gbps (97.4%)，以及第三章中 libfabric 自带的 fi_rma_bw 测试工具达到的 94.751 Gbps (94.8%)，我们的程序已经非常接近满速了。

至此，我们已经榨干了 AWS p5 集群的 3200 Gbps 带宽。在这个系列文章中，我们从零开始，认识了 RDMA、EFA、libfabric，感悟了高性能网络系统设计哲学，然后编写代码，一步步地实现了 RECV、SEND、GPUDirect RDMA WRITE、操作队列，接着扩展到多网卡，再进行了一系列的优化，最终达到了 3108.283 Gbps (97.1%) 的传输速度。希望这个系列文章对大家有所帮助。

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/15_lazy.cpp

驾驭3200Gbps网络(14): 批量提交操作

2024-12-31T00:00:00-08:00

在上一章中，我们对不同线程的状态进行分片，避免了线程之间的同步，将传输速度提升到了 1522.738 Gbps，可惜依然只达到了总带宽 3200 Gbps 的 47.6%。显然 CPU 仍然是瓶颈，我们需要进一步优化。

我们在 ContinuePostWrite() 函数中提交新的 WRITE 操作时，每次只会提交一个操作。随后 CPU 的控制流就会离开这个函数，进入到其他的函数中。这样 CPU 的控制流就会频繁地在不同的函数之间切换，导致 CPU 的缓存失效率增加。我们可以尝试将多个操作一起提交，减少 CPU 的控制流切换。我们把本章的程序命名为 14_batch.cpp。

批量提交操作

本章的修改非常简单。只需要在 ContinuePostWrite() 函数中增加一个循环，每次提交16个操作。

struct RandomFillRequestState {
  // ...

  void ContinuePostWrite(size_t gpu_idx) {
    constexpr int kBatchSize = 16;
    auto &s = write_states[gpu_idx];
    if (s.i_repeat == total_repeat)
      return;
    auto page_size = request_msg->page_size;
    auto num_pages = request_msg->num_pages;
    auto &group = (*net_groups)[gpu_idx];
    for (int i = 0; i < kBatchSize; ++i) {
      auto net_idx = group.GetNext();
      group.nets[net_idx]->PostWrite(...);
      ++posted_write_ops[gpu_idx];

      if (++s.i_page == num_pages) {
        s.i_page = 0;
        if (++s.i_buf == buf_per_gpu) {
          s.i_buf = 0;
          if (++s.i_repeat == total_repeat)
            return;
        }
      }
    }
  }
};

运行效果

从上面的视频中可以看到，我们在 ContinuePostWrite() 一次性提交16个操作之后，传输速度达到了 2589.488 Gbps，达到了总带宽 3200 Gbps 的 80.9%，相比起之前的 1522.567 Gbps 几乎翻倍。现在我们已经离满速 3200 Gbps 不远了，下一章我们将进一步榨干性能。

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/14_batch.cpp

驾驭3200Gbps网络(13): 状态分片

2024-12-31T00:00:00-08:00

在上一章中，我们为多线程绑定了 CPU 核心，将传输速度提升到了 1237.738 Gbps，可惜依然只达到了总带宽 3200 Gbps 的 38.7%。显然 CPU 仍然是瓶颈，我们需要进一步优化。

回想我们把程序由单线程改为多线程的过程中，我们将一些变量改为了原子类型（Atomic Type），以避免多个线程之间对状态读写的竞争（Race Condition）。这其中包含了 posted_write_ops 和 finished_write_ops 这两个变量。因为在一秒钟内我们会提交和完成上百万个操作，所以对这两个原子变量的操作会非常频繁，有可能会成为瓶颈。

注意到在我们的程序中，每个线程负责一张显卡，他们之间的进度彼此是独立的。我们没有必要对每一个操作的提交和完成都进行原子自增操作。在本章中，我们尝试将不同线程的状态分片，使得大部分时候每个线程都不需要读写其他线程的状态。我们把本章的程序命名为 13_shard.cpp。

状态分片

我们将 posted_write_ops 和 finished_write_ops 这两个原子变量改为了数组，每个线程只负责自己的状态。另外，我们增加一个原子变量 cnt_finished_gpus，用于同步已经完成的线程数。

struct RandomFillRequestState {
  // ...
  std::array<uint64_t, 8> posted_write_ops;
  std::array<uint64_t, 8> finished_write_ops;
  std::atomic<size_t> cnt_finished_gpus = 0;
  // ...
};

在状态机转变到 kWrite 之前，我们需要对这两个数组进行初始化。

struct RandomFillRequestState {
  // ...

  void HandleWarmupCompletion(Network &net, RdmaOp &op) {
    // ...

    // Prepare RDMA WRITE the data to remote GPU.
    printf("Started RDMA WRITE to the remote GPU memory.\n");
    // ...
    std::fill(posted_write_ops.begin(), posted_write_ops.end(), 0);
    std::fill(finished_write_ops.begin(), finished_write_ops.end(), 0);
    state = State::kWrite;
  }
};

在提交新的 WRITE 操作时，我们只需要对当前线程的状态进行操作。

struct RandomFillRequestState {
  // ...

  void ContinuePostWrite(size_t gpu_idx) {
    // ...
    group.nets[net_idx]->PostWrite(
        RdmaWriteOp{ ... },
        [this](Network &net, RdmaOp &op) { HandleWriteCompletion(net, op); });
    ++posted_write_ops[gpu_idx];
    // ...
  }
};

在 WRITE 操作的完成回调中，我们则需要稍微仔细一点处理。首先我们需要增加当前线程的 finished_write_ops。如果达到了输出进度的条件，那么我们就需要读取所有线程的进度，并进行求和。因为在输出进度的时候我们不要求完全准确的数值，所以我们没有必要加入额外的同步操作，完全可以直接脏读（Dirty Read）其他线程的状态。进一步地我们可以使用 AVX2 指令集来加速求和操作。

另一方面，如果当前线程完成的操作数达到了总操作数，那么我们就需要增加 cnt_finished_gpus，并检查是否所有线程都已经完成。如果所有线程都已经完成，那么我们就可以输出最终的进度并结束程序。因为 cnt_finished_gpus 是原子变量，所以我们这里可以保证正确性。而因为只有当当前线程完成了所有操作时才会读写 cnt_finished_gpus，所以这里也不会因为原子变量的竞争而导致性能下降。

uint64_t SumU64x8AVX2(const std::array<uint64_t, 8> &arr) {
  // https://www.perplexity.ai/search/c-sum-over-a-std-array-uint64-p8.MLN_mQeOwoik79acOxg
  __m256i sum_vec = _mm256_loadu_si256((const __m256i *)arr.data());
  sum_vec = _mm256_add_epi64(
      sum_vec, _mm256_loadu_si256((const __m256i *)(arr.data() + 4)));
  __m128i sum_128 = _mm_add_epi64(_mm256_extracti128_si256(sum_vec, 0),
                                  _mm256_extracti128_si256(sum_vec, 1));
  return _mm_extract_epi64(sum_128, 0) + _mm_extract_epi64(sum_128, 1);
}

struct RandomFillRequestState {
  // ...

  void HandleWriteCompletion(Network &net, RdmaOp &op) {
    auto gpu_finished_ops = ++finished_write_ops[op.write.buf->cuda_device];
    if (gpu_finished_ops % 16384 == 0) {
      auto now = std::chrono::high_resolution_clock::now();
      auto posted = SumU64x8AVX2(posted_write_ops);
      auto finished = SumU64x8AVX2(finished_write_ops);
      PrintProgress(now, posted, finished);
    }
    if (gpu_finished_ops == total_write_ops_per_gpu) {
      if (++cnt_finished_gpus == connect_msg->num_gpus) {
        state = State::kDone;
        PrintProgress(std::chrono::high_resolution_clock::now(),
                      total_write_ops, total_write_ops);
        printf("\nFinished all RDMA WRITEs to the remote GPU memory.\n");
      }
    }
  }
};

以上就是全部的改动。我们可以看到，通过状态分片，我们避免了多个线程之间对状态读写的竞争，提升了程序的性能。我们来看一下运行效果。

运行效果

从上面的视频中可以看到，我们对不同线程的状态进行分片之后，传输速度达到了 1522.567 Gbps，达到了总带宽 3200 Gbps 的 47.6%，相比起之前的 1237.738 Gbps，提升了 23%。这是一个不错的提升，不过离我们的目标还有一半的距离，我们还要继续努力。

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/13_shard.cpp

驾驭3200Gbps网络(12): 绑定CPU核心

2024-12-31T00:00:00-08:00

在上一章中，我们启用了多线程，为每一张显卡创建了一个线程。然而传输速度依然仅为 355.301 Gbps，只达到了总带宽 3200 Gbps 的 11.1%。

如果我们仔细对比上一章程序的运行过程与再上一章单线程的版本的运行过程，我们会发现一个问题：生成随机数这一步变慢了。尽管生成随机数的用时并不被计入传输速度的测试，但是这一现象本身仍然十分奇怪。尽管使用了多个线程，但是我们仍然只使用一个线程来生成随机数，那么为什么生成随机数的速度会变慢呢？

一个可能的原因是操作系统在多个线程之间切换的时候，会把线程切换到不同的 CPU 核心上。这样一来，线程在生成随机数的时候，会在不同的 CPU 缓存之间来回切换，导致了缓存的失效。甚至有可能，线程在生成随机数的时候，会在不同的 NUMA 节点之间来回切换，导致了内存的访问延迟，以及导致了 GPU 显存的访问延迟。

既然这一问题会使得随机数生成变慢，那么同样也很可能使网络完成队列的处理以及新操作的提交变慢。为了解决这一问题，我们可以尝试将每个线程绑定到一个 CPU 核心上。这样一来，线程在生成随机数以及进行网络传输时，就不会在不同的 CPU 核心之间来回切换，也不会在不同的 NUMA 节点之间来回切换。我们把本章的程序命名为 12_pin.cpp。

绑定CPU核心

在第八章的总线拓扑探测中，我们已经把所有的 CPU 物理核心按照 NUMA 平均分配给了8张显卡。这里我们只需要对于每张显卡任意挑出一个 CPU 核心，然后把显卡对应的线程绑定到这个 CPU 核心上即可。绑定 CPU 核心可以通过 pthread_setaffinity_np() 函数来实现。

int ServerMain(int argc, char **argv) {
    // ...

    // Multi-thread Poll completions
    std::vector<std::thread> threads;
    threads.reserve(num_gpus);
    for (size_t gpu_idx = 0; gpu_idx < net_groups.size(); ++gpu_idx) {
      const auto &cpus = topo_groups[gpu_idx].cpus;
      int preferred_cpu = cpus[cpus.size() / 2];
      threads.emplace_back([&s, preferred_cpu, gpu_idx] {
        // Pin CPU
        cpu_set_t cpuset;
        CPU_ZERO(&cpuset);
        CPU_SET(preferred_cpu, &cpuset);
        CHECK(pthread_setaffinity_np(
          pthread_self(), sizeof(cpu_set_t), &cpuset) == 0);

        // Poll completions
        // ...
      });
    }

    // ...
}

运行效果

从上面的视频中可以看到，我们的程序绑定了 CPU 核心之后，传输速度达到了 1237.738 Gbps，达到了总带宽 3200 Gbps 的 38.7%。这是一次不错的提升，不过离我们的目标还有一定的距离，我们还需要进一步优化。

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/12_pin.cpp

驾驭3200Gbps网络(11): 多线程

2024-12-31T00:00:00-08:00

在上一章中，解决了测速前的预热问题，然而传输速度依然仅为 293.461 Gbps，只达到了总带宽 3200 Gbps 的 9.2%。我们可以猜测，可能是因为 CPU 无法及时处理32张网卡的完成队列以及提交新的操作，导致了网卡的空闲。在这一章中，我们尝试使用多线程来解决这个问题。我们把本章的程序命名为 11_multithread.cpp。

多线程

我们打算使用8个线程，每个线程负责1张显卡及其对应的4张网卡。

为了避免多个线程之间对状态读写的竞争（Race Condition），我们需要将服务器端状态机中的一些变量改变成原子类型（Atomic Type）。

struct RandomFillRequestState {
  // ...
  std::atomic<State> state = State::kWaitRequest;
  std::atomic<size_t> posted_warmups = 0;
  std::atomic<size_t> cnt_warmups = 0;
  std::atomic<size_t> posted_write_ops = 0;
  std::atomic<size_t> finished_write_ops = 0;
  // ...
};

在服务器端的主循环中，我们为每一张显卡创建一个线程。

int ServerMain(int argc, char **argv) {
  // ...

  // Loop forever. Accept one client at a time.
  for (;;) {
    printf("------\n");
    // State machine
    RandomFillRequestState s(&nets, &net_groups, &cuda_bufs);
    // RECV for CONNECT
    nets[0].PostRecv(buf1, [&s](Network &net, RdmaOp &op) { s.OnRecv(net, op); });
    // RECV for RandomFillRequest
    nets[0].PostRecv(buf2, [&s](Network &net, RdmaOp &op) { s.OnRecv(net, op); });
    // Multi-thread Poll completions
    std::vector<std::thread> threads;
    threads.reserve(num_gpus);
    for (size_t gpu_idx = 0; gpu_idx < net_groups.size(); ++gpu_idx) {
      // Start a thread for each GPU
      threads.emplace_back([&s, gpu_idx] {
        auto nets = (*s.net_groups)[gpu_idx].nets;
        while (s.state != RandomFillRequestState::State::kDone) {
          for (auto *net : nets) {
            net->PollCompletion();
          }
          switch (s.state) {
          case RandomFillRequestState::State::kWaitRequest:
            break;
          case RandomFillRequestState::State::kPostWarmup:
            s.PostWarmup(gpu_idx);
            break;
          case RandomFillRequestState::State::kWaitWarmup:
            break;
          case RandomFillRequestState::State::kWrite:
            s.ContinuePostWrite(gpu_idx);
            break;
          case RandomFillRequestState::State::kDone:
            break;
          }
        }
      });
    }
    for (auto &t : threads) {
      t.join();
    }
  }

  return 0;
}

以上便是全部的修改。

运行效果

从上面的视频中可以看到，我们的程序在多线程的情况下，传输速度达到了 355.301 Gbps，占用了总带宽的 11.1%。比起单线程的 293.461 Gbps，提升了 21%。不过这个速度还是远远低于我们的预期，我们还要继续努力。

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/11_multithread.cpp

驾驭3200Gbps网络(10): 测试前预热

2024-12-31T00:00:00-08:00

在上一章中，我们确实能够使用8张显卡32张网卡，然而传输速度仅为 287.089 Gbps，只达到了总带宽 3200 Gbps 的 9.0%。

如果我们仔细观察上一章程序的运行过程，我们会看到程序一开始的传输速度在 100 Gbps 左右，然后迅速升高至 260 Gbps 左右，然后逐渐缓慢上升至 287 Gbps。这一现象十分可疑。这说明了程序一开始的时候受到了一个比较大的延迟，随后便能够一直保持相同的传输速度。这一章让我们来解决这个慢速启动的问题。我们把本章的程序命名为 10_warmup.cpp。

我们很容易能够猜到，这个现象可能是因为第一个 WRITE 操作花费了特别长的时间导致的。而之所以第一个操作花费了特别长的时间，是因为除了第一张网卡以外，其他的网卡都还没有与客户端建立连接。

预热

要解决这个问题也很简单，在正式进入速度测试之前，我们可以引入一个预热阶段。在预热阶段，我们让每一个网卡都向对方发送一个 WRITE 操作，这样就确保了两端的连接都已经建立。在预热阶段结束后，我们再开始正式的速度测试。

我们在服务器端的状态机里面加入一些关于预热的状态：

struct RandomFillRequestState {
  enum class State {
    kWaitRequest,
    kPostWarmup,  // Added
    kWaitWarmup,  // Added
    kWrite,
    kDone,
  };

  struct WriteState {
    bool warmup_posted = false;  // Added
    size_t i_repeat = 0;
    size_t i_buf = 0;
    size_t i_page = 0;
  };

  // ...
  size_t posted_warmups = 0;
  size_t cnt_warmups = 0;
  // ...
};

当服务器端收到 RANDOM_FILL 请求时，我们进入 kPostWarmup 状态：

struct RandomFillRequestState {
  // ...

  void HandleRequest(Network &net, RdmaOp &op) {
    // ...
    // Generate random data and copy to local GPU memory
    // ...

    // Prepare for warmup
    write_states.resize(connect_msg->num_gpus);
    state = State::kPostWarmup;
  }
};

然后我们增加一个 PostWarmup(gpu_idx) 函数，为这张显卡对应的所有网卡提交一个 WRITE 操作。当所有显卡都提交完了预热操作之后，我们让状态机进入 kWaitWarmup 状态：

struct RandomFillRequestState {
  // ...

  void PostWarmup(size_t gpu_idx) {
    // Warmup the connection.
    // Write 1 page via each network
    auto &s = write_states[gpu_idx];
    if (s.warmup_posted) {
      return;
    }

    auto page_size = request_msg->page_size;
    auto &group = (*net_groups)[gpu_idx];
    for (size_t k = 0; k < group.nets.size(); ++k) {
      auto net_idx = group.GetNext();
      const auto &mr =
          connect_msg->mr((gpu_idx * nets_per_gpu + net_idx) * buf_per_gpu);
      auto write = RdmaWriteOp{ ... };
      group.nets[net_idx]->PostWrite(std::move(write),
                                     [this](Network &net, RdmaOp &op) {
                                       HandleWarmupCompletion(net, op);
                                     });
    }
    s.warmup_posted = true;
    if (++posted_warmups == connect_msg->num_gpus) {
      state = State::kWaitWarmup;
    }
  }
};

然后在预热操作的回调函数中，我们检查是否所有的预热操作都已经完成。如果是，我们就进入 kWrite 状态：

struct RandomFillRequestState {
  // ...

  void HandleWarmupCompletion(Network &net, RdmaOp &op) {
    if (++cnt_warmups < connect_msg->num_nets) {
      return;
    }
    printf("Warmup completed.\n");

    // Prepare RDMA WRITE the data to remote GPU.
    printf("Started RDMA WRITE to the remote GPU memory.\n");
    total_write_ops = connect_msg->num_gpus * buf_per_gpu *
                      request_msg->num_pages * total_repeat;
    write_op_size = request_msg->page_size;
    write_states.resize(connect_msg->num_gpus);
    write_start_at = std::chrono::high_resolution_clock::now();
    state = State::kWrite;
  }
};

最后就是在服务器端的主循环中，我们判断如果状态机处于 kPostWarmup 状态，我们就调用 PostWarmup(gpu_idx) 函数：

int ServerMain(int argc, char **argv) {
    // ...
    while (s.state != RandomFillRequestState::State::kDone) {
      for (size_t gpu_idx = 0; gpu_idx < net_groups.size(); ++gpu_idx) {
        for (auto *net : net_groups[gpu_idx].nets) {
          net->PollCompletion();
        }
        switch (s.state) {
        case RandomFillRequestState::State::kWaitRequest:
          break;
        case RandomFillRequestState::State::kPostWarmup:  // Added
          s.PostWarmup(gpu_idx);
        case RandomFillRequestState::State::kWaitWarmup:  // Added
          break;
        case RandomFillRequestState::State::kWrite:
          s.ContinuePostWrite(gpu_idx);
          break;
        case RandomFillRequestState::State::kDone:
          break;
        }
      }
    }
    // ...
}

运行效果

在上面的视频中我们可以看到，在加入了预热之后，程序一开始的传输速度就能够达到 290 Gbps 左右，并且之后一直稳定在这个速度附近。最终的传输速度为 293.461 Gbps，达到了总带宽 3200 Gbps 的 9.2%。我们会继续在下一章中优化这个程序。

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/10_warmup.cpp

驾驭3200Gbps网络(9): 使用32张网卡

2024-12-25T00:00:00-08:00

在上一章中，我们已经搞清楚了系统的拓扑。而在第七章中我们已经能够在单张网卡上达到 97.844 Gbps 带宽。那么在这一章里，我们将使用8张显卡对应的32张网卡，看看我们的程序能达到多少带宽。我们把这个程序命名为 9_multinet.cpp。

网络

我们首先对 Network 类进行一个微小的修改。一是我们记录下来网卡对应的显卡 cuda_device。二是我们在打开网络的时候，让所有的网卡都共享同一个 fabric 对象。

struct Network {
  // ...
  int cuda_device;
};

Network Network::Open(struct fi_info *fi, int cuda_device,
                      struct fid_fabric *fabric) {
  if (!fabric) {
    FI_CHECK(fi_fabric(fi->fabric_attr, &fabric, nullptr));
  }

  // ...

  return Network{fi, fabric, domain, cq, av, ep, addr, cuda_device};
}

负载均衡

因为我们所有 WRITE 操作的大小都是相同的，所以负载均衡非常容易，只需要轮流使用显卡对应的四张网卡即可。为此，我们可以实现一个简单的 NetworkGroup 类，保存四张网卡的信息，并且使用 Round-Robin 来选择下一张网卡。

struct NetworkGroup {
  std::vector<Network *> nets;
  uint8_t rr_mask;
  uint8_t rr_idx = 0;

  NetworkGroup(std::vector<Network *> &&nets) {
    CHECK(nets.size() <= kMaxNetworksPerGroup);
    CHECK((nets.size() & (nets.size() - 1)) == 0); // power of 2
    this->rr_mask = nets.size() - 1;
    this->nets = std::move(nets);
  }
  NetworkGroup(const NetworkGroup &) = delete;
  NetworkGroup(NetworkGroup &&) = default;

  uint8_t GetNext() {
    rr_idx = (rr_idx + 1) & rr_mask;
    return rr_idx;
  }
};

`CONNECT` 消息

在之前的程序中，因为只使用了一张网卡，所以客户端只需要向服务器端发送一个地址即可。而现在，因为我们需要使用多张网卡，所以我们在 CONNECT 消息中需要发送显卡的数量、网卡的数量、以及每张网卡的地址和内存区域。和之前一样，我们将定长的数据放在结构体当中，而变长的数据放在结构体之后的内存中。

struct AppConnectMessage {
  AppMessageBase base;
  size_t num_gpus;
  size_t num_nets;
  size_t num_mr;

  EfaAddress &net_addr(size_t index) {
    CHECK(index < num_nets);
    return ((EfaAddress *)((uintptr_t)&base + sizeof(*this)))[index];
  }

  MemoryRegion &mr(size_t index) {
    CHECK(index < num_mr);
    return ((MemoryRegion *)((uintptr_t)&base + sizeof(*this) +
                             num_nets * sizeof(EfaAddress)))[index];
  }

  size_t MessageBytes() const {
    return sizeof(*this) + num_nets * sizeof(EfaAddress) +
           num_mr * sizeof(MemoryRegion);
  }
};

服务器端逻辑

从一张网卡变成8张显卡32张网卡，服务器端的状态机的一些变量需要由单个变量变成数组。

constexpr size_t kMaxNetworksPerGroup = 4;

struct RandomFillRequestState {
  std::vector<Network> *nets;
  std::vector<NetworkGroup> *net_groups;
  std::vector<Buffer> *cuda_bufs;
  std::vector<std::array<fi_addr_t, kMaxNetworksPerGroup>> remote_addrs;
  std::vector<WriteState> write_states;
  // ...
};

当收到 CONNECT 消息时，我们需要让服务器端和客户端的网卡一一对应起来，增加到对应的地址向量中。

struct RandomFillRequestState {
  // ...

  void HandleConnect(Network &net, RdmaOp &op) {
    // ...

    // Assuming remote has the same number of GPUs and NICs.
    CHECK(msg.num_gpus == cuda_bufs->size());
    CHECK(msg.num_nets == nets->size());

    // Add peer addresses
    nets_per_gpu = msg.num_nets / msg.num_gpus;
    buf_per_gpu = connect_msg->num_mr / connect_msg->num_nets;
    for (size_t i = 0; i < msg.num_gpus; ++i) {
      std::array<fi_addr_t, kMaxNetworksPerGroup> addrs = {};
      for (size_t j = 0; j < nets_per_gpu; ++j) {
        auto idx = i * nets_per_gpu + j;
        addrs[j] = nets->at(idx).AddPeerAddress(msg.net_addr(idx));
      }
      remote_addrs.push_back(addrs);
    }

    // Initialize write states
    write_states.resize(connect_msg->num_gpus);
  }
};

我们也需要对 ContinuePostWrite() 函数进行一些修改。我们传入一个 gpu_idx 参数，表示只提交这个显卡对应的网卡的 WRITE 操作。首先是先用 NetworkGroup 选择下一个网卡。然后我们要找到出来这个网卡对应的客户端地址 dest_addr、内存区域地址 mr.addr、以及内存区域的远程访问密钥 mr.rkey。在设置立即数上面，我们也要注意判断是否是当前网卡的最后一个 WRITE 操作（s.i_page + nets_per_gpu >= num_pages）。

struct RandomFillRequestState {
  // ...

  void ContinuePostWrite(size_t gpu_idx) {
    auto &s = write_states[gpu_idx];
    if (s.i_repeat == total_repeat)
      return;
    auto page_size = request_msg->page_size;
    auto num_pages = request_msg->num_pages;

    auto net_idx = (*net_groups)[gpu_idx].GetNext();
    uint32_t imm_data = 0;
    if (s.i_repeat + 1 == total_repeat && s.i_buf + 1 == buf_per_gpu &&
        s.i_page + nets_per_gpu >= num_pages) {
      // The last WRITE. Pass remote context back.
      imm_data = request_msg->remote_context;
    }
    const auto &mr = connect_msg->mr(
        (gpu_idx * nets_per_gpu + net_idx) * buf_per_gpu + s.i_buf);
    (*net_groups)[gpu_idx].nets[net_idx]->PostWrite(
        RdmaWriteOp{.buf = &(*cuda_bufs)[gpu_idx],
         .offset = s.i_buf * (page_size * num_pages) + s.i_page * page_size,
         .len = page_size,
         .imm_data = imm_data,
         .dest_ptr = mr.addr + request_msg->page_idx(s.i_page) * page_size,
         .dest_addr = remote_addrs[gpu_idx][net_idx],
         .dest_key = mr.rkey},
        [this](Network &net, RdmaOp &op) { HandleWriteCompletion(); });
    ++posted_write_ops;

    if (++s.i_page == num_pages) {
      s.i_page = 0;
      if (++s.i_buf == buf_per_gpu) {
        s.i_buf = 0;
        if (++s.i_repeat == total_repeat)
          return;
      }
    }
  }
};

在服务器端的主程序，我们首先检测系统拓扑，然后打开所有的网卡，并分组建立网络组。

int ServerMain(int argc, char **argv) {
  // Topology detection
  struct fi_info *info = GetInfo();
  auto topo_groups = DetectTopo(info);
  int num_gpus = topo_groups.size();
  int num_nets = topo_groups[0].fi_infos.size() * topo_groups.size();
  int nets_per_gpu = num_nets / num_gpus;

  // Open Netowrk
  std::vector<Network> nets;
  std::vector<NetworkGroup> net_groups;
  for (int cuda_device = 0; cuda_device < num_gpus; ++cuda_device) {
    std::vector<Network *> group_nets;
    for (auto *fi : topo_groups[cuda_device].fi_infos) {
      int cuda_device = nets.size() / nets_per_gpu;
      auto *fabric = nets.empty() ? nullptr : nets[0].fabric;
      nets.push_back(Network::Open(fi, cuda_device, fabric));
      group_nets.push_back(&nets.back());
    }
    net_groups.push_back(NetworkGroup(std::move(group_nets)));
  }
  PrintTopologyGroups(topo_groups);

  // ...
}

然后我们在每张显卡上都分配一个缓冲区。同时，向这张显卡对应的所有的网卡注册这个内存区域。

  // Allocate and register CUDA memory
  printf("Registered MR from");
  std::vector<Buffer> cuda_bufs;
  for (int i = 0; i < num_gpus; ++i) {
    CUDA_CHECK(cudaSetDevice(i));
    cuda_bufs.push_back(Buffer::AllocCuda(kMemoryRegionSize * 2, kBufAlign));
    for (int j = 0; j < nets_per_gpu; ++j) {
      nets[i * nets_per_gpu + j].RegisterMemory(cuda_bufs.back());
    }
    printf(" cuda:%d", i);
    fflush(stdout);
  }
  printf("\n");

我们打算使用第一张网卡来接收客户端连接，所以我们分配两个缓冲区并在第一张网卡上注册内存区域。

  // Allocate and register message buffer
  auto buf1 = Buffer::Alloc(kMessageBufferSize, kBufAlign);
  auto buf2 = Buffer::Alloc(kMessageBufferSize, kBufAlign);
  nets[0].RegisterMemory(buf1);
  nets[0].RegisterMemory(buf2);

在主循环中，我们先在第一张网卡上提交两个 RECV 操作。在状态机未完成前，我们每次遍历每一个显卡。对于一张显卡，我们先处理其对应的所有的网卡的完成队列。然后如果状态机的状态处于 kWrite，则继续提交更多的 WRITE 操作。

  // Loop forever. Accept one client at a time.
  for (;;) {
    printf("------\n");
    // State machine
    RandomFillRequestState s(&nets, &net_groups, &cuda_bufs);
    // RECV for CONNECT
    nets[0].PostRecv(buf1, [&s](Network &net, RdmaOp &op) { s.OnRecv(net, op); });
    // RECV for RandomFillRequest
    nets[0].PostRecv(buf2, [&s](Network &net, RdmaOp &op) { s.OnRecv(net, op); });
    // Wait for completion
    while (s.state != RandomFillRequestState::State::kDone) {
      for (size_t gpu_idx = 0; gpu_idx < net_groups.size(); ++gpu_idx) {
        for (auto *net : net_groups[gpu_idx].nets) {
          net->PollCompletion();
        }
        switch (s.state) {
        case RandomFillRequestState::State::kWaitRequest:
          break;
        case RandomFillRequestState::State::kWrite:
          s.ContinuePostWrite(gpu_idx);
          break;
        case RandomFillRequestState::State::kDone:
          break;
        }
      }
    }
  }

  return 0;

客户端逻辑

客户端逻辑的修改和服务器端类似。首先检测系统拓扑，然后打开所有的网卡，并分组建立网络组。然后在每张显卡上分配缓冲区并在其对应的网卡上注册内存区域。再在第一张网卡上分配缓冲区和注册内存区域。这里不再赘述。

然后我们需要在 CONNECT 消息中包含显卡的数量、网卡的数量、每张网卡的地址、每张网卡上的内存区域。然后使用第一张网卡向服务器端发送这个消息。

int ClientMain(int argc, char **argv) {
  // ...

  // Send address to server
  auto &connect_msg = *(AppConnectMessage *)buf1.data;
  connect_msg = {
      .base = {.type = AppMessageType::kConnect},
      .num_gpus = (size_t)num_gpus,
      .num_nets = nets.size(),
      .num_mr = nets.size() * 2,
  };
  for (size_t i = 0; i < nets.size(); i++) {
    connect_msg.net_addr(i) = nets[i].addr;
    int cuda_device = nets[i].cuda_device;
    connect_msg.mr(i * 2) = {
        .addr = (uint64_t)cuda_bufs1[cuda_device].data,
        .size = cuda_bufs1[cuda_device].size,
        .rkey = nets[i].GetMR(cuda_bufs1[cuda_device])->key,
    };
    connect_msg.mr(i * 2 + 1) = {
        .addr = (uint64_t)cuda_bufs2[cuda_device].data,
        .size = cuda_bufs2[cuda_device].size,
        .rkey = nets[i].GetMR(cuda_bufs2[cuda_device])->key,
    };
  }
  bool connect_sent = false;
  nets[0].PostSend(
      server_addr, buf1, connect_msg.MessageBytes(),
      [&connect_sent](Network &net, RdmaOp &op) { connect_sent = true; });
  while (!connect_sent) {
    nets[0].PollCompletion();
  }
}

在等待 REMOTE WRITE 回调函数的时候，之前的程序只需要等待一次回调，而现在我们需要在每张网卡上都等待一次回调。

  // Prepare to receive the last REMOTE WRITE from server
  int cnt_last_remote_write_received = 0;
  uint32_t remote_write_op_id = 0x123;
  for (auto &net : nets) {
    net.AddRemoteWrite(remote_write_op_id, [&cnt_last_remote_write_received](
                                               Network &net, RdmaOp &op) {
      ++cnt_last_remote_write_received;
    });
  }

  // Send message to server
  // ...

  // Wait for REMOTE WRITE from server
  while (cnt_last_remote_write_received != num_nets) {
    for (auto &net : nets) {
      net.PollCompletion();
    }
  }
  printf("Received RDMA WRITE to local GPU memory.\n");

到此我们就完成了所有的修改。

运行效果

在上面的视频中我们可以看到，我们确实能够使用8张显卡32张网卡，然而传输速度仅为 287.089 Gbps，只达到了总带宽 3200 Gbps 的 9.0%。接下来我们将一步一步地优化我们的程序，直到我们能打满 3200 Gbps 的带宽。

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/9_multinet.cpp

驾驭3200Gbps网络(8): 系统拓扑

2024-12-25T00:00:00-08:00

在上一章中，我们在单张网卡上达到了 97.844 Gbps 的传输速度。在第二章中我们已经介绍了，使用把显卡和临近的网卡配对在一起才能达到最佳的传输速度。在我们开始使用多张网卡之前，我们在这一章先搞清楚机器的系统拓扑。我们把这个程序命名为 8_topo.cpp。

获取系统拓扑

在 Linux 上要获取系统拓扑，可以使用 hwloc 库，其中包含了一个命令行工具 lstopo，可以生成系统拓扑的图。我们运行 lstopo，可以看到下面的拓扑图：

稍微简化一下，AWS p5 实例的系统拓扑如下图所示：

可以看到系统安装了两块 CPU，每块 CPU 下面连接着4个 PCIe 交换机。每个 PCIe 交换机下面挂载了4张 100 Gbps 的 EFA 网卡，1张 NVIDIA H100 显卡，以及一块 3.84 TB 的 NVMe SSD。

另一种获得系统拓扑的方法是使用 lspci -tv 命令。我们可以看到每个 PCIe 交换机的地址，以及每个 PCIe 交换机下面挂载的设备。例如：

-+-[0000:bb]---00.0-[bc-cb]----00.0-[bd-cb]--+-01.0-[c6]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
 |                                           +-01.1-[c7]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
 |                                           +-01.2-[c8]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
 |                                           +-01.3-[c9]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
 |                                           +-01.4-[ca]----00.0  NVIDIA Corporation GH100 [H100 SXM5 80GB]
 |                                           \-01.5-[cb]----00.0  Amazon.com, Inc. NVMe SSD Controller
 +-[0000:aa]---00.0-[ab-ba]----00.0-[ac-ba]--+-01.0-[b5]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
 |                                           +-01.1-[b6]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
 |                                           +-01.2-[b7]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
 |                                           +-01.3-[b8]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
 |                                           +-01.4-[b9]----00.0  NVIDIA Corporation GH100 [H100 SXM5 80GB]
 |                                           \-01.5-[ba]----00.0  Amazon.com, Inc. NVMe SSD Controller
...

在 Linux 上，PCI 地址具有如下格式：

XXXX:BB:DD.F

XXXX 是 PCI 域（Domain）的编号，是一个用4个字符的十六进制数表示的16位整数，通常是 0000。
BB 是总线（Bus）的编号，是一个用2个字符的十六进制数表示的8位整数。
DD 是设备（Device）的编号，是一个用2个字符的十六进制数表示的5位整数。
F 是功能（Function）的编号，是一个用1个字符的十六进制数表示的3位整数。

要在程序中获得系统拓扑，最常见的办法就是使用 hwloc库，这也是 aws-ofi-nccl 所使用的方法。但是因为我不想引入额外的依赖，所以我选择直接读取 Linux 的 /sys 文件系统。

Linux sysfs

在 /sys/bus/pci/devices/ 目录下，列出了所有的 PCI 设备。每一个 PCI 设备都是一个软链接，指向 /sys/devices/ 中的路径，而这个路径本身就反映了设备的拓扑结构。例如：

$ ls -la /sys/bus/pci/devices/
...
b5:00.0 -> ../../../devices/pci0000:aa/0000:aa:00.0/0000:ab:00.0/0000:ac:01.0/0000:b5:00.0
b6:00.0 -> ../../../devices/pci0000:aa/0000:aa:00.0/0000:ab:00.0/0000:ac:01.1/0000:b6:00.0
b7:00.0 -> ../../../devices/pci0000:aa/0000:aa:00.0/0000:ab:00.0/0000:ac:01.2/0000:b7:00.0
b8:00.0 -> ../../../devices/pci0000:aa/0000:aa:00.0/0000:ab:00.0/0000:ac:01.3/0000:b8:00.0
b9:00.0 -> ../../../devices/pci0000:aa/0000:aa:00.0/0000:ab:00.0/0000:ac:01.4/0000:b9:00.0
ba:00.0 -> ../../../devices/pci0000:aa/0000:aa:00.0/0000:ab:00.0/0000:ac:01.5/0000:ba:00.0
...
c6:00.0 -> ../../../devices/pci0000:bb/0000:bb:00.0/0000:bc:00.0/0000:bd:01.0/0000:c6:00.0
c7:00.0 -> ../../../devices/pci0000:bb/0000:bb:00.0/0000:bc:00.0/0000:bd:01.1/0000:c7:00.0
c8:00.0 -> ../../../devices/pci0000:bb/0000:bb:00.0/0000:bc:00.0/0000:bd:01.2/0000:c8:00.0
c9:00.0 -> ../../../devices/pci0000:bb/0000:bb:00.0/0000:bc:00.0/0000:bd:01.3/0000:c9:00.0
ca:00.0 -> ../../../devices/pci0000:bb/0000:bb:00.0/0000:bc:00.0/0000:bd:01.4/0000:ca:00.0
cb:00.0 -> ../../../devices/pci0000:bb/0000:bb:00.0/0000:bc:00.0/0000:bd:01.5/0000:cb:00.0
...

对比前面的 lspci -tv 结果可以看到，显卡 [ca] 和4张网卡 [c6] [c7] [c8] [c9] 都挂载在 PCIe 交换机 [bd] 下面。通过这个软链接结构，我们就可以构建出系统的拓扑结构。

获取显卡的 PCI 地址

通过 cudaGetDeviceProperties() 函数，我们可以获得显卡的 PCI 地址。

int num_gpus = 0;
CUDA_CHECK(cudaGetDeviceCount(&num_gpus));
for (int cuda_device = 0; cuda_devicei < num_gpus; ++cuda_device) {
  char pci_addr[16];
  cudaDeviceProp prop;
  CUDA_CHECK(cudaGetDeviceProperties(&prop, cuda_device));
  snprintf(pci_addr, sizeof(pci_addr), "%04x:%02x:%02x.0",
            prop.pciDomainID, prop.pciBusID, prop.pciDeviceID);
  // ...
}

获取网卡的 PCI 地址

在第三章末尾的 fi_info 输出中，我们可以看到 fi_info 结构体中有网卡的 PCI 地址。

struct fi_info *info = GetInfo();
for (auto* fi = info; fi; fi = fi->next) {
  char pci_addr[16];
  snprintf(pci_addr, sizeof(pci_addr), "%04x:%02x:%02x.%d",
            fi->nic->bus_attr->attr.pci.domain_id,
            fi->nic->bus_attr->attr.pci.bus_id,
            fi->nic->bus_attr->attr.pci.device_id,
            fi->nic->bus_attr->attr.pci.function_id);
  // ...
}

获取 CPU 拓扑

NUMA 节点信息

当机器上有多个 NUMA 节点时，编写高性能程序的时候经常要考虑到 NUMA 问题。Linux 系统中每个 NUMA 节点的信息存在下面这个正则表达式匹配的文件夹中：

/sys/devices/system/node/node[0-9]+/

例如：

$ ls -la /sys/devices/system/node/
total 0
drwxr-xr-x  5 root root    0 Oct 14 04:36 .
drwxr-xr-x 10 root root    0 Oct 14 04:36 ..
-r--r--r--  1 root root 4096 Dec 31 22:04 has_cpu
-r--r--r--  1 root root 4096 Dec 31 22:04 has_generic_initiator
-r--r--r--  1 root root 4096 Dec 31 22:04 has_memory
-r--r--r--  1 root root 4096 Dec 31 22:04 has_normal_memory
drwxr-xr-x  4 root root    0 Oct 14 04:36 node0
drwxr-xr-x  4 root root    0 Oct 14 04:36 node1
-r--r--r--  1 root root 4096 Dec  3 02:39 online
-r--r--r--  1 root root 4096 Dec 31 22:04 possible
drwxr-xr-x  2 root root    0 Dec 31 20:45 power
-rw-r--r--  1 root root 4096 Oct 14 04:36 uevent

这里我们可以看到系统中有2个 NUMA 节点。

CPU 核心信息

每个 NUMA 节点下面的 CPU 编号存在下面这个正则表达式匹配的文件夹中：

/sys/devices/system/node/node[0-9]+/cpu[0-9]+/

例如：

$ ls -la /sys/devices/system/node/node1/cpu*          
lrwxrwxrwx 1 root root     0 Dec 31 20:45 /sys/devices/system/node/node1/cpu48 -> ../../cpu/cpu48
lrwxrwxrwx 1 root root     0 Dec 31 20:45 /sys/devices/system/node/node1/cpu49 -> ../../cpu/cpu49
lrwxrwxrwx 1 root root     0 Dec 31 20:45 /sys/devices/system/node/node1/cpu50 -> ../../cpu/cpu50
lrwxrwxrwx 1 root root     0 Dec 31 20:45 /sys/devices/system/node/node1/cpu51 -> ../../cpu/cpu51
lrwxrwxrwx 1 root root     0 Dec 31 20:45 /sys/devices/system/node/node1/cpu52 -> ../../cpu/cpu52
...
-r--r--r-- 1 root root 28672 Dec 31 22:07 /sys/devices/system/node/node1/cpulist
-r--r--r-- 1 root root  4096 Dec  3 02:39 /sys/devices/system/node/node1/cpumap

CPU 超线程信息

对很多高性能程序来说，使用超线程反而会降低性能。我们这里可以进一步地筛选出 CPU 物理核心的编号。下面这个路径中的文件包含了与这个 CPU 核心共享同一个物理核心的所有逻辑核心的编号：

/sys/devices/system/node/node[0-9]+/cpu[0-9]+/topology/thread_siblings_list

如果 thread_siblings_list 中只有一个编号，那么这个 CPU 核心就是物理核心，例如：

$ cat /sys/devices/system/node/node1/cpu89/topology/thread_siblings_list 
89

因为 AWS p5 实例默认关闭了超线程，所以在这里我们看到 thread_siblings_list 中只有一个编号。

PCI 设备 NUMA 节点信息

PCI 设备挂载在哪个 NUMA 节点上，可以通过下面这个路径中的文件获得：

/sys/bus/pci/devices/XXXX:BB:DD.F/numa_node

例如:

$ cat '/sys/bus/pci/devices/0000:53:00.0/numa_node' 
0
$ cat '/sys/bus/pci/devices/0000:ca:00.0/numa_node' 
1

程序接口

struct TopologyGroup {
  int cuda_device;
  int numa;
  std::vector<struct fi_info *> fi_infos;
  std::vector<int> cpus;
};

std::vector<TopologyGroup> DetectTopo(struct fi_info *info);

我们编写一个函数 DetectTopo() 来获取系统拓扑，对 GPU、网卡、CPU 物理核心进行分组。这个函数对每个 GPU 返回一个 TopologyGroup 结构体。结构体中包含了 GPU 的编号、NUMA 节点编号、网卡的 fi_info 结构体指针、CPU 核心编号。我们把在同一个 NUMA 节点下的所有 CPU 物理核心平均分配给每个 GPU。因为这个代码用 C++ 编写过于丑陋，所以我就不在文章中展示了。

运行效果

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/8_topo.cpp

驾驭3200Gbps网络(7): 操作队列及带宽测试

2024-12-25T00:00:00-08:00

在上一章中，我们实现了 GPUDirect RDMA WRITE。但是在上一章的程序中，我们只是一次性提交了16个 WRITE 操作。如果我们传输更多数据，因为网络的拥塞，fi_writemsg() 就会返回 -EAGAIN，而我们的程序并没有处理这种情况。在这一章里，我们将实现一个操作队列，将暂时无法提交的操作放到队列里，等待网络拥塞消失后再提交。与此同时，我们可以顺便测试一下我们的程序能达到多少带宽。我们把这个程序命名为 7_queue.cpp。

操作队列

首先我们在 Network 类中添加一个操作队列，保存所有未完成的操作。

struct Network {
  // ...
  std::deque<RdmaOp *> pending_ops;
};

接下来，我们增加一个函数用来尽可能多地提交操作，直到提交完所有操作或者遇到 EAGAIN。

void Network::ProgressPendingOps() {
  while (!pending_ops.empty()) {
    auto *op = pending_ops.front();
    pending_ops.pop_front();
    ssize_t ret = 0;
    switch (op->type) {
    case RdmaOpType::kRecv:
      ret = fi_recvmsg(...);
      break;
    case RdmaOpType::kSend:
      ret = fi_sendmsg(...);
      break;
    case RdmaOpType::kWrite:
      ret = fi_writemsg(...);
      break;
    case RdmaOpType::kRemoteWrite:
      CHECK(false); // Unreachable
      break;
    }
    if (ret == -FI_EAGAIN) {
      // Put it back to the front of the queue
      pending_ops.push_front(op);
      break;
    }
    if (ret) {
      // Unexpected error. Don't put it back.
      // Delete the op since it's not going to be in the completion queue.
      delete op;
      fprintf(stderr, "Failed to ProgressPendingOps. ret=%ld (%s)\n",
              ret, fi_strerror(-ret));
      fflush(stderr);
      break;
    }
  }
}

因为提交操作的主体（即 fi_{recv,send,write}msg()）已经转移到了 ProgressPendingOps() 函数中，所以我们的 Post{Recv,Send,Write}() 函数要做的事情就很简单了，只需要把操作放到队列里，然后调用 ProgressPendingOps()。

void Network::PostRecv(...) {
  auto *op = new RdmaOp{ ... };
  pending_ops.push_back(op);
  ProgressPendingOps();
}

void Network::PostSend(...) {
  auto *op = new RdmaOp{ ... };
  pending_ops.push_back(op);
  ProgressPendingOps();
}

void Network::PostWrite(...) {
  auto *op = new RdmaOp{ ... };
  pending_ops.push_back(op);
  ProgressPendingOps();
}

最后，我们还需要修改 PollCompletion() 函数。当我们收到了一些完成事件时，这意味着网络可能腾出了一些空间，我们可以尝试提交更多操作。因此我们在处理完完成队列之后，再次调用 ProgressPendingOps()。

void Network::PollCompletion() {
  // Process completions
  struct fi_cq_data_entry cqe[kCompletionQueueReadCount];
  for (;;) {
    auto ret = fi_cq_read(cq, cqe, kCompletionQueueReadCount);
    // ...
  }

  // Try to make progress.
  ProgressPendingOps();
}

以上就是对 Network 类的全部修改。加上这些修改之后，我们的网络库就可以处理网络拥塞的情况了。

服务器端逻辑

在之前的程序里面，服务器端会一口气将所有的 WRITE 操作提交出去，再回到主循环中处理完成队列。这样的逻辑在网络拥塞的情况下是不可行的。当 EAGAIN 出现的时候，我们必须等待之前的一些操作完成，并且处理完成队列。如果我们不处理完成队列，哪怕网络拥塞消失了，我们也无法提交新的操作。因此我们需要修改服务器端的逻辑，交替提交新的操作和处理完成队列。

为了让我们的带宽测试时间更长，我们将每个 WRITE 操作重复500次。

现在，让我们开始修改服务器端的状态机。首先让我们重新定义状态机结构体中的成员。我们需要增加一个 State 枚举类型表示状态机的状态；增加一个 WriteState 结构体以保存循环变量；以及我们增加一些成员变量用记录带宽测试的开始时间和当前进度。

struct RandomFillRequestState {
  enum class State {
    kWaitRequest,
    kWrite,
    kDone,
  };

  struct WriteState {
    size_t i_repeat;
    size_t i_buf;
    size_t i_page;
  };

  Network *net;
  Buffer *cuda_buf;
  size_t total_bw = 0;
  State state = State::kWaitRequest;

  fi_addr_t client_addr = FI_ADDR_UNSPEC;
  AppConnectMessage *connect_msg = nullptr;
  AppRandomFillMessage *request_msg = nullptr;

  size_t total_repeat = 500;
  WriteState write_state;
  size_t total_write_ops = 0;
  size_t write_op_size = 0;
  size_t posted_write_ops = 0;
  size_t finished_write_ops = 0;
  std::chrono::time_point<std::chrono::high_resolution_clock> write_start_at;
};

当服务器端收到 RANDOM_FILL 请求时，我们不再一次性提交所有的 WRITE 操作。而是设置好相关的变量，将状态机转移到 kWrite 状态。

struct RandomFillRequestState {
  // ...

  void HandleRequest(Network &net, RdmaOp &op) {
    // ...
    // Generate random data and copy to local GPU memory
    // ...

    // Prepare RDMA WRITE the data to remote GPU.
    total_write_ops =
        connect_msg->num_mr * request_msg->num_pages * total_repeat;
    posted_write_ops = 0;
    finished_write_ops = 0;
    write_op_size = request_msg->page_size;
    write_state = {.i_repeat = 0, .i_buf = 0, .i_page = 0};
    write_start_at = std::chrono::high_resolution_clock::now();
    state = State::kWrite;
    printf("Started RDMA WRITE to the remote GPU memory.\n");
  }
};

我们把提交 WRITE 操作的代码放到一个新函数中：

struct RandomFillRequestState {
  // ...

  void ContinuePostWrite() {
    auto &s = write_state;
    if (s.i_repeat == total_repeat)
      return;
    auto page_size = request_msg->page_size;
    auto num_pages = request_msg->num_pages;

    uint32_t imm_data = 0;
    if (s.i_repeat + 1 == total_repeat && s.i_buf + 1 == connect_msg->num_mr &&
        s.i_page + 1 == num_pages) {
      // The last WRITE. Pass remote context back.
      imm_data = request_msg->remote_context;
    }
    net->PostWrite(
        RdmaWriteOp{ ... },
        [this](Network &net, RdmaOp &op) { HandleWriteCompletion(); });
    ++posted_write_ops;

    if (++s.i_page == num_pages) {
      s.i_page = 0;
      if (++s.i_buf == connect_msg->num_mr) {
        s.i_buf = 0;
        ++s.i_repeat;
      }
    }
  }
};

在 ContinuePostWrite() 中，如果还有 WRITE 操作没有提交，我们就提交一个新的 WRITE 操作。与上一章不同的是，我们对于每一个 WRITE 操作都设置了一个回调函数 HandleWriteCompletion()。在这个回调函数中，我们会输出当前的进度以及带宽。当最后一个 WRITE 操作完成时，我们将状态机转移到 kDone 状态。

struct RandomFillRequestState {
  // ...

  void PrintProgress(...) { ... }

  void HandleWriteCompletion() {
    ++finished_write_ops;
    if (finished_write_ops % 16384 == 0) {
      auto now = std::chrono::high_resolution_clock::now();
      PrintProgress(now, posted_write_ops, finished_write_ops);
    }
    if (finished_write_ops == total_write_ops) {
      auto now = std::chrono::high_resolution_clock::now();
      PrintProgress(now, posted_write_ops, finished_write_ops);
      printf("\nFinished all RDMA WRITEs to the remote GPU memory.\n");
      state = State::kDone;
    }
  }
};

最后，我们修改服务器端的主循环。除了每次处理完成队列之外，我们还要检查状态机的状态。如果状态机处于 kWrite 状态，我们就继续提交 WRITE 操作。

int ServerMain(int argc, char **argv) {
  // ...

  // Loop forever. Accept one client at a time.
  for (;;) {
    printf("------\n");
    // State machine
    RandomFillRequestState s(&net, &cuda_buf);
    // RECV for CONNECT
    net.PostRecv(buf1, [&s](Network &net, RdmaOp &op) { s.OnRecv(net, op); });
    // RECV for RandomFillRequest
    net.PostRecv(buf2, [&s](Network &net, RdmaOp &op) { s.OnRecv(net, op); });
    // Wait for completion
    while (s.state != RandomFillRequestState::State::kDone) {
      net.PollCompletion();
      switch (s.state) {
      case RandomFillRequestState::State::kWaitRequest:
        break;
      case RandomFillRequestState::State::kWrite:
        s.ContinuePostWrite();
        break;
      case RandomFillRequestState::State::kDone:
        break;
      }
    }
  }
  return 0;
}

客户端逻辑

客户端的逻辑不需要大的修改。唯一的一点小小的改动就是，我们增加 GPU 上缓冲区的大小，并且把默认的页面大小减小为 64 KiB，把默认传输的页面数量增加到1000。因为传输的页数变多了，所以 RANDOM_FILL 消息的大小也增大了，所以我们也增大了用于接收和发送的缓冲区的大小。

constexpr size_t kMessageBufferSize = 1 << 20;
constexpr size_t kMemoryRegionSize = 1UL << 30;

int ClientMain(int argc, char **argv) {
  size_t page_size = 65536;
  size_t num_pages = 1000;
  // ...
}

运行效果

可以看到，传输速度达到了 97.844 Gbps，几乎打满了带宽。

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/7_queue.cpp

驾驭3200Gbps网络(6): GPUDirect RDMA WRITE

2024-12-25T00:00:00-08:00

在上一章中，我们实现了双向的 RECV 和 SEND，这两个操作都是双侧 RDMA 操作（Two-sided RDMA）。在本章中，我们将拓展上一章的程序，实现 WRITE，即直接写入远端的内存。WRITE 是一个单侧 RDMA 操作（One-sided RDMA），无需远端 CPU 的参与。在本章中，我们直接写入到 GPU 的内存。因为从 libfabric API 的角度来看，写入到宿主机内存和写到显存是一样的，所以如果读者的需求是写入到宿主机内存，只需要将显存的地址转换为宿主机内存的地址即可。我们把本章的程序命名为 6_write.cpp。

业务逻辑

从本章开始，我们都将假设我们在构建一个具有如下业务逻辑的应用程序。

想象我们在存储一些值，每个值可以用一个页编码（Page Index）来索引。一个值被分成了 N*M 份，分散存储在 N 个 GPU 中，每个 GPU 有 M 个大缓冲区。每个值在每个缓冲区中的大小都是相同的（Page Size）。客户端向服务器发送一个请求，要求服务器端按照给定的随机种子，在指定的页上填充随机数。大致的业务逻辑如下：

def random_fill(seed: int,
                num_gpus: int,
                buf_addrs: list[int],
                page_size: int,
                page_indices: list[int]):
  bufs_per_gpu = len(buf_addrs) // num_gpus
  for gpu_idx in range(num_gpus):
    for gpu_buf_idx in range(bufs_per_gpu):
      buf_idx = gpu_idx * bufs_per_gpu + gpu_buf_idx
      rng = RandomGenerator(seed + buf_idx)
      for page_idx in page_indices:
        addr = buf_addrs[buf_idx] + page_idx * page_size
        rng.fill(addr, page_size)

在本章及后面的章节里面，我们都假设每个 GPU 上有两个缓冲区。

在本章中，我们侧重学习如何使用 RDMA WRITE，以及 GPUDirect RDMA（即网卡和 GPU 之间的直接数据传输）。因此，我们只考虑1个 GPU，以及传输非常少量的页。在后面的章节中，我们会让程序更健壮，能够传输更多的页面，以及支持多个 GPU、多个网卡。

GPU 上的缓冲区

libfabric 支持访问其他设备的内存，这一特性在 libfabric 中称作 HMEM。对于 CUDA 设备，libfabric 有两种访问方式：通过 GDRCopy 库，以及通过 Linux 内核的 dma-buf。我并没有找到关于这两者的文档，只能从 Github 上找到一些只言片语：

NCCL 在 CUDA 11.7 以上默认使用 dma-buf，对于不支持 dma-buf 的软硬件使用 GDR [1]。
libfabric 的核心开发者 Sean Hefty 提到 dma-buf 和 GDR 类似，但是用的是更上游的机制 [2]。
libfabric 的 efa Provider 在探测 CUDA 的 p2p 支持时，会优先尝试 dma-buf [3]。

我自己的使用体验也感受不出两者的性能差异。我猜测对于新的软硬件，可能可以优先使用 dma-buf。不过在本文中，我会把两者的代码路径都写出来，以便读者参考。

在前面的文章里面，我们实现了 Buffer 类，里面所使用的地址是宿主机内存的地址。为了支持 GPU 上面的内存，我们额外增加两个成员：cuda_device 和 dmabuf_fd。同时我们增加一个函数用来从显存中分配空间。

struct Buffer {
  void *data;
  size_t size;
  int cuda_device = -1;  // Added
  int dmabuf_fd = -1;    // Added

  bool is_cuda() const { return cuda_device >= 0; }

  static Buffer AllocCuda(size_t size, size_t align) {
    void *raw_data;
    struct cudaPointerAttributes attrs = {};
    CUDA_CHECK(cudaMalloc(&raw_data, size));
    CUDA_CHECK(cudaPointerGetAttributes(&attrs, raw_data));
    CHECK(attrs.type == cudaMemoryTypeDevice);
    int cuda_device = attrs.device;
    int fd = -1;
    CU_CHECK(cuMemGetHandleForAddressRange(
        &fd, (CUdeviceptr)align_up(raw_data, align), size,
        CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD, 0));
    return Buffer(raw_data, size, align, cuda_device, fd);
  }

  // ...
};

这里我们使用了 cudaMalloc() 来从显存分配空间。然后我们使用 cuMemGetHandleForAddressRange() 来获取 dmabuf_fd。

注册内存区域

在注册内存区域的时候，我们也要做出一些变化：

void Network::RegisterMemory(Buffer &buf) {
  struct fid_mr *mr;
  struct fi_mr_attr mr_attr = {
      .iov_count = 1,
      .access = FI_SEND | FI_RECV | FI_REMOTE_WRITE | FI_REMOTE_READ |
                FI_WRITE | FI_READ,
  };
  struct iovec iov = {.iov_base = buf.data, .iov_len = buf.size};
  struct fi_mr_dmabuf dmabuf = {
      .fd = buf.dmabuf_fd, .offset = 0, .len = buf.size, .base_addr = buf.data};
  uint64_t flags = 0;
  if (buf.is_cuda()) {
    mr_attr.iface = FI_HMEM_CUDA;
    mr_attr.device.cuda = buf.cuda_device;
    if (buf.dmabuf_fd != -1) {
      mr_attr.dmabuf = &dmabuf;
      flags = FI_MR_DMABUF;
    } else {
      mr_attr.mr_iov = &iov;
    }
  } else {
    mr_attr.mr_iov = &iov;
  }
  FI_CHECK(fi_mr_regattr(domain, &mr_attr, flags, &mr));
  this->mr[buf.data] = mr;
}

首先是 access。这里我们增加了本地及远程的单边 RDMA 读写权限。然后是根据不同的内存类型进行分别设置：

对于宿主机内存，我们需要使用 mr_iov。
对于使用 GDRCopy 的显存，我们使用 mr_iov，并且设置 iface 为 FI_HMEM_CUDA。
对于使用 dma-buf 的显存，我们使用 dmabuf，并且设置 iface 为 FI_HMEM_CUDA。因为 mr_iov 和 dmabuf 在同一个 union 里面，所以我们需要设置 flags 为 FI_MR_DMABUF，让 libfabric 知道我们使用的是 dmabuf。

`WRITE` 操作

在之前的章节里面，我们构建了 RdmaOp 类用来保存 RDMA 操作的上下文，其中包含了 RECV 和 SEND 两种操作类型。WRITE 操作和前面两种操作类似，当操作结束后，完成队列中会有相应的上下文，因此我们需要增加一个保存 WRITE 操作上下文的类型。

另外一方面，因为 WRITE 是单边 RDMA 操作，所以目标节点并不会产生一个完成事件。例外的是，当 WRITE 操作带上一个立即数（Immediate Data）时，目标节点会产生一个完成事件。我们可以把这个立即数当作一个标志，用这个标志来查找对应的上下文。在之前章节获取的 fi_info 信息中，我们可以看到 EFA 所支持的立即数大小为4字节。

constexpr size_t kEfaImmDataSize = 4;

enum class RdmaOpType : uint8_t {
  kRecv = 0,
  kSend = 1,
  kWrite = 2,        // Added
  kRemoteWrite = 3,  // Added
};

struct RdmaWriteOp {
  Buffer *buf;
  size_t offset;
  size_t len;
  uint32_t imm_data;
  uint64_t dest_ptr;
  fi_addr_t dest_addr;
  uint64_t dest_key;
};
static_assert(std::is_pod_v<RdmaWriteOp>);

struct RdmaRemoteWriteOp {
  uint32_t op_id;
};
static_assert(std::is_pod_v<RdmaRemoteWriteOp>);
static_assert(sizeof(RdmaRemoteWriteOp) <= kEfaImmDataSize);

struct RdmaOp {
  RdmaOpType type;
  union {
    RdmaRecvOp recv;
    RdmaSendOp send;
    RdmaWriteOp write;               // Added
    RdmaRemoteWriteOp remote_write;  // Added
  };
  std::function<void(Network &, RdmaOp &)> callback;
};

接下来我们实现 WRITE 操作。

void Network::PostWrite(RdmaWriteOp &&write,
                        std::function<void(Network &, RdmaOp &)> &&callback) {
  auto *op = new RdmaOp{
      .type = RdmaOpType::kWrite,
      .write = std::move(write),
      .callback = std::move(callback),
  };
  struct iovec iov = {
      .iov_base = (uint8_t *)write.buf->data + write.offset,
      .iov_len = write.len,
  };
  struct fi_rma_iov rma_iov = {
      .addr = write.dest_ptr,
      .len = write.len,
      .key = write.dest_key,
  };
  struct fi_msg_rma msg = {
      .msg_iov = &iov,
      .desc = &GetMR(*write.buf)->mem_desc,
      .iov_count = 1,
      .addr = write.dest_addr,
      .rma_iov = &rma_iov,
      .rma_iov_count = 1,
      .context = op,
      .data = write.imm_data,
  };
  uint64_t flags = 0;
  if (write.imm_data) {
    flags |= FI_REMOTE_CQ_DATA;
  }
  FI_CHECK(fi_writemsg(ep, &msg, flags)); // TODO: handle EAGAIN
}

因为 WRITE 操作的输入参数较多，为了使代码更清晰，我们使用了 RdmaWriteOp 结构体来保存输入参数。与 SEND 操作只需要指定发送端的 [buf, buf+size) 不同，WRITE 操作还需要指定目标端的内存地址。额外地，RDMA 为了防止未经授权的远程读写，还需要指定目标端内存区域的远程访问密钥（Remote Access Key）。我们在下文中会看到这个密钥是如何交换的。

对于带有立即数的 WRITE 操作，我们需要设置 FI_REMOTE_CQ_DATA 标志。这个标志会让目标端产生一个完成事件，我们可以通过这个完成事件来查找对应的上下文。

尽管 fi_writemsg() 支持多个发起端的 iov 以及多个目标端的 rma_iov，但是从 fi_info 中我们可以看到，EFA 只支持单个 rma_iov。因此，如果我们有多个不连续的内存区域需要写入，我们需要多次调用 fi_writemsg()。

另一个有趣的点是，在注册内存区域时，我们需要单独设置 dma-buf 的标志。但在 RDMA 操作中，我们并不需要关心内存区域在哪个设备上，也不关心是通过哪种机制访问的。

`REMOTE WRITE` 回调函数

当 WRITE 操作带有立即数时，目标端会产生一个完成事件。如果目标端期待一个带有立即数的 WRITE 操作，那么我们可以提前设置好回调函数，以便在完成事件发生时调用。

struct Network {
  // ...
  std::unordered_map<uint32_t, RdmaOp *> remote_write_ops;
};

void Network::AddRemoteWrite(
    uint32_t id, std::function<void(Network &, RdmaOp &)> &&callback) {
  CHECK(remote_write_ops.count(id) == 0);
  auto *op = new RdmaOp{
      .type = RdmaOpType::kRemoteWrite,
      .remote_write = RdmaRemoteWriteOp{.op_id = id},
      .callback = std::move(callback),
  };
  remote_write_ops[id] = op;
}

处理完成事件

在处理完成事件的时候，我们需要先判断是否是带有立即数的 WRITE 操作，即 FI_REMOTE_WRITE。如果是，我们需要从 remote_write_ops 中找到立即数 cqe.data 对应的上下文。如果是其他类型的操作，那么上下文就是 cqe.op_context。

void HandleCompletion(Network &net, const struct fi_cq_data_entry &cqe) {
  RdmaOp *op = nullptr;
  if (cqe.flags & FI_REMOTE_WRITE) {
    // REMOTE WRITE does not have op_context
    uint32_t op_id = cqe.data;
    if (!op_id)
      return;
    auto it = net.remote_write_ops.find(op_id);
    if (it == net.remote_write_ops.end())
      return;
    op = it->second;
    net.remote_write_ops.erase(it);
  } else {
    // RECV / SEND / WRITE
    op = (RdmaOp *)cqe.op_context;
    if (!op)
      return;
    if (cqe.flags & FI_RECV) {
      op->recv.recv_size = cqe.len;
    } else if (cqe.flags & FI_SEND) {
      // Nothing special
    } else if (cqe.flags & FI_WRITE) {
      // Nothing special
    } else {
      fprintf(stderr, "Unhandled completion type. cqe.flags=%lx\n", cqe.flags);
      std::exit(1);
    }
  }
  if (op->callback)
    op->callback(net, *op);
  delete op;
}

应用程序消息

首先我们要修改 AppConnectMessage，告诉服务器端关于客户端内存区域的信息，包括每个内存区域的虚拟地址、大小、以及远程访问密钥。

struct AppConnectMessage {
  struct MemoryRegion {
    uint64_t addr;
    uint64_t size;
    uint64_t rkey;
  };

  AppMessageBase base;
  EfaAddress client_addr;
  size_t num_mr;

  MemoryRegion &mr(size_t index) {
    CHECK(index < num_mr);
    return ((MemoryRegion *)((uintptr_t)&base + sizeof(*this)))[index];
  }

  size_t MessageBytes() const {
    return sizeof(*this) + num_mr * sizeof(MemoryRegion);
  }
};

然后我们增加一个新的消息类型 AppRandomFillMessage，用来告诉服务器端要在客户端的哪些页上填充随机数，以及页面大小和随机种子。另外，消息中还包含了 remote_context。之后服务器端会将 remote_context 作为立即数发送回客户端，以便客户端能够找到对应的上下文。

struct AppRandomFillMessage {
  AppMessageBase base;
  uint32_t remote_context;
  uint64_t seed;
  size_t page_size;
  size_t num_pages;

  uint32_t &page_idx(size_t index) {
    CHECK(index < num_pages);
    return ((uint32_t *)((uintptr_t)&base + sizeof(*this)))[index];
  }

  size_t MessageBytes() const {
    return sizeof(*this) + num_pages * sizeof(uint32_t);
  }
};

在上面的这两个类型里，我们都先将固定长度的部分定义在结构体的最前面，然后通过指针的方式访问变长的部分。

服务器端逻辑

类似于上一章中的做法，我们使用一个状态机来处理两个不同类型的消息。

struct RandomFillRequestState {
  Buffer *cuda_buf;
  fi_addr_t client_addr = FI_ADDR_UNSPEC;
  bool done = false;
  AppConnectMessage *connect_msg = nullptr;

  explicit RandomFillRequestState(Buffer *cuda_buf) : cuda_buf(cuda_buf) {}

  void OnRecv(Network &net, RdmaOp &op) {
    if (client_addr == FI_ADDR_UNSPEC) {
      HandleConnect(net, op);
    } else {
      HandleRequest(net, op);
    }
  }
};

当收到 CONNECT 消息时，我们把客户端的地址添加到服务器端的地址向量中，并且把这个 CONNECT 消息保存下来，以便后续使用。

struct RandomFillRequestState {
  // ...

  void HandleConnect(Network &net, RdmaOp &op) {
    auto *base_msg = (AppMessageBase *)op.recv.buf->data;
    CHECK(base_msg->type == AppMessageType::kConnect);
    CHECK(op.recv.recv_size >= sizeof(AppConnectMessage));
    auto &msg = *(AppConnectMessage *)base_msg;
    CHECK(op.recv.recv_size == msg.MessageBytes());
    CHECK(msg.num_mr > 0);

    // Save the message. Note that we don't reuse the buffer.
    connect_msg = &msg;

    // Add the client to AV
    client_addr = net.AddPeerAddress(msg.client_addr);

    printf("Received CONNECT message from client:\n");
    printf("  addr: %s\n", msg.client_addr.ToString().c_str());
    for (size_t i = 0; i < msg.num_mr; i++) {
      printf("  MR[%zu]: addr=0x%012lx size=%lu rkey=0x%016lx\n", i,
             msg.mr(i).addr, msg.mr(i).size, msg.mr(i).rkey);
    }
  }
};

当收到 RANDOM_FILL 消息时，我们首先先在服务器端自己的 GPU 缓冲区上填充随机数。

std::vector<uint8_t> RandomBytes(uint64_t seed, size_t size) { ... }

struct RandomFillRequestState {
  // ...

  void HandleRequest(Network &net, RdmaOp &op) {
    auto *base_msg = (const AppMessageBase *)op.recv.buf->data;
    CHECK(base_msg->type == AppMessageType::kRandomFill);
    CHECK(op.recv.recv_size >= sizeof(AppRandomFillMessage));
    auto &msg = *(AppRandomFillMessage *)base_msg;
    CHECK(op.recv.recv_size == msg.MessageBytes());

    printf("Received RandomFill request from client:\n");
    printf("  remote_context: 0x%08x\n", msg.remote_context);
    printf("  seed: 0x%016lx\n", msg.seed);
    printf("  page_size: %zu\n", msg.page_size);
    printf("  num_pages: %zu\n", msg.num_pages);

    // Generate random data and copy to local GPU memory
    printf("Generating random data");
    for (size_t i = 0; i < connect_msg->num_mr; ++i) {
      auto bytes = RandomBytes(msg.seed + i, msg.page_size * msg.num_pages);
      CUDA_CHECK(cudaMemcpy((uint8_t *)cuda_buf->data + i * bytes.size(),
                            bytes.data(), bytes.size(),
                            cudaMemcpyHostToDevice));
      printf(".");
      fflush(stdout);
    }
    printf("\n");

    // ...
  }
};

然后我们为每一页的数据提交一个 WRITE 操作。对于非最后一页的数据，我们跳过回调函数。对于最后一页的数据，我们在服务器端的回调函数中将状态机的 done 置为 true。并且我们在最后一页的 WRITE 操作中带上 remote_context 作为立即数，这样客户端的完成队列中就会有一个完成事件，并且客户端可以根据这个立即数找到对应的上下文。

struct RandomFillRequestState {
  void HandleRequest(Network &net, RdmaOp &op) {
    // ...

    // RDMA WRITE the data to remote GPU.
    for (size_t i = 0; i < connect_msg->num_mr; ++i) {
      for (size_t j = 0; j < msg.num_pages; j++) {
        uint32_t imm_data = 0;
        std::function<void(Network &, RdmaOp &)> callback;
        if (i + 1 == connect_msg->num_mr && j + 1 == msg.num_pages) {
          // The last WRITE.
          imm_data = msg.remote_context;
          callback = [this](Network &net, RdmaOp &op) {
            CHECK(op.type == RdmaOpType::kWrite);
            done = true;
            printf("Finished RDMA WRITE to the remote GPU memory.\n");
          };
        }
        net.PostWrite(
            {.buf = cuda_buf,
             .offset = i * (msg.page_size * msg.num_pages) + j * msg.page_size,
             .len = msg.page_size,
             .imm_data = imm_data,
             .dest_ptr =
                 connect_msg->mr(i).addr + msg.page_idx(j) * msg.page_size,
             .dest_addr = client_addr,
             .dest_key = connect_msg->mr(i).rkey},
            std::move(callback));
      }
    }
  }
};

完成了以上部分之后，剩下的服务器端的主程序逻辑就十分简单了：

constexpr size_t kMemoryRegionSize = 16 << 20;

int ServerMain(int argc, char **argv) {
  // Open Netowrk
  struct fi_info *info = GetInfo();
  auto net = Network::Open(info);

  // Allocate and register message buffer
  auto buf1 = Buffer::Alloc(kMessageBufferSize, kBufAlign);
  net.RegisterMemory(buf1);
  auto buf2 = Buffer::Alloc(kMessageBufferSize, kBufAlign);
  net.RegisterMemory(buf2);

  // Allocate and register CUDA memory
  auto cuda_buf = Buffer::AllocCuda(kMemoryRegionSize * 2, kBufAlign);
  net.RegisterMemory(cuda_buf);
  printf("Registered 1 buffer on cuda:%d\n", cuda_buf.cuda_device);

  // Loop forever. Accept one client at a time.
  for (;;) {
    printf("------\n");
    // State machine
    RandomFillRequestState s(&cuda_buf);
    // RECV for CONNECT
    net.PostRecv(buf1, [&s](Network &net, RdmaOp &op) { s.OnRecv(net, op); });
    // RECV for RandomFillRequest
    net.PostRecv(buf2, [&s](Network &net, RdmaOp &op) { s.OnRecv(net, op); });
    // Wait for completion
    while (!s.done) {
      net.PollCompletion();
    }
  }

  return 0;
}

客户端逻辑

客户端首先从命令行参数中获得服务器端的地址以及要写入的页面的大小和数量。

int ClientMain(int argc, char **argv) {
  auto server_addrname = EfaAddress::Parse(argv[1]);
  size_t page_size = std::stoull(argv[2]);
  size_t num_pages = std::stoull(argv[3]);
  size_t max_pages = kMemoryRegionSize / page_size;
  CHECK(page_size * num_pages <= kMemoryRegionSize);

  // ...
}

然后打开网络接口，注册一个缓冲区用户发送消息，注册两个 GPU 缓冲区用来存储随机数。

  // Open Netowrk
  struct fi_info *info = GetInfo();
  auto net = Network::Open(info);
  auto server_addr = net.AddPeerAddress(server_addrname);

  // Allocate and register message buffer
  auto buf1 = Buffer::Alloc(kMessageBufferSize, kBufAlign);
  net.RegisterMemory(buf1);

  // Allocate and register CUDA memory
  auto cuda_buf1 = Buffer::AllocCuda(kMemoryRegionSize, kBufAlign);
  net.RegisterMemory(cuda_buf1);
  auto cuda_buf2 = Buffer::AllocCuda(kMemoryRegionSize, kBufAlign);
  net.RegisterMemory(cuda_buf2);
  printf("Registered 2 buffers on cuda:%d\n", cuda_buf1.cuda_device);

接着，向服务器端发送 CONNECT 消息。

  // Send address and MR to server
  auto &connect_msg = *(AppConnectMessage *)buf1.data;
  connect_msg = {
      .base = {.type = AppMessageType::kConnect},
      .client_addr = net.addr,
      .num_mr = 2,
  };
  connect_msg.mr(0) = {.addr = (uint64_t)cuda_buf1.data,
                       .size = kMemoryRegionSize,
                       .rkey = net.GetMR(cuda_buf1)->key};
  connect_msg.mr(1) = {.addr = (uint64_t)cuda_buf2.data,
                       .size = kMemoryRegionSize,
                       .rkey = net.GetMR(cuda_buf2)->key};
  bool connect_sent = false;
  net.PostSend(
      server_addr, buf1, connect_msg.MessageBytes(),
      [&connect_sent](Network &net, RdmaOp &op) { connect_sent = true; });
  while (!connect_sent) {
    net.PollCompletion();
  }
  printf("Sent CONNECT message to server\n");

在发送 RANDOM_FILL 请求之前，我们先设置 REMOTE WRITE 的回调函数。

  // Prepare to receive the last REMOTE WRITE from server
  bool last_remote_write_received = false;
  uint32_t remote_write_op_id = 0x123;
  net.AddRemoteWrite(remote_write_op_id,
                     [&last_remote_write_received](Network &net, RdmaOp &op) {
                       last_remote_write_received = true;
                     });

然后我们选中一些页面，发送 RANDOM_FILL 请求。

  // Prepare request
  uint64_t req_seed = ...;
  std::vector<uint32_t> page_idx = ...;

  // Send message to server
  auto &req_msg = *(AppRandomFillMessage *)buf1.data;
  req_msg = {
      .base = {.type = AppMessageType::kRandomFill},
      .remote_context = remote_write_op_id,
      .seed = req_seed,
      .page_size = page_size,
      .num_pages = num_pages,
  };
  for (size_t i = 0; i < num_pages; i++) {
    req_msg.page_idx(i) = page_idx[i];
  }
  bool req_sent = false;
  net.PostSend(server_addr, buf1, req_msg.MessageBytes(),
               [&req_sent](Network &net, RdmaOp &op) { req_sent = true; });
  while (!req_sent) {
    net.PollCompletion();
  }
  printf("Sent RandomFillRequest to server. page_size: %zu, num_pages: %zu\n",
         page_size, num_pages);

等待最后一个 REMOTE WRITE 操作的完成，最后验证收到的数据是正确的。

  // Wait for REMOTE WRITE from server
  while (!last_remote_write_received) {
    net.PollCompletion();
  }
  printf("Received RDMA WRITE to local GPU memory.\n");

  // Verify data
  auto expected1 = RandomBytes(req_seed, page_size * num_pages);
  auto expected2 = RandomBytes(req_seed + 1, page_size * num_pages);
  auto actual1 = std::vector<uint8_t>(page_size * num_pages);
  auto actual2 = std::vector<uint8_t>(page_size * num_pages);
  for (size_t i = 0; i < num_pages; i++) {
    CUDA_CHECK(cudaMemcpy(actual1.data() + i * page_size,
                          (uint8_t *)cuda_buf1.data + page_idx[i] * page_size,
                          page_size, cudaMemcpyDeviceToHost));
    CUDA_CHECK(cudaMemcpy(actual2.data() + i * page_size,
                          (uint8_t *)cuda_buf2.data + page_idx[i] * page_size,
                          page_size, cudaMemcpyDeviceToHost));
  }
  CHECK(expected1 == actual1);
  CHECK(expected2 == actual2);
  printf("Data is correct\n");

  return 0;

运行效果

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/6_write.cpp

驾驭3200Gbps网络(5): 双向接收发送

2024-12-25T00:00:00-08:00

在上一章中，我们实现了单向的接收和发送。在本章中，我们将拓展上一章的程序，实现双向的接收和发送。在服务器端收到客户端的消息后，对消息进行翻转，然后再发回给客户端。我们把这一个程序命名为 5_reverse.cpp。

不同类型的消息

要让服务器端能给客户端发消息，首先要让服务器端知道客户端的地址。因此我们首先让客户端向服务器端发送自己的地址，然后再向服务器端发送消息，最后再接收服务器端的消息。为了区分不同类型的消息，我们在消息头部加上一个 type 字段进行区分。

enum class AppMessageType : uint8_t {
  kConnect = 0,
  kData = 1,
};

struct AppMessageBase {
  AppMessageType type;
};

struct AppConnectMessage {
  AppMessageBase base;
  EfaAddress client_addr;
};

struct AppDataMessage {
  AppMessageBase base;
  // Data follows
};

注意到 AppDataMessage 结构体中并没有定义长度，这是因为在发送时会指定缓冲区长度，在接收时会得到缓冲区长度。扣除掉消息头部的长度，就是数据的长度。

为什么不将 CONNECT 和 DATA 两个消息合并成一个消息呢？

服务器端并不需要向客户端发送 CONNECT 消息，因此不需要在 DATA 消息中包含 CONNECT 的内容。
对这两者进行拆分，让通信的逻辑更清晰。
客户端也可以多次发送 DATA 消息，而不需要每次都发送 CONNECT 消息。
在后面的章节中，我们会对 CONNECT 消息进行扩展，包含更多的信息。

服务器端逻辑

因为客户端会发送两个消息，所以服务器端需要维护一个状态机，来处理不同类型的消息。

struct ReverseRequestState {
  fi_addr_t client_addr = FI_ADDR_UNSPEC;
  bool done = false;

  // ...
};

当收到 CONNECT 消息时，客户端会将客户端的地址添加到自己的地址向量中，并将其对应的 fi_addr_t 保存在状态机中。

struct ReverseRequestState {
  // ...

  void HandleConnect(Network &net, RdmaOp &op) {
    auto *base_msg = (const AppMessageBase *)op.recv.buf->data;
    CHECK(base_msg->type == AppMessageType::kConnect);
    CHECK(op.recv.recv_size == sizeof(AppConnectMessage));
    auto *msg = (const AppConnectMessage *)base_msg;
    printf("Received CONNECT message from client: %s\n",
           msg->client_addr.ToString().c_str());
    client_addr = net.AddPeerAddress(msg->client_addr);
  }
};

当收到 DATA 消息时，服务器端会将消息翻转后再发回给客户端。并且将状态机的 done 置为 true。

struct ReverseRequestState {
  // ...

  void HandleData(Network &net, RdmaOp &op) {
    auto *base_msg = (const AppMessageBase *)op.recv.buf->data;
    CHECK(base_msg->type == AppMessageType::kData);
    auto *msg = (uint8_t *)op.recv.buf->data + sizeof(*base_msg);
    auto len = op.recv.recv_size - sizeof(*base_msg);
    printf("Received message (len=%zu): %.*s\n", len, (int)len, msg);
    for (size_t i = 0, j = len - 1; i < j; ++i, --j) {
      auto t = msg[i];
      msg[i] = msg[j];
      msg[j] = t;
    }
    net.PostSend(client_addr, *op.recv.buf, op.recv.recv_size,
                 [this](Network &net, RdmaOp &op) {
                   printf("Sent reversed message to client\n");
                   done = true;
                 });
  }
};

尽管我们一次只处理一个客户端的请求，但是我们仍然需要使用两个缓冲区，以及提前提交两个 RECV 操作。这是因为，如果我们等到收到 CONNECT 之后再提交 RECV，那么客户端发送的数据有可能在提交第二个 RECV 之前就已经发出了，这样就会因为服务器端没有等待中的 RECV 操作而导致数据丢失。

因为这两个 RECV 操作是同时提交的，无法在指定回调函数的时候区分状态机的状态，所以我们需要让两者使用同样的回调函数，然后在回调函数中根据状态机的状态来决定下一步的操作。

struct ReverseRequestState {
  // ...

  void OnRecv(Network &net, RdmaOp &op) {
    if (client_addr == FI_ADDR_UNSPEC) {
      HandleConnect(net, op);
    } else {
      HandleData(net, op);
    }
  }
};

最后就是服务器端的完整逻辑：

int ServerMain(int argc, char **argv) {
  struct fi_info *info = GetInfo();
  auto net = Network::Open(info);

  // Allocate and register memory
  auto buf1 = Buffer::Alloc(kMessageBufferSize, kBufAlign);
  net.RegisterMemory(buf1);
  auto buf2 = Buffer::Alloc(kMessageBufferSize, kBufAlign);
  net.RegisterMemory(buf2);

  // Loop forever. Accept one client at a time.
  for (;;) {
    // State machine
    ReverseRequestState s;
    // RECV for CONNECT and DATA
    net.PostRecv(buf1, [&s](Network &net, RdmaOp &op) { s.OnRecv(net, op); });
    net.PostRecv(buf2, [&s](Network &net, RdmaOp &op) { s.OnRecv(net, op); });
    // Wait for completion
    while (!s.done) {
      net.PollCompletion();
    }
  }

  return 0;
}

客户端逻辑

客户端的逻辑也很简单。首先我们打开网络接口，并注册两个缓冲区。

int ClientMain(int argc, char **argv) {
  CHECK(argc == 2 || argc == 3);
  auto server_addrname = EfaAddress::Parse(argv[1]);
  std::string message = argc == 3 ? argv[2] : "Hello, world!";

  struct fi_info *info = GetInfo();
  auto net = Network::Open(info);
  auto server_addr = net.AddPeerAddress(server_addrname);

  auto buf1 = Buffer::Alloc(kMessageBufferSize, kBufAlign);
  net.RegisterMemory(buf1);
  auto buf2 = Buffer::Alloc(kMessageBufferSize, kBufAlign);
  net.RegisterMemory(buf2);

  // ...
}

接着我们向服务器端发送 CONNECT 消息。

  // Send address to server
  auto *connect_msg = (AppConnectMessage *)buf1.data;
  connect_msg->base.type = AppMessageType::kConnect;
  connect_msg->client_addr = net.addr;
  bool connect_sent = false;
  net.PostSend(server_addr, buf1, sizeof(*connect_msg),
               [&connect_sent](Network &net, RdmaOp &op) {
                 printf("Sent CONNECT message to server\n");
                 connect_sent = true;
               });
  while (!connect_sent) {
    net.PollCompletion();
  }

接下来我们向服务器端发送 DATA 消息。为了保证能够接收到服务器端的回复，我们需要先提交一个 RECV 操作，再提交 SEND 操作。

  // Prepare to receive reversed message from server
  bool msg_received = false;
  net.PostRecv(buf2, [&msg_received](Network &net, RdmaOp &op) {
    auto *msg = (const char *)op.recv.buf->data;
    auto len = op.recv.recv_size;
    printf("Received message (len=%zu): %.*s\n", len, (int)len, msg);
    msg_received = true;
  });

  // Send message to server
  auto *data_msg = (AppDataMessage *)buf1.data;
  data_msg->base.type = AppMessageType::kData;
  memcpy((char *)buf1.data + sizeof(*data_msg), message.c_str(),
         message.size());
  net.PostSend(
      server_addr, buf1, sizeof(*data_msg) + message.size(),
      [](Network &net, RdmaOp &op) { printf("Sent message to server\n"); });

最后我们等待服务器端的回复。

  // Wait for message from server
  while (!msg_received) {
    net.PollCompletion();
  }

  return 0;

运行效果

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/5_reverse.cpp

驾驭3200Gbps网络(4): 单向接收发送

2024-12-25T00:00:00-08:00

在经过前面一系列章节的铺垫，本章终于可以开始写代码了。本章的目标是在两台机器之间使用单张网卡实现单向的 RECV 和 SEND。

虽然 libfabric 的语言是 C，但是因为我实在是不懂 C 语言，所以我接下来都会使用我熟悉的 C++。并且因为这一系列的文章是一个教程，所以我不会刻意封装 C++ 的类，而是尽量保持简单。同时，我也会省略许多操作的异常返回值的处理，只是加入一些断言来做检查。

#define CHECK(stmt)                                                            \
  do {                                                                         \
    if (!(stmt)) {                                                             \
      fprintf(stderr, "%s:%d %s\n", __FILE__, __LINE__, #stmt);                \
      std::exit(1);                                                            \
    }                                                                          \
  } while (0)

#define FI_CHECK(stmt)                                                         \
  do {                                                                         \
    int rc = (stmt);                                                           \
    if (rc) {                                                                  \
      fprintf(stderr, "%s:%d %s failed with %d (%s)\n", __FILE__, __LINE__,    \
              #stmt, rc, fi_strerror(-rc));                                    \
      std::exit(1);                                                            \
    }                                                                          \
  } while (0)

获取网络接口信息

在上一章中我们知道了 libfabric 支持非常多设备，并且 EFA 支持两种不同的协议（RDM 和 DGRAM）。在这一系列的文章中我们都专注于使用 RDM 协议的 EFA。所以我们首先要用 fi_getinfo() 筛选出来符合要求的 libfabric 网络接口。

#include <rdma/fabric.h>
#include <rdma/fi_cm.h>
#include <rdma/fi_domain.h>
#include <rdma/fi_endpoint.h>
#include <rdma/fi_errno.h>
#include <rdma/fi_rma.h>

struct fi_info *GetInfo() {
  struct fi_info *hints, *info;
  hints = fi_allocinfo();
  hints->ep_attr->type = FI_EP_RDM;
  hints->fabric_attr->prov_name = strdup("efa");
  FI_CHECK(fi_getinfo(FI_VERSION(2, 0), nullptr, nullptr, 0, hints, &info));
  fi_freeinfo(hints);
  return info;
}

在这一段程序里面，我们在 hints 里面指定了端点类型为 RDM，Provider 名称为 efa。fi_getinfo() 会将所有符合条件的网络接口以链表的形式返回。

接下来，让我们遍历这个链表，输出所有网络接口的基本信息，作为最基本的正确性检测。

int main() {
  struct fi_info *info = GetInfo();
  for (auto *fi = info; fi; fi = fi->next) {
    printf("domain: %14s", fi->domain_attr->name);
    printf(", nic: %10s", fi->nic->device_attr->name);
    printf(", fabric: %s", fi->fabric_attr->prov_name);
    printf(", link: %.0fGbps", fi->nic->link_attr->speed / 1e9);
    printf("\n");
  }
  return 0;
}

让我们把程序保存为 src/4_hello.cpp。在上一章我们把 libfabric 安装到了 build/libfabric/ 文件夹。我们可以使用下面的命令编译：

g++ -Wall -Werror -std=c++17 -O2 -g \
  -I./build/libfabric/include \
  -o build/4_hello src/4_hello.cpp \
  -L./build/libfabric/lib -lfabric

运行编译好的可执行文件，可以得到下面的输出，证明我们最基础的软硬件设置都是正确的：

domain:  rdmap79s0-rdm, nic:  rdmap79s0, fabric: efa, link: 100Gbps
domain:  rdmap80s0-rdm, nic:  rdmap80s0, fabric: efa, link: 100Gbps
domain:  rdmap81s0-rdm, nic:  rdmap81s0, fabric: efa, link: 100Gbps
domain:  rdmap82s0-rdm, nic:  rdmap82s0, fabric: efa, link: 100Gbps
domain:  rdmap96s0-rdm, nic:  rdmap96s0, fabric: efa, link: 100Gbps
domain:  rdmap97s0-rdm, nic:  rdmap97s0, fabric: efa, link: 100Gbps
domain:  rdmap98s0-rdm, nic:  rdmap98s0, fabric: efa, link: 100Gbps
domain:  rdmap99s0-rdm, nic:  rdmap99s0, fabric: efa, link: 100Gbps
domain: rdmap113s0-rdm, nic: rdmap113s0, fabric: efa, link: 100Gbps
domain: rdmap114s0-rdm, nic: rdmap114s0, fabric: efa, link: 100Gbps
domain: rdmap115s0-rdm, nic: rdmap115s0, fabric: efa, link: 100Gbps
domain: rdmap116s0-rdm, nic: rdmap116s0, fabric: efa, link: 100Gbps
domain: rdmap130s0-rdm, nic: rdmap130s0, fabric: efa, link: 100Gbps
domain: rdmap131s0-rdm, nic: rdmap131s0, fabric: efa, link: 100Gbps
domain: rdmap132s0-rdm, nic: rdmap132s0, fabric: efa, link: 100Gbps
domain: rdmap133s0-rdm, nic: rdmap133s0, fabric: efa, link: 100Gbps
domain: rdmap147s0-rdm, nic: rdmap147s0, fabric: efa, link: 100Gbps
domain: rdmap148s0-rdm, nic: rdmap148s0, fabric: efa, link: 100Gbps
domain: rdmap149s0-rdm, nic: rdmap149s0, fabric: efa, link: 100Gbps
domain: rdmap150s0-rdm, nic: rdmap150s0, fabric: efa, link: 100Gbps
domain: rdmap164s0-rdm, nic: rdmap164s0, fabric: efa, link: 100Gbps
domain: rdmap165s0-rdm, nic: rdmap165s0, fabric: efa, link: 100Gbps
domain: rdmap166s0-rdm, nic: rdmap166s0, fabric: efa, link: 100Gbps
domain: rdmap167s0-rdm, nic: rdmap167s0, fabric: efa, link: 100Gbps
domain: rdmap181s0-rdm, nic: rdmap181s0, fabric: efa, link: 100Gbps
domain: rdmap182s0-rdm, nic: rdmap182s0, fabric: efa, link: 100Gbps
domain: rdmap183s0-rdm, nic: rdmap183s0, fabric: efa, link: 100Gbps
domain: rdmap184s0-rdm, nic: rdmap184s0, fabric: efa, link: 100Gbps
domain: rdmap198s0-rdm, nic: rdmap198s0, fabric: efa, link: 100Gbps
domain: rdmap199s0-rdm, nic: rdmap199s0, fabric: efa, link: 100Gbps
domain: rdmap200s0-rdm, nic: rdmap200s0, fabric: efa, link: 100Gbps
domain: rdmap201s0-rdm, nic: rdmap201s0, fabric: efa, link: 100Gbps

打开网络接口

上一章中我们介绍了 libfabric 的软件对象模型，是一个树形的结构。但是因为在这里我们只打算使用一个端点，所以为了简单起见，我们可以把所有的对象都放在一个结构体里面：

struct Network {
  struct fi_info *fi;
  struct fid_fabric *fabric;
  struct fid_domain *domain;
  struct fid_cq *cq;
  struct fid_av *av;
  struct fid_ep *ep;
  // ... other fields

  static Network Open(struct fi_info *fi);
  // ... other methods
};

让我们实现 Open() 方法，打开 fi 所指代的网络接口：

Network Network::Open(struct fi_info *fi) {
  struct fid_fabric *fabric;
  FI_CHECK(fi_fabric(fi->fabric_attr, &fabric, nullptr));

  struct fid_domain *domain;
  FI_CHECK(fi_domain(fabric, fi, &domain, nullptr));

  struct fid_cq *cq;
  struct fi_cq_attr cq_attr = {};
  cq_attr.format = FI_CQ_FORMAT_DATA;
  FI_CHECK(fi_cq_open(domain, &cq_attr, &cq, nullptr));

  struct fid_av *av;
  struct fi_av_attr av_attr = {};
  FI_CHECK(fi_av_open(domain, &av_attr, &av, nullptr));

  struct fid_ep *ep;
  FI_CHECK(fi_endpoint(domain, fi, &ep, nullptr));
  FI_CHECK(fi_ep_bind(ep, &cq->fid, FI_SEND | FI_RECV));
  FI_CHECK(fi_ep_bind(ep, &av->fid, 0));

  FI_CHECK(fi_enable(ep));
  //...
}

在这一段代码中，我们首先创建了 fabric，然后根据 fi 创建了 domain。

接着我们创建完成队列 cq 和地址向量 av。对于完成队列，我们指定了 FI_CQ_FORMAT_DATA 格式，之后当我们轮询完成队列的时候，就可以得到一个 fi_cq_data_entry 结构体：

struct fi_cq_data_entry {
  void     *op_context;
  uint64_t flags;
  size_t   len;
  void     *buf;
  uint64_t data;
};

op_context 是我们在提交操作的时候传入的上下文，用来区分不同的操作。我们在后文会用到。
flags 是一个位掩码，表示了这个完成事件的类型。我们在后文会用到。
buf 和 len 表示了这个完成事件对应的数据缓冲区和长度。这对于 RECV 操作是有用的。
data 是 WRITE_IMM 操作所携带的立即数。在本章中我们不会用到，但是在后面的章节中会用到。

其他完成队列的格式可以参见 fi_cq(3)。

接着我们创建了端点 ep，并且绑定了完成队列和地址向量。最后我们启用了端点。

EFA 端点地址

让我们顺便在 Network::Open() 中保存一下端点的地址。根据 EFA RDM 协议文档，EFA 端点地址长度为 32 字节。我们可以使用 fi_getname() 来获取端点地址：

Network Network::Open(struct fi_info *fi) {
  // ...
  uint8_t addr[64];
  size_t addrlen = sizeof(addr);
  FI_CHECK(fi_getname(&ep->fid, addr, &addrlen));
  FI_CHECK(addrlen == 32);
  // ...
}

这个地址是一个二进制的地址。为了方便输入输出，我们可以把这个地址转换成十六进制的字符串：

struct EfaAddress {
  uint8_t bytes[32];

  explicit EfaAddress(uint8_t bytes[32]) { memcpy(this->bytes, bytes, 32); }

  std::string ToString() const {
    char buf[65];
    for (size_t i = 0; i < 32; i++) {
      snprintf(buf + 2 * i, 3, "%02x", bytes[i]);
    }
    return std::string(buf, 64);
  }

  static EfaAddress Parse(const std::string &str) {
    CHECK(str.size() == 64);
    uint8_t bytes[32];
    for (size_t i = 0; i < 32; i++) {
      sscanf(str.c_str() + 2 * i, "%02hhx", &bytes[i]);
    }
    return EfaAddress(bytes);
  }
};

在 Network 结构体中保存这个地址：

struct Network {
  // ...
  EfaAddress addr;
  // ...
};

Network Network::Open(struct fi_info *fi) {
  // ...
  return Network{fi, fabric, domain, cq, av, ep, EfaAddress(addr)};
}

添加目标地址

在前面的章节中我们提到了，为了减少地址解析的开销，libfabric 在发起网络操作之前就需要将目标地址添加到地址向量中。fi_av_insert() 函数可以用来添加地址，并会返回一个 fi_addr_t 类型的地址。让我们增加一个方法来添加目标地址：

fi_addr_t Network::AddPeerAddress(const EfaAddress &peer_addr) {
  fi_addr_t addr = FI_ADDR_UNSPEC;
  int ret = fi_av_insert(av, peer_addr.bytes, 1, &addr, 0, nullptr);
  if (ret != 1) {
    fprintf(stderr, "fi_av_insert failed: %d\n", ret);
    std::exit(1);
  }
  return addr;
}

缓冲区

在 libfabric 中，数据传输的缓冲区是由用户提供的。我们使用一个简单的结构体来保存缓冲区的基本信息。除了指针和长度，还有一个需要注意的是内存地址的对齐。在 aws-ofi-nccl 的代码中，我发现了他们使用128字节的对齐。虽然我没有查到 libfabric 或者 EFA 任何关于对齐的要求，但是为了保险起见，我们也使用128字节的对齐。

constexpr size_t kBufAlign = 128; // EFA alignment requirement

void *align_up(void *ptr, size_t align) {
  uintptr_t addr = (uintptr_t)ptr;
  return (void *)((addr + align - 1) & ~(align - 1));
}

struct Buffer {
  void *data;
  size_t size;

  static Buffer Alloc(size_t size, size_t align) {
    void *raw_data = malloc(size);
    CHECK(raw_data != nullptr);
    return Buffer(raw_data, size, align);
  }

  Buffer(Buffer &&other)
      : data(other.data), size(other.size), raw_data(other.raw_data) {
    other.data = nullptr;
    other.raw_data = nullptr;
  }

  ~Buffer() { free(raw_data); }

private:
  void *raw_data;

  Buffer(void *raw_data, size_t raw_size, size_t align) {
    this->raw_data = raw_data;
    this->data = align_up(raw_data, align);
    this->size = (size_t)((uintptr_t)raw_data + raw_size - (uintptr_t)data);
  }
  Buffer(const Buffer &) = delete;
};

注册内存区域

在前面章节我们提到，要让网卡能够访问用户虚拟地址空间中的缓冲区，我们需要使用 fi_mr_regattr() 方法注册这个内存区域。

struct Network {
  // ...
  std::unordered_map<void *, struct fid_mr *> mr;
  void RegisterMemory(Buffer &buf);
  struct fid_mr *GetMR(const Buffer &buf);
};

void Network::RegisterMemory(Buffer &buf) {
  struct fid_mr *mr;
  struct fi_mr_attr mr_attr = {};
  struct iovec iov = {.iov_base = buf.data, .iov_len = buf.size};
  mr_attr.mr_iov = &iov;
  mr_attr.iov_count = 1;
  mr_attr.access = FI_SEND | FI_RECV;
  uint64_t flags = 0;
  FI_CHECK(fi_mr_regattr(domain, &mr_attr, flags, &mr));
  this->mr[buf.data] = mr;
}

struct fid_mr *Network::GetMR(const Buffer &buf) {
  auto it = mr.find(buf.data);
  CHECK(it != mr.end());
  return it->second;
}

提交操作

接下来让我们看看如何提交一个 RDMA 网络操作，以及如何保存操作的上下文。

操作的上下文

当一个网络操作完成的时候，我们需要触发后续的应用逻辑。为了区分不同的操作，我们使用一个类 RdmaOp 来保存操作的类型、数据和回调函数。我们在提交操作的时候将这个 RdmaOp 作为上下文传入，当操作完成的时候，我们就可以根据这个上下文来调用回调函数。

enum class RdmaOpType : uint8_t {
  kRecv = 0,
  kSend = 1,
};

struct RdmaRecvOp {
  Buffer *buf;
  fi_addr_t src_addr; // Set after completion
  size_t recv_size;   // Set after completion
};
static_assert(std::is_pod_v<RdmaRecvOp>);

struct RdmaSendOp {
  Buffer *buf;
  size_t len;
  fi_addr_t dest_addr;
};
static_assert(std::is_pod_v<RdmaSendOp>);

struct RdmaOp {
  RdmaOpType type;
  union {
    RdmaRecvOp recv;
    RdmaSendOp send;
  };
  std::function<void(Network &, RdmaOp &)> callback;
};

`RECV` 操作

要接收数据，我们需要提交一个 RECV 操作。在 libfabric 中，我们使用 fi_recvmsg() 函数来提交一个 RECV 操作。

void Network::PostRecv(Buffer &buf,
                       std::function<void(Network &, RdmaOp &)> &&callback) {
  auto *op = new RdmaOp{
      .type = RdmaOpType::kRecv,
      .recv =
          RdmaRecvOp{.buf = &buf, .src_addr = FI_ADDR_UNSPEC, .recv_size = 0},
      .callback = std::move(callback),
  };
  struct iovec iov = {
      .iov_base = buf.data,
      .iov_len = buf.size,
  };
  struct fi_msg msg = {
      .msg_iov = &iov,
      .desc = &GetMR(buf)->mem_desc,
      .iov_count = 1,
      .addr = FI_ADDR_UNSPEC,
      .context = op,
  };
  FI_CHECK(fi_recvmsg(ep, &msg, 0)); // TODO: handle EAGAIN
}

在上面代码中，我们首先分配了一个 RdmaOp 对象。这里我们直接使用 new 来分配内存。之后，我们会在操作完成的时候释放这个内存。

在调用 fi_recvmsg() 的时候，我们传入了缓冲区及其内存区域的描述符、作为上下文的 RdmaOp 对象。我们没有指定源地址，因此这个 RECV 操作可以接收任何地址的数据。

特别注意到，我们没有处理 fi_recvmsg() 的返回值。当网络操作非常频繁的时候，fi_recvmsg() 可能会返回 FI_EAGAIN，表示当前的网络接口已经没有空间来接收新的数据了。在这种情况下，我们需要等待一段时间再次提交 RECV 操作。因为在本章中我们只接收和发送一个消息，所以大概率不会遇到这种情况。因此在本章中我们不会处理这种情况，但是在后面的章节中我们会处理这种情况。

`SEND` 操作

SEND 操作和 RECV 操作类似，只是我们需要指定目标地址。

void Network::PostSend(fi_addr_t addr, Buffer &buf, size_t len,
                       std::function<void(Network &, RdmaOp &)> &&callback) {
  CHECK(len <= buf.size);
  auto *op = new RdmaOp{
      .type = RdmaOpType::kSend,
      .send = RdmaSendOp{.buf = &buf, .len = len, .dest_addr = addr},
      .callback = std::move(callback),
  };
  struct iovec iov = {
      .iov_base = buf.data,
      .iov_len = len,
  };
  struct fi_msg msg = {
      .msg_iov = &iov,
      .desc = &GetMR(buf)->mem_desc,
      .iov_count = 1,
      .addr = addr,
      .context = op,
  };
  FI_CHECK(fi_sendmsg(ep, &msg, 0)); // TODO: handle EAGAIN
}

轮询完成队列

我们需要轮询完成队列来获取操作的完成事件。我们可以使用 fi_cq_read() 函数来读取完成队列。我们需要传入一个数组来保存完成事件，以及数组的长度。因为我们前面指定了完成队列的格式为 FI_CQ_FORMAT_DATA，所以我们可以在栈上分配一个 fi_cq_data_entry 数组来保存完成事件。如果完成事件的数量超过了数组的长度，那么我们需要多次调用 fi_cq_read() 来读取所有的完成事件。

constexpr size_t kCompletionQueueReadCount = 16;

void Network::PollCompletion() {
  struct fi_cq_data_entry cqe[kCompletionQueueReadCount];
  for (;;) {
    auto ret = fi_cq_read(cq, cqe, kCompletionQueueReadCount);
    // ...
  }
}

fi_cq_read() 的返回值有几种情况：

如果返回值大于 0，表示成功读取了完成事件。我们可以遍历 cqe 数组，处理每一个完成事件。
如果返回值是 -FI_EAGAIN，表示没有更多的完成事件。我们可以退出循环。
如果返回值是 -FI_EAVAIL，表示完成队列中有错误事件。我们可以使用 fi_cq_readerr() 函数来读取错误事件。
- 如果 fi_cq_readerr() 返回1，表示成功读取了一个错误事件。我们可以输出错误信息。
- 如果 fi_cq_readerr() 返回其他值，则表示出现了其他错误。我们可以直接退出程序。

    if (ret > 0) {
      for (ssize_t i = 0; i < ret; i++) {
        HandleCompletion(*this, cqe[i]);
      }
    } else if (ret == -FI_EAGAIN) {
      break;  // No more completions
    } else if (ret == -FI_EAVAIL) {
      struct fi_cq_err_entry err_entry;
      ret = fi_cq_readerr(cq, &err_entry, 0);
      if (ret > 0) {
        fprintf(stderr, "Failed libfabric operation: %s\n",
                fi_cq_strerror(cq, err_entry.prov_errno, err_entry.err_data,
                               nullptr, 0));
      } else {
        fprintf(stderr, "fi_cq_readerr error: %zd (%s)\n", ret,
                fi_strerror(-ret));
        std::exit(1);
      }
    } else {
      fprintf(stderr, "fi_cq_read error: %zd (%s)\n", ret, fi_strerror(-ret));
      std::exit(1);
    }

在处理完成事件的时候，我们首先读取 op_context。如果 op_context == nullptr，表示我们在提交操作的时候没有传入上下文，这种情况下我们直接忽略这个完成事件。

否则，我们就可以根据 flags 来判断这个完成事件的类型。如果是 FI_RECV，那么我们就可以读取 len 和 buf，并且调用回调函数。如果是 FI_SEND，那么我们就可以调用回调函数。

最后，我们需要释放 RdmaOp 对象。

void HandleCompletion(Network &net, const struct fi_cq_data_entry &cqe) {
  auto comp_flags = cqe.flags;
  auto op = (RdmaOp *)cqe.op_context;
  if (!op) {
    return;
  }
  if (comp_flags & FI_RECV) {
    op->recv.recv_size = cqe.len;
    if (op->callback)
      op->callback(net, *op);
  } else if (comp_flags & FI_SEND) {
    if (op->callback)
      op->callback(net, *op);
  } else {
    fprintf(stderr, "Unhandled completion type. comp_flags=%lx\n", comp_flags);
    std::exit(1);
  }
  delete op;
}

因为在本章中我们只有 RECV 和 SEND 操作，所以我们只处理这两种完成事件。在后面的章节中，我们会处理更多的操作类型。

服务器端逻辑

完成了上面所有的准备工作，我们就可以编写服务器端的逻辑了。

在服务器端，我们首先打开网络接口，然后注册一个缓冲区，然后提交一个 RECV 操作。当 RECV 操作完成的时候，我们就可以读取数据，并且再次提交一个 RECV 操作。

要如何让客户端知道服务器端的地址呢？很多其他的示例程序要么选择了另外启动一个 TCP socket 来交换地址，要么选择了 fork() 出一个子进程来运行客户端。在这里我们选择了一个更简单的方法，就是在服务器端打印出地址，然后让用户手动输入。

完整的服务器端代码如下：

constexpr size_t kMessageBufferSize = 8192;

int ServerMain(int argc, char **argv) {
  struct fi_info *info = GetInfo();
  auto net = Network::Open(info);
  printf("domain: %14s", info->domain_attr->name);
  printf(", nic: %10s", info->nic->device_attr->name);
  printf(", fabric: %s", info->fabric_attr->prov_name);
  printf(", link: %.0fGbps", info->nic->link_attr->speed / 1e9);
  printf("\n");
  printf("Run client with the following command:\n");
  printf("  %s %s\n", argv[0], net.addr.ToString().c_str());
  printf("  %s %s \"anytext\"\n", argv[0], net.addr.ToString().c_str());
  printf("------\n");

  auto buf_msg = Buffer::Alloc(kMessageBufferSize, kBufAlign);
  net.RegisterMemory(buf_msg);
  net.PostRecv(buf_msg, [](Network &net, RdmaOp &op) {
    auto *msg = (const char *)op.recv.buf->data;
    auto len = op.recv.recv_size;
    printf("Received message (len=%zu): %.*s\n", len, (int)len, msg);
    net.PostRecv(*op.recv.buf, std::move(op.callback));
  });

  for (;;) {
    net.PollCompletion();
  }

  return 0;
}

客户端逻辑

客户端逻辑与服务器端逻辑类似。我们首先打开网络接口，然后将服务器端的地址添加到地址向量中，接着注册一个缓冲区，然后提交一个 SEND 操作。当 SEND 操作完成的时候，客户端的逻辑就结束了。

int ClientMain(int argc, char **argv) {
  CHECK(argc == 2 || argc == 3);
  auto server_addrname = EfaAddress::Parse(argv[1]);
  std::string message = argc == 3 ? argv[2] : "Hello, world!";

  struct fi_info *info = GetInfo();
  auto net = Network::Open(info);
  printf("domain: %14s", info->domain_attr->name);
  printf(", nic: %10s", info->nic->device_attr->name);
  printf(", fabric: %s", info->fabric_attr->prov_name);
  printf(", link: %.0fGbps", info->nic->link_attr->speed / 1e9);
  printf("\n");
  auto server_addr = net.AddPeerAddress(server_addrname);
  auto buf_msg = Buffer::Alloc(kMessageBufferSize, kBufAlign);
  net.RegisterMemory(buf_msg);
  memcpy(buf_msg.data, message.data(), message.size());

  bool sent = false;
  net.PostSend(server_addr, buf_msg, message.size(),
               [&sent](Network &net, RdmaOp &op) {
                 auto *msg = (const char *)op.send.buf->data;
                 auto len = op.send.len;
                 printf("Sent message (len=%zu): %.*s\n", len, (int)len, msg);
                 sent = true;
               });
  while (!sent) {
    net.PollCompletion();
  }

  return 0;
}

运行效果

本章代码：https://github.com/abcdabcd987/libfabric-efa-demo/blob/master/src/4_hello.cpp

驾驭3200Gbps网络(3): libfabric

2024-12-25T00:00:00-08:00

libfabric 一套通用的高性能网络接口，它的风格与前文提到的 RDMA ibverbs 接口十分相像，但是更易用一点。应用程序只需要调用 libfabric 的上层接口，而具体的协议由不同的 Provider 实现。libfabric 官方的 Provider 包括了 tcp、udp、shm（共享内存）、verbs（即 RDMA ibverbs）以及对本文来说最重要的 efa。

概念及术语

libfabric 定义了非常多的概念和术语，并且包含了很多缩写。这里我用稍不严谨地语言概括一下。

软件对象模型

Fabric: 所有硬件资源和软件状态的的集合，类似于用于保存全局状态的一个结构。
Domain: 类似于一张网卡，例如 eth0。一个 Fabric 可以拥有多个 Domain。
(EP) Endpoint: 数据收发的端点，一个 Domain 可以拥有多个端点，并且每个端点也有不同的类型。例如在同一张以太网卡上，可以监听多个 ip:port，可以使用 TCP 和 UDP 这样不同类型的协议。
(EQ) Event Queue：事件队列，用来报告控制平面操作的完成。在本文中没有用到。
(CQ) Completion Queue: 完成队列，报告数据平面操作的完成（RECV / SEND / WRITE / READ）。CQ 属于 Domain，可以绑定到 Endpoint 上。
(AV) Address Vector：存储解析过的网络地址。需要先将通信的另一方加入 AV 才能对其发起通信。AV 属于 Domain，可以绑定到 Endpoint 上。
(MR) Memory Region：用于收发数据的缓冲区。无论是哪种 RDMA 操作，都需要指定 MR。注册 MR 需要经过操作系统内核，因为操作系统内核需要设置 CPU 页表以及其他 PCIe 设备的页表。

通信方式

端点类型（Endpoint types）
1. FI_EP_MSG（Reliable-connected）：类似于 RDMA 中的 RC QP，基于连接的可靠传输。EFA 不支持这个类型。
2. FI_EP_RDM（Reliable-unconnected）：基于数据报的可靠传输。
  - 包含重传
  - 支持某些类型的消息送达顺序保证。例如 EFA 支持 send-after-send。
  - 可以传输任意大小的数据。可以使用单边 RDMA（One-sided RDMA，WRITE / READ）和双边 RDMA（Two-sided RDMA，RECV / SEND）。
  - EFA 使用亚马逊自研的SRD（Scalable Reliable Datagram）协议，所以这是 EFA 的主力端点类型。
  - verbs Provider 则不原生支持这一端点类型，因为 RDMA 本身并没有类似的 QP 类型。不过 libfabric 上可以通过 RxM (RDM over MSG) Provider 来模拟这一端点类型。
3. FI_EP_DGRAM（Unreliable datagram）：类似于 RDMA 中的 UD QP，基于数据报的不可靠传输。在 EFA 上只能传输小于 MTU 的数据，并且只能使用双边 RDMA。在这里不多做介绍。
端点能力（Endpoint Capability）
1. FI_MSG 双边 RDMA：RECV / SEND，见 fi_msg(3)
2. FI_TAGGED：类似于 FI_MSG，但是每条消息带有一个标记，接收方可以根据这个标记选择缓冲区。本文不多做介绍。见 fi_tagged(3)
3. FI_RMA 单边 RDMA：WRITE / WRITE_IMM / READ，见 fi_rma(3)
4. FI_ATOMIC 单边 RDMA 原子操作。本文不多做介绍。见 fi_atomic(3)

安装 `libfabric`

libfabric 依赖于 rdma-core 来进行 RDMA 操作，依赖于 GDRCopy 来实现 GPUDirect DMA。另外 libfabric 提供了一套性能测试工具 fabtests。下面这一段脚本会在当前目录下面建立一个 build/ 目录，下载这些库的代码到 build/ 目录中，编译并安装这些库到 build/ 的子目录当中。

mkdir -p build
cd build
BUILD_DIR=$(pwd)

# RDMA Core
sudo apt-get install rdma-core

# GDRCopy
wget -O gdrcopy-2.4.4.tar.gz https://github.com/NVIDIA/gdrcopy/archive/refs/tags/v2.4.4.tar.gz
cd gdrcopy-2.4.4/
make prefix="$BUILD_DIR/gdrcopy" \
    CUDA=/usr/local/cuda \
    -j$(nproc --all) all install
cd ..
export LD_LIBRARY_PATH="$BUILD_DIR/gdrcopy/lib:$LD_LIBRARY_PATH"

# libfabric
wget https://github.com/ofiwg/libfabric/releases/download/v2.0.0/libfabric-2.0.0.tar.bz2
tar xf libfabric-2.0.0.tar.bz2
cd libfabric-2.0.0
./configure --prefix="$BUILD_DIR/libfabric" \
    --with-cuda=/usr/local/cuda \
    --with-gdrcopy="$BUILD_DIR/gdrcopy"
make -j$(nproc --all)
make install
cd ..
export LD_LIBRARY_PATH="$BUILD_DIR/libfabric/lib:$LD_LIBRARY_PATH"

# fabtests
wget https://github.com/ofiwg/libfabric/releases/download/v2.0.0/fabtests-2.0.0.tar.bz2
tar xf fabtests-2.0.0.tar.bz2
cd fabtests-2.0.0
./configure --prefix="$BUILD_DIR/fabtests" \
    --with-cuda=/usr/local/cuda \
    --with-libfabric="$BUILD_DIR/libfabric"
make -j$(nproc --all)
make install
cd ..

运行示例程序

现在我们可以运行 libfabric 自带的示例程序，一来可以验证软件已经安装到位，二来可以对我们的硬件有个初步的了解。

获得基本的网卡信息

./build/libfabric/bin/fi_info --verbose

---
fi_info:
    caps: [ FI_MSG, FI_RMA, FI_TAGGED, FI_ATOMIC, FI_READ, FI_WRITE, FI_RECV, FI_SEND, FI_REMOTE_READ, FI_REMOTE_WRITE, FI_MULTI_RECV, FI_LOCAL_COMM, FI_REMOTE_COMM, FI_SOURCE, FI_DIRECTED_RECV ]
    mode: [  ]
    addr_format: FI_ADDR_EFA
    src_addrlen: 32
    dest_addrlen: 0
    src_addr: fi_addr_efa://[fe80::8e7:efff:feee:e81d]:0:0
    dest_addr: (null)
    handle: (nil)
    fi_tx_attr:
        caps: [ FI_MSG, FI_RMA, FI_TAGGED, FI_ATOMIC, FI_READ, FI_WRITE, FI_SEND ]
        mode: [  ]
        op_flags: [ FI_COMPLETION, FI_INJECT, FI_TRANSMIT_COMPLETE, FI_DELIVERY_COMPLETE ]
        msg_order: [ FI_ORDER_SAS, FI_ORDER_ATOMIC_RAR, FI_ORDER_ATOMIC_RAW, FI_ORDER_ATOMIC_WAR, FI_ORDER_ATOMIC_WAW ]
        inject_size: 4096
        size: 4096
        iov_limit: 4
        rma_iov_limit: 1
        tclass: 0x0
    fi_rx_attr:
        caps: [ FI_MSG, FI_RMA, FI_TAGGED, FI_ATOMIC, FI_RECV, FI_REMOTE_READ, FI_REMOTE_WRITE, FI_MULTI_RECV, FI_SOURCE, FI_DIRECTED_RECV ]
        mode: [  ]
        op_flags: [ FI_MULTI_RECV, FI_COMPLETION ]
        msg_order: [ FI_ORDER_SAS, FI_ORDER_ATOMIC_RAR, FI_ORDER_ATOMIC_RAW, FI_ORDER_ATOMIC_WAR, FI_ORDER_ATOMIC_WAW ]
        size: 8192
        iov_limit: 4
    fi_ep_attr:
        type: FI_EP_RDM
        protocol: FI_PROTO_EFA
        protocol_version: 4
        max_msg_size: 18446744073709551615
        msg_prefix_size: 0
        max_order_raw_size: 8776
        max_order_war_size: 8776
        max_order_waw_size: 8776
        mem_tag_format: 0xaaaaaaaaaaaaaaaa
        tx_ctx_cnt: 1
        rx_ctx_cnt: 1
        auth_key_size: 0
    fi_domain_attr:
        domain: 0x0
        name: rdmap79s0-rdm
        threading: FI_THREAD_SAFE
        progress: FI_PROGRESS_AUTO
        resource_mgmt: FI_RM_ENABLED
        av_type: FI_AV_TABLE
        mr_mode: [ FI_MR_VIRT_ADDR, FI_MR_ALLOCATED, FI_MR_PROV_KEY, FI_MR_HMEM ]
        mr_key_size: 4
        cq_data_size: 4
        cq_cnt: 512
        ep_cnt: 256
        tx_ctx_cnt: 256
        rx_ctx_cnt: 256
        max_ep_tx_ctx: 1
        max_ep_rx_ctx: 1
        max_ep_stx_ctx: 0
        max_ep_srx_ctx: 0
        cntr_cnt: 0
        mr_iov_limit: 1
        caps: [ FI_LOCAL_COMM, FI_REMOTE_COMM ]
        mode: [  ]
        auth_key_size: 0
        max_err_data: 0
        mr_cnt: 262144
        tclass: 0x0
    fi_fabric_attr:
        name: efa
        prov_name: efa
        prov_version: 200.0
        api_version: 2.0
    nic:
        fi_device_attr:
            name: rdmap79s0
            device_id: 0xefa1
            device_version: 6
            vendor_id: 0x1d0f
            driver: efa
            firmware: 0.0.0.0
        fi_bus_attr:
            bus_type: FI_BUS_PCI
            fi_pci_attr:
                domain_id: 0
                bus_id: 79
                device_id: 0
                function_id: 0
        fi_link_attr:
            address: EFA-fe80::8e7:efff:feee:e81d
            mtu: 8760
            speed: 100000000000
            state: FI_LINK_UP
            network_type: Ethernet
---
...

fi_info 会输出所有网卡的信息，我在上面贴出来了第一张 EFA 的全部信息以供参考。从中我们可以看到一些关键点，在后面的几篇文章中会用到，我在下面特地摘出来：

src_addrlen: 32                         # EFA 地址长度为32字节
fi_tx_attr:
    size: 4096                          # 发送队列能容纳4096个操作
    iov_limit: 4                        # 发送操作最多能指定4个本地缓冲区
    rma_iov_limit: 1                    # 发送操作只能指定一个远端缓冲区
fi_rx_attr:
    size: 8192                          # 接收队列能容纳8192个操作
    iov_limit: 4                        # 接收操作最多能指定4个本地缓冲区
fi_ep_attr:
    max_msg_size: 18446744073709551615  # 可以发送任意大的数据
fi_domain_attr:
    name: rdmap79s0-rdm                 # domain 名称
    mr_key_size: 4                      # MR 远程密钥大小为4字节
    cq_data_size: 4                     # WRITE_IMM 携带的立即数为4字节
    mr_iov_limit: 1                     # 注册MR时一次只能指定一块缓冲区
nic:
    fi_device_attr:
        name: rdmap79s0                 # 网卡名称
        driver: efa                     # EFA 网卡
    fi_link_attr:
        mtu: 8760                       # 单个数据包最大为8760字节
        speed: 100000000000             # 网卡的速率为 100 Gbps

带宽测试

我们还可以运行 fabtests 当中的带宽测试。下面的这个命令可以在两台机器之间进行 GPUDirect RDMA WRITE 操作。可以看到，最高速度达到了 11843.86 MB/sec，也就是 94.751 Gbps，几乎打满了带宽。

ip-172-19-226-174$ ./build/fabtests/bin/fi_rma_bw -p efa -o write -E -D cuda -S all
ip-172-19-230-131$ ./build/fabtests/bin/fi_rma_bw -p efa -o write -E -D cuda -S all 172.19.226.174
bytes   iters   total       time     MB/sec    usec/xfer   Mxfers/sec
1       20k     19k         0.03s      0.75       1.33       0.75
2       20k     39k         0.03s      1.58       1.27       0.79
3       20k     58k         0.03s      2.36       1.27       0.79
4       20k     78k         0.03s      3.18       1.26       0.79
6       20k     117k        0.03s      4.72       1.27       0.79
8       20k     156k        0.03s      6.37       1.26       0.80
12      20k     234k        0.03s      9.48       1.27       0.79
16      20k     312k        0.03s     12.66       1.26       0.79
24      20k     468k        0.03s     19.18       1.25       0.80
32      20k     625k        0.03s     24.78       1.29       0.77
48      20k     937k        0.02s     38.75       1.24       0.81
64      20k     1.2m        0.02s     51.87       1.23       0.81
96      20k     1.8m        0.02s     77.48       1.24       0.81
128     20k     2.4m        0.02s    103.46       1.24       0.81
192     20k     3.6m        0.02s    155.55       1.23       0.81
256     20k     4.8m        0.02s    206.08       1.24       0.80
384     20k     7.3m        0.02s    307.40       1.25       0.80
512     20k     9.7m        0.02s    412.82       1.24       0.81
768     20k     14m         0.03s    614.11       1.25       0.80
1k      20k     19m         0.03s    814.12       1.26       0.80
1.5k    20k     29m         0.03s   1212.41       1.27       0.79
2k      20k     39m         0.03s   1590.49       1.29       0.78
3k      20k     58m         0.03s   2340.84       1.31       0.76
4k      20k     78m         0.03s   3068.05       1.34       0.75
6k      20k     117m        0.03s   4462.85       1.38       0.73
8k      20k     156m        0.03s   5440.12       1.51       0.66
12k     20k     234m        0.04s   6770.81       1.81       0.55
16k     20k     312m        0.04s   7921.86       2.07       0.48
24k     20k     468m        0.06s   8319.43       2.95       0.34
32k     20k     625m        0.07s   9470.25       3.46       0.29
48k     20k     937m        0.10s  10127.33       4.85       0.21
64k     2k      125m        0.01s  10465.67       6.26       0.16
96k     2k      187m        0.02s  10885.78       9.03       0.11
128k    2k      250m        0.03s  10215.66      12.83       0.08
192k    2k      375m        0.03s  11301.58      17.40       0.06
256k    2k      500m        0.05s  11501.58      22.79       0.04
384k    2k      750m        0.08s  10123.87      38.84       0.03
512k    2k      1000m       0.10s  10749.66      48.77       0.02
768k    2k      1.4g        0.13s  11843.86      66.40       0.02
1m      200     200m        0.02s  10968.94      95.60       0.01
1.5m    200     300m        0.03s  10837.99     145.12       0.01
2m      200     400m        0.04s  10204.87     205.51       0.00
3m      200     600m        0.06s  10522.41     298.95       0.00
4m      200     800m        0.08s  10826.94     387.39       0.00
6m      200     1.1g        0.11s  11488.20     547.65       0.00
8m      200     1.5g        0.16s  10504.60     798.57       0.00

现在我们已经验证了软硬件都已就绪，那么下一章我们就可以开始编写代码了。

驾驭3200Gbps网络(2): 高性能网络系统设计哲学

2024-12-25T00:00:00-08:00

在介绍 libfabric 之前，让我们先来从更高的层级思考高性能网络库的设计，这也方便我们理解 ibverbs 以及 libfabric 的接口和编程模型。在这个话题上，其实 libfabric 官方的介绍文章写得非常好，我建议感兴趣的读者也可以仔细阅读。我在2021年初的时候用 ibverbs 写过一个简易的 RDMA 通信程序，当时有很多类似的想法，因此当我时隔3年读到这篇文章的时候就产生了强烈的共鸣。

套接字

首先让我们回顾一下套接字（Socket）的 send() 和 recv() 接口：

ssize_t send(int sockfd, const void* buf, size_t len, int flags);
ssize_t recv(int sockfd, void* buf, size_t len, int flags);

应用程序最关心的是把 [buf, buf+size) 这一段数据传到 sockfd 代表的目的地，而这套接口就简单到几乎只包含这些信息。从1983年发明以来，这套接口就未经大改地保留到了现在。简洁和通用是其成功秘诀。这套接口隐藏了底下网络协议的所有细节，所有的脏活累活都包在操作系统内核中了。

这套接口在设计之初就考虑了阻塞式（Blocking）以及非阻塞式(Non-Blocking)的调用方式。从阻塞模式切换到非阻塞模式只需要一行代码：

fcntl(sockfd, F_SETFL, O_NONBLOCK);

在使用阻塞模式时，这套接口的调用直到有部分消息发送成功（或者接收到部分消息）才会返回。在使用非阻塞模式时，这套接口总是会立刻返回，返回值即是发送了（或者接收了）多少字节。搭配上 select() / poll() / epoll / io_uring 这些 I/O 复用（I/O Multiplexing）的工具，很容易地就能编写出能异步处理大量请求的网络程序。

网络的异步天性

异步是网络的天性。一个网络操作能否成功远不是 Linux 内核能控制的，比如硬件的缓冲区满了需要暂停，出现了拥塞需要控制，出现了丢包需要重传。而且这些网络操作是否成功，完全不是在一次函数调用内可以知道的。一个数据包可能要过几百毫秒才能到达大洋彼岸。想象一下如果让我们来实现 Linux 内核的网络模块，我们应该如何做呢？

如果应用使用的是阻塞模式，一切都好说。等一切尘埃落定再唤醒应用即可。但是如果应用使用的是非阻塞模式，内核需要立即返回 send() 成功发送了多少字节（或者 recv() 收到了多少字节）。可是在不挂起用户态进程的前提下，内核可能甚至连网卡的硬件资源还有空余没有都不知道，更别说知道能有多少字节在几百毫秒后漂洋过海。

在这样的前提下，我能想到的一个可行的实现方式，便是让内核自己维护一份缓冲区。当用户态进程调用 send() 时，内核将用户给的内容复制到自己的缓冲区。因为内核十分清楚自己的缓冲区还剩多大，所以内核可以立即返回到用户态进程。而在之后，内核可以从自己的缓冲区中慢慢地与网络发生异步交互。

显然在高性能网络中我们是想要尽量避免额外的复制的。有没有别的办法呢？比如让应用提供这个缓冲区，而不是在内核维护。这个想法如何呢？

缓冲区的复制

假设 Linux 内核允许用户态进程提供缓冲区，那么用户态进程必须保证在内核通知它网络操作完成之前，一直不去修改这个缓冲区，也不会释放掉这个缓冲区。否则，要么传输的数据产生损坏，要么进程崩溃，甚至还可能带来安全隐患。显然这与 Linux 的设计哲学是相悖的，内核不应该信任用户态。

我再进一步地想象了一下，硬要这么做也不是实现不了。禁止用户态进程访问某段内存区间倒不是什么难事，内核只需要在 CPU 的页表（Page Table）上动动手脚就行了。然而如果真的这么做了，反而会进一步带来更多的问题。比如要操作页表的话，就需要缓冲区对其到页表大小。另外，操作页表也不是一件那么快的事情。频繁地修改页表和页表缓存也会降低执行效率。

总之，我觉得让内核完全信任用户态进程提供的缓冲区还是一件不太可行的事情。让内核维护自己维护一份缓冲区反而容易很多，只是需要多一次复制而已。

因此，我认为这多出来的一次复制，正是套接字的接口设计带来的必然的结果。这一次额外的复制与套接字简单易用的设计哲学是一体两面的。

地址解析

libfabric 的介绍文章中还提到了一点我之前没有想到过的——地址解析。套接字拿到的地址是 IP 地址，而要进行通信，内核还需要查找更底层的地址，也就是 MAC 地址。这一步地址解析也是需要花费一定时间的。对于 TCP 来说，可以在建立连接的时候把底层的地址缓存好，因此影响不大。然而对于 UDP 来说，因为它是没有连接的，所以每一次的收发都需要经过地址转换，花费不少时间。

高性能网络

那么要实现高性能网络，应用程序、网络接口、操作系统以及硬件需要有哪些默契的配合呢？

缓冲区的所有权

前面我们提到，让应用程序提供缓冲区可以避免额外的复制，但是前提是应用程序在网络操作结束之前都不应该再去读写这段缓冲区。这种软件层面的契约，按照现在时髦的 Rust 语言话来说，就是个所有权的问题。在调用网络操作之后，缓冲区的所有权应该从应用程序转移到网络接口。

这种对用户态的高度信任无疑是不被 Linux 内核所接受的。那么我们一定要事事争得内核的同意吗？可以让用户态进程直接操纵网卡吗？有没有优雅的办法既保证系统的安全又能够达成这种契约？

绕过操作系统内核

套接字必须经过操作系统内核的原因是只有内核才能操纵网卡，整个网络栈都是在内核里面实现的。而如果反过来，用户态和网卡一起实现了网络栈，那么网络的调用就不一定要经过内核了。如果数据不用经过内核，那么性能肯定是会大大增加的，因为用户态和内核态的切换非常花费时间。

应用与网卡共享内存

假设操作系统内核允许应用程序和网卡一同共享某一段内存空间，那么所有的数据操作都可以变得非常快速。在发送数据时，网卡可以直接从内存读取数据；在接收数据时，网卡可以直接将数据写入到内存中；应用程序可以轮询（Poll）网卡的完成队列（Completion Queue）来知道网络操作是否完成，而这个轮询操作只不过是读取某段虚拟地址而已。

控制平面与数据平面分离

上面说的这些听起来很美好，但是如果一切都绕过操作系统内核，让用户态对硬件有完全的访问权限，那么操作系统的抽象就被破坏了，安全性更是无从保障。实际上，我们完全可以把网络操作分为两类：

控制平面（Control Plane）：
- 控制平面的网络操作可以经过操作系统内核，由内核来执行高权限的操作，以确保系统的安全性。这是一条慢速路径（Slow-Path）。
- 例如：开启监听状态、地址解析、注册缓冲区、注册完成队列
数据平面（Data Plane）：
- 数据平面的网络操作应该绕开操作系统内核，这是一条快速路径（Fast-Path）。
- 例如：接收数据、发送数据、轮询完成队列

既然我们要让应用程序提供缓冲区，那么我们可以要求应用程序提前告诉操作系统内核这是哪一段虚拟地址。内核此时可以在 CPU 的页表上做相应的标记。同时，内核也可以设置网卡的页表。这样一来，用户态进程和网卡硬件都可以访问同一段内存，并且内核已经提前设置好了双方的访问权限。

假设应用程序不遵守所有权的约定，那么破坏面最多只是使得应用程序本身崩溃，无法扩散到用户态以外，不会影响系统的其他进程。

接收先于发送

没有了内核提供缓冲区，当网卡收到一段消息的时候，网卡怎么知道应该写入应用程序的哪一段地址呢？

要解决这个问题，我们可以要求应用程序先提交几个接收操作。当新数据到来时，网卡就可以把数据写到其中一个接收操作指定的内存中。如果在新的数据到来时，没有等待中的接收操作，那么网卡就可以直接拒绝这段数据。反过来说，如果应用程序想要继续接收更多的数据，应用程序有责任在接收操作结束之后，再提交一个接收操作。

TCP 及 RDMA 的比较

上面这张图简单描述了 TCP 和 RDMA 传输数据的流程。

左侧是 TCP。消息首先要从用户态复制到内核，然后内核才能驱动网卡发送消息。在接收端，网卡将数据写到内核，内核再将数据复制到用户态。由于有内核的参与，以及多了额外的缓冲区复制，TCP 的这一套流程性能是比较低的。

右侧是 RDMA。用户态可以直接告诉网卡要发送的消息的内存地址，网卡可以直接读取这段数据发送到远端。同时，用户态程序可以读取完成队列，从而知道发送操作是否完成。在操作完成了之后，应用就可以重新使用这一段内存。在接收端，网卡直接将数据写入等待中的接收操作所指定的内存区域，然后网卡更新完成队列。当接收端的应用程序读取到完成队列时，它就知道数据到达了。RDMA 的这一套流程无疑是十分高效的。

如果考虑到显卡，那么 TCP 和 RDMA 的性能差别就更大了。上图便是从本地显卡到远端显卡的消息传输的示意图。

左侧是 TCP。因为操作系统没法直接读取显存，所以这里额外多了一步从显存到用户态内存的复制。在接收端，同理也是多了一次从用户态内存到显存的复制。这让 TCP 本就不高的性能更是雪上加霜。

右侧是 RDMA。当应用程序要发送数据的时候，用户态进程只需要告诉网卡这个消息所在的虚拟地址，网卡就能直接读取显存。这一技术叫做 GPUDirect RDMA。值得注意的是，用户态进程在发送消息的时候，并不需要刻意区分到底这个虚拟地址是在 CPU 内存上的，还是在显卡的显存上的，因为内核已经提前设置好了 CPU、显卡以及网卡的页表，使得他们能够理解同一个虚拟地址空间。接收端也是类似的，网卡直接将数据写入显存，无需经过任何额外的复制。

硬件总线的拓扑结构

在上面的例子中我们可以发现，RDMA 不仅节省了软件层面的开销，也节省了 PCIe 总线以及内存总线上的数据传输。当一台机器上有更多的显卡以及更多的网卡的时候，我们需要关心硬件总线的拓扑结构，从而避免不必要的跨总线数据传输，避免总线上的阻塞。

以 AWS p5 实例为例，上图画出了系统的拓扑结构。每台机器插了两块 CPU。每块 CPU 下面连接着4个 PCIe 交换机。每个 PCIe 交换机下面挂载了4张 100 Gbps 的 EFA 网卡，1张 NVIDIA H100 显卡，以及一块 3.84 TB 的 NVMe SSD。

如上图所示，如果我们将同一个 PCIe 交换机下面的显卡和网卡配对，那么在 GPUDirect RDMA 的帮助下，他们之间的通信就只需要经过这个 PCIe 交换机，可以充分地享受 PCIe Gen5 x16 的带宽。同时，他们也不会影响到其他 PCIe 交换机，也不会影响到与 CPU 内存的总线。

而如果使用 TCP 进行传输，如上图所示，则网卡和显卡之间的通信需要绕一大圈。数据要从 GPU 显存经过 PCIe 交换机，再经过 CPU 内部的 PCIe 通道，再经过 CPU 的内存总线，才能复制到用户态内存中。接着内核会对这段数据再进行一次复制。然后内核上的数据需要经过 CPU 内存总线，经过 CPU 内部的 PCIe 通道，经过 PCIe 交换机，最终才能到达网卡。而当8张显卡32块网卡同时工作时，TCP 这样的传输显然会带来 PCIe 总线和内存总线上严重的阻塞。

驾驭3200Gbps网络(1): RDMA和EFA

2024-12-25T00:00:00-08:00

在开始写代码操作 3200 Gbps 网络之前，让我们先了解一下驱动这个高性能网络所需要的硬件、软件以及系统设计哲学。

RDMA

我们日常使用的网络大多是基于 TCP/IP 协议的网络。应用程序通过套接字（Socket）与 Linux 内核交互。Linux 内核维护了完整的 TCP/IP 协议栈，以及负责操纵网卡（NIC，Network Interface Controller）。两个网卡之间通过以太网（Ethernet）进行网络层的通信。大家比较日常能见到的就是家用机中的千兆网卡（1 Gbps）以及服务器上的万兆网卡（10 Gbps）。

高性能网络则是另一套完全不同的软硬件技术栈，称为 RDMA（远程直接内存访问，Remote Direct Memory Access）。

应用层：应用程序一般通过 ibverbs API（也称为 verbs API）与 RDMA 网卡进行交互。verbs API 与套接字有非常大的区别，背后是完全不同的系统设计哲学，我们会在后文提到。
协议层：不同于 TCP/IP，RDMA 有一套自己的协议。因为我不是做这方面的研究的，所以具体的技术细节我也不了解。可能对我来说最直观的感受就是，RDMA 没有“IP 地址”和“MAC 地址”，而是另有一套地址系统。
网络层：RDMA 最常用的传输层是 InfiniBand 和 RoCE。InfiniBand 是私有协议，比较贵，也比较稳定。RoCE（基于融合以太网的远程直接内存访问，RDMA over Converged Ethernet）性价比比较高，但需要细致地调优，很多科技大厂也发表了不少关于 RoCE 的论文。
网卡：最常见的 RDMA 网卡就是 Mellanox 的 ConnectX 系列网卡了。在2020年英伟达收购了 Mellanox 之后，英伟达更是坐实了当年网络上“N卡网速快”的梗。ConnectX-7 这一代网卡达到了 400 Gbps 的网速，常常与 H100 这一代的显卡进行搭配。

顾名思义，RDMA 能直接对远程的（以及本地的）内存进行读写。而且这里指的内存也不止是绑定在 CPU 上的内存，也可以是在 PCIe 总线上其他设备的内存，例如 GPU 的显存。相比起套接字的 send() 和 recv() 两种操作，RDMA 可以做的操作更丰富一点。RDMA 的操作可以分为两类，双边 RDMA 和单边 RDMA：

双边 RDMA（Two-sided RDMA）：需要通信双方的 CPU 的参与
1. RECV：目标方准备接受消息
2. SEND：发起方将信息发送给目标方
单边 RDMA（One-sided RDMA）：只需要发起方的 CPU 参与，目标方 CPU 并不知情
1. WRITE：将数据从发起方的内存直接写入目标方的内存
  - WRITE_IMM：类似 WRITE，将数据从发起方的内存直接写入目标方的内存，但额外地在目标方的完成队列（Completion Queue）中插入一个整数，用来通知目标方 CPU。
2. READ：直接从目标方的内存读取数据并写入发起方的内存
3. ATOMIC：对目标方的内存进行原子操作，如 Compare-and-Swap、Fetch-Add

在 RDMA 中，一对通信节点被称为队列对（Queue Pair，QP）。类似于 IP 协议上常用的两种传输类型 TCP 和 UDP，RDMA 上有三种传输类型：

可靠连接队列对（Reliable Connected Queue Pair, RC QP）
- 基于连接，因此在通信之前需要先建立连接。通常来说连接数有限制。
- 可靠传输，包含重传，保证消息送达顺序。
- 可以一次性传输很大的消息，消息的大小可以远超 MTU。
- 支持所有 RDMA 操作：RECV、SEND、WRITE、READ、ATOMIC
不可靠连接队列对（Unreliable Connected Queue Pair, UC QP）
- 类似 RC，但是传输不可靠，不保证消息送达，不保证消息送达顺序。
- 只支持部分 RDMA 操作：RECV、SEND、WRITE
不可靠数据报队列对（Unreliable Datagram Queue Pair, UD QP）
- 无需建立连接。
- 不保证消息送达，不保证消息送达顺序。
- 消息必须小于 MTU
- 只支持 RECV、SEND

AWS EFA

显然大多数的个人甚至公司都是买不起 NVIDIA HGX 集群的，只能依赖于云服务提供商。云服务提供商一般不喜欢把硬件直接暴露给租户，而是会在上面加一层虚拟化。在 AWS p5 实例上，AWS 提供的虚拟高性能网卡称为 EFA（Elastic Fabric Adapter）。

至于 EFA 底下用的是亚马逊自研的网卡还是标准化的产品，再底下的网络协议是 Infiniband 还是 RoCE，再底下的物理介质是光纤还是铜缆，这些我们一概不知。当然我们也不需要知道，有问题找客服就行了。不过有意思的是，EFA 使用了亚马逊的自研协议 SRD（Scalable Reliable Datagram）。

那么应用程序应该如何使用 EFA 呢？EFA 的文档中有下面一张图：

从这张图以及文档中我们可以知道，我们需要通过一个叫做 libfabric来使用 EFA。

驾驭3200Gbps网络(0): 导言

2024-12-25T00:00:00-08:00

今年早些时候我有幸入职了 Perplexity AI，终于用上了最强配置的服务器——AWS p5 实例，上面搭载了8张 NVSwitch 互联的 NVIDIA H100 显卡。更令我兴奋的是，服务器之间搭载了 3200 Gbps 的超高速网络。我觉得要是我能写一个程序用上这 3200 Gbps 的带宽，一定是一件非常炫酷的事情！

最近我花了一周的时间，大概摸到了一些门道，写了个小小的概念验证程序，用上了 97% 的带宽。因为我觉得这个摸索的过程挺有意思的，再加上网上关于 RDMA、EFA、libfabric、高性能网络的文章和教程十分有限，所以我打算把我这一周学到的知识分享出来。既是一个记录，也可以当作一个入门教程来看。

对 MLSys 熟悉的朋友们可能要问了：这不是 PyTorch 或者 NCCL 一行代码就能搞定的事情吗？确实，NCCL 在集体通信（Collective Communication）方面已经非常成熟了，也是大语言模型的训练和推理的基石。然而在其他应用场景下，我觉得集体通信还是有一些不太适合的地方：

集体通信需要建立起全局通信域（MPI World）。如果要动态地增加、减少或者替换集群中的节点，那么就需要先让整个集群停下来。
集体通信采用了同步通信模型，不论实现方式是阻塞式的还是非阻塞式的，对我来说都是一种很强的心智负担。我更习惯的是像 gRPC 那样的异步通信模型。
最重要的是，能自己造一个轮子不是很好玩吗？

因为我的实验环境是 AWS p5 集群，所以本文提到的一些技术细节可能只适用于 AWS p5 集群。不过我希望本文还是能对其他的高性能网络环境有一定的参考价值。

因为内容比较多，所以我把内容拆成了几篇文章，欢迎大家点击阅读：

驾驭3200Gbps网络(0): 导言
驾驭3200Gbps网络(1): RDMA和EFA
驾驭3200Gbps网络(2): 高性能网络系统设计哲学
驾驭3200Gbps网络(3): libfabric
驾驭3200Gbps网络(4): 单向接收发送
驾驭3200Gbps网络(5): 双向接收发送
驾驭3200Gbps网络(6): GPUDirect RDMA WRITE
驾驭3200Gbps网络(7): 操作队列及带宽测试 [97.433 Gbps (97.4%)]
驾驭3200Gbps网络(8): 总线拓扑
驾驭3200Gbps网络(9): 使用32张网卡 [287.089 Gbps (9.0%)]
驾驭3200Gbps网络(10): 测试前预热 [293.461 Gbps (9.2%)]
驾驭3200Gbps网络(11): 多线程 [355.301 Gbps (11.1%)]
驾驭3200Gbps网络(12): 绑定CPU核心 [1237.738 Gbps (38.7%)]
驾驭3200Gbps网络(13): 状态分片 [1522.567 Gbps (47.6%)]
驾驭3200Gbps网络(14): 批量提交操作 [2589.488 Gbps (80.9%)]
驾驭3200Gbps网络(15): 惰性提交操作 [3108.283 Gbps (97.1%)]

在我的探索及写作过程中，得到了下面这些朋友的帮助，特此感谢：Brian Barrett、Shi Jin。同时我也非常感谢公司愿意让我花时间捣鼓这些技术，以及提供了这么强的硬件环境。

以下参考资料是我觉得非常有用的：

libfabric 官方文档，尤其是介绍性的 fi_intro(7) 和 fi_arch(7)。在使用具体的 API 细节时，可以参阅各个 API 的文档。
aws-ofi-nccl 的代码。这是 AWS 的 NCCL 插件，是所有 AWS 上运行的多机机器学习任务的通信基础。这个代码库中有很多关于 libfabric 及 efa Provider 的实际应用例子。
fabtests 的代码。这是 libfabric 的测试工具，可以用来测试 libfabric 的性能。

文章的代码可以在 GitHub 中找到：https://github.com/abcdabcd987/libfabric-efa-demo

多个大语言微调模型并行推断的潜力

2023-09-11T00:00:00-08:00

随着开源预训练大型语言模型（Large Language Model, LLM ）变得更加强大和开放，越来越多的开发者将大语言模型纳入到他们的项目中。其中一个关键的适应步骤是将领域特定的文档集成到预训练模型中，这被称为微调。

通常情况下，来自领域特定文档的额外知识与预训练模型已经知道的相比微不足道。在这种情况下，低秩适应（Low-Rank Adaptation，LoRA ）技术证明是有价值的。

通过 LoRA，微调模型仅向预训练模型添加不到0.1%的参数。具体来说，这意味着 LoRA 微调模型仅增加了10~200 MB 的存储，具体取决于配置。从计算角度来看，考虑到与预训练模型相比参数的增加极少，额外的计算负载相对较小。

基于存储和计算的额外开销都很小这一点，我相信构建一个多租户的大语言微调模型的推断服务具有很大潜力。这个服务可以托管成千上万个 LoRA 模型，它们都共享相同的预训练大语言模型。在每个批次的执行中，每个用户请求都会调用一个独立的微调模型，从而分摊存储和计算成本到各种不同的模型中。

在我的上一篇博客文章中，我深入探讨了大语言模型推断中的批处理效应。在这篇文章中，我将详细介绍为什么多租户 LoRA 推断服务具有巨大的潜力。

背景知识：文本生成

文本生成服务，如ChatGPT，接受用户文本输入并提供文本响应。这个输入被称为“提示”（prompt）。在内部，当大语言模型处理文本时，它在一系列“词元”（token）上操作。我们可以大致将词元视为几个字符或一个单词。文本生成过程有两个主要阶段：

预填充阶段（或称“编码”，“初始化”）接受整个提示并生成随后的词元以及一个“键值缓存”(KV Cache)。
解码阶段处理新生成的词元和键值缓存，然后生成下一个词元，同时更新键值缓存。这个阶段不断重复，直到模型完成其输出。

有趣的是，尽管预填充阶段处理的词元数量比解码阶段多100倍，但它们的计算延迟是在同一数量级的。由于解码阶段会重复执行上百次，因此在这篇文章中，我将集中讨论如何优化解码阶段。

背景知识：大语言模型的架构和批处理

在其核心，大语言模型的架构非常简单。它主要包括多个 Transformer 层，所有层都共享相同的架构。每一层包括四个计算密集型组件： QKV 投影、自注意力、输出投影和前馈网络（Feed-Forward Network, FFN ）。

概括地说，其中包含了两种算子：

自注意力（Self-Attention，黄色标出）涉及矩阵-矩阵乘法。
密集投影（Dense Projection，绿色标出）涉及向量-矩阵乘法。

考虑到每个批次中每个序列只有一个词元，密集投影计算也非常微小，不足以充分利用 GPU 。因此，扩大批处理大小几乎不会影响密集投影的延迟，这使得大批处理大小对于构建高吞吐、低延迟的推断服务至关重要。

关于大语言模型推断中的批处理的更详细分析，请查看我的以前的博客文章。

背景知识： LoRA

给定形状为[H1, H2]的预训练参数矩阵W，LoRA微调训练一个形状为[H1, R]的小矩阵A和形状为[R, H2]的B。我们使用(W+AB)作为微调模型的权重。这里的R是 LoRA 微调指定的秩，通常远小于原始维度（>= 4096），通常在8~32之间。

这种方法背后的逻辑是，与原始权重相比，新增加的知识只占一小部分，因此 LoRA 将增量压缩成两个低秩矩阵。

与完全微调相比，LoRA 预训练显著降低了微调模型的存储和内存需求。

在大语言模型中，由于所有参数都位于密集投影中，LoRA 可以集成到 Transformer 层中的任何位置。虽然HuggingFace PEFT库仅将 LoRA 加到q_proj和v_proj，但一些研究，如QLoRA，主张将其包含在所有的密集投影中。

LoRA 延迟和批处理效应

尽管在存储方面，LoRA 矩阵明显小于原始权重矩阵，但计算的延迟并不成比例地减少。我们可以使用以下代码对骨干模型和 LoRA 增量的延迟进行基准测试：

h1 = 4096
h2 = 11008
r = 16
for bs in range(1, 33):
  w = torch.randn(h1, h2, dtype=torch.float16, device="cuda:0")
  a = torch.randn(bs, h1, r, dtype=torch.float16, device="cuda:0")
  b = torch.randn(bs, r, h2, dtype=torch.float16, device="cuda:0")
  x = torch.randn(bs, 1, h1, dtype=torch.float16, device="cuda:0")
  bench(lambda: x @ w)
  bench(lambda: x @ a @ b)

上图表明，LoRA 增量仅比骨干模型快2.5倍。

然而，显然 LoRA 批处理效应与骨干模型的效应相似。批处理大小的增加仅在较小程度上影响延迟。这个特性使得多租户 LoRA 非常可行。

在骨干模型中，所有一个批次中的所有请求都针对同一模型。而在多租户 LoRA 服务中，一个批次中的请求可能会调用不同的 LoRA 微调模型。数学表达如下：

\[\begin{pmatrix} \vec{y_1} \\ \vec{y_2} \\ \vdots \\ \vec{y_b} \end{pmatrix} := \begin{pmatrix} \vec{x_1} \\ \vec{x_2} \\ \vdots \\ \vec{x_b} \end{pmatrix} W + \begin{pmatrix} \vec{x_1}A_1B_1 \\ \vec{x_2}A_2B_2 \\ \vdots \\ \vec{x_b}A_bB_b\end{pmatrix}\]

挑战在于以批处理的方式将不同的 LoRA 增量应用于一个批处理中的各个输入，与此同时还要维持“白吃的午餐”式的批处理效应。

批处理 LoRA 算子

我们想要的批处理 LoRA 算子具有以下函数签名：

def add_lora(
    y: torch.Tensor,      # (batch_size, 1, out_features)
    x: torch.Tensor,      # (batch_size, 1, in_features)
    A: torch.Tensor,      # (num_loras, in_features, lora_rank)
    B: torch.Tensor,      # (num_loras, lora_rank, out_features)
    I: torch.LongTensor,  # (batch_size,)
):
  """Semantics: y[i] += x[i] @ A[I[i]] @ B[I[i]]"""
  raise NotImplementedError()

一个最简单的实现方法是在批处理维度上进行循环：

def lora_loop(
    y: torch.Tensor,  # (batch_size, 1, out_features)
    x: torch.Tensor,  # (batch_size, 1, in_features)
    A: torch.Tensor,  # (num_loras, in_features, lora_rank)
    B: torch.Tensor,  # (num_loras, lora_rank, out_features)
    I: torch.LongTensor,  # (batch_size,)
):
  for i, idx in enumerate(I.cpu().numpy()):
    y[i] += x[i] @ A[idx] @ B[idx]

让我们对循环版本进行基准测试。为了进行比较，我们可以包括一个“作弊”实现，其中我们假设批处理中每个请求的 LoRA 矩阵已经被合并在一起。这样，我们只测量批量矩阵乘法（Batched Matrix Multiplication, bmm）的延迟。

def lora_cheat_bmm(
    y: torch.Tensor,  # (batch_size, 1, out_features)
    x: torch.Tensor,  # (batch_size, 1, in_features)
    cheat_A: torch.Tensor,  # (batch_size, in_features, lora_rank)
    cheat_B: torch.Tensor,  # (batch_size, lora_rank, out_features)
):
  y += x @ cheat_A @ cheat_B


num_loras = 50
h1 = 4096
h2 = 11008
r = 16
A = torch.randn(num_loras, h1, r, dtype=torch.float16, device="cuda:0")
B = torch.randn(num_loras, r, h2, dtype=torch.float16, device="cuda:0")
for bs in range(1, 33):
  x = torch.randn(bs, 1, h1, dtype=torch.float16, device="cuda:0")
  y = torch.randn(bs, 1, h2, dtype=torch.float16, device="cuda:0")
  I = torch.randint(num_loras, (bs,), dtype=torch.long, device="cuda:0")
  cheat_A = A[I, :, :]
  cheat_B = B[I, :, :]

  bench(lambda: lora_loop(y, x, A, B, I))
  bench(lambda: lora_cheat_bmm(y, x, cheat_A, cheat_B))

可预见的是，循环版本明显较慢，并且失去了批处理效应。这是因为它逐个处理输入，而不是利用为批处理数据设计的高效 CUDA 核心。

然而，bmm方法提供了一个很好的启发。我们的目标变得清晰起来：首先将所有 LoRA 矩阵汇总到一个临时的张量中，然后使用bmm。经过一番挖掘，我发现了torch.index_select()函数，它可以高效地执行批量汇总(Batched Gather)。于是我们可以如下一个gbmm（ gather-bmm ）实现：

def lora_gbmm(
    y: torch.Tensor,  # (batch_size, 1, out_features)
    x: torch.Tensor,  # (batch_size, 1, in_features)
    A: torch.Tensor,  # (num_loras, in_features, lora_rank)
    B: torch.Tensor,  # (num_loras, lora_rank, out_features)
    I: torch.LongTensor,  # (batch_size,)
):
  a = torch.index_select(A, 0, I) # (batch_size, in_features, lora_rank)
  b = torch.index_select(B, 0, I) # (batch_size, lora_rank, out_features)
  y += x @ a @ b

BGMV 算子

虽然gbmm非常有效，但它并不是最终解决方案。我们没有必要仅仅是因为 bmm 需要连续的存储而将 LoRA 增量而汇总到一个连续的空间中。理想情况下，聚合可以在 CUDA 核心内部进行，于bmm操作同事进行。如果可能的话，这将消除与torch.index_select()相关的 GPU 内存读写操作。

我请叶子豪大牛帮忙，他是精通高性能 CUDA 核心的编写。经过几轮迭代，子豪开发了一个非常快的 CUDA 程序，把 LoRA 所需的计算分成两半。我们将这个算子命名为 BGMV (Batched Gather Matrix-Vector Multiplication)：

有了 BGMV 算子，add_lora()函数可以使用两次 BGMV 调用来实现：

def lora_bgmv(
    y: torch.Tensor,  # (batch_size, 1, out_features)
    x: torch.Tensor,  # (batch_size, 1, in_features)
    A: torch.Tensor,  # (num_loras, in_features, lora_rank)
    B: torch.Tensor,  # (num_loras, lora_rank, out_features)
    I: torch.LongTensor,  # (batch_size,)
):
  tmp = torch.zeros((x.size(0), A.size(-1)), dtype=x.dtype, device=x.device)
  bgmv(tmp, x, A, I)
  bgmv(y, tmp, B, I)

这是gbmm和bgmv的基准结果：

如图所示，gbmm非常有效。聚合过程相对于bmm增加了约20%的延迟，但保持了令人满意的批处理效应。而子豪编写的 bgmv 算子更加高效，甚至超过了bmm。

多租户 LoRA 文本生成性能

在bgmv算子的基础上，我开发了一个名为Punica的实验项目，支持多个 LoRA 模型。 Punica 的独特能力是能把不同 LoRA 模型的请求合并在一个批处理中。我对比了 Punica 和一系列知名系统的性能，包括HuggingFace Transformers、DeepSpeed、Faster Transformer和vLLM。

在测试中，每个请求都针对不同的 LoRA 模型。鉴于其他系统没有明确针对多租户 LoRA 服务进行优化，它们的批处理大小为1。我使用了HuggingFace PEFT库来把 LoRA 加到 HuggingFace Transformers 和 DeepSpeed 中。我尚未调整 vLLM 和 Faster Transformer，所以它们在没有 LoRA 的情况下运行。以下是结果：

DeepSpeed、vLLM 和 Faster Transformers 都有高度优化的 Transformer 实现，这些系统的吞吐量比标准的 HuggingFace Transformers 高出3倍。但由于批处理大小被限制在1，它们在多租户 LoRA 服务效率方面表现不佳。

相反，Punica 在批处理大小为16时比这些系统高出8倍，与普通的 HuggingFace Transformers 相比，甚至高达23倍。值得注意的是，Punica 的吞吐量几乎与批处理大小呈线性关系。

值得注意的是，Punica 仍然是一个早期的研究原型。它目前使用与普通的 HuggingFace Transformers 相同的 Transformer 实现，除了 LoRA 和自注意力运算符。一些已知的 Transformer 层的优化在 Punica 中尚未实现。这解释了在批处理大小为1时 Punica 与其他高度优化的系统之间的性能差距。

吞吐量不错，那么从延迟方面表现如何？

如图所示，Punica 中的批处理不会引入显著的延迟。

示例用途

在展示了多租户 LoRA 服务的高效性后，让我们设想一些潜在的应用场景：

用一本新的小说对 LoRA 模型进行细调，以帮助读者总结每个角色的旅程。
针对迅速发展的新闻对 LoRA 模型进行调整，以让读者了解最新动态。
基于网页内容对 LoRA 模型进行优化，提高读者的理解能力。

我将这种方法称为“即时细调”(Just-in-time Fine-tuning)，因为 LoRA 的训练速度非常快（在我的试验中，每个训练周期不到一秒）。

总结

本文展示了用批处理加速多个 LoRA 微调模型并行推断的可行性。我实现的 Punica 项目展现出了关于批处理大小几乎线性的吞吐量扩展，并且增加批处理大小并不显著增加延迟。

在 Github 上可以查看 Punica 的开源代码：https://github.com/punica-ai/punica

你可以在这个 Jupyter Notebook中访问本文中使用的基准代码。

这项研究仍然在进行中。我正在积极开展这项研究项目，预计很快会发布一个在线演示。我欢迎任何反馈或想法，请随时在评论部分或下面的按钮中分享。

Lequn Chen || abcdabcd987

大语言模型系统中RDMA通信的一些探索

集合通信的痛点

RDMA 的痛点

解决痛点的思路

MoE Kernel

感慨万千

跨机秒传RL模型参数更新的一些后续

流水线优化

PyTorch Profiler

补充一些图

Rank0-based 和 P2P 参数更新的对比

Torch.Distributed + RPC 和 RDMA 参数更新的对比

跨机秒传RL模型参数更新的一些探索

起因

预想的解决方案

RDMA 传输参数的概念验证

传输失败

获取 PyTorch 内存区域

传输速度

真实的参数更新（惨败）

真实的参数更新（推倒重来）

DTensor

DeviceMesh

训推参数匹配

生成路由表

执行参数更新

秒传 Qwen3-235B FP8

为什么跑不满带宽

为什么现有框架跑不快

对现有框架的快速调研

一些猜测

你需要的或许是一个 RDMA 通信库？

更新

驾驭3200Gbps网络(15): 惰性提交操作

惰性提交操作

运行效果

驾驭3200Gbps网络(14): 批量提交操作

批量提交操作

运行效果

驾驭3200Gbps网络(13): 状态分片

状态分片

运行效果

驾驭3200Gbps网络(12): 绑定CPU核心

绑定CPU核心

运行效果

驾驭3200Gbps网络(11): 多线程

多线程

运行效果

驾驭3200Gbps网络(10): 测试前预热

预热

运行效果

驾驭3200Gbps网络(9): 使用32张网卡

网络

负载均衡

CONNECT 消息

服务器端逻辑

客户端逻辑

运行效果

驾驭3200Gbps网络(8): 系统拓扑

获取系统拓扑

Linux sysfs

获取显卡的 PCI 地址

获取网卡的 PCI 地址

获取 CPU 拓扑

NUMA 节点信息

CPU 核心信息

CPU 超线程信息

PCI 设备 NUMA 节点信息

程序接口

运行效果

驾驭3200Gbps网络(7): 操作队列及带宽测试

操作队列

服务器端逻辑

客户端逻辑

运行效果

驾驭3200Gbps网络(6): GPUDirect RDMA WRITE

业务逻辑

GPU 上的缓冲区

注册内存区域

`CONNECT` 消息

`WRITE` 操作

`REMOTE WRITE` 回调函数

`RECV` 操作

`SEND` 操作

安装 `libfabric`