Just-in-Time Tuning Guide

Just-in-time tuning enables automatic model tuning without modifying your existing code. You can enable it with an environment variable and run your script - AITune will automatically discover and tune modules during execution.

Overview

Just-in-time tuning provides a zero-code-change approach to model tuning:

• Automatic Discovery: Automatically detects PyTorch modules during execution
• Zero Code Changes: No need to modify your existing scripts
• Hierarchical Tuning: Recursively tunes modules from top to bottom
• Configurable: Fine-tune behavior through environment variables or configuration

Quick Start

Enabling Just-in-Time Tuning

The simplest way to enable JIT tuning is through an environment variable:

$
export AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning
$
python your_script.py

Your script will run with automatic tuning enabled.

Note: Setting the environment variable affects the entire shell session, which may impact other Python processes running in the same shell. We recommend either setting the environment variable immediately before running your script, or using import-based activation inside the script instead.

Alternative: Import-Based Activation

You can also enable just-in-time tuning by adding a single import at the beginning of your script:

1
import aitune.torch.jit.enable  # Enable JIT tuning
2
3
# Your existing code remains unchanged
4
import torch
5
from diffusers import DiffusionPipeline
6
7
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers")
8
pipe.to("cuda")
9
10
# Tuning happens automatically during inference
11
images = pipe("A beautiful landscape")

Using Annotation (Decorator)

For fine-grained control, you can use the @patch_for_jit_tuning decorator on specific functions:

1
from aitune.torch import patch_for_jit_tuning, jit_config
2
import timm
3
import torch
4
import logging
5
6
logging.basicConfig(level=logging.INFO)
7
8
jit_config.min_samples = 2
9
jit_config.batch_axis_required = False
10
11
@patch_for_jit_tuning
12
def create_resnet():
13
    """This function will have JIT tuning enabled."""
14
    return timm.create_model("resnet18", pretrained=False).to("cuda")
15
16
# Your model
17
model = create_resnet().cuda()
18
19
# Tuning happens automatically when the model is called
20
with torch.no_grad():
21
    output = model(torch.randn(1, 3, 224, 224, device="cuda"))
22
    output = model(torch.randn(2, 3, 224, 224, device="cuda"))

This approach allows you to:

• Enable just-in-time tuning for specific functions only
• Keep the rest of your code unchanged

How Just-in-Time Tuning Works

Just-in-time tuning follows this process:

1. Initial Runs: The first few inferences are used to detect model architecture and record input/output shapes until jit_config.min_samples is met

2. Module Discovery: Identifies all PyTorch modules in the execution path

3. Hierarchical Tuning: Attempts to tune modules starting from the top level:

   • If successful, the module is tuned
   • If a graph break is detected or tuning fails, AITune recursively tunes child modules

4. Depth Limiting: Stops at a configurable depth level

Graph of Execution Detection

AITune has two mechanisms to detect different graphs of execution:

• The first one is based on method signature. If inputs for a particular module change, different Graphs are created. Each of them has a separate backend.
• The second one uses the torch dynamo feature to detect graph breaks when a module contains conditional logic based on input data.

Here is an example of a graph break:

1
import torch
2
3
class DynamicModule(torch.nn.Module):
4
5
    def __init__(self):
6
        super().__init__()
7
        self.child_a = torch.nn.Linear(10, 20)
8
        self.child_b = torch.nn.Linear(10, 20)
9
10
    def forward(self, x):
11
        if x.sum() > 0:  # Graph break: conditional on input
12
            return self.child_a(x)
13
        else:
14
            return self.child_b(x)

When AITune detects a graph break, it skips tuning that module and attempts to tune this module’s children.

Tuning Modes

JIT tuning supports two modes that control when tuning is triggered.

Eager Mode (default)

In eager mode, AITune tunes a module automatically after each forward pass as soon as the required number of samples has been collected. This works well for pipelines where every module is called a predictable number of times per step (e.g. a simple classifier or a fixed-step inference loop).

1
from aitune.torch import jit_config
2
from aitune.torch.jit.config import JITMode
3
4
jit_config.mode = JITMode.TUNE_EAGER  # default, explicit assignment not required

