diff --git a/py/torch_tensorrt/dynamo/_tracer.py b/py/torch_tensorrt/dynamo/_tracer.py index 399559347b..0160d9c3d4 100644 --- a/py/torch_tensorrt/dynamo/_tracer.py +++ b/py/torch_tensorrt/dynamo/_tracer.py @@ -1,6 +1,8 @@ from __future__ import annotations +import importlib.util import logging +from contextlib import nullcontext from inspect import signature from typing import Any, Optional, Tuple, Union @@ -79,13 +81,22 @@ def trace( dim_registry = build_dim_registry(arg_inputs, kwarg_inputs) dynamic_shapes = get_dynamic_shapes_args(mod, arg_inputs, dim_registry) dynamic_shapes.update(get_dynamic_shapes_kwargs(kwarg_inputs, dim_registry)) - exp_program = export( - mod, - tuple(torch_arg_inputs), - kwargs=torch_kwarg_inputs, - dynamic_shapes=dynamic_shapes, - strict=kwargs.get("strict", False), - ) + + export_context = nullcontext() + if importlib.util.find_spec("modelopt") is not None: + from modelopt.torch.quantization.utils import export_torch_mode, is_quantized + + if is_quantized(mod): + export_context = export_torch_mode() + + with export_context: + exp_program = export( + mod, + tuple(torch_arg_inputs), + kwargs=torch_kwarg_inputs, + dynamic_shapes=dynamic_shapes, + strict=kwargs.get("strict", False), + ) return exp_program diff --git a/py/torch_tensorrt/dynamo/conversion/impl/quantize.py b/py/torch_tensorrt/dynamo/conversion/impl/quantize.py index 8dd32b11fc..8ec0895832 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/quantize.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/quantize.py @@ -40,7 +40,7 @@ def quantize( source_ir: Optional[SourceIR], name: str, input_tensor: TRTTensor, - amax: Union[np.ndarray, torch.Tensor], + amax: Union[np.ndarray, torch.Tensor, TRTTensor], num_bits: int, exponent_bits: int, ) -> TRTTensor: @@ -78,25 +78,27 @@ def quantize( max_bound = 448 axis = None - # int8 weight quantization is per-channel quantization(it can have one or multiple amax values) - if dtype == trt.DataType.INT8 and amax.numel() > 1: - # if the amax has more than one element, calculate the axis, otherwise axis value will be ignored - amax_init_shape = amax.shape - amax = amax.squeeze().data - assert ( - len(amax.shape) == 1 - ), f"TensorRT does not support multi-axis quantization. {name=} {amax_init_shape=} {amax.shape=} " - axis = list(amax_init_shape).index(list(amax.shape)[0]) - assert ( - axis == 0 - ), f"{name=} {amax=} is per-channel quantization, expected axis to be 0, but got {axis=}" - else: - # int8 activation and fp8 weight/activation quantization is per-tensor quantization, it can only have single amax value - assert ( - amax.numel() == 1 - ), f"{name=} is per-tensor quantization, expected amax is a singular value, but got {amax.shape=}" - + # Dynamic amax (TRT ITensor) is always treated as per-tensor; numel()/shape + # checks only apply to constant torch/numpy amax values. if not isinstance(amax, trt.ITensor): + # int8 weight quantization is per-channel quantization(it can have one or multiple amax values) + if dtype == trt.DataType.INT8 and amax.numel() > 1: + # if the amax has more than one element, calculate the axis, otherwise axis value will be ignored + amax_init_shape = amax.shape + amax = amax.squeeze().data + assert ( + len(amax.shape) == 1 + ), f"TensorRT does not support multi-axis quantization. {name=} {amax_init_shape=} {amax.shape=} " + axis = list(amax_init_shape).index(list(amax.shape)[0]) + assert ( + axis == 0 + ), f"{name=} {amax=} is per-channel quantization, expected axis to be 0, but got {axis=}" + else: + # int8 activation and fp8 weight/activation quantization is per-tensor quantization, it can only have single amax value + assert ( + amax.numel() == 1 + ), f"{name=} is per-tensor quantization, expected amax is a singular value, but got {amax.shape=}" + amax = to_torch(amax, None) scale = torch.divide(amax, max_bound) scale = get_trt_tensor(ctx, scale, name + "_scale", dtype=torch.float32)