lammps-mdi

MDI engine drivers for LAMMPS — run ML forcefields (MACE, and future models) on GPU via the MolSSI Driver Interface, communicating with a standard LAMMPS binary (no Kokkos compilation needed).

Designed for use with SEAMM, but works with any LAMMPS workflow that supports MDI.

How it works

LAMMPS acts as an MDI driver: it handles atom positions, neighbor lists (at the coarse level), and time integration. The lammps-mdi engine process acts as an MDI engine: it receives coordinates from LAMMPS each step, evaluates the ML model on GPU, and returns energies, forces, and (if periodic) the stress tensor.

mpirun -np 1  mace-mdi  -mdi "..."   ← GPU process: MACE on A100
         : -np 1  lmp  -mdi "..." -in input.dat  ← CPU process: time integration

The two processes communicate over MPI via the MDI protocol.

Supported engines

Installation

See INSTALL_CONDA.md for full HPC instructions.

Usage

As a console script (recommended)

SEAMM_FF=/path/to/model.model \
mpirun --mca mpi_yield_when_idle 1 \
    -np 1 mdi_bind.sh mace-mdi -mdi "-role ENGINE -name MACE -method MPI" \
    : -np 1 mdi_bind.sh lmp -mdi "-role DRIVER -name LAMMPS -method MPI" -in input.dat

Install the ML runtime (PyTorch, vesin, cuEquivariance, MACE) into the LAMMPS environment with:

lammps-mdi install-ml            # or --dry-run to see the plan first

It picks the PyTorch wheel from the machine's NVIDIA driver. A wheel built for a newer CUDA than the driver supports imports without complaint and then reports no GPU, so the first thing to touch the device fails far from the real cause; install-ml avoids that, and keeps mace-torch from pulling a different torch from PyPI. Use --tag to force a wheel tag, since the newest tag a driver supports does not always offer the newest torch for your Python version.

The mace-mdi command accepts several options:

mace-mdi --help

  -mdi MDI_STRING      MDI initialization string [required]
  --model PATH         Path to MACE model (overrides SEAMM_FF)
  --device DEVICE      PyTorch device (default: cuda:0)
  --dtype {float32,float64}
  --enable-cueq        Enable cuEquivariance acceleration
  --enable-oeq         Enable openEquivariance acceleration
  --max-pairs-per-point N
                       Neighbour-list capacity per point (default 256).
                       Raise for dense systems or long cutoffs; the periodic
                       timing line reports the largest value actually seen
                       against this limit.
  --profile-steps N    Profile the first N forward passes with torch.profiler
                       and write a Chrome trace (chrome://tracing, or
                       https://ui.perfetto.dev). Default 0 (disabled).
  --log-level LEVEL    DEBUG / INFO / WARNING / ERROR

From lammps.ini (SEAMM)

[local]
installation = conda   # or modules, or local

gpu-code = mpirun --mca mpi_yield_when_idle 1 \
    -np 1 ~/SEAMM/bin/mdi_bind.sh \
    mace-mdi -mdi "-role ENGINE -name MACE -method MPI" \
    : -np 1 ~/SEAMM/bin/mdi_bind.sh \
    lmp -mdi "-role DRIVER -name LAMMPS -method MPI"

As a Python library

from lammps_mdi import MACEEngine

engine = MACEEngine(
    model_path="/path/to/model.model",
    device="cuda:0",
    default_dtype="float32",
    enable_cueq=True,
)
engine.run("-role ENGINE -name MACE -method MPI")

Shell scripts

The package bundles four helper scripts, installed via lammps-mdi install-scripts:

The CPU/GPU mappings in mdi_bind.sh, gpu_bind.sh, and cpu_bind.sh are currently hard-coded for a dual-GPU EPYC 7763 system. They will be made configurable in a future release.

Requirements

Contributing

Issues and pull requests are welcome at https://github.com/molssi-seamm/lammps-mdi.

License

MIT — see LICENSE.