From 57fff822f428c06abada585ab7d6917cb3fb16c1 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 11 Sep 2026 20:24:49 +0000 Subject: [PATCH 1/5] [6463897] Fix narrow FP16 histogram calibration Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 1 + modelopt/onnx/quantization/ort_patching.py | 12 +++---- tests/gpu/onnx/test_ort_patching.py | 39 ++++++++++++++-------- 3 files changed, 33 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8c6bf86efdc..2e16cbedf05 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -25,6 +25,7 @@ Changelog **Bug Fixes** +- Fix ONNX histogram calibration failing or producing invalid bin edges for FP16 activations with narrow ranges. - Fix ``--use_fsdp2`` HuggingFace checkpoint export gathering the whole model onto rank 0, which made export the dominant phase of a PTQ run and could exhaust host memory on large models. The model is now split into per-decoder-layer units dealt round-robin across ranks; each rank gathers every unit but keeps, packs, and writes only the ones it owns, so a rank buffers roughly ``model / world_size`` instead of the whole checkpoint, and rank 0 writes the combined index. Export configurations that cannot be split this way now raise instead of producing a mismatched checkpoint: FSDP2 combined with another DTensor parallelism (for example FSDP2 + tensor parallel on a 2-D mesh; HSDP is supported), models whose decoder layers cannot be discovered, a decoder layer object reused across layers, and a module that holds the decoder layers while owning parameters of its own. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. - Add FP8 and INT8 recipes that quantize timm ResNet shortcut inputs immediately before residual adds. The torch ONNX example now accepts PTQ and AutoQuantize recipes through ``--recipe`` and uses ``--qformat`` when no recipe is provided. ResNet supports only FP8 and INT8 because TensorRT has limited convolution kernel support; AutoQuantize and other quantization formats are no longer supported for ResNet. diff --git a/modelopt/onnx/quantization/ort_patching.py b/modelopt/onnx/quantization/ort_patching.py index f10b97d13d1..f07f6443445 100755 --- a/modelopt/onnx/quantization/ort_patching.py +++ b/modelopt/onnx/quantization/ort_patching.py @@ -105,6 +105,9 @@ def _collect_value(histogram_collector, name_to_arr): curr_data_arr = curr_data_arr.flatten() concat_data_arr = np.concatenate((concat_data_arr, curr_data_arr)) + # NumPy may otherwise compute histogram edges in FP16 and collapse narrow ranges. + if concat_data_arr.dtype == np.float16: + concat_data_arr = concat_data_arr.astype(np.float32) data_arr = concat_data_arr # ========================================================== if data_arr.size > 0: @@ -130,9 +133,6 @@ def _collect_value(histogram_collector, name_to_arr): old_histogram, data_arr, min_value, max_value, threshold ) else: - # Cast range endpoints to Python float so numpy computes bin edges in - # float64. A fp16 threshold here can underflow the 128-bin linspace - # and trip "Too many bins for data range" on numpy >= 2.0. range_max = float(threshold) hist, hist_edges = np.histogram( data_arr, histogram_collector.num_bins, range=(-range_max, range_max) @@ -1126,6 +1126,9 @@ def _collect_value_histogram_collector_single_node_calibration(histogram_collect """Collect histogram on real value.""" for tensor, data_arr in name_to_arr.items(): data_arr = np.asarray(data_arr).flatten() + # NumPy may otherwise compute histogram edges in FP16 and collapse narrow ranges. + if data_arr.dtype == np.float16: + data_arr = data_arr.astype(np.float32) 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 @@ -1147,9 +1150,6 @@ def _collect_value_histogram_collector_single_node_calibration(histogram_collect threshold, ) else: - # Cast range endpoints to Python float so numpy computes bin edges in - # float64. A fp16 threshold here can underflow the 128-bin linspace - # and trip "Too many bins for data range" on numpy >= 2.0. range_max = float(threshold) hist, hist_edges = np.histogram( data_arr, histogram_collector.num_bins, range=(-range_max, range_max) diff --git a/tests/gpu/onnx/test_ort_patching.py b/tests/gpu/onnx/test_ort_patching.py index 26ca49f39b4..fdb35525cc8 100644 --- a/tests/gpu/onnx/test_ort_patching.py +++ b/tests/gpu/onnx/test_ort_patching.py @@ -153,20 +153,33 @@ def test_collect_value(self, mock_histogram_collector, sample_tensor_data): assert "tensor1" in mock_histogram_collector.histogram_dict assert "tensor2" in mock_histogram_collector.histogram_dict - def test_collect_value_fp16_narrow_range(self, mock_histogram_collector): - # fp16 activations with a small range (threshold ~1e-5) used to raise - # "Too many bins for data range" on numpy >= 2.0, because the fp16 range - # produced a fp16 linspace where consecutive bin edges rounded together. + @pytest.mark.parametrize( + ("collect_value", "batched_input"), + [ + (_collect_value, True), + (_collect_value_histogram_collector_single_node_calibration, False), + ], + ) + def test_collect_value_fp16_narrow_range(self, collect_value, batched_input): + collector = HistogramCollector( + method="entropy", + symmetric=False, + num_bins=128, + num_quantized_bins=128, + percentile=None, + scenario="same", + ) activations = np.zeros(1000, dtype=np.float16) - activations[0] = np.float16(1e-5) - name_to_arr = {"narrow_fp16_tensor": [activations]} - - _collect_value(mock_histogram_collector, name_to_arr) - - hist, edges, _, _, _ = mock_histogram_collector.histogram_dict["narrow_fp16_tensor"] - assert hist.sum() == activations.size - assert len(edges) == mock_histogram_collector.num_bins + 1 - assert not np.any(np.diff(edges) == 0), "fp16 bin edges collapsed" + for activation_max in (1e-6, 1e-6, 2e-6): + activations[0] = np.float16(activation_max) + name_to_arr = {"narrow_fp16_tensor": [activations] if batched_input else activations} + collect_value(collector, name_to_arr) + + hist, edges, _, _, threshold = collector.histogram_dict["narrow_fp16_tensor"] + assert hist.sum() == 3 * activations.size + assert len(edges) == len(hist) + 1 + assert np.all(np.diff(edges) > 0), "fp16 bin edges are not strictly increasing" + assert np.asarray(threshold).dtype.itemsize >= np.dtype(np.float32).itemsize def test_collect_absolute_value(self, mock_histogram_collector, sample_tensor_data): """Test _collect_absolute_value function.""" From a7fbd0af792b5846e53025b87ceade958be4189e Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 11 Sep 2026 20:29:06 +0000 Subject: [PATCH 2/5] Simplify histogram regression parameters Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- tests/gpu/onnx/test_ort_patching.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tests/gpu/onnx/test_ort_patching.py b/tests/gpu/onnx/test_ort_patching.py index fdb35525cc8..1e7d87a2ac3 100644 --- a/tests/gpu/onnx/test_ort_patching.py +++ b/tests/gpu/onnx/test_ort_patching.py @@ -154,13 +154,13 @@ def test_collect_value(self, mock_histogram_collector, sample_tensor_data): assert "tensor2" in mock_histogram_collector.histogram_dict @pytest.mark.parametrize( - ("collect_value", "batched_input"), + "collect_value", [ - (_collect_value, True), - (_collect_value_histogram_collector_single_node_calibration, False), + _collect_value, + _collect_value_histogram_collector_single_node_calibration, ], ) - def test_collect_value_fp16_narrow_range(self, collect_value, batched_input): + def test_collect_value_fp16_narrow_range(self, collect_value): collector = HistogramCollector( method="entropy", symmetric=False, @@ -172,8 +172,7 @@ def test_collect_value_fp16_narrow_range(self, collect_value, batched_input): activations = np.zeros(1000, dtype=np.float16) for activation_max in (1e-6, 1e-6, 2e-6): activations[0] = np.float16(activation_max) - name_to_arr = {"narrow_fp16_tensor": [activations] if batched_input else activations} - collect_value(collector, name_to_arr) + collect_value(collector, {"narrow_fp16_tensor": [activations]}) hist, edges, _, _, threshold = collector.histogram_dict["narrow_fp16_tensor"] assert hist.sum() == 3 * activations.size From 55d58e83d9f48e41f999487d94d0317414fd4530 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 11 Sep 2026 21:11:57 +0000 Subject: [PATCH 3/5] [6463897] Preserve FP16 calibration range types Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 2 +- modelopt/onnx/quantization/ort_patching.py | 35 ++++++++++++++--- tests/gpu/onnx/test_ort_patching.py | 44 ++++++++++++++++++++++ 3 files changed, 74 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2e16cbedf05..7459aed9c6f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -25,7 +25,7 @@ Changelog **Bug Fixes** -- Fix ONNX histogram calibration failing or producing invalid bin edges for FP16 activations with narrow ranges. +- Fix ONNX INT8 entropy calibration failing or producing invalid bin edges for FP16 activations with narrow ranges. - Fix ``--use_fsdp2`` HuggingFace checkpoint export gathering the whole model onto rank 0, which made export the dominant phase of a PTQ run and could exhaust host memory on large models. The model is now split into per-decoder-layer units dealt round-robin across ranks; each rank gathers every unit but keeps, packs, and writes only the ones it owns, so a rank buffers roughly ``model / world_size`` instead of the whole checkpoint, and rank 0 writes the combined index. Export configurations that cannot be split this way now raise instead of producing a mismatched checkpoint: FSDP2 combined with another DTensor parallelism (for example FSDP2 + tensor parallel on a 2-D mesh; HSDP is supported), models whose decoder layers cannot be discovered, a decoder layer object reused across layers, and a module that holds the decoder layers while owning parameters of its own. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. - Add FP8 and INT8 recipes that quantize timm ResNet shortcut inputs immediately before residual adds. The torch ONNX example now accepts PTQ and AutoQuantize recipes through ``--recipe`` and uses ``--qformat`` when no recipe is provided. ResNet supports only FP8 and INT8 because TensorRT has limited convolution kernel support; AutoQuantize and other quantization formats are no longer supported for ResNet. diff --git a/modelopt/onnx/quantization/ort_patching.py b/modelopt/onnx/quantization/ort_patching.py index f07f6443445..f23d5fba09c 100755 --- a/modelopt/onnx/quantization/ort_patching.py +++ b/modelopt/onnx/quantization/ort_patching.py @@ -94,6 +94,31 @@ def load_model_with_shape_infer(model_path: Path) -> onnx.ModelProto: return model +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()): @@ -105,9 +130,7 @@ def _collect_value(histogram_collector, name_to_arr): curr_data_arr = curr_data_arr.flatten() concat_data_arr = np.concatenate((concat_data_arr, curr_data_arr)) - # NumPy may otherwise compute histogram edges in FP16 and collapse narrow ranges. - if concat_data_arr.dtype == np.float16: - concat_data_arr = concat_data_arr.astype(np.float32) + concat_data_arr = _prepare_histogram_data(histogram_collector, tensor, concat_data_arr) data_arr = concat_data_arr # ========================================================== if data_arr.size > 0: @@ -1126,9 +1149,7 @@ def _collect_value_histogram_collector_single_node_calibration(histogram_collect """Collect histogram on real value.""" for tensor, data_arr in name_to_arr.items(): data_arr = np.asarray(data_arr).flatten() - # NumPy may otherwise compute histogram edges in FP16 and collapse narrow ranges. - if data_arr.dtype == np.float16: - data_arr = data_arr.astype(np.float32) + 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 @@ -1685,6 +1706,8 @@ def _quantize_static( 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) diff --git a/tests/gpu/onnx/test_ort_patching.py b/tests/gpu/onnx/test_ort_patching.py index 1e7d87a2ac3..e13153a5aa0 100644 --- a/tests/gpu/onnx/test_ort_patching.py +++ b/tests/gpu/onnx/test_ort_patching.py @@ -55,7 +55,9 @@ _init_calibrater_base, _merge_range_min_max_calibrater_single_node_calibration, _merge_range_minmax_calibrator, + _prepare_histogram_data, _quantize_static, + _restore_histogram_calibration_dtypes, _select_tensors_to_calibrate, load_model_with_shape_infer, ) @@ -177,9 +179,51 @@ def test_collect_value_fp16_narrow_range(self, collect_value): hist, edges, _, _, threshold = collector.histogram_dict["narrow_fp16_tensor"] assert hist.sum() == 3 * activations.size assert len(edges) == len(hist) + 1 + assert edges.dtype == np.float32 assert np.all(np.diff(edges) > 0), "fp16 bin edges are not strictly increasing" assert np.asarray(threshold).dtype.itemsize >= np.dtype(np.float32).itemsize + tensors_range = TensorsData( + CalibrationMethod.Entropy, collector.compute_collection_result() + ) + _restore_histogram_calibration_dtypes(collector, tensors_range) + tensor_range = tensors_range["narrow_fp16_tensor"] + assert tensor_range.lowest.dtype == np.float16 + assert tensor_range.highest.dtype == np.float16 + assert tensor_range.bins.dtype == np.float32 + + def test_restore_histogram_calibration_dtypes_clamps_fp16(self): + collector = HistogramCollector( + method="distribution", + symmetric=False, + num_bins=512, + num_quantized_bins=128, + percentile=None, + scenario="same", + ) + _prepare_histogram_data(collector, "tensor", np.array([], dtype=np.float16)) + + fp32_max = np.finfo(np.float32).max + tensor_data = TensorData( + lowest=np.float32(-fp32_max), + highest=np.float32(fp32_max), + avg=np.float32(fp32_max), + std=np.float32(fp32_max), + hist=np.array([1]), + hist_edges=np.array([-1, 1], dtype=np.float32), + ) + tensors_range = TensorsData(CalibrationMethod.Distribution, {"tensor": tensor_data}) + + _restore_histogram_calibration_dtypes(collector, tensors_range) + + fp16_limits = np.finfo(np.float16) + tensor_range = tensors_range["tensor"] + assert tensor_range.lowest == fp16_limits.min + assert tensor_range.highest == fp16_limits.max + assert tensor_range.avg == fp16_limits.max + assert tensor_range.std == fp16_limits.max + assert tensor_range.hist_edges.dtype == np.float32 + def test_collect_absolute_value(self, mock_histogram_collector, sample_tensor_data): """Test _collect_absolute_value function.""" # Convert to float32 to avoid the float64 assertion error From a9ab37a36b9a71481ec46e2bbd3383174de80161 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 11 Sep 2026 22:03:40 +0000 Subject: [PATCH 4/5] [6463897] Strengthen histogram calibration tests Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- tests/gpu/onnx/test_ort_patching.py | 2 +- tests/unit/onnx/test_autocast_quantize.py | 13 ++++++++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/gpu/onnx/test_ort_patching.py b/tests/gpu/onnx/test_ort_patching.py index e13153a5aa0..82f0eac6d67 100644 --- a/tests/gpu/onnx/test_ort_patching.py +++ b/tests/gpu/onnx/test_ort_patching.py @@ -178,7 +178,7 @@ def test_collect_value_fp16_narrow_range(self, collect_value): hist, edges, _, _, threshold = collector.histogram_dict["narrow_fp16_tensor"] assert hist.sum() == 3 * activations.size - assert len(edges) == len(hist) + 1 + assert len(hist) > collector.num_bins assert edges.dtype == np.float32 assert np.all(np.diff(edges) > 0), "fp16 bin edges are not strictly increasing" assert np.asarray(threshold).dtype.itemsize >= np.dtype(np.float32).itemsize diff --git a/tests/unit/onnx/test_autocast_quantize.py b/tests/unit/onnx/test_autocast_quantize.py index bc123aad9dc..930187980f8 100644 --- a/tests/unit/onnx/test_autocast_quantize.py +++ b/tests/unit/onnx/test_autocast_quantize.py @@ -55,7 +55,18 @@ def test_autocast_quantize_int8(tmp_path, keep_io_types, bias_add): assert os.path.isfile(output_onnx_path) # Load the output model and check QDQ node placements - graph = gs.import_onnx(onnx.load(output_onnx_path)) + quantized_model = onnx.load(output_onnx_path) + graph = gs.import_onnx(quantized_model) + + activation_scale_names = { + node.input[1] for node in quantized_model.graph.node if node.op_type == "QuantizeLinear" + } + activation_scale_types = { + initializer.data_type + for initializer in quantized_model.graph.initializer + if initializer.name in activation_scale_names + } + assert activation_scale_types == {onnx.TensorProto.FLOAT16} # Check that all MatMul nodes are quantized mm_nodes = [n for n in graph.nodes if n.op == "MatMul"] From 244f5cd232915721164402f608507b31d7b38289 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Fri, 11 Sep 2026 22:50:01 +0000 Subject: [PATCH 5/5] [6463897] Handle full-range FP16 calibration scales Retry scale calculation in FP32 only when ORT produces a non-finite FP16 result, then preserve the FP16 initializer dtype. Move the histogram regressions into the CPU unit suite. Co-Authored-By: Codex Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 2 +- modelopt/onnx/quantization/ort_patching.py | 26 ++- tests/gpu/onnx/test_ort_patching.py | 71 ------- .../test_ort_patching_histogram.py | 182 ++++++++++++++++++ 4 files changed, 208 insertions(+), 73 deletions(-) create mode 100644 tests/unit/onnx/quantization/test_ort_patching_histogram.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7459aed9c6f..65c03e42caf 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -25,7 +25,7 @@ Changelog **Bug Fixes** -- Fix ONNX INT8 entropy calibration failing or producing invalid bin edges for FP16 activations with narrow ranges. +- Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations. - Fix ``--use_fsdp2`` HuggingFace checkpoint export gathering the whole model onto rank 0, which made export the dominant phase of a PTQ run and could exhaust host memory on large models. The model is now split into per-decoder-layer units dealt round-robin across ranks; each rank gathers every unit but keeps, packs, and writes only the ones it owns, so a rank buffers roughly ``model / world_size`` instead of the whole checkpoint, and rank 0 writes the combined index. Export configurations that cannot be split this way now raise instead of producing a mismatched checkpoint: FSDP2 combined with another DTensor parallelism (for example FSDP2 + tensor parallel on a 2-D mesh; HSDP is supported), models whose decoder layers cannot be discovered, a decoder layer object reused across layers, and a module that holds the decoder layers while owning parameters of its own. - Speed up ``mtq.quantize`` on FSDP2-sharded fused-MoE models. Promoting static-block weight quantizers gathered each expert's slice of the fused weight across ranks even though only quantizer state is read, adding a collective per expert to calibration. - Add FP8 and INT8 recipes that quantize timm ResNet shortcut inputs immediately before residual adds. The torch ONNX example now accepts PTQ and AutoQuantize recipes through ``--recipe`` and uses ``--qformat`` when no recipe is provided. ResNet supports only FP8 and INT8 because TensorRT has limited convolution kernel support; AutoQuantize and other quantization formats are no longer supported for ResNet. diff --git a/modelopt/onnx/quantization/ort_patching.py b/modelopt/onnx/quantization/ort_patching.py index f23d5fba09c..d2957ec183f 100755 --- a/modelopt/onnx/quantization/ort_patching.py +++ b/modelopt/onnx/quantization/ort_patching.py @@ -52,7 +52,7 @@ import onnxruntime as ort import pynvml from onnx import onnx_pb -from onnxruntime.quantization import calibrate +from onnxruntime.quantization import calibrate, qdq_quantizer from onnxruntime.quantization.base_quantizer import BaseQuantizer from onnxruntime.quantization.calibrate import ( CalibraterBase, @@ -74,6 +74,7 @@ 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 @@ -94,6 +95,28 @@ def load_model_with_shape_infer(model_path: Path) -> onnx.ModelProto: 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: @@ -1818,4 +1841,5 @@ def patch_ort_modules(calibrate_per_node: bool = False): 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/tests/gpu/onnx/test_ort_patching.py b/tests/gpu/onnx/test_ort_patching.py index 82f0eac6d67..84224dcffa0 100644 --- a/tests/gpu/onnx/test_ort_patching.py +++ b/tests/gpu/onnx/test_ort_patching.py @@ -55,9 +55,7 @@ _init_calibrater_base, _merge_range_min_max_calibrater_single_node_calibration, _merge_range_minmax_calibrator, - _prepare_histogram_data, _quantize_static, - _restore_histogram_calibration_dtypes, _select_tensors_to_calibrate, load_model_with_shape_infer, ) @@ -155,75 +153,6 @@ def test_collect_value(self, mock_histogram_collector, sample_tensor_data): assert "tensor1" in mock_histogram_collector.histogram_dict assert "tensor2" in mock_histogram_collector.histogram_dict - @pytest.mark.parametrize( - "collect_value", - [ - _collect_value, - _collect_value_histogram_collector_single_node_calibration, - ], - ) - def test_collect_value_fp16_narrow_range(self, collect_value): - collector = HistogramCollector( - method="entropy", - symmetric=False, - num_bins=128, - num_quantized_bins=128, - percentile=None, - scenario="same", - ) - activations = np.zeros(1000, dtype=np.float16) - for activation_max in (1e-6, 1e-6, 2e-6): - activations[0] = np.float16(activation_max) - collect_value(collector, {"narrow_fp16_tensor": [activations]}) - - hist, edges, _, _, threshold = collector.histogram_dict["narrow_fp16_tensor"] - assert hist.sum() == 3 * activations.size - assert len(hist) > collector.num_bins - assert edges.dtype == np.float32 - assert np.all(np.diff(edges) > 0), "fp16 bin edges are not strictly increasing" - assert np.asarray(threshold).dtype.itemsize >= np.dtype(np.float32).itemsize - - tensors_range = TensorsData( - CalibrationMethod.Entropy, collector.compute_collection_result() - ) - _restore_histogram_calibration_dtypes(collector, tensors_range) - tensor_range = tensors_range["narrow_fp16_tensor"] - assert tensor_range.lowest.dtype == np.float16 - assert tensor_range.highest.dtype == np.float16 - assert tensor_range.bins.dtype == np.float32 - - def test_restore_histogram_calibration_dtypes_clamps_fp16(self): - collector = HistogramCollector( - method="distribution", - symmetric=False, - num_bins=512, - num_quantized_bins=128, - percentile=None, - scenario="same", - ) - _prepare_histogram_data(collector, "tensor", np.array([], dtype=np.float16)) - - fp32_max = np.finfo(np.float32).max - tensor_data = TensorData( - lowest=np.float32(-fp32_max), - highest=np.float32(fp32_max), - avg=np.float32(fp32_max), - std=np.float32(fp32_max), - hist=np.array([1]), - hist_edges=np.array([-1, 1], dtype=np.float32), - ) - tensors_range = TensorsData(CalibrationMethod.Distribution, {"tensor": tensor_data}) - - _restore_histogram_calibration_dtypes(collector, tensors_range) - - fp16_limits = np.finfo(np.float16) - tensor_range = tensors_range["tensor"] - assert tensor_range.lowest == fp16_limits.min - assert tensor_range.highest == fp16_limits.max - assert tensor_range.avg == fp16_limits.max - assert tensor_range.std == fp16_limits.max - assert tensor_range.hist_edges.dtype == np.float32 - def test_collect_absolute_value(self, mock_histogram_collector, sample_tensor_data): """Test _collect_absolute_value function.""" # Convert to float32 to avoid the float64 assertion error diff --git a/tests/unit/onnx/quantization/test_ort_patching_histogram.py b/tests/unit/onnx/quantization/test_ort_patching_histogram.py new file mode 100644 index 00000000000..e986454bb57 --- /dev/null +++ b/tests/unit/onnx/quantization/test_ort_patching_histogram.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for ONNX Runtime histogram quantization patches.""" + +import numpy as np +import onnx +import onnxruntime as ort +import pytest +from onnx import TensorProto, helper, numpy_helper +from onnxruntime.quantization.calibrate import ( + CalibrationDataReader, + CalibrationMethod, + HistogramCollector, + TensorData, + TensorsData, +) + +from modelopt.onnx.quantization.ort_patching 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, +) + + +def test_compute_scale_zp_fp16_overflow_fallback(): + zero_point, scale = _compute_scale_zp( + np.array(-65504, dtype=np.float16), + np.array(65504, dtype=np.float16), + np.array(-128, dtype=np.int8), + np.array(127, dtype=np.int8), + symmetric=True, + ) + + assert zero_point.dtype == np.int8 + assert zero_point == 0 + assert scale.dtype == np.float16 + assert scale == np.float16(514) + + +def test_quantize_static_fp16_high_range_scale(tmp_path): + class HighRangeDataReader(CalibrationDataReader): + def __init__(self): + self.rewind() + + def get_next(self): + return next(self.data, None) + + def rewind(self): + values = np.array([[-65504, 65504, -32752, 32752]], dtype=np.float16) + self.data = iter([{"input": values}]) + + model_path = tmp_path / "model.onnx" + output_path = tmp_path / "model.quant.onnx" + graph = helper.make_graph( + [helper.make_node("MatMul", ["input", "weight"], ["output"], name="matmul")], + "fp16_high_range", + [helper.make_tensor_value_info("input", TensorProto.FLOAT16, [1, 4])], + [helper.make_tensor_value_info("output", TensorProto.FLOAT16, [1, 4])], + [numpy_helper.from_array(np.eye(4, dtype=np.float16), name="weight")], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) + model.ir_version = min(model.ir_version, 10) + onnx.save(model, model_path) + + patch_ort_modules(False) + _quantize_static( + model_path, + output_path, + HighRangeDataReader(), + nodes_to_quantize=["matmul"], + op_types_to_quantize=["MatMul"], + calibrate_method=CalibrationMethod.Entropy, + extra_options={ + "ExecutionProviders": ["CPUExecutionProvider"], + "ActivationSymmetric": True, + "AddQDQPairToWeight": True, + }, + ) + + quantized_model = onnx.load(output_path) + initializers = { + initializer.name: initializer for initializer in quantized_model.graph.initializer + } + scale_initializers = [ + initializers[node.input[1]] + for node in quantized_model.graph.node + if node.op_type in {"QuantizeLinear", "DequantizeLinear"} and node.input[1] in initializers + ] + assert scale_initializers + assert {initializer.data_type for initializer in scale_initializers} == {TensorProto.FLOAT16} + assert all( + np.isfinite(numpy_helper.to_array(initializer)).all() for initializer in scale_initializers + ) + ort.InferenceSession(output_path, providers=["CPUExecutionProvider"]) + + +@pytest.mark.parametrize( + "collect_value", + [ + _collect_value, + _collect_value_histogram_collector_single_node_calibration, + ], +) +def test_collect_value_fp16_narrow_range(collect_value): + collector = HistogramCollector( + method="entropy", + symmetric=False, + num_bins=128, + num_quantized_bins=128, + percentile=None, + scenario="same", + ) + activations = np.zeros(1000, dtype=np.float16) + for activation_max in (1e-6, 1e-6, 2e-6): + activations[0] = np.float16(activation_max) + collect_value(collector, {"narrow_fp16_tensor": [activations]}) + + hist, edges, _, _, threshold = collector.histogram_dict["narrow_fp16_tensor"] + assert hist.sum() == 3 * activations.size + assert len(hist) > collector.num_bins + assert edges.dtype == np.float32 + assert np.all(np.diff(edges) > 0), "fp16 bin edges are not strictly increasing" + assert np.asarray(threshold).dtype.itemsize >= np.dtype(np.float32).itemsize + + tensors_range = TensorsData(CalibrationMethod.Entropy, collector.compute_collection_result()) + _restore_histogram_calibration_dtypes(collector, tensors_range) + tensor_range = tensors_range["narrow_fp16_tensor"] + assert tensor_range.lowest.dtype == np.float16 + assert tensor_range.highest.dtype == np.float16 + assert tensor_range.bins.dtype == np.float32 + + +def test_restore_histogram_calibration_dtypes_clamps_fp16(): + collector = HistogramCollector( + method="distribution", + symmetric=False, + num_bins=512, + num_quantized_bins=128, + percentile=None, + scenario="same", + ) + _prepare_histogram_data(collector, "tensor", np.array([], dtype=np.float16)) + _prepare_histogram_data(collector, "missing_tensor", np.array([], dtype=np.float16)) + + fp32_max = np.finfo(np.float32).max + tensor_data = TensorData( + lowest=np.float32(-fp32_max), + highest=np.float32(fp32_max), + avg=np.float32(fp32_max), + std=np.float32(fp32_max), + hist=np.array([1]), + hist_edges=np.array([-1, 1], dtype=np.float32), + ) + tensors_range = TensorsData(CalibrationMethod.Distribution, {"tensor": tensor_data}) + + _restore_histogram_calibration_dtypes(collector, tensors_range) + + fp16_limits = np.finfo(np.float16) + tensor_range = tensors_range["tensor"] + assert tensor_range.lowest == fp16_limits.min + assert tensor_range.highest == fp16_limits.max + assert tensor_range.avg == fp16_limits.max + assert tensor_range.std == fp16_limits.max + assert tensor_range.hist_edges.dtype == np.float32 + assert "missing_tensor" not in tensors_range