Ahead-of-time Inspect Guide

The inspect function is a powerful tool for analyzing PyTorch models and pipelines. It helps you understand model structure, identify tuneable modules, and gather execution statistics. It can also be a first step to pick modules before ahead-of-time tuning.

Overview

Inspection provides:

• Module Discovery: Automatically finds all PyTorch modules in your model or pipeline
• Execution Tracking: Identifies which modules are executed during inference
• Performance Profiling: Measures execution time for each module
• Input and output data types: records model input and its results

Basic Usage

1
import aitune.torch as ait
2
3
# Inspect a model
4
modules_info = ait.inspect(model, input_data)
5
6
# Display results
7
modules_info.describe()

Inspection Parameters

Complete Signature

1
modules_info = ait.inspect(
2
    obj=model,                      # Model or pipeline to inspect
3
    dataset=input_data,             # Representative input data
4
    inference_function=None,        # Optional custom inference function
5
    number_of_iterations=10,        # Iterations for profiling
6
    warmup_iterations=5,            # Warmup iterations
7
    min_depth=0,                    # Minimum depth for module search
8
    max_depth=5,                    # Maximum depth for module search
9
)

Parameter Details

obj (Required)

The object to inspect. Can be:

• torch.nn.Module: Any PyTorch module
• Callable: Any callable containing PyTorch modules (e.g., HuggingFace pipelines)

1
# PyTorch module
2
model = torchvision.models.resnet50()
3
modules_info = ait.inspect(model, input_data)
4
5
# Diffusion pipeline
6
from diffusers import StableDiffusionPipeline
7
pipe = StableDiffusionPipeline.from_pretrained("...")
8
modules_info = ait.inspect(pipe, input_data)

dataset (Required)

This is the source of data for your model. It can be:

• torch.Tensor
• sequence of strings, tensors, or dictionaries. The collate function is used to stack samples into batches.
• torch.utils.data.Dataset

1
# Single tensor with batch dimension
2
input_data = torch.randn(1, 3, 224, 224)
3
4
# List of strings
5
input_data = [
6
    "prompt1", "prompt2"
7
]
8
9
# List of tensors
10
input_data = [torch.randn(3, 224, 224) for _ in range(4)]  # 4 random images, notice there is no batch dimension
11
12
# List of dictionaries
13
input_data = [
14
    {"input_ids": torch.tensor([...]), "attention_mask": torch.tensor([...])},
15
]

For customization you can use ait.DataLoaderFactory class.

inference_function (Optional)

Custom function for running inference. Useful for complex execution logic:

1
def custom_inference(prompt, steps=50):
2
    """Function forces width, height and number of steps."""
3
    return pipe(
4
        prompt=prompt,
5
        width=1024,
6
        height=1024,
7
        num_inference_steps=steps,
8
    )
9
10
modules_info = ait.inspect(
11
    pipe,
12
    input_data=[{"prompt": "test"}],
13
    inference_function=custom_inference
14
)

If the custom function uses scalar arguments that should also be tuned, use the same function later in ait.tune() after wrapping the inspected modules. See Inspect and tune with the same workload wrapper.

number_of_iterations (Default: 10)

Number of iterations for profiling execution time:

1
# Quick inspection
2
modules_info = ait.inspect(model, input_data, number_of_iterations=3)

warmup_iterations (Default: 5)

Warmup iterations before profiling to stabilize measurements:

1
modules_info = ait.inspect(
2
    model,
3
    input_data,
4
    warmup_iterations=3  # Speed up inspection by reducing warmup iterations
5
)

min_depth (Default: 0)

Minimum depth level for module discovery. Increase if root-level modules don’t work:

1
# Start from root level (default)
2
modules_info = ait.inspect(model, input_data, min_depth=0)
3
4
# Skip root, inspect children
5
modules_info = ait.inspect(model, input_data, min_depth=1)
6
modules_info = ait.inspect(model, input_data, min_depth=2)

max_depth (Default: 5)

If a nested (child) module has a larger depth than max_depth it will be skipped from inspection.

Both min_depth and max_depth narrow the inspection search to a reasonable range.

Working with InspectedModulesInfo

The inspect function returns an InspectedModulesInfo object with several useful methods:

describe()

Display comprehensive information about discovered modules:

1
modules_info = ait.inspect(model, input_data)
2
modules_info.describe()

Output example for stable-diffusion-3-medium-diffusers:

Module Execution Summary:
==========================================================================================================================================
    Module Name        Calls    Total Time (s)    Avg Time (s)    % of Total    # of params      # of layers           precisions
------------------------------------------------------------------------------------------------------------------------------------------
decoder                      1           0.1760           0.1760        2.80%         49545475                6              torch.float16
text_encoder                 2           0.0057           0.0029        0.09%        123650304                2              torch.float16
text_encoder_2               2           0.0124           0.0062        0.20%        694659840                2              torch.float16
text_encoder_3               2           0.0783           0.0391        1.24%       4762310656                2  torch.float16, torch.float32
transformer                 28           5.9934           0.2140       95.17%       2028328000                6              torch.float16
------------------------------------------------------------------------------------------------------------------------------------------
Total execution time: 6.297374 seconds
Number of batch iterations: 1
==========================================================================================================================================

get_modules()

This function allows getting found modules. Its basic usage returns all of them:

1
# Get all executed modules
2
modules = modules_info.get_modules()

You can place additional criteria:

• min_execution_ratio - minimum ratio of total execution time e.g. .9
• limit - maximum number of modules to return, e.g., 5

If those criteria are not sufficient, you can manually filter modules:

1
modules_info = ait.inspect(model, input_data)
2
3
# Get only transformer blocks
4
transformer_modules = [
5
    m for m in modules_info.get_modules()
6
    if "transformer" in m.name.lower()
7
]

Troubleshooting

Issue: No modules found

1
# Solution: Increase min_depth or check object structure
2
modules_info = ait.inspect(model, input_data, min_depth=1)

Summary

Ahead-of-time inspection can be a first step in exploring your model’s structure and performance. It can also be used to select modules for tuning. For details on the tuning workflow, head to the AOT Tuning Guide.