Ahead-of-Time Tuning Guide

Ahead-of-time tuning is a mode where you explicitly control which modules to tune. This method provides precise control over the tuning process and is recommended for production environments.

Overview

Ahead-of-time tuning follows a four-step workflow:

1. Inspect: Analyze your model or pipeline to identify tuneable modules
2. Wrap: Wrap selected modules for tuning
3. Tune: Execute the tuning process across different backends
4. Persist: Save and load tuned models for later deployment

This approach offers several advantages:

• Control: Explicitly choose which modules to tune, pick strategies and backends, and mix different technologies
• Performance: Benchmark and select optimal configurations
• Speed: Save the tuned model to a deployable artifact to be loaded on the production environment
• Reproducibility: Deterministic tuning results

Quick Start

Here’s a complete example using Stable Diffusion:

1
import aitune.torch as ait
2
from diffusers import DiffusionPipeline
3
4
# Initialize pipeline
5
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers")
6
pipe.to("cuda")
7
8
# Prepare input data
9
input_data = [{"prompt": "A beautiful landscape with mountains and a lake"}]
10
11
# Step 1: Inspect pipeline to discover modules
12
modules_info = ait.inspect(pipe, input_data)
13
14
# Display discovered modules
15
modules_info.describe()
16
17
# Step 2: Wrap modules for tuning
18
modules = modules_info.get_modules()
19
pipe = ait.wrap(pipe, modules)
20
21
# Step 3: Tune the pipeline
22
ait.tune(pipe, input_data)
23
24
# Step 4: Save the tuned pipeline
25
ait.save(pipe, "tuned_pipe.ait")
26
27
# Use the tuned pipeline
28
images = pipe(["A beautiful landscape with mountains and a lake"])

Detailed Workflow

1. Inspection Phase

The inspect function analyzes your model or pipeline to identify PyTorch modules that can be tuned. For a detailed guide on inspection, see the AOT Inspect Guide.

2. Wrapping Phase

Given the list of modules from the previous step, you can wrap them for tuning. Under the hood, each torch.nn.Module is wrapped (imagine a proxy object) with AITune Module which intercepts all forward calls to get data, tune the module and serve the tuned version.

The following line shows how to wrap modules.

1
model = ait.wrap(model, modules)

You can also specify tuning strategies during wrapping:

1
import aitune.torch as ait
2
3
strategy = ait.OneBackendStrategy(backend=ait.backend.TensorRTBackend())
4
model = ait.wrap(model, modules, strategy=strategy)

If you would like to have more control over picking the modules, you can manually wrap torch.nn.Module. When wrapping, you can specify a strategy for each module separately; i.e., you can combine different strategies backends into one model.

1
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers")
2
pipe.to("cuda")
3
4
pipe.unet = ait.Module(pipe.unet, strategy=strategy_for_unet)
5
pipe.transformer = ait.Module(pipe.transformer, strategy=strategy_for_transformer)

User-provided dynamic shapes

When recorded samples do not cover every shape needed in production, pass an explicit shape contract to ait.Module. Each mapping key is a tensor’s forward parameter path, and each value describes the tensor’s full rank:

• Use an integer for a fixed dimension.
• Use BatchDim for the logical batch dimension.
• Use DynamicDim for any other bounded dynamic dimension.

min and max are inclusive. opt selects the preferred compilation shape and defaults to max when omitted. Configure dynamic_shapes when constructing ait.Module directly; ait.wrap does not accept per-module shape definitions.

1
import aitune.torch as ait
2
3
batch = ait.BatchDim("batch", min=1, opt=2, max=8)
4
height = ait.DynamicDim("spatial", min=224, opt=224, max=512)
5
width = ait.DynamicDim("spatial", min=224, opt=224, max=512)
6
7
model = ait.Module(
8
    model,
9
    dynamic_shapes={"x": (batch, 3, height, width)},
10
)

The integer 3 fixes the channel dimension. height and width are shared because they have the same name, even though they are separate DynamicDim objects. Shared definitions must have the same type and bounds.

