Improving Network Performance of HPC Systems Using NVIDIA Magnum IO NVSHMEM and GPUDirect Async

发表时间: 2022-11

速读

一句话结论 本文在 NVSHMEM 通信库中引入了 InfiniBand GPUDirect Async (IBGDA) 技术，通过让 GPU 直接向网卡提交通信请求并完全绕过 CPU 代理，彻底解除了细粒度小消息的节点间通信吞吐量瓶颈。

要解决什么问题 在包含数万个 GPU 的大规模高性能计算系统中，强缩放工作负载极度依赖 GPU 间高效的细粒度通信。NVIDIA Magnum IO NVSHMEM 是一个提供分区全局地址空间数据访问模型的通信库，它在节点内利用 NVLink 能实现极高的细粒度数据访问效率。然而，当跨节点进行横向扩展时，传统的 NVSHMEM InfiniBand 可靠连接（IBRC）传输机制存在严重的架构卡点：它必须依赖主机 CPU 上的代理线程来管理网络操作。当 GPU 上的 CUDA 内核发起通信时，只能将工作描述符写入主机内存的代理缓冲区，随后由 CPU 代理线程检测描述符、构建网络请求、更新主机内存中的门铃记录（DBR）缓冲区，并最终敲击网卡（NIC）的硬件门铃寄存器。现代网卡每秒能处理数亿个通信请求，GPU 也能以同等速率生成请求，但 CPU 代理线程处理这些请求的速率却要低几个数量级。这种严重的速率不匹配导致 CPU 周期被大量消耗，且 CPU 代理成为了通信控制路径上的绝对瓶颈。在小消息和高并发场景下，这种机制根本无法打满网络硬件的峰值吞吐量，导致节点间网络通信与节点内 NVLink 通信在小消息性能上存在巨大鸿沟。

怎么做的 核心思路是利用 GPUDirect Async 技术建立一条从 GPU 流式多处理器（SM）到网卡的直接控制路径，将 CPU 完全从通信关键路径中剔除。为了实现这一目标，IBGDA 对底层通信部件的物理位置和交互逻辑进行了重构。首先，传统机制中位于主机内存的工作队列（WQ）缓冲区和门铃记录（DBR）缓冲区被全部下放到了 GPU 内存中，这使得 GPU SM 能够以极低的延迟直接访问它们，同时保留了网卡通过 GPUDirect RDMA 技术对其进行访问的能力。当应用程序在 GPU 上调用 NVSHMEM 通信操作时，GPU 的 SM 会直接创建一个网卡工作描述符，并将其写入 GPU 内存中的 WQ 缓冲区，同时更新 DBR 缓冲区。随后，SM 直接向网卡硬件中的门铃寄存器写入数据来通知网卡。网卡收到通知后，利用 GPUDirect RDMA 技术直接从 GPU 内存读取工作描述符和待传输数据，并发送到远程节点。传输完成后，网卡同样使用 GPUDirect RDMA 将完成事件直接写入完成队列（CQ）缓冲区，从而通知 GPU 操作已结束。我们可以用通信延迟的构成来表达这一机制的转变，传统 IBRC 机制的单次通信延迟为：$$T_{IBRC} = T_{SM \to Proxy} + T_{CPU\_Process} + T_{CPU \to NIC} + T_{Transfer}$$ 而 IBGDA 机制彻底消除了 CPU 处理开销，将其简化为：$$T_{IBGDA} = T_{SM \to NIC} + T_{Transfer}$$ 此外，IBGDA 还设计了自动数据合并功能，当同一个线程束内的目标地址连续时，硬件能自动将 32 个小消息合并成一个大消息发送。

效果如何 实验在 4 台 DGX-A100 服务器上搭建，节点间通过 NVIDIA Quantum HDR 交换机连接 NVIDIA ConnectX-6 200 Gb/s InfiniBand 网卡。为凸显网络性能，所有实验均强制禁用了节点内的 NVLink 通信。对比基线为 NVSHMEM 标准发行版中的两种传输方式：代表传统 CPU 代理路线的 IBRC 传输，以及代表本文新路线的 IBGDA 传输。在单边写带宽测试中，当使用 64 个合作线程阵列（CTA）发起通信时，IBGDA 仅需 2 KiB 的小消息就能打满网络峰值带宽，而 IBRC 受限于 CPU 瓶颈，无论怎么增加 CTA 数量都无法降低达到峰值带宽所需的最小消息尺寸；对于小于 1 KiB 的消息，IBGDA 提供的吞吐量比 IBRC 高出多达 9.5 倍。在 8 字节标量小消息吞吐量测试中，IBRC 的发送速率死锁在约 1.7 MOPS（百万操作/秒），而 IBGDA 仅用 8 个 CTA 就逼近了网卡 215 MOPS 的硬件极限；当触发自动数据合并特性时，IBGDA 的发送速率甚至突破了网卡峰值，接近 2000 MOPS。在 Jacobi 求解器强缩放案例中，使用 IBRC 时细粒度标量发送版本的延迟是数据聚合发送版本的 8 倍以上，而换用 IBGDA 后，细粒度发送版本的延迟大幅降低至与聚合发送版本持平。在 32 个处理元件参与的 All-to-All 集体操作测试中，IBGDA 彻底消除了 IBRC 中 CPU 代理造成的序列化瓶颈，在小于 16 KiB 的消息场景下展现出更低且更稳定的总延迟。作者也指出，发送大量小消息的绝对网络开销依然略高于发送单个大消息，但 IBGDA 凭借并行提交海量请求的能力，让应用程序得以隐藏这些网络开销。

