Add unit test to show cpu compile error from foundation models

.github/workflows/tests.yml

@@ -18,6 +18,19 @@ concurrency:
 
 # yamllint disable rule:line-length
 jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Install pre-commit
+        run: pip install pre-commit
+      - name: Run pre-commit
+        run: pre-commit run --all-files --show-diff-on-failure
+
   run-tests:
     runs-on: ${{ github.event.repository.private && 'self-hosted' || 'ubuntu-latest' }}
     strategy:
@@ -32,12 +45,6 @@ jobs:
         run: |
           apt update && apt install -y python3 python3-pip git
           python3 -m pip install --upgrade pip
-      - name: Linting
-        run: |
-          pip install ruff yamllint
-          ruff check .
-          ruff format --check .
-          yamllint .
       - name: Install PyTorch
         env:
           TORCH: "${{ matrix.torch-version }}"

.pre-commit-config.yaml

@@ -1,15 +1,30 @@
 exclude: '.git|.tox'
-default_stages: [commit]
+default_stages: [pre-commit]
 fail_fast: true
 
+ci:
+  autoupdate_schedule: quarterly
+
 repos:
-  - repo: https://github.com/psf/black
-    rev: 24.4.2
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.12.2
+    hooks:
+      - id: ruff
+        args: [--fix]
+      - id: ruff-format
+
+  - repo: https://github.com/adrienverge/yamllint
+    rev: v1.37.1
     hooks:
-      - id: black
+      - id: yamllint
 
-  - repo: https://github.com/pycqa/flake8
-    rev: 7.1.0
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
     hooks:
-      - id: flake8
-        additional_dependencies: [Flake8-pyproject]
+      - id: check-case-conflict
+      - id: check-symlinks
+      - id: check-yaml
+      - id: destroyed-symlinks
+      - id: end-of-file-fixer
+      - id: mixed-line-ending
+      - id: trailing-whitespace

AUTHORS.md

@@ -25,4 +25,4 @@ Present and past contributors include, in alphabetical order by last name:
  - Menghang (David) Wang
  - William C. Witt
  - Albert Zhu
- - Laura Zichi
+ - Laura Zichi

docs/_static/custom.css

@@ -17,4 +17,4 @@
 /* Justify text in list items that contain multiple sentences */
 li p {
     text-align: justify;
-}
+}

docs/api/data_stats.rst

@@ -23,7 +23,7 @@ As an example, we show how one can configure the full :class:`~nequip.data.DataS
        - name: num_neighbors_mean
          field:
            _target_: nequip.data.NumNeighbors
-         metric: 
+         metric:
            _target_: nequip.data.Mean
        - name: per_atom_energy_mean
          field:
@@ -55,12 +55,12 @@ As an example, we show how one can configure the full :class:`~nequip.data.DataS
 
 .. autoclass:: nequip.data.StandardDeviation
     :members:
-    
+
 .. autoclass:: nequip.data.Min
     :members:
-    
+
 .. autoclass:: nequip.data.Max
     :members:
 
 .. autoclass:: nequip.data.Count
-    :members:
+    :members:

docs/api/data_transforms.rst

@@ -13,4 +13,4 @@ Data transforms convert the raw data from the :class:`~nequip.data.dataset.Atomi
     :members:
 
 .. autoclass:: nequip.data.transforms.StressSignFlipTransform
-    :members:
+    :members:

docs/api/datamodule.rst

@@ -32,4 +32,4 @@ For usage examples and configuration guidance, see the :doc:`../guide/configurat
      :members:
 
 .. autoclass:: nequip.data.datamodule.WaterDataModule
-     :members:
+     :members:

docs/api/lightning_module.rst

@@ -14,4 +14,4 @@ NequIPLightningModule
     :members:
 
 .. autoclass:: nequip.train.ScheduleFreeLightningModule
-    :members:
+    :members:

docs/api/model.rst

@@ -5,5 +5,3 @@ nequip.model
 
   nequip_model
   save_model
-
-

docs/api/nequip_model.rst

@@ -3,4 +3,4 @@ NequIP Message Passing GNN Model
 
 .. autofunction:: nequip.model.NequIPGNNModel
 
