CuEquivariance Acceleration

Warning

The NequIP CuEquivariance acceleration is still under development.

CuEquivariance, developed by NVIDIA, provides GPU-accelerated tensor product operations for equivariant neural networks. This integration accelerates NequIP models during inference, with training support currently a work in progress.

Warning

float64 model compatibility with CuEquivariance remains untested. For double precision models (i.e., model_dtype: float64), compatibility is not guaranteed.

Requirements:

pip install cuequivariance-torch cuequivariance-ops-torch-cu12

ASE Inference with CuEquivariance

For ASE inference, you can compile your trained model with CuEquivariance acceleration enabled using nequip-compile:

ASE AOT Inductor Compilation

nequip-compile \
    path/to/model.ckpt \
    path/to/compiled_model.nequip.pt2 \
    --device cuda \
    --mode aotinductor \
    --target ase \
    --modifiers enable_CuEquivariance

ASE TorchScript Compilation (PyTorch < 2.10 only)

nequip-compile \
    path/to/model.ckpt \
    path/to/compiled_model.nequip.pth \
    --device cuda \
    --mode torchscript \
    --target ase \
    --modifiers enable_CuEquivariance

To use the compiled model, you must import cuequivariance_torch before loading:

import cuequivariance_torch
from nequip.integrations.ase import NequIPCalculator

# Load the compiled model
calc = NequIPCalculator.from_compiled_model(
    "path/to/compiled_model.nequip.pt2/pth",
    device="cuda"
)

# Use with ASE atoms object
atoms.calc = calc
energy = atoms.get_potential_energy()
forces = atoms.get_forces()

LAMMPS ML-IAP Integration

CuEquivariance can also be used with LAMMPS through the ML-IAP interface. This provides a stable integration path for production molecular dynamics simulations with CuEquivariance acceleration.

To prepare a model for LAMMPS ML-IAP with CuEquivariance:

nequip-prepare-lmp-mliap \
  /path/to/model_file \
  /path/to/output.nequip.lmp.pt \
  --modifiers enable_CuEquivariance

Where model_file can be either a checkpoint file (.ckpt) or package file. The resulting .nequip.lmp.pt file can be used directly in LAMMPS scripts with the pair_style mliap command. See the ML-IAP documentation for complete usage instructions and examples.