fix(checkpoint): match device_map keys by dotted path

ChangeLog.md

@@ -5,6 +5,11 @@ SPDX-License-Identifier: Apache-2.0
 
 # Changelog
 
+## Unreleased
+
+- fix: `load` removes the entries of nested tuned modules before it loads the rest of the `state_dict`.
+- fix: `device_map` keys for nested tuned modules are full dotted paths, e.g. `inner.linear1`. The attribute name alone, e.g. `linear1`, now matches only a direct submodule.
+
 ## 0.4.0
 
 This is the first release of **Torch Tweak**. Torch Tweak is focused on Intel GPU (XPU) support.

tests/unit/torch/checkpoint/test_torch_checkpoint.py

@@ -149,6 +149,18 @@ def wrap_partial_model(model):
     _tune_save_load_helper(model_factory, sample, "partial_model_test.tt", wrap_partial_model, checkpoint)
 
 
+@requires_xpu
+def test_load_state_dict_for_module_consumes_tuned_submodule_keys(model_factory, sample):
+    model = model_factory()
+    model.linear1 = Module(model.linear1, "demo-simple1", strategy=OneBackendStrategy(TorchInductorBackend()))
+    tune(model, sample, batch_sizes=[1, 2], dry_run=False, disable_external_logging=False)
+    state_dict = model.state_dict()
+
+    TorchCheckpoint.load_state_dict_for_module(model_factory(), state_dict, device_map={})
+
+    assert list(state_dict) == ["linear2.weight", "linear2.bias"]
+
+
 @requires_xpu
 def test_tune_save_load_part_of_model_with_full_device_map(model_factory, sample, torch_device, checkpoint):
     """Test tuning, saving and loading a model with individual layers wrapped in Modules."""
@@ -207,6 +219,58 @@ def wrap_partial_model(model):
         )
 
 
+class _NestedToyModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.inner = ToyTorchModel(is_linear=True)
+
+    def forward(self, x):
+        return self.inner(x)
+
+
+def _wrap_inner_layers(model):
+    model.inner.linear1 = Module(
+        model.inner.linear1, "demo-simple1", strategy=OneBackendStrategy(TorchInductorBackend())
+    )
+    model.inner.linear2 = Module(
+        model.inner.linear2, "demo-simple2", strategy=OneBackendStrategy(TorchInductorBackend())
+    )
+    return model
+
+
+@requires_xpu
+@pytest.mark.parametrize(
+    "device_map_keys",
+    [
+        pytest.param(["inner.linear1", "inner.linear2"], id="full_device_map"),
+        pytest.param(["inner.linear1"], id="partial_device_map"),
+    ],
+)
+def test_tune_save_load_nested_part_of_model_with_device_map(sample, torch_device, checkpoint, device_map_keys):
+    """Test that a tuned module nested two levels deep is addressed by its dotted path."""
+    _tune_save_load_helper(
+        lambda: _NestedToyModel().xpu(),
+        sample,
+        "nested_model_test.tt",
+        _wrap_inner_layers,
+        checkpoint,
+        device_map=dict.fromkeys(device_map_keys, torch_device),
+    )
+
+
+@requires_xpu
+def test_tune_save_load_nested_part_of_model_rejects_leaf_name_in_device_map(sample, torch_device, checkpoint):
+    with pytest.raises(ValueError, match="Some modules in the device_map were not found: .*"):
+        _tune_save_load_helper(
+            lambda: _NestedToyModel().xpu(),
+            sample,
+            "nested_model_test.tt",
+            _wrap_inner_layers,
+            checkpoint,
+            device_map={"linear1": torch_device},
+        )
+
+
 @requires_xpu
 def test_tune_save_load_pipeline(pipeline_factory, sample, checkpoint):
     """Test tuning, saving and loading a pipeline."""

torch_tweak/torch/checkpoint/torch_checkpoint.py

@@ -79,7 +79,9 @@ def load(
             module_or_pipeline: The module or pipeline to load the state into. This can be
                 either a torch.nn.Module or a pipeline object with module attributes.
             path: The path from where the module state should be loaded.
-            device_map: The device map to load module to.
+            device_map: Maps dotted paths of tuned modules to devices, e.g. `"inner.linear1"` for `linear1`
+                inside `inner`. For a pipeline, a path starts with the pipeline attribute name.
+                The key `""` is the module itself.
 
         Returns:
             The loaded module, which may be either the original module type or a wrapped module
@@ -172,7 +174,7 @@ def load_state_dict_for_module(
                 The dictionary should have keys that correspond to the module structure,
                 with tuned module states stored in a format recognized by WrapperModule.
             device_map (dict): The device map to load module to.
-            module_name (str): The name of the module for device map matching.
+            module_name (str): The name of the module in its pipeline, prefixed to `device_map` keys of submodules.
 
         Returns:
             torch.nn.Module: The module with tuned components loaded from the state dictionary.
@@ -192,6 +194,7 @@ def replace_tuned_module(
             parent_module: TorchModule | None = None,
             prefix: str = "",
             module_name: str = "",
+            module_path: str = "",
         ) -> None:
             """Traverses the module and replaces torch modules with tuned modules.
 
@@ -203,12 +206,13 @@ def replace_tuned_module(
                 device_map: The device map to load modules to.
                 parent_module: The parent module.
                 prefix: The prefix of the module.
-                module_name: The name of the module.
+                module_name: The attribute name of the module in `parent_module`.
+                module_path: The dotted path of the module, used as the `device_map` key.
             """
             stored_object = local_state_dict.get(prefix, {})
             if WrapperModule.is_state_dict_valid(stored_object):
                 # This is a tuned module, replace with wrapper module
-                device = device_map.pop(module_name, None)
+                device = device_map.pop(module_path, None)
 
                 wrapper_module = WrapperModule.from_dict(module, stored_object, device)
                 if parent_module is None:
@@ -224,14 +228,13 @@ def replace_tuned_module(
                 # This is an original torch module
                 for child_name, child_module in module.named_children():
                     if child_module is not None:
-                        child_prefix = prefix + child_name + "."
-                        child_state_dict = {k: v for k, v in local_state_dict.items() if k.startswith(child_prefix)}
                         replace_tuned_module(
                             module=child_module,
-                            local_state_dict=child_state_dict,
+                            local_state_dict=local_state_dict,
                             parent_module=module,
-                            prefix=child_prefix,
+                            prefix=prefix + child_name + ".",
                             module_name=child_name,
+                            module_path=f"{module_path}.{child_name}" if module_path else child_name,
                             device_map=device_map,
                         )
 
@@ -240,6 +243,7 @@ def replace_tuned_module(
             local_state_dict=state_dict,
             device_map=device_map,
             module_name=module_name,
+            module_path=module_name,
         )
 
         # Handle original torch modules not tuned, strict=False because wrapped module keys removed.