Making stronglyTyped default for modelopt evaluation

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

Author: i-riyad

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> \
+    --quantize_mode=stronglyTyped \
     --model_name=vit_base_patch16_224
 ```
 

examples/onnx_ptq/evaluate.py

@@ -48,29 +48,16 @@ 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",
-        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",
-    )
     parser.add_argument(
         "--results_path", type=str, default=None, help="Save the results to the specified path"
     )
 
     args = parser.parse_args()
 
+    # Note. stronglyTyped is recommended, all other modes have been deprecated in TensorRT
     deployment = {
         "runtime": "TRT",
-        "precision": args.quantize_mode,
+        "precision": "stronglyTyped",
     }
 
     # 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/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']. 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,
     ]
 
 

modelopt/torch/_deploy/_runtime/trt_client.py

@@ -49,7 +49,6 @@ def deployment_table(self) -> DeploymentTable:
                 "bf16",
                 "fp8",
                 "int8",
-                "int8_iq",
                 "int4",
                 "stronglyTyped",
                 "best",

tests/_test_utils/onnx_quantization/utils.py

@@ -20,7 +20,11 @@ def _assert_nodes_are_quantized(nodes):
     for node in nodes:
         for inp_idx, inp in enumerate(node.inputs):
             if isinstance(inp, gs.Variable):
-                assert node.i(inp_idx).op == "DequantizeLinear", (
-                    f"Input '{inp.name}' of node '{node.name}' is not quantized but should be!"
+                qnode = node
+                # After quantization, the quantized node can be casted
+                if qnode.i(inp_idx).op == "Cast":
+                    qnode = qnode.i(inp_idx)
+                assert qnode.i(inp_idx).op == "DequantizeLinear", (
+                    f"Input '{inp.name}' of node '{qnode.name}' is not quantized but should be!"
                 )
     return True