Merge branch 'main' into dependabot/pip/setuptools-lt-84.0.0

Author: frostedoyster

.github/workflows/integration-test.yml

@@ -12,7 +12,7 @@ jobs:
     strategy:
       max-parallel: 10
       matrix:
-        python_version: ['3.10', '3.12']
+        python_version: ['3.11', '3.14']
 
     steps:
       - uses: actions/checkout@v7

.github/workflows/stale.yml

@@ -10,7 +10,7 @@ jobs:
       issues: write
       pull-requests: write
     steps:
-      - uses: actions/stale@v10.4.0
+      - uses: actions/stale@v11.0.0
         with:
           days-before-stale: 30
           days-before-close: 14

.github/workflows/test.yml

@@ -32,7 +32,7 @@ jobs:
     # Without it, GitHub auto-expands EVERY field of the object-typed
     # `shard` matrix entry — including `shard.paths` — into a
     # >100-character truncated string like
-    # `test (3.13, core, tests/core/models tests/core/graph ...`.
+    # `test (3.14, core, tests/core/models tests/core/graph ...`.
     # That diverges from the short custom status emitted by
     # multi-trigger-setup, forcing branch protection to pick one event
     # type's names and break the others.
@@ -54,7 +54,7 @@ jobs:
       max-parallel: 10
       fail-fast: false
       matrix:
-        python_version: ['3.11', '3.13']
+        python_version: ['3.11', '3.14']
         # tests/perf is run by perf-test.yml on a separate cadence.
         # tests/applications and tests/lammps are run by integration-test.yml
         # and test_lammps_gpu respectively. Both shards below intentionally
@@ -112,7 +112,7 @@ jobs:
             -r tests/requirements.txt # pin test packages
 
       - name: Install torchsim (Python 3.12+)
-        if: ${{ matrix.python_version == '3.12' || matrix.python_version == '3.13' }}
+        if: ${{ matrix.python_version == '3.12' || matrix.python_version == '3.13' || matrix.python_version == '3.14' }}
         run: |
           pip install packages/fairchem-core[torchsim]
 
@@ -145,7 +145,7 @@ jobs:
         run: |
           pytest -m "serial and not gpu" ${{ matrix.shard.paths }} -vv --ignore=tests/demo/ocpapi/tests/integration/ --exclude-models=uma-s-1p1,uma-s-1p2 --cov-report=xml --cov=fairchem --cov-append --junitxml=junit-serial-${{ matrix.shard.name }}.xml -o junit_family=legacy -c ./packages/fairchem-core/pyproject.toml --suppress-no-test-exit-code
 
-      - if: ${{ matrix.python_version == '3.13' }}
+      - if: ${{ matrix.python_version == '3.14' }}
         name: codecov-coverage
         uses: codecov/codecov-action@v7
         with:
@@ -154,7 +154,7 @@ jobs:
           token: ${{ secrets.CODECOV_TOKEN }}
           verbose: true
 
-      - if: ${{ matrix.python_version == '3.13' }}
+      - if: ${{ matrix.python_version == '3.14' }}
         name: codecov-test-results
         uses: codecov/test-results-action@v1
         with:
@@ -180,7 +180,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python_version: ['3.13']
+        python_version: ['3.14']
         sweep_model: [uma-s-1p1, uma-s-1p2]
         shard: *cpu_shards
 
@@ -225,7 +225,7 @@ jobs:
             -r tests/requirements.txt # pin test packages
 
       - name: Install torchsim (Python 3.12+)
-        if: ${{ matrix.python_version == '3.12' || matrix.python_version == '3.13' }}
+        if: ${{ matrix.python_version == '3.12' || matrix.python_version == '3.13' || matrix.python_version == '3.14' }}
         run: |
           pip install packages/fairchem-core[torchsim]
 
@@ -285,7 +285,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python_version: ['3.13']
+        python_version: ['3.14']
         # YAML anchor: shard list is reused by test_gpu_sweep below.
         # Do not reorder jobs (the alias must follow the anchor textually).
         shard: &gpu_shards
@@ -327,7 +327,7 @@ jobs:
           -r tests/requirements.txt # pin test packages
 
       - name: Install torchsim (Python 3.12+)
-        if: ${{ matrix.python_version == '3.12' || matrix.python_version == '3.13' }}
+        if: ${{ matrix.python_version == '3.12' || matrix.python_version == '3.13' || matrix.python_version == '3.14' }}
         run: |
           pip install packages/fairchem-core[torchsim]
 
@@ -375,7 +375,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python_version: ['3.13']
+        python_version: ['3.14']
         sweep_model: [uma-s-1p1, uma-s-1p2]
         shard: *gpu_shards
 
@@ -408,7 +408,7 @@ jobs:
           -r tests/requirements.txt # pin test packages
 
       - name: Install torchsim (Python 3.12+)
-        if: ${{ matrix.python_version == '3.12' || matrix.python_version == '3.13' }}
+        if: ${{ matrix.python_version == '3.12' || matrix.python_version == '3.13' || matrix.python_version == '3.14' }}
         run: |
           pip install packages/fairchem-core[torchsim]
 
@@ -450,15 +450,15 @@ jobs:
     strategy:
       max-parallel: 1
       matrix:
