feat(inspecting): measure module depth in executed modules

ChangeLog.md

@@ -5,6 +5,17 @@ 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`.
+  `inspect` warns when the walk stops before any module is registered.
+  `warn_if_no_modules=False` turns that warning off.
+
 ## 0.4.0
 
 This is the first release of **Torch Tweak**. Torch Tweak is focused on Intel GPU (XPU) support.

docs/guides/inspect.md

@@ -42,8 +42,8 @@ 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
+    warn_if_no_modules=True,  # Warn when the walk stops before any module is registered
 )
 ```
 
@@ -133,24 +133,48 @@ 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.
+
+For a model whose `forward` loops over `self.layers` (an `nn.ModuleList`) and then calls `self.head`:
+
+| Depth | Modules |
+|---|---|
+| 0 | `model` |
+| 1 | `layers.0`, `layers.1`, `head` |
+| 2 | `layers.0.mlp`, `layers.1.mlp` |
+
+#### depth (Default: 0)
+
+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`.
+
+Increase it when the outermost modules cannot be tuned:
 
 ```python
-# Start from root level (default)
-modules_info = tt.inspect(model, input_data, min_depth=0)
+# Cut at the outermost executed modules (default)
+modules_info = tt.inspect(model, input_data, depth=0)
 
-# Skip root, inspect children
-modules_info = tt.inspect(model, input_data, min_depth=1)
-modules_info = tt.inspect(model, input_data, min_depth=2)
+# Cut one level deeper
+modules_info = tt.inspect(model, input_data, depth=1)
 ```
 
-#### max_depth (Default: 5)
+#### warn_if_no_modules (Default: True)
 
-If a nested (child) module has a larger depth than `max_depth` it will be skipped from inspection.
+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.
 
-Both `min_depth` and `max_depth` narrow the inspection search to a reasonable range.
+```python
+modules_info = tt.inspect(model, input_data, warn_if_no_modules=False)
+```
 
 ## Working with InspectedModulesInfo
 
@@ -211,8 +235,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

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 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, 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

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)

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)

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)

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