From ae9a6743847f82783de1d55a53617081ac19a62b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?=
 <lukasz.slusarczyk@intel.com>
Date: Wed, 21 Oct 2026 16:39:00 +0200
Subject: [PATCH 1/4] feat(inspecting): measure module depth in executed
 modules
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

inspect() counted depth in object path segments.
nn.ModuleList, plain containers, and modules that never ran each added a level.
Two depths could yield the same modules, and a deeper one could yield none.
get_modules() also dropped executed modules shallower than min_depth,
so those parts of the model were never tuned.

Depth is now the number of executed modules above a module.
min_depth is renamed to depth.
get_modules() returns every executed module at depth,
plus shallower executed modules without executed children.
It no longer resets forward_called.

max_depth is removed.
No returned module is deeper than depth, so the traversal limit follows from depth.

A diffusers pipeline that used min_depth=1 to reach its components now uses depth=0.

Signed-off-by: Łukasz Ślusarczyk <lukasz.slusarczyk@intel.com>
---
 ChangeLog.md                                  |   9 ++
 docs/guides/inspect.md                        |  45 ++++---
 examples/E5Large/e5large/tune.py              |   4 +-
 examples/ParakeetCTC/parakeet_ctc/tune.py     |   2 +-
 examples/ParakeetRNNT/parakeet_rnnt/tune.py   |   2 +-
 .../stable_diffusion/benchmark_hybrid.py      |   1 -
 .../torch/inspecting/test_module_inspector.py | 113 +++++++++++++++---
 tests/unit/torch/inspecting/test_wrapping.py  |   2 +-
 torch_tweak/torch/inspecting/inspecting.py    |  14 ++-
 torch_tweak/torch/inspecting/module_info.py   |  15 ++-
 .../torch/inspecting/module_inspector.py      |  81 ++++++-------
 11 files changed, 197 insertions(+), 91 deletions(-)

diff --git a/ChangeLog.md b/ChangeLog.md
index 59c96761..43d3d09c 100644
--- a/ChangeLog.md
+++ b/ChangeLog.md
@@ -5,6 +5,15 @@ SPDX-License-Identifier: Apache-2.0
 
 # Changelog
 
+## Unreleased
+
+- 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`.
+  `depth` keeps shallower executed modules without executed children, so a cut covers the whole model.
+  `max_depth` is removed, the traversal limit follows from `depth`.
+  A diffusers pipeline that used `min_depth=1` to reach its components now uses `depth=0`.
+
 ## 0.4.0
 
 This is the first release of **Torch Tweak**. Torch Tweak is focused on Intel GPU (XPU) support.
diff --git a/docs/guides/inspect.md b/docs/guides/inspect.md
index 58f85008..77399f2f 100644
--- a/docs/guides/inspect.md
+++ b/docs/guides/inspect.md
@@ -42,8 +42,7 @@ modules_info = tt.inspect(
     inference_function=None,  # Optional custom inference function
     number_of_iterations=10,  # Iterations for profiling
     warmup_iterations=5,  # Warmup iterations
-    min_depth=0,  # Minimum depth for module search
-    max_depth=5,  # Maximum depth for module search
+    depth=0,  # Depth at which the model is cut into modules
 )
 ```
 
@@ -133,24 +132,38 @@ modules_info = tt.inspect(
 )
 ```
 
-#### min_depth (Default: 0)
+#### Module depth
 
-Minimum depth level for module discovery. Increase if root-level modules don't work:
+The depth of a module is the number of executed modules above it.
+Depth 0 holds the outermost executed modules.
+That is the model itself, or the executed components of a pipeline object.
+Containers such as `nn.ModuleList`, lists, and dicts do not add a level.
+Modules whose `forward` never ran do not add a level either.
 
-```python
-# Start from root level (default)
-modules_info = tt.inspect(model, input_data, min_depth=0)
+For a model whose `forward` loops over `self.layers` (an `nn.ModuleList`) and then calls `self.head`:
 
-# Skip root, inspect children
-modules_info = tt.inspect(model, input_data, min_depth=1)
-modules_info = tt.inspect(model, input_data, min_depth=2)
-```
+| Depth | Modules |
+|---|---|
+| 0 | `model` |
+| 1 | `layers.0`, `layers.1`, `head` |
+| 2 | `layers.0.mlp`, `layers.1.mlp` |
+
+#### depth (Default: 0)
 
-#### max_depth (Default: 5)
+Depth at which the model is cut into modules for tuning.
+Every executed module at `depth` is returned.
+A shallower executed module without executed children is returned too, so no part of the model is left out.
+In the example above, `depth=2` returns `layers.0.mlp`, `layers.1.mlp`, and `head`.
 
-If a nested (child) module has a larger depth than `max_depth` it will be skipped from inspection.
+Increase it when the outermost modules cannot be tuned:
 
-Both `min_depth` and `max_depth` narrow the inspection search to a reasonable range.
+```python
+# Cut at the outermost executed modules (default)
+modules_info = tt.inspect(model, input_data, depth=0)
+
+# Cut one level deeper
+modules_info = tt.inspect(model, input_data, depth=1)
+```
 
 ## Working with InspectedModulesInfo
 
@@ -211,8 +224,8 @@ transformer_modules = [m for m in modules_info.get_modules() if "transformer" in
 ### Issue: No modules found
 
 ```python
