diff --git a/ChangeLog.md b/ChangeLog.md index bbca75d..3c0b29b 100644 --- a/ChangeLog.md +++ b/ChangeLog.md @@ -7,8 +7,10 @@ SPDX-License-Identifier: Apache-2.0 ## Unreleased +- fix: `save` stores tuned modules nested in pipeline components under their dotted path, e.g. `vae.decoder`. Modules that are not tuned are not saved. +- fix: a tuned module that a pipeline holds under several paths is saved once and restored as one module. +- feat: `save` warns about tuned modules it does not save, e.g. a module the pipeline holds in a plain list. - fix: a passthrough module without parameters keeps its inputs on their device instead of moving them to CPU. - - feat (breaking): `inspect` measures module depth in executed modules. Containers such as `nn.ModuleList` and modules that never run no longer add a level. `min_depth` is renamed to `depth`. diff --git a/tests/unit/torch/checkpoint/test_torch_checkpoint.py b/tests/unit/torch/checkpoint/test_torch_checkpoint.py index fe95cd4..202ddab 100644 --- a/tests/unit/torch/checkpoint/test_torch_checkpoint.py +++ b/tests/unit/torch/checkpoint/test_torch_checkpoint.py @@ -5,6 +5,7 @@ # NOTE: This file has been modified by Intel Corporation. """Tests for model state storage functionality.""" +import logging from unittest.mock import Mock import pytest @@ -61,13 +62,12 @@ def sample(model_factory): yield model_factory().sample().xpu() -def _tune_save_load_helper(model_factory, samples, checkpoint_name, wrap_model_fn, checkpoint, device_map=None): +def _tune_save_load_helper(model_factory, samples, wrap_model_fn, checkpoint, device_map=None): """Tune a wrapped model, save it, load it into a fresh model and compare predictions. Args: model_factory: Function that creates a fresh model instance. samples (torch.Tensor): Input samples for model inference. - checkpoint_name (str): File name of the checkpoint inside the storage root. wrap_model_fn: Function that takes a model and returns a wrapped version ready for tuning. checkpoint (TorchCheckpoint): Checkpoint used to save and load the model. @@ -83,11 +83,11 @@ def _tune_save_load_helper(model_factory, samples, checkpoint_name, wrap_model_f wrapped_model = wrap_model_fn(model) tune(wrapped_model, samples, batch_sizes=[1, 2]) - checkpoint.save(wrapped_model, checkpoint_name) + checkpoint.save(wrapped_model, "checkpoint.tt") # create a new model instance and load the tuned model model = model_factory() - model_loaded = checkpoint.load(model, checkpoint_name, device_map=device_map) + model_loaded = checkpoint.load(model, "checkpoint.tt", device_map=device_map) with torch.no_grad(): preds = model_loaded(test_data) @@ -104,6 +104,31 @@ def _wrap_linear_layers(model): return model +def _wrap_net_layers(pipeline): + _wrap_linear_layers(pipeline.net) + return pipeline + + +def _alias_net_linear1(pipeline): + pipeline.linear1 = pipeline.net.linear1 + return pipeline + + +def _tune_aliased_pipeline(pipeline, sample): + pipeline = _alias_net_linear1(_wrap_net_layers(pipeline)) + tune(pipeline, sample, batch_sizes=[1, 2]) + return pipeline + + +class _PipelineWithLayerList(ToyPipeline): + def __init__(self): + super().__init__() + self.layers = [vars(self).pop("linear2")] + + def __call__(self, x): + return self.layers[0](self.linear1(self.preprocessing(x))) + + @requires_xpu @pytest.mark.parametrize( "backend", @@ -114,17 +139,13 @@ def _wrap_linear_layers(model): ) # cover both compiled and eager backends def test_tune_save_load_whole_model(model_factory, sample, backend, checkpoint): """Test tuning, saving and loading a whole model wrapped in a single Module.""" - _tune_save_load_helper( - model_factory, sample, "top_model_test.tt", lambda model: _wrap_whole_model(model, backend), checkpoint - ) + _tune_save_load_helper(model_factory, sample, lambda model: _wrap_whole_model(model, backend), checkpoint) @requires_xpu def test_tune_save_load_whole_model_with_device_map(model_factory, sample, torch_device, checkpoint): """Test loading a whole tuned model with a device_map entry for the top module.""" - _tune_save_load_helper( - model_factory, sample, "top_model_test.tt", _wrap_whole_model, checkpoint, device_map={"": torch_device} - ) + _tune_save_load_helper(model_factory, sample, _wrap_whole_model, checkpoint, device_map={"": torch_device}) @requires_xpu @@ -134,7 +155,6 @@ def test_tune_save_load_whole_model_with_invalid_device_map(model_factory, sampl _tune_save_load_helper( model_factory, sample, - "top_model_test.tt", _wrap_whole_model, checkpoint, device_map={"linear123": torch_device}, @@ -148,7 +168,6 @@ def test_tune_save_load_part_of_model(model_factory, sample, torch_device, check _tune_save_load_helper( model_factory, sample, - "partial_model_test.tt", _wrap_linear_layers, checkpoint, device_map=dict.fromkeys(device_map_keys, torch_device), @@ -162,7 +181,6 @@ def test_tune_save_load_part_of_model_with_invalid_device_map(model_factory, sam _tune_save_load_helper( model_factory, sample, - "partial_model_test.tt", _wrap_linear_layers, checkpoint, device_map={"linear123": torch_device}, @@ -176,7 +194,6 @@ def test_tune_save_load_pipeline(pipeline_factory, sample, torch_device, checkpo _tune_save_load_helper( pipeline_factory, sample, - "pipeline_test.tt", _wrap_linear_layers, checkpoint, device_map=dict.fromkeys(device_map_keys, torch_device), @@ -190,7 +207,6 @@ def test_tune_save_load_pipeline_with_invalid_device_map(pipeline_factory, sampl _tune_save_load_helper( pipeline_factory, sample, - "pipeline_test.tt", _wrap_linear_layers, checkpoint, device_map={"linear123": torch_device}, @@ -227,6 +243,141 @@ def test_tune_save_load_complex_pipeline_layers( torch.testing.assert_close(preds, expected, rtol=1e-4, atol=1e-4) +@requires_xpu +@pytest.mark.parametrize( + "device_map_keys", + [ + pytest.param([], id="no_device_map"), + pytest.param(["net.linear1", "net.linear2"], id="full_device_map"), + pytest.param(["net.linear1"], id="partial_device_map"), + ], +) +def test_tune_save_load_pipeline_with_nested_tuned_modules( + complex_pipeline_factory, sample, torch_device, checkpoint, device_map_keys +): + _tune_save_load_helper( + complex_pipeline_factory, + sample, + _wrap_net_layers, + checkpoint, + device_map=dict.fromkeys(device_map_keys, torch_device), + ) + + +@requires_xpu +def test_tune_save_load_pipeline_with_nested_tuned_modules_rejects_leaf_name_in_device_map( + complex_pipeline_factory, sample, torch_device, checkpoint +): + with expect_device_map_not_found(): + _tune_save_load_helper( + complex_pipeline_factory, + sample, + _wrap_net_layers, + checkpoint, + device_map={"linear1": torch_device}, + ) + + +@requires_xpu +def test_state_dict_from_pipeline_stores_only_nested_tuned_modules(complex_pipeline_factory, sample): + pipeline = complex_pipeline_factory() + pipeline.net.linear1 = Module( + pipeline.net.linear1, "demo-simple1", strategy=OneBackendStrategy(TorchInductorBackend()) + ) + tune(pipeline, sample, batch_sizes=[1, 2]) + + state_dict = TorchCheckpoint.state_dict_from_pipeline(pipeline) + + assert list(state_dict) == ["net.linear1"] + + +@requires_xpu +def test_state_dict_from_pipeline_stores_aliased_tuned_module_once(complex_pipeline_factory, sample): + pipeline = _tune_aliased_pipeline(complex_pipeline_factory(), sample) + + state_dict = TorchCheckpoint.state_dict_from_pipeline(pipeline) + + assert list(state_dict) == ["net.linear1", "net.linear2"] + + +@requires_xpu +def test_tune_save_load_pipeline_restores_aliased_tuned_module_as_one_module( + complex_pipeline_factory, sample, checkpoint +): + checkpoint.save(_tune_aliased_pipeline(complex_pipeline_factory(), sample), "checkpoint.tt") + + restored = checkpoint.load(_alias_net_linear1(complex_pipeline_factory()), "checkpoint.tt") + + assert isinstance(restored.linear1, Module) + assert restored.linear1 is restored.net.linear1 + + +@requires_xpu +def test_tune_save_load_pipeline_restores_tuned_module_aliased_by_its_parent( + complex_pipeline_factory, sample, checkpoint +): + pipeline = _wrap_net_layers(complex_pipeline_factory()) + pipeline.net.linear1_alias = pipeline.net.linear1 + tune(pipeline, sample, batch_sizes=[1, 2]) + checkpoint.save(pipeline, "checkpoint.tt") + restored = complex_pipeline_factory() + restored.net.linear1_alias = restored.net.linear1 + + restored = checkpoint.load(restored, "checkpoint.tt") + + assert isinstance(restored.net.linear1_alias, Module) + assert restored.net.linear1_alias is restored.net.linear1 + + +@requires_xpu +def test_load_state_dict_for_pipeline_restores_duplicate_alias_keys_as_one_module(complex_pipeline_factory, sample): + state_dict = TorchCheckpoint.state_dict_from_pipeline(_tune_aliased_pipeline(complex_pipeline_factory(), sample)) + state_dict["linear1"] = state_dict["net.linear1"] + restored = _alias_net_linear1(complex_pipeline_factory()) + + TorchCheckpoint.load_state_dict_for_pipeline(restored, state_dict, device_map={}) + + assert isinstance(restored.linear1, Module) + assert restored.linear1 is restored.net.linear1 + + +@requires_xpu +def test_load_state_dict_for_pipeline_consumes_device_map_for_duplicate_alias_keys( + complex_pipeline_factory, sample, torch_device +): + state_dict = TorchCheckpoint.state_dict_from_pipeline(_tune_aliased_pipeline(complex_pipeline_factory(), sample)) + state_dict["linear1"] = state_dict["net.linear1"] + device_map = {"linear1": torch_device, "net.linear1": torch_device} + + TorchCheckpoint.load_state_dict_for_pipeline(_alias_net_linear1(complex_pipeline_factory()), state_dict, device_map) + + assert device_map == {} + + +@requires_xpu +@pytest.mark.parametrize("key", ["missing", "net.missing", "missing.linear1", "net.linear1.missing"]) +def test_load_state_dict_for_pipeline_rejects_missing_path(complex_pipeline_factory, key): + with pytest.raises(ValueError, match="Some modules in the state_dict were not found in the pipeline"): + TorchCheckpoint.load_state_dict_for_pipeline(complex_pipeline_factory(), {key: {}}, device_map={}) + + +@requires_xpu +def test_state_dict_from_pipeline_warns_about_tuned_module_outside_module_attributes(sample, caplog): + pipeline = _PipelineWithLayerList() + pipeline.linear1 = Module(pipeline.linear1, "demo-simple1", strategy=OneBackendStrategy(TorchInductorBackend())) + pipeline.layers[0] = Module(pipeline.layers[0], "demo-simple2", strategy=OneBackendStrategy(TorchInductorBackend())) + tune(pipeline, sample, batch_sizes=[1, 2]) + # `tune` calls `setup_logging`, which removes the `caplog` handler from the root logger. + logging.getLogger().addHandler(caplog.handler) + caplog.clear() + + TorchCheckpoint.state_dict_from_pipeline(pipeline) + + assert caplog.messages == [ + "Tuned modules not saved, the pipeline does not hold them in module attributes: demo-simple2" + ] + + @requires_xpu def test_get_pipeline_modules(pipeline_factory): pipeline = pipeline_factory() @@ -252,6 +403,15 @@ def __init__(self): assert isinstance(modules, dict) +def test_get_pipeline_modules_does_not_descend_into_module_twice(): + pipeline = ToyComplexPipeline() + pipeline.net.self_ref = pipeline.net + + modules = TorchCheckpoint.get_pipeline_modules(pipeline) + + assert list(modules) == ["net", "net.linear1", "net.linear2", "net.self_ref"] + + @requires_xpu def test_state_dict_from_pipeline(pipeline_factory): pipeline = pipeline_factory() diff --git a/torch_tweak/torch/checkpoint/torch_checkpoint.py b/torch_tweak/torch/checkpoint/torch_checkpoint.py index 18eec19..a2ab840 100644 --- a/torch_tweak/torch/checkpoint/torch_checkpoint.py +++ b/torch_tweak/torch/checkpoint/torch_checkpoint.py @@ -4,6 +4,7 @@ # NOTE: This file has been modified by Intel Corporation. """Torch checkpoint module.""" +import logging from collections import OrderedDict from pathlib import Path from typing import Any @@ -14,8 +15,12 @@ from torch_tweak.torch.checkpoint.checkpoint import Checkpoint from torch_tweak.torch.checkpoint.storage import Storage from torch_tweak.torch.module.wrapper_module import Module as WrapperModule +from torch_tweak.torch.module.wrapper_module import ModuleState +from torch_tweak.torch.module_registry import MODULE_REGISTRY from torch_tweak.torch.utils.device import get_device +logger = logging.getLogger(__name__) + class TorchCheckpoint(Checkpoint): """Class for storing tuned module using torch module's state serialization. @@ -40,18 +45,16 @@ def save( This method handles both individual torch modules and pipelines containing modules. For individual modules, it directly uses the module's state_dict() method. - For pipelines, it extracts state dictionaries from all tuned modules in the pipeline. - - The pipeline is expected to be an object with attributes that are tuned modules (WrapperModule instances). - Only tuned modules are saved in the state dictionary, as regular torch modules don't have tuning information. - - If no tuned modules are found in the pipeline, a ValueError is raised. + For a pipeline, it saves each tuned module (WrapperModule instance) under its dotted path, + e.g. `"vae.decoder"`. Modules that are not tuned are not saved. Args: module_or_pipeline: The module or pipeline to save. This can be either a torch.nn.Module - or a pipeline object with module attributes. For pipelines, only tuned modules - (WrapperModule instances) will be saved. + or a pipeline object with module attributes. path: The path where the module state should be saved. + + Raises: + ValueError: If the pipeline contains no tuned modules. """ if isinstance(module_or_pipeline, TorchModule): state_dict = module_or_pipeline.state_dict() @@ -107,20 +110,36 @@ def load( return module_or_pipeline @staticmethod - def get_pipeline_modules(pipeline: Any) -> dict[str, torch.nn.Module]: + def get_pipeline_modules(pipeline: Any) -> dict[str, TorchModule | WrapperModule]: """Get the modules from the pipeline. Args: pipeline: The pipeline object to extract modules from. Returns: - dict: A dictionary mapping attribute names to Module objects. + dict: Maps the dotted path of every module the pipeline holds, e.g. `"vae.decoder"`, to the module. + The search does not descend into tuned modules or into a module it has already visited. """ - modules = {} - + modules: dict[str, TorchModule | WrapperModule] = {} + visited_module_ids = set() + + def add_module(path: str, module: TorchModule | WrapperModule) -> None: + # Every path to a module is kept, so that `load_state_dict_for_pipeline` replaces each alias. + # Children are listed under the first path only: replacing a child on a shared owner covers all paths. + modules[path] = module + if isinstance(module, WrapperModule) or id(module) in visited_module_ids: + return + visited_module_ids.add(id(module)) + # `named_children()` yields a module held twice by one parent only once. + for child_name, child in module._modules.items(): + if child is not None: + add_module(f"{path}.{child_name}", child) + + # TODO: search lists, dicts and plain objects as `ModuleInspector` does, so that every module + # `tune` tunes is also saved. for name, attr in vars(pipeline).items(): - if isinstance(attr, (torch.nn.Module, WrapperModule)): - modules[name] = attr + if isinstance(attr, (TorchModule, WrapperModule)): + add_module(name, attr) return modules @@ -128,6 +147,8 @@ def get_pipeline_modules(pipeline: Any) -> dict[str, torch.nn.Module]: def state_dict_from_pipeline(pipeline: Any) -> dict[str, Any]: """Extract the state_dict from a pipeline (e.g. HF Diffusers pipeline). + Each tuned module is stored once under its dotted path, also when the pipeline reaches it by several paths. + Args: pipeline: The pipeline to extract the state_dict from. @@ -135,11 +156,24 @@ def state_dict_from_pipeline(pipeline: Any) -> dict[str, Any]: dict: The state_dict from the pipeline. """ state_dict = OrderedDict() - pipeline_modules = TorchCheckpoint.get_pipeline_modules(pipeline) + saved_module_ids = set() - for name, module in pipeline_modules.items(): - if isinstance(module, WrapperModule): + for name, module in TorchCheckpoint.get_pipeline_modules(pipeline).items(): + if isinstance(module, WrapperModule) and id(module) not in saved_module_ids: state_dict[name] = module.state_dict() + saved_module_ids.add(id(module)) + + unsaved_module_names = [ + name + for name, module in MODULE_REGISTRY.modules.items() + if module.state == ModuleState.TUNED and id(module) not in saved_module_ids + ] + if unsaved_module_names: + logger.warning( + "Tuned modules not saved, the pipeline does not hold them in module attributes: %s", + ", ".join(unsaved_module_names), + ) + if not state_dict: raise ValueError("No tuned modules found in the pipeline") return state_dict @@ -253,9 +287,9 @@ def replace_tuned_module( def load_state_dict_for_pipeline(pipeline, state_dict, device_map: dict[str, torch.device]): """Load a state dictionary into a pipeline. - This method loads the state dictionary into the pipeline by matching module names - in the state dictionary with attributes in the pipeline. Each matched module is - replaced with the tuned module. + This method loads the state dictionary into the pipeline by matching dotted module paths + in the state dictionary with modules in the pipeline. Each matched module is replaced + with the tuned module in every place where the pipeline holds it. Args: pipeline: The pipeline to load the state dictionary into. @@ -269,19 +303,31 @@ def load_state_dict_for_pipeline(pipeline, state_dict, device_map: dict[str, tor ValueError: If a module in the state_dict is not found in the pipeline or if a matched attribute is not a torch.nn.Module or WrapperModule. """ + modules = TorchCheckpoint.get_pipeline_modules(pipeline) matched_modules = set() + restored_modules = {} for name, module_state_dict in state_dict.items(): - if hasattr(pipeline, name): - module = getattr(pipeline, name) - if isinstance(module, (torch.nn.Module, WrapperModule)): - matched_modules.add(name) - setattr( - pipeline, - name, - TorchCheckpoint.load_state_dict_for_module(module, module_state_dict, device_map, name), - ) - else: + module = modules.get(name) + if module is None: + if hasattr(pipeline, name): raise ValueError(f"Module {name} is not a torch.nn.Module or WrapperModule") + continue + + matched_modules.add(name) + if id(module) in restored_modules: + # Older checkpoints store an aliased module under each of its paths. + device_map.pop(name, None) + continue + + restored_modules[id(module)] = TorchCheckpoint.load_state_dict_for_module( + module, module_state_dict, device_map, name + ) + + for path, module in modules.items(): + if id(module) in restored_modules: + owner_path, _, attribute = path.rpartition(".") + owner = modules[owner_path] if owner_path else pipeline + setattr(owner, attribute, restored_modules[id(module)]) unmatched_keys = set(state_dict.keys()).difference(matched_modules) if unmatched_keys: