fix(checkpoint): drop entries of nested tuned modules from the loaded state_dict

ChangeLog.md

@@ -5,6 +5,10 @@ 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`.
+
 ## 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."""

torch_tweak/torch/checkpoint/torch_checkpoint.py

@@ -224,13 +224,11 @@ 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,
                             device_map=device_map,
                         )