AWS Sagemaker Hyperpod 上的高性能 GPU 内存传输

现代深度学习基础设施通常需要在机器间的 GPU 之间传输海量数据。在 Perplexity，我们遇到了一个独特的技术挑战：如何以尽可能高的速度在机器间高效传输非连续的 GPU 内存区域。我们的目标平台 AWS p5 实例通过 32 张网卡提供了令人惊叹的 3200 Gbps 网络带宽。本文将分享我们构建自定义高性能网络解决方案的历程，该方案实现了理论带宽的 97.1%。

我们的用例提出了几个关键技术要求：

在远程 GPU 之间对非连续内存块进行高带宽传输
能够动态添加或删除 Kubernetes 部署中的节点，且不中断正在进行的操作
支持对等通信模式

虽然 NVIDIA 的 NCCL 库是分布式深度学习的事实标准，但它并不适合我们的用例：

NCCL 在集合通信方面表现出色，但需要建立一个静态的“世界”（world），这要求在调整参与节点时重启整个集群。
NCCL 的同步通信模型增加了我们异步工作负载的复杂性
我们希望对内存传输模式进行直接控制以进行优化
构建我们自己的解决方案提供了宝贵的学习机会

为了理解我们的解决方案，首先让我们探索现代高性能网络与传统网络有何不同。

我们日常使用的大多数网络都依赖 TCP/IP 协议，其中应用程序通过使用套接字的操作系统内核与网卡进行通信。然而，高性能网络使用 RDMA（远程直接内存访问）——这是一种完全不同的硬件和软件栈，能够在不涉及 CPU 的情况下实现机器之间的直接内存访问。

AWS 提供 Elastic Fabric Adapter (EFA)，这是一种实现 Amazon 自定义协议（称为 Scalable Reliable Datagram (SRD)）的自定义网络接口。与传统 TCP/IP 网络（数据必须在用户空间、内核空间和网络缓冲区之间多次复制）不同，结合 RDMA 的 EFA 允许在 GPU 内存和网卡之间进行直接数据传输，完全绕过 CPU。

构建高性能网络系统需要重新思考几个基本假设：

缓冲区所有权：与内核管理网络缓冲区并需要在用户空间和内核空间之间进行复制的传统套接字不同，RDMA 要求应用程序自行管理缓冲区。当应用程序发起网络操作时，它会将缓冲区所有权转移给网卡，直到操作完成，从而消除了数据复制的需要。
内存注册：应用程序必须向操作系统内核注册内存区域。内核设置虚拟地址映射，允许 CPU、GPU 和网卡理解相同的虚拟地址。此注册是一次性操作，可实现随后的零拷贝数据传输。
控制平面与数据平面：高性能网络将操作分为两类：
1. 控制平面操作（如连接设置和内存注册）通过内核进行，以确保安全性
2. 数据平面操作（实际数据传输）绕过内核以获得最高性能
先接收后传输：在没有内核管理缓冲区的情况下，应用程序必须预先发布接收操作，指定传入数据应放置的位置。这与应用程序可以在任何时间接收数据的套接字模型发生了根本性的转变。
基于轮询的完成机制：应用程序不是通过 epoll 等机制等待内核通知，而是直接轮询硬件完成队列。这消除了系统调用开销，并允许对已完成的操作做出立即响应。
硬件拓扑感知：理解并针对硬件拓扑进行优化对于实现最高性能至关重要。

AWS p5 实例拥有复杂的内部架构。如下图所示，每个实例包含两个形成两个 NUMA 节点的 CPU 插槽，每个 NUMA 节点连接到四个 PCIe 交换机：

在每个 PCIe 交换机下，我们发现：

四个 100 Gbps EFA 网卡
一个 NVIDIA H100 GPU
一个 NVMe SSD

TCP/IP 和 RDMA 传输的数据路径展示了其方法上的根本差异：

