GTC 大会火热进行中!别错过 NVIDIA CEO 黄仁勋的最新发布,以及 AI 和加速计算的必听会议。
模拟/建模/设计

通过 GROMACS 大幅提高多节点 NVIDIA GPU 的可扩展性

GROMACS 是一种广泛用于模拟生物分子系统的科学软件包,在理解对疾病预防和治疗重要的重要生物学过程中发挥着至关重要的作用。 GROMACS 可以并行使用多个 GPU 以尽可能快地运行每个模拟。

在过去几年中, NVIDIA 和 主要 GROMACS 开发人员 合作进行了一系列多 GPU 和多节点优化。

在这篇文章中,我们展示了这些改进中的最新进展,通过启用 GPU 粒子网格 Ewald ( PME )分解和 GPU directcommunication :新 GROMACS 2023 发布版本中提供的一项功能。我们观察到,通过这项工作,性能提高了 21 倍。

实现改进的多节点性能

之前的文章 中,我们对单个节点内的多 GPU 可伸缩性进行了优化,包括 GPU direct 通信的开发。我们描述了 GROMACS 通常如何将一个 GPU 分配给 PME 长程力计算(通过到傅里叶空间的变换执行),而剩余的 GPU 用于短程粒子 – 粒子( PP )力计算(直接在真实空间中执行)。

我们还描述了直接在这些 GPU 之间进行通信的发展:多个 PP GPU + PP-PME 通信之间的光环交换。我们在论文 Heterogeneous parallelization and acceleration of molecular dynamics simulations in GROMACS 中更详细地描述了这项工作。

GPU 与 MPI 直接通信

在 GROMACS 2022 发布版本中,扩展了 GPU 直接通信功能,以支持与 CUDA 感知 MPI 的兼容性,从而实现跨多个节点(以及每个节点内)的 GPU 直接通信。有关如何激活此功能的更多信息,请参阅本文后面的 How to run 部分。

然而,仍然存在与 PME 的单个 GPU 的限制相关的可扩展性限制。虽然可以增加更多 GPU 来处理 PP 力计算,但不可能将大多数模拟扩展到几个节点之外。当放大时,单个 PME GPU 在某一点上成为限制因素。 PME 计算必须在多个 GPU 中进行分解,以实现进一步的缩放。

PME GPU 分解

通过引入 PME GPU 分解,在全新 GROMACS 2023 发布版本中,这一单一 PME GPU 限制已被解除。这利用了新的 NVIDIA cuFFTMp library ,该库能够以分布式方式跨计算节点内和跨计算节点的多个 GPU 执行所需的快速傅里叶变换( FFT )。

cuFFTMp 则使用 NVSHMEM ,这是一种并行编程接口,可实现快速的单边通信。它可以有效地利用节点内和节点间互连来执行分布式 3D FFT 所需的所有通信。

通过集成这一新功能, GROMACS 现在可以在模拟中计算多个 PME 等级,从而显著增强了可扩展性和性能。为了便于移植,该实现还支持 heFFTe ,这是 cuFFTMp 的替代方案。我们发现 cuFFTMp 的速度快了 2 倍,因此我们将在本篇文章的剩余部分重点关注这一点。

这项工作涉及到对 GROMACS 中处理 PME 通信的方式的几个复杂调整,以及与新 cuFFTMp 功能的接口。

为了实现这一新功能,进行了一些算法改进。这包括引入流水线以增强计算和通信的重叠,以及重新设计并行实现以使用网格重叠减少机制,而不是在 PME 计算之前重新分配数据。

有关更多信息,请参阅 GROMACS GitLab 问题 Support PME decomposition in CUDA backend Optimize CUDA PME Decomposition

绩效结果

在本节中,我们将展示使用新性能特性的性能结果。我们使用了以下测试用例:

Two charts show that PME GPU decomposition with GPU direct communication enables massively improved multi-node scalability and performance. The larger case can scale to more nodes than the smaller case.
图 1 。 GROMACS STMV 的性能 ( left ) 和 BenchPEP-h( 右 ) 节点数增加的基准(高即好)

在图 1 中,我们在 Selene DGX-A100 集群上每个节点使用了四个 NVIDIA A100 GPU 。有关更多信息,请参阅本文的下一节。红色方块、蓝色圆圈和绿色三角形分别表示使用传统代码路径、仅 GPU 直接通信和结合 GPU PME 分解的 GPU 直接通信的结果。

您可以看到, GPU 直接通信的引入比通过 CPU 内存进行通信的传统代码路径的速度提高了 2-3 倍。通过 GPU 直接通信,基准可以扩展到大约两个节点,但在这一点上,单个 PME GPU 成为限制因素,不可能进一步扩展。

结合 GPU 直接通信和 PME GPU 分解的结果使用每个节点一个 PME GPU ,我们发现这给出了最佳平衡。每个节点上的剩余三个 GPU 专用于 PP 短程力计算。

在单个节点上, PME GPU 分解没有影响,因为您仍然只有一个 PME GPU 。