-.. autofunction:: nequip.model.NequIPGNNEnergyModel
+.. autofunction:: nequip.model.NequIPGNNEnergyModel

docs/api/nn.rst

@@ -6,4 +6,4 @@ nequip.nn
     :members:
 
 .. autoclass:: nequip.nn.PerTypeScaleShift
-    :members:
+    :members:

docs/api/save_model.rst

@@ -18,8 +18,8 @@ These files can be loaded using the following :class:`~nequip.model.ModelFromChe
   - Use absolute paths instead of relative paths.
   - Do not change the directory structure or move your files when using the model loaders.
   - Ideally, store the original checkpoint/package files somewhere that makes their association with the new training run clear to you.
-    
-  Be aware that iterated nested use of :class:`~nequip.model.ModelFromCheckpoint` will result in a checkpoint chaining phenomenon where loading the checkpoint at the end of the chain requires successfully loading every intermediate checkpoint file in the chain. One can break this chain if necessary by using ``nequip-package`` to convert the checkpoint file into a packaged model, and then using :class:`~nequip.model.ModelFromPackage`. 
+
+  Be aware that iterated nested use of :class:`~nequip.model.ModelFromCheckpoint` will result in a checkpoint chaining phenomenon where loading the checkpoint at the end of the chain requires successfully loading every intermediate checkpoint file in the chain. One can break this chain if necessary by using ``nequip-package`` to convert the checkpoint file into a packaged model, and then using :class:`~nequip.model.ModelFromPackage`.
 
 .. autofunction:: nequip.model.ModelFromCheckpoint
 
@@ -36,4 +36,4 @@ This is typically used by the ASE calculator or other inference applications.
 Modifying Saved Models
 ######################
 
-.. autofunction:: nequip.model.modify
+.. autofunction:: nequip.model.modify

docs/api/train.rst

@@ -7,4 +7,4 @@ nequip.train
   lightning_module
   metrics
   callbacks
-  ddp
+  ddp

docs/dev/contributing.md

@@ -1,7 +1,7 @@
 # Contributing to NequIP
 
 ```{note}
-If you want to make a major change, add a new feature, or have any uncertainty about where the change should go or how it should be designed, please reach out to discuss with us first. Details on contacting the development team can be found in `README.md`. 
+If you want to make a major change, add a new feature, or have any uncertainty about where the change should go or how it should be designed, please reach out to discuss with us first. Details on contacting the development team can be found in `README.md`.
 ```
 
 ## Design Philosophy
@@ -60,7 +60,7 @@ ruff format --check .
 ```
 
 The formatter can be easily set up to run automatically on file save in various editors.
