Adding `NequIP`/`Allegro` Foundation Potentials

[kavanase]

Hi!
Thank you for putting together such a nice and useful package.
We were hoping to add our NequIP and Allegro models, for which we have now trained foundation potentials for (see nequip.net; https://matbench-discovery.materialsproject.org/).
Is there an established process for this? Do you need anything from our end?
Thanks!

[Lingyu-Kong]

Hello!

We’ve also noticed the release of the NequIP and Allegro foundation models. We’re very impressed by the accuracy and speed of both models, and we’ve started exploring how to elegantly integrate them into MatterTune without losing their performance advantages. Since we found that the NequIP repository is organized quite differently from the foundation model codebase that we are familiar with, we’ll need some time to analyze and adapt it. If you could help answer the following questions, it would greatly accelerate our development progress:

1. In MatterTune, model fine-tuning requires wrapping the backbone model into a Lightning LightningModule and training it using the Lightning Trainer. Could you please point us to the code locations in NequIP that implement the following functionalities?
   a. Loading a pretrained checkpoint
   b. The backbone model’s nn.Module or LightningModule class (where the full model forward is implemented)
   c. How to convert input ASE Atoms structures into model-acceptable inputs

2. We noticed that the NequIP repository also contains code for training models with Lightning (see nequip/train/lightning.py
   ). Is this part of the code still functional and recommended for reference?

3. We also observed that NequIP relies on nequip_compile for compilation to accelerate the model. We’re concerned that wrapping the model into a LightningModule for integration into MatterTune might break its compilability. Are there any details we should pay special attention to when doing this?

Again, thank you very much for your advice and help!

[kavanase]

Great!

I will defer to @cw-tan as the expert here, but try to answer these questions:

In MatterTune, model fine-tuning requires wrapping the backbone model into a Lightning LightningModule and training it using the Lightning Trainer. Could you please point us to the code locations in NequIP that implement the following functionalities?
a. Loading a pretrained checkpoint
b. The backbone model’s nn.Module or LightningModule class (where the full model forward is implemented)
c. How to convert input ASE Atoms structures into model-acceptable inputs

a. ModelFromPackage (https://github.com/mir-group/nequip/blob/main/nequip/model/saved_models/package.py#L70) to load nequip.zip package files (preferred for greater forwards & backwards compatibility, these are the pre-trained model files we provide at https://nequip.net), or alternatively ModelFromCheckpoint (https://github.com/mir-group/nequip/blob/main/nequip/model/saved_models/checkpoint.py#L26)
b. I think these are the relevant modules: https://github.com/mir-group/nequip/blob/main/nequip/train/lightning.py & https://github.com/mir-group/nequip/blob/main/nequip/model/nequip_models.py
c. from_ase; https://github.com/mir-group/nequip/blob/main/nequip/data/ase.py#L18 – example usage shown in some other parts of the code.

We noticed that the NequIP repository also contains code for training models with Lightning (see nequip/train/lightning.py
). Is this part of the code still functional and recommended for reference?

Yes NequIPLightningModule in that module is our core training module that we currently use and recommend.

We also observed that NequIP relies on nequip_compile for compilation to accelerate the model. We’re concerned that wrapping the model into a LightningModule for integration into MatterTune might break its compilability. Are there any details we should pay special attention to when doing this?

I think it will depend if you mean for training or for inference here. nequip-compile is not used for training, but we do usually use torch.compile to accelerate training (controlled by compile_mode: compile in NequIP's configs).
nequip-compile is used for AOTI (recommended, alternatively Torchscript) compilation of the models for significantly accelerated inference. We also support model modifiers for accelerated inference, namely accelerated tensor product kernels OpenEquivariance and CuEquivariance, with some details on the docs here: https://nequip.readthedocs.io/en/latest/guide/getting-started/workflow.html#compilation & https://nequip.readthedocs.io/en/latest/guide/accelerations/gpu_kernel_modifiers.html

[cw-tan]

What @kavanase said sounds about right. Let us know if you want any more information. Happy to provide support and advice. Please consider joining the NequIP Zulip, which might be more convenient for real time discussions and exchanges: https://docs.google.com/forms/d/e/1FAIpQLSfTKAjlpK-u1fqefMC2fl9rah1ICIZ_s58sj1I0dBs1_DskMQ/viewform

[Lingyu-Kong]

Hello @kavanase @cw-tan!

We’ve just integrated the NequIP/Allegro atomistic foundation models into the MatterTune. Our users can now fine-tune and train NequIP and Allegro foundation models on MatterTune. In addition, since the MatterTune platform refactors the model implementations to satisfy its modular design principles, models trained on MatterTune cannot be directly packaged using the nequip-package API provided in the original NequIP repository. Therefore, we have implemented an equivalent packaging function, nequip_model_package()link, following the design of the nequip-package API. NequIP/Allegro models trained with MatterTune, once packaged using this function, are fully compatible with the other NequIP APIs, including nequip-compile.

However, we have noticed a few issues with our current solution:

1. In the original nequip-package API implementation, the .nequip.zip file stores the model, example_data, pkg_metadata, and the training config.yaml. In our nequip_model_package function, model, example_data, and pkg_metadata are saved in the same way, but for config.yaml we currently substitute an empty dummy_config. We have observed that ```config.yaml`` does not appear to be used during the subsequent compile step, but we are not sure whether this implementation could cause any problems.

2. Our MatterTune platform also implements functionality similar to Per-Type Energy Shifts and Scales, but this is implemented on the MatterTune side rather than inside NequIP’s GraphModel. As a result, this functionality is effectively lost during model packaging and compiling. Although we can still obtain reasonably good fine-tuning results without using it, we would like to address this issue. We have examined NequIP’s implementation but have not been able to clearly identify where the Per-Type Energy Shifts and Scales module is created in a NequIP fine-tuning run, how it is integrated into the model, and how it participates in packaging and compiling. Could you please help clarify this?

3. We performed a simple speed comparison of our fine-tuned models before and after packaging and nequip-compile. We benchmarked ASE MD performance using the compile options --mode torchscript --device cuda --target ase on a single A40 GPU. We found that the speed difference before and after compilation is quite small (the compiled model is only about 10–20% faster), and when switching to aotinductor the speed actually becomes significantly slower. Is this behavior expected?

[cw-tan]

Hi @Lingyu-Kong , fantastic, thank you!

1. The config file is just an artefact for reference in case one seeks to reproduce the packaged model. \
2. The packaged file format contains the model that already has its own PerTypeScaleShift module (https://nequip.readthedocs.io/en/latest/api/nn.html#nequip.nn.PerTypeScaleShift). For shifting and scaling, we provide a model modifier to do so: https://nequip.readthedocs.io/en/latest/guide/training-techniques/fine_tuning.html#modifying-per-type-energy-shifts-and-scales. The idea is for the user to apply the model modifier as in our fine-tuning docs. During packaging and/or compiling, the checkpoint loading mechanism (https://github.com/mir-group/nequip/blob/b076b902fa1f4d3ca93c9dbeb24566d12b993532/nequip/model/saved_models/load_utils.py#L86) will construct the model in the same way, i.e. loading model from package, applying the model modifier, then loading in the checkpoint weights.
3. That's interesting. We only have benchmarks on A100, H100, H200s. Anyway, for NequIP models, using OpenEquivariance kernels should provide a substantial speedup: https://nequip.readthedocs.io/en/latest/guide/accelerations/openequivariance.html

[Lingyu-Kong]

Hi @cw-tan

Thank you for the clarification — this is very helpful. I agree that the issue is most likely related to the GPU and CPU architectures. We will further investigate this by running more systematic performance benchmarks and multi-GPU parallel tests on our A100 cluster. In addition, we will try the modifier-based approach you suggested to integrate PerTypeScaleShift into the model.

[kavanase]

Hi @Lingyu-Kong, thank you for working on this!
Just to add in response to this point:

We found that the speed difference before and after compilation is quite small (the compiled model is only about 10–20% faster), and when switching to aotinductor the speed actually becomes significantly slower. Is this behavior expected?

The AOTI slowdown is quite unexpected (assuming there is no difference in the use of OpenEquivariance – which I assume is not being used yet here). We generally see it giving a 1.5-3x speedup for our NequIP models