diff --git a/examples/windows/onnx_ptq/genai_llm/quantize.py b/examples/windows/onnx_ptq/genai_llm/quantize.py index 13f6ac80452..369532a1163 100644 --- a/examples/windows/onnx_ptq/genai_llm/quantize.py +++ b/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" diff --git a/modelopt/onnx/quantization/__main__.py b/modelopt/onnx/quantization/__main__.py index 1cbcaf857c0..edf05df30e3 100644 --- a/modelopt/onnx/quantization/__main__.py +++ b/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, ) diff --git a/modelopt/onnx/quantization/autotune/benchmark.py b/modelopt/onnx/quantization/autotune/benchmark.py index b87478a1572..ba5cf1142bf 100644 --- a/modelopt/onnx/quantization/autotune/benchmark.py +++ b/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)) diff --git a/modelopt/onnx/quantization/fp8.py b/modelopt/onnx/quantization/fp8.py index b7146173a25..b31d80fa70b 100755 --- a/modelopt/onnx/quantization/fp8.py +++ b/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]}" diff --git a/modelopt/onnx/quantization/graph_utils.py b/modelopt/onnx/quantization/graph_utils.py index 70e3ab0abd2..1087a8889aa 100755 --- a/modelopt/onnx/quantization/graph_utils.py +++ b/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, diff --git a/modelopt/onnx/quantization/int8.py b/modelopt/onnx/quantization/int8.py index 170446dab5a..c3f266d6193 100755 --- a/modelopt/onnx/quantization/int8.py +++ b/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]}") diff --git a/modelopt/onnx/quantization/ort_utils.py b/modelopt/onnx/quantization/ort_utils.py index f7799c634f0..a8d211b32c7 100755 --- a/modelopt/onnx/quantization/ort_utils.py +++ b/modelopt/onnx/quantization/ort_utils.py @@ -180,6 +180,10 @@ def _load_extra_cudnn_dlls(): This scans the nvidia-cudnn bin directory and loads any cudnn*.dll not already loaded in the process. """ + # Fix github code quality test failure + if sys.platform != "win32": + return + import ctypes import ctypes.wintypes @@ -195,7 +199,7 @@ def _load_extra_cudnn_dlls(): logger.debug("No cudnn*.dll files found in %s", cudnn_bin_dir) return - get_module_handle_w = ctypes.windll.kernel32.GetModuleHandleW # type: ignore[attr-defined] + get_module_handle_w = ctypes.windll.kernel32.GetModuleHandleW get_module_handle_w.argtypes = [ctypes.wintypes.LPCWSTR] get_module_handle_w.restype = ctypes.wintypes.HMODULE @@ -314,30 +318,54 @@ def _check_for_nv_tensorrt_rtx_libs(): return found -def _prepare_ep_list(calibration_eps: list[str]): +def _prepare_ep_list( + calibration_eps: list[str], + input_shapes_profile: Sequence[dict[str, str]] | None = None, +): """Prepares the EP list for ORT from the given user input.""" logger.debug(f"Preparing execution providers list from: {calibration_eps}") + if input_shapes_profile is not None: + assert len(input_shapes_profile) == len(calibration_eps), ( + "Number of calibration EPs and number of input-shapes-profile don't match" + ) + + def _append_provider( + providers: list[str | tuple[str, dict]], + ep_index: int, + provider_name: str, + provider_options: dict | None = None, + ) -> None: + if not isinstance(provider_name, str): + raise TypeError("provider_name must be a string") + if provider_options is not None and not isinstance(provider_options, dict): + raise TypeError("provider_options must be a dictionary") + + profile = input_shapes_profile[ep_index] if input_shapes_profile is not None else {} + options = dict(provider_options or {}) + options.update(profile) + providers.append((provider_name, options) if options else provider_name) + providers: list[str | tuple[str, dict]] = [] - for ep in calibration_eps: + for i, ep in enumerate(calibration_eps): if "cuda" in ep: try: _check_for_libcudnn() device_id = int(ep.split(":")[1]) if ":" in ep else 0 - providers.append(("CUDAExecutionProvider", {"device_id": device_id})) + _append_provider(providers, i, "CUDAExecutionProvider", {"device_id": device_id}) logger.debug(f"Added CUDA EP with device_id: {device_id}") except Exception as e: logger.warning(f"Failed to enable ORT with CUDA EP: '{e}'") elif "dml" in ep: device_id = int(ep.split(":")[1]) if ":" in ep else 0 - providers.append(("DmlExecutionProvider", {"device_id": device_id})) + _append_provider(providers, i, "DmlExecutionProvider", {"device_id": device_id}) logger.debug(f"Added DML EP with device_id: {device_id}") elif "cpu" in ep: - providers.append("CPUExecutionProvider") + _append_provider(providers, i, "CPUExecutionProvider") logger.debug("Added CPU EP") elif "NvTensorRtRtx" in ep: try: _check_for_nv_tensorrt_rtx_libs() - providers.append("NvTensorRTRTXExecutionProvider") + _append_provider(providers, i, "NvTensorRTRTXExecutionProvider") logger.debug("Added NvTensorRtRtx EP") except Exception as e: logger.warning(f"Failed to enable ORT with NvTensorRtRtx EP: '{e}'") @@ -345,7 +373,7 @@ def _prepare_ep_list(calibration_eps: list[str]): try: _check_for_tensorrt() _check_for_libcudnn() - providers.append("TensorrtExecutionProvider") + _append_provider(providers, i, "TensorrtExecutionProvider") logger.debug("Added TensorRT EP") except Exception as e: logger.warning(f"Failed to enable ORT with TensorRT EP: '{e}'") @@ -420,6 +448,100 @@ def _make_trt_ep_first_choice(calibration_eps, trt_plugins): return trt_plugins +def create_input_shapes_profile( + model_id: str, calibration_eps: list[str], trust_remote_code: bool = False +) -> list[dict[str, str]]: + """Create per-EP input shape profiles from a Hugging Face config. + + ``model_id`` can be a Hugging Face model ID, local config directory, or local + ``config.json`` path. + The returned list matches ``calibration_eps`` order. EPs that do not use input shape + profiles receive an empty dictionary. + """ + from transformers import AutoConfig + + def empty_profiles() -> list[dict[str, str]]: + return [{} for _ in calibration_eps] + + def warn_and_return_empty(reason: str) -> list[dict[str, str]]: + logger.warning( + f"Could not create input shape profiles from model_id={model_id!r}: {reason}. " + "Falling back to empty input shape profiles. Some TensorRT/NvTensorRtRtx EP " + "versions can infer shapes automatically; if session creation fails, pass " + "input_shapes_profile manually." + ) + return empty_profiles() + + def get_config_attr(config, aliases: list[str], field_name: str): + for alias in aliases: + value = getattr(config, alias, None) + if value is not None: + return value + raise ValueError( + f"missing required config field for {field_name}; tried aliases: {', '.join(aliases)}" + ) + + try: + config = AutoConfig.from_pretrained(model_id, trust_remote_code=trust_remote_code) + except Exception as exc: + return warn_and_return_empty(str(exc)) + + try: + num_attention_heads = get_config_attr( + config, ["num_attention_heads", "n_head", "num_heads"], "num_attention_heads" + ) + head_dim = getattr(config, "head_dim", None) + if head_dim is None: + hidden_size = get_config_attr( + config, ["hidden_size", "n_embd", "d_model"], "hidden_size" + ) + head_dim = hidden_size // num_attention_heads + num_kv_heads = getattr(config, "num_key_value_heads", None) or num_attention_heads + num_layers = get_config_attr( + config, ["num_hidden_layers", "n_layer", "num_layers"], "num_hidden_layers" + ) + except (AttributeError, TypeError, ValueError, ZeroDivisionError) as exc: + return warn_and_return_empty(str(exc)) + + def make_shapes(batch_size: int, seq_len: int, past_seq_len: int) -> str: + shapes = [f"input_ids:{batch_size}x{seq_len}", f"attention_mask:{batch_size}x{seq_len}"] + for layer_idx in range(num_layers): + shapes.extend( + [ + f"past_key_values.{layer_idx}.key:{batch_size}x{num_kv_heads}x{past_seq_len}x{head_dim}", + f"past_key_values.{layer_idx}.value:{batch_size}x{num_kv_heads}x{past_seq_len}x{head_dim}", + ] + ) + return ",".join(shapes) + + min_shapes = make_shapes(batch_size=1, seq_len=1, past_seq_len=0) + opt_shapes = make_shapes(batch_size=1, seq_len=512, past_seq_len=512) + max_shapes = make_shapes(batch_size=1, seq_len=1024, past_seq_len=1024) + + profiles: list[dict[str, str]] = [] + for ep in calibration_eps: + if "NvTensorRtRtx" in ep: + profiles.append( + { + "nv_profile_min_shapes": min_shapes, + "nv_profile_opt_shapes": opt_shapes, + "nv_profile_max_shapes": max_shapes, + } + ) + elif ep == "trt": + profiles.append( + { + "trt_profile_min_shapes": min_shapes, + "trt_profile_opt_shapes": opt_shapes, + "trt_profile_max_shapes": max_shapes, + } + ) + else: + profiles.append({}) + + return profiles + + def create_inference_session( onnx_path_or_model: str | bytes, calibration_eps: list[str], @@ -445,13 +567,12 @@ def create_inference_session( logger.debug( f"Input-Shapes-Profile: EP: {calibration_eps[i]}, key: {k}, value: {v}" ) - providers = _prepare_ep_list(calibration_eps) + providers = _prepare_ep_list(calibration_eps, input_shapes_profile) logger.debug(f"Creating session with providers: {providers}") return ort.InferenceSession( onnx_path_or_model, sess_options=sess_options, providers=providers, - provider_options=None if input_shapes_profile is None else input_shapes_profile, ) @@ -482,9 +603,12 @@ def configure_ort( calibrate_per_node: bool = False, custom_ops_to_quantize: list[str] = [], op_types_needing_output_quant: list[str] | None = None, + input_shapes_profile: Sequence[dict[str, str]] | None = None, ): """Configure and patches ORT to support ModelOpt ONNX quantization.""" logger.info("Configuring ORT for ModelOpt ONNX quantization") + if calibration_eps is None: + calibration_eps = ["cpu", "cuda:0", "trt"] # Register custom QDQ operators logger.debug("Registering custom QDQ operators") @@ -538,6 +662,8 @@ def configure_ort( ] if trt_extra_plugin_lib_paths is not None: trt_extra_plugin_lib_paths = ";".join(trt_extra_plugin_lib_paths) + execution_providers = _prepare_ep_list(calibration_eps, input_shapes_profile) + trt_guided_options = { "QuantizeBias": False, "ActivationSymmetric": True, @@ -554,7 +680,7 @@ def configure_ort( True ), "TrtExtraPluginLibraryPaths": trt_extra_plugin_lib_paths, - "ExecutionProviders": _prepare_ep_list(calibration_eps), # type: ignore[arg-type] + "ExecutionProviders": execution_providers, } quantizable_op_types = get_quantizable_op_types(op_types_to_quantize) diff --git a/modelopt/onnx/quantization/quantize.py b/modelopt/onnx/quantization/quantize.py index 6bbc3299129..d8b2471f127 100755 --- a/modelopt/onnx/quantization/quantize.py +++ b/modelopt/onnx/quantization/quantize.py @@ -63,7 +63,7 @@ ) from modelopt.onnx.quantization.int4 import quantize as quantize_int4 from modelopt.onnx.quantization.int8 import quantize as quantize_int8 -from modelopt.onnx.quantization.ort_utils import update_trt_ep_support +from modelopt.onnx.quantization.ort_utils import create_input_shapes_profile, update_trt_ep_support from modelopt.onnx.quantization.qdq_utils import ( qdq_to_dq, remove_graph_input_q, @@ -94,6 +94,25 @@ def _normalize_quantize_mode_for_opset(quantize_mode: str) -> str: return quantize_mode +def _realign_input_shapes_profile( + input_shapes_profile: Sequence[dict[str, str]], + original_calibration_eps: list[str], + calibration_eps: list[str], +) -> Sequence[dict[str, str]]: + """Keep per-EP profiles aligned after ``calibration_eps`` is updated.""" + assert len(input_shapes_profile) == len(original_calibration_eps), ( + "Number of calibration EPs and number of input-shapes-profile don't match" + ) + assert len(set(original_calibration_eps)) == len(original_calibration_eps), ( + "Calibration EPs must be unique when input_shapes_profile is provided" + ) + if original_calibration_eps == calibration_eps: + return input_shapes_profile + + profiles_by_ep = dict(zip(original_calibration_eps, input_shapes_profile, strict=True)) + return [profiles_by_ep.get(ep, {}) for ep in calibration_eps] + + def _preprocess_onnx( onnx_path: str, use_external_data_format: bool, @@ -358,6 +377,8 @@ def quantize( simplify: bool = False, calibrate_per_node: bool = False, input_shapes_profile: Sequence[dict[str, str]] | None = None, + model_id: str | None = None, + trust_remote_code: bool = False, direct_io_types: bool = False, opset: int | None = None, target_dla: bool = False, @@ -492,6 +513,12 @@ def quantize( If None of the calibration_eps require any such shapes profile for model inputs, then nothing needs to be set for this "input_shapes_profile" parameter. Default value is None. + model_id: + Hugging Face model ID, local config directory, or local ``config.json`` path used to infer input shape + profiles when ``input_shapes_profile`` is not provided. + trust_remote_code: + Whether to allow custom code to be loaded when resolving ``model_id`` with Hugging Face transformers. + Defaults to False. direct_io_types: If True, modify the I/O types in the quantized ONNX model to be lower precision whenever possible. If False, keep the I/O types in the quantized ONNX model the same as in the given ONNX model. @@ -603,8 +630,18 @@ def quantize( quantize_mode, opset, ) + original_calibration_eps = list(calibration_eps) trt_plugins = update_trt_ep_support(calibration_eps, has_dds_op, has_custom_op, trt_plugins) # type: ignore[arg-type] + if input_shapes_profile is not None: + input_shapes_profile = _realign_input_shapes_profile( + input_shapes_profile, original_calibration_eps, calibration_eps + ) + elif model_id: + input_shapes_profile = create_input_shapes_profile( + model_id, calibration_eps, trust_remote_code=trust_remote_code + ) + if calibration_data_reader is None: # Use random scales if calibration data is not supplied if calibration_data is None: @@ -635,6 +672,7 @@ def quantize( intermediate_generated_files, calibration_data_reader, calibration_eps, + input_shapes_profile, ) if calibrate_per_node and not calibration_shapes: @@ -697,6 +735,7 @@ def quantize( direct_io_types=direct_io_types, opset=opset, autotune=autotune, + input_shapes_profile=input_shapes_profile, **kwargs, ) diff --git a/tests/unit/onnx/quantization/test_autotune_quantization_integration.py b/tests/unit/onnx/quantization/test_autotune_quantization_integration.py index 52e65158066..aa6a27c1eb7 100644 --- a/tests/unit/onnx/quantization/test_autotune_quantization_integration.py +++ b/tests/unit/onnx/quantization/test_autotune_quantization_integration.py @@ -14,6 +14,7 @@ # limitations under the License. import importlib +import json import sys import pytest @@ -36,3 +37,72 @@ def test_quantization_cli_parser_imports_without_tensorrt(): args = parser.parse_args(["--onnx_path", "dummy.onnx"]) assert args.onnx_path == "dummy.onnx" assert args.quantize_mode == "int8" + + +def test_quantization_cli_parses_inline_input_shapes_profile(): + from modelopt.onnx.quantization.__main__ import get_parser + + profile = [{"nv_profile_min_shapes": "input_ids:1x1"}, {}] + args = get_parser().parse_args( + [ + "--onnx_path", + "dummy.onnx", + "--input_shapes_profile", + json.dumps(profile), + ] + ) + + assert args.input_shapes_profile == profile + + +def test_quantization_cli_parses_input_shapes_profile_file(tmp_path): + from modelopt.onnx.quantization.__main__ import get_parser + + profile = [{"trt_profile_min_shapes": "input_ids:1x1"}, {}] + profile_path = tmp_path / "profile.json" + profile_path.write_text(json.dumps(profile), encoding="utf-8") + + args = get_parser().parse_args( + [ + "--onnx_path", + "dummy.onnx", + "--input_shapes_profile", + str(profile_path), + ] + ) + + assert args.input_shapes_profile == profile + + +def test_quantization_cli_forwards_input_shapes_profile(monkeypatch, tmp_path): + import modelopt.onnx.quantization.__main__ as quantization_cli + + profile = [{"nv_profile_min_shapes": "input_ids:1x1"}, {}] + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"onnx") + captured = {} + + def fake_quantize(onnx_path_arg, **kwargs): + captured["onnx_path"] = onnx_path_arg + captured.update(kwargs) + + monkeypatch.setattr(quantization_cli, "quantize", fake_quantize) + monkeypatch.setattr( + sys, + "argv", + [ + "modelopt.onnx.quantization", + "--onnx_path", + str(onnx_path), + "--calibration_eps", + "NvTensorRtRtx", + "cpu", + "--input_shapes_profile", + json.dumps(profile), + ], + ) + + quantization_cli.main() + + assert captured["onnx_path"] == str(onnx_path) + assert captured["input_shapes_profile"] == profile diff --git a/tests/unit/onnx/quantization/test_ort_utils.py b/tests/unit/onnx/quantization/test_ort_utils.py new file mode 100644 index 00000000000..ee4f81f5645 --- /dev/null +++ b/tests/unit/onnx/quantization/test_ort_utils.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import sys +import types + +from modelopt.onnx.quantization import ort_utils +from modelopt.onnx.quantization.ort_utils import create_input_shapes_profile + + +def _raise_trt_unavailable(): + raise RuntimeError("trt unavailable") + + +def test_create_input_shapes_profile_forwards_trust_remote_code(monkeypatch): + calls = [] + + class FakeAutoConfig: + @staticmethod + def from_pretrained(model_id, **kwargs): + calls.append((model_id, kwargs)) + return types.SimpleNamespace( + hidden_size=128, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=1, + ) + + monkeypatch.setitem( + sys.modules, "transformers", types.SimpleNamespace(AutoConfig=FakeAutoConfig) + ) + + default_profiles = create_input_shapes_profile("local-config.json", ["NvTensorRtRtx", "cpu"]) + trusted_profiles = create_input_shapes_profile( + "custom-model", ["NvTensorRtRtx"], trust_remote_code=True + ) + + assert calls == [ + ("local-config.json", {"trust_remote_code": False}), + ("custom-model", {"trust_remote_code": True}), + ] + assert default_profiles[0]["nv_profile_min_shapes"].startswith("input_ids:1x1") + assert default_profiles[1] == {} + assert trusted_profiles[0]["nv_profile_opt_shapes"].startswith("input_ids:1x512") + + +def test_create_input_shapes_profile_returns_empty_for_bad_model_id(monkeypatch, caplog): + class FakeAutoConfig: + @staticmethod + def from_pretrained(model_id, **kwargs): + raise OSError(f"cannot load {model_id}") + + monkeypatch.setitem( + sys.modules, "transformers", types.SimpleNamespace(AutoConfig=FakeAutoConfig) + ) + caplog.set_level(logging.WARNING, logger="modelopt.onnx") + + profiles = create_input_shapes_profile("bad-model", ["NvTensorRtRtx", "cpu"]) + + assert profiles == [{}, {}] + assert "bad-model" in caplog.text + assert "input_shapes_profile manually" in caplog.text + + +def test_create_input_shapes_profile_uses_common_config_aliases(monkeypatch): + class FakeAutoConfig: + @staticmethod + def from_pretrained(model_id, **kwargs): + return types.SimpleNamespace(n_embd=96, n_head=6, n_layer=2) + + monkeypatch.setitem( + sys.modules, "transformers", types.SimpleNamespace(AutoConfig=FakeAutoConfig) + ) + + profiles = create_input_shapes_profile("gpt-style-config", ["NvTensorRtRtx"]) + + assert "past_key_values.1.key:1x6x0x16" in profiles[0]["nv_profile_min_shapes"] + assert "past_key_values.1.value:1x6x512x16" in profiles[0]["nv_profile_opt_shapes"] + + +def test_create_input_shapes_profile_head_dim_does_not_require_hidden_size(monkeypatch): + class FakeAutoConfig: + @staticmethod + def from_pretrained(model_id, **kwargs): + return types.SimpleNamespace( + head_dim=32, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=1, + ) + + monkeypatch.setitem( + sys.modules, "transformers", types.SimpleNamespace(AutoConfig=FakeAutoConfig) + ) + + profiles = create_input_shapes_profile("head-dim-config", ["trt"]) + + assert "past_key_values.0.key:1x2x0x32" in profiles[0]["trt_profile_min_shapes"] + + +def test_create_input_shapes_profile_returns_empty_for_missing_config_key(monkeypatch, caplog): + class FakeAutoConfig: + @staticmethod + def from_pretrained(model_id, **kwargs): + return types.SimpleNamespace(num_attention_heads=4, num_hidden_layers=1) + + monkeypatch.setitem( + sys.modules, "transformers", types.SimpleNamespace(AutoConfig=FakeAutoConfig) + ) + caplog.set_level(logging.WARNING, logger="modelopt.onnx") + + profiles = create_input_shapes_profile("missing-hidden-size", ["NvTensorRtRtx"]) + + assert profiles == [{}] + assert "missing-hidden-size" in caplog.text + assert "hidden_size" in caplog.text + + +def test_prepare_ep_list_keeps_profiles_aligned_when_ep_is_disabled(monkeypatch): + monkeypatch.setattr(ort_utils, "_check_for_tensorrt", _raise_trt_unavailable) + monkeypatch.setattr(ort_utils, "_check_for_nv_tensorrt_rtx_libs", lambda: True) + + providers = ort_utils._prepare_ep_list( + ["trt", "NvTensorRtRtx", "cpu"], + [ + {"trt_profile_min_shapes": "trt_profile"}, + {"nv_profile_min_shapes": "rtx_profile"}, + {}, + ], + ) + + assert providers == [ + ("NvTensorRTRTXExecutionProvider", {"nv_profile_min_shapes": "rtx_profile"}), + "CPUExecutionProvider", + ] + + +def test_create_inference_session_filters_profile_with_disabled_ep(monkeypatch): + captured_kwargs = {} + + class FakeInferenceSession: + def __init__(self, *args, **kwargs): + captured_kwargs.update(kwargs) + + monkeypatch.setattr(ort_utils, "_check_for_tensorrt", _raise_trt_unavailable) + monkeypatch.setattr(ort_utils.ort, "InferenceSession", FakeInferenceSession) + + ort_utils.create_inference_session( + "model.onnx", + ["trt", "cpu"], + [{"trt_profile_min_shapes": "trt_profile"}, {}], + ) + + assert captured_kwargs["providers"] == ["CPUExecutionProvider"] + assert "provider_options" not in captured_kwargs diff --git a/tests/unit/onnx/quantization/test_quantize_api.py b/tests/unit/onnx/quantization/test_quantize_api.py index 464fb1a88b9..f350d5d89f4 100644 --- a/tests/unit/onnx/quantization/test_quantize_api.py +++ b/tests/unit/onnx/quantization/test_quantize_api.py @@ -13,8 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for ONNX quantization opset handling.""" +"""Tests for ONNX quantization API handling.""" +import importlib import os import onnx @@ -50,6 +51,97 @@ ] +def test_realign_input_shapes_profile_after_calibration_eps_update(): + quantize_module = importlib.import_module("modelopt.onnx.quantization.quantize") + + profiles = quantize_module._realign_input_shapes_profile( + [{"cpu_profile": "cpu"}, {"trt_profile": "trt"}], + ["cpu", "trt"], + ["trt", "cpu"], + ) + + assert profiles == [{"trt_profile": "trt"}, {"cpu_profile": "cpu"}] + + +def test_realign_input_shapes_profile_rejects_duplicate_calibration_eps(): + quantize_module = importlib.import_module("modelopt.onnx.quantization.quantize") + + with pytest.raises(AssertionError, match="Calibration EPs must be unique"): + quantize_module._realign_input_shapes_profile( + [{"cpu_profile": "first"}, {"cpu_profile": "second"}], + ["cpu", "cpu"], + ["cpu"], + ) + + +def test_quantize_infers_input_profiles_after_ep_support_update(monkeypatch, tmp_path): + quantize_module = importlib.import_module("modelopt.onnx.quantization.quantize") + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"") + captured = {} + + def fake_preprocess( + onnx_path, + use_external_data_format, + output_path, + enable_shared_constants_duplication, + trt_plugins, + trt_plugins_precision, + override_shapes, + simplify, + quantize_mode, + opset, + ): + return onnx_path, object(), [], True, False, False, {}, {} + + def fake_update_trt_ep_support(calibration_eps, has_dds_op, has_custom_op, trt_plugins): + assert has_custom_op is True + calibration_eps.remove("trt") + calibration_eps.insert(0, "trt") + return trt_plugins + + def fake_create_input_shapes_profile(model_id, calibration_eps, trust_remote_code=False): + captured["profile_eps"] = list(calibration_eps) + captured["trust_remote_code"] = trust_remote_code + return [{"trt_profile_min_shapes": "trt_profile"}, {}] + + def fake_find_nodes_from_mha_to_exclude(*args): + captured["find_eps"] = list(args[-2]) + captured["find_profile"] = args[-1] + return [] + + def fake_quantize_int8(**kwargs): + captured["quantize_eps"] = list(kwargs["calibration_eps"]) + captured["quantize_profile"] = kwargs["input_shapes_profile"] + + monkeypatch.setattr(quantize_module, "_preprocess_onnx", fake_preprocess) + monkeypatch.setattr(quantize_module, "update_trt_ep_support", fake_update_trt_ep_support) + monkeypatch.setattr( + quantize_module, "create_input_shapes_profile", fake_create_input_shapes_profile + ) + monkeypatch.setattr( + quantize_module, "find_nodes_from_mha_to_exclude", fake_find_nodes_from_mha_to_exclude + ) + monkeypatch.setattr(quantize_module, "validate_op_types_spelling", lambda *args: None) + monkeypatch.setattr(quantize_module, "quantize_int8", fake_quantize_int8) + monkeypatch.setattr(quantize_module.onnx.checker, "check_model", lambda *args: None) + + quantize_module.quantize( + str(onnx_path), + calibration_eps=["cpu", "trt"], + calibration_data_reader=object(), + model_id="local-config", + trust_remote_code=True, + ) + + assert captured["profile_eps"] == ["trt", "cpu"] + assert captured["trust_remote_code"] is True + assert captured["find_eps"] == ["trt", "cpu"] + assert captured["quantize_eps"] == ["trt", "cpu"] + assert captured["find_profile"] == [{"trt_profile_min_shapes": "trt_profile"}, {}] + assert captured["quantize_profile"] == [{"trt_profile_min_shapes": "trt_profile"}, {}] + + @pytest.mark.parametrize("quant_mode", ["int8", "fp8", "int4"]) @pytest.mark.parametrize( ("scenario_name", "export_opset_offset", "request_opset_offset", "expected_opset_offset"),