[6463897] Fix narrow FP16 histogram calibration

CHANGELOG.rst

@@ -25,6 +25,7 @@ Changelog
 
 **Bug Fixes**
 
+- 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.

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,53 @@ 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:
+        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,6 +153,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))
 
+        concat_data_arr = _prepare_histogram_data(histogram_collector, tensor, concat_data_arr)
         data_arr = concat_data_arr
         # ==========================================================
         if data_arr.size > 0:
@@ -130,9 +179,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 +1172,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()
+        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
@@ -1147,9 +1194,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)
@@ -1685,6 +1729,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)
@@ -1795,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

tests/gpu/onnx/test_ort_patching.py

@@ -153,21 +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
 
-    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.
-        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"
-
     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

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

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"]