Preserve zero batches in histogram calibration

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.

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
 

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