B3 Build native LAMMPS with ML-IAP, Kokkos, and MPI

  • ML-IAP embeds Python: needs a prepared CPython 3.12 archive (Torch, MACE, CuPy) with matching Python.h and libpython3.12.so. The notebook venv will not do.

  • Pinned source and Python archives are private inputs, not in the lesson checkout; obtain and verify them first.

Native LAMMPS build pipeline

The Arrhenius job checks both archive identities, then loads the reviewed GCC/CUDA environment:

build-lammps-mpi.sbatch
26source /software/sse2/init/hpc_init_sse.sh >/dev/null 2>&1
27module load GPU/buildenv-gcccuda/2026.03-cu13.0 >/dev/null 2>&1
28test -x "$(command -v make)"

  • Target: Arrhenius GH200 (Grace CPU, Hopper GPU); CMake enables MPI, ML-IAP, Kokkos, Python, Kokkos_ARCH_ARMV9_GRACE, Kokkos_ARCH_HOPPER90.

  • mpicc/mpicxx must match the MPI stack used at run time.

  • Not portable: JUPITER needs its own GH200 build against its compiler/MPI stack.

build-lammps-mpi.sbatch
51cmake -S "$source_dir/cmake" -B "$scratch/mpi-build" -G 'Unix Makefiles' \
52  -D CMAKE_BUILD_TYPE=Release -D CMAKE_INSTALL_PREFIX="$prefix" \
53  -D CMAKE_C_COMPILER="$(command -v mpicc)" \
54  -D CMAKE_CXX_COMPILER="$source_dir/lib/kokkos/bin/nvcc_wrapper" \
55  -D MPI_C_COMPILER="$(command -v mpicc)" \
56  -D MPI_CXX_COMPILER="$(command -v mpicxx)" \
57  -D BUILD_MPI=ON -D BUILD_SHARED_LIBS=ON -D LAMMPS_EXCEPTIONS=ON \
58  -D PKG_KOKKOS=ON -D PKG_ML-IAP=ON -D PKG_ML-SNAP=ON \
59  -D PKG_PYTHON=ON -D MLIAP_ENABLE_PYTHON=ON \
60  -D Kokkos_ENABLE_SERIAL=ON -D Kokkos_ENABLE_OPENMP=ON \
61  -D Kokkos_ENABLE_CUDA=ON -D Kokkos_ENABLE_CUDA_LAMBDA=ON \
62  -D Kokkos_ARCH_ARMV9_GRACE=ON -D Kokkos_ARCH_HOPPER90=ON \
63  -D FFT_KOKKOS=KISS \
64  -D Python_EXECUTABLE="$python" \
65  -D Python_INCLUDE_DIR="$scratch/python/include/python3.12" \
66  -D Python_LIBRARY="$scratch/python/lib/libpython3.12.so" \

  • Runs in a reviewed Slurm build allocation; writes an MPI runtime candidate archive.

  • Does not overwrite an existing LAMMPS, submit itself, or prove multi-GPU MD correct.

  • Excerpts omit identity checks, private paths and validation; see the full build job.

  • One-rank smoke first, then the optional 1/2/4-GPU scaling run in B8.

Use the tested GPU path

The MACE ML-IAP guide recommends unified ML-IAP, CUDA Kokkos, Newton on, half neighbour lists. Our example sets these via the LAMMPS library, not lmp:

lammps_si.py
23    count = len(atoms)
24    lmp = lammps.lammps(cmdargs=["-log", "none", "-screen", "none", "-k", "on", "g", "1"])
25    try:
26        if not lmp.has_package("KOKKOS") or not lmp.has_style("pair", "mliap/kk"):
27            raise RuntimeError("LAMMPS needs KOKKOS and mliap/kk")
28        if not lmp.has_style("fix", "nve/kk"):
29            raise RuntimeError("LAMMPS needs nve/kk")
30        activate_mliappy_kokkos(lmp)
31        length = float(atoms.cell.lengths()[0])
32        for command in (
33            "package kokkos newton on neigh half",
34            "units metal", "atom_style atomic/kk", "newton on", "boundary p p p",

  • mliap/kk unified then applies the exported model; the complete example adds atoms, integrator, timing and cleanup.

  • These settings are the baseline for all comparisons here.

  • Kokkos does not imply active cuEquivariance kernels; check runtime and export before crediting them with a speedup.

Export the same checkpoint for ML-IAP

ALCHEMI reads the MACE file; LAMMPS reads an ML-IAP export. In the native Python environment, export to a fresh path outside Git. The exporter checks the model hash, needs one visible GPU and will not overwrite:

export MLIP_MLIAP_MODEL=/path/outside/git/mace-mp-0a-small-mliap.pt
python examples/export_mace_mliap.py \
  --model "$MLIP_MODEL" --output "$MLIP_MLIAP_MODEL"

  • Record the printed export hash with the runtime identity; on Arrhenius it should match reference/model.toml (a fresh export did, and passed a short one-rank run).

  • Elsewhere, recheck export and one-rank run; a matching hash is not a force or trajectory validation.

  • Exporter source: model check and overwrite refusal.

Arrhenius multi-rank profile (site specific, not portable):

  • Site MPICH, PMI2/CXI, gpu/aware on, peer-visible GPUs; hiding other GPUs from a rank caused an ML-IAP ghost-exchange failure.

  • Elsewhere, recheck compiler/MPI ABI, device assignment, transport and Slurm options. JUPITER’s native setup is separate; a future Leonardo profile should be too.

  • NCCL is for ALCHEMI’s DomainParallel path, not a replacement for this build’s MPI.