CuEquivariance Acceleration

CuEquivariance, developed by NVIDIA, provides GPU-accelerated tensor product operations for equivariant neural networks. This integration accelerates Allegro models during both training and inference.

Requirements:

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

Training with CuEquivariance

To enable CuEquivariance acceleration during training, use the model modifier in your configuration file. This replaces standard tensor product operations with optimized GPU kernels:

training_module:
  _target_: nequip.train.EMALightningModule

  # ... other training module configurations ...

  model:
    _target_: nequip.model.modify
    modifiers:
      - modifier: enable_CuEquivarianceContracter
    model:
      _target_: allegro.model.AllegroModel

      # ... your standard Allegro model configuration ...

      compile_mode: compile
      # ^ CuEquivariance composes with train-time compilation

CuEquivariance composes with torch.compile, and can be used in conjunction with train-time compilation.

Inference with CuEquivariance

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

AOT Inductor Compilation

Warning

AOT Inductor compilation with CuEquivariance is not supported for float64 models (i.e., model_dtype: float64). Use TorchScript compilation instead for double precision models.

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

TorchScript Compilation

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

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()

If cuequivariance is not imported before model loading, this error will be thrown: RuntimeError: Could not find schema for cuequivariance_ops::tensor_product_uniform_1d_jit