Deferred Mode

In deferred mode, AITune records samples during forward passes but does not tune automatically until you mark a safe synchronization point. Calling aitune.torch.jit.tune.deferred() after a full pipeline step enables tuning; the actual tuning then happens on the next normal forward pass.

This mode is intended for pipelines where different modules are called a variable number of times per step — for example, iterative denoising loops in text-to-image (Stable Diffusion, FLUX) or text-to-video models. In such cases, eager mode may attempt to tune a module before all modules in the pipeline have been recorded; deferred mode lets you choose a safe synchronisation point after a full pass while still compiling from inside the usual model flow.

1
import aitune.torch.jit.enable  # or set AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning
2
3
from aitune.torch import jit_config
4
from aitune.torch.jit.config import JITMode
5
from aitune.torch.jit.tune import deferred as jit_deferred
6
from diffusers import DiffusionPipeline
7
8
jit_config.mode = JITMode.TUNE_DEFERRED
9
10
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers")
11
pipe.to("cuda")
12
13
# First full pipeline step — records samples for every module encountered
14
pipe("A beautiful landscape")
15
16
# Mark that deferred tuning may start
17
jit_deferred()
18
19
# This call triggers tuning from the normal pipeline flow
20
pipe("A snowy mountain at sunset")
21
22
# Subsequent calls use the tuned pipeline
23
pipe("A snowy mountain at sunset")

Configuration

The main configuration of the tuning process is in the config object - this is the common configuration between ahead-of-time and just-in-time mode. On top of that, there are settings particular for the just-in-time mode in jit_config.

This example shows how to get those configuration objects:

1
import aitune.torch as ait
2
3
# main config
4
ait.config
5
6
# just in time config
7
ait.jit_config

Environment Variables

Control just-in-time tuning behavior with environment variables:

$
# Enable just-in-time tuning
$
export AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning
$
$
# Set cache directory
$
export AITUNE_JIT_CACHE_DIR=/path/to/cache
$
$
# Run your script
$
python your_script.py

Programmatic Configuration

For more control, configure just-in-time tuning in your code:

1
from aitune.torch import jit_config
2
3
# Minimum samples before tuning
4
jit_config.min_samples = 2
5
6
# Require batch axis detection
7
jit_config.batch_axis_required = True
8
9
# Maximum module depth
10
jit_config.max_depth_level = 2
11
12
# Minimum parameters to consider tuning
13
jit_config.min_parameters = 1000
14
15
# Enable/disable graph break detection
16
jit_config.detect_graph_breaks = True
17
18
# Skip specific module types
19
jit_config.skip_modules = ["BatchNorm2d", "LayerNorm"]
20
21
# Set device for tuning
22
jit_config.device = "cuda"
23
24
# Enable dry run mode
25
jit_config.dry_run = False
26
27
# Configure tune strategy (which backends to try and how)
28
from aitune.torch.backend import TensorRTBackend, TensorRTBackendConfig
29
from aitune.torch.tune_strategy import FirstWinsStrategy
30
jit_config.strategy = FirstWinsStrategy(
31
    backends=[TensorRTBackend(config=TensorRTBackendConfig(use_dynamo=True))],
32
)

Configuration Options

mode

Tuning mode — controls when tuning is triggered.

1
from aitune.torch.jit.config import JITMode
2
jit_config.mode = JITMode.TUNE_EAGER  # Default

Available values:

• JITMode.TUNE_EAGER — tunes automatically after each forward pass once the sample threshold is reached.
• JITMode.TUNE_DEFERRED — collects samples until aitune.torch.jit.tune.deferred() is called, then tunes on the next normal forward pass.
• JITMode.INSPECT — inspect-only mode; no tuning is performed (see JIT Inspect).

min_samples

Minimum number of samples to record before attempting tuning.

1
jit_config.min_samples = 1  # Default: 1