文章标题与作者

• 标题：使用 NVIDIA Magnum IO NVSHMEM 和 GPUDirect Async 提升 HPC 系统的网络性能
• 作者：Pak Markthub
• 日期：2022年11月22日

---

A1 主要贡献

本文旨在解决当前大规模高性能计算（HPC）系统中存在的细粒度通信效率低下的核心问题。随着HPC系统扩展到数万个GPU，特别是在强缩放（strong scaling）应用场景下，GPU间的通信效率成为性能瓶颈。传统的基于NVSHMEM的节点间通信方法依赖CPU作为代理来管理网络操作，这限制了细粒度通信的吞吐量，因为CPU的处理速度远低于现代网络接口控制器（NIC）和GPU。

为了解决这一问题，本文的研究目标是提升NVSHMEM在节点间的通信性能，尤其是针对小消息和高并发场景。

其核心创新点是引入了一种名为InfiniBand GPUDirect Async (IBGDA) 的新通信方法。IBGDA构建于GPUDirect Async技术之上，其主要贡献在于：
1. 绕过CPU：IBGDA使GPU的流式多处理器（SM）能够直接向NIC提交通信请求，完全绕过了CPU在通信控制路径中的参与。
2. 提升吞吐量：通过消除CPU代理瓶颈，IBGDA显著提高了细粒度通信的吞吐量，使得小消息通信也能接近网络硬件的峰值速率。
3. 简化优化：IBGDA弥合了节点内NVLink通信与节点间网络通信在小消息性能上的差距，使开发者能够更轻松地编写同时适用于纵向扩展（scale-up）和横向扩展（scale-out）的高效代码，而无需为两种场景设计不同的通信策略。

---

A3 背景知识与设计原则

引言与背景
当今领先的高性能计算（HPC）系统包含数万个GPU。在NVIDIA系统中，节点内的GPU通过NVLink互连实现纵向扩展，而节点间的GPU则通过InfiniBand等网络实现横向扩展。用于GPU通信、共享工作和高效并行操作的软件库统称为NVIDIA Magnum IO，它是并行、异步和智能数据中心IO的体系架构。对于许多应用而言，扩展到如此大的系统需要在GPU之间进行高效的细粒度通信。这对于旨在通过增加计算资源来减少解决特定问题时间的强缩放工作负载尤其关键。NVIDIA Magnum IO NVSHMEM是一个基于OpenSHMEM规范的通信库，它为HPC系统中所有GPU的内存提供了一个分区全局地址空间（PGAS）数据访问模型。由于支持GPU集成的通信，该库对于强缩放工作负载是一个特别合适且高效的工具。在此模型中，数据通过单边读、写和原子更新通信例程进行访问。这种通信模型通过与GPU架构的紧密集成，在NVLink上实现了细粒度数据访问的高效率。然而，由于需要主机CPU来管理通信操作，节点间数据访问的高效率一直是一个挑战。本文介绍了一种在NVSHMEM中称为InfiniBand GPUDirect Async（IBGDA）的新通信方法，它构建在GPUDirect Async技术家族之上。IBGDA在NVSHMEM 2.6.0中引入，并在NVSHMEM 2.7.0和2.8.0中得到显著改进。它使GPU在发出节点间NVSHMEM通信时能够绕过CPU，而无需对现有应用程序进行任何更改。如下文所示，这为使用NVSHMEM的应用程序带来了吞吐量和扩展性的显著提升。

传统CPU代理通信机制
使用NVLink进行节点内通信可以通过GPU流式多处理器（SM）发起的加载和存储指令来实现。然而，节点间通信涉及向网络接口控制器（NIC）提交工作请求以执行异步数据传输操作。在IBGDA引入之前，NVSHMEM InfiniBand可靠连接（IBRC）传输使用CPU上的代理线程来管理通信。当使用代理线程时，NVSHMEM执行以下操作序列：
1. 应用程序启动一个在GPU内存中生成数据的CUDA内核。
2. 应用程序调用一个NVSHMEM操作（如 nvshmem_put）与另一个处理元件（PE）进行通信。在执行细粒度或重叠通信时，此操作可以在CUDA内核内调用。NVSHMEM操作将一个工作描述符写入位于主机内存中的代理缓冲区。
3. NVSHMEM代理线程检测到该工作描述符并发起相应的网络操作。