-# Solution: Increase min_depth or check object structure
-modules_info = tt.inspect(model, input_data, min_depth=1)
+# Solution: Change depth or check object structure
+modules_info = tt.inspect(model, input_data, depth=1)
 ```
 
 ## Summary
diff --git a/examples/E5Large/e5large/tune.py b/examples/E5Large/e5large/tune.py
index c8ef63b5..c39e869f 100644
--- a/examples/E5Large/e5large/tune.py
+++ b/examples/E5Large/e5large/tune.py
@@ -80,10 +80,10 @@ def call_wrapper(*args, **kwargs):
             device="xpu",
         )
 
-    # NOTE: without min_depth=2, inspector finds a wrapper module `pipeline._modules["0"]` and fails with incorrect input in tunning
+    # NOTE: without depth=2, inspector finds a wrapper module `pipeline._modules["0"]` and fails with incorrect input in tunning
     #     with error: RuntimeError: Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment. @ torch/_tensor.py", line 136, in __deepcopy__
     #     during make_batch and deepcopy of arguments with all kwags arguments passed as dict in args[0]
-    inspected_modules_info = inspect(model, input_texts, inference_function=call_wrapper, min_depth=2)
+    inspected_modules_info = inspect(model, input_texts, inference_function=call_wrapper, depth=2)
     inspected_modules_info.describe()
 
     # NOTE: workaround for SentenceTransformer to work with tuned model see below first dirty hack
diff --git a/examples/ParakeetCTC/parakeet_ctc/tune.py b/examples/ParakeetCTC/parakeet_ctc/tune.py
index 910fb04b..dc8ad0c5 100644
--- a/examples/ParakeetCTC/parakeet_ctc/tune.py
+++ b/examples/ParakeetCTC/parakeet_ctc/tune.py
@@ -57,7 +57,7 @@ def call_wrapper(*args, **kwargs):
 
     input_data = [{"audio": str(audio_path)}]
 
-    inspected_modules_info = inspect(pipeline, input_data, inference_function=call_wrapper, min_depth=1)
+    inspected_modules_info = inspect(pipeline, input_data, inference_function=call_wrapper, depth=1)
     inspected_modules_info.describe()
 
     modules = inspected_modules_info.get_modules(min_execution_percentage=0.01)
diff --git a/examples/ParakeetRNNT/parakeet_rnnt/tune.py b/examples/ParakeetRNNT/parakeet_rnnt/tune.py
index 10a2b191..16d2f4f8 100644
--- a/examples/ParakeetRNNT/parakeet_rnnt/tune.py
+++ b/examples/ParakeetRNNT/parakeet_rnnt/tune.py
@@ -80,7 +80,7 @@ def call_wrapper(*args, **kwargs):
     call_wrapper(audio=str(audio_path))
 
     logger.info("Inspecting model...")
-    inspected_modules_info = inspect(pipeline, input_data, inference_function=call_wrapper, min_depth=1)
+    inspected_modules_info = inspect(pipeline, input_data, inference_function=call_wrapper, depth=1)
     inspected_modules_info.describe()
 
     modules = inspected_modules_info.get_modules(min_execution_percentage=0.01)
diff --git a/examples/StableDiffusion/stable_diffusion/benchmark_hybrid.py b/examples/StableDiffusion/stable_diffusion/benchmark_hybrid.py
index 3a313fcf..f23b1f89 100644
--- a/examples/StableDiffusion/stable_diffusion/benchmark_hybrid.py
+++ b/examples/StableDiffusion/stable_diffusion/benchmark_hybrid.py
@@ -101,7 +101,6 @@ def call_for_inspect(*args, **kwargs):
         inference_function=call_for_inspect,
         number_of_iterations=2,
         warmup_iterations=1,
-        min_depth=1,
     )
     modules = modules_info.get_modules(min_execution_percentage=min_execution_pct)
     logger.info("Wrapping %d modules (min_execution_percentage=%.2f)", len(modules), min_execution_pct)
diff --git a/tests/unit/torch/inspecting/test_module_inspector.py b/tests/unit/torch/inspecting/test_module_inspector.py
index 8b330e99..7e390805 100644
--- a/tests/unit/torch/inspecting/test_module_inspector.py
+++ b/tests/unit/torch/inspecting/test_module_inspector.py
@@ -101,7 +101,7 @@ def test_initialization():
     assert inspector._module_info == {}
     assert inspector._original_forward == {}
     assert inspector._inspected_objects == set()
-    assert inspector._max_recursion_depth == 5
+    assert inspector._depth == 0
 
 
 def test_reset(simple_model):
@@ -149,15 +149,6 @@ def test_inspect_nested_model(nested_model):
     assert root_module.name == "Sequential"
 
 
-def test_inspect_nested_model_max_depth(nested_model):
-    """Test inspecting a nested model."""
-    inspector = ModuleInspector(max_depth=1)
-    inspector.inspect(nested_model)
-
-    # Check that all modules were registered
-    assert len(inspector._module_info) == 4  # Sequential, Linear, NestedModule, Linear
-
-
 def test_inspect_model_with_dict(model_with_dict):
     """Test inspecting a model with dictionary attributes."""
     inspector = ModuleInspector()
@@ -278,15 +269,15 @@ def test_inspect_recursion_depth_limit():
     # Create a deeply nested structure
     deep_dict = {}
     current = deep_dict
-    for _ in range(10):  # More than max_recursion_depth
+    for _ in range(10):
         current["nested"] = {"module": nn.Linear(10, 10)}
         current = current["nested"]
 
-    inspector = ModuleInspector(max_depth=4)
+    inspector = ModuleInspector()
     inspector.inspect(deep_dict)
 
     # Should still work without infinite recursion
-    assert len(inspector._module_info) == 3  # dict is not included as it is not a nn.Module
+    assert len(inspector._module_info) == 5  # dict is not included as it is not a nn.Module
 
 
 def test_inspect_duplicate_objects():
@@ -343,6 +334,96 @@ def test_get_modules_with_dict_model(model_with_dict):
     assert all(module.forward_called for module in executed_modules)
 
 
+class _Block(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mlp = nn.Linear(8, 8)
+
+    def forward(self, x):
+        return self.mlp(x)
+
+
+class _EncoderWithHead(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layers = nn.ModuleList([_Block(), _Block()])
+        self.head = nn.Linear(8, 8)
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return self.head(x)
+
+
+def _executed_encoder_inspector(depth: int = 0) -> ModuleInspector:
+    model = _EncoderWithHead()
+    inspector = ModuleInspector(depth=depth)
+    inspector.inspect(model)
+    model(torch.randn(1, 8))
+    return inspector
+
+
+def test_depth_counts_executed_modules_only():
+    """`nn.ModuleList` is never executed, so it does not add a depth level."""
+    depths = {m.object_path: m.depth for m in _executed_encoder_inspector()._module_info.values() if m.forward_called}
+
+    assert depths == {"": 0, ".layers.0": 1, ".layers.1": 1, ".head": 1, ".layers.0.mlp": 2, ".layers.1.mlp": 2}
+
+
+@pytest.mark.parametrize(
+    ("depth", "expected"),
+    [
+        (0, [""]),
+        (1, [".layers.0", ".layers.1", ".head"]),
+        (2, [".layers.0.mlp", ".layers.1.mlp", ".head"]),
+        (5, [".layers.0.mlp", ".layers.1.mlp", ".head"]),
+    ],
+)
+def test_get_modules_keeps_shallower_leaves(depth, expected):
+    """`head` has no executed children, so it stays in every cut below its own depth."""
+    modules = _executed_encoder_inspector(depth).get_modules()
+
+    assert sorted(m.object_path for m in modules) == sorted(expected)
+
+
+def test_get_modules_can_be_called_twice():
+    inspector = _executed_encoder_inspector(depth=2)
+    first = inspector.get_modules()
+
+    assert inspector.get_modules() == first
+
+
+@pytest.mark.parametrize(
+    ("depth", "expected"),
+    [
+        (0, [".encoder", ".decoder"]),
+        (1, [".encoder.mlp", ".decoder.mlp"]),
+    ],
+)
+def test_get_modules_from_non_module_root(depth, expected):
+    """Executed members of a plain object are at depth 0."""
+
+    class Pipeline:
+        def __init__(self):
+            self.encoder = _Block()
+            self.decoder = _Block()
+
+        def __call__(self, x):
+            return self.decoder(self.encoder(x))
+
+    pipeline = Pipeline()
+    inspector = ModuleInspector(depth=depth)
+    inspector.inspect(pipeline)
+    pipeline(torch.randn(1, 8))
+
+    assert [m.object_path for m in inspector.get_modules()] == expected
+
+
+def test_negative_depth():
+    with pytest.raises(ValueError):
+        ModuleInspector(depth=-1)
+
+
 def test_wrap_forward_methods(simple_model):
     """Test wrapping forward methods."""
     inspector = ModuleInspector()
@@ -554,7 +635,7 @@ def test_get_modules_with_parakeet_rnnt_like_model(parakeet_rnnt_like_model):
 def test_get_modules_with_parakeet_rnnt_like_model_depth_1(parakeet_rnnt_like_model):
     """Test getting executed modules from ParakeetRNNTLikeModel."""
 
-    inspector = ModuleInspector(min_depth=1)
+    inspector = ModuleInspector(depth=1)
     inspector.inspect(parakeet_rnnt_like_model)
 
     parakeet_rnnt_like_model(torch.randn(1, 10))
@@ -617,7 +698,7 @@ def run(self, x):
 def test_second_references_to_module_custom_class():
     model = SecondReferencesToModuleCustomClass()
 
-    inspector = ModuleInspector(min_depth=1)
+    inspector = ModuleInspector(depth=1)
     inspector.inspect(model)
 
     model.run(torch.randn(1, 10))
@@ -649,7 +730,7 @@ def test_second_references_to_module_in_module():
     assert "unet2" not in set(vars(model).keys())
     assert "unet2_unused" not in set(vars(model).keys())
 
-    inspector = ModuleInspector(min_depth=1)
+    inspector = ModuleInspector(depth=1)
     inspector.inspect(model)
 
     model(torch.randn(1, 10))
diff --git a/tests/unit/torch/inspecting/test_wrapping.py b/tests/unit/torch/inspecting/test_wrapping.py
index 4ad798ab..d29d6968 100644
--- a/tests/unit/torch/inspecting/test_wrapping.py
+++ b/tests/unit/torch/inspecting/test_wrapping.py
@@ -451,7 +451,7 @@ def test_tune_reaches_every_module_with_a_repeated_name(tmp_path, monkeypatch):
     monkeypatch.setenv("TORCH_TWEAK_CACHE_DIR", str(tmp_path))
     model = _EncoderDecoder().eval()
     sample = torch.randn(8)
-    modules_info = inspect(model, sample, number_of_iterations=1, warmup_iterations=1, min_depth=3)
+    modules_info = inspect(model, sample, number_of_iterations=1, warmup_iterations=1, depth=2)
     strategy = FirstWinsStrategy(backends=[TorchEagerBackend()])
 
     wrap(model, modules_info.get_modules(), strategy=strategy)
diff --git a/torch_tweak/torch/inspecting/inspecting.py b/torch_tweak/torch/inspecting/inspecting.py
index 425d8326..8773a4ae 100644
--- a/torch_tweak/torch/inspecting/inspecting.py
+++ b/torch_tweak/torch/inspecting/inspecting.py
@@ -31,26 +31,30 @@ def inspect(
     inference_function: Callable | None = None,
     number_of_iterations: int = DEFAULT_INSPECT_ITERATIONS,
     warmup_iterations: int = DEFAULT_WARMUP_ITERATIONS,
-    min_depth: int = 0,
-    max_depth: int = 5,
+    depth: int = 0,
 ) -> InspectedModulesInfo:
     """Inspect provided callable object searching for nn.Module members executed as part of forward pass.
 
+    The depth of a module is the number of executed modules above it. Depth 0 holds the outermost
+    executed modules: `obj` itself when it is an executed module, otherwise its executed members.
+
     Args:
         obj: Callable object to inspect.
         dataset: List of tuples with batch size and input.
         inference_function: Custom inference function to use for inspection, obj is used by default.
         number_of_iterations: Number of iterations to run for inference.
         warmup_iterations: Number of iterations to run for warmup.
-        min_depth: Minimum depth of the modules to inspect, if root level modules is not working, try to increase this value
-        max_depth: Maximum depth of the modules to inspect
+        depth: Depth at which the model is cut into modules. Executed modules at this depth are returned,
+            together with shallower executed modules that have no executed children.
+            Increase it when the outermost modules cannot be tuned.
+
     Returns:
         InspectedModulesInfo object.
     """
     setup_logging(format_string=LOG_FORMAT)
 
     logger.info("Inspecting object searching for executed nn.Module members.")
-    model_inspector = ModuleInspector(min_depth=min_depth, max_depth=max_depth)
+    model_inspector = ModuleInspector(depth=depth)
     model_inspector.inspect(obj)
 
     # If no inference function is provided, use the obj
diff --git a/torch_tweak/torch/inspecting/module_info.py b/torch_tweak/torch/inspecting/module_info.py
index 83e61958..ee706e03 100644
--- a/torch_tweak/torch/inspecting/module_info.py
+++ b/torch_tweak/torch/inspecting/module_info.py
@@ -24,7 +24,20 @@ class ModuleInfo:
     input_types: list[dict[str, Any]] = field(default_factory=list)  # unused, even in tests
     output_types: list[dict[str, Any]] = field(default_factory=list)  # used in tests only
     object_path: str | None = None
-    depth: int = 0
+
+    @property
+    def executed_parent(self) -> Optional["ModuleInfo"]:
+        """The closest ancestor in the object tree whose forward was called."""
+        parent = self.parent
+        while parent is not None and not parent.forward_called:
+            parent = parent.parent
+        return parent
+
+    @property
+    def depth(self) -> int:
+        """The number of executed ancestors in the object tree."""
+        parent = self.executed_parent
+        return 0 if parent is None else parent.depth + 1
 
     @property
     def module_type(self) -> type:
diff --git a/torch_tweak/torch/inspecting/module_inspector.py b/torch_tweak/torch/inspecting/module_inspector.py
index da975163..330ec074 100644
--- a/torch_tweak/torch/inspecting/module_inspector.py
+++ b/torch_tweak/torch/inspecting/module_inspector.py
@@ -20,6 +20,10 @@
 dot notation, even for dictionaries and lists. (e.g., '.list.0.layers`,
 '.dict.key.layers`, `.encoder.layers.0.self_attn.q_proj.weight`).
 
