Can you guys optimize your model speed?

[usccolumbia]

Model	Energy/atom (eV)	Max force (eV/A)	Steps	Model load (s)	Relax total (s)	Cold single-CIF total (s)
MatterSim	-5.081119	0.000025	0	2.945	0.491	3.436
EquFlashV2	-5.087032	0.002869	27	13.384	123.136	136.520
EquiformerV3	-5.087420	0.006962	29	4.734	3.632	8.366
NequIP-OAM-L	-5.096622	0.019583	24	4.399	6.796	11.195
NequIP-OAM-XL	-5.089541	0.001770	27	5.204	10.596	15.799

it costs 123 seconds compared to other's 3-10 seconds for relaxing a single same cif!

[usccolumbia]

EquFlashV2 uses similar GNN as others, why it takes so much time so slow?
8:19 PM
Yes, EquFlashV2 is “GNN-like,” but the slowness here is mostly implementation/runtime overhead, not just model type.

From our strict fmax=0.0001 run:

EquFlashV2 total_wall_s = 135.91 s
initial_force_eval_s = 113.76 s
relaxation_s = 22.12 s
optimizer steps = 89
So the main problem is not the 89 relaxation steps. The huge cost is the first force evaluation.

Likely reasons:

CUDA/Triton/torch compile overhead
EquFlashV2 is probably compiling kernels or building optimized graph code on the first call. That first atoms.get_forces() costs ~114 s.

GGNN/FairChem stack overhead
The Samsung GGNN wrapper loads through FairChem-style infrastructure, with lots of config/model machinery. It is heavier than the NequIP compiled TorchScript path.

Model architecture is heavy
The log showed EquFlashV2 has about 44.9M parameters, with high-order equivariant features up to lmax=4 and many interaction blocks. “GNN” does not mean cheap; equivariant GNNs can be very expensive.

[ghijkim]

@usccolumbia
Thank you for reporting this.

Although EquFlashV2 is also an equivariant GNN, it is designed with a larger model capacity and uses a higher lmax (l = 4) than many other e3nn-based models. In particular, the higher lmax significantly increases the computational cost of the Clebsch–Gordan tensor product operations, which are one of the dominant computational bottlenecks during inference.

That said, the reported runtime still appears to be unusually high. While EquFlashV2 is expected to be slower than the comparison models, it should not be slower by this margin. We recommend using the cuEquivariance 0.6.0 kernels if you are currently using a version other than 0.6.0.

If you could share your environment (e.g., package versions or an environment file), we'd be happy to investigate the issue and identify the cause of the performance discrepancy.

[gaoispingisme-dotcom]

Hello developers,
In our tests, EquFlashV2 is also much slower than the other models. Small systems often take around 180-210 s per case, and the runtime is almost constant for small systems, suggesting a large fixed cold-start or first-forward overhead. @ghijkim
Thank you very much.

Env:

• EquFlashV2 checkpoint: equflashv2_oam.pt
• Model size: about 44.9M parameters
• Checkpoint size: about 858 MB
• Environment includes:
  • torch 2.9.1+cu126
  • fairchem-core 1.10.0
  • e3nn 0.5.6
  • cuequivariance 0.6.0
  • cuequivariance-torch 0.6.0
  • cuequivariance-ops-torch-cu12 0.6.0
  • ASE 3.26.0
  • GGNN 0.1

[choant2003]

@usccolumbia @gaoispingisme-dotcom
Thanks for reporting this.

We checked the inference behavior on an A100 40GB using a simple Si diamond structure with small random displacements:

import numpy as np
from ase.build import bulk
import ase.io
from GGNN.common.calculator import UCalculator
import time
import torch

atoms = bulk(
    name="Si",
    crystalstructure="diamond",
    a=5.43,
    cubic=True,
)

atoms = atoms.repeat((2, 2, 2))

sigma = 0.002  # Angstrom
seed = 42
rng = np.random.default_rng(seed)

noise = rng.normal(0.0, sigma, size=atoms.positions.shape)
atoms.positions += noise
atoms.wrap()

calc = UCalculator(checkpoint_path="eqv2_oam.pt", cpu=False)
atoms.calc = calc

for i in range(200):
    noise = rng.normal(0.0, sigma, size=atoms.positions.shape)
    atoms.positions += noise
    atoms.wrap()

    torch.cuda.synchronize()
    begin = time.time()

    atoms.get_forces()

    torch.cuda.synchronize()
    end = time.time()

    print(end - begin)

In this test, the first step takes around 100 seconds, but after the third step, most force evaluations take about 0.2 seconds.

EquFlashV2 uses the cuequivariance library. The large latency in the first few calls mainly comes from the internal compilation process after the model is first loaded. After this warm-up/compile stage, the per-step inference speed should be much closer to the other models.

For this reason, in our relaxation and MD workflows, we usually keep the calculator/model instance alive and reuse it across configurations instead of repeatedly loading a new model for each structure.

I have also seen recent attempts in newer versions of cuequivariance to cache compiled kernels and reduce this first-step overhead. We will try this and let you know if it improves the cold-start time.

If you still observe a persistent speed gap after the initial warm-up steps, please share the approximate configuration and code you used to run the benchmark. I will take a closer look on my side as well.

[gaoispingisme-dotcom]

@ghijkim
Thank you very much for your reply. After warm-up, EquFlashV2 is indeed not slow.
Also, EquFlashV2 appears to be very memory efficient compared with other universal MLIPs. Great work!