Aitune Tune

AITune Model Tuning (Agent Skill)

Tuning is the process of finding the best backend for a model/pipeline. It iterates backends from highest to lowest performance, stopping at the first that compiles correctly and meets the speedup threshold.

Tuning Modes

Choose the mode based on user requirements — these are parallel paths, not sequential steps.

The workflow below (Phases 1–5) is the AOT path. For JIT (eager or deferred), load how-to-jit-tune.md instead.

Backends

TensorRT Backend Configuration

from aitune.torch.backend import TensorRTBackend, TensorRTBackendConfig

config = TensorRTBackendConfig(quantization_config=ONNXAutoCastConfig(), use_cuda_graphs=True, use_dynamo=True)
backend = TensorRTBackend(config)

Tuning Strategies

Common Issues

Key Source Paths

• aitune/torch/ — Main AOT tuning API (inspect, wrap, tune, save, load)
• aitune/torch/jit/ — JIT tuning implementation and config
• aitune/torch/backend/ — All backend implementations
• aitune/torch/tune_strategy/ — Strategy implementations
• examples/ — Working examples: LLM, ResNet, StableDiffusion, FLUX, Parakeet, ESM2, E5Large

Workflow - Execution Patterns

Prerequisites

• Create and use one directory in current working directory i.e. tuning/<model_name> for the tuning process.
• Find and install dependencies for the model/pipeline.
• Use input data provided by the user, use huggingface examples and as a last resort use the default input data.
• Try to construct the input of batch size 2 for the model/pipeline to be able to detect the batch dimensions.
• Try not to disaggregate the pipeline - do not create submodules glue code - keep the pipeline as one unit.
• To make output concise try block progress bars for pipelines with pipe.set_progress_bar_config(disable=True) or TQDM_DISABLE=1 environment variable etc.
• To get tuning logs try to set AITUNE_CONSOLE_OUTPUT=1 environment variables
• Create a new cache directory for the tuning process and set it as AITUNE_CACHE_DIR environment variable.
• Create a new triton cache directory and set it as TRITON_CACHE_DIR environment variable.
• Do the environment setup checklist before running the tuning script.

Environment Setup Checklist

Before running any tuning script, verify:

# GPU is visible
nvidia-smi

# CUDA available in Python and GPU SM version (needed for fp8 gating)
python -c "import torch; print(torch.cuda.is_available(), torch.version.cuda, 'SM', torch.cuda.get_device_capability())"

# AITune importable
python -c "import aitune.torch as ait; print('ok')"

# TensorRT importable (for TRT backends)
python -c "import tensorrt; print(tensorrt.__version__)"

If any other check fails, stop and report the blocker before attempting tuning.

Phase 1: Inspecting a model

Use the aitune-inspect skill to inspect the model provided by the user and get the names and utilization of the submodules.

Capture the output and use it to populate the Tuning Summary.

Phase 2: Wrapping with a specific backend

Example configurations

Below, there are all backends configurations for the all optimization steps.

from aitune.torch.backend import (
    TensorRTBackend, TensorRTBackendConfig, ONNXAutoCastConfig,
    TorchTensorRTJitBackend, TorchTensorRTAotBackend,
    TorchAOBackend, TorchAOBackendConfig, TorchInductorJitBackend, TorchInductorAotBackend,
)
from aitune.torch.tune_strategy import OneBackendStrategy, MaxThroughputStrategy

# TensorRT fp16 (default production choice)
cfg = TensorRTBackendConfig(quantization_config=ONNXAutoCastConfig(), use_cuda_graphs=True, use_dynamo=True)
strategy = OneBackendStrategy(TensorRTBackend(cfg))

# TensorRT fp32 (fallback if fp16 correctness fails)
cfg = TensorRTBackendConfig()
strategy = OneBackendStrategy(TensorRTBackend(cfg))

# TorchTRT AOT
strategy = OneBackendStrategy(TorchTensorRTAotBackend())

# TorchTRT JIT (torch.compile path)
strategy = OneBackendStrategy(TorchTensorRTJitBackend())

# TorchAO
strategy = OneBackendStrategy(TorchAOBackend(TorchAOBackendConfig(quantization="int8wo")))

# Torch Inductor JIT and AOT
strategy = OneBackendStrategy(TorchInductorAotBackend())
strategy = OneBackendStrategy(TorchInductorJitBackend())


modules = ait.inspect(model, input_data).get_modules()
model = ait.wrap(model, modules, strategy=strategy)
ait.tune(model, input_data)

Depth-Scoped Wrapping

The caller specifies the module scope. Use the appropriate inspect() call:

# Root scope — entire pipeline as one unit
modules_info = ait.inspect(model, input_data, min_depth=0)

# Depth=1 scope — immediate submodules only
modules_info = ait.inspect(model, input_data, min_depth=1)

modules = modules_info.get_modules()
model = ait.wrap(model, modules, strategy=strategy)
ait.tune(model, input_data)

If the caller (e.g. tuning-assistant agent) specifies a backend trial order, follow that. The Default Backend Trial Order below applies when running standalone.

Backend Selection

Use the inspection output to determine trial order. Do not invent an order — use the table below.

Declare the trial order explicitly before starting Phase 3.

Default Backend Trial Order

When no specific requirements are given, try in this order and stop at the first GO result:

1. TensorRTBackend(fp16, dynamo=On/Off) — highest performance
2. TensorRTBackend(fp32, dynamo=On/Off) — if fp16 correctness fails
3. TorchTensorRTAotBackend — good perf, better graph break tolerance
4. TorchTensorRTJitBackend — torch.compile path, most PyTorch-compatible
5. TorchAOBackend — PyTorch-native, no TRT required
6. TorchInductorAotBackend — Torch Inductor AOT backend
7. TorchInductorJitBackend — Torch Inductor JIT backend
8. TorchEagerBackend — fallback baseline, no compilation

If TRT correctness fails with use_dynamo=True, retry with use_dynamo=False before moving to the next backend.

Phase 3 — Optimization Loop (iterate until working backend is found)

This is the core loop. Run it for each backend in the declared trial order. Stop as soon as a backend produces a GO result. If all backends are exhausted, produce a NO-GO report and move to Phase 4.

Step 3a — Tune with the selected backend

Check examples and how-to-aot-tune.md to create the tuning script. Capture the output and use it to populate the Tuning Summary.

Distinguish failure types:

• Compile failure (backend rejects the model): advance to next backend in trial order
• Environment failure (read-only triton cache, missing permissions, writable path issues): fix the environment condition and retry the same backend before advancing — do not count environment failures as backend failures
• Unhandled exception / script crash (OOM, CUDA error, import error): diagnose using common-errors.md before advancing

Step 3b — Check correctness validation

Use the aitune-validate skill to validate the output of the tuned model against the baseline model.

Step 3c — Benchmark

Use the aitune-benchmark skill to benchmark the tuned model.

Loop stopping criteria

After each backend attempt, print a one-line status:

[Backend: TensorRTBackend-fp16] compile=OK  speedup=2.1x  correctness=PASS  → WINNER
[Backend: TorchTRT-AOT]        compile=FAIL (graph break in layer X)        → SKIP

Phase 4 — Save Checkpoint

If a winning backend was found, save the tuned model and verify the file exists:

import json, os, aitune.torch as ait

checkpoint_path = "tuned_model.ait"
ait.save(model, checkpoint_path)

print(json.dumps({
    "checkpoint": checkpoint_path,
    "saved": os.path.exists(checkpoint_path),
    "size_mb": os.path.getsize(checkpoint_path) / 1024**2 if os.path.exists(checkpoint_path) else 0,
}))

Phase 5 — Deployment Readiness Report

Populate from actual captured numbers. Do not estimate or approximate — every value must come from a script output.

## Tuning & Deployment Readiness Report

**Model**: [from inspection]
**Winning backend**: [backend name + precision]
**Backends attempted**: [count and names]
**Analysis date**: [date]


### Performance Results

| Metric          | Baseline (Eager) | Tuned   | Delta    |
|-----------------|------------------|---------|----------|
| Throughput (r/s)| [value]          | [value] | [+X.Xx]  |
| Avg latency (ms)| [value]          | [value] | [-Xms]   |
| p50 latency (ms)| [value]          | [value] | [-Xms]   |
| p95 latency (ms)| [value]          | [value] | [-Xms]   |
| p99 latency (ms)| [value]          | [value] | [-Xms]   |
| GPU memory (MB) | [value]          | [value] | [+-XMB]  |

**Speedup**: [X.Xx]
**Compilation time**: [Xs]


### Correctness Validation

**Max absolute difference**: [value]
**Max relative error**: [value]
**Tolerance threshold**: [atol used]
**Result**: PASS / FAIL


### Backend Trial Log

| Backend | Compile | Speedup | Correctness | Outcome |
|---------|---------|---------|-------------|---------|
| [name]  | OK/FAIL | [X.Xx]  | PASS/FAIL   | WINNER/SKIP |


### Deployment Recommendation

**Status**: GO / CONDITIONAL GO / NO-GO

**Reasoning**: [grounded in the numbers above]

**Conditions** (if CONDITIONAL GO): [specific, actionable list]

**Checkpoint path**: [absolute path to .ait file]

Common Errors and Diagnostics

There are some common errors and diagnostics that you can find in the attached file common-errors.md load if needed.