fix: support ITensor amex in quantize and add export_torch_mode() wrapper

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
 

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)