From b30aa5100874cd85e63800ccde020bcf8e66402d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?= Date: Fri, 23 Oct 2026 11:59:57 +0200 Subject: [PATCH 1/2] fix(checkpoint): drop entries of nested tuned modules from the loaded state_dict MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `load` handed each child a filtered copy of the state_dict. The entry of a nested tuned module was removed from that copy only. It stayed in the state_dict that `load_state_dict(strict=False)` receives, which ignored it. The `if state_dict` skip for a fully tuned module never applied below the top level. Children now share the same state_dict, which also drops one copy per module. Signed-off-by: Łukasz Ślusarczyk --- ChangeLog.md | 4 ++++ tests/unit/torch/checkpoint/test_torch_checkpoint.py | 12 ++++++++++++ torch_tweak/torch/checkpoint/torch_checkpoint.py | 6 ++---- 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/ChangeLog.md b/ChangeLog.md index 59c96761..d4f7dba5 100644 --- a/ChangeLog.md +++ b/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. diff --git a/tests/unit/torch/checkpoint/test_torch_checkpoint.py b/tests/unit/torch/checkpoint/test_torch_checkpoint.py index d8957c83..1b9027f8 100644 --- a/tests/unit/torch/checkpoint/test_torch_checkpoint.py +++ b/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.""" diff --git a/torch_tweak/torch/checkpoint/torch_checkpoint.py b/torch_tweak/torch/checkpoint/torch_checkpoint.py index 18eec190..db082571 100644 --- a/torch_tweak/torch/checkpoint/torch_checkpoint.py +++ b/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, ) From c8ad2dff22f7c20b293a6591afef6ce01157d911 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?= Date: Fri, 23 Oct 2026 12:29:08 +0200 Subject: [PATCH 2/2] fix(checkpoint): match device_map keys by dotted path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `load` looked up a nested tuned module in `device_map` by its attribute name. Two tuned modules named `linear1` in different submodules shared one entry. The key is now the dotted path, e.g. `inner.linear1`. A direct submodule keeps its key, so `linear1` still matches `model.linear1`. Signed-off-by: Łukasz Ślusarczyk --- ChangeLog.md | 1 + .../torch/checkpoint/test_torch_checkpoint.py | 52 +++++++++++++++++++ .../torch/checkpoint/torch_checkpoint.py | 14 +++-- 3 files changed, 63 insertions(+), 4 deletions(-) diff --git a/ChangeLog.md b/ChangeLog.md index d4f7dba5..da0b3bad 100644 --- a/ChangeLog.md +++ b/ChangeLog.md @@ -8,6 +8,7 @@ SPDX-License-Identifier: Apache-2.0 ## 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 diff --git a/tests/unit/torch/checkpoint/test_torch_checkpoint.py b/tests/unit/torch/checkpoint/test_torch_checkpoint.py index 1b9027f8..647b7382 100644 --- a/tests/unit/torch/checkpoint/test_torch_checkpoint.py +++ b/tests/unit/torch/checkpoint/test_torch_checkpoint.py @@ -219,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.""" diff --git a/torch_tweak/torch/checkpoint/torch_checkpoint.py b/torch_tweak/torch/checkpoint/torch_checkpoint.py index db082571..1cdb0eec 100644 --- a/torch_tweak/torch/checkpoint/torch_checkpoint.py +++ b/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: @@ -230,6 +234,7 @@ def replace_tuned_module( parent_module=module, 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, ) @@ -238,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.