Add: support input_shape_profile for trt-rtx ep

examples/windows/onnx_ptq/genai_llm/quantize.py

@@ -27,65 +27,13 @@
 from transformers import AutoConfig, AutoTokenizer
 
 from modelopt.onnx.quantization.int4 import quantize as quantize_int4
+from modelopt.onnx.quantization.ort_utils import create_input_shapes_profile
 
 logging.getLogger().setLevel(logging.INFO)
 
 pt_to_np = {"torch.int64": np.int64, "torch.float32": np.float32, "torch.float16": np.float16}
 
 
-def prepare_input_shapes_string(
-    batch_size, seq_len, past_seq_len, num_layers, num_kv_heads, head_dim
-):
-    shapes = ""
-
-    shapes += f"input_ids:{batch_size}x{seq_len}"
-    shapes += f",attention_mask:{batch_size}x{seq_len}"
-
-    for i in range(num_layers):
-        key_name = f"past_key_values.{i}.key"
-        value_name = f"past_key_values.{i}.value"
-        shapes += f",{key_name}:{batch_size}x{num_kv_heads}x{past_seq_len}x{head_dim}"
-        shapes += f",{value_name}:{batch_size}x{num_kv_heads}x{past_seq_len}x{head_dim}"
-
-    return shapes
-
-
-def get_input_shapes_profile(model_name_or_path):
-    config = AutoConfig.from_pretrained(model_name_or_path)
-
-    head_dim = config.hidden_size // config.num_attention_heads
-    if hasattr(config, "head_dim") and config.head_dim is not None:
-        head_dim = config.head_dim
-    num_kv_heads = config.num_key_value_heads
-    num_layers = config.num_hidden_layers
-
-    min_shapes = prepare_input_shapes_string(1, 1, 0, num_layers, num_kv_heads, head_dim)
-    max_shapes = prepare_input_shapes_string(1, 1024, 1024, num_layers, num_kv_heads, head_dim)
-    opt_shapes = prepare_input_shapes_string(1, 512, 512, num_layers, num_kv_heads, head_dim)
-
-    return min_shapes, max_shapes, opt_shapes
-
-
-def make_input_shapes_profile_for_ep_list(ep_list, model_name_or_path):
-    # Input-shapes-profile will be used in provider-options for ORT session creation.
-    # Provider options (even if {}) are needed for all EPs when we provide for any one of them.
-    # Using empty shapes_profile for non-NvTensorRtRtx EPs.
-    input_shapes_profile_sequence = []
-    for ep in ep_list:
-        if ep == "NvTensorRtRtx":
-            min_shapes, max_shapes, opt_shapes = get_input_shapes_profile(model_name_or_path)
-            input_shapes_profile = {
-                "nv_profile_min_shapes": min_shapes,
-                "nv_profile_max_shapes": max_shapes,
-                "nv_profile_opt_shapes": opt_shapes,
-            }
-            input_shapes_profile_sequence.append(input_shapes_profile)
-        else:
-            input_shapes_profile_sequence.append({})
-
-    return input_shapes_profile_sequence
-
-
 def make_model_input(
     config,
     input_ids_arg,
@@ -414,10 +362,14 @@ def main(args):
     )
 
     input_shapes_profile_data = None
