Optimization of MACE in LAMMPS

[braydenbanks323]

Hi all,

I am running LAMMPS with the MACE package installed using some pretrained foundational models. I noticed some slow simulation times, so I am trying to parallelize my run to improve performance. I am running on a cluster where I submit jobs through a slurm script.
My slurm script contains

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=2
#SBATCH --cpus-per-task=16
#SBATCH --gres=gpu:2

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export OMP_NUM_THREADS=16

ulimit -l unlimited
mpirun -np 2 lmp -k on g 2 -sf kk -e both -in water.in

My input file has the lines:
atom_style atomic
atom_modify map yes
pair_style mace/kk no_domain_decomposition
pair_coeff * * /path/to/mace/models/MACE-matpes-r2scan-omat-ft.model-lammps.pt O H

Yet when I try to run this script, I get this error:
Exception: Specified device cuda:1 does not match device of data cuda:0
Exception raised from make_tensor at aten/src/ATen/Functions.cpp:25 (most recent call first):

Has anybody done some optimizing work like this? What have you found to work?

[ilyes319]

Did you try using cuequivariance with the MLIAP lammps interface ?