add train=true/false flag to freeze checkpoint weights

Author: dallasfoster

nvalchemi/models/uma.py

@@ -130,6 +130,10 @@ class UMAWrapper(nn.Module, BaseModelMixin):
         UMA task: one of ``_UMA_TASKS``. Determines which per-task head
         in the multi-task model is used and which inputs (charge, spin)
         must be populated.
+    train : bool
+        ``False`` (default) freezes all weights for inference (lossless for
+        autograd forces); ``True`` keeps fairchem's trainable/frozen split
+        so weights stay exposed for fine-tuning.
 
     Attributes
     ----------
@@ -139,7 +143,9 @@ class UMAWrapper(nn.Module, BaseModelMixin):
         The UMA task this wrapper is pinned to.
     """
 
-    def __init__(self, predict_unit: Any, task_name: str = "omol") -> None:
+    def __init__(
+        self, predict_unit: Any, task_name: str = "omol", train: bool = False
+    ) -> None:
         super().__init__()
         if task_name not in _UMA_TASKS:
             raise ValueError(
@@ -198,6 +204,16 @@ def __init__(self, predict_unit: Any, task_name: str = "omol") -> None:
             active_outputs=active_outputs,
         )
 
+        # Inference (train=False): freeze all weights — conservative forces
+        # come from autograd w.r.t. positions, so this is lossless and avoids
+        # building a weight-grad graph each forward. Training: leave fairchem's
+        # loaded trainable/frozen split intact so weights stay exposed.
+        self._train = train
+        if not train:
+            for p in self.predict_unit.model.parameters():
+                p.requires_grad_(False)
+        self.train(train)
+
     # ------------------------------------------------------------------
     # Construction
     # ------------------------------------------------------------------
@@ -210,6 +226,7 @@ def from_checkpoint(
         device: str | torch.device = "cpu",
         inference_settings: Any = "default",
         overrides: dict | None = None,
+        train: bool = False,
     ) -> "UMAWrapper":
         """Resolve and load a UMA checkpoint.
 
@@ -242,6 +259,14 @@ def from_checkpoint(
         overrides : dict | None
             Optional overrides forwarded to fairchem's inference-settings
             builder.
+        train : bool
+            If ``False`` (default), freeze all weights for inference —
+            lossless, since conservative forces come from autograd on
+            positions. If ``True``, keep fairchem's loaded trainable/frozen
+            split so weights remain exposed for fine-tuning. Note: the
+            ``forward`` path goes through fairchem's inference ``predict``
+            (eval mode, detached forces); gradient-based training requires a
+            separate path through the raw model.
         """
         import os as _os
 
@@ -272,7 +297,7 @@ def from_checkpoint(
                 f"local file path. Known names: "
                 f"{sorted(pretrained_mlip.available_models)[:6]}..."
             )
-        return cls(predict_unit, task_name=task_name)
+        return cls(predict_unit, task_name=task_name, train=train)
 
     def _extract_cutoff(self) -> float:
         """Pull the radial cutoff from the loaded backbone.

test/models/test_uma.py

@@ -74,6 +74,8 @@ def __init__(self) -> None:
         self.r_max = torch.tensor(6.0)
         self.sph_feature_size = 16  # (lmax+1)² with lmax=3
         self.sphere_channels = 128
+        # A trainable weight so the train/freeze flag is observable in tests.
+        self.weight = nn.Parameter(torch.zeros(1))
 
     def forward(self, data: FCAtomicData) -> dict:
         n = data.pos.shape[0]
@@ -198,6 +200,18 @@ def test_task_stored(self, mock_omol):
     def test_cutoff_from_backbone(self, mock_omol):
         assert math.isclose(mock_omol.cutoff, 6.0, abs_tol=1e-6)
 
+    def test_inference_freezes_weights(self, mock_omol):
+        """train=False (default) freezes the underlying weights for inference."""
+        params = list(mock_omol.predict_unit.model.parameters())
+        assert params and all(not p.requires_grad for p in params)
+        assert mock_omol.training is False
+
+    def test_train_keeps_weights_trainable(self, mock_pu):
+        """train=True leaves weights trainable/exposed for fine-tuning."""
+        w = UMAWrapper(mock_pu, task_name="omol", train=True)
+        assert any(p.requires_grad for p in w.predict_unit.model.parameters())
+        assert w.training is True
+
 
 class TestModelConfig:
     def test_omol_active_outputs(self, mock_omol):