Update composition and scaler to use mts.learn.nn.Module

This simplify the forward that no longer needs to move data to a
different dtype. Instead the model is moved once in the trainer
and then stays on the same device

Author: Luthaf

src/metatrain/composition/__init__.py

@@ -31,6 +31,7 @@ def train_or_load_composition_model(
     atomic_baseline: FixedCompositionWeights | str,
     train_datasets: List[Union[Dataset, Subset]],
     other_additive_models: List[nn.Module],
+    device: torch.device,
     batch_size: int,
     is_distributed: bool,
     checkpoint_dir: str = "",
@@ -46,6 +47,8 @@ def train_or_load_composition_model(
     :param train_datasets: Training datasets
     :param other_additive_models: Other additive models (e.g. ZBL) to
         subtract before fitting
+    :param device: Device on which to train the composition model. This should
+        be a device that supports float64 (e.g. CPU or CUDA, but not MPS).
     :param batch_size: Batch size for data loading
     :param is_distributed: Whether training is distributed
     :param checkpoint_dir: Directory to save the composition model checkpoint
@@ -86,7 +89,6 @@ def train_or_load_composition_model(
                     f"'{target_info.unit}'."
                 )
         composition_model.load_state_dict(loaded.state_dict())
-        composition_model.sync_tensor_maps()
     else:
         assert isinstance(atomic_baseline, dict)
         logging.info("Calculating composition weights")
@@ -98,12 +100,10 @@ def train_or_load_composition_model(
         hypers["distributed"] = is_distributed
         trainer = Trainer(hypers=hypers)
         trainer._additive_models = other_additive_models
-        # The trainer fits on devices[0]; pass the model's current device so
-        # embedded training stays where the parent architecture put the model.
         trainer.train(
             model=composition_model,
             dtype=torch.float64,
-            devices=[composition_model.dummy_buffer.device],
+            devices=[device],
             train_datasets=train_datasets,
             val_datasets=train_datasets,
             checkpoint_dir=checkpoint_dir,

src/metatrain/composition/_base_composition.py

@@ -10,12 +10,13 @@
 import metatensor.torch as mts
 import torch
 from metatensor.torch import Labels, LabelsEntry, TensorBlock, TensorMap
+from metatensor.torch.learn import nn
 from metatomic.torch import ModelOutput, System
 
 from .documentation import FixedCompositionWeights  # noqa: F401
 
 
-class BaseCompositionModel(torch.nn.Module):
+class BaseCompositionModel(nn.Module):
     """
     Fits a composition model for a dict of targets.
 
@@ -88,9 +89,11 @@ def __init__(
         self.atomic_types = torch.as_tensor(atomic_types, dtype=torch.int32)
         self.target_names = []
         self.sample_kinds = {}
-        self.XTX = {}
-        self.XTY = {}
-        self.weights = {}
+        # `XTX` and `XTY` are only used during fitting, not at inference, so they
+        # are registered as non-persistent buffers to keep them out of the state_dict.
+        self.register_buffer("XTX", {}, persistent=False)
+        self.register_buffer("XTY", {}, persistent=False)
+        self.register_buffer("weights", {})
 
         # go from an atomic type to its position in `self.atomic_types`
         self.register_buffer(
@@ -167,7 +170,6 @@ def add_output(self, target_name: str, layout: TensorMap) -> None:
                     values=torch.zeros(
                         len(self.atomic_types),
                         len(self.atomic_types),
-                        dtype=torch.float64,
                     ),
                     samples=Labels(["center_type"], self.atomic_types.reshape(-1, 1)),
                     components=[],
@@ -186,7 +188,6 @@ def add_output(self, target_name: str, layout: TensorMap) -> None:
                         len(self.atomic_types),
                         *[len(c) for c in block.components],
                         len(block.properties),
-                        dtype=torch.float64,
                     ),
                     samples=Labels(["center_type"], self.atomic_types.reshape(-1, 1)),
                     components=block.components,
@@ -203,7 +204,6 @@ def add_output(self, target_name: str, layout: TensorMap) -> None:
                         len(self.atomic_types),
                         *[len(c) for c in block.components],
                         len(block.properties),
-                        dtype=torch.float64,
                     ),
                     samples=Labels(["center_type"], self.atomic_types.reshape(-1, 1)),
                     components=block.components,
@@ -239,10 +239,15 @@ def accumulate(
         :param targets: Dict of target names to :py:class:`TensorMap` containing
             the target values for each system in the batch.
         """
-
         device = systems[0].positions.device
         dtype = systems[0].positions.dtype
-        self._sync_device_dtype(device, dtype)
+
+        if dtype != torch.float64:
+            raise ValueError(
+                "Composition model accumulation must be done in float64. "
+                f"Got systems with dtype {dtype}. Please move the systems to "
+                "float64 before accumulating."
+            )
 
         # check that the systems contain no unexpected atom types
         for system in systems:
@@ -386,6 +391,15 @@ def fit(
         if targets_to_fit is None:
             targets_to_fit = self.target_names
 
+        if len(targets_to_fit) > 0:
+            dtype = self.XTX[targets_to_fit[0]].block(0).values.dtype
+            if dtype != torch.float64:
+                raise ValueError(
+                    "Composition model fitting must be done in float64. "
+                    f"Got dtype {dtype}. Please move the model to "
+                    "float64 before fitting (e.g. `model.to(dtype=torch.float64)`)."
+                )
+
         sanitized_fixed_weights = self._sanitize_fixed_weights(fixed_weights)
 
         # fit
@@ -486,10 +500,7 @@ def forward(
         :raises ValueError: If no weights have been computed or if `outputs` keys
             contain unsupported keys.
         """
-
         device = systems[0].positions.device
-        dtype = systems[0].positions.dtype
-        self._sync_device_dtype(device, dtype)
 
         # Build the sample labels that are required
         _, sample_labels = _get_system_indices_and_labels(systems, device)
@@ -616,32 +627,6 @@ def _compute_X_per_atom(
         )
         return one_hot_encoding.to(dtype)
 
-    def _sync_device_dtype(self, device: torch.device, dtype: torch.dtype) -> None:
-        """
-        Move the accumulated quantities and the fitted weights to the given
-        device and dtype.
-
-        Needed because they are stored as ``TensorMap`` dicts, which
-        ``torch.nn.Module.to`` does not move.
-
-        :param device: Device to move the quantities to.
-        :param dtype: Dtype to convert the quantities to.
-        """
-        self.atomic_types = self.atomic_types.to(device=device)
-        self.type_to_index = self.type_to_index.to(device=device)
-        self.XTX = {
-            target_name: tm.to(device=device, dtype=dtype)
-            for target_name, tm in self.XTX.items()
-        }
-        self.XTY = {
-            target_name: tm.to(device=device, dtype=dtype)
-            for target_name, tm in self.XTY.items()
-        }
-        self.weights = {
-            target_name: tm.to(device=device, dtype=dtype)
-            for target_name, tm in self.weights.items()
-        }
-
 
 def _include_key(key: LabelsEntry) -> bool:
     """

src/metatrain/composition/checkpoints.py

@@ -1,3 +1,51 @@
+import torch
+
+
+def model_update_v1_v2(checkpoint: dict, prefix: str = "") -> None:
+    """
+    Update model checkpoint from version 1 to version 2.
+
+    The model now uses metatensor's ``nn.Module`` and ``register_buffer`` instead of
+    manually tracking the data.
+
+    :param checkpoint: The checkpoint to update.
+    :param prefix: Prefix prepended to the state_dict keys of the composition model
+        (e.g. ``"additive_models.0."`` when the composition model is nested inside
+        another architecture's ``additive_models[0]``).
+    """
+    scaler_key = f"{prefix}model"
+
+    for key in ["model_state_dict", "best_model_state_dict"]:
+        if (state_dict := checkpoint.get(key)) is None:
+            continue
+
+        # If both model_state_dict and best_model_state_dict point to the same
+        # dict, the upgrade was already applied in the first iteration.
+        if f"{scaler_key}._mts_helper" in state_dict:
+            continue
+
+        dummy_buffer = state_dict[f"{prefix}dummy_buffer"]
+        empty_tensor = torch.zeros(
+            0, dtype=dummy_buffer.dtype, device=dummy_buffer.device
+        )
+
+        extra_state: dict[str, dict] = {"weights": {}}
+
+        for target_name in checkpoint["model_data"]["dataset_info"].targets:
+            buffer_key = f"{prefix}{target_name}_composition_buffer"
+            if buffer_key not in state_dict:
+                continue
+
+            extra_state["weights"][target_name] = (
+                "metatensor.TensorMap",
+                state_dict.pop(buffer_key),
+                empty_tensor,
+            )
+
+        state_dict[f"{scaler_key}._mts_helper"] = empty_tensor
+        state_dict[f"{scaler_key}._extra_state"] = extra_state
+
+
 def trainer_update_v1_v2(checkpoint: dict) -> None:
     """
     Update a v1 Trainer checkpoint to v2.