MPI并行的NEP训练

biogon

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

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

复制代码

下载mpi分支命令

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

复制代码

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

1. make nep_mpi

复制代码

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

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

1. 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网络）

1. #!/bin/bash
   
2. #SBATCH --ntasks=4
   
3. #SBATCH --cpus-per-task=2
   
4. #SBATCH -p gpu
   
5. #SBATCH -w n1,n2
   
6. #SBATCH -J nep_mpi
   
7. 
   
8. CUDA openmpi 环境变量加载
   
9. export OMPI_MCA_coll=^hcoll,han
   
10. export OMPI_MCA_pml=ob1
    
11. export OMPI_MCA_btl=tcp,self,vader
    
12. 
    
13. mpirun /yourdir/GPUMD/src/nep_mpi

复制代码

使用roce/ib网络的脚本

1. #!/bin/bash
   
2. #SBATCH --ntasks=4
   
3. #SBATCH --cpus-per-task=2
   
4. #SBATCH -p gpu
   
5. #SBATCH -w n1,n2
   
6. #SBATCH -J nep_mpi
   
7. 
   
8. CUDA openmpi 环境变量加载
   
9. 
   
10. export OMP_NUM_THREADS=2
    
11. export OMPI_MCA_pml=ucx
    
12. export OMPI_MCA_osc=ucx
    
13. 
    
14. mpirun /yourdir/GPUMD/src/nep_mpi

复制代码

使用CUDA-Aware MPI

1. #!/bin/bash
   
2. #SBATCH --ntasks=4
   
3. #SBATCH --cpus-per-task=2
   
4. #SBATCH -p gpu
   
5. #SBATCH -w n1,n2
   
6. #SBATCH -J nep_mpi
   
7. 
   
8. CUDA openmpi 环境变量加载
   
9. 
   
10. export OMP_NUM_THREADS=2
    
11. export OMPI_MCA_pml=ucx
    
12. export OMPI_MCA_osc=ucx
    
13. export OMPI_MCA_opal_cuda_support=true
    
14. export UCX_MEMTYPE_CACHE=n
    
15. export UCX_TLS=rc,cuda_copy,cuda_ipc,gdr_copy
    
16. 
    
17. mpirun /yourdir/GPUMD/src/nep_mpi

复制代码

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