aitune.torch.backend.torchao_backend | NVIDIA AITune Documentation

Torchao backend.

Module Contents

Classes

Name	Description
`TorchAOBackend`	Backend that does torch quantization.
`TorchAOBackendConfig`	Configuration for TorchAOBackend.

Data

DEFAULT_QUANTIZATION

MXFP8DQ_BLOCK_SIZE_DIVISIBILITY

MX_FORMATS_AVAILABLE

NVFP4DQ_BLOCK_SIZE_DIVISIBILITY

QuantizationType

_BASE_QUANTIZATION_CONFIGS

_HW_DEPENDENT_QUANTIZATIONS

logger

API

class aitune.torch.backend.torchao_backend.TorchAOBackend(
    config: aitune.torch.backend.torchao_backend.TorchAOBackendConfig | None = None
)

Bases: Backend

Backend that does torch quantization.

If you would like to use customize quantization, you can pass in a quantization config.

STATE_COMPILE_DYNAMIC

= 'compile_dynamic'

STATE_CONFIG

= 'config'

STATE_DATA

= 'data'

STATE_DEVICE

= 'device'

STATE_ORIG_MODULE

= 'orig_module'

STATE_TYPE

= 'type'

_compile_dynamic

= self._config.dynamic

_config

aitune.torch.backend.torchao_backend.TorchAOBackend._activate()

Activates backend.

aitune.torch.backend.torchao_backend.TorchAOBackend._build(
    module: torch.nn.Module,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec,
    data: list[aitune.torch.module.recording_module.Sample],
    cache_dir: pathlib.Path
) -> aitune.torch.backend.backend.Backend

Builds the model with torchao quantization and torch.compile.

aitune.torch.backend.torchao_backend.TorchAOBackend._check_hardware_compatibility(
    module: torch.nn.Module
) -> None

Check hardware and model constraints for the selected quantization format.

Parameters:

module

nn.Module

The PyTorch module to validate.

Raises:

RuntimeError: If hardware or model constraints are not met.

aitune.torch.backend.torchao_backend.TorchAOBackend._check_mxfp8dq_compatibility(
    module: torch.nn.Module
) -> None

Validate sm100+ GPU, bfloat16 dtype, and block-size alignment for mxfp8dq.

Pre-flight mirror of the bfloat16 assertion in torchao.prototype.mx_formats.inference_workflow._mx_inference_linear_transform and the block_size=32 field of MXDynamicActivationMXWeightConfig. The sm100+ and library requirement is captured by the module-level MX_FORMATS_AVAILABLE flag (evaluated at import time). Dtype is bfloat16-only — torchao asserts this explicitly; float32 is not accepted despite what older comments suggested.

Parameters:

module

nn.Module

The PyTorch module to validate.

Raises:

RuntimeError: If sm100+ GPU is unavailable, any parameter is not bfloat16, or a weight’s last dimension is not divisible by MXFP8DQ_BLOCK_SIZE_DIVISIBILITY.

aitune.torch.backend.torchao_backend.TorchAOBackend._check_nvfp4dq_compatibility(
    module: torch.nn.Module
) -> None

Validate sm100+ GPU, bfloat16 dtype, and weight shape alignment for nvfp4dq.

Pre-flight mirror of the assertions in torchao.prototype.mx_formats.inference_workflow._nvfp4_inference_linear_transform (is_sm_at_least_100() assert, bfloat16 assert, and weight.shape[-2] % 16 check). The sm100+ and library requirement is captured by the module-level MX_FORMATS_AVAILABLE flag (evaluated at import time).

Parameters:

module

nn.Module

The PyTorch module to validate.

Raises:

RuntimeError: If sm100+ GPU is unavailable, the model dtype is unsupported, or any weight’s last two dimensions are not divisible by the block size.

aitune.torch.backend.torchao_backend.TorchAOBackend._deactivate()

Deactivates backend.

aitune.torch.backend.torchao_backend.TorchAOBackend._deploy()

Deploys the backend.

aitune.torch.backend.torchao_backend.TorchAOBackend._do_torchao_quantization()

Apply TorchAO quantization and warm up the compiled model.

aitune.torch.backend.torchao_backend.TorchAOBackend._infer(
    args: typing.Any = (),
    kwargs: typing.Any = {}
) -> typing.Any

Runs inference with the given arguments.

Parameters:

*args

AnyDefaults to ()

inference arguments

**kwargs

AnyDefaults to {}

inference keyword arguments

Returns: Any

The result of the inference.

aitune.torch.backend.torchao_backend.TorchAOBackend._iter_quantized_parameters(
    module: torch.nn.Module
)

Iterate parameters from modules selected by the TorchAO filter function.

TorchAO applies filter_fn at the module level and skips modules where it returns False. The compatibility preflight must mirror that behavior, otherwise it can reject parameters in layers that TorchAO would not quantize anyway, such as embeddings, output projections, or small Linear layers excluded by a model-specific predicate.