Forward input paths are shown in the input_spec sample-metadata table in tuning logs. Use the Semantic Path column as the dynamic_shapes key. Top-level tensors use their forward parameter name, while tensors nested in dictionaries, sequences, dataclasses, or supported custom objects include every key, index, or attribute:

Tensors:
Access Path       Semantic Path          Shape
----------------  ---------------------  ----------------
x                 x                      [1, 3, 224, 224]
options["mask"]   ('options', 'mask')    [1, 224, 224]
items[0]          ('items', 0)           [1, 128]
request.image     ('request', 'image')   [1, 3, 512, 512]

Use a string such as "x" for a top-level path and a tuple such as ("options", "mask"), ("items", 0), or ("request", "image") for a nested path. Inputs omitted from the mapping keep using shapes inferred from recorded samples.

Each dynamic shape definition must match the corresponding input tensor’s rank and fixed dimensions, and its ranges must include all dimension sizes present in the tuning samples. The explicit definitions then determine the ranges used by supported AOT backends.

See the runnable ResNet dynamic-shapes example, which records (1, 3, 224, 224) and runs the tuned model at (2, 3, 256, 256).

3. Tuning Phase

The tune function executes the actual tuning:

1
ait.tune(
2
    func=model,                   # The wrapped callable module or pipeline to tune
3
    dataset=input_data,           # Dataset to use for tuning (list, Dataset, DataLoaderFactory, or Tensor)
4
    batch_sizes=[1, 4, 8],         # Optional: Multiple batch sizes. Defaults to [1, 2]
5
    max_num_batches_per_batch_size=10,  # Max batches per size. Defaults to None (all)
6
    device="cuda",                  # Device for tuning. Defaults to "cuda:0"
7
    dry_run=False,                  # Set True to test without tuning
8
    disable_external_logging=False, # Disable third-party logs
9
    clear_cache=False,              # Clear AITune cache before tuning
10
    ignore_failing_modules=True,    # Keep tuning remaining modules when one fails
11
)

Tuning Parameters

• func: The wrapped callable (model or pipeline)
• dataset: Dataset for tuning. It can be a list of samples, torch.utils.data.Dataset, DataLoaderFactory, Tensor or sequence of tensors, dictionaries, strings
• batch_sizes: List of batch sizes to tune against. If not specified, values [1, 2] will be used
• max_num_batches_per_batch_size: Maximum number of batches per batch size. If None, all batches will be used
• device: Device to use for tuning. Defaults to “cuda:0”
• dry_run: If True, performs a dry run without actual tuning
• disable_external_logging: Disable logging from external libraries
• clear_cache: Clear AITune cache before tuning
• ignore_failing_modules: If True, modules that fail tuning fall back to eager execution and tuning continues

Tuning time depends on the tuned modules’ size, used strategy, and number of backends. Modules are tuned one by one. If a strategy has many backends to pick from, it takes the one that fulfills specific strategy criteria. Each backend is validated against returning proper numeric results (check against NANs and infinity) and output shapes.

Note: If you specify a batch size that is not a power of 2, it will be used to gather samples but the actual search for maximum throughput will round it up to the nearest power of 2.

4. Persistence Phase

Once tuned, you can save your model for later use. This is crucial for production deployments to avoid re-tuning every time:

1
# Save the tuned model/pipeline
2
ait.save(pipe, "tuned_model.ait")

The tuned artifact will be saved in the checkpoints folder. The save function creates several files:

• tuned_model.ait: The compressed checkpoint containing tuned and original weights
• tuned_model_sha256_sums.txt: SHA256 hashes for verification

To do inference, you can load the tuned model/pipeline:

1
# Note: Initializing the original object is required before loading
2
pipe = DiffusionPipeline.from_pretrained(...)
3
pipe = ait.load(pipe, "tuned_model.ait")
4
# pipe is ready for use

Custom Inference Functions

For complex pipelines, you can provide a custom inference function:

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

Inspect and tune with the same workload wrapper

When scalar arguments change module execution or input shapes, put those options in a workload wrapper and use that same wrapper for both inspect() and tune(). This is common for diffusion pipelines where you want to tune for several image sizes, step counts, guidance scales, or sequence-length settings.

