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] 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 59c9676..d4f7dba 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 d8957c8..1b9027f 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 18eec19..db08257 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, )