AITune needs at minimum 1 sample to perform tuning. Multiple samples are required to detect that the model has dynamic axes. Based on the data seen, it detects input and output shapes with dynamic dimensions and minimum/maximum shapes which are required by some backends. If you cannot control data feed to the model, you can increase this setting so that AITune has enough samples to detect proper edge shapes.

max_depth_level

Maximum depth of a module in module hierarchy to be considered for tuning.

1
jit_config.max_depth_level = 2  # Default: 1

Example:

• Depth 0: Root module only
• Depth 1: Root and immediate children
• Depth 2: Root, children, and grandchildren

min_parameters

Minimum number of parameters for a module to be considered for tuning.

1
jit_config.min_parameters = 1000  # Default: 0

Why it matters: Small modules may not benefit from tuning and can be skipped to save time.

detect_graph_breaks

Enable graph break detection.

1
jit_config.detect_graph_breaks = True  # Default: False

Why it matters: Graph breaks prevent static optimization. When disabled, AITune may attempt to tune modules with dynamic control flow (which will likely fail).

skip_modules

List of module class names to skip during tuning.

1
jit_config.skip_modules = ["BatchNorm2d", "LayerNorm", "Dropout"]

Why it matters: Some modules (like normalization layers) typically don’t benefit from tuning.

cache_dir

Directory for per-run build artifacts and logs.

1
from pathlib import Path
2
jit_config.cache_dir = Path("/path/to/cache")

Note: JIT mode does not reuse this directory as a tuned checkpoint cache. A new Python process starts tuning from scratch.

strategy

Tune strategy to use during tuning. Decides which backends to try and how to choose between them.

1
from aitune.torch.backend import (
2
    TensorRTBackend,
3
    TensorRTBackendConfig,
4
    TorchInductorJitBackend,
5
)
6
from aitune.torch.tune_strategy import FirstWinsStrategy
7
8
jit_config.strategy = FirstWinsStrategy(
9
    backends=[
10
        TensorRTBackend(config=TensorRTBackendConfig(use_dynamo=True)),
11
        TorchInductorJitBackend(),
12
    ],
13
)

Accepts any TuneStrategy (e.g. FirstWinsStrategy, MaxThroughputStrategy, OneBackendStrategy). Leave it as None (default) to use a FirstWinsStrategy over TensorRT (with and without dynamo) and TorchInductorJitBackend.

This setting is common for all tuned modules.

patch_exclude

Extra package prefixes or fully qualified class names the JIT patcher will not intercept, in addition to built-in defaults.

1
jit_config.patch_exclude = ("my_package.MyCustomLayer", "third_party.SpecialModule")

Why it matters: Some third-party modules may conflict with JIT patching. Add their class names or package prefixes here to prevent interception beyond the built-in exclusion list.

Limitations and Considerations

1. No Persistent Checkpoint Reuse

Just-in-time tuning writes build artifacts and logs to jit_config.cache_dir, but it does not reuse tuned results across runs. Each time you start a new Python interpreter, tuning starts from scratch.

2. No Benchmarking

Dynamic axes are detected; however, they cannot be matched against real batch sizes. This is due to missing explicit data source, and hence, AITune cannot control batch size. Without this information, it cannot extrapolate batches to any size, which is required by benchmarking functionality.

3. Dynamic Shape Detection Requires Multiple Samples

The default min_samples is 1, which is enough to attempt tuning. If you need detection of dynamic axes and min/max shape, feed multiple input shapes through the model and raise min_samples so AITune records those shapes before tuning.

Just-in-Time Tuning vs Ahead-of-Time Tuning

The following table summarizes the difference between those two modes:

Debugging Just-in-Time Tuning

Just-in-time tuning does not require modification of the original python script. To assist or debug the tuning process, there are several features that might be helpful.

Enable Logging

Make sure your logging level is at least INFO. If tuning happens you should be able to see appropriate log entries.

1
import logging
2
logging.basicConfig(level=logging.INFO)

Dry Run Mode

To test which modules are captured without an actual tuning, change the following config option:

1
from aitune.torch import jit_config
2
jit_config.dry_run = True

It will log what is about to happen in a real scenario.

Model Hierarchy

If you would like to see the hierarchy of a model AITune discovered, you can use the following code:

1
from aitune.torch import PatchedModule
2
3
PatchedModule.print_hierarchy()

Here is an example from ResNet:

JIT Tuning Hierarchy:
├─ ResNet 📊11.7M level=0🪜 state=tuned🎯 (TensorRTBackend) call_count=4
  ├─ Conv2d 📊9.4K level=1🪜 state=detached☑️ call_count=4
  ├─ BatchNorm2d 📊128 level=1🪜 state=detached☑️ call_count=4
  ├─ BasicBlock 📊74.0K level=1🪜 state=detached☑️ call_count=4
    ├─ Conv2d 📊36.9K level=2🪜 state=detached☑️ call_count=4
    ├─ BatchNorm2d 📊128 level=2🪜 state=detached☑️ call_count=4
    ├─ Conv2d 📊36.9K level=2🪜 state=detached☑️ call_count=4
    ├─ BatchNorm2d 📊128 level=2🪜 state=detached☑️ call_count=4
... rest of the hierarchy is abbreviated ...

The hierarchy output is a tree view of the model modules that AITune discovered while tracing. Each line represents a module instance, with indentation showing parent-child relationships. The markers provide context about how AITune treats each module:

• Module name (e.g., ResNet, Conv2d) identifies the layer type or submodule.
• 📊 shows the parameter count for that module (e.g., 11.7M, 9.4K).
• level is the depth in the module tree (root is 0).
• state indicates how AITune handled the module during tuning (for example tuned or detached).
• Backend in parentheses (e.g., TensorRTBackend) is shown for tuned modules.
• call_count is the number of times the module was observed during collection.

Internally a PatchedModule can be in the following states:

• INIT: ”⏳” before the first forward call; hierarchy is not fully resolved yet.
• RECORDING: ”🔴” after the first forward call; hierarchy resolved and collecting samples.
• TUNED: ”🎯” tuning succeeded; the module forwards through a tuned backend.
• EAGER: “⚠️” tuning failed or was not possible; module falls back to the original unmodified model.
• SKIPPED: ”🚫” explicitly skipped (e.g., in skip_modules) and not tuned.
• DETACHED: “☑️” detached because a parent module was tuned, so children are unpatched.

Tuning History

To investigate what may have failed in the process, you can see the history of just-in-time tuning:

1
from aitune.torch import PatchedModule
2
3
PatchedModule.print_history()

Here is abbreviated example from ResNet

 'New top module: ResNet 📊11.7M level=0🪜 state=init⏳ call_count=1',
 ...
 'New child module: Linear 📊513.0K level=1🪜 state=init⏳ call_count=1',
 'No graph breaks in ResNet 📊11.7M. Checking took 4.13s',
 'Tuning ResNet 📊11.7M took 8.60s',
 'Unpatching child module: Linear 📊513.0K level=1🪜 state=detached☑️ call_count=4',
 ...

Basically, history is a list of steps just-in-time tuning took during the process.

Common Issues and Solutions

Issue: Modules Not Being Tuned

Possible causes:

• Not enough samples (min_samples not met)
• Module too small (min_parameters threshold)
• Module in skip_modules list
• Graph breaks detected
• In some environments the export AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning does not start just-in-time tuning. This is a known issue. If it happens, try using a decorator or import to start tuning.

Solution:

1
from aitune.torch import jit_config
2
import logging
3
4
logging.basicConfig(level=logging.DEBUG)
5
jit_config.min_samples = 2
6
jit_config.min_parameters = 0

Issue: Strange Errors or Recompilations

Check if the following scenario has happened:

• AITune got enough samples, it tuned a module
• AITune got another sample but with larger min/max shapes

In such a case - backend was tuned for shapes say (1, 10) but it may later get data with shapes (2, 20) - which are out of bound. This may result in backend failure depending on the technology used or a recompilation (e.g. TorchInductor backend).

Next Steps

• Learn about just-in-time Inspect for detailed module analysis
• Learn about AOT Tuning as an alternative approach
• Explore Backend Configuration
• Review Deployment Guide for production use