+The depth of a module is the number of executed modules above it in the object
+tree, see `ModuleInfo.depth`. Containers such as lists, dicts, `nn.ModuleList`,
+and modules whose forward never ran do not count.
+
 Environment variables that could be used:
 
     TORCH_TWEAK_INSPECT_DEBUG:
@@ -34,7 +38,7 @@
 
 import time
 from collections.abc import Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from functools import cache, wraps
 from logging import getLogger
 from typing import Any
@@ -52,7 +56,7 @@
 
 logger = getLogger(__name__)
 
-DEFAULT_MAX_RECURSION_DEPTH = 5
+EXTRA_PATH_SEGMENTS = 4
 
 DEFAULT_INSPECT_DEBUG = False
 
@@ -63,7 +67,7 @@
 class InspectContext:
     """Context for inspecting modules."""
 
-    depth: int = 0
+    path_depth: int = 0
     object_path: str = ""
     module_parent: ModuleInfo | None = None
 
@@ -87,36 +91,40 @@ def create_module_info(self, module: torch.nn.Module | list | dict | Any) -> Mod
             module=module,
             parent=self.module_parent,
             object_path=self.object_path,
-            depth=self.depth,
         )
 
     def next(self, name: str = "", parent: ModuleInfo | None = None) -> "InspectContext":
         """Get the next inspect context - increment depth and add name to object path."""
         return InspectContext(
-            depth=self.depth + 1,
+            path_depth=self.path_depth + 1,
             object_path=f"{self.object_path}.{name}" if name else self.object_path,
             module_parent=parent,
         )
 
     def clone(self, **kwargs) -> "InspectContext":
         """Get the inspect context, with optional changes."""
