Issue with compiling mode in ASE calculator for map0/mpb0 model in MACE v0.3.10

[yixian929]

** Description: I encountered an issue while trying to compile the map0/mpb0 model using MACE v0.3.10. The error seems to be related to recent updates in the newest version. Note that this issue does not occur in MACE v0.3.9.

**Environment:
torchaudio: 2.4.1 (py311_cu124, pytorch)
torchmetrics: 1.6.1 (pypi_0, pypi)
torchtriton: 3.0.0 (py311, pytorch)
torchvision: 0.19.1 (py311_cu124, pytorch)
CUDA Version: 12.5
NVIDIA Driver Version: 556.12
GPU: NVIDIA RTX 5000 Ada Generation
Installed CUDA Packages:

cuda-cudart: 12.4.127 (nvidia)
cuda-cupti: 12.4.127 (nvidia)
cuda-libraries: 12.4.1 (nvidia)
cuda-nvrtc: 12.4.127 (nvidia)
cuda-nvtx: 12.4.127 (nvidia)
cuda-opencl: 12.8.55 (nvidia)
cuda-runtime: 12.4.1 (nvidia)
cuda-version: 12.8 (nvidia)
cuequivariance: 0.1.0 (pypi_0, pypi)
cuequivariance-ops-torch-cu12: 0.1.0 (pypi_0, pypi)
cuequivariance-torch: 0.1.0 (pypi_0, pypi)

** Error Message:
Error(s) in loading state_dict for ScaleShiftMACE:
Unexpected key(s) in state_dict: "pair_repulsion_fn.r_max", "pair_repulsion_fn.cutoff.p", "pair_repulsion_fn.cutoff.r_max".

**Possible Cause: This issue seems to be related to the recent update in the newest version: "Clean up unused Polynomial Cutoff Class from ZBLBasis, remove r_max argument" by @CompRhys in #569.

**Steps to Reproduce:
Use the following script:
calculator = mace_mp(model='medium-mpa-0', default_dtype='float32', device='cuda', dispersion=False, compile_mode="default")
Run the script and observe the error.

**Additional Context:
MACE Version: 0.3.10
Date: Tue Feb 4 10:16:11 2025

[ilyes319]

@yixian929, thank you, I ll look at that. Just so you know, you will get much more speed up with cueq "enable_cueq=True" than with compilation atm.

[yixian929]

@ilyes319
Thank you for suggesting the use of cueq. I actually tried it before and encountered the same issue as in #740. Using cueq in ASE MD simulation only resulted in a 30% speedup for a medium model with a system of 1000 atoms. However, I achieved a 2x speedup with the compiled calculator. I will give it another try based on your suggestion. When I did fine-tuning with cueq in my environment, I did notice a significant speedup, so I think the issue might be due to the model/system size or something specific to ASE. I will give it another try based on your suggestion.

[ilyes319]

What GPU were you using for your cueq simulation?

[yixian929]

I have tested on NVIDIA RTX 5000 Ada, other enviroment details are described on the top