在两个节点上,使用两个 PME GPU 的性能与使用单个 PME GPU 的性能类似。前者对于 STMV 稍慢,对于 benchPEP-h 稍快。

在这种规模下,关键路径仍然对 PP 工作负载敏感,因此需要进行实验以找到性能最佳的配置。例如,我们牺牲了一个 PP GPU 来启用额外的 PME GPU 。然而,随着节点数量的进一步扩大,关键路径完全由 PME 性能决定, PME GPU 分解实现了更好的可伸缩性。

STMV 案例可以扩展到八个节点,在这种情况下,结果比传统代码路径快 3 倍,比没有 GPU PME 分解的结果快 1.6 倍。

对于 STMV ,在 16 个节点上,您可以看到比 8 个节点稍有改进,但并行效率要低得多。较大的benchPEP-h案例可以扩展到测试的最大 64 节点( 256-GPU )配置,比传统代码路径快 21 倍,比没有 GPU PME 分解的结果快 11 倍。

如何构建和运行 GROMACS

在本节中,我们将详细介绍如何使用 PME GPU 分解运行 GROMACS 以获得性能结果。

我们按照网站上的下载和安装说明安装了 NVIDIA HPC SDK 22.11 。对于 GROMACS 2023 ,我们不推荐任何更高版本的 HPC SDK ,因为兼容性问题将在未来的开发中解决。

获取 GROMACS 2023 发布版本如下:

git clone  https://gitlab.com/gromacs/gromacs.git 
cd gromacs
git checkout v2023

按如下方式构建 GROMACS :

# make a new directory to build in
mkdir build
cd build

# set the location of the math_libs directory in the NVIDIA HPC installation
HPCSDK_LIBDIR=/lustre/fsw/devtech/hpc-devtech/alang/packages/nvhpc/nvhpc_2022_2211_Linux_x86_64_cuda_11.8-install/Linux_x86_64/2022/math_libs

# build the code with PME GPU decomposition with cuFFTMp enabled,
# in an environment with a CUDA-aware OpenMPI installation 
# (see https://manual.gromacs.org/current/install-guide/index.html)
 cmake \
     ../ \
 -DGMX_OPENMP=ON -DGMX_MPI=ON -DGMX_BUILD_OWN_FFTW=ON \
 -DGMX_GPU=CUDA  -DCMAKE_BUILD_TYPE=Release -DGMX_DOUBLE=off \
 -DGMX_USE_CUFFTMP=ON -DcuFFTMp_ROOT=$HPCSDK_LIBDIR 
 
make -j 8 #Build using 8 CPU threads. Can increase this if you have more CPU cores available.

HPCSDK_LIBDIR变量设置为 HPC SDK 安装的Linux_x86_64/2022/math_libs子目录。

获取输入文件如下:

wget https://zenodo.org/record/3893789/files/GROMACS_heterogeneous_parallelization_benchmark_info_and_systems_JCP.tar.gz
tar zxvf GROMACS_heterogeneous_parallelization_benchmark_info_and_systems_JCP.tar.gz
ln -s GROMACS_heterogeneous_parallelization_benchmark_info_and_systems_JCP/stmv/topol.tpr .

我们在 NVIDIA Selene 集群上进行了基准测试,该集群由多个 DGX-A100 服务器组成,每个服务器都有八个 A100-SXM4 GPU ,通过 Infiniband 连接。然而,根据更典型的 HPC 安装,每个节点只使用了四个 GPU 。

我们在将 STMV topol.tpr输入文件复制到工作目录后,使用标准作业调度方法(本例中为 Slurm )运行代码。我们提供了 full submission script used on Selene ,其中包括详细的评论。

要使用 PME GPU 分解,请设置以下变量:

# Specify that GPU direct communication should be used
export GMX_ENABLE_DIRECT_GPU_COMM=1

# Specify that GPU PME decomposition should be used
export GMX_GPU_PME_DECOMPOSITION=1

通过 -npme <N>标志指定 PME 的总数 GPU 至mdrun。如果没有 PME GPU 分解,N为 1 ,因为您只能使用单个 PME GPU 。

通过分解,您可以将N设置为正在使用的节点数,以便为每个节点指定一个 PME GPU ,而每个节点中的其他三个 GPU 专用于 PP 。考虑到 PP 和 PME 工作负载的相对计算开销,这种划分通常会提供良好的平衡,但建议在任何特定情况下进行实验。

总结

在这篇文章中,我们描述了我们在增强 GROMACS 多节点可扩展性方面的工作,我们展示了高达 21 倍的性能改进,并详细介绍了如何使用这些新功能。这使研究人员能够继续突破关键生物过程知识的界限。

我们继续通过一系列开发改进 GROMACS ,不仅针对进一步增强的可扩展性,还针对每个计算节点的性能和效率。

GROMACS 2023 的另一个令人兴奋的发展是,尖端的 CUDA 图形技术已被用于提高性能,特别是对于小型科学系统。

按照这些说明,尝试将您自己的 GROMACS 案例扩展到许多 GPU ,并加速您的科学研究。

你在哪里可以学到更多?加入 GROMACS 论坛。

 

Tags