JIT Tuning Example

Aim

This example demonstrates how to use NVIDIA AITune to do JIT tuning without any code changes required i.e. there is no additional import to the user’s code.

Environment Setup

Install dependencies:

uv sync

Usage

If you would like to run a model with JIT tuning enabled set the following environment variable:

export AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning

or set it inplace

AUTOWRAPT_BOOTSTRAP=aitune_enable_jit_tuning uv run resnet

If the environment variable is not set, the model will run without any tuning.

Resnet

uv run resnet

Diffusion models

Stable Diffusion 3

uv run diffusion --model stabilityai/stable-diffusion-3-medium-diffusers

Stable Diffusion XL

uv run diffusion --model stabilityai/stable-diffusion-xl-base-1.0

Flux

uv run diffusion --model black-forest-labs/FLUX.1-dev