diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7654acab550..ad0e3d910d3 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -105,7 +105,10 @@ repos: modelopt/torch/quantization/utils/calib_utils.py| modelopt/torch/quantization/ggml/codebooks.py| modelopt/onnx/quantization/operators.py| - modelopt/onnx/quantization/ort_patching.py| + modelopt/onnx/quantization/ort_calibration.py| + modelopt/onnx/quantization/ort_calibration_per_node.py| + modelopt/onnx/quantization/ort_quantization.py| + modelopt/onnx/quantization/ort_session.py| modelopt/torch/_deploy/utils/onnx_utils.py| modelopt/torch/export/transformer_engine.py| modelopt/torch/puzzletron/anymodel/models/gpt_oss/gpt_oss_pruned_to_mxfp4.py| diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c15cd53075d..ea40dd66b70 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -23,6 +23,10 @@ Changelog **Backward Breaking Changes** +- The ``modelopt.onnx.quantization.ort_patching`` module has been removed with no + compatibility shim; update direct imports to use ``ort_session`` for model loading + and session setup, ``ort_calibration`` or ``ort_calibration_per_node`` for calibration, + ``ort_quantization`` for static Q/DQ quantization, and ``ort_patches`` for patch composition. - The ``modelopt.onnx.quantization.graph_utils`` module has been removed with no compatibility shim; update direct imports using this migration map: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index f895ac1d93b..3f344153850 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -145,6 +145,8 @@ See [`modelopt/torch/quantization/utils/calib_utils.py`](./modelopt/torch/quanti We use [pytest](https://docs.pytest.org/) for all tests. For any new features / examples, make sure to add tests and that the coverage check in your PR passes. The tests are organized into the following directories: +Test filenames must mirror the production module they primarily cover: tests for `foo.py` belong in `test_foo.py`. When a production module is renamed, split, or removed, rename or split its tests in the same change so test filenames do not preserve obsolete module names. A cross-module integration test may use a feature-oriented name only when no single production module owns the behavior. + - `tests/unit`: Fast cpu-based unit tests for the core ModelOpt library. They should not take more than a few seconds to run. - `tests/gpu`: Fast GPU-based unit tests for the core ModelOpt library. In most cases, they should not take more than a few seconds to run. - `tests/gpu_megatron`: Fast GPU-based unit tests for the core ModelOpt library for Megatron-Core features. In most cases, they should not take more than a few seconds to run. diff --git a/modelopt/onnx/quantization/fp8.py b/modelopt/onnx/quantization/fp8.py index ddd0dc2021f..7159988b961 100755 --- a/modelopt/onnx/quantization/fp8.py +++ b/modelopt/onnx/quantization/fp8.py @@ -39,7 +39,7 @@ find_nodes_to_exclude, ) from modelopt.onnx.quantization.int8 import _find_nodes_to_quantize -from modelopt.onnx.quantization.ort_patching import _quantize_static as quantize_static +from modelopt.onnx.quantization.ort_quantization import _quantize_static as quantize_static from modelopt.onnx.quantization.ort_utils import configure_ort from modelopt.onnx.quantization.precision_utils import ( _convert_to_runtime_precision, diff --git a/modelopt/onnx/quantization/int8.py b/modelopt/onnx/quantization/int8.py index 61ea5274898..059a577724c 100755 --- a/modelopt/onnx/quantization/int8.py +++ b/modelopt/onnx/quantization/int8.py @@ -35,7 +35,7 @@ find_nodes_from_matmul_to_exclude, find_nodes_to_exclude, ) -from modelopt.onnx.quantization.ort_patching import _quantize_static as quantize_static +from modelopt.onnx.quantization.ort_quantization import _quantize_static as quantize_static from modelopt.onnx.quantization.ort_utils import configure_ort from modelopt.onnx.quantization.partitioning import ( find_fusible_partitions, diff --git a/modelopt/onnx/quantization/ort_calibration.py b/modelopt/onnx/quantization/ort_calibration.py new file mode 100644 index 00000000000..df12afbf89b --- /dev/null +++ b/modelopt/onnx/quantization/ort_calibration.py @@ -0,0 +1,506 @@ +# Adapted from https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/calibrate.py +# +# MIT License +# +# Copyright (c) Microsoft Corporation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# 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. + +"""Ordinary ONNX Runtime calibration patches.""" + +__all__ = [] + +from collections.abc import Sequence +from pathlib import Path + +import numpy as np +import onnx +import onnxruntime as ort +import pynvml +from onnx import onnx_pb +from onnxruntime.quantization.calibrate import ( + CalibrationDataReader, + CalibrationMethod, + HistogramCollector, + TensorData, + TensorsData, +) +from tqdm import tqdm + +from modelopt.onnx.logging_config import logger +from modelopt.onnx.quantization.ort_session import load_model_with_shape_infer + + +def _select_tensors_to_calibrate(calibrator, model: onnx.ModelProto): + """Select input/output tensors of candidate nodes to calibrate. + + Returns: + tensors (set): set of tensor name. + value_infos (dict): tensor name to value info. + """ + value_infos = {vi.name: vi for vi in model.graph.value_info} + value_infos.update({ot.name: ot for ot in model.graph.output}) + value_infos.update({it.name: it for it in model.graph.input}) + initializer = {init.name for init in model.graph.initializer} + + tensors_to_calibrate = set() + tensor_type_to_calibrate = {onnx_pb.TensorProto.FLOAT, onnx_pb.TensorProto.FLOAT16} + + for node in model.graph.node: + # Hack: in calibrator.op_types_to_calibrate we pass nodes_to_quantize + if node.name in calibrator.op_types_to_calibrate: + for tensor_name in node.input: + if tensor_name in value_infos: + vi = value_infos[tensor_name] + if ( + vi.type.HasField("tensor_type") + and (vi.type.tensor_type.elem_type in tensor_type_to_calibrate) + and (tensor_name not in initializer) + ): + tensors_to_calibrate.add(tensor_name) + for tensor_name in node.output: + if tensor_name in value_infos: + vi = value_infos[tensor_name] + if vi.type.HasField("tensor_type") and ( + vi.type.tensor_type.elem_type in tensor_type_to_calibrate + ): + tensors_to_calibrate.add(tensor_name) + + return tensors_to_calibrate, value_infos + + +def _init_calibrater_base( + calibrater, + model_path: str | Path, + op_types_to_calibrate: Sequence[str] | None = None, + augmented_model_path="augmented_model.onnx", + symmetric=False, + use_external_data_format=False, + per_channel=False, +): + """Initialize calibrater base class. + + :param model_path: ONNX model to calibrate. It should be a model file path + :param op_types_to_calibrate: operator types to calibrate. By default, calibrate all the float32/float16 tensors. + :param augmented_model_path: save augmented model to this path. + :param symmetric: make range of tensor symmetric (central point is 0). + :param use_external_data_format: use external data format to store model which size is >= 2Gb + + Modification: Additional members including single_node_model_path_map, providers, and trt_extra_plugin_lib_paths + were added and initialized to support calibration per node feature. + """ + if isinstance(model_path, str): + calibrater.model = load_model_with_shape_infer(Path(model_path)) + elif isinstance(model_path, Path): + calibrater.model = load_model_with_shape_infer(model_path) + else: + raise ValueError("model_path should be model path.") + + calibrater.op_types_to_calibrate = op_types_to_calibrate + calibrater.augmented_model_path = augmented_model_path + calibrater.symmetric = symmetric + calibrater.use_external_data_format = use_external_data_format + calibrater.per_channel = per_channel + calibrater.augment_model = None + calibrater.infer_session = None + calibrater.execution_providers = [] + + # Add single node calibration members + calibrater.single_node_model_path_map = {} # {path: ([inputs], [outputs])} + calibrater.providers = [] + calibrater.trt_extra_plugin_lib_paths = None + + +def _prepare_histogram_data(histogram_collector, tensor, data_arr): + """Use FP32 for histogram math while remembering the source dtype.""" + if data_arr.dtype != np.float16: + return data_arr + + original_dtypes = getattr(histogram_collector, "_modelopt_original_dtypes", {}) + original_dtypes[tensor] = data_arr.dtype + histogram_collector._modelopt_original_dtypes = original_dtypes + return data_arr.astype(np.float32) + + +def _restore_histogram_calibration_dtypes(histogram_collector, tensors_range): + """Restore source dtypes at the calibration-to-quantization boundary.""" + original_dtypes = getattr(histogram_collector, "_modelopt_original_dtypes", {}) + for tensor, dtype in original_dtypes.items(): + if tensor not in tensors_range: + continue + tensor_data = tensors_range[tensor] + dtype_limits = np.finfo(dtype) + for attribute in ("lowest", "highest", "avg", "std"): + if hasattr(tensor_data, attribute): + value = np.clip(getattr(tensor_data, attribute), dtype_limits.min, dtype_limits.max) + setattr(tensor_data, attribute, np.asarray(value, dtype=dtype)) + + +def _collect_value(histogram_collector, name_to_arr): + """Collect histogram on real value.""" + for tensor, data_arr in tqdm(name_to_arr.items()): + # ====================== Modification ====================== + concat_data_arr = np.asarray(data_arr[0]) + concat_data_arr = concat_data_arr.flatten() + for i in range(1, len(data_arr)): + curr_data_arr = np.asarray(data_arr[i]) + curr_data_arr = curr_data_arr.flatten() + concat_data_arr = np.concatenate((concat_data_arr, curr_data_arr)) + + concat_data_arr = _prepare_histogram_data(histogram_collector, tensor, concat_data_arr) + data_arr = concat_data_arr + # ========================================================== + if data_arr.size > 0: + min_value = np.min(data_arr) + max_value = np.max(data_arr) + else: + min_value = np.array(0, dtype=data_arr.dtype) + max_value = np.array(0, dtype=data_arr.dtype) + + # Change the inf and nan values to meaningful min/max + min_value = ( + np.finfo(np.float32).tiny if np.isinf(min_value) or np.isnan(min_value) else min_value + ) + max_value = ( + np.finfo(np.float32).max if np.isinf(max_value) or np.isnan(max_value) else max_value + ) + + threshold = max(abs(min_value), abs(max_value)) + + if tensor in histogram_collector.histogram_dict: + old_histogram = histogram_collector.histogram_dict[tensor] + histogram_collector.histogram_dict[tensor] = histogram_collector.merge_histogram( + old_histogram, data_arr, min_value, max_value, threshold + ) + else: + range_max = float(threshold) + hist, hist_edges = np.histogram( + data_arr, histogram_collector.num_bins, range=(-range_max, range_max) + ) + histogram_collector.histogram_dict[tensor] = ( + hist, + hist_edges, + min_value, + max_value, + threshold, + ) + + +def _collect_absolute_value(histogram_collector, name_to_arr): + """Collect histogram on absolute value.""" + for tensor, data_arr in name_to_arr.items(): + if isinstance(data_arr, list): + for arr in data_arr: + assert isinstance(arr, np.ndarray), ( + f"Unexpected type {type(arr)} for tensor={tensor!r}" + ) + dtypes = {a.dtype for a in data_arr} + assert len(dtypes) == 1, ( + f"The calibration expects only one element type but got {dtypes} for tensor={tensor!r}" + ) + # ====================== Modification ====================== + concat_data_arr = np.asarray(data_arr[0]) + concat_data_arr = concat_data_arr.flatten() + for i in range(1, len(data_arr)): + curr_data_arr = np.asarray(data_arr[i]) + curr_data_arr = curr_data_arr.flatten() + concat_data_arr = np.concatenate((concat_data_arr, curr_data_arr)) + data_arr_np = concat_data_arr + # ========================================================== + elif not isinstance(data_arr, np.ndarray): + raise ValueError(f"Unexpected type {type(data_arr)} for tensor={tensor!r}") + else: + data_arr_np = data_arr + data_arr_np = data_arr_np.flatten() + if data_arr_np.size > 0: + min_value = np.min(data_arr_np) + max_value = np.max(data_arr_np) + else: + min_value = np.array(0, dtype=data_arr_np.dtype) + max_value = np.array(0, dtype=data_arr_np.dtype) + + data_arr_np = np.absolute(data_arr_np) # only consider absolute value + + if tensor not in histogram_collector.histogram_dict: + # first time it uses num_bins to compute histogram. + hist, hist_edges = np.histogram(data_arr_np, bins=histogram_collector.num_bins) + hist_edges = hist_edges.astype(data_arr_np.dtype) + assert data_arr_np.dtype != np.float64, ( + "only float32 or float16 is supported, every constant must be explicitly typed" + ) + histogram_collector.histogram_dict[tensor] = (hist, hist_edges, min_value, max_value) + else: + old_histogram = histogram_collector.histogram_dict[tensor] + old_min = old_histogram[2] + old_max = old_histogram[3] + assert hasattr(old_min, "dtype"), ( + f"old_min should be a numpy array but is {type(old_min)}" + ) + assert hasattr(old_max, "dtype"), ( + f"old_min should be a numpy array but is {type(old_max)}" + ) + old_hist = old_histogram[0] + old_hist_edges = old_histogram[1] + temp_amax = np.max(data_arr_np) + if temp_amax > old_hist_edges[-1]: + # increase the number of bins + width = old_hist_edges[1] - old_hist_edges[0] + # NOTE: np.arange may create an extra bin after the one containing temp_amax + new_bin_edges = np.arange(old_hist_edges[-1] + width, temp_amax + width, width) + old_hist_edges = np.hstack((old_hist_edges, new_bin_edges)) + hist, hist_edges = np.histogram(data_arr_np, bins=old_hist_edges) + hist_edges = hist_edges.astype(data_arr_np.dtype) + hist[: len(old_hist)] += old_hist + assert data_arr_np.dtype != np.float64, ( + "only float32 or float16 is supported, every constant must be explicitly typed" + ) + histogram_collector.histogram_dict[tensor] = ( + hist, + hist_edges, + min(old_min, min_value), + max(old_max, max_value), + ) + + +def _compute_data_minmax_calibrator(calibrator): + """Compute the min-max range of tensor. + + :returns: dictionary mapping: {added node names: (ReduceMin, ReduceMax) pairs } + """ + if len(calibrator.intermediate_outputs) == 0: + return calibrator.calibrate_tensors_range + + output_names = [ + calibrator.infer_session.get_outputs()[i].name + for i in range(len(calibrator.intermediate_outputs[0])) + ] + + output_dicts_list = [ + dict(zip(output_names, intermediate_output), strict=True) + for intermediate_output in calibrator.intermediate_outputs + ] + + merged_output_dict = {} + for d in output_dicts_list: + for k, v in d.items(): + merged_output_dict.setdefault(k, []).append(v) + + # ====================== Modification ====================== + # Group qdq tensors should have the same scaling factor. Each tensor in group should add + # other tensors in its merged_dict value. In this way, calibrator will generate the same + # scaling factor. + if calibrator.group_qdq_tensors: + for cur, group in calibrator.group_qdq_tensors.items(): + for other in group: + if cur == other: + continue + for d in output_dicts_list: + for k, v in d.items(): + cur_min = cur + "_" + "ReduceMin" + cur_max = cur + "_" + "ReduceMax" + other_min = other + "_" + "ReduceMin" + other_max = other + "_" + "ReduceMax" + if k == other_min: + merged_output_dict[cur_min].append(v) + elif k == other_max: + merged_output_dict[cur_max].append(v) + # ============================================================ + + added_output_names = output_names[calibrator.num_model_outputs :] + calibrate_tensor_names = [ + added_output_names[i].rpartition("_")[0] for i in range(0, len(added_output_names), 2) + ] # output names + + merged_added_output_dict = { + i: merged_output_dict[i] + for i in merged_output_dict + if i not in calibrator.model_original_outputs + } + + pairs = [] + for i in range(0, len(added_output_names), 2): + if calibrator.moving_average: + min_value_array = np.mean(merged_added_output_dict[added_output_names[i]], axis=0) + max_value_array = np.mean(merged_added_output_dict[added_output_names[i + 1]], axis=0) + else: + min_value_array = np.min(merged_added_output_dict[added_output_names[i]], axis=0) + max_value_array = np.max(merged_added_output_dict[added_output_names[i + 1]], axis=0) + + if calibrator.symmetric: + max_absolute_value = np.max([np.abs(min_value_array), np.abs(max_value_array)], axis=0) + pairs.append((-max_absolute_value, max_absolute_value)) + else: + pairs.append((min_value_array, max_value_array)) + + new_calibrate_tensors_range = TensorsData( + CalibrationMethod.MinMax, dict(zip(calibrate_tensor_names, pairs, strict=False)) + ) + if calibrator.calibrate_tensors_range: + calibrator.calibrate_tensors_range = calibrator.merge_range( + calibrator.calibrate_tensors_range, new_calibrate_tensors_range + ) + else: + calibrator.calibrate_tensors_range = new_calibrate_tensors_range + + return calibrator.calibrate_tensors_range + + +def _collect_data_minmax_calibrator(calibrator, data_reader: CalibrationDataReader): + """This function overwrite is needed to solve OOM issue due to the unlimited accumulation of intermediate_outputs. + + Support for: MinMax Calibrator. + Modification: indented the last lines of code inside the while loop in order to run compute_data for each sample + batch individually instead of the entire data at once. The assumption here is that the ONNX file has bs=N + and the calibration data size is M (where M is a multiple of N). So the calibrator is a sequence of M/N + samples with bs=N. + """ + run_options = ort.RunOptions() + try: + pynvml.nvmlInit() + gpu_count = pynvml.nvmlDeviceGetCount() + pynvml.nvmlShutdown() + except Exception as e: + logger.error(f"Failed to get GPU count: {e}") + gpu_count = 0 + gpu_str = ";".join([f"gpu:{i}" for i in range(gpu_count)]) + run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", f"cpu:0;{gpu_str}") + while True: + inputs = data_reader.get_next() + if not inputs: + break + run_options = ort.RunOptions() + + calibrator.intermediate_outputs.append( + calibrator.infer_session.run(None, inputs, run_options=run_options) + ) + + # ======== Modification: block is indentend in ======== + if len(calibrator.intermediate_outputs) == 0: + raise ValueError("No data is collected.") + + t = calibrator.compute_data() + if not isinstance(t, TensorsData): + raise TypeError(f"compute_data must return a TensorsData not {type(t)}.") + calibrator.clear_collected_data() + + +def _merge_range_minmax_calibrator(calibrator, old_range: TensorsData, new_range: TensorsData): + """This function is an auxiliary function of collect_data to solve the OOM issue in the MinMax Calibrator. + + Issue fixed with this function: old_range is not a dictionary, but old_range.data is. + TODO: create an MR in the ORT repository for this function. Alternatively, we can also file the MR fixing + TensorData (need to at least add items() function there). + """ + if not old_range: + return new_range + + for key, value in old_range.data.items(): + value_tuple = value.range_value + new_range_tuple = new_range.data[key].range_value + if calibrator.moving_average: + min_value = value_tuple[0] + calibrator.averaging_constant * ( + new_range_tuple[0] - value_tuple[0] + ) + max_value = value_tuple[1] + calibrator.averaging_constant * ( + new_range_tuple[1] - value_tuple[1] + ) + else: + min_value = min(value_tuple[0], new_range_tuple[0]) + max_value = max(value_tuple[1], new_range_tuple[1]) + new_range.data[key] = TensorData(lowest=min_value, highest=max_value) + + return new_range + + +def _collect_data_histogram_calibrator(calibrator, data_reader: CalibrationDataReader): + """This function overwrite is needed to solve OOM issue due to the unlimited accumulation of intermediate_outputs. + + Support for: Histogram Calibrator (which affects Entropy, Percentile, and DIstribution Calibrators). + Modification: indented the last lines of code inside the while loop in order to run compute_data for each sample + batch individually instead of the entire data at once. + """ + while True: + inputs = data_reader.get_next() + if not inputs: + break + calibrator.intermediate_outputs.append(calibrator.infer_session.run(None, inputs)) + + # ======== Modification: block is indentend in ======== + # Here, compute_date is calculated for every sample batch instead of the entire data at once. + if len(calibrator.intermediate_outputs) == 0: + raise ValueError("No data is collected.") + + output_names = [ + calibrator.infer_session.get_outputs()[i].name + for i in range(len(calibrator.intermediate_outputs[0])) + ] + output_dicts_list = [ + dict(zip(output_names, intermediate_output)) + for intermediate_output in calibrator.intermediate_outputs + ] + + merged_dict = {} + for d in output_dicts_list: + for k, v in d.items(): + merged_dict.setdefault(k, []).append(v) + + # Group qdq tensors should have the same scaling factor. Each tensor in group should add + # other tensors in its merged_dict value. In this way, calibrator will generate the same + # scaling factor. + if calibrator.group_qdq_tensors: + for cur, group in calibrator.group_qdq_tensors.items(): + for other in group: + if cur == other: + continue + for d in output_dicts_list: + for k, v in d.items(): + if k == other: + merged_dict[cur].append(v) + + clean_merged_dict = { + i: merged_dict[i] for i in merged_dict if i in calibrator.tensors_to_calibrate + } + + if not calibrator.collector: + calibrator.collector = HistogramCollector( + method=calibrator.method, + symmetric=calibrator.symmetric, + num_bins=calibrator.num_bins, + num_quantized_bins=calibrator.num_quantized_bins, + percentile=calibrator.percentile, + scenario=calibrator.scenario, + ) + calibrator.collector.collect(clean_merged_dict) + + calibrator.clear_collected_data() diff --git a/modelopt/onnx/quantization/ort_calibration_per_node.py b/modelopt/onnx/quantization/ort_calibration_per_node.py new file mode 100644 index 00000000000..aa89a04014d --- /dev/null +++ b/modelopt/onnx/quantization/ort_calibration_per_node.py @@ -0,0 +1,904 @@ +# Adapted from https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/calibrate.py +# +# MIT License +# +# Copyright (c) Microsoft Corporation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# 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. + +"""Per-node ONNX Runtime calibration patches.""" + +__all__ = [] + +import gc +import uuid + +import numpy as np +import onnx +import onnxruntime as ort +from onnxruntime.quantization.calibrate import ( + CalibrationDataReader, + CalibrationMethod, + HistogramCollector, + TensorData, + TensorsData, +) +from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference +from tqdm import tqdm + +from modelopt.onnx.logging_config import logger +from modelopt.onnx.quantization.ort_calibration import _prepare_histogram_data + + +def _compute_data_min_max_calibrater_single_node_calibration(calibrater) -> TensorData: + """Compute the min-max range of tensor. + + :return: dictionary mapping: {added node names: (ReduceMin, ReduceMax) pairs } + + Modification: Instead of aggregating two consecutive outputs to a MinMax pair, retrieve a MinMax pair from + outputs of Concat. + """ + if not calibrater.intermediate_outputs: + return calibrater.calibrate_tensors_range + + # Get output names and merge all intermediate outputs + output_names = [out.name for out in calibrater.infer_session.get_outputs()] + + # Merge outputs across all batches, filtering out original model outputs + merged_outputs = {} + for intermediate_output in calibrater.intermediate_outputs: + for name, value in zip(output_names, intermediate_output): + if name not in calibrater.model_original_outputs: + merged_outputs.setdefault(name, []).append(value) + + # Compute min/max pairs for each tensor + pairs = [] + tensor_names = [] + + for output_name, values in merged_outputs.items(): + tensor_names.append(output_name.rpartition("_")[0]) + + if calibrater.moving_average: + min_val, max_val = np.mean(values, axis=0) + else: + stacked_values = np.stack(values, axis=0) + min_val = np.min(stacked_values, axis=0)[0] + max_val = np.max(stacked_values, axis=0)[1] + + if calibrater.symmetric: + max_abs = max(np.abs(min_val), np.abs(max_val)) + pairs.append((-max_abs, max_abs)) + else: + pairs.append((min_val, max_val)) + + # Create and merge tensor range data + new_range = TensorsData(CalibrationMethod.MinMax, dict(zip(tensor_names, pairs))) + + calibrater.calibrate_tensors_range = ( + calibrater.merge_range(calibrater.calibrate_tensors_range, new_range) + if calibrater.calibrate_tensors_range + else new_range + ) + + return calibrater.calibrate_tensors_range + + +def _merge_range_min_max_calibrater_single_node_calibration( + calibrater, old_range: TensorsData, new_range: TensorsData +): + """This function is an auxiliary function of collect_data to solve the OOM issue in the MinMax Calibrator. + + Issue fixed with this function: old_range is not a dictionary, but old_range.data is. + TODO: create an MR in the ORT repository for this function. Alternatively, we can also file the MR fixing + TensorData (need to at least add items() function there). + """ + if not old_range: + return new_range + + def _merge_ranges(old_min, old_max, new_min, new_max): + if calibrater.moving_average: + alpha = calibrater.averaging_constant + return (old_min + alpha * (new_min - old_min), old_max + alpha * (new_max - old_max)) + return min(old_min, new_min), max(old_max, new_max) + + old_data = old_range.data + for key, new_tensor in new_range.data.items(): + if key in old_data: + old_min, old_max = old_data[key].range_value + new_min, new_max = new_tensor.range_value + merged_min, merged_max = _merge_ranges(old_min, old_max, new_min, new_max) + old_data[key] = TensorData(lowest=merged_min, highest=merged_max) + else: + old_data[key] = new_tensor + + return old_range + + +def _collect_data_min_max_calibrater_single_node_calibration( + calibrater, data_reader: CalibrationDataReader +): + """Collects calibration data (min/max) for a MinMax Calibrator by processing single-node models batch by batch. + + This function addresses an OOM issue by computing calibration data for each batch individually, + rather than accumulating all intermediate outputs across the entire dataset. It assumes the ONNX model + has a batch size of N, and the calibration data size M is a multiple of N, processing M/N batches. + + Args: + calibrater: The calibrater object managing model inference and data collection. + data_reader: Provides batches of input data for calibration. + """ + input_counter = 0 + while True: + inputs = data_reader.get_next() + if not inputs: + break + logger.debug(f"Collecting tensor data and finding min & max for input #{input_counter}") + + # We are using single node model scheme. Set up model to input dependency map + model_to_input_dep_map = {} + for model_path, io_tensors in calibrater.single_node_model_path_map.items(): + model_to_input_dep_map[model_path] = io_tensors[0].copy() # List of input names + + # Setup input queues + input_queue = [model_input.name for model_input in calibrater.model.graph.input] + + pbar = tqdm(total=len(model_to_input_dep_map.keys())) + # Resolve nodes are independent from inputs to add their outputs as inputs + inferred_model_list = [] + for model_path, input_deps in model_to_input_dep_map.items(): + if len(input_deps) == 0: + calibrater.create_inference_session( + execution_providers=calibrater.providers, + trt_extra_plugin_lib_paths=calibrater.trt_extra_plugin_lib_paths, + model_path=model_path, + ) + outputs = calibrater.infer_session.run(None, {}) + + # Add output to inputs + need_calibration = False + for output_idx, output in enumerate(outputs): + output_name = calibrater.infer_session.get_outputs()[output_idx].name + inputs[output_name] = output + input_queue.append(output_name) + if ( + output_name + in [output_tensor.name for output_tensor in calibrater.model.graph.output] + and output_name not in calibrater.model_original_outputs + ): + need_calibration = True + + # Mark model path to remove it from dependency map + inferred_model_list.append(model_path) + + # For each inference, compute data before moving to other nodes if tensor is to be calibrated + if need_calibration: + calibrater.intermediate_outputs.append(outputs) + if len(calibrater.intermediate_outputs) == 0: + raise ValueError("No data is collected.") + + t = calibrater.compute_data() + if not isinstance(t, TensorsData): + raise TypeError(f"compute_data must return a TensorsData not {type(t)}.") + calibrater.clear_collected_data() + + gc.collect() + + pbar.update(1) + + # Remove inferred model from dependency map + for model_path in inferred_model_list: + model_to_input_dep_map.pop(model_path) + + gc.collect() + + # Process topological inference + input_ref_count = {} + while input_queue: + current_input_name = input_queue.pop(0) + + # Initialize input reference count + input_ref_count[current_input_name] = sum( + current_input_name in input_deps for input_deps in model_to_input_dep_map.values() + ) + + # Perform inference + inferred_model_list = [] + for model_path, input_deps in model_to_input_dep_map.items(): + if current_input_name in input_deps: + input_deps.remove(current_input_name) + + # If all dependencies are met, perform inference for the node. + if len(input_deps) == 0: + # Make dictionary of only needed inputs. + inputs_to_feed = {} + for input_name_to_feed in calibrater.single_node_model_path_map[model_path][ + 0 + ]: + inputs_to_feed[input_name_to_feed] = inputs[input_name_to_feed] + + calibrater.create_inference_session( + execution_providers=calibrater.providers, + trt_extra_plugin_lib_paths=calibrater.trt_extra_plugin_lib_paths, + model_path=model_path, + ) + outputs = calibrater.infer_session.run(None, inputs_to_feed) + + # Mark model path to remove it from dependency map + inferred_model_list.append(model_path) + + # Decrease reference count for used inputs and remove if no reference + for input_name in calibrater.single_node_model_path_map[model_path][0]: + input_ref_count[input_name] -= 1 + if input_ref_count[input_name] == 0: + del inputs[input_name] + del input_ref_count[input_name] + + gc.collect() + + # Add outputs to inputs + need_calibration = False + for output_idx, output in enumerate(outputs): + output_name = calibrater.infer_session.get_outputs()[output_idx].name + inputs[output_name] = output + input_queue.append(output_name) + if ( + output_name + in [ + output_tensor.name + for output_tensor in calibrater.model.graph.output + ] + and output_name not in calibrater.model_original_outputs + ): + need_calibration = True + + # For each inference, compute data before moving to other nodes if tensor is to be calibrated + if need_calibration: + calibrater.intermediate_outputs.append(outputs) + if len(calibrater.intermediate_outputs) == 0: + raise ValueError("No data is collected.") + + t = calibrater.compute_data() + if not isinstance(t, TensorsData): + raise TypeError( + f"compute_data must return a TensorsData not {type(t)}." + ) + calibrater.clear_collected_data() + + gc.collect() + + pbar.update(1) + + # Remove inferred model from dependency map + for model_path in inferred_model_list: + model_to_input_dep_map.pop(model_path) + + gc.collect() + pbar.close() + input_counter += 1 + + +def _collect_data_histogram_calibrater_single_node_calibration(calibrator, data_reader): + """Collects histogram data for single-node calibration, processing batches to avoid OOM. + + Args: + calibrator: Histogram calibrator instance. + data_reader: CalibrationDataReader providing input data. + """ + input_counter = 0 + while True: + inputs = data_reader.get_next() + if not inputs: + break + logger.debug(f"Collecting tensor data for input #{input_counter}") + + # We are using single node model scheme. Set up model to input dependency map + model_to_input_dep_map = {} + for model_path, io_tensors in calibrator.single_node_model_path_map.items(): + model_to_input_dep_map[model_path] = io_tensors[0].copy() # List of input names + + # Compute data for input tensors + input_only_model = onnx.helper.make_model( + onnx.helper.make_graph( + [], + f"{calibrator.augmented_model_path[:-5]}_input_only", + calibrator.model.graph.input, + calibrator.model.graph.input, + ), + opset_imports=calibrator.model.opset_import, + functions=calibrator.model.functions, + ir_version=calibrator.model.ir_version, + ) + calibrator.infer_session = ort.InferenceSession(input_only_model.SerializeToString()) + calibrator.intermediate_outputs.append( + [ + inputs[calibrator.infer_session.get_outputs()[i].name] + for i in range(len(calibrator.infer_session.get_outputs())) + ] + ) + if len(calibrator.intermediate_outputs) == 0: + raise ValueError("No data is collected.") + + output_names = [ + calibrator.infer_session.get_outputs()[i].name + for i in range(len(calibrator.intermediate_outputs[0])) + ] + output_dicts_list = [ + dict(zip(output_names, intermediate_output)) + for intermediate_output in calibrator.intermediate_outputs + ] + + merged_dict = {} + for d in output_dicts_list: + for k, v in d.items(): + merged_dict.setdefault(k, []).append(v) + + clean_merged_dict = { + i: merged_dict[i] for i in merged_dict if i in calibrator.tensors_to_calibrate + } + + if not calibrator.collector: + calibrator.collector = HistogramCollector( + method=calibrator.method, + symmetric=calibrator.symmetric, + num_bins=calibrator.num_bins, + num_quantized_bins=calibrator.num_quantized_bins, + percentile=calibrator.percentile, + scenario=calibrator.scenario, + ) + calibrator.collector.collect(clean_merged_dict) + + calibrator.clear_collected_data() + + gc.collect() + + # Setup input queues + input_queue = [model_input.name for model_input in calibrator.model.graph.input] + + pbar = tqdm(total=len(model_to_input_dep_map.keys())) + # Resolve nodes are independent from inputs to add their outputs as inputs + inferred_model_list = [] + for model_path, input_deps in model_to_input_dep_map.items(): + if len(input_deps) == 0: + calibrator.create_inference_session( + execution_providers=calibrator.providers, + trt_extra_plugin_lib_paths=calibrator.trt_extra_plugin_lib_paths, + model_path=model_path, + ) + outputs = calibrator.infer_session.run(None, {}) + + # Add output to inputs + need_calibration = False + for output_idx, output in enumerate(outputs): + output_name = calibrator.infer_session.get_outputs()[output_idx].name + inputs[output_name] = output + input_queue.append(output_name) + if output_name in calibrator.tensors_to_calibrate: + need_calibration = True + + # Mark model path to remove it from dependency map + inferred_model_list.append(model_path) + + # For each inference, compute data before moving to other nodes if tensor is to be calibrated + if need_calibration: + calibrator.intermediate_outputs.append(outputs) + if len(calibrator.intermediate_outputs) == 0: + raise ValueError("No data is collected.") + + output_names = [ + calibrator.infer_session.get_outputs()[i].name + for i in range(len(calibrator.intermediate_outputs[0])) + ] + output_dicts_list = [ + dict(zip(output_names, intermediate_output)) + for intermediate_output in calibrator.intermediate_outputs + ] + + merged_dict = {} + for d in output_dicts_list: + for k, v in d.items(): + merged_dict.setdefault(k, []).append(v) + + clean_merged_dict = { + i: merged_dict[i] + for i in merged_dict + if i in calibrator.tensors_to_calibrate + } + + if not calibrator.collector: + calibrator.collector = HistogramCollector( + method=calibrator.method, + symmetric=calibrator.symmetric, + num_bins=calibrator.num_bins, + num_quantized_bins=calibrator.num_quantized_bins, + percentile=calibrator.percentile, + scenario=calibrator.scenario, + ) + calibrator.collector.collect(clean_merged_dict) + + calibrator.clear_collected_data() + + gc.collect() + pbar.update(1) + + # Remove inferred model from dependency map + for model_path in inferred_model_list: + model_to_input_dep_map.pop(model_path) + + gc.collect() + + # Process topological inference + input_ref_count = {} + while input_queue: + current_input_name = input_queue.pop(0) + + # Initialize input reference count + input_ref_count[current_input_name] = sum( + current_input_name in input_deps for input_deps in model_to_input_dep_map.values() + ) + + # Perform inference + inferred_model_list = [] + for model_path, input_deps in model_to_input_dep_map.items(): + if current_input_name in input_deps: + input_deps.remove(current_input_name) + + # If all dependencies are met, perform inference for the node. + if len(input_deps) == 0: + # Make dictionary of only needed inputs. + inputs_to_feed = {} + for input_name_to_feed in calibrator.single_node_model_path_map[model_path][ + 0 + ]: + inputs_to_feed[input_name_to_feed] = inputs[input_name_to_feed] + + calibrator.create_inference_session( + execution_providers=calibrator.providers, + trt_extra_plugin_lib_paths=calibrator.trt_extra_plugin_lib_paths, + model_path=model_path, + ) + outputs = calibrator.infer_session.run(None, inputs_to_feed) + + # Mark model path to remove it from dependency map + inferred_model_list.append(model_path) + + # Decrease reference count for used inputs and remove if no reference + for input_name in calibrator.single_node_model_path_map[model_path][0]: + input_ref_count[input_name] -= 1 + if input_ref_count[input_name] == 0: + del inputs[input_name] + del input_ref_count[input_name] + + gc.collect() + + # Add outputs to inputs + need_calibration = False + for output_idx, output in enumerate(outputs): + output_name = calibrator.infer_session.get_outputs()[output_idx].name + inputs[output_name] = output + input_queue.append(output_name) + if output_name in calibrator.tensors_to_calibrate: + need_calibration = True + + # For each inference, compute data before moving to other nodes if tensor is to be calibrated + if need_calibration: + calibrator.intermediate_outputs.append(outputs) + if len(calibrator.intermediate_outputs) == 0: + raise ValueError("No data is collected.") + + output_names = [ + calibrator.infer_session.get_outputs()[i].name + for i in range(len(calibrator.intermediate_outputs[0])) + ] + output_dicts_list = [ + dict(zip(output_names, intermediate_output)) + for intermediate_output in calibrator.intermediate_outputs + ] + + merged_dict = {} + for d in output_dicts_list: + for k, v in d.items(): + merged_dict.setdefault(k, []).append(v) + + clean_merged_dict = { + i: merged_dict[i] + for i in merged_dict + if i in calibrator.tensors_to_calibrate + } + + if not calibrator.collector: + calibrator.collector = HistogramCollector( + method=calibrator.method, + symmetric=calibrator.symmetric, + num_bins=calibrator.num_bins, + num_quantized_bins=calibrator.num_quantized_bins, + percentile=calibrator.percentile, + scenario=calibrator.scenario, + ) + calibrator.collector.collect(clean_merged_dict) + + calibrator.clear_collected_data() + + gc.collect() + pbar.update(1) + + # Remove inferred model from dependency map + for model_path in inferred_model_list: + model_to_input_dep_map.pop(model_path) + + gc.collect() + pbar.close() + input_counter += 1 + + +def _collect_histogram_collector_single_node_calibration(histogram_collector, name_to_arr): + """Collect tensor data and make histogram. + + Modification: Remove print line to make calibration per node log output cleaner. + """ + # TODO: Currently we have different collect() for entropy and percentile method respectively. + # Need unified collect in the future. + if histogram_collector.method in {"distribution", "entropy"}: + return histogram_collector.collect_value(name_to_arr) + elif histogram_collector.method == "percentile": + if histogram_collector.symmetric: + return histogram_collector.collect_absolute_value(name_to_arr) + else: + return histogram_collector.collect_value(name_to_arr) + else: + raise ValueError("Only 'entropy', 'percentile' or 'distribution' methods are supported") + + +def _collect_value_histogram_collector_single_node_calibration(histogram_collector, name_to_arr): + """Collect histogram on real value.""" + for tensor, data_arr in name_to_arr.items(): + data_arr = np.asarray(data_arr).flatten() + data_arr = _prepare_histogram_data(histogram_collector, tensor, data_arr) + min_value, max_value = (np.min(data_arr), np.max(data_arr)) if data_arr.size > 0 else (0, 0) + + # Replace inf/nan with float32 min/max + min_value = ( + np.finfo(np.float32).tiny if np.isinf(min_value) or np.isnan(min_value) else min_value + ) + max_value = ( + np.finfo(np.float32).max if np.isinf(max_value) or np.isnan(max_value) else max_value + ) + + threshold = max(abs(min_value), abs(max_value)) + + if tensor in histogram_collector.histogram_dict: + histogram_collector.histogram_dict[tensor] = histogram_collector.merge_histogram( + histogram_collector.histogram_dict[tensor], + data_arr, + min_value, + max_value, + threshold, + ) + else: + range_max = float(threshold) + hist, hist_edges = np.histogram( + data_arr, histogram_collector.num_bins, range=(-range_max, range_max) + ) + histogram_collector.histogram_dict[tensor] = ( + hist, + hist_edges, + min_value, + max_value, + threshold, + ) + + +def _augment_graph_min_max_calibrater_single_node_calibration(calibrater): + """Augment outputs to retrieve MinMax pair. + + Adds ReduceMin and ReduceMax nodes to all quantization_candidates op type nodes in + model and ensures their outputs are stored as part of the graph output. + + :return: augmented ONNX model + + Modification: Add an additional Concat after Reshaped output to not rely on error-prone indexing. + Create multiple single node ONNX models to be used to calibrate per node. + """ + tensors, _ = calibrater.select_tensors_to_calibrate(calibrater.model) + reshape_shape_name = str(uuid.uuid4()) + reshape_shape = onnx.numpy_helper.from_array(np.array([1], dtype=np.int64), reshape_shape_name) + calibrater.model.graph.initializer.append(reshape_shape) + + def add_reduce_min_max(tensor_name): + keepdims = 1 + minmax_output = tensor_name + "_MinMax" + + # Create reduce nodes + reduce_nodes = [ + onnx.helper.make_node( + op_name, + [tensor_name], + [tensor_name + "_" + op_name + "_Reshape"], + keepdims=keepdims, + name=tensor_name + "_" + op_name, + ) + for op_name in ["ReduceMin", "ReduceMax"] + ] + + # Create reshape nodes + reshape_nodes = [ + onnx.helper.make_node( + "Reshape", + inputs=[node.output[0], reshape_shape_name], + outputs=[tensor_name + "_" + op_name], + name=node.output[0], + ) + for node, op_name in zip(reduce_nodes, ["ReduceMin", "ReduceMax"]) + ] + + # Create concat node + concat_node = onnx.helper.make_node( + "Concat", + inputs=[tensor_name + "_ReduceMin", tensor_name + "_ReduceMax"], + outputs=[minmax_output], + name=tensor_name + "_ReduceMin_ReduceMax_Concat", + axis=0, + ) + + calibrater.model.graph.node.extend(reduce_nodes + reshape_nodes + [concat_node]) + + # Get tensor type + value_infos = {vi.name: vi for vi in calibrater.model.graph.value_info} + value_infos.update({o.name: o for o in calibrater.model.graph.output}) + value_infos.update({i.name: i for i in calibrater.model.graph.input}) + + if tensor_name not in value_infos: + raise ValueError( + f"Unable to guess tensor type for tensor {tensor_name!r}, " + f"running shape inference before quantization may resolve this issue." + ) + + calibrater.model.graph.output.append( + onnx.helper.make_tensor_value_info( + minmax_output, value_infos[tensor_name].type.tensor_type.elem_type, [2] + ) + ) + + # Make sure all shapes are resolved before adding min max nodes + calibrater.model = SymbolicShapeInference.infer_shapes(calibrater.model) + + for tensor in tensors: + add_reduce_min_max(tensor) + + # Make sure all shapes are resolved after adding min max nodes + calibrater.model = SymbolicShapeInference.infer_shapes(calibrater.model) + + onnx.save( + calibrater.model, + calibrater.augmented_model_path, + save_as_external_data=calibrater.use_external_data_format, + ) + + # Build single node models and save them + model_counter = 0 + initializer_name_map = { + initializer.name: initializer for initializer in calibrater.model.graph.initializer + } + value_info_name_map = { + value_info.name: value_info for value_info in calibrater.model.graph.value_info + } + input_name_map = {input.name: input for input in calibrater.model.graph.input} + output_name_map = {output.name: output for output in calibrater.model.graph.output} + for node in calibrater.model.graph.node: + single_node_model_name = ( + f"{calibrater.augmented_model_path[:-5]}_single_node_{model_counter}" + ) + single_node_model_node = [] + single_node_model_inputs = [] + single_node_model_outputs = [] + single_node_model_initializers = [] + single_node_model_input_names = [] + single_node_model_output_names = [] + + # Add node + single_node_model_node.append(node) + + # Process each input for node + for input_name in node.input: + # Skip empty tensors + if input_name == "": + continue + + is_input_initializer = False + # If a node input is an initializer, add it to initializer list + if input_name in initializer_name_map: + single_node_model_initializers.append(initializer_name_map[input_name]) + is_input_initializer = True + + value_info_found = False + # If a node input is not an initializer, set it as a model input + if not is_input_initializer: + for name_map in [value_info_name_map, input_name_map, output_name_map]: + if input_name in name_map: + single_node_model_inputs.append(name_map[input_name]) + single_node_model_input_names.append(input_name) + value_info_found = True + break + + if not value_info_found: + raise ValueError( + f"{calibrater.augmented_model_path} is not properly shape inferenced." + ) + + # Process each output for node + for output_name in node.output: + value_info_found = False + for name_map in [value_info_name_map, output_name_map]: + if output_name in name_map: + single_node_model_outputs.append(name_map[output_name]) + single_node_model_output_names.append(output_name) + value_info_found = True + break + + if not value_info_found: + raise ValueError( + f"{calibrater.augmented_model_path} is not properly shape inferenced." + ) + + # Create a new onnx model + single_node_model = onnx.helper.make_model( + onnx.helper.make_graph( + single_node_model_node, + single_node_model_name, + single_node_model_inputs, + single_node_model_outputs, + single_node_model_initializers, + ), + opset_imports=calibrater.model.opset_import, + functions=calibrater.model.functions, + ir_version=calibrater.model.ir_version, + ) + + # Save it to a new onnx file + onnx.save(single_node_model, f"{single_node_model_name}.onnx") + + # Save model info and increase model counter + calibrater.single_node_model_path_map[f"{single_node_model_name}.onnx"] = ( + single_node_model_input_names, + single_node_model_output_names, + ) + model_counter += 1 + + +def _augment_graph_histogram_calibrater_single_node_calibration(calibrater): + """Make all quantization_candidates op type nodes as part of the graph output. + + :return: augmented ONNX model + """ + calibrater.tensors_to_calibrate, value_infos = calibrater.select_tensors_to_calibrate( + calibrater.model + ) + for tensor in calibrater.tensors_to_calibrate: + if tensor not in calibrater.model_original_outputs: + calibrater.model.graph.output.append(value_infos[tensor]) + + onnx.save( + calibrater.model, + calibrater.augmented_model_path, + save_as_external_data=calibrater.use_external_data_format, + ) + + # Build single node models and save them + initializer_name_map = { + initializer.name: initializer for initializer in calibrater.model.graph.initializer + } + value_info_name_map = { + value_info.name: value_info for value_info in calibrater.model.graph.value_info + } + input_name_map = {input.name: input for input in calibrater.model.graph.input} + output_name_map = {output.name: output for output in calibrater.model.graph.output} + model_counter = 0 + for node in calibrater.model.graph.node: + single_node_model_name = ( + f"{calibrater.augmented_model_path[:-5]}_single_node_{model_counter}" + ) + single_node_model_nodes = [] + single_node_model_inputs = [] + single_node_model_outputs = [] + single_node_model_initializers = [] + single_node_model_input_names = [] + single_node_model_output_names = [] + + # Add node + single_node_model_nodes.append(node) + + # Process each input for node + for input_name in node.input: + # Skip empty tensors + if input_name == "": + continue + + is_input_initializer = False + # If a node input is an initializer, add it to initializer list + if input_name in initializer_name_map: + single_node_model_initializers.append(initializer_name_map[input_name]) + is_input_initializer = True + + # If a node input is not an initializer, set it as a model input + if not is_input_initializer: + value_info_found = False + for name_map in [value_info_name_map, input_name_map, output_name_map]: + if input_name in name_map: + single_node_model_inputs.append(name_map[input_name]) + single_node_model_input_names.append(input_name) + value_info_found = True + break + + if not value_info_found: + raise ValueError( + f"{calibrater.augmented_model_path} is not properly shape inferenced." + ) + + # Process each output for node + for output_name in node.output: + value_info_found = False + for name_map in [value_info_name_map, output_name_map]: + if output_name in name_map: + single_node_model_outputs.append(name_map[output_name]) + single_node_model_output_names.append(output_name) + value_info_found = True + break + + if not value_info_found: + raise ValueError( + f"{calibrater.augmented_model_path} is not properly shape inferenced." + ) + + # Create a new onnx model + single_node_model = onnx.helper.make_model( + onnx.helper.make_graph( + single_node_model_nodes, + single_node_model_name, + single_node_model_inputs, + single_node_model_outputs, + single_node_model_initializers, + ), + opset_imports=calibrater.model.opset_import, + functions=calibrater.model.functions, + ir_version=calibrater.model.ir_version, + ) + + # Save it to a new onnx file + onnx.save(single_node_model, f"{single_node_model_name}.onnx") + + # Save model info and increase model counter + calibrater.single_node_model_path_map[f"{single_node_model_name}.onnx"] = ( + single_node_model_input_names, + single_node_model_output_names, + ) + model_counter += 1 diff --git a/modelopt/onnx/quantization/ort_patches.py b/modelopt/onnx/quantization/ort_patches.py new file mode 100644 index 00000000000..9cd68077641 --- /dev/null +++ b/modelopt/onnx/quantization/ort_patches.py @@ -0,0 +1,92 @@ +# 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. + +"""Composition of ModelOpt's ONNX Runtime monkey patches.""" + +__all__ = [] + +from onnxruntime.quantization import calibrate, qdq_quantizer +from onnxruntime.quantization.base_quantizer import BaseQuantizer +from onnxruntime.quantization.calibrate import ( + CalibraterBase, + HistogramCalibrater, + HistogramCollector, + MinMaxCalibrater, +) +from onnxruntime.quantization.qdq_quantizer import QDQQuantizer + +from modelopt.onnx.logging_config import logger +from modelopt.onnx.quantization.ort_calibration import ( + _collect_absolute_value, + _collect_data_histogram_calibrator, + _collect_data_minmax_calibrator, + _collect_value, + _compute_data_minmax_calibrator, + _init_calibrater_base, + _merge_range_minmax_calibrator, + _select_tensors_to_calibrate, +) +from modelopt.onnx.quantization.ort_calibration_per_node import ( + _augment_graph_histogram_calibrater_single_node_calibration, + _augment_graph_min_max_calibrater_single_node_calibration, + _collect_data_histogram_calibrater_single_node_calibration, + _collect_data_min_max_calibrater_single_node_calibration, + _collect_histogram_collector_single_node_calibration, + _collect_value_histogram_collector_single_node_calibration, + _compute_data_min_max_calibrater_single_node_calibration, + _merge_range_min_max_calibrater_single_node_calibration, +) +from modelopt.onnx.quantization.ort_quantization import ( + _adjust_tensor_ranges, + _check_opset_version, + _compute_scale_zp, + _create_calibrator_with_extra_options, +) +from modelopt.onnx.quantization.ort_session import _create_inference_session_with_ep_config + + +def patch_ort_modules(calibrate_per_node: bool = False): + """Patches the ORT modules.""" + logger.debug("Patching ORT modules") + if calibrate_per_node: + MinMaxCalibrater.augment_graph = _augment_graph_min_max_calibrater_single_node_calibration + MinMaxCalibrater.collect_data = _collect_data_min_max_calibrater_single_node_calibration + MinMaxCalibrater.compute_data = _compute_data_min_max_calibrater_single_node_calibration + MinMaxCalibrater.merge_range = _merge_range_min_max_calibrater_single_node_calibration + HistogramCalibrater.augment_graph = ( + _augment_graph_histogram_calibrater_single_node_calibration + ) + HistogramCalibrater.collect_data = ( + _collect_data_histogram_calibrater_single_node_calibration + ) + HistogramCollector.collect = _collect_histogram_collector_single_node_calibration + HistogramCollector.collect_value = ( + _collect_value_histogram_collector_single_node_calibration + ) + else: + HistogramCollector.collect_value = _collect_value + HistogramCollector.collect_absolute_value = _collect_absolute_value + MinMaxCalibrater.compute_data = _compute_data_minmax_calibrator + MinMaxCalibrater.collect_data = _collect_data_minmax_calibrator + MinMaxCalibrater.merge_range = _merge_range_minmax_calibrator + HistogramCalibrater.collect_data = _collect_data_histogram_calibrator + + calibrate.create_calibrator = _create_calibrator_with_extra_options + CalibraterBase.create_inference_session = _create_inference_session_with_ep_config + CalibraterBase.select_tensors_to_calibrate = _select_tensors_to_calibrate + QDQQuantizer.check_opset_version = _check_opset_version + BaseQuantizer.adjust_tensor_ranges = _adjust_tensor_ranges + qdq_quantizer.compute_scale_zp = _compute_scale_zp + CalibraterBase.__init__ = _init_calibrater_base diff --git a/modelopt/onnx/quantization/ort_patching.py b/modelopt/onnx/quantization/ort_patching.py deleted file mode 100755 index d2957ec183f..00000000000 --- a/modelopt/onnx/quantization/ort_patching.py +++ /dev/null @@ -1,1845 +0,0 @@ -# Adapted from https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/quant_utils.py -# and https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/calibrate.py -# and https://github.com/microsoft/onnxruntime/blob/2ac381c55397dffff327cc6efecf6f95a70f90a1/onnxruntime/python/tools/quantization/onnx_quantizer.py -# -# MIT License -# -# Copyright (c) Microsoft Corporation -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# -# The above copyright notice and this permission notice shall be included in all -# copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. - -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 AND MIT -# -# 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. - -"""This module contains all the patched functions from ORT.""" - -import gc -import tempfile -import uuid -from collections.abc import Sequence -from pathlib import Path - -import numpy as np -import onnx -import onnxruntime as ort -import pynvml -from onnx import onnx_pb -from onnxruntime.quantization import calibrate, qdq_quantizer -from onnxruntime.quantization.base_quantizer import BaseQuantizer -from onnxruntime.quantization.calibrate import ( - CalibraterBase, - CalibrationDataReader, - CalibrationMethod, - DistributionCalibrater, - EntropyCalibrater, - HistogramCalibrater, - HistogramCollector, - MinMaxCalibrater, - PercentileCalibrater, - TensorData, - TensorsData, -) -from onnxruntime.quantization.qdq_quantizer import QDQQuantizer -from onnxruntime.quantization.quant_utils import ( - QuantFormat, - QuantizationMode, - QuantType, - add_infer_metadata, -) -from onnxruntime.quantization.quant_utils import compute_scale_zp as _ort_compute_scale_zp -from onnxruntime.quantization.quantize import check_static_quant_arguments -from onnxruntime.quantization.registry import QDQRegistry, QLinearOpsRegistry -from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference -from tqdm import tqdm - -import modelopt.onnx.utils as onnx_utils -from modelopt.onnx.logging_config import logger - - -def load_model_with_shape_infer(model_path: Path) -> onnx.ModelProto: - """Load model while performing symbolic shape infer and ONNX shape inference.""" - model = onnx.load(str(model_path), load_external_data=True) - try: - model = onnx_utils.infer_shapes(model) - add_infer_metadata(model) - except Exception as e: - logger.info(f"Failed to infer shapes for model {model_path}: {e}") - return model - - -def _compute_scale_zp(rmin, rmax, qmin, qmax, symmetric=False, min_real_range=None): - """Retry FP16 scale calculation in FP32 when range subtraction overflows.""" - range_dtype = np.asarray(rmax).dtype - if range_dtype != np.float16: - return _ort_compute_scale_zp(rmin, rmax, qmin, qmax, symmetric, min_real_range) - - with np.errstate(over="ignore", invalid="ignore"): - zero_point, scale = _ort_compute_scale_zp(rmin, rmax, qmin, qmax, symmetric, min_real_range) - if np.all(np.isfinite(scale)): - return zero_point, scale - - zero_point, scale = _ort_compute_scale_zp( - np.asarray(rmin, dtype=np.float32), - np.asarray(rmax, dtype=np.float32), - qmin, - qmax, - symmetric, - min_real_range, - ) - return zero_point, np.asarray(scale, dtype=range_dtype) - - -def _prepare_histogram_data(histogram_collector, tensor, data_arr): - """Use FP32 for histogram math while remembering the source dtype.""" - if data_arr.dtype != np.float16: - return data_arr - - original_dtypes = getattr(histogram_collector, "_modelopt_original_dtypes", {}) - original_dtypes[tensor] = data_arr.dtype - histogram_collector._modelopt_original_dtypes = original_dtypes - return data_arr.astype(np.float32) - - -def _restore_histogram_calibration_dtypes(histogram_collector, tensors_range): - """Restore source dtypes at the calibration-to-quantization boundary.""" - original_dtypes = getattr(histogram_collector, "_modelopt_original_dtypes", {}) - for tensor, dtype in original_dtypes.items(): - if tensor not in tensors_range: - continue - tensor_data = tensors_range[tensor] - dtype_limits = np.finfo(dtype) - for attribute in ("lowest", "highest", "avg", "std"): - if hasattr(tensor_data, attribute): - value = np.clip(getattr(tensor_data, attribute), dtype_limits.min, dtype_limits.max) - setattr(tensor_data, attribute, np.asarray(value, dtype=dtype)) - - -def _collect_value(histogram_collector, name_to_arr): - """Collect histogram on real value.""" - for tensor, data_arr in tqdm(name_to_arr.items()): - # ====================== Modification ====================== - concat_data_arr = np.asarray(data_arr[0]) - concat_data_arr = concat_data_arr.flatten() - for i in range(1, len(data_arr)): - curr_data_arr = np.asarray(data_arr[i]) - curr_data_arr = curr_data_arr.flatten() - concat_data_arr = np.concatenate((concat_data_arr, curr_data_arr)) - - concat_data_arr = _prepare_histogram_data(histogram_collector, tensor, concat_data_arr) - data_arr = concat_data_arr - # ========================================================== - if data_arr.size > 0: - min_value = np.min(data_arr) - max_value = np.max(data_arr) - else: - min_value = np.array(0, dtype=data_arr.dtype) - max_value = np.array(0, dtype=data_arr.dtype) - - # Change the inf and nan values to meaningful min/max - min_value = ( - np.finfo(np.float32).tiny if np.isinf(min_value) or np.isnan(min_value) else min_value - ) - max_value = ( - np.finfo(np.float32).max if np.isinf(max_value) or np.isnan(max_value) else max_value - ) - - threshold = max(abs(min_value), abs(max_value)) - - if tensor in histogram_collector.histogram_dict: - old_histogram = histogram_collector.histogram_dict[tensor] - histogram_collector.histogram_dict[tensor] = histogram_collector.merge_histogram( - old_histogram, data_arr, min_value, max_value, threshold - ) - else: - range_max = float(threshold) - hist, hist_edges = np.histogram( - data_arr, histogram_collector.num_bins, range=(-range_max, range_max) - ) - histogram_collector.histogram_dict[tensor] = ( - hist, - hist_edges, - min_value, - max_value, - threshold, - ) - - -def _collect_absolute_value(histogram_collector, name_to_arr): - """Collect histogram on absolute value.""" - for tensor, data_arr in name_to_arr.items(): - if isinstance(data_arr, list): - for arr in data_arr: - assert isinstance(arr, np.ndarray), ( - f"Unexpected type {type(arr)} for tensor={tensor!r}" - ) - dtypes = {a.dtype for a in data_arr} - assert len(dtypes) == 1, ( - f"The calibration expects only one element type but got {dtypes} for tensor={tensor!r}" - ) - # ====================== Modification ====================== - concat_data_arr = np.asarray(data_arr[0]) - concat_data_arr = concat_data_arr.flatten() - for i in range(1, len(data_arr)): - curr_data_arr = np.asarray(data_arr[i]) - curr_data_arr = curr_data_arr.flatten() - concat_data_arr = np.concatenate((concat_data_arr, curr_data_arr)) - data_arr_np = concat_data_arr - # ========================================================== - elif not isinstance(data_arr, np.ndarray): - raise ValueError(f"Unexpected type {type(data_arr)} for tensor={tensor!r}") - else: - data_arr_np = data_arr - data_arr_np = data_arr_np.flatten() - if data_arr_np.size > 0: - min_value = np.min(data_arr_np) - max_value = np.max(data_arr_np) - else: - min_value = np.array(0, dtype=data_arr_np.dtype) - max_value = np.array(0, dtype=data_arr_np.dtype) - - data_arr_np = np.absolute(data_arr_np) # only consider absolute value - - if tensor not in histogram_collector.histogram_dict: - # first time it uses num_bins to compute histogram. - hist, hist_edges = np.histogram(data_arr_np, bins=histogram_collector.num_bins) - hist_edges = hist_edges.astype(data_arr_np.dtype) - assert data_arr_np.dtype != np.float64, ( - "only float32 or float16 is supported, every constant must be explicitly typed" - ) - histogram_collector.histogram_dict[tensor] = (hist, hist_edges, min_value, max_value) - else: - old_histogram = histogram_collector.histogram_dict[tensor] - old_min = old_histogram[2] - old_max = old_histogram[3] - assert hasattr(old_min, "dtype"), ( - f"old_min should be a numpy array but is {type(old_min)}" - ) - assert hasattr(old_max, "dtype"), ( - f"old_min should be a numpy array but is {type(old_max)}" - ) - old_hist = old_histogram[0] - old_hist_edges = old_histogram[1] - temp_amax = np.max(data_arr_np) - if temp_amax > old_hist_edges[-1]: - # increase the number of bins - width = old_hist_edges[1] - old_hist_edges[0] - # NOTE: np.arange may create an extra bin after the one containing temp_amax - new_bin_edges = np.arange(old_hist_edges[-1] + width, temp_amax + width, width) - old_hist_edges = np.hstack((old_hist_edges, new_bin_edges)) - hist, hist_edges = np.histogram(data_arr_np, bins=old_hist_edges) - hist_edges = hist_edges.astype(data_arr_np.dtype) - hist[: len(old_hist)] += old_hist - assert data_arr_np.dtype != np.float64, ( - "only float32 or float16 is supported, every constant must be explicitly typed" - ) - histogram_collector.histogram_dict[tensor] = ( - hist, - hist_edges, - min(old_min, min_value), - max(old_max, max_value), - ) - - -def _check_opset_version(onnx_quantizer): - ai_onnx_domain = [ - opset - for opset in onnx_quantizer.model.model.opset_import - if not opset.domain or opset.domain in ["ai.onnx", "ai.onnx.contrib"] - ] - opset_version = ai_onnx_domain[0].version - - if opset_version == 10: - return 10 - - if opset_version < 10: - onnx_quantizer.model.model.opset_import.remove(ai_onnx_domain[0]) - onnx_quantizer.model.model.opset_import.extend([onnx.helper.make_opsetid("", 11)]) - opset_version = 11 - - if opset_version < 19 and onnx_quantizer.weight_qType == onnx_pb.TensorProto.FLOAT8E4M3FN: - onnx_quantizer.model.model.opset_import.remove(ai_onnx_domain[0]) - onnx_quantizer.model.model.opset_import.extend([onnx.helper.make_opsetid("", 19)]) - # Set ir_version to 10, remove it once ORT supports ir_version 11 - onnx_quantizer.model.model.ir_version = 10 - opset_version = 19 - - onnx_quantizer.fuse_dynamic_quant = True - return opset_version - - -def _select_tensors_to_calibrate(calibrator, model: onnx.ModelProto): - """Select input/output tensors of candidate nodes to calibrate. - - Returns: - tensors (set): set of tensor name. - value_infos (dict): tensor name to value info. - """ - value_infos = {vi.name: vi for vi in model.graph.value_info} - value_infos.update({ot.name: ot for ot in model.graph.output}) - value_infos.update({it.name: it for it in model.graph.input}) - initializer = {init.name for init in model.graph.initializer} - - tensors_to_calibrate = set() - tensor_type_to_calibrate = {onnx_pb.TensorProto.FLOAT, onnx_pb.TensorProto.FLOAT16} - - for node in model.graph.node: - # Hack: in calibrator.op_types_to_calibrate we pass nodes_to_quantize - if node.name in calibrator.op_types_to_calibrate: - for tensor_name in node.input: - if tensor_name in value_infos: - vi = value_infos[tensor_name] - if ( - vi.type.HasField("tensor_type") - and (vi.type.tensor_type.elem_type in tensor_type_to_calibrate) - and (tensor_name not in initializer) - ): - tensors_to_calibrate.add(tensor_name) - for tensor_name in node.output: - if tensor_name in value_infos: - vi = value_infos[tensor_name] - if vi.type.HasField("tensor_type") and ( - vi.type.tensor_type.elem_type in tensor_type_to_calibrate - ): - tensors_to_calibrate.add(tensor_name) - - return tensors_to_calibrate, value_infos - - -def _configure_session_providers( - sess_options: ort.SessionOptions, - providers: list[str | tuple[str, dict]], - trt_rtx_backend: str, -) -> dict[str, list[str | tuple[str, dict]]]: - """Configure providers using the mechanism required by the selected EP. - - ``providers`` contains provider names or ``(name, options)`` pairs in priority order. - ABI EPs are exposed as devices and must be added to ``sess_options``; passing them through - ``InferenceSession(providers=...)`` overrides that configuration. This helper preserves the - ABI device path while returning normal provider arguments for other EPs. - """ - if trt_rtx_backend != "abi": - return {"providers": providers} - - available_providers = set(ort.get_available_providers()) - ep_devices = ort.get_ep_devices() - plugin_provider_names = {device.ep_name for device in ep_devices} - available_providers - provider_names = { - provider[0] if isinstance(provider, tuple) else provider for provider in providers - } - if not plugin_provider_names.intersection(provider_names): - return {"providers": providers} - - for provider in providers: - provider_name, provider_options = ( - provider if isinstance(provider, tuple) else (provider, {}) - ) - if provider_name in plugin_provider_names: - selected_devices = [device for device in ep_devices if device.ep_name == provider_name] - sess_options.add_provider_for_devices(selected_devices, provider_options) - else: - sess_options.add_provider(provider_name, provider_options) - return {} - - -def _create_inference_session_with_ep_config(calibrator, **kwargs): - """Create an ORT InferenceSession.""" - model_path = kwargs.get("model_path") - logger.debug("Creating inference session with Execution Provider configuration") - - trt_rtx_backend = kwargs.get("trt_rtx_backend", "legacy") - if trt_rtx_backend not in ("legacy", "abi"): - raise ValueError(f"trt_rtx_backend must be 'legacy' or 'abi', got {trt_rtx_backend!r}") - - sess_options = ort.SessionOptions() - sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL - sess_options.add_session_config_entry("session.use_device_allocator_for_initializers", "1") - sess_options.enable_cpu_mem_arena = False - - providers = kwargs.get("execution_providers", []) - logger.debug(f"Execution providers: {providers}") - - # Note. This path can be an empty string, which denotes that the model has custom ops and TRT EP is needed. - calibrator.trt_extra_plugin_lib_paths = kwargs.get("trt_extra_plugin_lib_paths") - if calibrator.trt_extra_plugin_lib_paths is not None: - logger.debug(f"TRT extra plugin paths: {calibrator.trt_extra_plugin_lib_paths}") - if "TensorrtExecutionProvider" not in ort.get_available_providers(): - raise RuntimeError( - f"Could not find `TensorrtExecutionProvider`, only {ort.get_available_providers()}" - ) - trt_ep_options = ( - {"trt_extra_plugin_lib_paths": calibrator.trt_extra_plugin_lib_paths} - if calibrator.trt_extra_plugin_lib_paths - else {} - ) - - # Set GPU memory usage limit - trt_ep_options["trt_max_workspace_size"] = 80 * (1024**3) # 80GB - logger.debug(f"TRT EP options: {trt_ep_options}") - - if "TensorrtExecutionProvider" in providers: - providers.remove("TensorrtExecutionProvider") - providers.insert(0, ("TensorrtExecutionProvider", trt_ep_options)) - - def _update_provider_config(provider, config): - if isinstance(provider, tuple) and len(provider) > 1 and isinstance(provider[1], dict): - provider[1].update(config) - else: - provider = (provider, config) - return provider - - for i in range(len(providers)): - if any(p in providers[i] for p in ["CPUExecutionProvider", "CUDAExecutionProvider"]): - providers[i] = _update_provider_config( - providers[i], {"arena_extend_strategy": "kSameAsRequested"} - ) - - session_path = calibrator.augmented_model_path if model_path is None else model_path - provider_kwargs = _configure_session_providers(sess_options, providers, trt_rtx_backend) - calibrator.infer_session = ort.InferenceSession( - session_path, - sess_options=sess_options, - **provider_kwargs, - ) - - # Group qdq tensors will have the same scaling factor. - calibrator.group_qdq_tensors = kwargs.get("group_qdq_tensors") - if calibrator.group_qdq_tensors: - logger.debug(f"Group QDQ tensors: {calibrator.group_qdq_tensors}") - - -def _compute_data_minmax_calibrator(calibrator): - """Compute the min-max range of tensor. - - :returns: dictionary mapping: {added node names: (ReduceMin, ReduceMax) pairs } - """ - if len(calibrator.intermediate_outputs) == 0: - return calibrator.calibrate_tensors_range - - output_names = [ - calibrator.infer_session.get_outputs()[i].name - for i in range(len(calibrator.intermediate_outputs[0])) - ] - - output_dicts_list = [ - dict(zip(output_names, intermediate_output), strict=True) - for intermediate_output in calibrator.intermediate_outputs - ] - - merged_output_dict = {} - for d in output_dicts_list: - for k, v in d.items(): - merged_output_dict.setdefault(k, []).append(v) - - # ====================== Modification ====================== - # Group qdq tensors should have the same scaling factor. Each tensor in group should add - # other tensors in its merged_dict value. In this way, calibrator will generate the same - # scaling factor. - if calibrator.group_qdq_tensors: - for cur, group in calibrator.group_qdq_tensors.items(): - for other in group: - if cur == other: - continue - for d in output_dicts_list: - for k, v in d.items(): - cur_min = cur + "_" + "ReduceMin" - cur_max = cur + "_" + "ReduceMax" - other_min = other + "_" + "ReduceMin" - other_max = other + "_" + "ReduceMax" - if k == other_min: - merged_output_dict[cur_min].append(v) - elif k == other_max: - merged_output_dict[cur_max].append(v) - # ============================================================ - - added_output_names = output_names[calibrator.num_model_outputs :] - calibrate_tensor_names = [ - added_output_names[i].rpartition("_")[0] for i in range(0, len(added_output_names), 2) - ] # output names - - merged_added_output_dict = { - i: merged_output_dict[i] - for i in merged_output_dict - if i not in calibrator.model_original_outputs - } - - pairs = [] - for i in range(0, len(added_output_names), 2): - if calibrator.moving_average: - min_value_array = np.mean(merged_added_output_dict[added_output_names[i]], axis=0) - max_value_array = np.mean(merged_added_output_dict[added_output_names[i + 1]], axis=0) - else: - min_value_array = np.min(merged_added_output_dict[added_output_names[i]], axis=0) - max_value_array = np.max(merged_added_output_dict[added_output_names[i + 1]], axis=0) - - if calibrator.symmetric: - max_absolute_value = np.max([np.abs(min_value_array), np.abs(max_value_array)], axis=0) - pairs.append((-max_absolute_value, max_absolute_value)) - else: - pairs.append((min_value_array, max_value_array)) - - new_calibrate_tensors_range = TensorsData( - CalibrationMethod.MinMax, dict(zip(calibrate_tensor_names, pairs, strict=False)) - ) - if calibrator.calibrate_tensors_range: - calibrator.calibrate_tensors_range = calibrator.merge_range( - calibrator.calibrate_tensors_range, new_calibrate_tensors_range - ) - else: - calibrator.calibrate_tensors_range = new_calibrate_tensors_range - - return calibrator.calibrate_tensors_range - - -def _compute_data_min_max_calibrater_single_node_calibration(calibrater) -> TensorData: - """Compute the min-max range of tensor. - - :return: dictionary mapping: {added node names: (ReduceMin, ReduceMax) pairs } - - Modification: Instead of aggregating two consecutive outputs to a MinMax pair, retrieve a MinMax pair from - outputs of Concat. - """ - if not calibrater.intermediate_outputs: - return calibrater.calibrate_tensors_range - - # Get output names and merge all intermediate outputs - output_names = [out.name for out in calibrater.infer_session.get_outputs()] - - # Merge outputs across all batches, filtering out original model outputs - merged_outputs = {} - for intermediate_output in calibrater.intermediate_outputs: - for name, value in zip(output_names, intermediate_output): - if name not in calibrater.model_original_outputs: - merged_outputs.setdefault(name, []).append(value) - - # Compute min/max pairs for each tensor - pairs = [] - tensor_names = [] - - for output_name, values in merged_outputs.items(): - tensor_names.append(output_name.rpartition("_")[0]) - - if calibrater.moving_average: - min_val, max_val = np.mean(values, axis=0) - else: - stacked_values = np.stack(values, axis=0) - min_val = np.min(stacked_values, axis=0)[0] - max_val = np.max(stacked_values, axis=0)[1] - - if calibrater.symmetric: - max_abs = max(np.abs(min_val), np.abs(max_val)) - pairs.append((-max_abs, max_abs)) - else: - pairs.append((min_val, max_val)) - - # Create and merge tensor range data - new_range = TensorsData(CalibrationMethod.MinMax, dict(zip(tensor_names, pairs))) - - calibrater.calibrate_tensors_range = ( - calibrater.merge_range(calibrater.calibrate_tensors_range, new_range) - if calibrater.calibrate_tensors_range - else new_range - ) - - return calibrater.calibrate_tensors_range - - -def _collect_data_minmax_calibrator(calibrator, data_reader: CalibrationDataReader): - """This function overwrite is needed to solve OOM issue due to the unlimited accumulation of intermediate_outputs. - - Support for: MinMax Calibrator. - Modification: indented the last lines of code inside the while loop in order to run compute_data for each sample - batch individually instead of the entire data at once. The assumption here is that the ONNX file has bs=N - and the calibration data size is M (where M is a multiple of N). So the calibrator is a sequence of M/N - samples with bs=N. - """ - run_options = ort.RunOptions() - try: - pynvml.nvmlInit() - gpu_count = pynvml.nvmlDeviceGetCount() - pynvml.nvmlShutdown() - except Exception as e: - logger.error(f"Failed to get GPU count: {e}") - gpu_count = 0 - gpu_str = ";".join([f"gpu:{i}" for i in range(gpu_count)]) - run_options.add_run_config_entry("memory.enable_memory_arena_shrinkage", f"cpu:0;{gpu_str}") - while True: - inputs = data_reader.get_next() - if not inputs: - break - run_options = ort.RunOptions() - - calibrator.intermediate_outputs.append( - calibrator.infer_session.run(None, inputs, run_options=run_options) - ) - - # ======== Modification: block is indentend in ======== - if len(calibrator.intermediate_outputs) == 0: - raise ValueError("No data is collected.") - - t = calibrator.compute_data() - if not isinstance(t, TensorsData): - raise TypeError(f"compute_data must return a TensorsData not {type(t)}.") - calibrator.clear_collected_data() - # ===================================================== - - -def _merge_range_minmax_calibrator(calibrator, old_range: TensorsData, new_range: TensorsData): - """This function is an auxiliary function of collect_data to solve the OOM issue in the MinMax Calibrator. - - Issue fixed with this function: old_range is not a dictionary, but old_range.data is. - TODO: create an MR in the ORT repository for this function. Alternatively, we can also file the MR fixing - TensorData (need to at least add items() function there). - """ - if not old_range: - return new_range - - for key, value in old_range.data.items(): - value_tuple = value.range_value - new_range_tuple = new_range.data[key].range_value - if calibrator.moving_average: - min_value = value_tuple[0] + calibrator.averaging_constant * ( - new_range_tuple[0] - value_tuple[0] - ) - max_value = value_tuple[1] + calibrator.averaging_constant * ( - new_range_tuple[1] - value_tuple[1] - ) - else: - min_value = min(value_tuple[0], new_range_tuple[0]) - max_value = max(value_tuple[1], new_range_tuple[1]) - new_range.data[key] = TensorData(lowest=min_value, highest=max_value) - - return new_range - - -def _merge_range_min_max_calibrater_single_node_calibration( - calibrater, old_range: TensorsData, new_range: TensorsData -): - """This function is an auxiliary function of collect_data to solve the OOM issue in the MinMax Calibrator. - - Issue fixed with this function: old_range is not a dictionary, but old_range.data is. - TODO: create an MR in the ORT repository for this function. Alternatively, we can also file the MR fixing - TensorData (need to at least add items() function there). - """ - if not old_range: - return new_range - - def _merge_ranges(old_min, old_max, new_min, new_max): - if calibrater.moving_average: - alpha = calibrater.averaging_constant - return (old_min + alpha * (new_min - old_min), old_max + alpha * (new_max - old_max)) - return min(old_min, new_min), max(old_max, new_max) - - old_data = old_range.data - for key, new_tensor in new_range.data.items(): - if key in old_data: - old_min, old_max = old_data[key].range_value - new_min, new_max = new_tensor.range_value - merged_min, merged_max = _merge_ranges(old_min, old_max, new_min, new_max) - old_data[key] = TensorData(lowest=merged_min, highest=merged_max) - else: - old_data[key] = new_tensor - - return old_range - - -def _collect_data_histogram_calibrator(calibrator, data_reader: CalibrationDataReader): - """This function overwrite is needed to solve OOM issue due to the unlimited accumulation of intermediate_outputs. - - Support for: Histogram Calibrator (which affects Entropy, Percentile, and DIstribution Calibrators). - Modification: indented the last lines of code inside the while loop in order to run compute_data for each sample - batch individually instead of the entire data at once. - """ - while True: - inputs = data_reader.get_next() - if not inputs: - break - calibrator.intermediate_outputs.append(calibrator.infer_session.run(None, inputs)) - - # ======== Modification: block is indentend in ======== - # Here, compute_date is calculated for every sample batch instead of the entire data at once. - if len(calibrator.intermediate_outputs) == 0: - raise ValueError("No data is collected.") - - output_names = [ - calibrator.infer_session.get_outputs()[i].name - for i in range(len(calibrator.intermediate_outputs[0])) - ] - output_dicts_list = [ - dict(zip(output_names, intermediate_output)) - for intermediate_output in calibrator.intermediate_outputs - ] - - merged_dict = {} - for d in output_dicts_list: - for k, v in d.items(): - merged_dict.setdefault(k, []).append(v) - - # Group qdq tensors should have the same scaling factor. Each tensor in group should add - # other tensors in its merged_dict value. In this way, calibrator will generate the same - # scaling factor. - if calibrator.group_qdq_tensors: - for cur, group in calibrator.group_qdq_tensors.items(): - for other in group: - if cur == other: - continue - for d in output_dicts_list: - for k, v in d.items(): - if k == other: - merged_dict[cur].append(v) - - clean_merged_dict = { - i: merged_dict[i] for i in merged_dict if i in calibrator.tensors_to_calibrate - } - - if not calibrator.collector: - calibrator.collector = HistogramCollector( - method=calibrator.method, - symmetric=calibrator.symmetric, - num_bins=calibrator.num_bins, - num_quantized_bins=calibrator.num_quantized_bins, - percentile=calibrator.percentile, - scenario=calibrator.scenario, - ) - calibrator.collector.collect(clean_merged_dict) - - calibrator.clear_collected_data() - # ===================================================== - - -def _collect_data_min_max_calibrater_single_node_calibration( - calibrater, data_reader: CalibrationDataReader -): - """Collects calibration data (min/max) for a MinMax Calibrator by processing single-node models batch by batch. - - This function addresses an OOM issue by computing calibration data for each batch individually, - rather than accumulating all intermediate outputs across the entire dataset. It assumes the ONNX model - has a batch size of N, and the calibration data size M is a multiple of N, processing M/N batches. - - Args: - calibrater: The calibrater object managing model inference and data collection. - data_reader: Provides batches of input data for calibration. - """ - input_counter = 0 - while True: - inputs = data_reader.get_next() - if not inputs: - break - logger.debug(f"Collecting tensor data and finding min & max for input #{input_counter}") - - # We are using single node model scheme. Set up model to input dependency map - model_to_input_dep_map = {} - for model_path, io_tensors in calibrater.single_node_model_path_map.items(): - model_to_input_dep_map[model_path] = io_tensors[0].copy() # List of input names - - # Setup input queues - input_queue = [model_input.name for model_input in calibrater.model.graph.input] - - pbar = tqdm(total=len(model_to_input_dep_map.keys())) - # Resolve nodes are independent from inputs to add their outputs as inputs - inferred_model_list = [] - for model_path, input_deps in model_to_input_dep_map.items(): - if len(input_deps) == 0: - calibrater.create_inference_session( - execution_providers=calibrater.providers, - trt_extra_plugin_lib_paths=calibrater.trt_extra_plugin_lib_paths, - model_path=model_path, - ) - outputs = calibrater.infer_session.run(None, {}) - - # Add output to inputs - need_calibration = False - for output_idx, output in enumerate(outputs): - output_name = calibrater.infer_session.get_outputs()[output_idx].name - inputs[output_name] = output - input_queue.append(output_name) - if ( - output_name - in [output_tensor.name for output_tensor in calibrater.model.graph.output] - and output_name not in calibrater.model_original_outputs - ): - need_calibration = True - - # Mark model path to remove it from dependency map - inferred_model_list.append(model_path) - - # For each inference, compute data before moving to other nodes if tensor is to be calibrated - if need_calibration: - calibrater.intermediate_outputs.append(outputs) - if len(calibrater.intermediate_outputs) == 0: - raise ValueError("No data is collected.") - - t = calibrater.compute_data() - if not isinstance(t, TensorsData): - raise TypeError(f"compute_data must return a TensorsData not {type(t)}.") - calibrater.clear_collected_data() - - gc.collect() - - pbar.update(1) - - # Remove inferred model from dependency map - for model_path in inferred_model_list: - model_to_input_dep_map.pop(model_path) - - gc.collect() - - # Process topological inference - input_ref_count = {} - while input_queue: - current_input_name = input_queue.pop(0) - - # Initialize input reference count - input_ref_count[current_input_name] = sum( - current_input_name in input_deps for input_deps in model_to_input_dep_map.values() - ) - - # Perform inference - inferred_model_list = [] - for model_path, input_deps in model_to_input_dep_map.items(): - if current_input_name in input_deps: - input_deps.remove(current_input_name) - - # If all dependencies are met, perform inference for the node. - if len(input_deps) == 0: - # Make dictionary of only needed inputs. - inputs_to_feed = {} - for input_name_to_feed in calibrater.single_node_model_path_map[model_path][ - 0 - ]: - inputs_to_feed[input_name_to_feed] = inputs[input_name_to_feed] - - calibrater.create_inference_session( - execution_providers=calibrater.providers, - trt_extra_plugin_lib_paths=calibrater.trt_extra_plugin_lib_paths, - model_path=model_path, - ) - outputs = calibrater.infer_session.run(None, inputs_to_feed) - - # Mark model path to remove it from dependency map - inferred_model_list.append(model_path) - - # Decrease reference count for used inputs and remove if no reference - for input_name in calibrater.single_node_model_path_map[model_path][0]: - input_ref_count[input_name] -= 1 - if input_ref_count[input_name] == 0: - del inputs[input_name] - del input_ref_count[input_name] - - gc.collect() - - # Add outputs to inputs - need_calibration = False - for output_idx, output in enumerate(outputs): - output_name = calibrater.infer_session.get_outputs()[output_idx].name - inputs[output_name] = output - input_queue.append(output_name) - if ( - output_name - in [ - output_tensor.name - for output_tensor in calibrater.model.graph.output - ] - and output_name not in calibrater.model_original_outputs - ): - need_calibration = True - - # For each inference, compute data before moving to other nodes if tensor is to be calibrated - if need_calibration: - calibrater.intermediate_outputs.append(outputs) - if len(calibrater.intermediate_outputs) == 0: - raise ValueError("No data is collected.") - - t = calibrater.compute_data() - if not isinstance(t, TensorsData): - raise TypeError( - f"compute_data must return a TensorsData not {type(t)}." - ) - calibrater.clear_collected_data() - - gc.collect() - - pbar.update(1) - - # Remove inferred model from dependency map - for model_path in inferred_model_list: - model_to_input_dep_map.pop(model_path) - - gc.collect() - pbar.close() - input_counter += 1 - - -def _collect_data_histogram_calibrater_single_node_calibration(calibrator, data_reader): - """Collects histogram data for single-node calibration, processing batches to avoid OOM. - - Args: - calibrator: Histogram calibrator instance. - data_reader: CalibrationDataReader providing input data. - """ - input_counter = 0 - while True: - inputs = data_reader.get_next() - if not inputs: - break - logger.debug(f"Collecting tensor data for input #{input_counter}") - - # We are using single node model scheme. Set up model to input dependency map - model_to_input_dep_map = {} - for model_path, io_tensors in calibrator.single_node_model_path_map.items(): - model_to_input_dep_map[model_path] = io_tensors[0].copy() # List of input names - - # Compute data for input tensors - input_only_model = onnx.helper.make_model( - onnx.helper.make_graph( - [], - f"{calibrator.augmented_model_path[:-5]}_input_only", - calibrator.model.graph.input, - calibrator.model.graph.input, - ), - opset_imports=calibrator.model.opset_import, - functions=calibrator.model.functions, - ir_version=calibrator.model.ir_version, - ) - calibrator.infer_session = ort.InferenceSession(input_only_model.SerializeToString()) - calibrator.intermediate_outputs.append( - [ - inputs[calibrator.infer_session.get_outputs()[i].name] - for i in range(len(calibrator.infer_session.get_outputs())) - ] - ) - if len(calibrator.intermediate_outputs) == 0: - raise ValueError("No data is collected.") - - output_names = [ - calibrator.infer_session.get_outputs()[i].name - for i in range(len(calibrator.intermediate_outputs[0])) - ] - output_dicts_list = [ - dict(zip(output_names, intermediate_output)) - for intermediate_output in calibrator.intermediate_outputs - ] - - merged_dict = {} - for d in output_dicts_list: - for k, v in d.items(): - merged_dict.setdefault(k, []).append(v) - - clean_merged_dict = { - i: merged_dict[i] for i in merged_dict if i in calibrator.tensors_to_calibrate - } - - if not calibrator.collector: - calibrator.collector = HistogramCollector( - method=calibrator.method, - symmetric=calibrator.symmetric, - num_bins=calibrator.num_bins, - num_quantized_bins=calibrator.num_quantized_bins, - percentile=calibrator.percentile, - scenario=calibrator.scenario, - ) - calibrator.collector.collect(clean_merged_dict) - - calibrator.clear_collected_data() - - gc.collect() - - # Setup input queues - input_queue = [model_input.name for model_input in calibrator.model.graph.input] - - pbar = tqdm(total=len(model_to_input_dep_map.keys())) - # Resolve nodes are independent from inputs to add their outputs as inputs - inferred_model_list = [] - for model_path, input_deps in model_to_input_dep_map.items(): - if len(input_deps) == 0: - calibrator.create_inference_session( - execution_providers=calibrator.providers, - trt_extra_plugin_lib_paths=calibrator.trt_extra_plugin_lib_paths, - model_path=model_path, - ) - outputs = calibrator.infer_session.run(None, {}) - - # Add output to inputs - need_calibration = False - for output_idx, output in enumerate(outputs): - output_name = calibrator.infer_session.get_outputs()[output_idx].name - inputs[output_name] = output - input_queue.append(output_name) - if output_name in calibrator.tensors_to_calibrate: - need_calibration = True - - # Mark model path to remove it from dependency map - inferred_model_list.append(model_path) - - # For each inference, compute data before moving to other nodes if tensor is to be calibrated - if need_calibration: - calibrator.intermediate_outputs.append(outputs) - if len(calibrator.intermediate_outputs) == 0: - raise ValueError("No data is collected.") - - output_names = [ - calibrator.infer_session.get_outputs()[i].name - for i in range(len(calibrator.intermediate_outputs[0])) - ] - output_dicts_list = [ - dict(zip(output_names, intermediate_output)) - for intermediate_output in calibrator.intermediate_outputs - ] - - merged_dict = {} - for d in output_dicts_list: - for k, v in d.items(): - merged_dict.setdefault(k, []).append(v) - - clean_merged_dict = { - i: merged_dict[i] - for i in merged_dict - if i in calibrator.tensors_to_calibrate - } - - if not calibrator.collector: - calibrator.collector = HistogramCollector( - method=calibrator.method, - symmetric=calibrator.symmetric, - num_bins=calibrator.num_bins, - num_quantized_bins=calibrator.num_quantized_bins, - percentile=calibrator.percentile, - scenario=calibrator.scenario, - ) - calibrator.collector.collect(clean_merged_dict) - - calibrator.clear_collected_data() - - gc.collect() - pbar.update(1) - - # Remove inferred model from dependency map - for model_path in inferred_model_list: - model_to_input_dep_map.pop(model_path) - - gc.collect() - - # Process topological inference - input_ref_count = {} - while input_queue: - current_input_name = input_queue.pop(0) - - # Initialize input reference count - input_ref_count[current_input_name] = sum( - current_input_name in input_deps for input_deps in model_to_input_dep_map.values() - ) - - # Perform inference - inferred_model_list = [] - for model_path, input_deps in model_to_input_dep_map.items(): - if current_input_name in input_deps: - input_deps.remove(current_input_name) - - # If all dependencies are met, perform inference for the node. - if len(input_deps) == 0: - # Make dictionary of only needed inputs. - inputs_to_feed = {} - for input_name_to_feed in calibrator.single_node_model_path_map[model_path][ - 0 - ]: - inputs_to_feed[input_name_to_feed] = inputs[input_name_to_feed] - - calibrator.create_inference_session( - execution_providers=calibrator.providers, - trt_extra_plugin_lib_paths=calibrator.trt_extra_plugin_lib_paths, - model_path=model_path, - ) - outputs = calibrator.infer_session.run(None, inputs_to_feed) - - # Mark model path to remove it from dependency map - inferred_model_list.append(model_path) - - # Decrease reference count for used inputs and remove if no reference - for input_name in calibrator.single_node_model_path_map[model_path][0]: - input_ref_count[input_name] -= 1 - if input_ref_count[input_name] == 0: - del inputs[input_name] - del input_ref_count[input_name] - - gc.collect() - - # Add outputs to inputs - need_calibration = False - for output_idx, output in enumerate(outputs): - output_name = calibrator.infer_session.get_outputs()[output_idx].name - inputs[output_name] = output - input_queue.append(output_name) - if output_name in calibrator.tensors_to_calibrate: - need_calibration = True - - # For each inference, compute data before moving to other nodes if tensor is to be calibrated - if need_calibration: - calibrator.intermediate_outputs.append(outputs) - if len(calibrator.intermediate_outputs) == 0: - raise ValueError("No data is collected.") - - output_names = [ - calibrator.infer_session.get_outputs()[i].name - for i in range(len(calibrator.intermediate_outputs[0])) - ] - output_dicts_list = [ - dict(zip(output_names, intermediate_output)) - for intermediate_output in calibrator.intermediate_outputs - ] - - merged_dict = {} - for d in output_dicts_list: - for k, v in d.items(): - merged_dict.setdefault(k, []).append(v) - - clean_merged_dict = { - i: merged_dict[i] - for i in merged_dict - if i in calibrator.tensors_to_calibrate - } - - if not calibrator.collector: - calibrator.collector = HistogramCollector( - method=calibrator.method, - symmetric=calibrator.symmetric, - num_bins=calibrator.num_bins, - num_quantized_bins=calibrator.num_quantized_bins, - percentile=calibrator.percentile, - scenario=calibrator.scenario, - ) - calibrator.collector.collect(clean_merged_dict) - - calibrator.clear_collected_data() - - gc.collect() - pbar.update(1) - - # Remove inferred model from dependency map - for model_path in inferred_model_list: - model_to_input_dep_map.pop(model_path) - - gc.collect() - pbar.close() - input_counter += 1 - - -def _collect_histogram_collector_single_node_calibration(histogram_collector, name_to_arr): - """Collect tensor data and make histogram. - - Modification: Remove print line to make calibration per node log output cleaner. - """ - # TODO: Currently we have different collect() for entropy and percentile method respectively. - # Need unified collect in the future. - if histogram_collector.method in {"distribution", "entropy"}: - return histogram_collector.collect_value(name_to_arr) - elif histogram_collector.method == "percentile": - if histogram_collector.symmetric: - return histogram_collector.collect_absolute_value(name_to_arr) - else: - return histogram_collector.collect_value(name_to_arr) - else: - raise ValueError("Only 'entropy', 'percentile' or 'distribution' methods are supported") - - -def _collect_value_histogram_collector_single_node_calibration(histogram_collector, name_to_arr): - """Collect histogram on real value.""" - for tensor, data_arr in name_to_arr.items(): - data_arr = np.asarray(data_arr).flatten() - data_arr = _prepare_histogram_data(histogram_collector, tensor, data_arr) - min_value, max_value = (np.min(data_arr), np.max(data_arr)) if data_arr.size > 0 else (0, 0) - - # Replace inf/nan with float32 min/max - min_value = ( - np.finfo(np.float32).tiny if np.isinf(min_value) or np.isnan(min_value) else min_value - ) - max_value = ( - np.finfo(np.float32).max if np.isinf(max_value) or np.isnan(max_value) else max_value - ) - - threshold = max(abs(min_value), abs(max_value)) - - if tensor in histogram_collector.histogram_dict: - histogram_collector.histogram_dict[tensor] = histogram_collector.merge_histogram( - histogram_collector.histogram_dict[tensor], - data_arr, - min_value, - max_value, - threshold, - ) - else: - range_max = float(threshold) - hist, hist_edges = np.histogram( - data_arr, histogram_collector.num_bins, range=(-range_max, range_max) - ) - histogram_collector.histogram_dict[tensor] = ( - hist, - hist_edges, - min_value, - max_value, - threshold, - ) - - -def _augment_graph_min_max_calibrater_single_node_calibration(calibrater): - """Augment outputs to retrieve MinMax pair. - - Adds ReduceMin and ReduceMax nodes to all quantization_candidates op type nodes in - model and ensures their outputs are stored as part of the graph output. - - :return: augmented ONNX model - - Modification: Add an additional Concat after Reshaped output to not rely on error-prone indexing. - Create multiple single node ONNX models to be used to calibrate per node. - """ - tensors, _ = calibrater.select_tensors_to_calibrate(calibrater.model) - reshape_shape_name = str(uuid.uuid4()) - reshape_shape = onnx.numpy_helper.from_array(np.array([1], dtype=np.int64), reshape_shape_name) - calibrater.model.graph.initializer.append(reshape_shape) - - def add_reduce_min_max(tensor_name): - keepdims = 1 - minmax_output = tensor_name + "_MinMax" - - # Create reduce nodes - reduce_nodes = [ - onnx.helper.make_node( - op_name, - [tensor_name], - [tensor_name + "_" + op_name + "_Reshape"], - keepdims=keepdims, - name=tensor_name + "_" + op_name, - ) - for op_name in ["ReduceMin", "ReduceMax"] - ] - - # Create reshape nodes - reshape_nodes = [ - onnx.helper.make_node( - "Reshape", - inputs=[node.output[0], reshape_shape_name], - outputs=[tensor_name + "_" + op_name], - name=node.output[0], - ) - for node, op_name in zip(reduce_nodes, ["ReduceMin", "ReduceMax"]) - ] - - # Create concat node - concat_node = onnx.helper.make_node( - "Concat", - inputs=[tensor_name + "_ReduceMin", tensor_name + "_ReduceMax"], - outputs=[minmax_output], - name=tensor_name + "_ReduceMin_ReduceMax_Concat", - axis=0, - ) - - calibrater.model.graph.node.extend(reduce_nodes + reshape_nodes + [concat_node]) - - # Get tensor type - value_infos = {vi.name: vi for vi in calibrater.model.graph.value_info} - value_infos.update({o.name: o for o in calibrater.model.graph.output}) - value_infos.update({i.name: i for i in calibrater.model.graph.input}) - - if tensor_name not in value_infos: - raise ValueError( - f"Unable to guess tensor type for tensor {tensor_name!r}, " - f"running shape inference before quantization may resolve this issue." - ) - - calibrater.model.graph.output.append( - onnx.helper.make_tensor_value_info( - minmax_output, value_infos[tensor_name].type.tensor_type.elem_type, [2] - ) - ) - - # Make sure all shapes are resolved before adding min max nodes - calibrater.model = SymbolicShapeInference.infer_shapes(calibrater.model) - - for tensor in tensors: - add_reduce_min_max(tensor) - - # Make sure all shapes are resolved after adding min max nodes - calibrater.model = SymbolicShapeInference.infer_shapes(calibrater.model) - - onnx.save( - calibrater.model, - calibrater.augmented_model_path, - save_as_external_data=calibrater.use_external_data_format, - ) - - # Build single node models and save them - model_counter = 0 - initializer_name_map = { - initializer.name: initializer for initializer in calibrater.model.graph.initializer - } - value_info_name_map = { - value_info.name: value_info for value_info in calibrater.model.graph.value_info - } - input_name_map = {input.name: input for input in calibrater.model.graph.input} - output_name_map = {output.name: output for output in calibrater.model.graph.output} - for node in calibrater.model.graph.node: - single_node_model_name = ( - f"{calibrater.augmented_model_path[:-5]}_single_node_{model_counter}" - ) - single_node_model_node = [] - single_node_model_inputs = [] - single_node_model_outputs = [] - single_node_model_initializers = [] - single_node_model_input_names = [] - single_node_model_output_names = [] - - # Add node - single_node_model_node.append(node) - - # Process each input for node - for input_name in node.input: - # Skip empty tensors - if input_name == "": - continue - - is_input_initializer = False - # If a node input is an initializer, add it to initializer list - if input_name in initializer_name_map: - single_node_model_initializers.append(initializer_name_map[input_name]) - is_input_initializer = True - - value_info_found = False - # If a node input is not an initializer, set it as a model input - if not is_input_initializer: - for name_map in [value_info_name_map, input_name_map, output_name_map]: - if input_name in name_map: - single_node_model_inputs.append(name_map[input_name]) - single_node_model_input_names.append(input_name) - value_info_found = True - break - - if not value_info_found: - raise ValueError( - f"{calibrater.augmented_model_path} is not properly shape inferenced." - ) - - # Process each output for node - for output_name in node.output: - value_info_found = False - for name_map in [value_info_name_map, output_name_map]: - if output_name in name_map: - single_node_model_outputs.append(name_map[output_name]) - single_node_model_output_names.append(output_name) - value_info_found = True - break - - if not value_info_found: - raise ValueError( - f"{calibrater.augmented_model_path} is not properly shape inferenced." - ) - - # Create a new onnx model - single_node_model = onnx.helper.make_model( - onnx.helper.make_graph( - single_node_model_node, - single_node_model_name, - single_node_model_inputs, - single_node_model_outputs, - single_node_model_initializers, - ), - opset_imports=calibrater.model.opset_import, - functions=calibrater.model.functions, - ir_version=calibrater.model.ir_version, - ) - - # Save it to a new onnx file - onnx.save(single_node_model, f"{single_node_model_name}.onnx") - - # Save model info and increase model counter - calibrater.single_node_model_path_map[f"{single_node_model_name}.onnx"] = ( - single_node_model_input_names, - single_node_model_output_names, - ) - model_counter += 1 - - -def _augment_graph_histogram_calibrater_single_node_calibration(calibrater): - """Make all quantization_candidates op type nodes as part of the graph output. - - :return: augmented ONNX model - """ - calibrater.tensors_to_calibrate, value_infos = calibrater.select_tensors_to_calibrate( - calibrater.model - ) - for tensor in calibrater.tensors_to_calibrate: - if tensor not in calibrater.model_original_outputs: - calibrater.model.graph.output.append(value_infos[tensor]) - - onnx.save( - calibrater.model, - calibrater.augmented_model_path, - save_as_external_data=calibrater.use_external_data_format, - ) - - # Build single node models and save them - initializer_name_map = { - initializer.name: initializer for initializer in calibrater.model.graph.initializer - } - value_info_name_map = { - value_info.name: value_info for value_info in calibrater.model.graph.value_info - } - input_name_map = {input.name: input for input in calibrater.model.graph.input} - output_name_map = {output.name: output for output in calibrater.model.graph.output} - model_counter = 0 - for node in calibrater.model.graph.node: - single_node_model_name = ( - f"{calibrater.augmented_model_path[:-5]}_single_node_{model_counter}" - ) - single_node_model_nodes = [] - single_node_model_inputs = [] - single_node_model_outputs = [] - single_node_model_initializers = [] - single_node_model_input_names = [] - single_node_model_output_names = [] - - # Add node - single_node_model_nodes.append(node) - - # Process each input for node - for input_name in node.input: - # Skip empty tensors - if input_name == "": - continue - - is_input_initializer = False - # If a node input is an initializer, add it to initializer list - if input_name in initializer_name_map: - single_node_model_initializers.append(initializer_name_map[input_name]) - is_input_initializer = True - - # If a node input is not an initializer, set it as a model input - if not is_input_initializer: - value_info_found = False - for name_map in [value_info_name_map, input_name_map, output_name_map]: - if input_name in name_map: - single_node_model_inputs.append(name_map[input_name]) - single_node_model_input_names.append(input_name) - value_info_found = True - break - - if not value_info_found: - raise ValueError( - f"{calibrater.augmented_model_path} is not properly shape inferenced." - ) - - # Process each output for node - for output_name in node.output: - value_info_found = False - for name_map in [value_info_name_map, output_name_map]: - if output_name in name_map: - single_node_model_outputs.append(name_map[output_name]) - single_node_model_output_names.append(output_name) - value_info_found = True - break - - if not value_info_found: - raise ValueError( - f"{calibrater.augmented_model_path} is not properly shape inferenced." - ) - - # Create a new onnx model - single_node_model = onnx.helper.make_model( - onnx.helper.make_graph( - single_node_model_nodes, - single_node_model_name, - single_node_model_inputs, - single_node_model_outputs, - single_node_model_initializers, - ), - opset_imports=calibrater.model.opset_import, - functions=calibrater.model.functions, - ir_version=calibrater.model.ir_version, - ) - - # Save it to a new onnx file - onnx.save(single_node_model, f"{single_node_model_name}.onnx") - - # Save model info and increase model counter - calibrater.single_node_model_path_map[f"{single_node_model_name}.onnx"] = ( - single_node_model_input_names, - single_node_model_output_names, - ) - model_counter += 1 - - -def _adjust_tensor_ranges(base_quantizer): - if base_quantizer.tensors_range is None: - return - - for node in base_quantizer.model.nodes(): - # adjust tensor_ranges for input of Clip and Relu node - if node.op_type in ["Clip", "Relu"]: - if base_quantizer.is_activation_symmetric: - continue - if not base_quantizer.should_quantize_node(node): - continue - if len(base_quantizer.model.input_name_to_nodes()[node.input[0]]) != 1: - continue - if ( - node.input[0] not in base_quantizer.tensors_range - or node.output[0] not in base_quantizer.tensors_range - ): - continue - td = base_quantizer.tensors_range[node.output[0]] - if not isinstance(td, TensorData): - raise TypeError(f"Unexpected type {type(td)} for {node.output[0]!r}.") - base_quantizer.tensors_range[node.input[0]] = td - # Adjust Softmax to range from 0.0 to 1.0 - elif node.op_type == "Softmax": - if node.output[0] not in base_quantizer.tensors_range: - continue - base_quantizer.tensors_range[node.output[0]] = TensorData( - lowest=np.float32(0.0), - highest=np.float32(1.0), - avg=np.float32(0.0), - std=np.float32(1.0), - ) - - # Patching nan values in TensorData - # These nan values should not appear in calibration with real inputs - for tensor_name in base_quantizer.tensors_range: - td = base_quantizer.tensors_range[tensor_name] - if np.isnan(td.range_value).any(): - base_quantizer.tensors_range[tensor_name] = TensorData( - lowest=np.float32(0.0), - highest=np.float32(448.0), - ) - - -def _create_calibrator_with_extra_options( - model: str | Path, - op_types_to_calibrate: Sequence[str] | None = None, - augmented_model_path="augmented_model.onnx", - calibrate_method=CalibrationMethod.MinMax, - use_external_data_format=False, - extra_options={}, -): - """This function overwrite is needed to pass the TRT plugin path and EP list to the inference session.""" - calibrator = None - if calibrate_method == CalibrationMethod.MinMax: - # default settings for min-max algorithm - symmetric = extra_options.get("symmetric", False) - moving_average = extra_options.get("moving_average", False) - averaging_constant = extra_options.get("averaging_constant", 0.01) - max_intermediate_outputs = extra_options.get("max_intermediate_outputs", None) - calibrator = MinMaxCalibrater( - model, - op_types_to_calibrate, - augmented_model_path, - use_external_data_format=use_external_data_format, - symmetric=symmetric, - moving_average=moving_average, - averaging_constant=averaging_constant, - max_intermediate_outputs=max_intermediate_outputs, - ) - elif calibrate_method == CalibrationMethod.Entropy: - # default settings for entropy algorithm - num_bins = extra_options.get("num_bins", 128) - num_quantized_bins = extra_options.get("num_quantized_bins", 128) - symmetric = extra_options.get("symmetric", False) - calibrator = EntropyCalibrater( - model, - op_types_to_calibrate, - augmented_model_path, - use_external_data_format=use_external_data_format, - symmetric=symmetric, - num_bins=num_bins, - num_quantized_bins=num_quantized_bins, - ) - elif calibrate_method == CalibrationMethod.Percentile: - # default settings for percentile algorithm - num_bins = extra_options.get("num_bins", 2048) - percentile = extra_options.get("percentile", 99.999) - symmetric = extra_options.get("symmetric", True) - calibrator = PercentileCalibrater( - model, - op_types_to_calibrate, - augmented_model_path, - use_external_data_format=use_external_data_format, - symmetric=symmetric, - num_bins=num_bins, - percentile=percentile, - ) - - elif calibrate_method == CalibrationMethod.Distribution: - # default settings for percentile algorithm - num_bins = extra_options.get("num_bins", 2048) - scenario = extra_options.get("scenario", "same") - - calibrator = DistributionCalibrater( - model, - op_types_to_calibrate, - augmented_model_path, - use_external_data_format=use_external_data_format, - num_bins=num_bins, - scenario=scenario, - ) - - if calibrator: - calibrator.augment_graph() - # ======== Modification: additional parameter with TRT plugin path ======== - calibrator.create_inference_session(**extra_options) - # ========================================================================= - return calibrator - - raise ValueError(f"Unsupported calibration method {calibrate_method}") - - -def _quantize_static( - model_input: str | Path | onnx.ModelProto, - model_output: str | Path, - calibration_data_reader: CalibrationDataReader, - quant_format=QuantFormat.QDQ, - op_types_to_quantize=None, - per_channel=False, - reduce_range=False, - activation_type=QuantType.QInt8, - weight_type=QuantType.QInt8, - nodes_to_quantize=None, - nodes_to_exclude=None, - use_external_data_format=False, - calibrate_method=CalibrationMethod.MinMax, - extra_options=None, -): - """Modification: enables TRT custom ops in the calibrator via 'TrtExtraPluginLibraryPaths' in extra_options. - - See ort.quantization.quantize.quantize_static for full function description. Additional info: - - extra_options: - key value pair dictionary for various options in different case. Current used: - ... - TrtExtraPluginLibraryPaths = string : - Default is None. Set TensorRT plugin paths if required. - ExecutionProviders = list[string] : - Default is [("CUDAExecutionProvider", {"device_id": 0}), "CPUExecutionProvider", - "TensorrtExecutionProvider"] - TrtRtxBackend = string : - Selects the legacy or ABI TensorRT-RTX execution provider implementation. - """ - logger.info("Starting static quantization") - logger.debug(f"Quantization format: {quant_format}") - logger.debug(f"Activation type: {activation_type}") - logger.debug(f"Weight type: {weight_type}") - logger.debug(f"Calibration method: {calibrate_method}") - if ( - QuantType.QFLOAT8E4M3FN in (activation_type, weight_type) - and calibrate_method != CalibrationMethod.Distribution - ): - raise ValueError( - "Only Distribution calibration method is supported for float quantization." - ) - - extra_options = extra_options or {} - nodes_to_exclude = nodes_to_exclude or [] - nodes_to_quantize = nodes_to_quantize or [] - op_types_to_quantize = op_types_to_quantize or [] - mode = QuantizationMode.QLinearOps - - if not op_types_to_quantize or len(op_types_to_quantize) == 0: - q_linear_ops = list(QLinearOpsRegistry.keys()) - qdq_ops = list(QDQRegistry.keys()) - op_types_to_quantize = list(set(q_linear_ops + qdq_ops)) - - model = ( - onnx_utils.infer_shapes(model_input) - if isinstance(model_input, onnx.ModelProto) - else load_model_with_shape_infer(Path(model_input)) - ) - - calib_extra_options_keys = [ - ("CalibTensorRangeSymmetric", "symmetric"), - ("CalibMovingAverage", "moving_average"), - ("CalibMovingAverageConstant", "averaging_constant"), - ("CalibMaxIntermediateOutputs", "max_intermediate_outputs"), - # ====================== Modification ====================== - ("TrtExtraPluginLibraryPaths", "trt_extra_plugin_lib_paths"), - ("ExecutionProviders", "execution_providers"), - ("TrtRtxBackend", "trt_rtx_backend"), - ("group_qdq_tensors", "group_qdq_tensors"), - ("QDQDisableWeightAdjustForInt32Bias", "disable_int32_weight_adjustment"), - # ========================================================== - ] - calib_extra_options = { - key: extra_options.get(name) - for (name, key) in calib_extra_options_keys - if name in extra_options - } - logger.debug(f"Calibration extra options: {calib_extra_options}") - - with tempfile.TemporaryDirectory(prefix="ort.quant.") as quant_tmp_dir: - if isinstance(model_input, onnx.ModelProto): - output_path = str(Path(quant_tmp_dir) / "model_input.onnx") - logger.debug(f"Saving model to temporary path: {output_path}") - onnx.save_model( - model_input, - output_path, - save_as_external_data=True, - ) - model_input = output_path - - logger.debug("Creating calibrator") - calibrator = calibrate.create_calibrator( - Path(model_input), - # ======== Modification ======== - nodes_to_quantize, - # ======== Modification ======== - augmented_model_path=Path(quant_tmp_dir).joinpath("augmented_model.onnx").as_posix(), - calibrate_method=calibrate_method, - use_external_data_format=use_external_data_format, - extra_options=calib_extra_options, - ) - - logger.debug("Collecting calibration data") - calibrator.collect_data(calibration_data_reader) - logger.debug("Computing tensor ranges") - tensors_range = calibrator.compute_data() - if not isinstance(tensors_range, TensorsData): - logger.error(f"Unexpected type {type(tensors_range)} for tensors_range") - raise TypeError( - f"Unexpected type {type(tensors_range)} for tensors_range and calibrator={type(calibrator)}." - ) - if isinstance(calibrator, HistogramCalibrater): - _restore_histogram_calibration_dtypes(calibrator.collector, tensors_range) - del calibrator - - check_static_quant_arguments(quant_format, activation_type, weight_type) - - if quant_format is QuantFormat.QOperator: - quantizer = QDQQuantizer( - model, - per_channel, - reduce_range, - mode, - True, # static - weight_type, - activation_type, - tensors_range, - nodes_to_quantize, - nodes_to_exclude, - op_types_to_quantize, - extra_options, - ) - else: - quantizer = QDQQuantizer( - model, - per_channel, - reduce_range, - weight_type, - activation_type, - tensors_range, - nodes_to_quantize, - nodes_to_exclude, - op_types_to_quantize, - extra_options, - ) - - quantizer.quantize_model() - quantizer.model.save_model_to_file(model_output, use_external_data_format) - - -def _init_calibrater_base( - calibrater, - model_path: str | Path, - op_types_to_calibrate: Sequence[str] | None = None, - augmented_model_path="augmented_model.onnx", - symmetric=False, - use_external_data_format=False, - per_channel=False, -): - """Initialize calibrater base class. - - :param model_path: ONNX model to calibrate. It should be a model file path - :param op_types_to_calibrate: operator types to calibrate. By default, calibrate all the float32/float16 tensors. - :param augmented_model_path: save augmented model to this path. - :param symmetric: make range of tensor symmetric (central point is 0). - :param use_external_data_format: use external data format to store model which size is >= 2Gb - - Modification: Additional members including single_node_model_path_map, providers, and trt_extra_plugin_lib_paths - were added and initialized to support calibration per node feature. - """ - if isinstance(model_path, str): - calibrater.model = load_model_with_shape_infer(Path(model_path)) - elif isinstance(model_path, Path): - calibrater.model = load_model_with_shape_infer(model_path) - else: - raise ValueError("model_path should be model path.") - - calibrater.op_types_to_calibrate = op_types_to_calibrate - calibrater.augmented_model_path = augmented_model_path - calibrater.symmetric = symmetric - calibrater.use_external_data_format = use_external_data_format - calibrater.per_channel = per_channel - calibrater.augment_model = None - calibrater.infer_session = None - calibrater.execution_providers = [] - - # Add single node calibration members - calibrater.single_node_model_path_map = {} # {path: ([inputs], [outputs])} - calibrater.providers = [] - calibrater.trt_extra_plugin_lib_paths = None - - -def patch_ort_modules(calibrate_per_node: bool = False): - """Patches the ORT modules.""" - logger.debug("Patching ORT modules") - if calibrate_per_node: - MinMaxCalibrater.augment_graph = _augment_graph_min_max_calibrater_single_node_calibration - MinMaxCalibrater.collect_data = _collect_data_min_max_calibrater_single_node_calibration - MinMaxCalibrater.compute_data = _compute_data_min_max_calibrater_single_node_calibration - MinMaxCalibrater.merge_range = _merge_range_min_max_calibrater_single_node_calibration - HistogramCalibrater.augment_graph = ( - _augment_graph_histogram_calibrater_single_node_calibration - ) - HistogramCalibrater.collect_data = ( - _collect_data_histogram_calibrater_single_node_calibration - ) - HistogramCollector.collect = _collect_histogram_collector_single_node_calibration - HistogramCollector.collect_value = ( - _collect_value_histogram_collector_single_node_calibration - ) - else: - HistogramCollector.collect_value = _collect_value - HistogramCollector.collect_absolute_value = _collect_absolute_value - MinMaxCalibrater.compute_data = _compute_data_minmax_calibrator - MinMaxCalibrater.collect_data = _collect_data_minmax_calibrator - MinMaxCalibrater.merge_range = _merge_range_minmax_calibrator - HistogramCalibrater.collect_data = _collect_data_histogram_calibrator - - calibrate.create_calibrator = _create_calibrator_with_extra_options - CalibraterBase.create_inference_session = _create_inference_session_with_ep_config - CalibraterBase.select_tensors_to_calibrate = _select_tensors_to_calibrate - QDQQuantizer.check_opset_version = _check_opset_version - BaseQuantizer.adjust_tensor_ranges = _adjust_tensor_ranges - qdq_quantizer.compute_scale_zp = _compute_scale_zp - CalibraterBase.__init__ = _init_calibrater_base diff --git a/modelopt/onnx/quantization/ort_quantization.py b/modelopt/onnx/quantization/ort_quantization.py new file mode 100644 index 00000000000..65f266a96ce --- /dev/null +++ b/modelopt/onnx/quantization/ort_quantization.py @@ -0,0 +1,399 @@ +# Adapted from https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/quant_utils.py +# and https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/calibrate.py +# and https://github.com/microsoft/onnxruntime/blob/2ac381c55397dffff327cc6efecf6f95a70f90a1/onnxruntime/python/tools/quantization/onnx_quantizer.py +# and https://github.com/microsoft/onnxruntime/blob/2ac381c55397dffff327cc6efecf6f95a70f90a1/onnxruntime/python/tools/quantization/quantize.py +# +# MIT License +# +# Copyright (c) Microsoft Corporation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# 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. + +"""Static Q/DQ quantization built on ONNX Runtime.""" + +__all__ = [] + +import tempfile +from collections.abc import Sequence +from pathlib import Path + +import numpy as np +import onnx +from onnx import onnx_pb +from onnxruntime.quantization import calibrate +from onnxruntime.quantization.calibrate import ( + CalibrationDataReader, + CalibrationMethod, + DistributionCalibrater, + EntropyCalibrater, + HistogramCalibrater, + MinMaxCalibrater, + PercentileCalibrater, + TensorData, + TensorsData, +) +from onnxruntime.quantization.qdq_quantizer import QDQQuantizer +from onnxruntime.quantization.quant_utils import QuantFormat, QuantizationMode, QuantType +from onnxruntime.quantization.quant_utils import compute_scale_zp as _ort_compute_scale_zp +from onnxruntime.quantization.quantize import check_static_quant_arguments +from onnxruntime.quantization.registry import QDQRegistry, QLinearOpsRegistry + +import modelopt.onnx.utils as onnx_utils +from modelopt.onnx.logging_config import logger +from modelopt.onnx.quantization.ort_calibration import _restore_histogram_calibration_dtypes +from modelopt.onnx.quantization.ort_session import load_model_with_shape_infer + + +def _compute_scale_zp(rmin, rmax, qmin, qmax, symmetric=False, min_real_range=None): + """Retry FP16 scale calculation in FP32 when range subtraction overflows.""" + range_dtype = np.asarray(rmax).dtype + if range_dtype != np.float16: + return _ort_compute_scale_zp(rmin, rmax, qmin, qmax, symmetric, min_real_range) + + with np.errstate(over="ignore", invalid="ignore"): + zero_point, scale = _ort_compute_scale_zp(rmin, rmax, qmin, qmax, symmetric, min_real_range) + if np.all(np.isfinite(scale)): + return zero_point, scale + + zero_point, scale = _ort_compute_scale_zp( + np.asarray(rmin, dtype=np.float32), + np.asarray(rmax, dtype=np.float32), + qmin, + qmax, + symmetric, + min_real_range, + ) + return zero_point, np.asarray(scale, dtype=range_dtype) + + +def _check_opset_version(onnx_quantizer): + ai_onnx_domain = [ + opset + for opset in onnx_quantizer.model.model.opset_import + if not opset.domain or opset.domain in ["ai.onnx", "ai.onnx.contrib"] + ] + opset_version = ai_onnx_domain[0].version + + if opset_version == 10: + return 10 + + if opset_version < 10: + onnx_quantizer.model.model.opset_import.remove(ai_onnx_domain[0]) + onnx_quantizer.model.model.opset_import.extend([onnx.helper.make_opsetid("", 11)]) + opset_version = 11 + + if opset_version < 19 and onnx_quantizer.weight_qType == onnx_pb.TensorProto.FLOAT8E4M3FN: + onnx_quantizer.model.model.opset_import.remove(ai_onnx_domain[0]) + onnx_quantizer.model.model.opset_import.extend([onnx.helper.make_opsetid("", 19)]) + # Set ir_version to 10, remove it once ORT supports ir_version 11 + onnx_quantizer.model.model.ir_version = 10 + opset_version = 19 + + onnx_quantizer.fuse_dynamic_quant = True + return opset_version + + +def _adjust_tensor_ranges(base_quantizer): + if base_quantizer.tensors_range is None: + return + + for node in base_quantizer.model.nodes(): + # adjust tensor_ranges for input of Clip and Relu node + if node.op_type in ["Clip", "Relu"]: + if base_quantizer.is_activation_symmetric: + continue + if not base_quantizer.should_quantize_node(node): + continue + if len(base_quantizer.model.input_name_to_nodes()[node.input[0]]) != 1: + continue + if ( + node.input[0] not in base_quantizer.tensors_range + or node.output[0] not in base_quantizer.tensors_range + ): + continue + td = base_quantizer.tensors_range[node.output[0]] + if not isinstance(td, TensorData): + raise TypeError(f"Unexpected type {type(td)} for {node.output[0]!r}.") + base_quantizer.tensors_range[node.input[0]] = td + # Adjust Softmax to range from 0.0 to 1.0 + elif node.op_type == "Softmax": + if node.output[0] not in base_quantizer.tensors_range: + continue + base_quantizer.tensors_range[node.output[0]] = TensorData( + lowest=np.float32(0.0), + highest=np.float32(1.0), + avg=np.float32(0.0), + std=np.float32(1.0), + ) + + # Patching nan values in TensorData + # These nan values should not appear in calibration with real inputs + for tensor_name in base_quantizer.tensors_range: + td = base_quantizer.tensors_range[tensor_name] + if np.isnan(td.range_value).any(): + base_quantizer.tensors_range[tensor_name] = TensorData( + lowest=np.float32(0.0), + highest=np.float32(448.0), + ) + + +def _create_calibrator_with_extra_options( + model: str | Path, + op_types_to_calibrate: Sequence[str] | None = None, + augmented_model_path="augmented_model.onnx", + calibrate_method=CalibrationMethod.MinMax, + use_external_data_format=False, + extra_options={}, +): + """This function overwrite is needed to pass the TRT plugin path and EP list to the inference session.""" + calibrator = None + if calibrate_method == CalibrationMethod.MinMax: + # default settings for min-max algorithm + symmetric = extra_options.get("symmetric", False) + moving_average = extra_options.get("moving_average", False) + averaging_constant = extra_options.get("averaging_constant", 0.01) + max_intermediate_outputs = extra_options.get("max_intermediate_outputs", None) + calibrator = MinMaxCalibrater( + model, + op_types_to_calibrate, + augmented_model_path, + use_external_data_format=use_external_data_format, + symmetric=symmetric, + moving_average=moving_average, + averaging_constant=averaging_constant, + max_intermediate_outputs=max_intermediate_outputs, + ) + elif calibrate_method == CalibrationMethod.Entropy: + # default settings for entropy algorithm + num_bins = extra_options.get("num_bins", 128) + num_quantized_bins = extra_options.get("num_quantized_bins", 128) + symmetric = extra_options.get("symmetric", False) + calibrator = EntropyCalibrater( + model, + op_types_to_calibrate, + augmented_model_path, + use_external_data_format=use_external_data_format, + symmetric=symmetric, + num_bins=num_bins, + num_quantized_bins=num_quantized_bins, + ) + elif calibrate_method == CalibrationMethod.Percentile: + # default settings for percentile algorithm + num_bins = extra_options.get("num_bins", 2048) + percentile = extra_options.get("percentile", 99.999) + symmetric = extra_options.get("symmetric", True) + calibrator = PercentileCalibrater( + model, + op_types_to_calibrate, + augmented_model_path, + use_external_data_format=use_external_data_format, + symmetric=symmetric, + num_bins=num_bins, + percentile=percentile, + ) + + elif calibrate_method == CalibrationMethod.Distribution: + # default settings for percentile algorithm + num_bins = extra_options.get("num_bins", 2048) + scenario = extra_options.get("scenario", "same") + + calibrator = DistributionCalibrater( + model, + op_types_to_calibrate, + augmented_model_path, + use_external_data_format=use_external_data_format, + num_bins=num_bins, + scenario=scenario, + ) + + if calibrator: + calibrator.augment_graph() + # ======== Modification: additional parameter with TRT plugin path ======== + calibrator.create_inference_session(**extra_options) + # ========================================================================= + return calibrator + + raise ValueError(f"Unsupported calibration method {calibrate_method}") + + +def _quantize_static( + model_input: str | Path | onnx.ModelProto, + model_output: str | Path, + calibration_data_reader: CalibrationDataReader, + quant_format=QuantFormat.QDQ, + op_types_to_quantize=None, + per_channel=False, + reduce_range=False, + activation_type=QuantType.QInt8, + weight_type=QuantType.QInt8, + nodes_to_quantize=None, + nodes_to_exclude=None, + use_external_data_format=False, + calibrate_method=CalibrationMethod.MinMax, + extra_options=None, +): + """Modification: enables TRT custom ops in the calibrator via 'TrtExtraPluginLibraryPaths' in extra_options. + + See ort.quantization.quantize.quantize_static for full function description. Additional info: + + extra_options: + key value pair dictionary for various options in different case. Current used: + ... + TrtExtraPluginLibraryPaths = string : + Default is None. Set TensorRT plugin paths if required. + ExecutionProviders = list[string] : + Default is [("CUDAExecutionProvider", {"device_id": 0}), "CPUExecutionProvider", + "TensorrtExecutionProvider"] + TrtRtxBackend = string : + Selects the legacy or ABI TensorRT-RTX execution provider implementation. + """ + logger.info("Starting static quantization") + logger.debug(f"Quantization format: {quant_format}") + logger.debug(f"Activation type: {activation_type}") + logger.debug(f"Weight type: {weight_type}") + logger.debug(f"Calibration method: {calibrate_method}") + if ( + QuantType.QFLOAT8E4M3FN in (activation_type, weight_type) + and calibrate_method != CalibrationMethod.Distribution + ): + raise ValueError( + "Only Distribution calibration method is supported for float quantization." + ) + + extra_options = extra_options or {} + nodes_to_exclude = nodes_to_exclude or [] + nodes_to_quantize = nodes_to_quantize or [] + op_types_to_quantize = op_types_to_quantize or [] + mode = QuantizationMode.QLinearOps + + if not op_types_to_quantize or len(op_types_to_quantize) == 0: + q_linear_ops = list(QLinearOpsRegistry.keys()) + qdq_ops = list(QDQRegistry.keys()) + op_types_to_quantize = list(set(q_linear_ops + qdq_ops)) + + model = ( + onnx_utils.infer_shapes(model_input) + if isinstance(model_input, onnx.ModelProto) + else load_model_with_shape_infer(Path(model_input)) + ) + + calib_extra_options_keys = [ + ("CalibTensorRangeSymmetric", "symmetric"), + ("CalibMovingAverage", "moving_average"), + ("CalibMovingAverageConstant", "averaging_constant"), + ("CalibMaxIntermediateOutputs", "max_intermediate_outputs"), + # ====================== Modification ====================== + ("TrtExtraPluginLibraryPaths", "trt_extra_plugin_lib_paths"), + ("ExecutionProviders", "execution_providers"), + ("TrtRtxBackend", "trt_rtx_backend"), + ("group_qdq_tensors", "group_qdq_tensors"), + ("QDQDisableWeightAdjustForInt32Bias", "disable_int32_weight_adjustment"), + # ========================================================== + ] + calib_extra_options = { + key: extra_options.get(name) + for (name, key) in calib_extra_options_keys + if name in extra_options + } + logger.debug(f"Calibration extra options: {calib_extra_options}") + + with tempfile.TemporaryDirectory(prefix="ort.quant.") as quant_tmp_dir: + if isinstance(model_input, onnx.ModelProto): + output_path = str(Path(quant_tmp_dir) / "model_input.onnx") + logger.debug(f"Saving model to temporary path: {output_path}") + onnx.save_model( + model_input, + output_path, + save_as_external_data=True, + ) + model_input = output_path + + logger.debug("Creating calibrator") + calibrator = calibrate.create_calibrator( + Path(model_input), + # ======== Modification ======== + nodes_to_quantize, + # ======== Modification ======== + augmented_model_path=Path(quant_tmp_dir).joinpath("augmented_model.onnx").as_posix(), + calibrate_method=calibrate_method, + use_external_data_format=use_external_data_format, + extra_options=calib_extra_options, + ) + + logger.debug("Collecting calibration data") + calibrator.collect_data(calibration_data_reader) + logger.debug("Computing tensor ranges") + tensors_range = calibrator.compute_data() + if not isinstance(tensors_range, TensorsData): + logger.error(f"Unexpected type {type(tensors_range)} for tensors_range") + raise TypeError( + f"Unexpected type {type(tensors_range)} for tensors_range and calibrator={type(calibrator)}." + ) + if isinstance(calibrator, HistogramCalibrater): + _restore_histogram_calibration_dtypes(calibrator.collector, tensors_range) + del calibrator + + check_static_quant_arguments(quant_format, activation_type, weight_type) + + if quant_format is QuantFormat.QOperator: + quantizer = QDQQuantizer( + model, + per_channel, + reduce_range, + mode, + True, # static + weight_type, + activation_type, + tensors_range, + nodes_to_quantize, + nodes_to_exclude, + op_types_to_quantize, + extra_options, + ) + else: + quantizer = QDQQuantizer( + model, + per_channel, + reduce_range, + weight_type, + activation_type, + tensors_range, + nodes_to_quantize, + nodes_to_exclude, + op_types_to_quantize, + extra_options, + ) + + quantizer.quantize_model() + quantizer.model.save_model_to_file(model_output, use_external_data_format) diff --git a/modelopt/onnx/quantization/ort_session.py b/modelopt/onnx/quantization/ort_session.py new file mode 100644 index 00000000000..b6c0d8841ae --- /dev/null +++ b/modelopt/onnx/quantization/ort_session.py @@ -0,0 +1,165 @@ +# Adapted from https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/quant_utils.py +# and https://github.com/microsoft/onnxruntime/blob/baeece44ba075009c6bfe95891a8c1b3d4571cb3/onnxruntime/python/tools/quantization/calibrate.py +# +# MIT License +# +# Copyright (c) Microsoft Corporation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 AND MIT +# +# 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. + +"""ONNX Runtime model loading and inference-session setup.""" + +__all__ = [] + +from pathlib import Path + +import onnx +import onnxruntime as ort +from onnxruntime.quantization.quant_utils import add_infer_metadata + +import modelopt.onnx.utils as onnx_utils +from modelopt.onnx.logging_config import logger + + +def load_model_with_shape_infer(model_path: Path) -> onnx.ModelProto: + """Load model while performing symbolic shape infer and ONNX shape inference.""" + model = onnx.load(str(model_path), load_external_data=True) + try: + model = onnx_utils.infer_shapes(model) + add_infer_metadata(model) + except Exception as e: + logger.info(f"Failed to infer shapes for model {model_path}: {e}") + return model + + +def _configure_session_providers( + sess_options: ort.SessionOptions, + providers: list[str | tuple[str, dict]], + trt_rtx_backend: str, +) -> dict[str, list[str | tuple[str, dict]]]: + """Configure providers using the mechanism required by the selected EP. + + ``providers`` contains provider names or ``(name, options)`` pairs in priority order. + ABI EPs are exposed as devices and must be added to ``sess_options``; passing them through + ``InferenceSession(providers=...)`` overrides that configuration. This helper preserves the + ABI device path while returning normal provider arguments for other EPs. + """ + if trt_rtx_backend != "abi": + return {"providers": providers} + + available_providers = set(ort.get_available_providers()) + ep_devices = ort.get_ep_devices() + plugin_provider_names = {device.ep_name for device in ep_devices} - available_providers + provider_names = { + provider[0] if isinstance(provider, tuple) else provider for provider in providers + } + if not plugin_provider_names.intersection(provider_names): + return {"providers": providers} + + for provider in providers: + provider_name, provider_options = ( + provider if isinstance(provider, tuple) else (provider, {}) + ) + if provider_name in plugin_provider_names: + selected_devices = [device for device in ep_devices if device.ep_name == provider_name] + sess_options.add_provider_for_devices(selected_devices, provider_options) + else: + sess_options.add_provider(provider_name, provider_options) + return {} + + +def _create_inference_session_with_ep_config(calibrator, **kwargs): + """Create an ORT InferenceSession.""" + model_path = kwargs.get("model_path") + logger.debug("Creating inference session with Execution Provider configuration") + + trt_rtx_backend = kwargs.get("trt_rtx_backend", "legacy") + if trt_rtx_backend not in ("legacy", "abi"): + raise ValueError(f"trt_rtx_backend must be 'legacy' or 'abi', got {trt_rtx_backend!r}") + + sess_options = ort.SessionOptions() + sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + sess_options.add_session_config_entry("session.use_device_allocator_for_initializers", "1") + sess_options.enable_cpu_mem_arena = False + + providers = kwargs.get("execution_providers", []) + logger.debug(f"Execution providers: {providers}") + + # Note. This path can be an empty string, which denotes that the model has custom ops and TRT EP is needed. + calibrator.trt_extra_plugin_lib_paths = kwargs.get("trt_extra_plugin_lib_paths") + if calibrator.trt_extra_plugin_lib_paths is not None: + logger.debug(f"TRT extra plugin paths: {calibrator.trt_extra_plugin_lib_paths}") + if "TensorrtExecutionProvider" not in ort.get_available_providers(): + raise RuntimeError( + f"Could not find `TensorrtExecutionProvider`, only {ort.get_available_providers()}" + ) + trt_ep_options = ( + {"trt_extra_plugin_lib_paths": calibrator.trt_extra_plugin_lib_paths} + if calibrator.trt_extra_plugin_lib_paths + else {} + ) + + # Set GPU memory usage limit + trt_ep_options["trt_max_workspace_size"] = 80 * (1024**3) # 80GB + logger.debug(f"TRT EP options: {trt_ep_options}") + + if "TensorrtExecutionProvider" in providers: + providers.remove("TensorrtExecutionProvider") + providers.insert(0, ("TensorrtExecutionProvider", trt_ep_options)) + + def _update_provider_config(provider, config): + if isinstance(provider, tuple) and len(provider) > 1 and isinstance(provider[1], dict): + provider[1].update(config) + else: + provider = (provider, config) + return provider + + for i in range(len(providers)): + if any(p in providers[i] for p in ["CPUExecutionProvider", "CUDAExecutionProvider"]): + providers[i] = _update_provider_config( + providers[i], {"arena_extend_strategy": "kSameAsRequested"} + ) + + session_path = calibrator.augmented_model_path if model_path is None else model_path + provider_kwargs = _configure_session_providers(sess_options, providers, trt_rtx_backend) + calibrator.infer_session = ort.InferenceSession( + session_path, + sess_options=sess_options, + **provider_kwargs, + ) + + # Group qdq tensors will have the same scaling factor. + calibrator.group_qdq_tensors = kwargs.get("group_qdq_tensors") + if calibrator.group_qdq_tensors: + logger.debug(f"Group QDQ tensors: {calibrator.group_qdq_tensors}") diff --git a/modelopt/onnx/quantization/ort_utils.py b/modelopt/onnx/quantization/ort_utils.py index c79e29c78ff..150448c49a8 100755 --- a/modelopt/onnx/quantization/ort_utils.py +++ b/modelopt/onnx/quantization/ort_utils.py @@ -36,7 +36,8 @@ from modelopt.onnx.logging_config import logger from modelopt.onnx.quantization.operators import QDQConvTranspose, QDQCustomOp, QDQNormalization -from modelopt.onnx.quantization.ort_patching import _configure_session_providers, patch_ort_modules +from modelopt.onnx.quantization.ort_patches import patch_ort_modules +from modelopt.onnx.quantization.ort_session import _configure_session_providers def _check_lib_in_ld_library_path(ld_library_path, lib_pattern): diff --git a/tests/gpu/onnx/test_ort_patching.py b/tests/gpu/onnx/test_ort_calibration.py similarity index 97% rename from tests/gpu/onnx/test_ort_patching.py rename to tests/gpu/onnx/test_ort_calibration.py index 84224dcffa0..372e62d3ac0 100644 --- a/tests/gpu/onnx/test_ort_patching.py +++ b/tests/gpu/onnx/test_ort_calibration.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for modelopt.onnx.quantization.ort_patching module.""" +"""Unit tests for the ONNX Runtime quantization capabilities.""" from unittest.mock import Mock, patch @@ -37,26 +37,32 @@ from onnxruntime.quantization.quant_utils import QuantType from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference -from modelopt.onnx.quantization.ort_patching import ( - _adjust_tensor_ranges, - _augment_graph_min_max_calibrater_single_node_calibration, - _check_opset_version, +from modelopt.onnx.quantization.ort_calibration import ( _collect_absolute_value, _collect_data_histogram_calibrator, - _collect_data_min_max_calibrater_single_node_calibration, _collect_data_minmax_calibrator, - _collect_histogram_collector_single_node_calibration, _collect_value, - _collect_value_histogram_collector_single_node_calibration, - _compute_data_min_max_calibrater_single_node_calibration, _compute_data_minmax_calibrator, - _create_calibrator_with_extra_options, - _create_inference_session_with_ep_config, _init_calibrater_base, - _merge_range_min_max_calibrater_single_node_calibration, _merge_range_minmax_calibrator, - _quantize_static, _select_tensors_to_calibrate, +) +from modelopt.onnx.quantization.ort_calibration_per_node import ( + _augment_graph_min_max_calibrater_single_node_calibration, + _collect_data_min_max_calibrater_single_node_calibration, + _collect_histogram_collector_single_node_calibration, + _collect_value_histogram_collector_single_node_calibration, + _compute_data_min_max_calibrater_single_node_calibration, + _merge_range_min_max_calibrater_single_node_calibration, +) +from modelopt.onnx.quantization.ort_quantization import ( + _adjust_tensor_ranges, + _check_opset_version, + _create_calibrator_with_extra_options, + _quantize_static, +) +from modelopt.onnx.quantization.ort_session import ( + _create_inference_session_with_ep_config, load_model_with_shape_infer, ) @@ -123,9 +129,9 @@ def test_load_model_with_shape_infer(self, simple_onnx_model, tmp_path): model_path = tmp_path / "test_model.onnx" onnx.save(simple_onnx_model, str(model_path)) - with patch("modelopt.onnx.quantization.ort_patching.onnx_utils") as mock_ssi: + with patch("modelopt.onnx.quantization.ort_session.onnx_utils") as mock_ssi: mock_ssi.infer_shapes.return_value = simple_onnx_model - with patch("modelopt.onnx.quantization.ort_patching.add_infer_metadata") as mock_aim: + with patch("modelopt.onnx.quantization.ort_session.add_infer_metadata") as mock_aim: result = load_model_with_shape_infer(model_path) mock_ssi.infer_shapes.assert_called_once() @@ -491,9 +497,11 @@ def test_quantize_static_basic(self, simple_onnx_model, tmp_path): with ( patch( - "modelopt.onnx.quantization.ort_patching.calibrate.create_calibrator" + "modelopt.onnx.quantization.ort_quantization.calibrate.create_calibrator" ) as mock_create, - patch("modelopt.onnx.quantization.ort_patching.QDQQuantizer") as mock_quantizer_class, + patch( + "modelopt.onnx.quantization.ort_quantization.QDQQuantizer" + ) as mock_quantizer_class, ): mock_calibrator = Mock() mock_calibrator.collect_data = Mock() @@ -528,7 +536,7 @@ def test_init_calibrater_base(self, simple_onnx_model, tmp_path): mock_calibrater = Mock() with patch( - "modelopt.onnx.quantization.ort_patching.load_model_with_shape_infer" + "modelopt.onnx.quantization.ort_calibration.load_model_with_shape_infer" ) as mock_load: mock_load.return_value = simple_onnx_model diff --git a/tests/unit/onnx/quantization/test_calibrated_quantization.py b/tests/unit/onnx/quantization/test_calibrated_quantization.py index 2ec732ab0cf..e726f667d39 100644 --- a/tests/unit/onnx/quantization/test_calibrated_quantization.py +++ b/tests/unit/onnx/quantization/test_calibrated_quantization.py @@ -316,36 +316,44 @@ def test_legacy_graph_utils_module_is_removed(): assert importlib.util.find_spec("modelopt.onnx.quantization.graph_utils") is None +_FUTURE_LEGACY_IMPORT_XFAIL = pytest.mark.xfail( + strict=True, + raises=AssertionError, + reason="Legacy calibrated implementation imports are expected to be unavailable", +) + + +def _future_legacy_import(module_name, removed_symbol, *, id): + return pytest.param(module_name, removed_symbol, id=id, marks=_FUTURE_LEGACY_IMPORT_XFAIL) + + @pytest.mark.parametrize( ("module_name", "removed_symbol"), [ - pytest.param("modelopt.onnx.quantization.int8", "quantize", id="int8-mode-function"), - pytest.param("modelopt.onnx.quantization.fp8", "quantize", id="fp8-mode-function"), + _future_legacy_import( + "modelopt.onnx.quantization.int8", "quantize", id="int8-mode-function" + ), + _future_legacy_import("modelopt.onnx.quantization.fp8", "quantize", id="fp8-mode-function"), pytest.param("modelopt.onnx.quantization.ort_patching", None, id="ort-patching-module"), - pytest.param("modelopt.onnx.quantization.qdq_utils", None, id="qdq-utils-module"), - pytest.param( + _future_legacy_import("modelopt.onnx.quantization.qdq_utils", None, id="qdq-utils-module"), + _future_legacy_import( "modelopt.onnx.quantization.qdq_utils", "quantize_weights_to_int4", id="int4-exporter-helper", ), - pytest.param( + _future_legacy_import( "modelopt.onnx.quantization.qdq_utils", "quantize_weights_to_mxfp8", id="mxfp8-exporter-helper", ), - pytest.param( + _future_legacy_import( "modelopt.onnx.quantization.qdq_utils", "fp4qdq_to_2dq", id="fp4-to-2dq-exporter-helper", ), ], ) -@pytest.mark.xfail( - strict=True, - raises=AssertionError, - reason="Legacy calibrated implementation imports are expected to be unavailable", -) -def test_future_legacy_calibrated_imports_are_removed(module_name, removed_symbol): +def test_legacy_calibrated_imports_are_removed(module_name, removed_symbol): module_spec = importlib.util.find_spec(module_name) if removed_symbol is None: assert module_spec is None diff --git a/tests/unit/onnx/quantization/test_ort_patching_histogram.py b/tests/unit/onnx/quantization/test_ort_quantization.py similarity index 80% rename from tests/unit/onnx/quantization/test_ort_patching_histogram.py rename to tests/unit/onnx/quantization/test_ort_quantization.py index e986454bb57..fdd75c5c1d7 100644 --- a/tests/unit/onnx/quantization/test_ort_patching_histogram.py +++ b/tests/unit/onnx/quantization/test_ort_quantization.py @@ -28,15 +28,51 @@ TensorsData, ) -from modelopt.onnx.quantization.ort_patching import ( +from modelopt.onnx.quantization import ort_patches +from modelopt.onnx.quantization.ort_calibration import ( _collect_value, - _collect_value_histogram_collector_single_node_calibration, - _compute_scale_zp, _prepare_histogram_data, - _quantize_static, _restore_histogram_calibration_dtypes, - patch_ort_modules, ) +from modelopt.onnx.quantization.ort_calibration_per_node import ( + _collect_value_histogram_collector_single_node_calibration, +) +from modelopt.onnx.quantization.ort_patches import patch_ort_modules +from modelopt.onnx.quantization.ort_quantization import _compute_scale_zp, _quantize_static + + +def test_patch_composition_uses_capability_owners(monkeypatch): + patch_targets = ( + ( + ort_patches.MinMaxCalibrater, + ("augment_graph", "collect_data", "compute_data", "merge_range"), + ), + (ort_patches.HistogramCalibrater, ("augment_graph", "collect_data")), + ( + ort_patches.HistogramCollector, + ("collect", "collect_value", "collect_absolute_value"), + ), + (ort_patches.calibrate, ("create_calibrator",)), + ( + ort_patches.CalibraterBase, + ("__init__", "create_inference_session", "select_tensors_to_calibrate"), + ), + (ort_patches.QDQQuantizer, ("check_opset_version",)), + (ort_patches.BaseQuantizer, ("adjust_tensor_ranges",)), + (ort_patches.qdq_quantizer, ("compute_scale_zp",)), + ) + for owner, attributes in patch_targets: + for attribute in attributes: + monkeypatch.setattr(owner, attribute, getattr(owner, attribute, None), raising=False) + + patch_ort_modules(False) + assert HistogramCollector.collect_value is _collect_value + + patch_ort_modules(True) + assert ( + HistogramCollector.collect_value + is _collect_value_histogram_collector_single_node_calibration + ) def test_compute_scale_zp_fp16_overflow_fallback():