Update composition and scaler to use mts.nn.Module

Author: Luthaf

src/metatrain/composition/__init__.py

@@ -86,7 +86,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")

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(
@@ -242,7 +245,6 @@ def accumulate(
 
         device = systems[0].positions.device
         dtype = systems[0].positions.dtype
-        self._sync_device_dtype(device, dtype)
 
         # check that the systems contain no unexpected atom types
         for system in systems:
@@ -486,10 +488,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 +615,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,48 @@
+import torch
+
+
+def model_update_v1_v2(checkpoint: dict) -> 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.
+    """
+    scaler_key = "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["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"{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.

src/metatrain/composition/model.py

@@ -2,9 +2,8 @@
 import warnings
 from typing import Dict, List, Literal, Optional, Union
 
-import metatensor.torch as mts
 import torch
-from metatensor.torch import Labels, TensorBlock, TensorMap
+from metatensor.torch import Labels, TensorMap
 from metatomic.torch import (
     AtomisticModel,
     ModelCapabilities,
@@ -28,7 +27,6 @@
 from ._base_composition import (
     BaseCompositionModel,
     FixedCompositionWeights,
-    _include_key,
 )
 from .documentation import ModelHypers
 
@@ -46,7 +44,7 @@ class CompositionModel(ModelInterface[ModelHypers]):
         model.
     """
 
-    __checkpoint_version__ = 1
+    __checkpoint_version__ = 2
     __supported_devices__ = ["cuda", "cpu"]
     __supported_dtypes__ = [torch.float64]
     __default_metadata__ = ModelMetadata(
@@ -266,9 +264,8 @@ def restart(self, dataset_info: DatasetInfo) -> "CompositionModel":
         self.dataset_info = merged_info
 
         self._new_outputs = []
-        buffer_names = [n for n, _ in self.named_buffers()]
         for target_name, target_info in self.target_infos.items():
-            if target_name + "_composition_buffer" in buffer_names:
+            if target_name in self.model.weights:
                 continue
             self._new_outputs.append(target_name)
             self.model.add_output(target_name, target_info.layout)
@@ -298,11 +295,6 @@ def forward(
         :return: A dictionary mapping each requested output name to the
             corresponding ``TensorMap`` containing the computed values.
         """
-        dtype = systems[0].positions.dtype
-        device = systems[0].positions.device
-
-        self.weights_to(device, dtype)
-
         for output_name in outputs.keys():
             if output_name not in self.outputs:
                 raise ValueError(
@@ -349,43 +341,6 @@ def _add_output(self, target_name: str, target_info: TargetInfo) -> None:
             description=target_info.description,
         )
 
-        layout = mts.filter_blocks(
-            target_info.layout,
-            Labels(
-                target_info.layout.keys.names,
-                torch.vstack(
-                    [key.values for key in target_info.layout.keys if _include_key(key)]
-                ),
-                assume_unique=True,
-            ),
-        )
-
-        fake_weights = TensorMap(
-            keys=layout.keys,
-            blocks=[
-                TensorBlock(
-                    values=torch.zeros(
-                        (len(self.atomic_types),) + b.values.shape[1:],
-                        dtype=torch.float64,
-                    ),
-                    samples=Labels(
-                        names=["center_type"],
-                        values=torch.tensor(self.atomic_types, dtype=torch.int).reshape(
-                            -1, 1
-                        ),
-                        assume_unique=True,
-                    ),
-                    components=b.components,
-                    properties=b.properties,
-                )
-                for b in layout.blocks()
-            ],
-        )
-        self.register_buffer(
-            target_name + "_composition_buffer",
-            mts.save_buffer(mts.make_contiguous(fake_weights)),
-        )
-
     def remove_output(self, target_name: str) -> None:
         """
         Remove a previously registered output target, mirroring ``_add_output``.
@@ -395,32 +350,6 @@ def remove_output(self, target_name: str) -> None:
         self.outputs.pop(target_name, None)
         self.dataset_info.targets.pop(target_name, None)
         self.model.remove_output(target_name)
-        buffer_name = target_name + "_composition_buffer"
-        if hasattr(self, buffer_name):
-            delattr(self, buffer_name)
-
-    def weights_to(self, device: torch.device, dtype: torch.dtype) -> None:
-        """
-        Move the fitted weights and the accumulated quantities to the given
-        device and dtype.
-
-        Needed because they are stored as ``TensorMap`` attributes, which
-        ``torch.nn.Module.to`` does not move.
-
-        :param device: Device to move the weights to.
-        :param dtype: Dtype to convert the weights to.
-        """
-        if len(self.model.weights) != 0:
-            if self.model.weights[list(self.model.weights.keys())[0]].device != device:
-                self.model.weights = {
-                    k: v.to(device) for k, v in self.model.weights.items()
-                }
-            if self.model.weights[list(self.model.weights.keys())[0]].dtype != dtype:
-                self.model.weights = {
-                    k: v.to(dtype) for k, v in self.model.weights.items()
-                }
-
-        self.model._sync_device_dtype(device, dtype)
 
     @staticmethod
     def is_valid_target(target_name: str, target_info: TargetInfo) -> bool:
@@ -472,21 +401,6 @@ def is_valid_target(target_name: str, target_info: TargetInfo) -> bool:
 
         return True
 
-    def sync_tensor_maps(self) -> None:
-        """
-        Reload the weight ``TensorMap`` objects from the registered buffers.
-
-        Must be called after the buffers change through means that bypass the
-        model, e.g. ``load_state_dict``.
-        """
-        for k in self.dataset_info.targets:
-            buffer = self.__getattr__(k + "_composition_buffer")
-            # ``mts.load_buffer`` dereferences the buffer on the host, so it
-            # segfaults on a GPU buffer: deserialize on the CPU and move the
-            # weights back to the buffer's device.
-            weights = mts.load_buffer(buffer.to(device="cpu"))
-            self.model.weights[k] = weights.to(device=buffer.device)
-
     def get_checkpoint(self) -> Dict:
         """
         Get the checkpoint of the model.
@@ -543,7 +457,6 @@ def load_checkpoint(
         )
 
         model.load_state_dict(model_state_dict)
-        model.sync_tensor_maps()
 
         model.metadata = merge_metadata(model.metadata, checkpoint.get("metadata"))
 
@@ -587,14 +500,11 @@ def export(self, metadata: Optional[ModelMetadata] = None) -> AtomisticModel:
             raise ValueError(f"unsupported dtype {dtype} for composition model")
 
         self.to(dtype)
-        self.weights_to(torch.device("cpu"), torch.float64)
-
-        interaction_range = 0.0
 
         capabilities = ModelCapabilities(
             outputs=self.outputs,
             atomic_types=self.atomic_types,
-            interaction_range=interaction_range,
+            interaction_range=0.0,
             length_unit=self.dataset_info.length_unit,
             supported_devices=self.__supported_devices__,
             dtype=dtype_to_str(dtype),

src/metatrain/composition/tests/checkpoints/model-v2_trainer-v2.ckpt.gz

@@ -77,6 +77,7 @@ class TestCheckpoints(CheckpointTests, CompositionTests):
     incompatible_trainer_checkpoints = [
         "checkpoints/model-v1_trainer-v1.ckpt.gz",
         "checkpoints/model-v1_trainer-v2.ckpt.gz",
+        "checkpoints/model-v2_trainer-v2.ckpt.gz",
     ]
 
 

src/metatrain/composition/tests/test_basic.py

@@ -236,6 +236,7 @@ def test_checkpoint_roundtrip_predictions():
 def test_regression_checkpoint():
     with gzip.open("regression_checkpoint.ckpt.gz", "rb") as fd:
         checkpoint = torch.load(fd, weights_only=False)
+    checkpoint = CompositionModel.upgrade_checkpoint(checkpoint)
     model = CompositionModel.load_checkpoint(checkpoint, context="export")
     model.eval()
 

src/metatrain/composition/tests/test_regression.py

@@ -3,7 +3,6 @@
 from pathlib import Path
 from typing import Any, Dict, List, Literal, Union
 
-import metatensor.torch as mts
 import torch
 
 from metatrain.utils.abc import ModelInterface, TrainerInterface
@@ -80,7 +79,8 @@ def train(
         else:
             device = devices[0]
             logging.info(f"Training on device {device} with dtype {dtype}")
-        model.to(device=device)
+
+        model.to(device=device, dtype=torch.float64)
 
         # Targets with fixed weights don't need data accumulation, only fit().
         targets_to_accumulate = [
@@ -190,7 +190,7 @@ def train(
             # A rank whose shard of some dataset was empty never accumulated
             # the corresponding targets, so its XTX/XTY are still on the CPU,
             # while NCCL needs them on the GPU for the all_reduce.
-            model.model._sync_device_dtype(device, torch.float64)
+            model.to(device=device, dtype=torch.float64)
             handles = []
             for target_name in targets_to_accumulate:
                 for XTX_block, XTY_block in zip(
@@ -209,16 +209,6 @@ def train(
 
         model.model.fit(fixed_weights, targets_to_fit=model._new_outputs)
 
-        for target_name in model.model.weights.keys():
-            model.register_buffer(
-                target_name + "_composition_buffer",
-                mts.save_buffer(
-                    mts.make_contiguous(
-                        model.model.weights[target_name].to("cpu", torch.float64)
-                    )
-                ).to(device),
-            )
-
         if checkpoint_dir and (not is_distributed or torch.distributed.get_rank() == 0):
             ckpt_path = Path(checkpoint_dir) / "composition_model.ckpt"
             self.save_checkpoint(model, ckpt_path)

src/metatrain/composition/trainer.py

@@ -463,8 +463,6 @@ def load_checkpoint(
         dtype = next(model.model.parameters()).dtype
 
         model.to(dtype).load_state_dict(model_state_dict)
-        model.additive_models[0].sync_tensor_maps()
-        model.scaler.sync_tensor_maps()
 
         # Loading the metadata from the checkpoint
         metadata = checkpoint.get("metadata", None)
@@ -483,11 +481,6 @@ def export(self, metadata: Optional[ModelMetadata] = None) -> AtomisticModel:
         # float64
         self.to(dtype)
 
-        # Additionally, the composition model contains some `TensorMap`s that cannot
-        # be registered correctly with Pytorch. This function moves them:
-
-        self.additive_models[0].weights_to(torch.device("cpu"), torch.float64)
-
         interaction_ranges = [self.hypers["descriptor"]["repflow"]["e_rcut"]]
         for additive_model in self.additive_models:
             if hasattr(additive_model, "cutoff_radius"):