代理线程与HCA的交互细节
代理线程在与NVIDIA InfiniBand主机通道适配器（HCA），例如ConnectX-6 HCA，交互时执行以下操作序列：
1. CPU创建一个工作描述符，并将其排入位于主机内存中的工作队列（WQ）缓冲区。
2. 该描述符指明了请求的操作，如RDMA写，并包含源地址、目标地址、大小和其他必要的网络信息。
3. CPU更新位于主机内存中的门铃记录（DBR）缓冲区。该缓冲区用于在NIC丢失对其门铃（DB）的写入时的恢复路径中。
4. CPU通过写入NIC的门铃（一个位于NIC硬件中的寄存器）来通知NIC。
5. NIC从WQ缓冲区读取工作描述符。
6. NIC使用GPUDirect RDMA直接从GPU内存复制数据。
7. NIC将数据传输到远程节点。
8. NIC通过向主机内存上的完成队列（CQ）缓冲区写入一个事件来指示网络操作已完成。
9. CPU轮询CQ缓冲区以检测网络操作的完成。
10. CPU通知GPU操作已完成。如果存在GDRCopy，它会直接向GPU内存写入一个通知标志。否则，它会将该标志写入代理缓冲区。GPU轮询相应的内存以获取工作请求的状态。

传统方法的瓶颈
虽然这种方法具有可移植性，并且可以为大块数据传输提供高带宽，但它有两个主要缺点：首先，CPU周期被代理线程持续消耗。其次，由于代理线程的瓶颈，你无法为细粒度传输达到峰值NIC吞吐量。现代NIC每秒可以处理数亿个通信请求。虽然GPU可以以这个速率生成请求，但CPU代理的处理速率要低几个数量级，从而为细粒度通信模式创造了瓶颈。

---

A2 方法细节

IBGDA：GPU直接控制NIC的通信路径
与代理发起的通信相比，IBGDA使用GPUDirect Async–KernelInitiated (GPUDirect Async–KI)技术，使GPU SM能够直接与NIC交互。这一过程包含以下步骤：
1. 应用程序启动一个在GPU内存中生成数据的CUDA内核。
2. 应用程序调用一个NVSHMEM操作（例如 nvshmem_put）与另一个PE通信。NVSHMEM操作利用SM创建一个NIC工作描述符，并将其直接写入WQ缓冲区。与CPU代理方法不同，这个WQ缓冲区位于GPU内存中。
3. SM更新同样位于GPU内存中的DBR缓冲区。
4. SM通过写入NIC的DB寄存器来通知NIC。
5. NIC使用GPUDirect RDMA读取WQ缓冲区中的工作描述符。
6. NIC使用GPUDirect RDMA读取GPU内存中的数据。
7. NIC将数据传输到远程节点。
8. NIC通过使用GPUDirect RDMA向CQ缓冲区写入数据，来通知GPU网络操作已完成。

IBGDA的核心优势
如上所示，IBGDA将CPU从通信控制路径中消除。在使用IBGDA时，GPU和NIC直接交换通信所需的信息。WQ和DBR缓冲区也被移至GPU内存，以提高SM访问时的效率，同时通过GPUDirect RDMA保留NIC的访问能力。

---

A4 实验环境

• 硬件配置:
  • 计算平台: 4台 DGX-A100 服务器。
  • 网络: NVIDIA ConnectX-6 200 Gb/s InfiniBand 网卡，通过一台 NVIDIA Quantum HDR 交换机连接。
• 软件配置:
  • 通信库: 对比了NVSHMEM标准发行版中的两种传输方式：IBGDA传输和使用CPU代理线程的IBRC传输。
• 实验设置:
  • 为了突出IBGDA的效果，所有实验均禁用了NVLink通信。这强制所有数据传输（即使是位于同一节点上的PE之间）都通过InfiniBand网络进行。

---

A4 实验结果

实验一：单边put带宽测试

• 实验内容:
  运行NVSHMEM性能测试套件中的shmem_put_bw基准测试，该测试使用nvshmem_double_put_nbi_block函数发起数据传输，衡量在不同通信参数下传输固定总量数据时实现的单边写操作带宽。实验在两个不同的DGX-A100节点上启动了两个PE，每个线程块设置一个线程，每个线程块分配一个QP（NIC队列对）。

