diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 46168b7fe0f..cd741e792a2 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -38,6 +38,7 @@ 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 ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` storing the MoE router at Megatron's ``moe_router_dtype``, which is a routing *compute* dtype, not a storage one. The router now exports at the export ``dtype`` like every other unquantized weight, matching what ``hf_ptq.py`` and the released NVFP4 checkpoints contain; pass ``moe_router_dtype`` to ``export_mcore_gpt_to_hf`` explicitly if you want the old fp32 storage. - Fix unified Megatron export writing a second, unreferenced copy of the vocab embedding when a model with MTP layers is exported with pipeline parallelism. The duplicate was never loaded but inflated the checkpoint by the size of the embedding (about 1 GB for Qwen3.6-35B-A3B); re-export to reclaim the space. - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations. 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])