-        return InspectContext(
-            depth=kwargs.get("depth", self.depth),
-            object_path=kwargs.get("object_path", self.object_path),
-            module_parent=kwargs.get("module_parent", self.module_parent),
-        )
+        return replace(self, **kwargs)
 
 
 class ModuleInspector:
     """Class for inspecting PyTorch modules and tracking their execution."""
 
-    def __init__(self, min_depth: int = 0, max_depth: int = 5):
-        """Initialize the module inspector."""
+    def __init__(self, depth: int = 0):
+        """Initialize the module inspector.
+
+        Args:
+            depth: Depth at which `get_modules` cuts the model.
+        """
+        if depth < 0:
+            raise ValueError(f"Expected depth >= 0, got {depth}")
         self._module_info: dict[Any, ModuleInfo] = {}
         self._original_forward: dict[Any, Any] = {}
         self._inspected_objects: set = set()  # Track objects that have been inspected
-        self._min_recursion_depth = min_depth
-        self._max_recursion_depth = max_depth
+        self._depth = depth
+        # Depth is known only after execution, so traversal is bounded by the object path.
+        # One depth level allows two path segments: the module and a container, as in `layers.0`.
+        # EXTRA_PATH_SEGMENTS covers further containers and modules that never run.
+        self._max_path_depth = 2 * (depth + 1) + EXTRA_PATH_SEGMENTS
         self._execution_parent_module: list[ModuleInfo] = []  # used for logging the execution path only
         if _get_inspect_debug():
             self._debug_inspecting_functions()
