Uma base inference speedups (#2111)

Speed up UMA inference and DDP training without changing numerics

(updated Aug 3rd)

# perf(inference): freeze prepared model parameters
## Summary
Freeze model parameters after inference-specific module preparation.
Inference
may differentiate outputs with respect to positions and cells for
forces,
stress, and Hessians, but it does not need parameter gradients. Freezing
after
MOLE and backend replacement avoids retaining unnecessary
weight-gradient
state while preserving input derivatives.
The policy now covers both `general` and `umas_fast_gpu`. Main's
folded-batch
linear path (`9547a5b3c`) removes the former general-backend
frozen-weight
regression. Prediction retains the full `AtomicData.clone()` boundary
because
graph parallelism, MOLE preparation, and conservative gradients can
mutate
model inputs.
This PR also adds frozen/unfrozen energy, force, stress, and Hessian
parity
coverage for both execution modes, singleton graph-parallel regression
coverage, and pretrained-model test routing markers.

## Performance

Compared current `origin/main` at `9547a5b3c` with the
parameter-freezing
change. Both sides include the new fused Wigner/SO2 edgewise kernels and
folded-batch linear operations.

All GPU runs used UMA-S-1.2.1 (`uma-s-1p2p1`), PyTorch 2.13.0+cu130,
H100
80 GB HBM3 GPUs, and the `h100_ocp_high` QoS.

### Fast GPU backend

Each atom count is one global FCC structure partitioned across a
graph-parallel group containing every allocated GPU; atom counts are not
per
GPU. The primary throughput metric is median synchronized QPS. Peak
allocated
memory is the maximum rank value; total allocated memory is summed
across
ranks.

| GPUs | Atoms | Main QPS | PR QPS | Speedup | Allocated GiB/GPU, main
-> PR | Allocated delta | Reserved delta | Total allocated GiB, main ->
PR |
| ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| 8 | 100 | 30.60 | 32.43 | 1.060x | 0.18 -> 0.17 | -8.9% | -3.0% | 1.44
-> 1.31 |
| 8 | 1,000 | 29.53 | 32.05 | 1.086x | 0.88 -> 0.70 | -20.6% | -20.3% |
7.06 -> 5.60 |
| 8 | 10,000 | 21.08 | 23.08 | 1.095x | 7.32 -> 5.64 | -22.9% | -21.0% |
58.48 -> 45.08 |
| 32 | 100 | 26.86 | 28.42 | 1.058x | 0.13 -> 0.13 | -2.9% | -5.2% |
4.10 -> 4.00 |
| 32 | 1,000 | 26.17 | 27.58 | 1.054x | 0.33 -> 0.27 | -17.2% | -13.0% |
10.41 -> 8.63 |
| 32 | 10,000 | 25.97 | 27.05 | 1.042x | 2.15 -> 1.63 | -24.0% | -21.7%
| 68.19 -> 51.97 |

Freezing remains beneficial after the fused fast-backend kernels: median
QPS
improves in all six cases by 4.2-9.5%, while the 10,000-atom cases save
23-24%
peak allocated memory and 21-22% peak reserved memory.

The fast-backend benchmark used candidate snapshot `a2dca9379`. Its fast
path
is behaviorally identical to final commit `3af44001e`; the final commit
only
extends the same freeze operation to other execution modes and broadens
tests.

### General backend

The general-backend check used one H100 without graph-parallel
initialization.
It applies freezing after MOLE preparation and before compilation,
exactly where
the final implementation applies it.

| Atoms | Main QPS | Frozen QPS | Speedup | Allocated GiB, main ->
frozen | Allocated delta | Reserved delta |
| ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| 100 | 47.53 | 55.52 | 1.168x | 1.04 -> 0.76 | -27.0% | -23.4% |
| 1,000 | 16.98 | 19.55 | 1.151x | 11.80 -> 8.36 | -29.1% | -25.0% |
| 2,000 | 10.24 | 11.76 | 1.149x | 21.40 -> 15.16 | -29.2% | -25.1% |

The folded-batch linear change reverses the former general-backend
regression:
freezing is now 14.9-16.8% faster and saves 27-29% allocated memory.

A local 16-thread CPU check with the general backend and compilation
disabled
measured `1.040x` at 32 atoms and `1.002x` at 1,000 atoms. The new path
is
positive for the small case and neutral for the larger case, rather than
the
previous CPU slowdown.

## Methodology

- Fast settings: FP32 base precision, TF32 enabled, `umas_fast_gpu`,
`merge_mole=True`, `compile=True`, and `activation_checkpointing=False`.
- General settings: the same settings with `execution_mode=general`; the
  one-GPU check did not initialize graph parallelism.
- Order on each GPU allocation: main A, candidate A, candidate B, main B
  (ABBA).
- Each GPU branch/size result: 5 warmups followed by 3 repeats of 10
  predictions.
- Reported QPS: median of 6 synchronized repeats across the two runs.
- Multi-GPU timing: maximum elapsed time across ranks after CUDA
  synchronization and a distributed barrier.
- Memory: CUDA peaks reset after warmup; maximum-rank and aggregate
allocated
  and reserved memory were recorded.
- Inputs: the same seed-42 FCC structures were used by every branch and
GPU
  count; all structure checksums match exactly.

## Fidelity

Energy, force, and stress checksums were captured for every run. Across
the
general and fast comparisons, the maximum main-to-frozen energy-sum
relative
difference was below `4.5e-5`, the maximum-force checksum absolute
difference
was `2.54e-4`, and the maximum-stress checksum absolute difference was
`3.24e-6`. These differences are on the same scale as repeated
distributed GPU
runs of the same commit.

Focused regression tests compare frozen and unfrozen energy, force,
stress, and
Hessian tensors for both execution modes. The fast backend uses
`hessian_vmap=False` because its custom backward operators do not
implement
`vmap` batching.

## Validation

- Pre-commit passed for every modified file after merging `origin/main`.
- `tests/core/common/test_gp_utils.py`: 18 passed on PyTorch 2.13.
- `test_untrained_hessian_cpu`: passed with all model parameters frozen.
- Frozen/unfrozen GPU derivative parity: 2 passed (`general` and
  `umas_fast_gpu`) in Slurm job `9976221`.
- General one-GPU ABBA benchmark: completed without failures in job
`9974069`.
- Fast 8- and 32-GPU ABBA benchmarks: completed without failures in jobs
  `9973417` and `9973418`.
- `graph_parallel_group_size=None` remains the no-GP path; group size 1
remains
  an intentional singleton-GP test path.

---------

Co-authored-by: Ray Gao <7001989+rayg1234@users.noreply.github.com>

Author: bkmi

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

src/fairchem/core/units/mlip_unit/predict.py

@@ -37,9 +37,7 @@
 )
 from fairchem.core.components.batch_server import get_app_handle_with_retry
 from fairchem.core.datasets.atomic_data import AtomicData, warn_if_upcasting
