Fix HF export crash when a dynamic-block quantizer has zero amax

CHANGELOG.rst

@@ -66,6 +66,7 @@ Changelog
 - 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.
 - Fail fast on non-finite AutoQuantize output gradients with an actionable error before accumulating sensitivity scores, without changing attention backend settings.
 - Fix ONNX INT8 entropy calibration failing or producing invalid quantization parameters for FP16 activations.
+- Fix HuggingFace checkpoint export failing with ``activation scaling factor 0.0 not positive`` when a dynamic-block quantizer (such as an NVFP4 input quantizer) ends calibration with an amax of zero because the calibration data never activated that layer or expert. Such a quantizer now exports a positive fallback scale and warns instead of crashing, matching what static quantizers already did; if you see the warning, check whether the layer is expected to be inactive and consider a larger calibration size.
 - 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/torch/quantization/nn/modules/tensor_quantizer.py

@@ -1084,10 +1084,38 @@ def _check_per_channel_block_sizes(block_sizes):
             # remove block_sizes
             self._block_sizes = None
 
+    def _sanitize_export_amax(self, amax: torch.Tensor) -> torch.Tensor:
+        """Replace zero/NaN amax entries with ``maxbound`` so exported scales stay positive.
+
+        A zero amax means calibration never activated the layer; downstream exporters divide by
+        the exported amax, so a zero would fail export or produce inf at inference. Kept
+        branch-free because export flows may pass a meta ``amax``; only the warning reads values.
+        """
+        sanitized = torch.nan_to_num(
+            torch.where(amax == 0, torch.full_like(amax, self.maxbound), amax),
+            nan=self.maxbound,
+        )
+
+        if not amax.is_meta:
+            num_invalid = int((torch.isnan(amax) | (amax == 0)).sum())
+            if num_invalid:
+                warnings.warn(
+                    f"{num_invalid}/{amax.numel()} amax entries of this "
+                    f"{type(self).__name__} are zero or NaN at export time, which means "
+                    "calibration never activated the corresponding layer or expert (or saw NaN "
+                    "activations). Substituting maxbound so the exported scaling factor stays "
+                    "positive. Consider increasing the calibration size if the layer is expected "
+                    "to be active.",
+                    stacklevel=3,
+                )
+        return sanitized
+
     def export_amax(self) -> torch.Tensor | None:
         """Export correctly formatted/shaped amax."""
         if self.block_sizes is not None and self.block_sizes.get("type", None) == "dynamic":
-            return self.amax
+            # Dynamic block quantizers keep a per-tensor amax (the NVFP4 second-level scale) that
+            # needs no reshaping, but it still has to be positive for the exporters.
+            return None if self.amax is None else self._sanitize_export_amax(self.amax)
 
         if self.amax is None:
             return None
@@ -1096,8 +1124,7 @@ def export_amax(self) -> torch.Tensor | None:
             amax = self.amax
         else:
             amax = self.amax.reshape(self._amax_shape_for_export)
-        amax[amax == 0] = self.maxbound
-        amax = torch.nan_to_num(amax, nan=self.maxbound)
+        amax = self._sanitize_export_amax(amax)
         clamp_min, clamp_max = torch.finfo(amax.dtype).tiny, torch.finfo(amax.dtype).max
         amax = amax.clamp(min=clamp_min, max=clamp_max)
 

tests/_test_utils/torch/quantization/tensor_quantizer_common.py

@@ -417,6 +417,49 @@ def test_amax_export(self):
         amax = quantizer.export_amax()
         assert amax.shape == (1,)
 
+    # One config per ``export_amax`` branch.
+    UNUSABLE_AMAX_CFGS = {
+        "dynamic_nvfp4": QuantizerAttributeConfig(
+            num_bits=(2, 1), block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)}
+        ),
+        "static_per_tensor": QuantizerAttributeConfig(num_bits=4),
+    }
+
+    @pytest.mark.parametrize("cfg_name", [*UNUSABLE_AMAX_CFGS])
+    @pytest.mark.parametrize("bad_value", [0.0, float("nan")])
+    def test_amax_export_unusable_amax(self, bad_value, cfg_name):
+        """An unusable amax must export as a positive scale without mutating the quantizer.
+
+        Regression test for NVBug 6768300. The NaN case pins the ``nan_to_num`` half, which the
+        zero case alone would not catch.
+        """
+        quantizer = TensorQuantizer(self.UNUSABLE_AMAX_CFGS[cfg_name]).to(self.device)
+        quantizer.amax = torch.full((1,), bad_value).to(self.device)
+
+        amax = quantizer.export_amax()
+
+        assert torch.all(amax > 0), amax
+        assert torch.all(amax == quantizer.maxbound), amax
+        # export must leave the calibrated state alone
+        stored = quantizer.amax
+        if bad_value == 0.0:
+            assert torch.all(stored == 0), stored
+        else:
+            assert torch.all(torch.isnan(stored)), stored
+
+    @pytest.mark.parametrize("cfg_name", [*UNUSABLE_AMAX_CFGS])
+    def test_amax_export_meta_amax(self, cfg_name):
+        """``export_amax()`` must stay usable when amax is on the meta device."""
+        quantizer = TensorQuantizer(self.UNUSABLE_AMAX_CFGS[cfg_name])
+        quantizer.amax = torch.zeros(1, device="meta")
+
+        amax = quantizer.export_amax()
+
+        # Shape differs per branch (the per-tensor path unsqueezes), so pin only that it stays
+        # meta instead of raising.
+        assert amax.is_meta, amax
+        assert amax.numel() == 1, amax.shape
+
     def test_save_restore(self):
         ref_quantizer = TensorQuantizer(QuantizerAttributeConfig(num_bits=4, axis=0))