ML-IAP inference performs substantially worse than Symmetrix on ARM64 and CUDA13

[MicPellegrino]

Hi,

We are simulating high-entropy alloys crystals of various size with a fine-tuned version of OMAT (medium) on the Arrhenius system in Sweden. Even after numerous interactions with the support team, we are unable to understand why MACE simulations with LAMMPS+MLIAP are substantially less performant than the ones with LAMMPS+Symmetrix:

------------------------------------------
#atoms | 54	    686    2662   6750   13718
Symmetrix [ns/day]
1 GPU  | 22.9   15.1   5.78   2.56   1.28
2 GPUs | 17.4   15.2   8.45   4.51   2.35
4 GPUs [ 16.5   16.0   12.2   7.18   4.10
8 GPUs | 15.4   15.3   13.1   10.3   6.40
ML-IAP + cuEq [ns/day]
1 GPU  | 3.16   3.12   3.14   2.95   2.38
2 GPUs | 2.60   2.52   2.50   2.30   2.15
4 GPUs | 2.72   2.77   2.68   2.60   2.58
8 GPUs | 2.63   2.65   2.46   2.50   2.41
------------------------------------------

The particularity of the Arrhenius system is that nodes mount GH200 superchips, with NVidia GPUs and CPUs, the latter having a ARM64 architecture. There doesn't seem to be any obvious MPI or I/O chokepoint. This is the timing breakdown for a 8 GPUs job on 2 nodes (13718 atoms) for MLIAP:

Section |  min time  |  avg time  |  max time  |%varavg| %total
---------------------------------------------------------------
Pair    | 66.143     | 67.203     | 69.341     |  14.4 | 93.66
Neigh   | 0.0088742  | 0.0093421  | 0.010292   |   0.5 |  0.01
Comm    | 2.1662     | 4.2939     | 5.3503     |  56.7 |  5.98
Output  | 0.001877   | 0.0024008  | 0.0027991  |   0.6 |  0.00
Modify  | 0.16845    | 0.17819    | 0.18421    |   1.3 |  0.25
Other   |            | 0.06382    |            |       |  0.09

And this is the Symmetrix counterpart:

Section |  min time  |  avg time  |  max time  |%varavg| %total
---------------------------------------------------------------
Pair    | 25.533     | 25.615     | 25.68      |   0.9 | 94.85
Neigh   | 0.0076401  | 0.0080515  | 0.0083644  |   0.2 |  0.03
Comm    | 1.081      | 1.1425     | 1.2162     |   3.8 |  4.23
Output  | 0.0021162  | 0.0030452  | 0.0050268  |   2.2 |  0.01
Modify  | 0.1571     | 0.17325    | 0.18442    |   2.5 |  0.64
Other   |            | 0.06507    |            |       |  0.24

Mind that we are using the same model, with single-precision weigths, and we are also using cuEquivariance when simulating with MLIAP.

Is it very odd that MLIAP shows almost no strong or weak scaling (no. atoms or no. GPUs).

We also noticed a ton of KOKKOS warnings, to the point that I would not even know where to start listing them, and this issue post would become pages long. I am not sure if that's a hint, but we see warning both when building MLAIP and Symmetrix. So I doubt it's the culprit.

EDIT - I think I forgot to state why this is an issue. At the current state, Symmetrix only supports models with l_max=3 and does not support per-atom virials. In our project, we plan to use an ensemble of models with l_max=2, and plan to explore thermo-mechanical properies of the alloys (hence the virials).

All the details for reproducing the builds are below.

CMake options for ML-IAP:

cmake \
    -C kokkos-cuda-nowrapper_arrhenius.cmake \
    -D CMAKE_BUILD_TYPE=Release \
    -D CMAKE_INSTALL_PREFIX=$BUILD_DIR \
    -D BUILD_MPI=ON \
    -D BUILD_OMP=ON \
    -D BUILD_SHARED_LIBS=ON \
    -D PKG_KOKKOS=ON \
    -D Kokkos_ARCH_NATIVE=ON \
    -D Kokkos_ARCH_HOPPER90=ON \
    -D Kokkos_ENABLE_AGGRESSIVE_VECTORIZATION=ON \
    -D Kokkos_ENABLE_SERIAL=ON \
    -D Kokkos_ENABLE_OPENMP=ON \
    -D Kokkos_ENABLE_NEIGH_HALF=ON \
    -D Kokkos_ENABLE_CUDA=ON \
    -D PKG_ML-IAP=ON \
    -D PKG_ML-MACE=ON \
    -D PKG_ML-SNAP=ON \
    -D MLIAP_ENABLE_PYTHON=ON \
    -D PKG_PYTHON=ON \
    -D PKG_OPENMP=ON \
    -D PKG_OPT=ON \
    -D PKG_KSPACE=ON \
    -D PKG_MANYBODY=ON \
    -D PKG_MOLECULE=ON \
    -D PKG_RIGID=ON \
    -D PKG_MC=ON \
    -D PKG_MISC=ON \
    -D PKG_MEAM=ON \
    -D PKG_EXTRA-COMPUTE=ON \
    -D PKG_EXTRA-DUMP=ON \
    -D PKG_EXTRA-FIX=ON \
    -D PKG_EXTRA-MOLECULE=ON \
    -D PKG_EXTRA-PAIR=ON \
    ../cmake

CMake options for Symmetrix:

cmake \
    -D CMAKE_BUILD_TYPE=Release \
    -D CMAKE_CXX_STANDARD=20 \
    -D CMAKE_CXX_STANDARD_REQUIRED=ON \
    -D CMAKE_CXX_FLAGS="${CMAKE_CXX_FLAGS} -march=native -ffast-math" \
    -D BUILD_SHARED_LIBS=ON \
    -D PKG_KOKKOS=ON \
    -D BUILD_OMP=ON \
    -D BUILD_MPI=ON \
    -D PKG_FEP=ON \
    -D PKG_MOLECULE=ON \
    -D PKG_MANYBODY=ON \
    -D PKG_KSPACE=ON \
    -D PKG_RIGID=ON \
    -D PKG_REAXFF=ON \
    -D PKG_MISC=ON \
    -D PKG_ASPHERE=ON \
    -D PKG_CLASS2=ON \
    -D PKG_MC=ON \
    -D PKG_MEAM=ON \
    -D PKG_EXTRA-COMPUTE=ON \
    -D PKG_EXTRA-DUMP=ON \
    -D PKG_EXTRA-FIX=ON \
    -D BUILD_SHARED_LIBS=ON \
    -D SYMMETRIX_KOKKOS=ON \
    -D SYMMETRIX_SPHERICART_CUDA=ON \
    -D Kokkos_ARCH_NATIVE=ON \
    -D Kokkos_ARCH_HOPPER90=ON \
    -D Kokkos_ENABLE_AGGRESSIVE_VECTORIZATION=ON \
    -D Kokkos_ENABLE_SERIAL=ON \
    -D Kokkos_ENABLE_OPENMP=ON \
    -D Kokkos_ENABLE_NEIGH_HALF=ON \
    -D Kokkos_ENABLE_CUDA=ON \
    -D CMAKE_INSTALL_PREFIX=$BUILD_DIR \
    ../cmake

Compilers:

GCC 14.3.0
MPICH 4.3.2
CUDA 13.0.3

Library versions (for MLIAP):

Package                       Version
----------------------------- -----------
ase                           3.29.0
ConfigArgParse                1.7.5
contourpy                     1.3.3
cuda-pathfinder               1.5.6
cuequivariance                0.9.1
cuequivariance-ops-cu13       0.9.1
cuequivariance-ops-torch-cu13 0.9.1
cuequivariance-torch          0.9.1
cupy-cuda13x                  14.1.1
cycler                        0.12.1
Cython                        3.2.0
e3nn                          0.4.4
filelock                      3.29.0
fonttools                     4.63.0
fsspec                        2026.4.0
gitdb                         4.0.12
GitPython                     3.1.53
h5py                          3.16.0
Jinja2                        3.1.6
kiwisolver                    1.5.0
lammps                        2026.2.11
lightning-utilities           0.15.3
lmdb                          2.3.0
mace-torch                    0.3.16
MarkupSafe                    3.0.3
matplotlib                    3.11.1
matscipy                      1.2.0
mpmath                        1.3.0
networkx                      3.6.1
numpy                         2.1.0
nvidia-cublas                 13.0.0.19
nvidia-cuda-cupti             13.0.48
nvidia-cuda-nvrtc             13.0.48
nvidia-cuda-runtime           13.0.48
nvidia-cudnn-cu13             9.13.0.50
nvidia-cufft                  12.0.0.15
nvidia-cufile                 1.15.0.42
nvidia-curand                 10.4.0.35
nvidia-cusolver               12.0.3.29
nvidia-cusparse               12.6.2.49
nvidia-cusparselt-cu13        0.8.0
nvidia-ml-py                  13.610.43
nvidia-nccl-cu13              2.27.7
nvidia-nvjitlink              13.0.39
nvidia-nvshmem-cu13           3.3.24
nvidia-nvtx                   13.0.39
opt_einsum                    3.4.0
opt-einsum-fx                 0.1.4
orjson                        3.11.9
packaging                     26.2
pandas                        3.0.3
pillow                        12.2.0
pip                           25.1.1
platformdirs                  4.11.0
prettytable                   3.18.0
pyparsing                     3.3.2
python-dateutil               2.9.0.post0
python_hostlist               2.3.0
PyYAML                        6.0.3
scipy                         1.18.0
setuptools                    78.1.0
six                           1.17.0
smmap                         5.0.3
sympy                         1.14.0
torch                         2.9.0+cu130
torch-ema                     0.3
torchaudio                    2.9.0
torchmetrics                  1.9.0
torchvision                   0.24.0
tqdm                          4.69.0
triton                        3.5.0
typing_extensions             4.15.0
wcwidth                       0.8.2

[alinelena]

while mlip being slower than symmetrix is known I am somehow surprised by some of the results.

can you share some input (preferably with a foundation model) and output(with warnings)

also you do not need
-D PKG_ML-MACE=ON \

[MicPellegrino]

The input file is the following (for both ML-IAP and Symmetrix):

units metal
atom_style atomic
atom_modify map yes
boundary p p p
newton on

variable inputfile string "${conffolder}/AlTiMoWCr_${nbox}_bcc.data"
variable outputfile2 string "${conffolder}/AlTiMoWCr_${nbox}_bcc_final.data"

read_data ${inputfile}

# Importing force field (ML-IAP)
pair_style mliap unified ${modelfile} 0
pair_coeff * * Al Ti Mo W Cr

# Importing force field (Symmetrix)
# pair_style symmetrix/mace
# pair_coeff * * ${modelfile} Al Ti Mo W Cr

# Preparation and run
velocity all create 300.0 4928459 rot yes dist gaussian
neighbor 1.0 bin
neigh_modify every 2 delay 10 check yes
fix nvt1 all nvt temp 300.0 300.0 1.0
timestep 0.001
thermo 500
run 500
reset_timestep 0
run 2000

write_data ${outputfile2}

The .pt file for ML-IAP has been created on the GPU nodes and in the same virtual environment used to launch simulations.

We get this warning when running (form LAMMPS .log file):

Kokkos::OpenMP::initialize WARNING: OMP_PROC_BIND environment variable not set
  In general, for best performance with OpenMP 4.0 or better set OMP_PROC_BIND=spread and OMP_PLACES=threads
  For best performance with OpenMP 3.1 set OMP_PROC_BIND=true
  For unit testing set OMP_PROC_BIND=false

MPI detected: For OpenMP binding to work as intended, MPI ranks must be bound to exclusive CPU sets.

/home/fabiano/venvs/lammps_11Feb2026-mliap/lib/python3.13/site-packages/e3nn/o3/_wigner.py:10: UserWarning: Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.
  _Jd, _W3j_flat, _W3j_indices = torch.load(os.path.join(os.path.dirname(__file__), 'constants.pt'))
W0727 12:30:40.100000 2513665 torch/fx/_symbolic_trace.py:52] is_fx_tracing will return true for both fx.symbolic_trace and torch.export. Please use is_fx_tracing_symbolic_tracing() for specifically fx.symbolic_trace or torch.compiler.is_compiling() for specifically torch.export/compile.

The problem is that if we set OMP_PROC_BIND=spread and OMP_PLACES=threads we get even worse performance (up to 10 times slower). We have tried any possible combination of the two env. variables, and at the end the only option that gives decent performance is:

unset OMP_PROC_BIND
unset OMP_PLACES
export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK

(I also remind that we get the same KOKKOS warnings when running with Symmetrix, so probably not the culprit?).
Regarding the Torch warning, even after 1 year of running MACE we still haven't figured out how to get rid of it.

I am happy to sare more information on input and log files. I cannot however share the fine-tuned model publicly. I can only say it's a version of this: https://github.com/ACEsuit/mace-mp/releases/download/mace_omat_0/mace-omat-0-medium.model, fine-tuned for the WMoAlTiCr alloy we are simulating.

[MicPellegrino]

P.S. here are the configuration files, in case someone wants to try replicate the simulations with OMAT:
confs-altimowcr.tar.gz