计算化学公社

标题: MPI并行的NEP训练 [打印本页]

作者
Author: biogon 时间: 2025-12-26 10:40
标题: MPI并行的NEP训练
本帖最后由 biogon 于 2025-12-26 10:57 编辑

GPUMD中原有的NEP模块没有MPI并行只有CUDA下的单机多卡并行，如果想训练结构很多的NEP力场对于没有单机多卡的用户来说不够快
所以在原有的NEP基础上增加了MPI并行的功能，可以让有很多单卡机器的用户通过MPI并行快速的训练力场
代码可以从我的GPUMD的fork的mpi分支获得，该分支基于目前最新的GPUMD 4.7版本，NEP和qNEP力场都可以用于训练

https://github.com/hsianghualite/GPUMD/tree/mpi

复制代码

下载mpi分支命令

git clone -b mpi https://github.com/hsianghualite/GPUMD.git

复制代码

执行如下命令即可编译mpi并行的NEP可执行文件

make nep_mpi

复制代码

编译要有CUDA > 9.0 版本和基于该CUDA作为后端的 openmpi > 4.0 版本

运行方法（使用前必须保证节点之间能免密码SSH访问）
基于TCP网络纯测试命令，使用n1和n2各两个gpu共4个

mpirun -np 4 --host n1:2,n2:2 --mca coll ^hcoll,han --mca pml ob1 --mca btl tcp,self,vader -x PATH -x LD_LIBRARY_PATH bash -lc 'export CUDA_VISIBLE_DEVICES=$OMPI_COMM_WORLD_LOCAL_RANK; exec /yourdir/GPUMD/src/nep_mpi'

复制代码

slurm提交脚本（使用TCP网络）

#!/bin/bash
#SBATCH --ntasks=4
#SBATCH --cpus-per-task=2
#SBATCH -p gpu
#SBATCH -w n1,n2
#SBATCH -J nep_mpi
CUDA openmpi 环境变量加载
export OMPI_MCA_coll=^hcoll,han
export OMPI_MCA_pml=ob1
export OMPI_MCA_btl=tcp,self,vader
mpirun /yourdir/GPUMD/src/nep_mpi

复制代码

使用roce/ib网络的脚本

#!/bin/bash
#SBATCH --ntasks=4
#SBATCH --cpus-per-task=2
#SBATCH -p gpu
#SBATCH -w n1,n2
#SBATCH -J nep_mpi
CUDA openmpi 环境变量加载
export OMP_NUM_THREADS=2
export OMPI_MCA_pml=ucx
export OMPI_MCA_osc=ucx
mpirun /yourdir/GPUMD/src/nep_mpi

复制代码

使用CUDA-Aware MPI

#!/bin/bash
#SBATCH --ntasks=4
#SBATCH --cpus-per-task=2
#SBATCH -p gpu
#SBATCH -w n1,n2
#SBATCH -J nep_mpi
CUDA openmpi 环境变量加载
export OMP_NUM_THREADS=2
export OMPI_MCA_pml=ucx
export OMPI_MCA_osc=ucx
export OMPI_MCA_opal_cuda_support=true
export UCX_MEMTYPE_CACHE=n
export UCX_TLS=rc,cuda_copy,cuda_ipc,gdr_copy
mpirun /yourdir/GPUMD/src/nep_mpi

复制代码

PS1：CUDA-Aware MPI编译方法后续会添加
PS2：目前的并行卡数目上限是nep.in中设置的population数目，后续可能会改进
使用过程中出现问题可以回帖或者在github直接提issue

欢迎光临计算化学公社 (http://bbs.keinsci.com/)