diff --git a/ChangeLog.md b/ChangeLog.md index 59c96761..da0b3bad 100644 --- a/ChangeLog.md +++ b/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. diff --git a/tests/unit/torch/checkpoint/test_torch_checkpoint.py b/tests/unit/torch/checkpoint/test_torch_checkpoint.py index d8957c83..647b7382 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.""" @@ -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.""" diff --git a/torch_tweak/torch/checkpoint/torch_checkpoint.py b/torch_tweak/torch/checkpoint/torch_checkpoint.py index 18eec190..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: @@ -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.