nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/onnx/quantization/__main__.py

@@ -0,0 +1,291 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Command-line entrypoint for ONNX PTQ."""
+
+import argparse
+
+import numpy as np
+
+from modelopt.onnx.quantization.quantize import quantize
+
+__all__ = ["main"]
+
+
+def get_parser() -> argparse.ArgumentParser:
+    """Get the argument parser for ONNX PTQ."""
+    argparser = argparse.ArgumentParser("python -m modelopt.onnx.quantization")
+    group = argparser.add_mutually_exclusive_group(required=False)
+    argparser.add_argument(
+        "--onnx_path", required=True, type=str, help="Input onnx model without Q/DQ nodes."
+    )
+    argparser.add_argument(
+        "--quantize_mode",
+        type=str,
+        choices=["fp8", "int8", "int4"],
+        default="int8",
+        help=("Quantization mode for the given ONNX model."),
+    )
+    argparser.add_argument(
+        "--calibration_method",
+        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}."
+        ),
+    )
+    group.add_argument(
+        "--calibration_data_path",
+        type=str,
+        help="Calibration data in npz/npy format. If None, random data for calibration will be used.",
+    )
+    group.add_argument(
+        "--calibration_cache_path",
+        type=str,
+        help="Pre-calculated activation tensor scaling factors aka calibration cache path.",
+    )
+    argparser.add_argument(
+        "--calibration_shapes",
+        type=str,
+        required=False,
+        help=(
+            "Optional model input shapes for calibration."
+            "Users should provide the shapes specifically if the model has non-batch dynamic dimensions."
+            "Example input shapes spec: input0:1x3x256x256,input1:1x3x128x128"
+        ),
+    )
+    argparser.add_argument(
+        "--calibration_eps",
+        type=str,
+        default=["cpu", "cuda:0", "trt"],
+        nargs="+",
+        help=(
+            "Priority order for the execution providers (EP) to calibrate the model. "
+            "Any subset of ['trt', 'cuda:x', dml:x, 'cpu'], where 'x' is the device id."
+            "If a custom op is detected in the model, 'trt' will automatically be added to the EP list."
+        ),
+    )
+    argparser.add_argument(
+        "--override_shapes",
+        type=str,
+        required=False,
+        help=(
+            "Override model input shapes with static shapes."
+            "Example input shapes spec: input0:1x3x256x256,input1:1x3x128x128"
+        ),
+    )
+    argparser.add_argument(
+        "--op_types_to_quantize",
+        type=str,
+        default=[],
+        nargs="+",
+        help="A space-separated list of node types to quantize.",
+    )
+    argparser.add_argument(
+        "--op_types_to_exclude",
+        type=str,
+        default=[],
+        nargs="+",
+        help="A space-separated list of node types to exclude from quantization.",
+    )
+    argparser.add_argument(
+        "--nodes_to_quantize",
+        type=str,
+        default=[],
+        nargs="+",
+        help="A space-separated list of node names to quantize. Regular expressions are supported.",
+    )
+    argparser.add_argument(
+        "--nodes_to_exclude",
+        type=str,
+        default=[],
+        nargs="+",
+        help="A space-separated list of node names to exclude from quantization. Regular expressions are supported.",
+    )
+    argparser.add_argument(
+        "--use_external_data_format",
+        action="store_true",
+        help=(
+            "If True or model size is larger than 2GB, "
+            "<MODEL_NAME>.onnx_data will be used to write weights and constants."
+        ),
+    )
+    argparser.add_argument(
+        "--keep_intermediate_files",
+        action="store_true",
+        help=(
+            "If True, keep the files generated during the ONNX models' conversion/calibration. "
+            "Otherwise, only the converted ONNX file is kept for the user."
+        ),
+    )
+    argparser.add_argument(
+        "--output_path",
+        type=str,
+        help=(
+            "Output filename to save the converted ONNX model. If None, save it in the same dir as "
+            "the original ONNX model with an appropriate suffix."
+        ),
+    )
+    argparser.add_argument(
+        "--log_level",
+        type=str,
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "debug", "info", "warning", "error"],
+        default="INFO",
+        help="Set the logging level for the quantization process.",
+    )
+    argparser.add_argument(
+        "--log_file",
+        type=str,
+        default=None,
+        help="Path to the log file for the quantization process.",
+    )
+    argparser.add_argument(
+        "--trt_plugins",
+        type=str,
+        default=None,
+        nargs="+",
+        help=(
+            "A space-separated list with the custom TensorRT plugin library paths in .so format (compiled shared "
+            "library). If this is not None, the TensorrtExecutionProvider is invoked, so make sure that the TensorRT "
+            "libraries are in the PATH or LD_LIBRARY_PATH variables."
+        ),
+    )
+    argparser.add_argument(
+        "--trt_plugins_precision",
+        type=str,
+        default=None,
+        nargs="+",
+        help=(
+            "A space-separated list indicating the precision for each custom op. "
+            "Each item should have the format <op_type>:<precision> (all inputs and outputs have the same precision) "
+            "or <op_type>:[<inp1_precision>,<inp2_precision>,...]:[<out1_precision>,<out2_precision>,...] "
+            "(inputs and outputs can have different precisions), where precision can be fp32 (default), "
+            "fp16, int8, or fp8. Note that int8/fp8 should be the same as the quantization mode. "
+            "For example: op_type_1:fp16 op_type_2:[int8,fp32]:[int8]."
+        ),
+    )
+    argparser.add_argument(
+        "--high_precision_dtype",
+        type=str,
+        default=None,
+        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."
+        ),
+    )
+    argparser.add_argument(
+        "--mha_accumulation_dtype",
+        type=str,
+        default="fp16",
+        help=(
+            "Accumulation dtype of MHA. This flag will only take effect when mha_accumulation_dtype == 'fp32' "
+            "and quantize_mode == 'fp8'. One of ['fp32', 'fp16']"
+        ),
+    )
+    argparser.add_argument(
+        "--disable_mha_qdq",
+        action="store_true",
+        help="If True, Q/DQ will not be added to MatMuls in MHA pattern.",
+    )
+    argparser.add_argument(
+        "--dq_only",
+        action="store_true",
+        help=(
+            "If True, FP32/FP16 weights will be converted to INT8/FP8 weights. Q nodes will get removed from the "
+            "weights and have only DQ nodes with those converted INT8/FP8 weights in the output model."
+        ),
+    )
+    argparser.add_argument(
+        "--use_zero_point",
+        type=bool,
+        default=False,
+        help=(
+            "If True, zero-point based quantization will be used - currently, applicable for awq_lite algorithm."
+        ),
+    )
+    argparser.add_argument(
+        "--passes",
+        type=str,
+        choices=["concat_elimination"],
+        default=["concat_elimination"],
+        nargs="+",
+        help=(
+            "A space-separated list of optimization passes name, if set, appropriate pre/post-processing passes will "
+            "be invoked."
+        ),
+    )
+    argparser.add_argument(
+        "--simplify",
+        action="store_true",
+        help="If True, the given ONNX model will be simplified before quantization is performed.",
+    )
+    argparser.add_argument(
+        "--calibrate_per_node",
+        action="store_true",
+        help=(
+            "If set, performs calibration per node instead of running inference over the entire network. "
+            "Useful for reducing memory consumption during large model inference."
+        ),
+    )
+    return argparser
+
+
+def main():
+    """Command-line entrypoint for ONNX PTQ."""
+    args = get_parser().parse_args()
+    calibration_data = None
+    if args.calibration_data_path:
+        calibration_data = np.load(args.calibration_data_path, allow_pickle=True)
+        if args.calibration_data_path.endswith(".npz"):
+            # 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,
+        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,
+        override_shapes=args.override_shapes,
+        op_types_to_quantize=args.op_types_to_quantize,
+        op_types_to_exclude=args.op_types_to_exclude,
+        nodes_to_quantize=args.nodes_to_quantize,
+        nodes_to_exclude=args.nodes_to_exclude,
+        use_external_data_format=args.use_external_data_format,
+        keep_intermediate_files=args.keep_intermediate_files,
+        output_path=args.output_path,
+        log_level=args.log_level,
+        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,
+        mha_accumulation_dtype=args.mha_accumulation_dtype,
+        disable_mha_qdq=args.disable_mha_qdq,
+        dq_only=args.dq_only,
+        use_zero_point=args.use_zero_point,
+        passes=args.passes,
+        simplify=args.simplify,
+        calibrate_per_node=args.calibrate_per_node,
+    )
+
+
+if __name__ == "__main__":
+    main()

modelopt/onnx/quantization/calib_utils.py

@@ -0,0 +1,178 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Provides basic calibration utils."""
+
+import struct
+from typing import Union
+
+import numpy as np
+import onnx
+from onnxruntime.quantization.calibrate import CalibrationDataReader
+
+from modelopt.onnx.logging_config import logger
+from modelopt.onnx.utils import (
+    gen_random_inputs,
+    get_input_names,
+    get_input_shapes,
+    parse_shapes_spec,
+)
+
+CalibrationDataType = Union[np.ndarray, dict[str, np.ndarray]]  # noqa: UP007
+
+
+class CalibrationDataProvider(CalibrationDataReader):
+    """Calibration data provider class."""
+
+    def __init__(
+        self,
+        onnx_path: str,
+        calibration_data: CalibrationDataType,
+        calibration_shapes: str | None = None,
+    ):
+        """Initializes the data provider class with the calibration data iterator.
+
+        Args:
+            onnx_path: Path to the ONNX model.
+            calibration_data: Numpy data to calibrate the model.
+                Ex. If a model has input shapes like {"sample": (2, 4, 64, 64), "timestep": (1,),
+                "encoder_hidden_states": (2, 16, 768)}, the calibration data should have dictionary
+                of tensors with shapes like {"sample": (1024, 4, 64, 64), "timestep": (512,),
+                "encoder_hidden_states": (1024, 16, 768)} to calibrate with 512 samples.
+            calibration_shapes: A string representing the shape of each input tensors for one calibration step.
+                If the shape is not provided for an input tensor, the shape is inferred from the onnx model directly,
+                with all the unknown dims filled with 1.
+        """
+        logger.info("Setting up CalibrationDataProvider for calibration")
+        # Tensor data is not required to generate the calibration data
+        # So even if the model has external data, we don't need to load them here
+        onnx_model = onnx.load(onnx_path)
+        input_names = get_input_names(onnx_model)
+        input_shapes = {} if calibration_shapes is None else parse_shapes_spec(calibration_shapes)
+        inferred_input_shapes = get_input_shapes(onnx_model)
+        for name in input_names:
+            if name not in input_shapes:
+                input_shapes[name] = inferred_input_shapes[name]
+                logger.debug(f"Inferred shape for {name}: {inferred_input_shapes[name]}")
+
+        # Validate calibration data against expected inputs by the model
+        if isinstance(calibration_data, np.ndarray):
+            assert len(input_names) == 1, "Calibration data has only one tensor."
+            calibration_data = {input_names[0]: calibration_data}
+            logger.debug(
+                f"Single tensor calibration data shape: {calibration_data[input_names[0]].shape}"
+            )
+        elif isinstance(calibration_data, dict):
+            assert len(input_names) == len(calibration_data), (
+                "Model input count and calibration data doesn't match."
+            )
+            for input_name in input_names:
+                assert input_name in calibration_data
+            logger.debug(f"Multi-tensor calibration data with {len(calibration_data)} inputs")
+        else:
+            raise ValueError(
+                f"calibration data must be numpy array or dict, got {type(calibration_data)}"
+            )
+
+        # Create list of model inputs with appropriate batch size
+        n_itr = int(calibration_data[input_names[0]].shape[0] / input_shapes[input_names[0]][0])
+        logger.debug(f"Creating {n_itr} calibration iterations")
+        self.calibration_data_list = [{}] * n_itr
+        for input_name in input_names:
+            for idx, calib_data in enumerate(
+                np.array_split(calibration_data[input_name], n_itr, axis=0)
+            ):
+                self.calibration_data_list[idx][input_name] = calib_data
+
+        self.calibration_data_reader = iter(self.calibration_data_list)
+
+    def get_next(self):
+        """Returns the next available calibration input from the reader."""
+        return next(self.calibration_data_reader, None)
+
+    def get_first(self):
+        """Returns the first calibration input from the reader without incrementing the iterator.
+
+        This is useful when doing a test run for the session.
+        """
+        assert len(self.calibration_data_list) > 0, "Calibration data list is empty!"
+        return self.calibration_data_list[0]
+
+    def rewind(self):
+        """Rewinds the data reader to the first index."""
+        self.calibration_data_reader = iter(self.calibration_data_list)
+
+
+class RandomDataProvider(CalibrationDataReader):
+    """Calibration data reader class with random data provider."""
+
+    def __init__(self, onnx_model: str | onnx.ModelProto, calibration_shapes: str | None = None):
+        """Initializes the data reader class with random calibration data."""
+        logger.info("Initializing RandomDataProvider")
+        if isinstance(onnx_model, str):
+            onnx_path = onnx_model
+            logger.debug(
+                f"Loading ONNX model from: {onnx_path} to read the input shapes for RandomDataProvider"
+            )
+            # Tensor data is not required to generate the calibration data
+            # So even if the model has external data, we don't need to load them here
+            onnx_model = onnx.load(onnx_path)
+        self.calibration_data_list: list[dict[str, np.ndarray]] = [
+            gen_random_inputs(onnx_model, calibration_shapes)
+        ]
+        self.calibration_data_reader = iter(self.calibration_data_list)
+
+    def get_next(self):
+        """Returns the next available calibration input from the reader."""
+        return next(self.calibration_data_reader, None)
+
+    def get_first(self):
+        """Returns the first calibration input from the reader without incrementing the iterator.
+
+        This is useful when doing a test run for the session.
+        """
+        assert len(self.calibration_data_list) > 0, "Calibration data list is empty!"
+        return self.calibration_data_list[0]
+
+    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/extensions.py

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Module to load C++ extensions."""
+
+import os
+import sys
+
+import cppimport
+
+from modelopt.onnx.logging_config import logger
+
+try:
+    logger.info("Loading extension modelopt_round_and_pack_ext...")
+    path = os.path.join(os.path.dirname(__file__), "src")
+    sys.path.append(path)
+    round_and_pack_ext = cppimport.imp("modelopt_round_and_pack_ext")
+    sys.path.remove(path)
+except Exception as e:
+    logger.warning(
+        f"{e}\nUnable to load `modelopt_round_and_pack_ext', falling back to python based optimized version"
+    )
+    round_and_pack_ext = None