GPU、用户空间应用程序、Linux 内核和网络接口之间的 TCP 和 GPUDirect RDMA 数据传输路径。

使用 TCP/IP（左侧）时，数据必须被多次复制：

从 GPU 内存到用户空间中的应用程序缓冲区
从用户空间缓冲区到内核缓冲区
从内核缓冲区到网卡
接收端发生相反的过程

每次复制操作都会消耗 CPU 周期和内存带宽。应用程序还必须为每次网络操作在用户空间和内核空间之间进行上下文切换。

相比之下，RDMA（右侧）提供了真正的零拷贝数据传输：

网卡直接从 GPU 内存读取数据
数据直接传输到远程网卡
远程网卡直接写入目标 GPU 内存
应用程序只需要在用户空间检查完成队列 (CQ) 即可知道传输何时完成

借助 RDMA 和适当的硬件配对，在两个 GPU 之间传输数据仅需要遍历本地 PCIe 交换机和网络：

通过 PCIe Gen5 x16 和四条 100Gbps EFA 连接链路相连的配对本地和远程 GPU。

相比之下，TCP/IP 传输必须多次通过主内存复制数据，从而导致严重的 PCIe 总线拥塞：

TCP 传输导致主内存、CPU、四个 PCIe 交换机和四块 GPU 之间的系统总线拥塞。

我们使用了 libfabric，这是一个为 Fabric 服务提供通用接口的框架。我们的实现使用了两种类型的 RDMA 操作：

用于承载有关内存区域元数据（例如偏移量和大小）的控制消息的双向 RDMA (SEND/RECV)。
用于实际数据传输的单向 RDMA WRITE，其中每个 WRITE 操作处理一个连续的内存块。

我们的开发经历了几个阶段：

使用 SEND/RECV 实现了基本的单向消息传输
扩展到双向通信
添加了用于 GPU-GPU 直接传输的 GPUDirect RDMA WRITE
扩展以处理多个并发传输
引入了用于鲁棒性的操作排队
在单张网卡上实现了 97.4% 的带宽利用率

当扩展到 32 张网卡时，我们实施了几个关键优化：

操作排队：我们没有直接向网卡提交操作，而是维护了一个应用层队列。这提供了抵御网络拥塞的鲁棒性，并简化了编程模型。
网络预热：预先建立连接提高了启动性能。
多线程：为每个 GPU 的网络操作设置专用问题。
CPU 核心绑定：将问题绑定到特定的 CPU 核心以避免 NUMA 影响和缓存未命中。
状态分片：减少问题之间的争用。
操作批处理：将多个操作一起提交。
延迟操作发布：操作首先在应用程序中排队。在轮询完成队列后，我们尝试将挂起的操作提交到网卡，以确保高效利用网络资源。
NUMA 感知资源分配：在正确的 NUMA 节点上分配完成队列等 libfabric 资源，以最大限度地减少内存访问延迟。

通过这些优化，我们最终达到了 3,108 Gbps 的性能，即理论最大带宽的 97.1%。

下面的视频展示了我们的命令行程序运行情况。它将不连续的 GPU 内存块传输到远程节点，传输速度达到 3108.283 Gbps，证明了对网络基础设施理论带宽的近乎完美的利用：

构建高性能网络系统需要了解硬件架构和系统设计原则。虽然 NCCL 等库为常见模式提供了出色的解决方案，但有时针对特定需求定制解决方案是必要的。

我们的历程证明，通过仔细关注系统架构、硬件拓扑和各种优化技术，实现接近理论值的网络性能是可能的。关键不仅在于理解各个组件，还在于它们如何交互以形成一个完整的系统。

此次旅程的完整技术深度剖析（包括实现细节和代码示例）可在我们的开源存储库中获取。

BibTeX

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}

技术挑战

现代高性能网络

高性能网络设计理念

理解硬件拓扑

使用 libfabric 构建

追求峰值性能之旅

结论

引用