使用ROCm的GPU感知MPI

GPU-aware MPI with ROCm — ROCm Blogs (amd.com)

注意: 此博客之前是 AMD Lab Notes博客系列的一部分。

MPI（消息传递接口）是高性能计算中进程间通信的事实标准。MPI进程在其本地数据上进行计算，同时进行大量的相互通信。这使得MPI程序可以在具有分布式内存空间的系统（例如集群）上执行。MPI支持不同类型的通信，包括点对点通信和集体通信。

点对点通信是基本的通信机制，其中发送进程和接收进程都参与通信。发送方有一个包含消息的缓冲区和一个包含接收方将使用的信息（例如，消息标签、发送方的排名编号等）的信封。接收方使用信封中的信息选择指定的消息并将其存储在接收缓冲区中。

在集体通信中，消息可以在一组进程之间交换，而不仅仅是两个进程之间。集体通信提供了一种方便、可移植和优化的方式，使进程能进行一对多和多对多的通信。一些集体通信的例子包括广播（broadcast）、全部收集（allgather）、全部到全部（alltoall）和全部归约（allreduce）。

GPU感知的MPI

如今，许多MPI应用程序支持在GPU集群上执行。在这些应用程序中，计算密集型部分代码被卸载并加速到GPU上执行，也被称为设备。当涉及到MPI通信时，MPI进程需要通信驻留在GPU缓冲区中的数据。GPU感知的MPI提供了一种将GPU缓冲区传递给MPI调用的机会。这减轻了程序员通过主机内存加载GPU缓冲区的负担，使他们能够开发出更具可读性和简洁性的应用程序。此外，它还可以利用ROCm RDMA（远程直接内存访问）等加速技术使应用程序运行得更加高效。ROCm RDMA使得第三方设备（如Mellanox Infiniband HCA，主机通道适配器）可以与GPU内存直接进行点对点数据路径传输，而无需主机干预。大多知名的MPI实现，包括OpenMPI、MVAPICH2和Cray MPICH都支持GPU感知通信。

下面的代码展示了一个简单的GPU感知的点对点通信示例：

#include <stdio.h>
#include <hip/hip_runtime.h>
#include <mpi.h>int main(int argc, char **argv) {int i,rank,size,bufsize;int *h_buf;int *d_buf;MPI_Status status;bufsize=100;MPI_Init(&argc,&argv);MPI_Comm_rank(MPI_COMM_WORLD, &rank);MPI_Comm_size(MPI_COMM_WORLD, &size);//allocate buffersh_buf=(int*) malloc(sizeof(int)*bufsize);hipMalloc(&d_buf, bufsize*sizeof(int));//initialize buffersif(rank==0) {for(i=0;i<bufsize;i++)h_buf[i]=i;}if(rank==1) {for(i=0;i<bufsize;i++)h_buf[i]=-1;}hipMemcpy(d_buf, h_buf, bufsize*sizeof(int), hipMemcpyHostToDevice);//communicationif(rank==0)MPI_Send(d_buf, bufsize, MPI_INT, 1, 123, MPI_COMM_WORLD);if(rank==1)MPI_Recv(d_buf, bufsize, MPI_INT, 0, 123, MPI_COMM_WORLD, &status);//validate resultsif(rank==1) {hipMemcpy(h_buf, d_buf, bufsize*sizeof(int), hipMemcpyDeviceToHost);for(i=0;i<bufsize;i++) {if(h_buf[i] != i)printf("Error: buffer[%d]=%d but expected %d\n", i, h_buf[i], i);}fflush(stdout);}//free buffersfree(h_buf);hipFree(d_buf);MPI_Finalize();
}

如代码所示，我们将GPU缓冲区（`d_buf`）传递给`MPI_Send`和`MPI_Recv`调用。这个缓冲区是用hipMalloc在GPU上分配的。
要编译和运行此代码，您需要在系统上有ROCm以及GPU感知的MPI实现。您可以在[AMD ROCm™安装](AMD ROCm™ installation — ROCm Blogs)中找到ROCm安装说明。不同MPI实现下构建和运行上述代码（gpu-aware.cpp）的说明将在本文档后面讨论。

使用 OpenMPI 进行 GPU 感知通信

如前所述，大多数知名的 MPI 实现都支持 GPU 感知通信。在本节中，我们提供了使用 ROCm 支持构建 GPU 感知 OpenMPI 的说明。

要构建具有 ROCm 支持的 GPU 感知 OpenMPI，首先需要安装 Unified Communication X （[UCX](Home - UCX-Unified Communication X UCX-Unified Communication X)）。UCX 是一种用于高带宽和低延迟网络的通信框架。请使用以下命令来构建 UCX（版本 1.14）：

git clone https://github.com/openucx/ucx.git
cd ucx
git checkout v1.14.x
./autogen.sh
./configure --prefix=$HOME/.local --with-rocm=/opt/rocm --without-knem --without-cuda --enable-gtest --enable-examples
make -j
make install

成功安装后，UCX 将会在 $HOME/.local 目录中可用。现在，我们可以使用以下命令安装具有 ROCm 支持的 GPU 感知 OpenMPI：