-  
+
 You can also use ``pre-commit install`` to install a [pre-commit](https://pre-commit.com/) hook.
 
 ```{tip}
@@ -73,4 +73,4 @@ PRs should generally come from feature branches.
 
 All PRs should be onto the `develop` branch or other feature branches, and **NOT** `main` (unless they fix a significant bug with immediate impact).
 
-On feature branches, it is preferable to [rebase](https://docs.github.com/en/get-started/using-git/about-git-rebase) wherever possible. We strive towards a clean and easily-readable commit history. Note that this does not apply to `develop` or `main` -- the commit history on these core branches are sacred.
+On feature branches, it is preferable to [rebase](https://docs.github.com/en/get-started/using-git/about-git-rebase) wherever possible. We strive towards a clean and easily-readable commit history. Note that this does not apply to `develop` or `main` -- the commit history on these core branches are sacred.

docs/dev/dev.rst

@@ -6,4 +6,4 @@ Developer Guide
    contributing
    understanding_nequip
    extension_packages
-   ../README
+   ../README

docs/dev/extension_packages.rst

@@ -22,4 +22,4 @@ Example Extension Packages
 Feel free to use the following extension packages for inspiration and reference:
 
 - **Allegro** (`GitHub <https://github.com/mir-group/allegro>`__, `Docs <https://nequip.readthedocs.io/projects/allegro/en/latest/?badge=latest>`__, `Paper <https://www.nature.com/articles/s41467-023-36329-y>`__): Strictly local equivariant models with excellent scalability for multirank molecular dynamics simulations.
-- **NequIP-LES** (`Github <https://github.com/ChengUCB/NequIP-LES>`__, `Paper <https://arxiv.org/abs/2507.14302>`__): An extension of NequIP and Allegro that adds long-range electrostatics via the Latent Ewald Summation (LES) algorithm.
+- **NequIP-LES** (`Github <https://github.com/ChengUCB/NequIP-LES>`__, `Paper <https://arxiv.org/abs/2507.14302>`__): An extension of NequIP and Allegro that adds long-range electrostatics via the Latent Ewald Summation (LES) algorithm.

docs/dev/extension_packages/data.rst

@@ -12,7 +12,7 @@ Field Types
 ~~~~~~~~~~~
 
 - **Graph fields**: per-frame (e.g., ``total_energy``, ``stress``)
-- **Node fields**: per-atom (e.g., ``forces``, ``positions``) 
+- **Node fields**: per-atom (e.g., ``forces``, ``positions``)
 - **Edge fields**: per-edge (e.g., ``edge_vectors``)
 - **Long fields**: integer dtype (e.g., ``atomic_numbers``)
 - **Cartesian tensors**: physical tensors (e.g., ``stress``)
@@ -26,7 +26,7 @@ Registration Example
 
    register_fields(
        graph_fields=["custom_energy"],
-       node_fields=["custom_forces"], 
+       node_fields=["custom_forces"],
        edge_fields=["custom_edge_attr"],
        long_fields=["custom_indices"],
        cartesian_tensor_fields={"custom_tensor": "ij=ji"}
@@ -66,4 +66,4 @@ Data Transforms
 
 Extension packages can implement custom data transforms to preprocess data during loading. Transforms are classes that implement a ``__call__`` method to modify ``AtomicDataDict`` objects.
 
-See the :doc:`transforms API documentation <../../api/data_transforms>` for available transform classes and their patterns.
+See the :doc:`transforms API documentation <../../api/data_transforms>` for available transform classes and their patterns.

docs/dev/extension_packages/getting_started.rst

@@ -19,4 +19,4 @@ When the NequIP framework is imported, it automatically discovers and loads all
 - Register custom OmegaConf resolvers
 - Perform any other initialization needed for your package
 
-Your package's ``__init__.py`` will be automatically imported whenever the NequIP framework is used, ensuring your extensions are always available.
+Your package's ``__init__.py`` will be automatically imported whenever the NequIP framework is used, ensuring your extensions are always available.

docs/dev/extension_packages/models.rst

@@ -53,7 +53,7 @@ Train-time via config:
            H: 0.12
      model:
        _target_: nequip.model.ModelFromPackage
-       # ... 
+       # ...
 
 Compile-time for accelerations:
 
@@ -64,4 +64,4 @@ Compile-time for accelerations:
 Implementation Tips
 -------------------
 
-When implementing custom :class:`torch.nn.Module` subclasses, use :meth:`torch.nn.Module.extra_repr` to provide crucial model information for debugging (visible when ``_NEQUIP_LOG_LEVEL=DEBUG`` is set).
+When implementing custom :class:`torch.nn.Module` subclasses, use :meth:`torch.nn.Module.extra_repr` to provide crucial model information for debugging (visible when ``_NEQUIP_LOG_LEVEL=DEBUG`` is set).

docs/dev/extension_packages/training.rst

@@ -22,5 +22,5 @@ Examples include :class:`~nequip.train.EMALightningModule` (exponential moving a
 When composing multiple LightningModule subclasses, follow these rules to avoid the "deadly diamond of death" multiple inheritance pattern:
 
 - Use unique, distinguishable names for class attributes to prevent overwrites during inheritance
-- Be very careful of inheritance order when creating "diamond" subclasses 
-- Use assertions to verify the composed module behaves as intended
+- Be very careful of inheritance order when creating "diamond" subclasses
+- Use assertions to verify the composed module behaves as intended

docs/dev/understanding_nequip.rst

@@ -6,4 +6,4 @@ This section provides comprehensive documentation for understanding NequIP's int
 .. toctree::
    :maxdepth: 2
 
-   understanding_nequip/workflows/index
+   understanding_nequip/workflows/index

docs/dev/understanding_nequip/workflows/index.rst

@@ -8,4 +8,4 @@ This section covers key workflows and processes within the NequIP framework.
 .. toctree::
    :maxdepth: 2
 
-   packaging
+   packaging

docs/guide/accelerations/cuequivariance.md

@@ -83,4 +83,4 @@ nequip-prepare-lmp-mliap \
 
 Where `model_file` can be either a [checkpoint file](../getting-started/files.md#checkpoint-files) (`.ckpt`) or [package file](../getting-started/files.md#package-files).
 The resulting `.nequip.lmp.pt` file can be used directly in LAMMPS scripts with the `pair_style mliap` command.
-See the [ML-IAP documentation](../../integrations/lammps/mliap.md) for complete usage instructions and examples.
+See the [ML-IAP documentation](../../integrations/lammps/mliap.md) for complete usage instructions and examples.

docs/guide/accelerations/ddp_training.md

@@ -6,7 +6,7 @@ There are two ways to set up multi-rank distributed data parallel (DDP) training
 
 If train-time compilation is not used, you can use PyTorch Lightning's {class}`~lightning.pytorch.strategies.DDPStrategy`.
 
-See [Lightning's docs](https://lightning.ai/docs/pytorch/stable/accelerators/gpu_intermediate.html#distributed-data-parallel) for how to set it up through Lightning's {class}`~lightning.pytorch.trainer.trainer.Trainer`. 
+See [Lightning's docs](https://lightning.ai/docs/pytorch/stable/accelerators/gpu_intermediate.html#distributed-data-parallel) for how to set it up through Lightning's {class}`~lightning.pytorch.trainer.trainer.Trainer`.
 
 **NOTE** that it is usually not necessary to explicitly set the {class}`~lightning.pytorch.strategies.DDPStrategy` as an input to the Lightning {class}`~lightning.pytorch.trainer.trainer.Trainer` if the cluster environment is set up to facilitate Lightning's automatic detection of cluster variables. The main aspect that deserves user attention is configuring the job submission script and the relevant `Trainer` arguments (`num_nodes` and sometimes `devices`) correctly.
 
@@ -28,7 +28,7 @@ A minimal SLURM example for doing DDP training with 2 nodes with 4 GPUs per node
 
 # ... set up (e.g. module load, activate Python env, etc)
 
-# ... cluster specific set up such as network interface 
+# ... cluster specific set up such as network interface
 # (e.g. MASTER_PORT, MASTER_ADDR, NCCL_SOCKET_IFNAME)
 
 srun nequip-train -cn config.yaml ++trainer.num_nodes=${SLURM_NNODES}
@@ -62,7 +62,7 @@ The `batch_size` configured under the dataloaders in the `data` [section of the
 
 As increasing the number of ranks (while holding the per-rank batch size constant) increases the effective batch size, you should consider adjusting other hyperparameters that you would typically adjust when raising the batch size, such as the learning rate (see [Lightning's docs](https://lightning.ai/docs/pytorch/stable/accelerators/gpu_faq.html#how-should-i-adjust-the-learning-rate-when-using-multiple-devices) for similar advice).
 
-It may be helpful to use a combination of {mod}`omegaconf`'s [variable interpolation](https://omegaconf.readthedocs.io/en/latest/usage.html#variable-interpolation), [environment variable resolver](https://omegaconf.readthedocs.io/en/latest/custom_resolvers.html#oc-env) and NequIP's custom arithmetic resolver `int_div` to dynamically configure these parameters based on the runtime environment. 
+It may be helpful to use a combination of {mod}`omegaconf`'s [variable interpolation](https://omegaconf.readthedocs.io/en/latest/usage.html#variable-interpolation), [environment variable resolver](https://omegaconf.readthedocs.io/en/latest/custom_resolvers.html#oc-env) and NequIP's custom arithmetic resolver `int_div` to dynamically configure these parameters based on the runtime environment.
 
 For example, to get the world size as a SLURM environment variable and set the per-rank batch size as the desired effective global batch size divided by the world size, you can use:
 
@@ -78,4 +78,4 @@ When using DDP, the {class}`torch.utils.data.distributed.DistributedSampler` may
 
 Be very careful when reporting validation or test metrics in DDP settings, as data duplication can lead to incorrect metrics. Either ensure data samples can be evenly distributed to all ranks, or perform validation/testing on a single rank.
 
-For more details on handling validation in distributed settings, see [Lightning's documentation](https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#test-loop).
+For more details on handling validation in distributed settings, see [Lightning's documentation](https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#test-loop).

docs/guide/accelerations/gpu_kernel_modifiers.md

@@ -18,4 +18,4 @@ NequIP GNN models support several GPU kernel modifiers that can significantly sp
 
 openequivariance
 cuequivariance
-```
+```

docs/guide/accelerations/precision.md

@@ -38,7 +38,7 @@ callbacks:
 callbacks:
   - _target_: nequip.train.callbacks.TF32Scheduler
     schedule:
-      0: true      # Enable TF32 for faster early training  
+      0: true      # Enable TF32 for faster early training
       100: false   # Disable TF32 at epoch 100 for precise convergence
 ```
 
@@ -58,4 +58,4 @@ nequip-compile model.ckpt compiled_model.pt --tf32 ...
 nequip-compile model.ckpt compiled_model.pt --no-tf32 ...
 ```
 
-The default behavior is to compile without TF32, regardless of training settings.
+The default behavior is to compile without TF32, regardless of training settings.

docs/guide/accelerations/pt2_compilation.md

@@ -13,7 +13,7 @@ or
 model:
     _target_: allegro.model.AllegroModel
     compile_mode: compile
-    # other model hyperparameters    
+    # other model hyperparameters
 ```
 Note that `compile_mode` can only be `compile` (use [`torch.compile`](https://pytorch.org/docs/stable/generated/torch.compile.html)), or `eager` (no compilation used). If `compile_mode` is unspecified, it defaults to `eager`.
 It will take a bit of time (around a minute or more) for the model to be compiled with [`torch.compile`](https://pytorch.org/docs/stable/generated/torch.compile.html) before training proceeds, but the speed-ups are worth it.
@@ -26,4 +26,4 @@ Train-time compilation will not work if any of the batch dimensions are all neve
 
 ```{warning}
 At present we advise against using train-time compilation on CPUs. As of PyTorch 2.6.0, there are known cases of **CPU**-specific train-time compilation issues for certain configurations of NequIP and Allegro models. Be cautious when trying to use train-time compilation with CPUs. If you encounter such issues, please open a GitHub issue.
-```
+```

docs/guide/configuration/config.md

@@ -70,7 +70,7 @@ The full set of options are found in the documentation of the [underlying object
 `training_module` defines the {class}`~nequip.train.NequIPLightningModule` (or its subclasses). Users are directed to the [`nequip.train.NequIPLightningModule` API documentation](../../api/lightning_module.rst) to learn how to configure it. Usually the {class}`~nequip.train.EMALightningModule` is the right choice.
 
 The following important objects are configured as part of the `training_module`:
- 
+
  ### `model`
   This section configures the model itself, including hyperparameters and the choice of architecture (for example, the NequIP message-passing E(3)-equivariant GNN, or the Allegro architecture). Refer to the [model documentation page](../../api/model.rst) to learn how to configure this section.
 
@@ -79,7 +79,7 @@ The following important objects are configured as part of the `training_module`:
 
  ### `optimizer` and `lr_scheduler`
 
-  The `optimizer` can be any PyTorch-compatible optimizer. Options from PyTorch can be found in {mod}`torch.optim`. The {class}`~torch.optim.Adam` optimizer, for example, can be configured as follows: 
+  The `optimizer` can be any PyTorch-compatible optimizer. Options from PyTorch can be found in {mod}`torch.optim`. The {class}`~torch.optim.Adam` optimizer, for example, can be configured as follows:
 ```yaml
 optimizer:
   _target_: torch.optim.Adam
@@ -99,4 +99,3 @@ lr_scheduler:
   frequency: 1
 ```
   The `scheduler` is a PyTorch-compatible learning rate scheduler. Options from PyTorch can be found in {mod}`torch.optim.lr_scheduler`.
-