Tuning Workflow
This guide provides an in-depth look at AITune’s tuning process, explaining how samples are gathered, how modules are tuned, and how strategies and backends work together to optimize your models.
Overview
The AITune tuning workflow is structured as follows:
- Sample Gathering: Execute the model with different batch sizes to collect metadata
- Graph Detection: Identify unique computational graphs based on input characteristics
- Module Tuning: Optimize each wrapped module sequentially
- Strategy Execution: Apply a tuning strategy to select the best backend
- Backend Activation: Set up the optimized backends for inference
The Tuning Process
1. Sample Gathering Phase
When you call ait.tune(), AITune first enters the sample gathering phase:
During this phase:
- The function/model is executed with samples from the dataset
- Each execution uses a different batch size (from
batch_sizesparameter) - Metadata is recorded for each execution:
- Input/Output shapes, dtypes, and structure
- Batch size information (stored in global context)
- Wrapped modules automatically detect and record this metadata
Key Point: At least 2 different batch sizes are required to detect batch dimensions. If only one batch size is provided, AITune assumes static shapes.
Sample Generation
AITune uses samples_generator to iterate through the dataset:
batch_size: Current batch size being processedargs,kwargs: Actual data samples for this batchmax_num_batches_per_batch_size: Limits how many model executions are done per batch size (useful to limit large datasets)
The global context tracks the current batch size, allowing wrapped modules to correlate shape changes with batch size changes.
2. Graph Detection
As samples are collected, wrapped modules detect unique computational graphs:
Graph Identity Rules:
- The same forward parameters passed positionally or by keyword → same graph
- Additional or omitted parameters may produce different graphs
- Tensors with different ranks → different graphs
- Different tensor shapes but same rank → same graph (dynamic shapes)
- Different non-tensor arguments (in strict mode) → different graphs
Note:
- If strict mode is turned off, only tensor data is taken into account when detecting graphs. The strict mode can be turned off with:
- For a particular graph, there is only a limited number of samples collected to limit memory usage. This threshold can be set with:
Each unique graph is represented by a GraphSpec containing:
- Name: Unique identifier (e.g., “0”, “1”, “2”)
- Input Spec:
SampleMetadatadescribing expected inputs - Output Spec:
SampleMetadatadescribing expected outputs
Batch and Dynamic Dimensions:
After seeing multiple samples, AITune identifies:
- Batch dimensions (e.g.,
batch0): Scale proportionally with batch size - Dynamic dimensions (e.g.,
dim0): Vary independently of batch size - Static dimensions: Never change
Example:
See also Execution Graphs - in depth explanation of detecting graphs and sample metadata.
3. Module-by-Module Tuning
After sample gathering, AITune tunes each wrapped module sequentially:
Why Sequential?
- Memory Management: Deactivating other modules frees GPU memory for tuning
- Isolation: Prevents interference between module optimizations
- Predictability: Each module gets full system resources
For each module, each graph is tuned separately, i.e., a strategy is called for data captured for the specific graph only and:
- Strategy tries to build a backend or backends and select the best one
- Each backend is validated against outputs, i.e., tensor shapes, values, NaNs (not a number)
- Strategies also profile a Torch eager baseline and reject or fall back from correct backends that do not beat the baseline
4. Strategy Execution
The tuning strategy determines which backend(s) to try and how to select the best one. AITune provides five built-in strategies.
Performance Validation
AITune validates backend performance against Torch eager by default. Most users should keep this behavior unchanged: AITune profiles an eager baseline, compares correctness-passing backends against it, and avoids selecting an optimization that does not improve performance.
Change the performance validation mode only when collecting comparison data, diagnosing backend performance, or intentionally skipping the eager comparison.
Performance validation has three modes:
Configure the mode on a tune strategy:
The method also accepts serialized string values:
Boolean values remain available as a shorthand. True maps to ENABLED, and False maps to DISABLED:
Use DIAGNOSTIC when you want eager and backend performance measurements for analytics but need backend selection to
remain independent of that comparison. This is useful for evaluating thresholds, investigating benchmark noise, or
collecting data before enabling a performance gate in a new environment.
Use DISABLED when the eager comparison is not meaningful or when you intentionally want to avoid its tuning-time
cost. This skips eager-baseline profiling, comparison, and speedup reporting. Profiling strategies still measure user
backends because those measurements are required to select a winner.
For OneBackendStrategy and FirstWinsStrategy, ENABLED rejects a correctness-passing backend when its throughput
does not exceed the eager baseline by the configured threshold. The default threshold is 1%, controlled by
min_speedup_threshold_percent.
For MaxThroughputStrategy, MinLatencyStrategy, and LatencyBudgetStrategy, ENABLED includes eager in the final
selection decision. DIAGNOSTIC records the eager comparison but selects the best successful user backend.
DISABLED skips eager profiling and also selects among successful user backends only.
In DIAGNOSTIC mode, profiling strategies can still use eager as a fallback when no user backend succeeds. The mode
prevents performance comparison from affecting selection; it does not remove the available fallback.
ENABLED and DIAGNOSTIC add the cost of building and profiling the eager baseline. Tuning telemetry records the
baseline and backend performance metrics, from which speedup can be derived; AITune also writes the calculated speedup
to its logs. DISABLED avoids that additional baseline work. See Observability for
telemetry output and Profiling and Hardware Metrics for profiling tools.
FirstWinsStrategy
Tries backends in priority order and returns the first one that builds, validates correctness, and meets the performance threshold against the Torch eager baseline. Not every backend can handle every model (e.g., TensorRT may fail during ONNX export, Torch Inductor may hit graph breaks), so this strategy provides automatic fallback instead of aborting.
Workflow:
- Try each backend in order
- Build the backend with the module and graph spec
- Validate correctness by comparing outputs
- Profile against the Torch eager baseline
- Return the first backend that passes the performance threshold
Use Case: Fast tuning with automatic fallback, especially for models you haven’t validated against every backend.
OneBackendStrategy
Uses exactly one backend, failing immediately with the original error if it cannot build or validate correctness. Unlike FirstWinsStrategy with a single backend, OneBackendStrategy surfaces build and correctness exceptions rather than catching them. If the backend is correct but does not pass the eager performance gate, it falls back to TorchEagerBackend.
Workflow:
- Build the specified backend
- Validate correctness
- Profile against the Torch eager baseline
- Return the backend when it passes the performance threshold, or fall back to the eager baseline when it is correct but slower
Use Case: Production with a validated backend where you want deterministic behavior.
MaxThroughputStrategy
Before actual tuning, this strategy tries to estimate max_batch_size. It does so by incrementing batch_size (in powers of 2) and measuring throughput using the original module. The max_batch_size is picked for the best throughput and then is used for selecting the best backend:
Workflow:
- Estimate
max_batch_size - Try each backend in order
- Build and validate each working backend
- Profile throughput for each backend given
max_batch_size - Return the fastest user backend when it beats the Torch eager baseline, otherwise fall back to eager
Use Case: When performance is critical and you want the absolute fastest backend.
5. Backend Building and Validation
A backend represents a different technology for tuning a torch module, e.g., TensorRT, TorchInductor, and it is used by a strategy. Before it is used, it acts as a blueprint, i.e., it is copied, and each copy is used to build, validate, and activate a particular tuned module. This is done so that it has no side-effects on different modules or graphs.
Each backend is a small state machine that enforces safe usage:
INIT→ACTIVE:build()succeeds. A build is allowed only once and must set up a runnable backend.ACTIVE→INACTIVE:deactivate()releases resources (and clears compiler/runtime caches).INACTIVE→ACTIVE:activate()restores the backend for inference.CHECKPOINT_LOADED→ACTIVEorDEPLOYED: A backend created from a checkpoint can be activated for tinkering or deployed for final use.ACTIVE→DEPLOYED:deploy()finalizes the backend. After this, state changes are not allowed.
The backend’s state is governed by the strategy and the user must not change it. After a module is successfully tuned, it can be used to do inference - the backend will be in ACTIVE or DEPLOYED states.
Validation Phase
After building a backend, the strategy tries to validate it. This is enabled by default and can be turned off with strategy.enable_correctness_check(False).
AITune validates correctness by running the tuned backend on sample data. These checks are required to ensure the backend is correctly built:
- Python basic types
int,floatmust be finite. - Tensors values must be finite.
- All nested structures are checked against points 1 and 2.
- Tensor shapes must match against original data i.e. static, dynamic and batch axes must match.
This ensures tuning does not break the module.
6. Backend Activation
After tuning is complete, all tuned modules are activated:
Activation means:
- Load the optimized backend into memory
- Prepare for inference
- Route future calls to the optimized backend
Now your model is ready for optimized inference!
Workflow Customization
The workflow can be adapted by implementing a custom strategy or a backend.
Custom Strategy
If you would like to write a custom strategy, extend the TuneStrategy class and implement the _tune method:
If needed, you can also extend the_pre_tune and _post_tune methods, which are invoked before and after the _tune.
Custom Backend
To write a custom backend, extend the Backend base class and BackendConfig data class which serves as a config for your custom backend. The following methods are required to be implemented:
key()- Returns a stable identifier for the backend/config combination; used for caching and lookup.describe()- Returns a short human-readable description of the backend/config changes.to_dict()- Serializes backend state; includesArtifactPathobjects for files or directories to bundle in checkpoints.from_dict()- Reconstructs a backend instance from the serialized state._build()- Builds backend artifacts for a specific module/graph and returns a ready backend._activate()- Loads/initializes the backend for inference after it was inactive or checkpoint-loaded._deactivate()- Releases runtime resources and makes the backend inactive._deploy()- Finalizes the backend for deployment; after this it cannot change state._infer()- Executes inference with the backend for the provided inputs.
Every custom backend must declare its supported execution modes and build mode. Include ExecutionMode.SINGLE_GPU for
ordinary modules and ExecutionMode.MULTI_GPU when the backend can preserve distributed module execution. If
activation compiles or otherwise requires the original PyTorch module, declare the backend as just-in-time so AITune
retains that module:
For serialization with ait.load and ait.save, the to_dict and from_dict methods are used. Anything placed in
the dictionary will be saved and restored as checkpoint state. Files and directories are copied only when explicitly
represented by ArtifactPath; ordinary Path values remain regular serialized configuration values.
Create a planned artifact from its owning cache directory and relative path. Both arguments accept strings and
Path objects, while the path property returns the effective filesystem location used by the backend:
If another API returns the path of an artifact it has already created, use
ArtifactPath.from_existing(engine_path, root=cache_dir). The explicit root preserves the intended directory
structure and prevents accidentally including a file outside that directory. A backend’s from_dict() receives
ready-to-use ArtifactPath objects, so it should not construct or rebase them itself.
AITune preserves tensorrt/model.engine inside the checkpoint and rebases the artifact root when loading it from a
different location.
Monitoring the Workflow
AITune provides detailed logging throughout the tuning process.
Example logs from tuning ResNet (an example is placed in the examples/ResNet folder):
Now the logs will show the inputs and outputs of the module:
As you can see, AITune detected the batch axis as the first one, hence the name batch0 and input shapes 3x224x224, i.e., batch of images, and output shape 1000, i.e., batch of categories.
Next, you can see the Max Throughput Strategy builds backends one by one:
Finally the winning backend is presented:
If you do not see logs, make sure the logger is configured to at least INFO level:
See Also
- Execution Graphs - Understanding graph detection
- Tune Strategies - Strategy reference