-        python_version: ['3.13']
+        python_version: ['3.14']
 
     steps:
       - uses: actions/checkout@v7
       - uses: mamba-org/setup-micromamba@v3
         with:
           generate-run-shell: true # This is the default behavior
           environment-name: test_env # Specifies a new environment to be created
-          create-args: python=3.13 # Specifies packages to install into the new environment
+          create-args: python=3.14 # Specifies packages to install into the new environment
 
       - name: Checkout code
         uses: actions/checkout@v7

CLAUDE.md

@@ -242,6 +242,53 @@ configs/                 # Hydra YAML configs (datasets, tasks, backbone, optimi
 - `lmdb` - Dataset storage format
 - `ray[serve]>=2.53.0` - Distributed computing
 
+## Testing Gotchas
+
+- Tests that download registered checkpoints must declare their models with a
+  `pretrained` marker. This lets base CI deselect them with `--exclude-models`
+  and routes them to the matching model-sweep job.
+- Freeze inference parameters after inference-specific module replacement.
+  Main's folded-batch linear path removes the former general-backend regression:
+  on one H100, freezing improved compiled general inference by 15-17% and cut
+  peak allocated memory by 27-29% at 100-2,000 atoms. PyTorch 2.13 CPU checks
+  improved by 4% at 32 atoms and were neutral at 1,000 atoms. Custom backward
+  paths must preserve input derivatives independently of parameter gradients.
+- `umas_fast_gpu` custom backward operators do not implement `vmap` batching.
+  Compute Hessians through the per-component loop (`hessian_vmap=False`) when
+  exercising that backend, and ensure inference settings forward that option
+  into the backbone configuration.
+- Set `CI=true` when reproducing CPU CI shards locally. Some multi-GPU graph
+  parallel tests rely on that environment variable for skipping instead of the
+  `gpu` marker, so the CI marker expression alone will still collect them.
+- `graph_parallel_group_size=None` disables graph-parallel setup. A value of
+  `1` intentionally initializes singleton graph- and data-parallel groups and
+  is used to exercise those paths in tests; do not treat it as disabled.
+- Keep the full `AtomicData.clone()` boundary in prediction unless benchmarks
+  justify changing it and every model-side mutation has been audited. Graph
+  parallelism, MOLE preparation, and conservative gradients can replace or
+  mutate input fields, so a selective shallow copy is brittle.
+
+## Numerical Precision
+
+- Model constructors must not mutate process-wide PyTorch precision settings
+  such as `torch.set_float32_matmul_precision`. Precision is caller-owned;
+  inference applies TF32 temporarily through `InferenceSettings.tf32` and
+  restores the prior settings afterward.
+- TF32 policy belongs to the training/evaluation unit config or
+  `InferenceSettings.tf32`, never to a model config or model attribute.
+  Execution callers scope and restore the policy outside compiled `forward`
+  methods because precision getters cannot be traced by fullgraph.
+- Training and evaluation units default TF32 to disabled. Configs should set
+  `tf32` only when overriding that default. Hydra CLI overrides for configs
+  that omit the key must use the add syntax, such as
+  `+runner.train_eval_unit.tf32=true`.
+- Keep one configurable TF32 context manager for scoped matmul precision and
+  cuDNN state instead of introducing overlapping context managers.
+- Use the `tf32_context_manager` name for that policy; it controls both matmul
+  precision and cuDNN TF32, so `matmul_context` is too narrow.
+- Training FLOPs profiling invokes the model from `on_train_start`; scoped
+  execution settings must cover profiling as well as train/eval step methods.
+
 ## Cluster Validation Gotchas
 
 - H100 compute nodes do not have PyPI egress. Provision Python environments on

configs/allscaip/omol_4m.yml

@@ -106,6 +106,7 @@ runner:
   eval_dataloader: ${eval_dataloader}
   train_eval_unit:
     _target_: fairchem.core.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    tf32: True
     job_config: ${job}
     tasks: ${tasks}
     model:

configs/allscaip/omol_4m_ft.yml

@@ -111,6 +111,7 @@ runner:
   eval_dataloader: ${eval_dataloader}
   train_eval_unit:
     _target_: fairchem.core.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    tf32: True
     job_config: ${job}
     tasks: ${tasks}
     model:

configs/allscaip/omol_all.yml

@@ -106,6 +106,7 @@ runner:
   eval_dataloader: ${eval_dataloader}
   train_eval_unit:
     _target_: fairchem.core.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    tf32: True
     job_config: ${job}
     tasks: ${tasks}
     model:

configs/allscaip/omol_all_ft.yml

@@ -111,6 +111,7 @@ runner:
   eval_dataloader: ${eval_dataloader}
   train_eval_unit:
     _target_: fairchem.core.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    tf32: True
     job_config: ${job}
     tasks: ${tasks}
     model:

configs/escaip/training/mptrj_direct_escaip_fair.yml

@@ -114,6 +114,7 @@ runner:
   eval_dataloader: ${eval_dataloader}
   train_eval_unit:
     _target_: fairchem.core.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    tf32: True
     job_config: ${job}
     tasks: ${tasks}
     model:

configs/escaip/training/mptrj_finetune_escaip_fair.yml

@@ -115,6 +115,7 @@ runner:
   eval_dataloader: ${eval_dataloader}
   train_eval_unit:
     _target_: fairchem.core.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    tf32: True
     job_config: ${job}
     tasks: ${tasks}
     model: