Normalize calibrated ONNX request contracts

CHANGELOG.rst

@@ -23,6 +23,12 @@ Changelog
 
 **Backward Breaking Changes**
 
+- ONNX ``quantize`` now accepts only the exact ``int8``, ``fp8``, and ``int4`` mode
+  names, and INT8/FP8 require exactly one explicit ``calibration_data`` or
+  ``calibration_data_reader`` source. Implicit random calibration and
+  ``calibration_cache_path`` / ``--calibration_cache_path`` were removed; omitted
+  calibration methods default to ``entropy`` for INT8 and ``max`` for FP8.
+
 - The ``modelopt.onnx.quantization.graph_utils`` module has been removed with no
   compatibility shim; update direct imports using this migration map:
 

modelopt/onnx/quantization/__main__.py

@@ -31,6 +31,12 @@
 __all__ = ["main"]
 
 
+def _removed_calibration_cache_path(value: str) -> str:
+    raise argparse.ArgumentTypeError(
+        "--calibration_cache_path was removed; use --calibration_data_path"
+    )
+
+
 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:
@@ -127,14 +133,14 @@ def get_parser() -> argparse.ArgumentParser:
         type=str,
         choices=["max", "entropy", "awq_clip", "rtn_dq"],
         help=(
-            "Calibration method choices for int8/fp8: {entropy (default), max}, "
-            "int4: {awq_clip (default), rtn_dq}."
+            "Calibration method choices for int8: {entropy (default), max}; "
+            "fp8: {max (default), entropy}; int4: {awq_clip (default), rtn_dq}."
         ),
     )
     group.add_argument(
         "--calibration_data_path",
         type=str,
-        help="Calibration data in npz/npy format. If None, random data for calibration will be used.",
+        help="Calibration data in npz/npy format. Required for int8 and fp8.",
     )
     group.add_argument(
         "--trust_calibration_data",
@@ -143,9 +149,10 @@ def get_parser() -> argparse.ArgumentParser:
     )
     group.add_argument(
         "--calibration_cache_path",
-        type=str,
-        help="Pre-calculated activation tensor scaling factors aka calibration cache path.",
+        type=_removed_calibration_cache_path,
+        help=argparse.SUPPRESS,
     )
+
     argparser.add_argument(
         "--calibration_shapes",
         type=str,
@@ -544,7 +551,6 @@ def main():
         quantize_mode=args.quantize_mode,
         calibration_data=calibration_data,
         calibration_method=args.calibration_method,
-        calibration_cache_path=args.calibration_cache_path,
         calibration_shapes=args.calibration_shapes,
         calibration_eps=args.calibration_eps,
         trt_rtx_backend=args.trt_rtx_backend,

modelopt/onnx/quantization/calib_utils.py

@@ -15,7 +15,6 @@
 
 """Provides basic calibration utils."""
 
-import struct
 from typing import TypeAlias
 
 import numpy as np
@@ -149,30 +148,3 @@ def get_first(self):
     def rewind(self):
         """Rewinds the data reader to the first index."""
         self.calibration_data_reader = iter(self.calibration_data_list)
-
-
-def import_scales_from_calib_cache(cache_path: str) -> dict[str, float]:
-    """Reads TensorRT calibration cache and returns as dictionary.
-
-    Args:
-        cache_path: Calibration cache path.
-
-    Returns:
-        Dictionary with scales in the format {tensor_name: float_scale}.
-    """
-    logger.info(f"Importing scales from calibration cache: {cache_path}")
-    with open(cache_path) as f:
-        scales_dict = {}
-        lines = f.readlines()
-        for i, line in enumerate(lines):
-            if i > 0:  # Skips the first line (i.e., TRT-8501-EntropyCalibration2)
-                layer_name, hex_value = line.replace("\n", "").split(": ")
-                try:
-                    scale = struct.unpack("!f", bytes.fromhex(hex_value))[0]
-                    scales_dict[layer_name + "_scale"] = scale
-                    logger.debug(f"Imported scale for {layer_name}: {scale}")
-                except Exception as e:
-                    logger.error(f"Failed to parse scale for tensor {layer_name}: {e!s}")
-                    raise ValueError(f"Scale value for tensor {layer_name} was not found!")
-
-        return scales_dict

modelopt/onnx/quantization/fp8.py

@@ -140,7 +140,6 @@ def quantize(
     onnx_path: str,
     calibration_method: str = "entropy",
     calibration_data_reader: CalibrationDataReader = None,
-    calibration_cache_path: str | None = None,
     calibration_shapes: str | dict | None = None,
     calibration_eps: list[str] = ["cpu", "cuda:0", "trt"],
     op_types_to_quantize: list[str] | None = None,
@@ -168,6 +167,10 @@ def quantize(
 
     Currently, ['Conv', 'Gemm', 'MatMul', 'Residual-Add'] quantization is supported.
     """
+    if "calibration_cache_path" in kwargs:
+        raise TypeError(
+            "calibration_cache_path was removed; provide calibration_data_reader instead"
+        )
     configure_logging(level=log_level.upper())
     logger.info("Starting FP8 quantization process")
     t_start = time.time()

modelopt/onnx/quantization/int8.py

@@ -28,7 +28,6 @@
 from onnxruntime.quantization.calibrate import CalibrationDataReader
 
 from modelopt.onnx.logging_config import configure_logging, logger
-from modelopt.onnx.quantization.calib_utils import import_scales_from_calib_cache
 from modelopt.onnx.quantization.graph_indexing import expand_node_names_from_patterns
 from modelopt.onnx.quantization.graph_selection import (
     find_nodes_from_convs_to_exclude,
@@ -53,7 +52,7 @@
     get_concat_eliminated_tensors,
     remove_partial_input_qdq,
 )
-from modelopt.onnx.quantization.qdq_utils import has_qdq_nodes, replace_scale_values
+from modelopt.onnx.quantization.qdq_utils import has_qdq_nodes
 
 
 def _find_nodes_to_quantize(
@@ -122,7 +121,6 @@ def quantize(
     onnx_path: str,
     calibration_method: str = "entropy",
     calibration_data_reader: CalibrationDataReader = None,
-    calibration_cache_path: str | None = None,
     calibration_shapes: str | dict | None = None,
     calibration_eps: list[str] = ["cpu", "cuda:0", "trt"],
     op_types_to_quantize: list[str] | None = None,
@@ -150,6 +148,10 @@ def quantize(
     Quantization of ['Add', 'AveragePool', 'BatchNormalization', 'Clip', 'Conv', 'ConvTranspose',
     'Gemm', 'GlobalAveragePool', 'MatMul', 'MaxPool', 'Mul'] op types are supported.
     """
+    if "calibration_cache_path" in kwargs:
+        raise TypeError(
+            "calibration_cache_path was removed; provide calibration_data_reader instead"
+        )
     configure_logging(level=log_level.upper())
     logger.info(f"Starting INT8 quantization with method: {calibration_method}")
     t_start = time.time()
@@ -230,23 +232,6 @@ def quantize(
                 )
                 nodes_to_quantize = [node.name for node in quantizable_nodes]
 
-    # Read the calibration cache and quantize nodes for which activation scale values are cached
-    if calibration_cache_path:
-        act_scales_dict = import_scales_from_calib_cache(calibration_cache_path)
-        logger.info(f"Using calibration cache from {calibration_cache_path}")
-        iq_quantized_nodes = []
-        quantized_tensors = [tensor_name.replace("_scale", "") for tensor_name in act_scales_dict]
-        for node in graph.nodes:
-            iq_quantized_nodes.extend(
-                [node.name for node_input in node.inputs if node_input.name in quantized_tensors]
-            )
-
-        logger.info(
-            f"Skipping quantization of nodes: {set(nodes_to_quantize) - set(iq_quantized_nodes)}"
-        )
-        if not autotune:
-            nodes_to_quantize = list(set(nodes_to_quantize).intersection(iq_quantized_nodes))
-
     # Update the list of nodes to quantize
     nodes_to_quantize = [
         node_name for node_name in nodes_to_quantize if node_name not in nodes_to_exclude
@@ -301,9 +286,6 @@ def quantize(
         remove_partial_input_qdq(graph, no_quantize_inputs)
         onnx_model = gs.export_onnx(graph)
 
-        if calibration_cache_path:
-            replace_scale_values(onnx_model.graph, act_scales_dict)
-
     onnx_model = _convert_to_runtime_precision(
         onnx_model,
         quantize_mode="int8",

modelopt/onnx/quantization/qdq_utils.py

@@ -479,37 +479,6 @@ def _insert_helper(
     )
 
 
-def replace_scale_values(graph: onnx.GraphProto, act_scales_dict: dict[str, float]) -> None:
-    """Replace scale values in the graph with values from calibration cache.
-
-    Args:
-        graph: ONNX graph to modify
-        act_scales_dict: Dictionary mapping scale tensor names to their new values
-    """
-    logger.debug(f"Replacing scale values for {len(act_scales_dict)} tensors")
-    initializer_indices = {init.name: idx for idx, init in enumerate(graph.initializer)}
-
-    for node in graph.node:
-        if node.op_type != "QuantizeLinear":
-            continue
-
-        scale_name = node.input[1]
-        if scale_name in act_scales_dict:
-            if scale_name not in initializer_indices:
-                raise ValueError(f"Scale tensor '{scale_name}' not found in graph initializers")
-
-            scale = onnx.numpy_helper.from_array(
-                np.float32(act_scales_dict[scale_name]), scale_name
-            )
-            graph.initializer[initializer_indices[scale_name]].CopyFrom(scale)
-            logger.debug(f"Updated scale value for {scale_name}")
-        else:
-            # For weight quantizers, verify the weight tensor exists
-            weight_name = node.input[0]
-            if weight_name not in initializer_indices:
-                raise ValueError(f"Weight tensor '{weight_name}' not found in graph initializers")
-
-
 def has_qdq_nodes(onnx_model: onnx.ModelProto):
     """Check if the onnx graph already has QDQ nodes."""
     qdq_ops = {QUANTIZE_NODE_NAME, DEQUANTIZE_NODE_NAME}

modelopt/onnx/quantization/quantize.py

@@ -468,10 +468,9 @@ def quantize(
     quantize_mode: str = "int8",
     calibration_data: CalibrationDataType = None,
     calibration_method: str | None = None,
-    calibration_cache_path: str | None = None,
     calibration_data_reader: CalibrationDataReader | None = None,
     calibration_shapes: str | None = None,
-    calibration_eps: list[str] = ["cpu", "cuda:0", "trt"],
+    calibration_eps: Sequence[str] = ("cpu", "cuda:0", "trt"),
     override_shapes: str | None = None,
     op_types_to_quantize: list[str] | None = None,
     op_types_to_exclude: list[str] | None = None,
@@ -491,7 +490,7 @@ def quantize(
     dq_only: bool = False,
     block_size: int | None = None,
     use_zero_point: bool = False,
-    passes: list[str] = ["concat_elimination"],
+    passes: Sequence[str] | None = ("concat_elimination",),
     simplify: bool = False,
     calibrate_per_node: bool = False,
     input_shapes_profile: Sequence[dict[str, str]] | None = None,
@@ -526,10 +525,9 @@ def quantize(
         calibration_data:
             Calibration data, either a numpy array or list/dict of numpy arrays.
         calibration_method:
-            Calibration method choices. Options are int8/fp8: {'entropy' (default), 'max'}
-            and int4: {'awq_clip' (default), 'awq_lite', 'awq_full', 'rtn_dq'}.
-        calibration_cache_path:
-            Path to pre-calculated activation tensor ranges, also known as calibration cache.
+            Calibration method choices. Options are int8: {'entropy' (default), 'max'},
+            fp8: {'max' (default), 'entropy'}, and int4:
+            {'awq_clip' (default), 'awq_lite', 'awq_full', 'rtn_dq'}.
         calibration_data_reader:
             Instance of a CalibrationDataReader object to provide calibration data.
         calibration_shapes:
@@ -693,13 +691,53 @@ def quantize(
         None, writes the quantized onnx model in the supplied output_path
         or writes to the same directory with filename like "<model_name>.quant.onnx".
     """
+    if "calibration_cache_path" in kwargs:
+        raise TypeError(
+            "calibration_cache_path was removed; provide calibration_data or calibration_data_reader instead"
+        )
+    calibration_eps = list(calibration_eps)
+    passes = [] if passes is None else list(passes)
+    op_types_to_quantize = list(op_types_to_quantize) if op_types_to_quantize is not None else None
+    op_types_to_exclude = list(op_types_to_exclude) if op_types_to_exclude is not None else None
+    op_types_to_exclude_fp16 = (
+        list(op_types_to_exclude_fp16) if op_types_to_exclude_fp16 is not None else None
+    )
+    nodes_to_quantize = list(nodes_to_quantize) if nodes_to_quantize is not None else None
+    nodes_to_exclude = list(nodes_to_exclude) if nodes_to_exclude is not None else None
+    trt_plugins = list(trt_plugins) if trt_plugins is not None else None
+    trt_plugins_precision = (
+        list(trt_plugins_precision) if trt_plugins_precision is not None else None
+    )
+    autotune_node_filter_list = (
+        list(autotune_node_filter_list) if autotune_node_filter_list is not None else None
+    )
+    input_shapes_profile = (
+        [dict(profile) for profile in input_shapes_profile]
+        if input_shapes_profile is not None
+        else None
+    )
+
+    if quantize_mode not in {"int8", "fp8", "int4"}:
+        raise ValueError(f"Unsupported quantization mode: {quantize_mode!r}")
+    if quantize_mode in {"int8", "fp8"}:
+        if calibration_method is None:
+            calibration_method = "entropy" if quantize_mode == "int8" else "max"
+        if calibration_method not in {"entropy", "max"}:
+            raise ValueError(
+                f"Unsupported calibration method {calibration_method!r} for {quantize_mode}"
+            )
+
     if trt_rtx_backend not in ("legacy", "abi"):
         raise ValueError(f"trt_rtx_backend must be 'legacy' or 'abi', got {trt_rtx_backend!r}")
     if trt_plugins and "NvTensorRtRtx" in calibration_eps:
         raise ValueError(
             "TensorRT plugin paths are not supported with the TensorRT-RTX backend. "
             "Remove --trt_plugins or select the classic TensorRT EP."
         )
+    if quantize_mode in {"int8", "fp8"} and (calibration_data is None) == (
+        calibration_data_reader is None
+    ):
+        raise ValueError("Provide exactly one of calibration_data or calibration_data_reader")
 
     configure_logging(log_level.upper(), log_file)
     logger.info(f"Starting quantization process for model: {onnx_path}")
@@ -856,9 +894,8 @@ def quantize(
             autotune_context,
             quantize_func,
             onnx_path=onnx_path,
-            calibration_method=calibration_method or "entropy",
+            calibration_method=calibration_method,
             calibration_data_reader=calibration_data_reader,
-            calibration_cache_path=calibration_cache_path,
             calibration_shapes=calibration_shapes,
             calibration_eps=calibration_eps,
             op_types_to_quantize=op_types_to_quantize,