aitune.torch.backend.tensorrt.onnx_autocast | NVIDIA AITune Documentation

NVIDIA ModelOpt ONNX quantization module for TensorRT backend.

Module Contents

Classes

Name	Description
`ONNXAutoCast`	NVIDIA ModelOpt ONNX autocast for TensorRT backend.
`ONNXAutoCastConfig`	Configuration for mixed precision quantization.

Functions

Name	Description
`_validate_autocast_precision`	Validate that precision is a supported value for autocast.

Data

MixedPrecision

logger

API

class aitune.torch.backend.tensorrt.onnx_autocast.ONNXAutoCast()

NVIDIA ModelOpt ONNX autocast for TensorRT backend.

This class provides functionality to autocast ONNX models using NVIDIA ModelOpt.

aitune.torch.backend.tensorrt.onnx_autocast.ONNXAutoCast.autocast(
    input_onnx_path: str | pathlib.Path,
    output_path: str | pathlib.Path,
    config: aitune.torch.backend.tensorrt.onnx_autocast.ONNXAutoCastConfig,
    samples: list[aitune.torch.module.recording_module.Sample] | None = None,
    graph_spec: aitune.torch.module.graph_spec.GraphSpec | None = None
) -> pathlib.Path

Autocast the ONNX model using NVIDIA ModelOpt.

Parameters:

input_onnx_path

str | Path

Path to the input ONNX model

output_path

str | Path

Path to the output ONNX model

config

ONNXAutoCastConfig

samples

list[Sample] | NoneDefaults to None

Optional calibration dataset for quantization. If None, random data will be used for calibration

graph_spec

GraphSpec | NoneDefaults to None

Graph specification containing input names mapping

Returns: Path

Path to the autocasted ONNX file

Raises:

ValueError: If unsupported precision is specified
RuntimeError: If autocast fails

class aitune.torch.backend.tensorrt.onnx_autocast.ONNXAutoCastConfig(
    precision: aitune.torch.backend.tensorrt.onnx_autocast.MixedPrecision = 'fp16',
    keep_io_types: bool = True
)

Dataclass

Configuration for mixed precision quantization.

Parameters:

precision

MixedPrecisionDefaults to 'fp16'

Mixed precision quantization precision (“fp16”, “bf16”)

keep_io_types

bool = True

precision

MixedPrecision = 'fp16'

aitune.torch.backend.tensorrt.onnx_autocast.ONNXAutoCastConfig.__post_init__()

Validate precision after initialization.

aitune.torch.backend.tensorrt.onnx_autocast.ONNXAutoCastConfig.from_dict(
    state_dict: dict
)

classmethod

Convert dict to MixedPrecisionConfig.

aitune.torch.backend.tensorrt.onnx_autocast._validate_autocast_precision(
    precision: str
) -> None

Validate that precision is a supported value for autocast.

aitune.torch.backend.tensorrt.onnx_autocast.MixedPrecision = Literal['fp16', 'bf16']

aitune.torch.backend.tensorrt.onnx_autocast.logger = logging.getLogger(__name__)