@@ -141,40 +149,19 @@ def reset(self) -> None:
         self._module_info.clear()
         self._original_forward.clear()
 
-    # fixme: this is wrong if depth is smaller than _min_recursion_depth - some branches are skipped, without even warning
-    # in case of min_depth=0, it will skip all modules except the root one (or in unusual cases all roots)
     def get_modules(self) -> list[ModuleInfo]:
-        """Get list of top-level executed modules or their first executed children.
+        """Get executed modules that cut the model at `depth`.
+
+        The result holds every executed module at `depth`, plus every shallower executed module
+        without executed children. No returned module is nested in another one.
 
         Returns:
-            List of modules that were executed at the top level or their first executed children
-            if the parent wasn't executed.
+            Modules in inspection order.
         """
-        executed_modules = []
-
-        # removing parents calls
-        for module_info in self._module_info.values():
-            if module_info.depth < self._min_recursion_depth:
-                module_info.forward_called = False
-
-        def parent_called(module_info: ModuleInfo) -> bool:
-            if module_info.parent is None:
-                return False
-
-            if module_info.parent.forward_called:
-                return True
-
-            return parent_called(module_info.parent)
-
-        for module_info in self._module_info.values():
-            if (
-                module_info.depth >= self._min_recursion_depth
-                and module_info.forward_called
-                and not parent_called(module_info)
-            ):
-                executed_modules.append(module_info)
-
-        return executed_modules
+        depth = self._depth
+        executed = [m for m in self._module_info.values() if m.forward_called]
+        with_executed_children = {id(m.executed_parent) for m in executed if m.executed_parent is not None}
+        return [m for m in executed if m.depth == depth or (m.depth < depth and id(m) not in with_executed_children)]
 
     def _inspect_object(self, obj: Any, *, context: InspectContext) -> None:
         """Inspect an object and its members for PyTorch modules.
