[BUG] AOT inductor compilation fails

[mjrs33]

Describe the bug
When I attempted to compile the pre-trained model from nequip.net, the following error occurred:

AssertionError: Compilation check MaxAbsError: 1.24863 (tol: 5e-05) for field `atomic_energy`.

To Reproduce

nequip-compile \
  nequip.net:mir-group/NequIP-OAM-L:0.1 \
  mir-group__NequIP-OAM-L__0.1.nequip.pt2 \
  --mode aotinductor \
  --device cuda \
  --target ase \
  --modifiers enable_OpenEquivariance

Expected behavior
MaxAbsError must be less than tol: 5e-05

Environment (please complete the following information):

• OS: Rocky Linux 9.5 (Blue Onyx)
• python version: Python 3.12.9
• python environment:
  • nequip==0.16.0
  • e3nn==0.5.8
  • openequivariance==0.4.1
  • torch==2.9.1 (2.9.1+cu128)

$ nvcc --version
Cuda compilation tools, release 12.8, V12.8.61
Build cuda_12.8.r12.8/compiler.35404655_0

$ python -c "import torch; print(torch.version.cuda)"
12.8

[kavanase]

Hi @mjrs33, thank you for reporting this.
Is your OpenEquivariance installation the latest official release, or the latest commit on their GitHub (which has the same numeric version)? I'm not sure, but this could potentially be related.

Regardless, would you be able to try compiling with other options for mode, device and target and let us know if this issue is only limited to this combination of choices?

[mjrs33]

Thank you for your response. The previous example was using the latest release ("pip install openequivariance"), but when I tried using the latest commit (clone repository → installation), I got the same result. Additionally, when setting mode=torchscript, compilation succeeded, but setting target to either batch or pair_nequip resulted in MaxAbsError: nan.

[cw-tan]

Hi @mjrs33 , could you try playing around with different PyTorch versions? And potentially try with a different machine. The reason that we even have that compilation check is to catch potential PyTorch compilation errors that could be related to specific devices and PyTorch versions, build environments, etc. What kind of GPU are you running it on?

[cw-tan]

Coincidentally I just experienced the same compile check error on an A100 with PyTorch 2.9.1. Downgrading to PyTorch 2.9.0 seemed to work for me. However, I am unsure if it's also tied to the Triton version in the env.

[mjrs33]

Thank you for checking. When I tried it with torch==2.9.0, it compiled successfully. Using it as a calculator also worked without any issues. For now, I’ll run it with the downgraded version.