-    if "NvTensorRtRtx" in args.calibration_eps and (args.algo not in ["rtn", "rtn_dq"]):
-        # NvTensorRtRtx EP uses (min, max, opt) profile for dynamic shapes in the model's inputs.
-        input_shapes_profile_data = make_input_shapes_profile_for_ep_list(
-            args.calibration_eps, args.model_name
+    if any(ep in ["NvTensorRtRtx", "trt"] for ep in args.calibration_eps) and (
+        args.algo not in ["rtn", "rtn_dq"]
+    ):
+        # TRT EPs use min/opt/max profiles for dynamic model inputs.
+        input_shapes_profile_data = create_input_shapes_profile(
+            args.model_name,
+            args.calibration_eps,
+            trust_remote_code=args.trust_remote_code,
         )
     print(
         f"\n--Quantize-Script-- input_shapes_profile is None? - {input_shapes_profile_data is None}\n"

modelopt/onnx/quantization/__main__.py

@@ -16,6 +16,7 @@
 """Command-line entrypoint for ONNX PTQ."""
 
 import argparse
+import json
 import os
 
 import numpy as np
@@ -30,6 +31,32 @@
 __all__ = ["main"]
 
 
+def parse_input_shapes_profile(value: str) -> list[dict[str, str]]:
+    """Parse input shapes profile from an inline JSON value or a JSON file path."""
+    try:
+        if os.path.exists(value):
+            validate_file_size(value, 1024 * 1024)
+            with open(value, encoding="utf-8") as profile_file:
+                profile = json.load(profile_file)
+        else:
+            profile = json.loads(value)
+    except (OSError, json.JSONDecodeError) as e:
+        raise argparse.ArgumentTypeError(
+            "input_shapes_profile must be a JSON file path or an inline JSON list of dictionaries"
+        ) from e
+
+    if not isinstance(profile, list) or not all(isinstance(item, dict) for item in profile):
+        raise argparse.ArgumentTypeError("input_shapes_profile must be a JSON list of dictionaries")
+
+    for item in profile:
+        if not all(isinstance(key, str) and isinstance(val, str) for key, val in item.items()):
+            raise argparse.ArgumentTypeError(
+                "input_shapes_profile dictionaries must contain only string keys and string values"
+            )
+
+    return profile
+
+
 def validate_file_size(file_path: str, max_size_bytes: int) -> None:
     """Validate that a file exists and does not exceed the maximum allowed size.
 
@@ -62,6 +89,32 @@ def get_parser() -> argparse.ArgumentParser:
     argparser.add_argument(
         "--onnx_path", required=True, type=str, help="Input onnx model without Q/DQ nodes."
     )
+    argparser.add_argument(
+        "--model_id",
+        required=False,
+        type=str,
+        help=(
+            "Hugging Face model ID, local config directory, or local config.json path used "
+            "to infer EP input shape profiles when --input_shapes_profile is not provided."
+        ),
+    )
+    argparser.add_argument(
+        "--input_shapes_profile",
+        required=False,
+        type=parse_input_shapes_profile,
+        help=(
+            "Input shape profile provider options as an inline JSON list or a path to a JSON file. "
+            "The list must have one dictionary per --calibration_eps entry, in the same order. "
+            'Example: \'[{"nv_profile_min_shapes":"input_ids:1x1",'
+            '"nv_profile_opt_shapes":"input_ids:1x512",'
+            '"nv_profile_max_shapes":"input_ids:1x1024"},{}]\'.'
+        ),
+    )
+    argparser.add_argument(
+        "--trust_remote_code",
+        action="store_true",
+        help="Allow custom code when resolving --model_id with Hugging Face transformers.",
+    )
     argparser.add_argument(
         "--quantize_mode",
         type=str,
@@ -516,6 +569,9 @@ def main():
         autotune_warmup_runs=args.autotune_warmup_runs,
         autotune_timing_runs=args.autotune_timing_runs,
         autotune_trtexec_args=args.autotune_trtexec_args,
+        input_shapes_profile=args.input_shapes_profile,
+        model_id=args.model_id,
+        trust_remote_code=args.trust_remote_code,
     )
 
 

modelopt/onnx/quantization/autotune/benchmark.py

@@ -387,10 +387,10 @@ def _load_plugin_libraries(self):
             self.logger.info(f"Loading TensorRT plugin: {plugin_path}")
 
             try:
-                if hasattr(os, "RTLD_LAZY") and hasattr(os, "RTLD_GLOBAL"):
-                    plugin_handle = ctypes.CDLL(
-                        str(plugin_path), mode=os.RTLD_LAZY | os.RTLD_GLOBAL
-                    )
+                rtld_lazy = getattr(os, "RTLD_LAZY", None)
+                rtld_global = getattr(os, "RTLD_GLOBAL", None)
+                if rtld_lazy is not None and rtld_global is not None:
+                    plugin_handle = ctypes.CDLL(str(plugin_path), mode=rtld_lazy | rtld_global)
                 else:
                     # Fallback for platforms without RTLD flags (e.g., Windows)
                     plugin_handle = ctypes.CDLL(str(plugin_path))

modelopt/onnx/quantization/fp8.py

@@ -18,6 +18,7 @@
 import os
 import tempfile
 import time
+from collections.abc import Sequence
 
 import numpy as np
 import onnx
@@ -184,6 +185,7 @@ def quantize(
     direct_io_types: bool = False,
     opset: int | None = None,
     autotune: bool = False,
+    input_shapes_profile: Sequence[dict[str, str]] | None = None,
     **kwargs,
 ) -> onnx.ModelProto:
     """Applies FP8 GEMM only quantization to an ONNX file.
@@ -230,6 +232,7 @@ def quantize(
             calibration_data_reader,
             calibration_eps,
             calibration_shapes,
+            input_shapes_profile,
         )
         nodes_to_exclude.extend(matmul_nodes_to_exclude)  # type: ignore[union-attr]
         logger.debug(f"Excluding {len(matmul_nodes_to_exclude)} MatMul nodes due to GEMV pattern")
@@ -249,6 +252,7 @@ def quantize(
         calibrate_per_node,
         custom_ops_to_quantize,
         kwargs.get("op_types_needing_output_quant"),
+        input_shapes_profile,
     )
     logger.info(
         f"Quantizable op types in the model: {[t for t in op_types_to_quantize if t in op_types]}"

modelopt/onnx/quantization/graph_utils.py

@@ -17,6 +17,7 @@
 
 import re
 from collections import defaultdict
+from collections.abc import Sequence
 from functools import reduce
 from typing import Any, cast
 
@@ -1038,6 +1039,7 @@ def get_extended_model_outputs(
     intermediate_generated_files: list[str],
     calibration_data_reader: CalibrationDataReader,
     calibration_eps: list[str],
+    input_shapes_profile: Sequence[dict[str, str]] | None = None,
 ) -> dict[str, np.ndarray]:
     """Run one inference step on an onnx model which has some intermediate tensor marked as model outputs.
 
@@ -1069,9 +1071,13 @@ def get_extended_model_outputs(
         extended_onnx_path = onnx_path.replace(".onnx", "_extended.onnx")
         save_onnx(extended_model, extended_onnx_path, save_as_external_data=True)
         intermediate_generated_files.append(extended_onnx_path)
-        session = create_inference_session(extended_onnx_path, calibration_eps)
+        session = create_inference_session(
+            extended_onnx_path, calibration_eps, input_shapes_profile
+        )
     else:
-        session = create_inference_session(extended_model.SerializeToString(), calibration_eps)
+        session = create_inference_session(
+            extended_model.SerializeToString(), calibration_eps, input_shapes_profile
+        )
 
     # Run extended model's inference.
     extended_model_output_names = [output.name for output in session.get_outputs()]
@@ -1088,6 +1094,7 @@ def find_nodes_from_matmul_to_exclude(
     calibration_data_reader: CalibrationDataReader = None,
     calibration_eps: list[str] = ["cpu", "cuda:0", "trt"],
     calibration_shapes: str | dict | None = None,
+    input_shapes_profile: Sequence[dict[str, str]] | None = None,
 ) -> list[str]:
     """Find MatMul nodes that meet gemv or small-gemm conditions and should be excluded.
 
@@ -1139,6 +1146,7 @@ def find_nodes_from_matmul_to_exclude(
             intermediate_generated_files or [],
             calibration_data_reader,
             calibration_eps,
+            input_shapes_profile,
         )
 
     logger.debug(f"Matmul nodes to exclude: {nodes_to_exclude}")
@@ -1356,6 +1364,7 @@ def _exclude_matmuls_by_inference(
     intermediate_generated_files: list[str],
     calibration_data_reader: CalibrationDataReader,
     calibration_eps: list[str],
+    input_shapes_profile: Sequence[dict[str, str]] | None = None,
 ) -> list[str]:
     """Use actual inference to find MatMuls with dimension 1 or small K/N."""
     # Add matmul outputs and second-input outputs to model outputs
@@ -1379,6 +1388,7 @@ def _exclude_matmuls_by_inference(
         intermediate_generated_files,
         calibration_data_reader,
         calibration_eps,
+        input_shapes_profile,
     )
 
     nodes_to_exclude = []
@@ -1421,6 +1431,7 @@ def find_nodes_from_mha_to_exclude(
     intermediate_generated_files: list[str] | None = None,
     calibration_data_reader: CalibrationDataReader = None,
     calibration_eps: list[str] = ["cpu", "cuda:0", "trt"],
+    input_shapes_profile: Sequence[dict[str, str]] | None = None,
 ) -> list[str]:
     """Find MatMul nodes in MHA pattern to exclude.
 
@@ -1481,6 +1492,7 @@ def find_nodes_from_mha_to_exclude(
             intermediate_generated_files,  # type: ignore[arg-type]
             calibration_data_reader,
             calibration_eps,
+            input_shapes_profile,
         )
 
         # For each MHA block,

modelopt/onnx/quantization/int8.py

@@ -18,6 +18,7 @@
 import os
 import tempfile
 import time
+from collections.abc import Sequence
 
 import onnx
 import onnx_graphsurgeon as gs
@@ -139,6 +140,7 @@ def quantize(
     direct_io_types: bool = False,
     opset: int | None = None,
     autotune: bool = False,
+    input_shapes_profile: Sequence[dict[str, str]] | None = None,
     **kwargs,
 ) -> onnx.ModelProto:
     """Applies INT8 quantization to an ONNX file using the compiler friendly heuristics.
@@ -177,6 +179,7 @@ def quantize(
             calibration_data_reader,
             calibration_eps,
             calibration_shapes,
+            input_shapes_profile,
         )
         nodes_to_exclude.extend(matmul_nodes_to_exclude)  # type: ignore[union-attr]
         logger.debug(f"Excluding {len(matmul_nodes_to_exclude)} MatMul nodes due to GEMV pattern")
@@ -201,6 +204,7 @@ def quantize(
         calibrate_per_node,
         custom_ops_to_quantize,
         kwargs.get("op_types_needing_output_quant"),
+        input_shapes_profile,
     )
     logger.info(f"Quantizable op types: {[t for t in quantizable_op_types if t in op_types]}")