-from fairchem.core.models.uma.nn.execution_backends import (
-    maybe_update_settings_backend,
-)
+from fairchem.core.models.uma.nn.execution_backends import maybe_update_settings_backend
 from fairchem.core.units.mlip_unit import InferenceSettings
 from fairchem.core.units.mlip_unit.mlip_unit import OutputSpec, Task
 from fairchem.core.units.mlip_unit.single_atom_patch import (
@@ -461,6 +459,8 @@ def _lazy_init(self, data: AtomicData) -> None:
         """
         # Model handles its own preparation (MOLE merge, eval mode, etc.)
         self.model.module.prepare_for_inference(data, self.inference_settings)
+        # Inference differentiates outputs with respect to inputs, not weights.
+        self.model.requires_grad_(False)
 
         self.model.to(self.inference_settings.base_precision_dtype)
 

tests/core/common/test_gp_utils.py

@@ -31,6 +31,22 @@ def _dummy_call(x):
     return x
 
 
+def _group_state():
+    return (
+        gp_utils.initialized(),
+        gp_utils.get_gp_world_size(),
+        gp_utils.get_dp_world_size(),
+    )
+
+
+def test_graph_parallel_size_one_initializes_groups():
+    config = PGConfig(backend="gloo", world_size=1, gp_group_size=1, use_gp=True)
+
+    output = spawn_multi_process(config, _group_state, init_pg_and_rank_and_launch_test)
+
+    assert output == [(True, 1, 1)]
+
+
 @pytest.mark.parametrize(
     "world_size, input, expected_output",
     [(1, 5, [5]), (3, 0, [0, 0, 0])],

tests/core/components/test_uma_speed_benchmark.py

@@ -10,8 +10,12 @@
 import tempfile
 from pathlib import Path
 
+import pytest
+
 from tests.core.testing_utils import launch_main
 
+pytestmark = pytest.mark.pretrained("uma-s-1p2")
+
 COMMON_ARGS = [
     "job.device_type=CPU",
     "runner.timeiters=1",

tests/core/units/mlip_unit/test_predict.py

@@ -1796,6 +1796,10 @@ def _test_untrained_hessian(checkpoint_path, device):
     # Get predictions
     preds = predictor.predict(batch)
 
+    assert all(
+        not parameter.requires_grad for parameter in predictor.model.parameters()
+    )
+
     # Verify energy, forces, and hessian are present
     assert "energy" in preds, "Energy prediction missing"
     assert "forces" in preds, "Forces prediction missing"
@@ -1816,6 +1820,77 @@ def _test_untrained_hessian(checkpoint_path, device):
     assert torch.allclose(hessian, hessian.T, atol=1e-5), "Hessian is not symmetric"
 
 
+@pytest.mark.gpu()
+@pytest.mark.parametrize("execution_mode", ["general", "umas_fast_gpu"])
+def test_frozen_parameters_preserve_input_derivatives(
+    conserving_mole_checkpoint, monkeypatch, execution_mode
+):
+    _test_frozen_parameters_preserve_input_derivatives(
+        conserving_mole_checkpoint[0], monkeypatch, execution_mode
+    )
+
+
+def _test_frozen_parameters_preserve_input_derivatives(
+    checkpoint_path, monkeypatch, execution_mode
+):
+    settings = InferenceSettings(
+        predict_untrained_forces={"omol"},
+        predict_untrained_stress={"omol"},
+        predict_untrained_hessian={"omol"},
+        activation_checkpointing=False,
+        merge_mole=True,
+        execution_mode=execution_mode,
+        hessian_vmap=False,
+    )
+
+    seed_everywhere(42)
+    frozen_predictor = MLIPPredictUnit(
+        checkpoint_path, device="cuda", inference_settings=settings
+    )
+    seed_everywhere(42)
+    unfrozen_predictor = MLIPPredictUnit(
+        checkpoint_path, device="cuda", inference_settings=settings
+    )
+    monkeypatch.setattr(
+        unfrozen_predictor.model,
+        "requires_grad_",
+        lambda requires_grad=True: unfrozen_predictor.model,
+    )
+
+    def make_batch():
+        atoms = molecule("H2O")
+        atoms.info.update({"charge": 0, "spin": 1})
+        data = AtomicData.from_ase(
+            atoms,
+            task_name="omol",
+            r_data_keys=["spin", "charge"],
+            molecule_cell_size=120,
+        )
+        return atomicdata_list_to_batch([data])
+
+    seed_everywhere(42)
+    frozen = frozen_predictor.predict(make_batch())
+    seed_everywhere(42)
+    unfrozen = unfrozen_predictor.predict(make_batch())
+
+    assert all(
+        not parameter.requires_grad for parameter in frozen_predictor.model.parameters()
+    )
+    assert any(
+        parameter.requires_grad for parameter in unfrozen_predictor.model.parameters()
+    )
+    assert frozen.keys() == unfrozen.keys()
+    for name in frozen:
+        atol = 1e-5 if name == "hessian" else 1e-6
+        torch.testing.assert_close(
+            frozen[name],
+            unfrozen[name],
+            rtol=1e-5,
+            atol=atol,
+            msg=lambda message, name=name: f"{name}: {message}",
+        )
+
+
 @pytest.mark.gpu()
 def test_hessian_activation_checkpointing(conserving_mole_checkpoint):
     """
@@ -1993,6 +2068,7 @@ def test_execution_mode_not_set_when_conditions_not_met(pretrained_model_name):
 # ---------------------------------------------------------------------------
 
 
+@pytest.mark.pretrained("uma-s-1p1")
 def test_uma_1p1_predict_unit_has_model_id():
     """UMA 1.1 checkpoints have no `model_id` on disk; the compat fixup
     back-fills it to `"UMA-1.1"` at load time."""
@@ -2008,6 +2084,7 @@ def test_uma_1p1_predict_unit_has_model_id():
     assert pu.model.module.backbone.model_id == UMA_1P1_MODEL_ID
 
 
+@pytest.mark.pretrained("uma-s-1p1")
 def test_uma_1p1_finetune_propagates_model_id():
     """When finetuning starts from UMA 1.1, the back-filled `model_id` is
     stashed onto `model.finetune_model_full_config` (the fixup runs inside