Merge branch 'main' into batching-enhancements

Author: lbluque

.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/sync-fork.yml

@@ -0,0 +1,33 @@
+name: sync fork
+
+on:
+  push:
+    branches:
+      - main
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: sync-fork-main
+  cancel-in-progress: true
+
+jobs:
+  push-to-private:
+    if: github.repository == facebookresearch/fairchem
+    runs-on: ubuntu-latest
+    environment: sync-private
+    steps:
+      - name: Checkout public repository
+        uses: actions/checkout@v7
+        with:
+          fetch-depth: 0
+          token: ${{ secrets.SYNC_PAT}}
+
+      - name: Push to private repository
+        run: |
+          git push \
+            https://github.com/fairinternal/fairchem.git \
+            HEAD:refs/heads/main \
+            --force

.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,32 @@ 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

docs/core/common_tasks/ase_calculator.md

@@ -60,11 +60,21 @@ os.environ['HF_TOKEN'] = 'MY_TOKEN'
 
 ## Default mode
 
-UMA is designed for both general-purpose usage (single or batched systems) and single-system long rollout (MD simulations, relaxations, etc.). For general-purpose use, we suggest using the [default settings](https://github.com/facebookresearch/fairchem/blob/main/src/fairchem/core/units/mlip_unit/api/inference.py#L92). This is a good trade-off between accuracy, speed, and memory consumption and should suffice for most applications. In this setting, on a single 80GB H100 GPU, we expect a user should be able to compute on systems as large as 50k-100k neighbors (depending on their atomic density). Batching is also supported in this mode.
+UMA defaults to the `merge_mole + compile` fast mode with TF32 disabled. This fast path requires fixed composition, task, charge, and spin across repeated evaluations. If a later evaluation changes any of these, the calculator prints a warning and permanently falls back to the unmerged, uncompiled model. Batching is supported; a mixed batch across any of the same parameters triggers the same fallback.
+
+## Batch mode
+
+Use batch mode for heterogeneous batches whose systems differ in composition, task, charge, or spin. It currently keeps MOLE unmerged and leaves compilation disabled. The named mode provides a stable entry point for future batch-specific optimizations, such as compilation without MOLE merging.
+
+```{code-cell} python3
+predictor = pretrained_mlip.get_predict_unit(
+    "uma-s-1p2", device="cuda", inference_settings="batch"
+)
+```
 
 ## Turbo mode
 
-For long rollout trajectory use-cases, such as molecular dynamics (MD) or relaxations, we provide a special mode called **turbo**, which optimizes for speed but restricts the user to using a single system where the atomic composition is held constant. Turbo mode is approximately 1.5-2x faster than default mode, depending on the situation. However, batching is not supported in this mode. It can be easily activated as shown below.
+Turbo mode uses the same `merge_mole + compile` fast path as default mode and additionally enables TF32. TF32 can improve performance on compatible hardware at a small precision trade-off. Similar to default mode, any changes in composition, task, charge, and spin across different evaluations trigger a fallback to the unoptimized execution path.
 
 ```{code-cell} python3
 predictor = pretrained_mlip.get_predict_unit(
@@ -74,7 +84,7 @@ predictor = pretrained_mlip.get_predict_unit(
 
 ## Custom modes for advanced users
 
-The advanced user might quickly see that **default** mode and **turbo** mode are special cases of our [inference settings api](https://github.com/facebookresearch/fairchem/blob/main/src/fairchem/core/units/mlip_unit/api/inference.py#L47). You can customize it for your application if you understand what you are doing. The following table provides more information.
+The advanced user might quickly see that **default**, **batch**, and **turbo** modes are special cases of our [inference settings api](https://github.com/facebookresearch/fairchem/blob/main/src/fairchem/core/units/mlip_unit/api/inference.py#L47). You can customize it for your application if you understand what you are doing. The following table provides more information.
 
 | Setting Flag  | Description |
 | ----- | ----- |

docs/core/common_tasks/workflows.md

@@ -25,7 +25,7 @@ Below is an example that uses the default `elastic_tensor_flow` flow:
 from __future__ import annotations
 
 from ase.build import bulk
-from quacc.recipes.mlp.elastic import elastic_tensor_flow
+from quacc.recipes.mlip.elastic import elastic_tensor_flow
 
 # Make an Atoms object of a bulk Cu structure
 atoms = bulk("Cu")
@@ -35,7 +35,7 @@ result = elastic_tensor_flow(
     atoms,
     job_params={
         "all": dict(
-            method="fairchem",
+            library="fairchem",
             name_or_path="uma-s-1p2",
             task_name="omat",
         ),

packages/fairchem-core/pyproject.toml

@@ -7,7 +7,7 @@ name = "fairchem-core"
 description = "Machine learning models for chemistry and materials science by the FAIR Chemistry team"
 license = {text = "MIT License"}
 dynamic = ["version", "readme"]
-requires-python = ">=3.11, <3.14"
+requires-python = ">=3.11, <3.15"
 dependencies = [
     "torch~=2.13.0",
     "ray[serve]>=2.53.0",

packages/fairchem-lammps/pyproject.toml

@@ -7,7 +7,7 @@ name = "fairchem-lammps"
 description = "Fairchem extension package to use LAMMPs with fairchem models"
 license = {text = "GnuV2 License"}
 dynamic = ["version", "readme"]
-requires-python = ">=3.10,<3.14"
+requires-python = ">=3.10,<3.15"
 dependencies = [
   "fairchem.core>=2.4.0",
 ]

src/fairchem/core/calculate/ase_calculator.py

@@ -134,9 +134,11 @@ def from_model_checkpoint(
             name_or_path: A model name from fairchem.core.pretrained.available_models or a path to the checkpoint
                 file
             task_name: Task name
-            inference_settings: Settings for inference. Can be "default" (general purpose) or "turbo"
-                (optimized for speed but requires fixed atomic composition). Advanced use cases can
-                use a custom InferenceSettings object.
+            inference_settings: Settings for inference. Both "default" and "turbo" use the
+                merge_mole + compile fast path, with automatic fallback if its fixed-input
+                contract is broken. "turbo" additionally enables TF32. "batch" keeps MOLE
+                unmerged for heterogeneous inputs. More advanced use cases can use a custom
+                InferenceSettings object.
             overrides: Optional dictionary of settings to override default inference settings.
             device: Optional torch device to load the model onto.
             seed: Random seed for reproducibility.