Making stronglyTyped default for modelopt evaluation (#287)

Signed-off-by: Riyad Islam <[email protected]>

Author: i-riyad

CHANGELOG.rst

@@ -5,10 +5,12 @@ Model Optimizer Changelog (Linux)
 ^^^^^^^^^^^^^^^^^
 
 **Deprecations**
+- Deprecated ``quantize_mode`` argument in ``examples/onnx_ptq/evaluate.py`` to support strongly typing. Use ``engine_precision`` instead.
 
 **Bug Fixes**
 
 **New Features**
+- ``high_precision_dtype`` default to fp16 in ONNX quantization, i.e. quantized output model weights are now FP16 by default.
 
 0.35 (2025-09-04)
 ^^^^^^^^^^^^^^^^^

examples/onnx_ptq/README.md

@@ -120,7 +120,7 @@ The following evaluation requires the `val` directory of the [ImageNet dataset](
 python evaluate.py \
     --onnx_path=<path to classification model> \
     --imagenet_path=<path to the ImageNet dataset> \
-    --quantize_mode=<fp8|int8|int4> \
+    --engine_precision=stronglyTyped \
     --model_name=vit_base_patch16_224
 ```
 
@@ -165,7 +165,7 @@ If the input model is of type image classification, use the following script to
 python evaluate.py \
     --onnx_path=<path to the exported ONNX model> \
     --imagenet_path=<path to the ImageNet dataset> \
-    --quantize_mode=stronglyTyped \
+    --engine_precision=stronglyTyped \
     --model_name=vit_base_patch16_224
 ```
 

examples/onnx_ptq/evaluate.py

@@ -48,29 +48,22 @@ def main():
     parser.add_argument(
         "--eval_data_size", type=int, default=None, help="Number of examples to evaluate"
     )
-    # By default, TensorRT autotunes tensor types to generate the fastest engine. When you specify
-    # to TensorRT that a network is strongly typed, it infers a type for each intermediate and
-    # output tensor using the rules in the operator type specification. For networks quantized in
-    # INT4 or FP8 mode, stronglyTyped as the mode is recommended for TensorRT deployment. Though
-    # INT8 networks are generally compiled with int8 mode, certain INT8 ViT networks compiled with
-    # stronglyTyped precision have shown better performance.
     parser.add_argument(
-        "--quantize_mode",
+        "--engine_precision",
         type=str,
         default="stronglyTyped",
-        choices=["fp8", "fp16", "fp32", "int4", "int8", "int8_iq", "bf16", "best", "stronglyTyped"],
-        help="Quantization mode for the TensorRT engine. \
-            Supported options: fp8, fp16, fp32, int8, int8_iq(implicit quantization), bf16, best, stronglyTyped",
+        choices=["best", "fp16", "stronglyTyped"],
+        help="Precision mode for the TensorRT engine. \
+            stronglyTyped is recommended, all other modes have been deprecated in TensorRT",
     )
     parser.add_argument(
         "--results_path", type=str, default=None, help="Save the results to the specified path"
     )
 
     args = parser.parse_args()
-
     deployment = {
         "runtime": "TRT",
-        "precision": args.quantize_mode,
+        "precision": args.engine_precision,
     }
 
     # Create an ONNX bytes object with the specified path

examples/onnx_ptq/evaluation.py

@@ -29,7 +29,7 @@
 deployment = {
     "runtime": "TRT",
     "accelerator": "GPU",
-    "precision": "fp32",
+    "precision": "stronglyTyped",
     "onnx_opset": "21",
 }
 

modelopt/onnx/quantization/__main__.py

@@ -180,11 +180,11 @@ def get_parser() -> argparse.ArgumentParser:
     argparser.add_argument(
         "--high_precision_dtype",
         type=str,
-        default=None,
+        default="fp16",
         choices=["fp32", "fp16", "bf16"],
         help=(
-            "High precision data type, one of ['fp32', 'fp16', 'bf16']. For int8 quantization, the default value is "
-            "'fp32' and 'fp16' for other quantization modes."
+            "High precision data type of the output model. If the input model is of dtype fp32, "
+            "it will be converted to fp16 dtype by default."
         ),
     )
     argparser.add_argument(
@@ -262,8 +262,6 @@ def main():
             # Convert the NpzFile object to a Python dictionary
             calibration_data = {key: calibration_data[key] for key in calibration_data.files}
 
-    default_high_precision_dtype = "fp32" if args.quantize_mode == "int8" else "fp16"
-
     quantize(
         args.onnx_path,
         quantize_mode=args.quantize_mode,
@@ -284,7 +282,7 @@ def main():
         log_file=args.log_file,
         trt_plugins=args.trt_plugins,
         trt_plugins_precision=args.trt_plugins_precision,
-        high_precision_dtype=args.high_precision_dtype or default_high_precision_dtype,
+        high_precision_dtype=args.high_precision_dtype,
         mha_accumulation_dtype=args.mha_accumulation_dtype,
         disable_mha_qdq=args.disable_mha_qdq,
         dq_only=args.dq_only,

modelopt/onnx/quantization/int8.py

@@ -124,7 +124,7 @@ def quantize(
     use_external_data_format: bool = False,
     intermediate_generated_files: list[str] = [],
     trt_extra_plugin_lib_paths: list[str] | None = None,
-    high_precision_dtype: str = "fp32",
+    high_precision_dtype: str = "fp16",
     passes: list[str] = ["concat_elimination"],
     log_level: str = "INFO",
     calibrate_per_node: bool = False,

modelopt/onnx/quantization/qdq_utils.py

@@ -529,6 +529,11 @@ def _get_successive_consumers(
     quantized_node = tensor_consumers.get(dq_node.output[0], [None])[0]
     if not quantized_node:
         raise ValueError(f"No consumer found for {dq_node.name}")
+    if quantized_node.op_type == "Cast":
+        next_node = tensor_consumers.get(quantized_node.output[0], [None])[0]
+        if not next_node:
+            raise ValueError(f"No consumer found after Cast for {quantized_node.name}")
+        quantized_node = next_node
 
     return dq_node, quantized_node
 
@@ -1030,7 +1035,7 @@ def is_pre_quant_scale_node(node: onnx.NodeProto) -> bool:
 
     # Remove transpose and reshape nodes
     new_nodes = [node for node in graph.node if node.name not in nodes_to_remove]
-    graph.node.clear()
+    del graph.node[:]
     graph.node.extend(new_nodes)
 
     def is_fp32_cast(node: onnx.NodeProto) -> bool:

modelopt/onnx/quantization/quantize.py

@@ -219,7 +219,7 @@ def quantize(
     log_file: str | None = None,
     trt_plugins: list[str] | None = None,
     trt_plugins_precision: list[str] | None = None,
-    high_precision_dtype: str | None = None,
+    high_precision_dtype: str = "fp16",
     mha_accumulation_dtype: str = "fp16",
     disable_mha_qdq: bool = False,
     dq_only: bool = True,
@@ -286,12 +286,13 @@ def quantize(
             Each item should have the format <op_type>:<precision>, where precision can be fp32 (default) or fp16.
             For example: op_type_1:fp16 op_type_2:fp32.
         high_precision_dtype:
-            High precision data type, one of ['fp32', 'fp16']. If high_precision_dtype == 'fp16', model's weight and
-            activation will be converted to fp16.
+            High precision data type of the output model. If high_precision_dtype is 'fp16' or 'bf16'
+            and the input model is of dtype fp32, model's weight and activation will be converted to
+            'fp16' or 'bf16'.
         mha_accumulation_dtype:
-            MHA accumulation dtype. One of ['fp32', 'fp16']. 'fp16' by default.
-            If quantize_mode == 'fp8' and mha_accumulation_dtype == 'fp32', Cast nodes will be added to
-            MHA's bmm1 and bmm2's input and output tensors.
+            MHA accumulation dtype. One of ['fp32', 'fp16']. 'fp16' by default. If quantize_mode == 'fp8' and
+            mha_accumulation_dtype == 'fp32', Cast nodes will be added to MHA's bmm1 and bmm2's input
+            and output tensors.
         disable_mha_qdq:
             Don't add Q/DQ layers to MatMuls in MHA pattern.
         dq_only:
@@ -461,7 +462,7 @@ def quantize(
             use_external_data_format=use_external_data_format,
             intermediate_generated_files=intermediate_generated_files,
             trt_extra_plugin_lib_paths=trt_plugins,
-            high_precision_dtype=high_precision_dtype,  # type: ignore[arg-type]
+            high_precision_dtype=high_precision_dtype,
             mha_accumulation_dtype=mha_accumulation_dtype,
             passes=passes,
             log_level=log_level,

modelopt/torch/_deploy/_runtime/tensorrt/constants.py

@@ -86,7 +86,6 @@ class TRTMode:
     BFLOAT16 = "bf16"
     FLOAT8 = "fp8"
     INT8 = "int8"
-    INT8_IQ = "int8_iq"
     INT4 = "int4"
     STRONGLY_TYPED = "stronglyTyped"
     BEST = "best"
@@ -98,7 +97,6 @@ class TRTMode:
     TRTMode.BFLOAT16: ["--bf16"],
     TRTMode.FLOAT8: ["--fp16", "--fp8"],
     TRTMode.INT8: ["--fp16", "--int8"],
-    TRTMode.INT8_IQ: ["--int8"],
     TRTMode.INT4: ["--fp16", "--int4"],
     TRTMode.STRONGLY_TYPED: ["--stronglyTyped"],
     TRTMode.BEST: ["--best"],

modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py

@@ -103,10 +103,10 @@ def _get_trtexec_params(
 def _is_low_bit_mode(trt_mode: str) -> bool:
     return trt_mode in [
         TRTMode.INT8,
-        TRTMode.INT8_IQ,
         TRTMode.INT4,
         TRTMode.FLOAT8,
         TRTMode.BEST,
+        TRTMode.STRONGLY_TYPED,
     ]