文章标题与作者

  • 标题：使用 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代理线程检测到该工作描述符并发起相应的网络操作。

图1. 使用CPU代理来发起NIC通信的GPU通信方式会导致通信瓶颈
图1. 使用CPU代理来发起NIC通信的GPU通信方式会导致通信瓶颈

代理线程与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网络操作已完成。

图2. 使用IBGDA的GPU通信建立了一条从GPU SM到NIC的直接控制路径，并将CPU从关键路径中移除
图2. 使用IBGDA的GPU通信建立了一条从GPU SM到NIC的直接控制路径，并将CPU从关键路径中移除

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倍。

  • 相关图表:

    图3. shmem_put_bw在IBRC下的带宽表现，显示了当扩展QP和CTA时，CPU代理对小消息尺寸造成的带宽上限
    图3. shmem_put_bw在IBRC下的带宽表现，显示了当扩展QP和CTA时，CPU代理对小消息尺寸造成的带宽上限

    图4. shmem_put_bw在IBGDA下的带宽表现，展示了小消息尺寸的带宽可以随着CTA和QP数量的增加而扩展
    图4. shmem_put_bw在IBGDA下的带宽表现，展示了小消息尺寸的带宽可以随着CTA和QP数量的增加而扩展

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

  • 实验内容:
    运行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。

  • 相关图表:

    图5. shmem_p_bw在IBRC和IBGDA下的put速率对比，显示了IBGDA在每秒发送数百万条小消息方面的性能优势
    图5. shmem_p_bw在IBRC和IBGDA下的put速率对比，显示了IBGDA在每秒发送数百万条小消息方面的性能优势

实验三：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通过并行提交大量小消息传输请求，使得应用程序能够隐藏这些开销。

  • 相关图表:

    图6. 使用P和PUT的Jacobi实现在IBRC和IBGDA下的延迟对比
    图6. 使用P和PUT的Jacobi实现在IBRC和IBGDA下的延迟对比

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

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

  • 实验结果:

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

  • 相关图表:

    图7. 32个PE的All-to-All传输在IBRC和IBGDA下的延迟对比
    图7. 32个PE的All-to-All传输在IBRC和IBGDA下的延迟对比

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系统中纵向扩展和横向扩展并存的复杂环境。