How to setup cuEquivariance fronted for speed up MACE with kokkos and GPU-domain balancing for memory efficiency

[so2koo]

Hi everyone,

I want to share how I improved speed and memory efficiency for large, dynamic systems (mine is 2k atoms) and ask if there is relevant guideline or advanced ideas.

1. Match cuEquivariance with cuEquivariance-ops

Initial compilation was done with cuequivariance 0.10.0 was probably installed automatically as the newest frontend when MACE was installed, while the older compartible ops package remained at 0.4.0. Current version 0.6.1 was selected because its precompiled operations still contain native V100 sm_70 kernels.

Previous (Incompatible) stack: 0.562 steps/s
Current (Coherent) 0.6.1 stack: 8.700 steps/s
Improvement attributable mainly to coherent native-kernel execution: approximately 15.5×.

Pip did not enforce matching versions because these packages were independently versioned without a strict equality requirement.

2. GPU domain balance

General idea on memory allocation

• LAMMPS divides atoms spatially among MPI ranks, normally one rank per GPU.
• Every GPU loads a complete copy of the MACE model. Model weights are not divided.
• Each GPU constructs a graph for its local atoms plus neighboring ghost atoms.
• MACE evaluates local energies and forces using PyTorch on that GPU.
• Ghost information is exchanged between ranks at domain boundaries.

However if your system includes large deformation, interface, void, density redistribution,causing memory imbalance from LAMMPS' spatial decomposition,
You can try static balancing after warm-up run to obtain an initial estimate of a better domain boundary. Because MACE pair timing can include MPI communication and synchronization, this automatically determined boundary should be treated as a starting point rather than the final optimum. Benchmark several nearby fixed domain boundaries and compare sustained MD speed, per-GPU memory usage, and workload distribution, then select the fastest decomposition that maintains sufficient GPU-memory headroom for the production run.

run 200
balance 1.0 shift xyz 10 1.05 weight time 0.8