git clone --recursive -b v5.0.x git@github.com:open-mpi/ompi.git
cd ompi/
./autogen.pl
./configure --prefix=$HOME/.local --with-ucx=$HOME/.local
make -j
make install

从 OpenMPI 5.0 开始，我们还可以在配置命令中添加 --with-rocm=/opt/rocm，以利用 Open MPI 中的一些 ROCm 功能，例如派生数据类型、MPI I/O 等。成功安装后，具有 ROCm 支持的 OpenMPI 将可在 $HOME/.local 中找到。接下来，我们可以按以下方式设置 PATH 和 LD_LIBRARY_PATH：

export PATH=$HOME/.local/bin:$PATH
export LD_LIBRARY_PATH=$HOME/.local/lib:$LD_LIBRARY_PATH

要使用 OpenMPI 编译一个 GPU 感知的 MPI 程序（如 gpu-aware.cpp），需要将 OMPI_CC 环境变量设置为 hipcc，以更改 mpicc 包装器的编译器。然后可以使用 mpicc 编译代码并使用 mpirun 运行代码：

export OMPI_CC=hipcc
mpicc -o ./gpu-aware ./gpu-aware.cpp
mpirun -n 2 ./gpu-aware

GPU感知通信与Cray MPICH

在本节中，我们讨论如何使用Cray MPICH构建和运行GPU感知的MPI程序。首先，确保您的系统上已加载ROCm、Cray MPICH和`craype-accel-amd-gfx90a/craype-accel-amd-gfx908`模块。您有两种编译代码的选择：

选项1：使用Cray编译器包装器编译代码并链接ROCm：

cc -o ./gpu-aware ./gpu-aware.cpp -I/opt/rocm/include/ -L/opt/rocm/lib  -lamdhip64 -lhsa-runtime64

选项2：使用hipcc编译代码并链接Cray MPICH：

hipcc -o ./gpu-aware ./gpu-aware.cpp -I/opt/cray/pe/mpich/8.1.18/ofi/cray/10.0/include/ -L/opt/cray/pe/mpich/8.1.18/ofi/cray/10.0/lib -lmpi

成功编译后，可以使用以下命令运行代码：

export MPICH_GPU_SUPPORT_ENABLED=1
srun -n 2 ./gpu-aware

请注意，将`MPICH_GPU_SUPPORT_ENABLED`设置为1，以启用GPU感知的通信。

使用OSU微基准测试进行性能测量

OSU微基准测试（OMB）提供了一系列MPI基准测试，用于测量各种MPI操作的性能，包括点对点、集合、基于主机和基于设备的通信。在本节中，我们讨论如何使用OSU微基准测试测量设备到设备的通信带宽。我们使用前面讨论的OpenMPI安装进行本节的实验。

您可以使用以下命令来构建支持ROCm的OSU微基准测试：

wget https://mvapich.cse.ohio-state.edu/download/mvapich/osu-micro-benchmarks-7.0.1.tar.gz
tar -xvf osu-micro-benchmarks-7.0.1.tar.gz
cd osu-micro-benchmarks-7.0.1
./configure --prefix=$HOME/.local/ CC=$HOME/.local/bin/mpicc CXX=$HOME/.local/bin/mpicxx --enable-rocm --with-rocm=/opt/rocm
make -j
make install

安装成功后，OSU微基准测试将位于`$HOME/.local/`。您可以使用以下命令运行带宽测试：

mpirun -n 2 $HOME/.local/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw  D D

在带宽测试中，发送进程将固定数量的消息背靠背发送给接收进程。接收进程在接收到所有这些消息后会发送一个回复。这个过程会重复多次。带宽是根据经过的时间和传输的字节数计算的。上述命令末尾的`D D`指定我们希望发送和接收缓冲区分配在设备上。

如果您没有获得预期的带宽，可能是OpenMPI默认没有使用UCX。要强制OpenMPI使用UCX，可以在mpirun命令中添加以下参数：

mpirun --mca pml ucx --mca pml_ucx_tls ib,sm,tcp,self,cuda,rocm -np 2 $HOME/.local/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw D D

您还可以使用以下命令运行集合通信测试：

mpirun -n 4 $HOME/.local/libexec/osu-micro-benchmarks/mpi/collective/osu_allreduce -d rocm

上述命令使用四个进程运行MPI_Allreduce延迟测试。在此测试中，基准测试会测量在大量迭代中，不同消息长度下四个进程间MPI_Allreduce集合操作的平均延迟。
-d rocm指定进程应将通信缓冲区分配在GPU设备上。

GPU-to-GPU 通信选项

