Poor performance on LAMMPS+KOKKOS, can I do something about it?

[MicPellegrino]

Hi,
I am trying out MACE on a small system of 500 Cu atoms, based on LAMMPS EAM benchmark (https://www.lammps.org/bench.html#eam).
The build options are based on https://mace-docs.readthedocs.io/en/latest/guide/lammps.html:

cmake \
    -D CMAKE_BUILD_TYPE=Release 
    -D CMAKE_INSTALL_PREFIX=$(pwd) 
    -D CMAKE_CXX_STANDARD=17 
    -D CMAKE_CXX_STANDARD_REQUIRED=ON 
    -D PKG_KOKKOS=ON 
    -D Kokkos_ENABLE_CUDA=ON 
    -D CMAKE_CXX_COMPILER=$(pwd)/../lib/kokkos/bin/nvcc_wrapper 
    -D Kokkos_ARCH_NATIVE=ON 
    -D Kokkos_ARCH_AMPERE86=ON 
    -D MKL_INCLUDE_DIR=$(pwd) 
    -D CMAKE_PREFIX_PATH=$(pwd)/../../libtorch-gpu 
    -D PKG_ML-MACE=ON 
    ../cmake

I am running on an NVIDIA RTX A4000 GPU on my workstation, getting a whopping performance of

• 0.038ns/day with MACE_MPtrj_2022.9.model;
• 0.171ns/day with 2023-12-10-mace-128-L0_energy_epoch-249.model.

I wonder if I am stuck with this performance or if there's a way to possibly improve it a bit. The only two things that I could notice are the following:

• When loading the MACE model, LAMMPS output The torch_float_dtype is: Double, is this expected? Does it mean I am running in double precision?
• Performance doesn't improve when building with -D Kokkos_ENABLE_OPENMP=yes -D BUILD_OMP=yes and running on multiple threads, while the examples in https://mace-docs.readthedocs.io/en/latest/guide/lammps.html seem to suggest multi-threading could help performance.
• The good-old neighbour list: based on what I've read on ML potentials it could very well not matter at all, but it would be good to get some tips about it.

Minor issues:

• For any system with more than ~1000 atoms I get a memory error similar to Multi-GPU run on LAMMPS facing "CUDA out of memory" error #777 and multi-gpu lammps issue #322, but this could simply be due to me being stuck with 16GB.
• If I don't specify a mass command in the LAMMPS script I get the following error: ERROR: Not all per-type masses are set. Type 1 is missing. (src/atom.cpp:763); not the end of the world, but this is a bit annoying as I have no idea what mass has been used when training the MACE model. I haven't seen anyone else reporting this error for MACE, so I don't know if it is expected or not.

[ilyes319]

Hey @MicPellegrino,

• The first thing is I recommend you use a smaller model, including the new MPA-0 model https://github.com/ACEsuit/mace-mp/releases/tag/mace_mpa_0 which is a medium model, and will provide a speed up compared to the large model you are using.

• If you can get away with using ASE instead of LAMMPS, you can use the cuequivariance library to accelerate MACE (https://mace-docs.readthedocs.io/en/latest/guide/cuda_acceleration.html). It is not yet compatible with LAMMPS.

• Use a better GPU :) A4000 are 4 years old, any more recent GPU will give you a big speed up. The recent gamer GPUs have a great ratio of speed/cost (4080/4090). On modern hardware (for example H100s) you get closer to several ns/day for your system size, so x10-x100 speed up.

• If you have CPUs at hand, we will release soon some CPU evaluator for MACE by @wcwitt. Stay tuned.

• Indeed the default is double precision, you can turn the model to single precision in Python before exporting it to lammps. It is usually a x2 speed uo.

[MicPellegrino]

Thanks for the quick reply. The fact that

any more recent GPU will give you a big speed up

is at least reassuring, since there's room for improvement.

[MicPellegrino]

Indeed the default is double precision, you can turn the model to single precision in Python before exporting it to lammps. It is usually a x2 speed uo.

You mean using the MACE create_lammps_model.py script?

[ilyes319]

yep indeed, if you install the latest version of the main branch, you have an extra input to the script call dtype.

[MicPellegrino]

I see. I tried but unfortunately the simulation is unstable when running with the single-precision model.

[ilyes319]

Not surprising, we see that improved a lot with CUDA kernels.