@@ -184,7 +171,7 @@ def _inspect_object(self, obj: Any, *, context: InspectContext) -> None:
             context: The context of the inspection
         """
         # Check recursion depth
-        if context.depth > self._max_recursion_depth:
+        if context.path_depth > self._max_path_depth:
             return
 
         # Skip Python built-in objects
@@ -399,8 +386,8 @@ def _wrap_debug_inspect(self, method: Callable, collection: bool = False) -> Cal
         @wraps(method)
         def inspect_debug(*args, **kwargs):
             context = kwargs["context"]
-            indent = "│" * context.depth
-            logger.debug("%s%s(%d)%s for %s", indent, "├", context.depth, method.__name__, context.object_path)
+            indent = "│" * context.path_depth
+            logger.debug("%s%s(%d)%s for %s", indent, "├", context.path_depth, method.__name__, context.object_path)
             try:
                 return method(*args, **kwargs)
             finally:

From 7fccdc2538a4fde85c187bf8513b6a4e54e9d04b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?=
 <lukasz.slusarczyk@intel.com>
Date: Fri, 23 Oct 2026 16:12:09 +0200
Subject: [PATCH 2/4] docs(examples): correct a typo in the E5Large depth note
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Łukasz Ślusarczyk <lukasz.slusarczyk@intel.com>
---
 examples/E5Large/e5large/tune.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/E5Large/e5large/tune.py b/examples/E5Large/e5large/tune.py
index c39e869f..32f92bff 100644
--- a/examples/E5Large/e5large/tune.py
+++ b/examples/E5Large/e5large/tune.py
@@ -80,7 +80,7 @@ def call_wrapper(*args, **kwargs):
             device="xpu",
         )
 
-    # NOTE: without depth=2, inspector finds a wrapper module `pipeline._modules["0"]` and fails with incorrect input in tunning
+    # NOTE: without depth=2, inspector finds a wrapper module `pipeline._modules["0"]` and fails with incorrect input in tuning
     #     with error: RuntimeError: Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment. @ torch/_tensor.py", line 136, in __deepcopy__
     #     during make_batch and deepcopy of arguments with all kwags arguments passed as dict in args[0]
     inspected_modules_info = inspect(model, input_texts, inference_function=call_wrapper, depth=2)

From b7784fa7221f410730c7c2dc3c0f94248a385cbd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?=
 <lukasz.slusarczyk@intel.com>
Date: Tue, 27 Oct 2026 12:20:14 +0100
Subject: [PATCH 3/4] applied Patryk comments
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Łukasz Ślusarczyk <lukasz.slusarczyk@intel.com>
---
 .../torch/inspecting/test_module_inspector.py | 60 ++++++++++++++++---
 .../torch/inspecting/module_inspector.py      |  8 ++-
 2 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/tests/unit/torch/inspecting/test_module_inspector.py b/tests/unit/torch/inspecting/test_module_inspector.py
index 7e390805..0639d460 100644
--- a/tests/unit/torch/inspecting/test_module_inspector.py
+++ b/tests/unit/torch/inspecting/test_module_inspector.py
@@ -264,20 +264,62 @@ def test_inspect_list_with_non_module_items():
     assert len(inspector._module_info) == 1
 
 
-def test_inspect_recursion_depth_limit():
+@pytest.fixture
+def deep_linears():
+    """Linears placed along one nested path."""
+    m1 = nn.Linear(10, 10)
+    m2 = nn.Linear(10, 10)
+    m3 = nn.Linear(10, 10)
+    m4 = nn.Linear(10, 10)
+    m5 = nn.Linear(10, 10)
+    m6 = nn.Linear(10, 10)
+    return m1, m2, m3, m4, m5, m6
+
+
+@pytest.fixture
+def deep_dict(deep_linears):
+    """One nested dict. Comments are `path_depth` of each linear."""
+    m1, m2, m3, m4, m5, m6 = deep_linears
+    return {
+        "nested": {  # path_depth 1
+            "module": m1,  # path_depth 2
+            "nested": {  # path_depth 2
+                "module": m2,  # path_depth 3
+                "nested": {  # path_depth 3
+                    "module": m3,  # path_depth 4
+                    "nested": {  # path_depth 4
+                        "module": m4,  # path_depth 5
+                        "nested": {  # path_depth 5
+                            "module": m5,  # path_depth 6
+                            "nested": {  # path_depth 6
+                                "module": m6,  # path_depth 7
+                            },
+                        },
+                    },
+                },
+            },
+        },
+    }
+
+
+def test_inspect_recursion_depth_limit(deep_dict, deep_linears):
     """Test that recursion depth is properly limited."""
-    # Create a deeply nested structure
-    deep_dict = {}
-    current = deep_dict
-    for _ in range(10):
-        current["nested"] = {"module": nn.Linear(10, 10)}
-        current = current["nested"]
+    m1, m2, m3, m4, m5, _m6 = deep_linears
+    inspector = ModuleInspector()
+    inspector.inspect(deep_dict)
+
+    assert inspector._max_path_depth == 6
+    assert {id(module) for module in inspector._module_info} == {id(m1), id(m2), id(m3), id(m4), id(m5)}
+
 
+def test_get_modules_is_empty_when_module_is_past_max_path_depth(deep_dict, deep_linears):
+    """A module past the calculated path limit is not returned."""
+    m6 = deep_linears[5]
     inspector = ModuleInspector()
     inspector.inspect(deep_dict)
+    m6(torch.randn(1, 10))
 
-    # Should still work without infinite recursion
-    assert len(inspector._module_info) == 5  # dict is not included as it is not a nn.Module
+    assert inspector.get_modules() == []
 
 
 def test_inspect_duplicate_objects():
diff --git a/torch_tweak/torch/inspecting/module_inspector.py b/torch_tweak/torch/inspecting/module_inspector.py
index 330ec074..d56a4dab 100644
--- a/torch_tweak/torch/inspecting/module_inspector.py
+++ b/torch_tweak/torch/inspecting/module_inspector.py
@@ -23,6 +23,8 @@
 The depth of a module is the number of executed modules above it in the object
 tree, see `ModuleInfo.depth`. Containers such as lists, dicts, `nn.ModuleList`,
 and modules whose forward never ran do not count.
+The walk stops at a path depth calculated from `depth`.
+Modules past that path are not registered.
 
 Environment variables that could be used:
 
@@ -158,10 +160,10 @@ def get_modules(self) -> list[ModuleInfo]:
         Returns:
             Modules in inspection order.
         """
