diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a20db406973..2124663a3ad 100755 --- a/CHANGELOG.rst +++ b/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. diff --git a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py b/modelopt/torch/quantization/nn/modules/tensor_quantizer.py index 18b97ac2774..a2349c6eb3f 100644 --- a/modelopt/torch/quantization/nn/modules/tensor_quantizer.py +++ b/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) diff --git a/tests/_test_utils/torch/quantization/tensor_quantizer_common.py b/tests/_test_utils/torch/quantization/tensor_quantizer_common.py index dd8d790ee3f..77d3acb7df3 100644 --- a/tests/_test_utils/torch/quantization/tensor_quantizer_common.py +++ b/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))