The important part is that inspection and tuning must execute the same workload. If you inspect through a wrapper that uses multiple scalar arguments but later call ait.tune(pipe, input_data) directly, tuning records a different execution path and may miss graph variants or shape ranges.

Note for FLUX and Stable Diffusion pipelines: image size, step count, guidance scale, and sequence-length settings are usually scalar keyword arguments on the pipeline call, not tensor samples in the dataset. Put every option you want to tune, such as height, width, num_inference_steps, guidance_scale, or max_sequence_length, inside the workload wrapper and pass that wrapper to both ait.inspect() and ait.tune(). Calling ait.tune(pipe, input_data) directly after inspecting through a wrapper will not exercise those same variants during tuning.

1
pipe = get_pipeline(model_name=model_name)
2
3
sizes = [(256,256), (512,512)]
4
5
def call_wrapper(*args, **kwargs):
6
    for height, width in sizes:
7
        pipe(
8
            *args,
9
            height=height,
10
            width=width,
11
            num_inference_steps=28,
12
            guidance_scale=1.0,
13
            max_sequence_length=512,
14
            **kwargs,
15
        )
16
17
18
input_data = [{"prompt": prompt}]
19
20
# Inspect the workload you intend to tune.
21
modules_info = ait.inspect(
22
    pipe,
23
    input_data,
24
    inference_function=call_wrapper,
25
)
26
27
# Wrap modules selected from that inspection.
28
pipe = ait.wrap(pipe, modules_info.get_modules())
29
30
# Tune through the same wrapper so recorded graphs and scalar options match inspection.
31
ait.tune(call_wrapper, input_data)

With config.strict_mode=True (the default), different non-tensor argument values can create separate graphs. Keep scalar arguments fixed when you want one graph, or exercise each option in the wrapper when those variants should be tuned.

Configuration Options

AITune has configuration for the tuning process, and each backend has its configuration.

Global Configuration

You can configure AITune globally:

1
from aitune.torch import config
2
3
# Set cache directory
4
config.cache_dir = "/path/to/cache"
5
6
# Set minimum samples for tuning
7
config.min_num_samples = 5
8
9
# Set maximum stored samples per graph
10
config.max_num_samples_stored = 100
11
12
# Device to move model after tuning
13
config.device_after_tuning = "cuda"
14
15
# Enable/disable strict mode for input validation
16
config.strict_mode = True
17
18
# Enable/disable HuggingFace integrations
19
config.enable_transformers_integration = True
20
config.enable_diffusers_integration = False

Backend-Specific Configuration

Each backend has its own corresponding configuration:

1
from aitune.torch.backend import TensorRTBackendConfig, TensorRTBackend
2
3
config = TensorRTBackendConfig(
4
    use_cuda_graphs=True,
5
    workspace_size=1 << 30,  # 1GB
6
)
7
backend = TensorRTBackend(config)

See backend-specific documentation:

• TensorRT Backend
• Torch-Inductor JIT Backend
• TorchAO Backend
• Torch TensorRT AOT Backend
• Torch TensorRT JIT Backend

Dry Run Mode

You can run tuning in dry-run mode. It records samples of data, detects batch and dynamic axes, and detects graphs of execution but does not call the actual backend to tune. This allows debugging if everything is working as expected.

The dry-run mode can be turned on with the proper argument:

1
import logging
2
3
# make sure logging if configured
4
logging.basicConfig(level=logging.INFO, force=True)
5
# invoke dry-run tuning
6
ait.tune(pipe, input_data, dry_run=True)

Example output from dry-run

2026-01-26 16:23:44,360 - INFO - ════════════════════════════════════════════════════════════════
2026-01-26 16:23:44,360 - INFO - 🎯 Tuning module: `transformer` (all graphs)
2026-01-26 16:23:44,367 - INFO - ------------------------------------------------------------
2026-01-26 16:23:44,367 - INFO - 🚀 Tuning graph `0` for module `transformer` (DRY RUN):
2026-01-26 16:23:44,368 - INFO -   number of parameters: 2028328000
2026-01-26 16:23:44,368 - INFO -   number of layers: 6
2026-01-26 16:23:44,369 - INFO -   precisions: torch.float16
2026-01-26 16:23:44,369 - INFO -   graph_spec:
2026-01-26 16:23:44,369 - INFO -     input_spec:
 Tensors:
╒═══════════════════════╤═══════════════════════╤══════════════════════════╤═══════════════════╤═══════════════════╤═══════════════╕
│ Access Path           │ Semantic Path         │ Shape                    │ Min Shape         │ Max Shape         │ Dtype         │
╞═══════════════════════╪═══════════════════════╪══════════════════════════╪═══════════════════╪═══════════════════╪═══════════════╡
│ encoder_hidden_states │ encoder_hidden_states │ ['batch0', 333, 4096]    │ [2, 333, 4096]    │ [4, 333, 4096]    │ torch.float16 │
├───────────────────────┼───────────────────────┼──────────────────────────┼───────────────────┼───────────────────┼───────────────┤
│ hidden_states         │ hidden_states         │ ['batch0', 16, 128, 128] │ [2, 16, 128, 128] │ [4, 16, 128, 128] │ torch.float16 │
├───────────────────────┼───────────────────────┼──────────────────────────┼───────────────────┼───────────────────┼───────────────┤
│ pooled_projections    │ pooled_projections    │ ['batch0', 2048]         │ [2, 2048]         │ [4, 2048]         │ torch.float16 │
├───────────────────────┼───────────────────────┼──────────────────────────┼───────────────────┼───────────────────┼───────────────┤
│ timestep              │ timestep              │ ['batch0']               │ [2]               │ [4]               │ torch.float32 │
╘═══════════════════════╧═══════════════════════╧══════════════════════════╧═══════════════════╧═══════════════════╧═══════════════╛
Other:
╒════════════════════════╤════════════════════════╤═════════╕
│ Access Path            │ Semantic Path          │ Value   │
╞════════════════════════╪════════════════════════╪═════════╡
│ joint_attention_kwargs │ joint_attention_kwargs │ None    │
├────────────────────────┼────────────────────────┼─────────┤
│ return_dict            │ return_dict            │ False   │
╘════════════════════════╧════════════════════════╧═════════╛
2026-01-26 16:23:44,370 - INFO -     output_spec:
 Tensors:
╒═══════════════╤═════════════════╤══════════════════════════╤═══════════════════╤═══════════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape                    │ Min Shape         │ Max Shape         │ Dtype         │
╞═══════════════╪═════════════════╪══════════════════════════╪═══════════════════╪═══════════════════╪═══════════════╡
│ output[0]     │ 0               │ ['batch0', 16, 128, 128] │ [2, 16, 128, 128] │ [4, 16, 128, 128] │ torch.float16 │
╘═══════════════╧═════════════════╧══════════════════════════╧═══════════════════╧═══════════════════╧═══════════════╛
2026-01-26 16:23:44,370 - INFO -   num samples: 1
2026-01-26 16:23:44,370 - INFO -   device: cuda:0
2026-01-26 16:23:44,370 - INFO -   cache_dir: /home/pbazan/.cache/aitune/transformer/0
2026-01-26 16:23:44,371 - INFO -   strategy:
2026-01-26 16:23:44,371 - INFO -     name: First Wins Strategy
2026-01-26 16:23:44,371 - INFO -     description: evaluate backends in order, return first working backend
2026-01-26 16:23:44,371 - INFO -     backends:
2026-01-26 16:23:44,371 - INFO -       TensorRTBackend(quantization_config=None)
2026-01-26 16:23:44,371 - INFO -       TorchInductorJitBackend()
2026-01-26 16:23:44,372 - INFO -       TorchEagerBackend()
2026-01-26 16:23:44,372 - INFO - ✅ Tuning module: `transformer` (all graphs) completed.

Next Steps

• Learn about Ahead-of-time Inspect for detailed module analysis
• Learn about Just-in-Time (JIT) Tuning as an alternative approach
• Explore Backend Configuration
• Review Tune Strategies
• See Deployment Guide