From 446fe6a8a503285bdf8ae9a2863c82b8cc09c4da Mon Sep 17 00:00:00 2001 From: Chenghao Liu Date: Thu, 10 Sep 2026 02:05:40 +0800 Subject: [PATCH] fix(quantization): preserve zero batches in histogram calibration Signed-off-by: Chenghao Liu --- CHANGELOG.rst | 2 + .../torch/quantization/calib/histogram.py | 16 ++++--- .../torch/quantization/test_calibrator.py | 44 +++++++++++++++++++ 3 files changed, 57 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 335678eac5a..3f74ffc7fa0 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,8 @@ Changelog **Bug Fixes** +- Fix histogram calibration failing on a nonzero batch after initial all-zero batches, while preserving the zero-valued samples in the calibrated distribution. + - 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/torch/quantization/calib/histogram.py b/modelopt/torch/quantization/calib/histogram.py index e27a5471e30..c9190ef58d9 100644 --- a/modelopt/torch/quantization/calib/histogram.py +++ b/modelopt/torch/quantization/calib/histogram.py @@ -114,18 +114,24 @@ def collect(self, x): # Because we collect histogram on absolute value, setting min=0 simplifying the rare case where # minimum value is not exactly 0 and first batch collected has larger min value than later batches x_max = x.max() - if self._calib_bin_edges is None and self._calib_hist is None: - self._calib_hist = torch.histc(x, bins=self._num_bins, min=0, max=x_max) + if self._calib_bin_edges is None or self._calib_bin_edges[-1] == 0: + # A zero-only prefix has no bin width yet. Preserve its counts + # in the first bin when the first nonzero batch defines the range. + zero_count = 0 if self._calib_hist is None else self._calib_hist.sum() + self._calib_hist = torch.histc( + x, bins=self._num_bins, min=0, max=x_max if x_max > 0 else 1 + ) + self._calib_hist[0] += zero_count self._calib_bin_edges = torch.linspace(0, x_max, self._num_bins + 1) else: - if x_max > self._calib_bin_edges[-1]: # type: ignore[index] - width = self._calib_bin_edges[1] - self._calib_bin_edges[0] # type: ignore[index] + if x_max > self._calib_bin_edges[-1]: + width = self._calib_bin_edges[1] - self._calib_bin_edges[0] self._num_bins = int((x_max / width).ceil().item()) self._calib_bin_edges = torch.arange( 0, x_max + width, width, device=x.device ) - hist = torch.histc(x, bins=self._num_bins, min=0, max=self._calib_bin_edges[-1]) # type: ignore[index] + hist = torch.histc(x, bins=self._num_bins, min=0, max=self._calib_bin_edges[-1]) hist[: self._calib_hist.numel()] += self._calib_hist # type: ignore[union-attr] self._calib_hist = hist diff --git a/tests/unit/torch/quantization/test_calibrator.py b/tests/unit/torch/quantization/test_calibrator.py index 9f7d77ce6f7..cc19d867a9a 100644 --- a/tests/unit/torch/quantization/test_calibrator.py +++ b/tests/unit/torch/quantization/test_calibrator.py @@ -23,6 +23,7 @@ from modelopt.torch.quantization import calib from modelopt.torch.quantization import nn as qnn from modelopt.torch.quantization import utils as quant_utils +from modelopt.torch.quantization.config import QuantizerAttributeConfig class TestMaxCalibrator: @@ -89,6 +90,49 @@ def test_track_amax_raises(self): class TestHistogramCalibrator: + @pytest.mark.parametrize("scale", [1e-6, 1.0, 100.0]) + @pytest.mark.parametrize("zero_batches", [1, 3]) + def test_initial_zero_batches(self, scale, zero_batches): + calibrator = calib.HistogramCalibrator(num_bins=256) + batches = [torch.zeros(17)] * zero_batches + batches += [torch.linspace(-scale, scale, 101), torch.zeros(13)] + for batch in batches: + calibrator.collect(batch) + + expected_hist, expected_edges = np.histogram( + torch.cat(batches).abs().numpy(), bins=256, range=(0, scale) + ) + np.testing.assert_array_equal(calibrator._calib_hist.numpy(), expected_hist) + np.testing.assert_allclose(calibrator._calib_bin_edges.numpy(), expected_edges) + assert calibrator.compute_amax("percentile", percentile=99.9) == expected_edges[-2] + + def test_only_zero_batches(self): + calibrator = calib.HistogramCalibrator(num_bins=256) + for count in [17, 23]: + calibrator.collect(torch.zeros(count)) + assert calibrator._calib_hist[0] == 40 + assert calibrator._calib_hist.sum() == 40 + assert calibrator.compute_amax("percentile") == 0 + calibrator.reset() + calibrator.collect(torch.tensor([0.0, 2.0])) + assert calibrator._calib_hist.sum() == 2 + assert calibrator._calib_bin_edges[-1] == 2 + + def test_quantizer_calibration_after_zero_batch(self): + quantizer = qnn.TensorQuantizer( + QuantizerAttributeConfig(calibrator="histogram", axis=None), + if_quant=False, + if_calib=True, + ) + quantizer(torch.zeros(16)) + quantizer(torch.linspace(0, 1, 256)) + quantizer.load_calib_amax(method="percentile") + quantizer.disable_calib() + quantizer.enable_quant() + output = quantizer(torch.linspace(0, 1, 256)) + assert torch.isfinite(output).all() + torch.testing.assert_close(output, torch.linspace(0, 1, 256), atol=0.01, rtol=0) + @pytest.mark.skip(reason="TODO: Fix assertions in test_grow") def test_grow(self, verbose): x_1 = torch.tensor([0, 255, 255, 255, 255, 255])