Accelerators#

This document describes available accelerator integrations in SevenNet and their installation guide.

Caution

We do not support CuEquivariance for LAMMPS: Torch. You must use LAMMPS: ML-IAP for CuEquivariance.

FlashTP, CuEquivariance, and OpenEquivariance provide acceleration for both SevenNet training and inference.

Tip

For small systems, FlashTP with LAMMPS: Torch shows performance advantage over cuEquivariance with LAMMPS: ML-IAP. A performance crossover occurs at around 10³ atoms, beyond which cuEquivariance becomes more efficient. (For more information, check section 2.7 of the SevenNet-Omni paper)

FlashTP with LAMMPS: Torch is generally faster than FlashTP with LAMMPS: ML-IAP.

CuEquivariance#

CuEquivariance is an NVIDIA Python library designed to facilitate the construction of high-performance geometric neural networks using segmented polynomials and triangular operations. CuEquivariance accelerates SevenNet during training, inference with ASE and LAMMPS via ML-IAP.

Requirements#

  • Python >= 3.10

  • cuEquivariance >= 0.6.1

Installation#

After installation of SevenNet, install cuEquivariance via the pip extras:

pip install sevenn[cueq12]  # For CUDA 12.x
pip install sevenn[cueq13]  # For CUDA 13.x

This will install cuequivariance, cuequivariance-torch, and the corresponding cuequivariance-ops-torch-cu{12,13}.

Alternatively, you can install cuEquivariance manually following their guideline.

Note

Some GeForce GPUs do not support pynvml, causing pynvml.NVMLError_NotSupported: Not Supported. Then try a lower cuEquivariance version, such as 0.6.1.

Check your installation:

python -c 'from sevenn.nn.cue_helper import is_cue_available; print(is_cue_available())'
True

FlashTP#

FlashTP, presented in FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials, is a high-performance Tensor-Product library for Machine Learning Interatomic Potentials (MLIPs).

FlashTP accelerates SevenNet during both training and inference, achieving up to ~4× speedup.

Requirements#

  • Python >= 3.10

  • flashTP >= 0.1.0

  • CUDA toolkit >= 12.0

Installation#

Choose CUDA_ARCH_LIST for your GPU(s) (see compute compatibility)

git clone https://github.com/SNU-ARC/flashTP.git
cd flashTP
pip install -r requirements.txt
CUDA_ARCH_LIST="80;90" pip install . --no-build-isolation

Note

During installation of FlashTP, subprocess.CalledProcessError: ninja ... exit status 137 typically indicates out-of-memory during compilation. Try reducing the build parallelism:

export MAX_JOBS=1

Check your installation:

python -c 'from sevenn.nn.flash_helper import is_flash_available; print(is_flash_available())'
True

For more information, see FlashTP.

OpenEquivariance (experimental)#

Caution

OpenEquivariance support in SevenNet is currently experimental.

OpenEquivariance (oEq) is a library for acceleration of equivariant tensor products. Unlike cuEquivariance, its kernels are compiled once per irreps configuration, eliminating runtime recompilation overhead during training on diverse datasets.

Requirements#

  • Python >= 3.10

  • openequivariance

Installation#

pip install sevenn[oeq]

Check your installation:

python -c 'from sevenn.nn.oeq_helper import is_oeq_available; print(is_oeq_available())'
True

For more information, see OpenEquivariance.

Usage#

After the installation, you can leverage the accelerator with appropriate flag (--enable_cueq), (--enable_flash), or (--enable_oeq) options