@@ -7,6 +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 (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`.
 
@@ -73,6 +73,9 @@ def _tune_save_load_helper(model_factory, samples, checkpoint_name, wrap_model_f
         checkpoint (TorchCheckpoint): Checkpoint used to save and load the model.
         device_map (dict | None): Device map passed to `checkpoint.load`.
 
+    Returns:
+        The loaded model.
+
     Raises:
         AssertionError: If model predictions after load don't match expected results.
     """
@@ -92,6 +95,7 @@ def _tune_save_load_helper(model_factory, samples, checkpoint_name, wrap_model_f
     with torch.no_grad():
         preds = model_loaded(test_data)
     torch.testing.assert_close(preds, expected, rtol=1e-3, atol=1e-3)
+    return model_loaded
 
 
 def _wrap_whole_model(model, backend=None):
@@ -104,6 +108,16 @@ 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
+
+
 @requires_xpu
 @pytest.mark.parametrize(
     "backend",
@@ -227,6 +241,113 @@ 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
+):
+    """Test tuning, saving and loading a pipeline whose attribute contains tuned modules."""
+    _tune_save_load_helper(
+        complex_pipeline_factory,
+        sample,
+        "nested_pipeline_test.tt",
+        _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
+):
+    """Test that a nested tuned module is addressed by its full path, not by its attribute name."""
+    with expect_device_map_not_found():
+        _tune_save_load_helper(
+            complex_pipeline_factory,
+            sample,
+            "nested_pipeline_test.tt",
+            _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):
+    """Test that only the tuned module is saved, not the rest of its pipeline component."""
+    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"]
+    assert list(state_dict["net.linear1"]) == [""]
+
+
+@requires_xpu
+def test_tune_save_load_pipeline_with_aliased_nested_tuned_module(complex_pipeline_factory, sample, checkpoint):
+    """Test that a tuned module reachable by two pipeline paths is saved once and restored as one module."""
+    pipeline = _alias_net_linear1(_wrap_net_layers(complex_pipeline_factory()))
+    tune(pipeline, sample, batch_sizes=[1, 2])
+    assert list(TorchCheckpoint.state_dict_from_pipeline(pipeline)) == ["net.linear1", "net.linear2"]
+
+    pipeline = _tune_save_load_helper(
+        lambda: _alias_net_linear1(complex_pipeline_factory()),
+        sample,
+        "aliased_pipeline_test.tt",
+        lambda pipeline: _alias_net_linear1(_wrap_net_layers(pipeline)),
+        checkpoint,
+    )
+
+    assert isinstance(pipeline.net.linear1, Module)
+    assert pipeline.linear1 is pipeline.net.linear1
+
+
+class _Vae(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.decoder = torch.nn.Linear(2, 2)
+
+
+class _PipelineWithDecoderAlias:
+    def __init__(self):
+        self.vae = _Vae()
+        self.decoder = self.vae.decoder
+
+
+def test_load_state_dict_for_pipeline_restores_duplicate_alias_keys_once(monkeypatch):
+    """Test loading a checkpoint that stores an aliased module under each of its paths."""
+    restored = _PipelineWithDecoderAlias()
+    replacement = torch.nn.Linear(2, 2)
+    calls = []
+
+    def load_module(module, state_dict, device_map, module_name):
+        calls.append(module_name)
+        device_map.pop(module_name, None)
+        return replacement
+
+    monkeypatch.setattr(TorchCheckpoint, "load_state_dict_for_module", staticmethod(load_module))
+    device_map = {"decoder": torch.device("cpu"), "vae.decoder": torch.device("cpu")}
+
+    TorchCheckpoint.load_state_dict_for_pipeline(restored, {"decoder": {}, "vae.decoder": {}}, device_map)
+
+    assert calls == ["decoder"]
+    assert restored.decoder is replacement
+    assert restored.vae.decoder is replacement
+    assert device_map == {}
+
+
 @requires_xpu
 def test_get_pipeline_modules(pipeline_factory):
     pipeline = pipeline_factory()
 
@@ -4,7 +4,7 @@
 # NOTE: This file has been modified by Intel Corporation.
 """Torch checkpoint module."""
 
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
 from pathlib import Path
 from typing import Any
 
@@ -40,18 +40,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,39 +105,55 @@ 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: A dictionary mapping attribute names to modules. Tuned modules nested in a module
+                attribute are added under their dotted path, e.g. `"vae.decoder"`.
         """
-        modules = {}
+        modules: dict[str, TorchModule | WrapperModule] = {}
+
+        def add_wrapped_children(module: TorchModule, prefix: str) -> None:
+            # `named_modules()` cannot be used: a wrapper forwards it to the wrapped module, which yields itself.
+            for child_name, child in module.named_children():
+                child_path = f"{prefix}.{child_name}"
+                if isinstance(child, WrapperModule):
+                    modules[child_path] = child
+                else:
+                    add_wrapped_children(child, child_path)
 
         for name, attr in vars(pipeline).items():
-            if isinstance(attr, (torch.nn.Module, WrapperModule)):
+            if isinstance(attr, WrapperModule):
                 modules[name] = attr
+            elif isinstance(attr, TorchModule):
+                modules[name] = attr
+                add_wrapped_children(attr, name)
 
         return modules
 
     @staticmethod
     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.
 
         Returns:
             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))
         if not state_dict:
             raise ValueError("No tuned modules found in the pipeline")
         return state_dict
@@ -249,13 +263,35 @@ def replace_tuned_module(
             final_module.load_state_dict(state_dict, strict=False)
         return final_module
 
+    @staticmethod
+    def get_pipeline_module_paths(pipeline: Any) -> dict[int, list[tuple[Any, str]]]:
+        """Get every place where the pipeline holds a module.
+
+        Args:
+            pipeline: The pipeline to search.
+
+        Returns:
+            dict: Maps `id()` of a module to the `(owner, attribute name)` pairs that hold it.
+        """
+        module_paths: dict[int, list[tuple[Any, str]]] = defaultdict(list)
+
+        for name, module in vars(pipeline).items():
+            if isinstance(module, (TorchModule, WrapperModule)):
+                module_paths[id(module)].append((pipeline, name))
+            if isinstance(module, TorchModule):
+                for _, parent in module.named_modules(remove_duplicate=False):
+                    for child_name, child in parent.named_children():
+                        module_paths[id(child)].append((parent, child_name))
+
+        return module_paths
+
     @staticmethod
     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.
@@ -270,18 +306,30 @@ def load_state_dict_for_pipeline(pipeline, state_dict, device_map: dict[str, tor
                         if a matched attribute is not a torch.nn.Module or WrapperModule.
         """
         matched_modules = set()
+        module_paths = TorchCheckpoint.get_pipeline_module_paths(pipeline)
+        restored_module_ids = set()
         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:
-                    raise ValueError(f"Module {name} is not a torch.nn.Module or WrapperModule")
+            *parent_path, attribute = name.split(".")
+            parent = pipeline
+            for component in parent_path:
+                parent = getattr(parent, component, None)
+            if parent is None or not hasattr(parent, attribute):
+                continue
+
+            module = getattr(parent, attribute)
+            if not isinstance(module, (torch.nn.Module, WrapperModule)):
+                raise ValueError(f"Module {name} is not a torch.nn.Module or WrapperModule")
+
+            matched_modules.add(name)
+            if id(module) in restored_module_ids:
+                # Older checkpoints store an aliased module under each of its paths.
+                device_map.pop(name, None)
+                continue
+
+            restored_module = TorchCheckpoint.load_state_dict_for_module(module, module_state_dict, device_map, name)
+            for owner, child_name in module_paths[id(module)]:
+                setattr(owner, child_name, restored_module)
+            restored_module_ids.add(id(restored_module))
 
         unmatched_keys = set(state_dict.keys()).difference(matched_modules)
         if unmatched_keys: