Inspect for AOT tuning

The inspect function allows you to analyze PyTorch models and pipelines to understand their structure, parameters, and execution flow. It provides detailed insights into model architecture and helps identify tuning opportunities.

1
import aitune.torch as ait
2
import torch.nn as nn
3
4
class SimpleModel(nn.Module):
5
    def __init__(self):
6
        super().__init__()
7
        self.linear = nn.Linear(100, 10)
8
9
    def forward(self, x):
10
        return self.linear(x)
11
12
model = SimpleModel()
13
14
# Inspect the model
15
ait.inspect(model, dataset)

Inspect for JIT tuning

JIT tuning also has a corresponding inspect mode which gathers information about the model/pipeline and allows checking model input and output arguments, hierarchy of the model, etc.

Here is a short snippet how to use it:

1
# required imports
2
import aitune.torch.jit.enable_inspection as inspection
3
4
# your code goes here
5
# ...
6
7
# you can export report to html file
8
inspection.save_report("filename.html", "YOUR_MODEL_NAME")

Tune

The tune function is the core functionality that automatically tunes your PyTorch models and pipelines for optimal inference performance. It supports various backends and automatically selects the best performing configuration.

1
import aitune.torch as ait
2
import torch
3
4
# Define your model
5
model = SimpleModel()
6
7
# Wrap the model
8
model = ait.Module(model)
9
10
# Define inference function
11
def inference_fn(x):
12
    return model(x)
13
14
# Tune the model
15
ait.tune(
16
    func=inference_fn,
17
    dataset=torch.randn(1, 100),
18
)

Save

The save function allows you to persist tuned models for later use. It stores tuned and original module weights together in a single file with a .ait extension. Apart from the checkpoint file, there is also a SHA hash file.

1
# Save the tuned model
2
import aitune.torch as ait
3
ait.save(model, "tuned_model.ait")

Example output:

$
checkpoints/
$
├── tuned_model
$
├── tuned_model.ait
$
└── tuned_model_sha256_sums.txt

You can copy the checkpoint file tuned_model.ait and SHA sums file to a target host or folder to use it for inference.

Note: We recommend deploying the *.ait package on the same hardware used for tuning to ensure functional and performance compatibility.

Load

The load function enables you to load previously tuned models from a checkpoint file.

1
# Load the tuned model
2
import aitune.torch as ait
3
tuned_model = ait.load(model, "tuned_model.ait")

On first load, the checkpoint file is decompressed and the tuned and original module weights are loaded. Subsequent loads will use the decompressed weights from the same folder.

Tune Strategies

NVIDIA AITune provides different strategies for selecting the optimal backend configuration. The strategies align with a common interface for the tuning process.

Not every backend can tune every model — each relies on different compilation technology with its own limitations (e.g., ONNX export for TensorRT, graph breaks in Torch Inductor, unsupported layers in TorchAO). Strategies control how AITune handles this.

Strategies also validate performance against a Torch eager baseline. Correct backends that do not beat eager by the configured threshold are rejected by OneBackendStrategy and FirstWinsStrategy; profiling strategies such as MaxThroughputStrategy, MinLatencyStrategy, and LatencyBudgetStrategy can compare candidates against a profiled eager baseline. Use strategy.enable_performance_validation(False) to skip Torch eager baseline profiling, performance checks, and speedup reporting.

FirstWinsStrategy

Tries backends in priority order and returns the first one that builds, validates correctness, and beats the Torch eager baseline by the configured threshold. If a backend fails or is slower than baseline, the strategy moves on to the next candidate instead of aborting.

1
from aitune.torch.backend import TensorRTBackend, TorchInductorJitBackend
2
from aitune.torch.tune_strategy import FirstWinsStrategy
3
4
strategy = FirstWinsStrategy(backends=[TensorRTBackend(), TorchInductorJitBackend()])

OneBackendStrategy

Uses exactly one backend, failing immediately with the original error if it cannot build. Use this when you have already validated that a backend works and want deterministic behavior. Unlike FirstWinsStrategy with a single backend, OneBackendStrategy surfaces the original exception rather than catching it.

1
from aitune.torch.backend import TensorRTBackend
2
from aitune.torch.tune_strategy import OneBackendStrategy
3
4
strategy = OneBackendStrategy(backend=TensorRTBackend())

MaxThroughputStrategy

Profiles all compatible backends and selects the fastest one that beats the Torch eager baseline, falling back to eager when no user backend is faster. Use this when maximum throughput matters and you can afford longer tuning time.

1
from aitune.torch.backend import TensorRTBackend, TorchInductorJitBackend, TorchEagerBackend
2
from aitune.torch.tune_strategy import MaxThroughputStrategy
3
4
strategy = MaxThroughputStrategy(backends=[TensorRTBackend(), TorchInductorJitBackend(), TorchEagerBackend()])

MinLatencyStrategy

Profiles all compatible backends and selects the one with the lowest latency that beats the Torch eager baseline, falling back to eager when no user backend is faster. Use this when response time matters more than throughput (e.g. interactive or real-time workloads).

1
from aitune.torch.backend import TensorRTBackend, TorchInductorJitBackend, TorchEagerBackend
2
from aitune.torch.tune_strategy import MinLatencyStrategy
3
4
strategy = MinLatencyStrategy(backends=[TensorRTBackend(), TorchInductorJitBackend(), TorchEagerBackend()])

LatencyBudgetStrategy

Profiles all compatible backends across the configured batch sizes, filters out results whose latency exceeds the budget, and selects the highest-throughput compliant backend. If no user backend satisfies the budget, tuning raises.

1
from aitune.torch.backend import TensorRTBackend, TorchInductorJitBackend
2
from aitune.torch.tune_strategy import LatencyBudgetStrategy
3
4
strategy = LatencyBudgetStrategy(
5
    latency_budget_ms=50.0,
6
    backends=[TensorRTBackend(), TorchInductorJitBackend()],
7
)