aitune.torch.backend.torchao_backend.TorchAOBackend._save_config(
    cache_dir: pathlib.Path
)

Store the backend configuration to a file.

aitune.torch.backend.torchao_backend.TorchAOBackend.describe() -> str

Returns the description of the backend.

aitune.torch.backend.torchao_backend.TorchAOBackend.from_dict(
    module: torch.nn.Module | None,
    state_dict: dict
)

classmethod

Creates a backend from a state_dict.

aitune.torch.backend.torchao_backend.TorchAOBackend.key() -> str

Returns the key of the backend.

aitune.torch.backend.torchao_backend.TorchAOBackend.to_dict()

Returns the state_dict of the backend.

class aitune.torch.backend.torchao_backend.TorchAOBackendConfig(
    fullgraph: bool = False,
    dynamic: bool | None = None,
    mode: aitune.torch.libs.torch_compile.TorchCompileMode | None = 'max-autotune',
    quantization: aitune.torch.backend.torchao_backend.QuantizationType | None = None,
    quantization_config: torchao.core.config.AOBaseConfig | None = None,
    filter_fn: collections.abc.Callable[[nn.Module, str], bool] | None = None
)

Dataclass

Bases: BackendConfig

Configuration for TorchAOBackend.

Parameters:

fullgraph

boolDefaults to False

Passed to torch.compile to require a single graph.

dynamic

bool | NoneDefaults to None

Passed to torch.compile. When None, TorchAOBackend enables dynamic compilation only for graph specs with detected dynamic axes, without mutating this config.

mode

TorchCompileMode | NoneDefaults to 'max-autotune'

Passed through to torch.compile(mode=...). Valid presets are "default", "reduce-overhead", "max-autotune", and "max-autotune-no-cudagraphs".

quantization

QuantizationType | NoneDefaults to None

Name of a built-in TorchAO quantization preset.

quantization_config

AOBaseConfig | NoneDefaults to None

Custom TorchAO quantization config. Use this instead of quantization for advanced TorchAO options.

filter_fn

Callable[[nn.Module, str], bool] | NoneDefaults to None

Optional TorchAO quantization predicate. It receives (module, fqn) and should return True for modules that should be quantized. Compatibility preflight checks use the same predicate.

dynamic

bool | None = None

filter_fn

Callable[[nn.Module, str], bool] | None = None

fullgraph

bool = False

mode

TorchCompileMode | None = 'max-autotune'

quantization

QuantizationType | None = None

quantization_config

AOBaseConfig | None = None

aitune.torch.backend.torchao_backend.TorchAOBackendConfig.__post_init__()

Post init for TorchAOBackendConfig.

aitune.torch.backend.torchao_backend.TorchAOBackendConfig._get_quantization_config(
    quantization
)

aitune.torch.backend.torchao_backend.TorchAOBackendConfig.describe() -> str

Returns the description of the backend.

aitune.torch.backend.torchao_backend.TorchAOBackendConfig.from_dict(
    data: dict
) -> aitune.torch.backend.torchao_backend.TorchAOBackendConfig

classmethod

Initialise config from a plain dict (e.g. parsed from YAML).

Supports two forms:

{"quantization": "int8wo"} — user-facing YAML form; quantization must be one of the supported QuantizationType literals.
{"quantization_config": <bytes>} — internal checkpoint form produced by to_dict(), where quantization_config is pickle-serialised.

aitune.torch.backend.torchao_backend.TorchAOBackendConfig.key() -> str

Returns the key of the backend configuration.

aitune.torch.backend.torchao_backend.TorchAOBackendConfig.to_dict()

Convert TorchAOBackendConfig to dict.

aitune.torch.backend.torchao_backend.DEFAULT_QUANTIZATION = 'fp8wo'

aitune.torch.backend.torchao_backend.MXFP8DQ_BLOCK_SIZE_DIVISIBILITY = 32

aitune.torch.backend.torchao_backend.MX_FORMATS_AVAILABLE = torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)

aitune.torch.backend.torchao_backend.NVFP4DQ_BLOCK_SIZE_DIVISIBILITY = 16

aitune.torch.backend.torchao_backend.QuantizationType = Literal['int8wo', 'int8dq', 'fp8wo', 'fp8dq', 'mxfp8dq', 'nvfp4dq']

aitune.torch.backend.torchao_backend._BASE_QUANTIZATION_CONFIGS: dict = {'int8wo': Int8WeightOnlyConfig(), 'int8dq': Int8DynamicActivationInt8WeightConf...

aitune.torch.backend.torchao_backend._HW_DEPENDENT_QUANTIZATIONS = frozenset({'nvfp4dq', 'mxfp8dq'})

aitune.torch.backend.torchao_backend.logger = getLogger(__name__)