在本节中，我们讨论了在使用 [AMD Instinct™ MI250](AMD Instinct™ MI250 Accelerators) 和 [AMD Instinct™ MI250X](https://www.amd.com/en/products/server-accelerators/instinct-mi250x) GPU 的系统中，GPU 对 GPU 的通信选项。每个 MI250(X) GPU 由两个图形计算芯片（GCD）组成。下图显示了一个包含 4 个 MI250 GPU（8 个 GCD）的节点的图表。每个绿色框表示一个具有两个 GCD 的 MI250 GPU。GCD 通过 Infinity Fabric 链接连接。每个 Infinity Fabric 链接的峰值带宽达到 50 GB/s。正如图中所示，不同的 GCD 之间存在不同数量的 Infinity Fabric 链接。例如，同一个 GPU 上的 GCD 通过 4 个链接连接，而不同 GPU 上的 GCD 通过 1 个或 2 个链接连接。从图中还能看出，不同的 GCD 之间有不同数量的跳数。例如，GCD 1 通过一次跳跃连接到 GCD 3，而 GCD 0 至少通过两次跳跃连接到 GCD 3。GCD 之间的最大可实现带宽取决于 GCD 之间的 Infinity Fabric 链接数量和跳数。

图 1：包含 4 个 MI250 GPU（8 个 GCD）的节点的图表。每个绿色框是一个具有两个 GCD 的 MI250 GPU。GCD 通过 Infinity Fabric 链接连接。

GPU 对 GPU 的通信有两个选项：
1. 使用系统直接内存访问 (SDMA) 引擎
2. 启动一个内核来处理通信

SDMA 引擎提供了与计算重叠进行通信的机会。然而，其缺点是它在 GCD 之间提供最大带宽为 50 GB/s。相比之下，内核提供了更高的通信带宽，但它们需要 CU 来移动数据，因此使用内核时与计算重叠通信的机会较少。可以使用环境变量 HSA_ENABLE_SDMA 来选择使用 SDMA 引擎还是内核。在 ROCM 5.4 及以下版本中，默认使用 SDMA 引擎。

以下实验 [1]展示了 GCD 0 与其对等设备间的通信带宽，对于大小为 16MB 的消息。我们将 HSA_ENABLE_SDMA 设置为 0，因此会启动一个内核来处理通信。我们设置 HIP_VISIBLE_DEVICES 来为每个实验选择对等的 GCD。使用 “-m ((16*1024*1024)):((16*1024*1024))” 我们指定我们感兴趣的消息大小，在此示例中为 16MiB。

export HSA_ENABLE_SDMA=0
export HIP_VISIBLE_DEVICES=0,1
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216           142235.39export HIP_VISIBLE_DEVICES=0,2
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            38963.65export HIP_VISIBLE_DEVICES=0,3
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            36903.57export HIP_VISIBLE_DEVICES=0,4
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            36908.74export HIP_VISIBLE_DEVICES=0,5
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            34986.54export HIP_VISIBLE_DEVICES=0,6
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            76276.14export HIP_VISIBLE_DEVICES=0,7
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            68788.80

上述实验表明，不同的 GCD 对之间存在不同的通信带宽。如前所述，这取决于 GCD 之间的 Infinity Fabric 链路数量以及跳数。例如，GCD 0 和 1 之间的通信带宽约为 142 GB/s，它们通过四条 Infinity Fabric 链接相连。每条 Infinity Fabric 链接的理论峰值带宽为 50 GB/s，因此 GCD 0 和 1 之间的理论峰值带宽为 200 GB/s。假设可实现的带宽约为理论带宽的 70%，那么预期 GCD 0 和 1 之间的通信带宽约为 142 GB/s。对于 GCD 0 和 2，它们通过一条 Infinity Fabric 链接相连，带宽为 38 GB/s。GCD 0 和 3 通过两跳相连，所以实现的带宽略低（约为 36 GB/s）。我们可以使用相同的逻辑来推理其他 GCD 对之间的通信带宽。

使用 SDMA 引擎时，我们得到的最大带宽约为 50 GB/s，正如前面讨论的那样：

export HSA_ENABLE_SDMA=1
export HIP_VISIBLE_DEVICES=0,1
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            49396.52export HIP_VISIBLE_DEVICES=0,2
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            41925.10export HIP_VISIBLE_DEVICES=0,3
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            41019.50export HIP_VISIBLE_DEVICES=0,4
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            42243.36export HIP_VISIBLE_DEVICES=0,5
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            41870.39export HIP_VISIBLE_DEVICES=0,6
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            49386.80export HIP_VISIBLE_DEVICES=0,7
mpirun -n 2 $OMB_DIR/libexec/osu-micro-benchmarks/mpi/pt2pt/osu_bw -m $((16*1024*1024)):$((16*1024*1024)) D D
# OSU MPI-ROCM Bandwidth Test v7.0
# Send Buffer on DEVICE (D) and Receive Buffer on DEVICE (D)
# Size      Bandwidth (MB/s)
16777216            49369.06

[附代码示例]

作者要感谢 Edgar Gabriel 和 Maria Ruiz Varela 对本文的有益评论和反馈。如果您有任何问题或意见，请在 GitHub 上联系我们 [讨论区](https://github.com/ROCm/rocm-blogs/discussions)

[1] 测试使用 ROCm 5.4、UCX 1.14 和 OpenMPI 5.0 进行。本文件中的测试结果不代表官方性能数据，但反映了不同通信选项的影响。实际性能取决于系统配置和环境设置。