LAMMPS and ML-IAP with KOKKOS issue with mliappy

[Shuggurath]

Hi, I installed NequIP using the ML-IAP (with kokkos) pairing when I built lammps, after few attempts on a HPC I manage to do the installation and built lammps with all the ML-IAP related flags. Then after the training I used nequip-prepare- Unfortunately when I was running a MD simulation using lammps I recieved a mliappy failure that blocked the simulation at the zeroth step.
(I know there is similar discussion of this issue without using kokkos but here I managed to build it with all the kokkos flags)

Lammps input file:

units metal
atom_style      atomic
newton on
boundary        p p p
read_data       data.co2

pair_style mliap unified nequip-large.nequip.lmp.pt 0
pair_coeff * * C O

mass 1 12.011
mass 2 15.999

timestep        0.001

neighbor        2.5 bin
neigh_modify    delay 0 every 1 check yes one 2000

minimize 1.0e-4 1.0e-6 10000 100000

velocity        all create 400.0 94221

fix 1 all nvt temp 400.0 400.0 $(100.0*dt)

with the following bash script to submit the job:

#! /bin/bash
#SBATCH --job-name=lammps
###SBATCH —-nodes=1                   
#SBATCH --ntasks-per-node=4     
#SBATCH --cpus-per-task=8       
#SBATCH --gpus-per-node=4     
###SBATCH --mem=480000MB
#SBATCH --time 00:29:00          


export OMP_NUM_THREADS=8

module load profile/deeplrn
module load spack
module load cuda/12.2
module load cudnn
module load openmpi/4.1.6--gcc--12.2.0-cuda-12.2

source nequip-allegro-lammps-new-cudakit118/new-nequip-allegro-lammps-env/bin/activate

srun -n 1 /software/nequip-allegro-lammps-new-cudakit118/lammps/build-mliap/lmp -in in.co2 -k on g 1 -sf kk -pk kokkos newton on neigh half > log.out

deactivate

The error :

Setting up cg/kk style minimization ...
  Unit style    : metal
  Current step  : 0
WARNING: Fix MINIMIZE/kk not compatible with sending data in Kokkos communication (src/KOKKOS/comm_kokkos.cpp:741)
WARNING: Fix with atom-based arrays not compatible with sending data in Kokkos communication, switching to legacy exchange/border communication (src/KOKKOS/comm_kokkos.cpp:754)
WARNING: Fix MINIMIZE/kk not compatible with Kokkos sorting on device (src/KOKKOS/atom_kokkos.cpp:217)
WARNING: Fix with atom-based arrays not compatible with Kokkos sorting on device, switching to legacy host sorting (src/KOKKOS/atom_kokkos.cpp:223)
ERROR: Running mliappy unified compute_forces failure. (src/KOKKOS/mliap_unified_kokkos.cpp:82)
Last input line: minimize 1.0e-4 1.0e-6 10000 100000

Notes:

the cmake I used to compile lammps :

cmake -B build-mliap -S cmake -D CMAKE_BUILD_TYPE=Release -D CMAKE_EXE_LINKER_FLAGS="-pthread -L/prod/opt/compilers/cuda/12.2/none/compat -Wl,-rpath=/prod/opt/compilers/cuda/12.2/none/compat -lcuda" -D CMAKE_CXX_COMPILER=$(pwd)/lib/kokkos/bin/nvcc_wrapper -D CMAKE_CXX_STANDARD=17 -D BUILD_SHARED_LIBS=ON -D BUILD_MPI=ON -D PKG_PYTHON=ON -D PKG_KOKKOS=ON -D PKG_GPU=ON -D GPU_API=cuda -D Kokkos_ENABLE_CUDA=ON -D Kokkos_ENABLE_CUDA_LAMBDA=ON -D Kokkos_ARCH_AMPERE80=ON -D Kokkos_ENABLE_IMPL_MDSPAN=ON -D Kokkos_ENABLE_STD_MDSPAN=OFF -D Kokkos_ENABLE_DEPRECATED_CODE=ON -D PKG_ML-IAP=ON -D PKG_MLIP=ON -D PKG_ML-SNAP=ON -D MLIAP_ENABLE_PYTHON=ON -D PKG_EXTRA-PAIR=on

the nequip-prepare command :

srun nequip-prepare-lmp-mliap $PWD/results/checkpoints/last.ckpt $PWD/nequip-large.nequip.lmp.pt --modifiers enable_OpenEquivariance

(if I remove the openequivariance modifiers I get the warnings but the simulation will run on the CPU only)

[wcwitt]

Can you please modify your code so these warnings don't appear? Maybe just try simple NVE dynamics? That will help us understand what to try next.

WARNING: Fix MINIMIZE/kk not compatible with sending data in Kokkos communication (src/KOKKOS/comm_kokkos.cpp:741)
WARNING: Fix with atom-based arrays not compatible with sending data in Kokkos communication, switching to legacy exchange/border communication (src/KOKKOS/comm_kokkos.cpp:754)
WARNING: Fix MINIMIZE/kk not compatible with Kokkos sorting on device (src/KOKKOS/atom_kokkos.cpp:217)
WARNING: Fix with atom-based arrays not compatible with Kokkos sorting on device, switching to legacy host sorting (src/KOKKOS/atom_kokkos.cpp:223)

[Shuggurath]

Ok, I'm gonna try on a simpler input file to isolate the problem

[cw-tan]

(if I remove the openequivariance modifiers I get the warnings but the simulation will run on the CPU only)

just clarifying this statement --- the simulation should be able to run on GPU without OpenEquivariance. OpenEquivariance provides kernels to accelerate GPU inference, but is not a prerequisite for doing inference on GPUs.

[Shuggurath]

I meant that if I use nequip prepare in this way:
srun nequip-prepare-lmp-mliap $PWD/results/checkpoints/last.ckpt $PWD/nequip-large.nequip.lmp.pt

(Without the --modifiers enable_OpenEquivariance option) the potential still give me the warnings but in the output of lammps I see at the end

Performance: 1.613 ns/day, 14.883 hours/ns, 18.664 timesteps/s, 14.334 katom-step/s
99.6% CPU use with 1 MPI tasks x 1 OpenMP threads