perf(inference): freeze all prepared model parameters

Author: bkmi

CLAUDE.md

@@ -247,12 +247,12 @@ configs/                 # Hydra YAML configs (datasets, tasks, backbone, optimi
 - 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.
-- Only freeze inference parameters for the `umas_fast_gpu` backend, after its
-  module replacement. On the H100 perf check, freezing the general backend cut
-  1000-atom QPS by about 50% with activation checkpointing either enabled or
-  disabled. The fast backend retained its throughput while saving 11-13% peak
-  memory. Its custom backward must preserve input derivatives independently of
-  parameter gradients.
+- 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

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

@@ -37,10 +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 (
-    ExecutionMode,
-    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 (
@@ -462,9 +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)
-        if self.inference_settings.execution_mode == ExecutionMode.UMAS_FAST_GPU:
-            # The fast backward only differentiates positions and cells.
-            self.model.requires_grad_(False)
+        # 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/units/mlip_unit/test_predict.py

@@ -1796,7 +1796,9 @@ def _test_untrained_hessian(checkpoint_path, device):
     # Get predictions
     preds = predictor.predict(batch)
 
-    assert any(parameter.requires_grad for parameter in predictor.model.parameters())
+    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"
@@ -1819,22 +1821,25 @@ def _test_untrained_hessian(checkpoint_path, device):
 
 
 @pytest.mark.gpu()
-def test_fast_gpu_frozen_parameters_preserve_input_derivatives(
-    conserving_mole_checkpoint, monkeypatch
+@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
+        conserving_mole_checkpoint[0], monkeypatch, execution_mode
     )
 
 
-def _test_frozen_parameters_preserve_input_derivatives(checkpoint_path, monkeypatch):
+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="umas_fast_gpu",
+        execution_mode=execution_mode,
         hessian_vmap=False,
     )