-        depth = self._depth
         executed = [m for m in self._module_info.values() if m.forward_called]
-        with_executed_children = {id(m.executed_parent) for m in executed if m.executed_parent is not None}
-        return [m for m in executed if m.depth == depth or (m.depth < depth and id(m) not in with_executed_children)]
+        candidates = [m for m in executed if m.depth <= self._depth]
+        parents = {id(m.executed_parent) for m in candidates if m.executed_parent is not None}
+        return [m for m in candidates if id(m) not in parents]
 
     def _inspect_object(self, obj: Any, *, context: InspectContext) -> None:
         """Inspect an object and its members for PyTorch modules.

From ab9c9460e44d77843b08c95cd1922ffc26b8ceb3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20=C5=9Alusarczyk?=
 <lukasz.slusarczyk@intel.com>
Date: Tue, 27 Oct 2026 13:44:17 +0100
Subject: [PATCH 4/4] feat(inspecting): warn when the walk stops before any
 module
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The path limit can hide every module and inspect still returns an empty list.
warn_if_no_modules=False turns that warning off.

Signed-off-by: Łukasz Ślusarczyk <lukasz.slusarczyk@intel.com>
---
 ChangeLog.md                                  |  2 +
 docs/guides/inspect.md                        | 11 +++++
 .../unit/torch/inspecting/test_inspecting.py  | 49 +++++++++++++++++++
 torch_tweak/torch/inspecting/inspecting.py    | 15 ++++++
 .../torch/inspecting/module_inspector.py      |  3 ++
 5 files changed, 80 insertions(+)

diff --git a/ChangeLog.md b/ChangeLog.md
index 43d3d09c..55ed3b3c 100644
--- a/ChangeLog.md
+++ b/ChangeLog.md
@@ -13,6 +13,8 @@ SPDX-License-Identifier: Apache-2.0
   `depth` keeps shallower executed modules without executed children, so a cut covers the whole model.
   `max_depth` is removed, the traversal limit follows from `depth`.
   A diffusers pipeline that used `min_depth=1` to reach its components now uses `depth=0`.
+  `inspect` warns when the walk stops before any module is registered.
+  `warn_if_no_modules=False` turns that warning off.
 
 ## 0.4.0
 
diff --git a/docs/guides/inspect.md b/docs/guides/inspect.md
index 77399f2f..f88bd259 100644
--- a/docs/guides/inspect.md
+++ b/docs/guides/inspect.md
@@ -43,6 +43,7 @@ modules_info = tt.inspect(
     number_of_iterations=10,  # Iterations for profiling
     warmup_iterations=5,  # Warmup iterations
     depth=0,  # Depth at which the model is cut into modules
+    warn_if_no_modules=True,  # Warn when the walk stops before any module is registered
 )
 ```
 
@@ -165,6 +166,16 @@ modules_info = tt.inspect(model, input_data, depth=0)
 modules_info = tt.inspect(model, input_data, depth=1)
 ```
 
+#### warn_if_no_modules (Default: True)
+
+Warn when the walk stops before any module is registered.
+The message says how far the walk went and to increase `depth`.
+Pass `False` to silence that warning.
+
+```python
+modules_info = tt.inspect(model, input_data, warn_if_no_modules=False)
+```
+
 ## Working with InspectedModulesInfo
 
 The `inspect` function returns an `InspectedModulesInfo` object with several useful methods:
diff --git a/tests/unit/torch/inspecting/test_inspecting.py b/tests/unit/torch/inspecting/test_inspecting.py
index 15350d06..d40382bf 100644
--- a/tests/unit/torch/inspecting/test_inspecting.py
+++ b/tests/unit/torch/inspecting/test_inspecting.py
@@ -357,6 +357,55 @@ def test_get_modules_after_inspection_with_min_execution_percentage(custom_objec
     assert executed_modules[0].execution_count == TEST_NUMBER_OF_ITERATIONS * custom_object.num_iterations
 
 
+@pytest.fixture
+def module_past_default_path_limit():
+    """A linear at path depth 7. Default depth=0 stops at path depth 6."""
+    return {
+        "nested": {
+            "nested": {
+                "nested": {
+                    "nested": {
+                        "nested": {
+                            "nested": {"module": nn.Linear(4, 4)},
+                        },
+                    },
+                },
+            },
+        },
+    }
+
+
+def _echo(tensor):
+    return tensor
+
+
+def test_inspect_warns_when_the_walk_stops_before_any_module(module_past_default_path_limit, capsys):
+    modules_info = inspect(module_past_default_path_limit, torch.zeros(1, 4), inference_function=_echo)
+
+    assert modules_info.get_modules() == []
+    assert (
+        "No modules found. The walk stopped at path depth 6, calculated from depth=0. Increase depth to walk further."
+    ) in capsys.readouterr().out
+
+
+def test_inspect_warning_can_be_disabled(module_past_default_path_limit, capsys):
+    modules_info = inspect(
+        module_past_default_path_limit,
+        torch.zeros(1, 4),
+        inference_function=_echo,
+        warn_if_no_modules=False,
+    )
+
+    assert modules_info.get_modules() == []
+    assert "No modules found" not in capsys.readouterr().out
+
+
+def test_inspect_does_not_warn_when_a_module_is_found(simple_model, sample_dataset, capsys):
+    inspect(simple_model, sample_dataset, number_of_iterations=1, warmup_iterations=1)
+
+    assert "No modules found" not in capsys.readouterr().out
+
+
 def test_get_modules_after_inspection_with_limit(custom_object, sample_dataset):
     """Test getting executed modules with a limit."""
     # When inspecting custom object with various attributes
diff --git a/torch_tweak/torch/inspecting/inspecting.py b/torch_tweak/torch/inspecting/inspecting.py
index 8773a4ae..56eb8171 100644
--- a/torch_tweak/torch/inspecting/inspecting.py
+++ b/torch_tweak/torch/inspecting/inspecting.py
@@ -32,6 +32,7 @@ def inspect(
     number_of_iterations: int = DEFAULT_INSPECT_ITERATIONS,
     warmup_iterations: int = DEFAULT_WARMUP_ITERATIONS,
     depth: int = 0,
+    warn_if_no_modules: bool = True,
 ) -> InspectedModulesInfo:
     """Inspect provided callable object searching for nn.Module members executed as part of forward pass.
 
