Prepare UMA inference gradients before compilation (#2143)

This standalone PR replaces
[#2126](#2126). The
ghstack PR was marked merged into its temporary `gh/mlazos/4/base`
branch, but its patch did not land on `main`.

UMA computes forces and stress by differentiating energy with respect to
positions and cells. The backbone previously enabled those gradients
with `requires_grad_` inside its compiled forward. Dynamo cannot
represent that tensor metadata mutation, so it stopped and resumed
tracing at each occurrence.

This change prepares position and cell gradient metadata in
`MLIPPredictUnit` after device and dtype conversion but before entering
the model. Guarded fallback calls remain in the backbone for direct
users; predictor-prepared inputs already satisfy those guards, so normal
compiled inference performs no metadata mutation inside the captured
region.

The fix is independent of topology source. It applies with both
`external_graph_gen=True` and the final validation configuration using
`external_graph_gen=False`, `internal_graph_gen_version=3`,
`merge_mole=True`, and `compile_dynamic_shapes=False`. No independent
flag enables it:

```python
settings = InferenceSettings(compile=True)
```

A compiled UMA-S-1p2 run previously reported five `requires_grad_`
graph-break sites and none after this change.

**Test plan**

```bash
PYTHONPATH=$PWD/src:$PYTHONPATH /home/mlazos/.conda/envs/pytorch-3.12/bin/python -m pytest -q tests/core/units/mlip_unit/test_predict.py -k test_prepare_inference_gradients
/home/mlazos/.conda/envs/pytorch-3.12/bin/pre-commit run --files src/fairchem/core/models/uma/escn_md.py src/fairchem/core/units/mlip_unit/predict.py tests/core/units/mlip_unit/test_predict.py
PYTHONPATH=$PWD/src:$PYTHONPATH /home/mlazos/.conda/envs/pytorch-3.12/bin/python -m compileall -q src/fairchem/core/models/uma/escn_md.py src/fairchem/core/units/mlip_unit/predict.py tests/core/units/mlip_unit/test_predict.py
```

Authored with assistance from Codex.

Author: mlazos

src/fairchem/core/models/uma/escn_md.py

@@ -780,9 +780,11 @@ def forward(self, data_dict: AtomicData) -> dict[str, torch.Tensor]:
         # Must be set before graph generation so the computation graph
         # tracks positions and cell through edge distance calculations.
         if not self.regress_config.direct_forces:
-            if self.regress_config.forces or self.regress_config.stress:
+            if (
+                self.regress_config.forces or self.regress_config.stress
+            ) and not data_dict["pos"].requires_grad:
                 data_dict["pos"].requires_grad_(True)
-            if self.regress_config.stress:
+            if self.regress_config.stress and not data_dict["cell"].requires_grad:
                 data_dict["cell"].requires_grad_(True)
 
         with record_function("generate_graph"):

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

@@ -90,6 +90,16 @@ def collated_predict(
     return collated_predict
 
 
+def _prepare_inference_gradients(backbone, data: AtomicData) -> None:
+    regress_config = getattr(backbone, "regress_config", None)
+    if regress_config is None or regress_config.direct_forces:
+        return
+    if regress_config.forces or regress_config.stress:
+        data["pos"].requires_grad_(True)
+    if regress_config.stress:
+        data["cell"].requires_grad_(True)
+
+
 class MLIPPredictUnitProtocol(Protocol):
     def predict(self, data: AtomicData, undo_element_references: bool) -> dict: ...
 
@@ -494,6 +504,9 @@ def predict(
             if torch.is_tensor(val) and val.is_floating_point():
                 data_device[key] = val.to(dtype)
 
+        backbone = self.model.module.backbone
+        _prepare_inference_gradients(backbone, data_device)
+
         # Model handles any per-prediction checks (e.g., MOLE consistency)
         try:
             self.model.module.on_predict_check(data_device)

tests/core/units/mlip_unit/test_predict.py

@@ -43,7 +43,10 @@
 from fairchem.core.models.uma.nn.execution_backends import UMASFastGPUBackend
 from fairchem.core.units.mlip_unit import InferenceSettings, MLIPPredictUnit
 from fairchem.core.units.mlip_unit.mlip_unit import initialize_finetuning_model
-from fairchem.core.units.mlip_unit.predict import ParallelMLIPPredictUnit
+from fairchem.core.units.mlip_unit.predict import (
+    ParallelMLIPPredictUnit,
+    _prepare_inference_gradients,
+)
 from fairchem.core.units.mlip_unit.single_atom_patch import (
     single_atom_prediction_from_lookup,
 )
@@ -67,6 +70,32 @@ def _resolve_checkpoint_path(name_or_path: str) -> str:
 ATOL = 5e-4
 
 
+@pytest.mark.parametrize(
+    ("forces", "stress", "pos_grad", "cell_grad"),
+    [
+        (False, False, False, False),
+        (True, False, True, False),
+        (True, True, True, True),
+    ],
+)
+def test_prepare_inference_gradients(forces, stress, pos_grad, cell_grad):
+    backbone = SimpleNamespace(
+        regress_config=SimpleNamespace(
+            direct_forces=False, forces=forces, stress=stress
+        )
+    )
+    data = {
+        "pos": torch.randn(4, 3),
+        "cell": torch.randn(1, 3, 3),
+    }
+
+    _prepare_inference_gradients(backbone, data)
+
+    assert data["pos"].requires_grad is pos_grad
+    assert data["cell"].requires_grad is cell_grad
+    _prepare_inference_gradients(SimpleNamespace(), data)
+
+
 _REPRESENTATIVE_ELEMENTS = [
     (1, 0, 2),  # H:  charge=0, spin=2
     (6, 0, 3),  # C:  charge=0, spin=3