aitune.torch.backend.tensorrt.torch_quantization

NVIDIA ModelOpt PyTorch quantization module for TensorRT backend.

Module Contents

Classes

Name	Description
`TorchQuantizationConfig`	Configuration for Torch quantization.
`TorchQuantizer`	NVIDIA ModelOpt PyTorch quantizer for TensorRT backend.

Functions

Name	Description
`_validate_torch_quantization_config`	Validate that quantization_config is a supported value.

Data

HAS_NVFP4_FP8_MHA_CONFIG

QuantizationConfig

logger

API

class aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizationConfig(
    quantization_config: aitune.torch.backend.tensorrt.torch_quantization.QuantizationConfig = 'FP8_DEFAULT_CFG',
    device: str = 'cuda'
)

Dataclass

Configuration for Torch quantization.

Parameters:

quantization_config

QuantizationConfigDefaults to 'FP8_DEFAULT_CFG'

ModelOpt quantization configuration name.

device

strDefaults to 'cuda'

Target device for quantization (default: “cuda”)

verbose

Enable verbose logging

device

str = 'cuda'

quantization_config

QuantizationConfig = 'FP8_DEFAULT_CFG'

aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizationConfig.__post_init__()

Validate quantization_config after initialization.

aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizationConfig.from_dict(
    state_dict: dict
)

classmethod

Convert dict to TorchQuantizationConfig.

class aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizer()

NVIDIA ModelOpt PyTorch quantizer for TensorRT backend.

This class provides functionality to quantize PyTorch models using NVIDIA ModelOpt for deployment with TensorRT. It supports various quantization modes including INT8, FP8, and INT4 quantization with different calibration methods.

aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizer._clean_memory()

Clean up memory.

aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizer._configure_quantizers_for_onnx_export(
    model: torch.nn.Module
)

Configure quantizers in the model for ONNX export.

Sets appropriate attributes on input and weight quantizers to ensure proper ONNX export. For input quantizers, sets dynamic quantization mode and half precision. For weight quantizers, sets static quantization mode.

Parameters:

model

nn.Module

The quantized PyTorch model with quantizers to configure

aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizer._create_forward_loop(
    calibration_samples: list[typing.Any]
)

Create forward loop function for model calibration.

Parameters:

calibration_samples

list[Any]

List of calibration samples

Returns:

Forward loop function for calibration

aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizer._get_quantization_config(
    config_name: str
) -> dict[str, typing.Any]

Get quantization configuration by name.

Parameters:

config_name

str

Name of the quantization configuration

Returns: dict[str, Any]

Quantization configuration dictionary

Raises:

ImportError: If ModelOpt is not available
ValueError: If configuration name is not supported

aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizer._prepare_calibration_sample(
    sample: aitune.torch.module.recording_module.Sample,
    device: str
) -> aitune.torch.module.recording_module.Sample

Move sample to target device and return as Sample tuple.

Parameters:

sample

Sample

Example input sample

device

str

Target device for calibration

Returns: Sample

List containing the sample converted to target device

aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizer.quantize(
    module: torch.nn.Module,
    sample: aitune.torch.module.recording_module.Sample,
    config: aitune.torch.backend.tensorrt.torch_quantization.TorchQuantizationConfig
) -> torch.nn.Module

Quantize the module using NVIDIA ModelOpt.

This method performs quantization of a PyTorch model using NVIDIA ModelOpt. It supports various quantization algorithms including FP8, INT8, and INT4.

Parameters:

module

nn.Module

PyTorch module to export

sample

Sample

Example input for the model

config

TorchQuantizationConfig

Returns: nn.Module

Quantized model

Raises:

ImportError: If ModelOpt is not installed
RuntimeError: If quantization or export fails

aitune.torch.backend.tensorrt.torch_quantization._validate_torch_quantization_config(
    config_name: str
) -> None

Validate that quantization_config is a supported value.

aitune.torch.backend.tensorrt.torch_quantization.HAS_NVFP4_FP8_MHA_CONFIG = Version(modelopt.__version__) >= Version('0.35.0')

aitune.torch.backend.tensorrt.torch_quantization.QuantizationConfig = Literal['NVFP4_DEFAULT_CFG', 'NVFP4_FP8_MHA_CONFIG', 'FP8_DEFAULT_CFG', 'INT8_DE...

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