aitune.torch.jit.config | NVIDIA AITune Documentation

Configuration for JIT module.

Module Contents

Classes

Name	Description
`Config`	Configuration for JIT module.
`JITMode`	Mode for JIT execution.

Data

config

API

class aitune.torch.jit.config.Config(
    mode: aitune.torch.jit.config.JITMode = JITMode.TUNE_EAGER,
    dry_run: bool = False,
    dry_run_failure_probability: float = 0.2,
    device: str | torch.device | None = None,
    min_samples: int = 1,
    batch_axis_required: bool = True,
    max_depth_level: int = 1,
    min_parameters: int = 0,
    detect_graph_breaks: bool = False,
    skip_modules: list[str] = list(),
    patch_exclude: tuple[str, ...] = (),
    cache_dir: pathlib.Path = (lambda: _AITUNE_JIT_CACHE_...,
    strategy: aitune.torch.tune_strategy.tune_strategy.TuneStrategy | None = None
)

Dataclass

Configuration for JIT module.

batch_axis_required

bool = True

cache_dir

Path

detect_graph_breaks

bool = False

device

str | device | None = None

dry_run

bool = False

dry_run_failure_probability

float = 0.2

max_depth_level

int = 1

min_parameters

int = 0

min_samples

int = 1

mode

JITMode = JITMode.TUNE_EAGER

patch_exclude

tuple[str, ...] = ()

skip_modules

list[str] = field(default_factory=list)

strategy

TuneStrategy | None = None

aitune.torch.jit.config.Config.__post_init__()

Post init.

aitune.torch.jit.config.Config.reset_to_defaults() -> None

Reset all options to their default values (e.g. for test isolation).

aitune.torch.jit.config.Config.resolve_strategy() -> aitune.torch.tune_strategy.tune_strategy.TuneStrategy

Return the tune strategy to use for JIT tuning.

When strategy is set explicitly it is returned as-is. Otherwise the default is a MaxThroughputStrategy. Ordinary modules profile TensorRT (with and without dynamo) and TorchInductor JIT. Distributed modules profile TorchInductor AOT and TorchInductor JIT. Candidates are resolved when the module is available.

Strategy and backend modules are imported lazily to keep the JIT config a thin data layer that doesn’t pull runtime modules at import time.

class aitune.torch.jit.config.JITMode

Bases: enum.Enum

Mode for JIT execution.

INSPECT

= 'inspect'

TUNE_DEFERRED

= 'tune_deferred'

TUNE_EAGER

= 'tune_eager'

aitune.torch.jit.config.config = Config()