@@ -47,6 +48,8 @@ def inspect(
         depth: Depth at which the model is cut into modules. Executed modules at this depth are returned,
             together with shallower executed modules that have no executed children.
             Increase it when the outermost modules cannot be tuned.
+        warn_if_no_modules: Warn when the walk stops at the path limit before any module is registered.
+            The path limit is calculated from `depth`. Increase `depth` to walk further.
 
     Returns:
         InspectedModulesInfo object.
@@ -82,6 +85,18 @@ def inspect(
     modules = model_inspector.get_modules()
 
     logger.info("Inspection done. Found %d candidate modules for tuning.", len(modules))
+    if (
+        warn_if_no_modules
+        and not modules
+        and not model_inspector._module_info
+        and model_inspector._stopped_at_path_limit
+    ):
+        logger.warning(
+            "No modules found. The walk stopped at path depth %d, calculated from depth=%d. "
+            "Increase depth to walk further.",
+            model_inspector._max_path_depth,
+            depth,
+        )
 
     inspected_modules_info = InspectedModulesInfo(total_execution_time, number_of_iterations)
     for module in modules:
diff --git a/torch_tweak/torch/inspecting/module_inspector.py b/torch_tweak/torch/inspecting/module_inspector.py
index d56a4dab..7266c499 100644
--- a/torch_tweak/torch/inspecting/module_inspector.py
+++ b/torch_tweak/torch/inspecting/module_inspector.py
@@ -127,6 +127,7 @@ def __init__(self, depth: int = 0):
         # One depth level allows two path segments: the module and a container, as in `layers.0`.
         # EXTRA_PATH_SEGMENTS covers further containers and modules that never run.
         self._max_path_depth = 2 * (depth + 1) + EXTRA_PATH_SEGMENTS
+        self._stopped_at_path_limit = False
         self._execution_parent_module: list[ModuleInfo] = []  # used for logging the execution path only
         if _get_inspect_debug():
             self._debug_inspecting_functions()
@@ -150,6 +151,7 @@ def reset(self) -> None:
 
         self._module_info.clear()
         self._original_forward.clear()
+        self._stopped_at_path_limit = False
 
     def get_modules(self) -> list[ModuleInfo]:
         """Get executed modules that cut the model at `depth`.
@@ -174,6 +176,7 @@ def _inspect_object(self, obj: Any, *, context: InspectContext) -> None:
         """
         # Check recursion depth
         if context.path_depth > self._max_path_depth:
+            self._stopped_at_path_limit = True
             return
 
         # Skip Python built-in objects