• 实验结果:

  • 大消息性能：对于大消息的粗粒度通信，IBGDA和IBRC都能达到网络峰值带宽。当应用从至少4个合作线程阵列（CTA）发起通信时，IBRC可以用小至16 KiB的消息使网络饱和（图3）。
  • 小消息瓶颈（IBRC）：对于IBRC，进一步增加CTA数量并不能减小达到峰值带宽所需的最小消息尺寸。小消息带宽的瓶颈在于CPU代理线程（图3）。
  • 小消息性能（IBGDA）：IBGDA通过移除代理瓶颈，在使用64个CTA发起通信时，能用小至2 KiB的消息就达到峰值带宽（图4）。
  • 扩展性对比：IBGDA的带宽随着执行通信的CTA数量持续扩展，而IBRC的扩展性在4个CTA时达到极限。因此，对于小于1 KiB的消息，IBGDA提供的吞吐量比IBRC高出多达9.5倍。

• 相关图表:

  

  

实验二：小消息吞吐量提升测试

• 实验内容:
  运行shmem_p_bw基准测试，该测试使用标量操作nvshmem_double_p将数据直接从GPU寄存器发送到远程PE。此操作是线程范围的，即每个调用此操作的线程传输一个8字节的数据。实验中，每个CTA启动1024个线程，并同步增加CTA和QP的数量。

• 实验结果:

  • IBRC性能上限：IBRC的put速率（以百万操作/秒，MOPS计）被限制在约1.7 MOPS，无论CTA和QP的数量如何增加（图5）。
  • IBGDA性能扩展：IBGDA的消息速率随CTA数量的增加而增加，仅用8个CTA就接近了NVIDIA ConnectX-6 InfiniBand NIC的硬件极限215 MOPS（图5）。
  • 数据合并特性：IBGDA提供了自动数据合并功能，当同一warp内的目标地址是连续的时，它能将32个小消息合并成一个大消息发送。
  • 超高吞吐量：如图5所示，当数据合并条件满足时，put速率可以超过NIC的峰值消息速率，接近2000 MOPS。

• 相关图表:

  

实验三：Jacobi方法案例研究

• 实验内容:
  分析NVSHMEM Jacobi基准测试的性能，对比IBGDA与IBRC在真实应用中的表现。该测试包含两种实现：

  1. nvshmem_p 版本：每个线程在数据可用时立即使用标量nvshmem_p操作发送数据。
  2. nvshmem_put 版本：每个CTA先将数据聚合到连续的GPU缓冲区，然后调用nvshmem_put_nbi_block发起通信。
     实验为强缩放场景，即固定矩阵大小，增加PE数量。

• 实验结果:

  • IBRC下的性能差异：使用IBRC时，nvshmem_p版本的Jacobi求解器延迟是nvshmem_put版本的8倍以上（图6）。
  • IBGDA下的性能一致性：使用IBGDA时，nvshmem_p和nvshmem_put两个版本都表现出良好的扩展性，并且其效率与nvshmem_p在NVLink上的表现相当。nvshmem_p与IBGDA的组合，其延迟与nvshmem_put与IBRC的组合相匹配（图6）。
  • 网络开销分析：nvshmem_p与IBGDA的延迟略高于nvshmem_put，这是因为发送一个大消息的网络开销低于发送许多小消息。然而，IBGDA通过并行提交大量小消息传输请求，使得应用程序能够隐藏这些开销。

• 相关图表:

  

实验四：All-to-All集体操作延迟测试

• 实验内容:
  测量32个PE参与的NVSHMEM all-to-all集体操作在IBRC和IBGDA下的延迟。

• 实验结果:

  • IBRC的序列化瓶颈：在IBRC中，CPU代理线程成为所有来自设备的操作的序列化点，这会引入额外的延迟，导致性能不稳定，尤其是在消息尺寸小于16 KiB时（图7）。
  • IBGDA的性能一致性：IBGDA消除了序列化瓶颈，展现出更稳定且更低的总延迟，特别是在小消息场景下（图7）。

• 相关图表:

  

---

A5 结论

本文展示了NVIDIA Magnum IO如何提升小消息网络性能，这对于部署在HPC数据中心成百上千个节点上的大规模应用尤为重要。自NVSHMEM 2.6.0版本起引入的InfiniBand GPUDirect Async (IBGDA)技术，使得GPU的SM能够直接向NIC提交通信请求，从而在NVIDIA InfiniBand网络上绕过CPU进行网络通信。

与依赖CPU代理管理通信的传统方法相比，IBGDA能够在更小的消息尺寸下维持显著更高的吞吐率。这些性能改进对于需要强缩放的应用至关重要，因为在这类应用中，随着工作负载扩展到更多的GPU，消息尺寸往往会减小。

此外，IBGDA弥合了节点内NVLink通信和节点间网络通信在小消息吞吐量上的性能差距，使得开发者可以更便捷地优化代码，以适应当前GPU加速的HPC系统中纵向扩展和横向扩展并存的复杂环境。