[NVBug: 5987078] Fix unified HF export of compressed NVFP4 weights (--low_memory_mode)

CHANGELOG.rst

@@ -85,6 +85,7 @@ Changelog
 - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. ``clear_stale_value_info`` now reconciles stale output shapes via symbolic shape inference (keeping every output's shape field populated), and AutoCast runs ONNX shape inference in strict mode and falls back to schema-based standalone type inference when it fails, so unresolved ops no longer leave tensors untyped.
 - Fused MoE expert auto-detection (``register_fused_experts_on_the_fly``) no longer requires an ``act_fn`` attribute. Some fused-expert modules (e.g. ``MiniMaxM3VLExperts``) apply a custom gated activation between the two ``F.linear`` calls instead of exposing ``act_fn``; they were silently skipped, leaving routed experts unquantized (an experts-only recipe matched nothing) and failing HF export with ``NotImplementedError``. ``_QuantFusedExperts`` is activation-agnostic (it only intercepts the two ``F.linear`` calls), so the requirement was unnecessary. This enables NVFP4/FP8 quantization and export for MiniMax-M2 / MiniMax-M3.
 - Fix unified HF export emitting transformers' *in-memory* (post-``conversion_mapping``) tensor names instead of the original model-hub names, breaking the unified-checkpoint contract (observed on MiniMax-M3: exported ``model.language_model.*`` / ``mlp.experts.*.gate_proj`` instead of hub ``language_model.model.*`` / ``block_sparse_moe.experts.*.w{1,2,3}``). transformers' own save-side ``revert_weight_conversion`` is disabled by ModelOpt because it raises ``RuntimeError`` on 0-d scalar scale tensors, so a new quant-aware reverse conversion (``modelopt/torch/export/quant_aware_conversion.py``) derives rename/split rules from the model's conversion mapping via transformers' ``reverse_transform()`` and carries each weight's companion scale tensors (``weight_scale``, ``weight_scale_2``, ``input_scale``, ``weight_scale_inv``, ``bias``) through the renames and un-fusions, so quantized exports round-trip to the hub names. Any mapping op that cannot be reversed quant-aware yet (e.g. still-stacked fused experts) falls back to the previous in-memory names instead of aborting the export.
+- Fix unified HF export of already-compressed NVFP4 weights, i.e. ``mtq.compress`` and ``examples/llm_ptq/hf_ptq.py --low_memory_mode`` (NVBug 5987078). Compressed weights are packed NVFP4 nibbles, so the per-block scale could not be recomputed from them: the export derived the block count from the *packed* last dim (half the logical one) and computed amax over nibble-pair bytes, writing a ``weight_scale`` of half the required size with meaningless values (on DeepSeek-R1-Distill-Llama-70B, ``[8192, 256]`` instead of ``[8192, 512]``). The export now reuses the per-block scale captured at compression time, rescaled into the exported ``weight_scale_2`` convention so the ``weight_scale * weight_scale_2`` dequantization product is preserved. The internal ``_scale`` / ``_double_scale`` quantizer buffers are also removed after use; they previously leaked into the checkpoint as ``*.weight_quantizer._double_scale`` entries, which made downstream loaders (vLLM, TensorRT-LLM PyTorch backend) fail with ``KeyError`` before any weight was loaded.
 - Fix ONNX FP16/BF16 conversion (``--high_precision_dtype fp16``) producing inconsistent tensor types on models with control-flow subgraphs (e.g. a ``Gemm`` inside an ``If`` branch reading an outer-scope activation alongside converted weights, or ``Resize`` ``scales`` that must stay FP32). Subgraph nodes now only run in low precision when all their float inputs are subgraph initializers; outer-scope captures and low-to-high-precision boundaries inside subgraphs are reconciled with ``Cast`` nodes, and ``Constant`` folding refreshes the constant's ``value_info`` so strongly-typed parsers (TensorRT) no longer reject the model. This is a behavioral change: previously a low-precision control-flow parent converted *every* float subgraph initializer, so a weight inside a branch that also reads an outer-scope FP32 activation could become FP16; such weights now stay FP32 so each node's inputs keep a single precision.
 - Nested submodel reverse mappings are now scoped against registered model namespaces, preventing text-only mappings from capturing an already nested VLM's ``model.visual.*`` namespace or double-prefixing ``model.language_model.*`` (observed on Qwen3.5).
 

modelopt/torch/export/unified_export_hf.py

@@ -57,6 +57,8 @@
 from modelopt.torch.quantization import set_quantizer_by_cfg_context
 from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer
 from modelopt.torch.quantization.qtensor import MXFP8QTensor, NVFP4QTensor
+from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper
+from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8
 from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, quantizer_attr_names
 from modelopt.torch.utils.dataset_utils import _disable_use_cache
 from modelopt.torch.utils.distributed import is_fsdp2_model
@@ -539,6 +541,30 @@ def llm_dummy_forward():
                 expert_id += 1
 
 
+def _compressed_per_block_scale(
+    weight_quantizer: TensorQuantizer, weight: QTensorWrapper
+) -> torch.Tensor | None:
+    """Per-block scale captured at compression time, in the modelopt E4M3 layout.
+
+    ``NVFP4QTensor.quantize(..., try_tensorrt=True)`` returns a cutlass-swizzled 1-D uint8 scale
+    when TensorRT-LLM is available on an FP4-capable device, so normalize it the way
+    ``NVFP4QTensor.dequantize`` does before it is used as an exported ``weight_scale``.
+    """
+    scale = getattr(weight_quantizer, "_scale", None)
+    if scale is None or not (scale.dtype == torch.uint8 and scale.ndim == 1):
+        return scale
+    try:
+        from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import (
+            cutlass_fp4_scale_to_modelopt_fp4_scale,
+        )
+    except ImportError as e:
+        raise ImportError(
+            "This weight was compressed by TensorRT-LLM, so its NVFP4 block scale is "
+            "cutlass-swizzled, but tensorrt_llm cannot be imported to convert it for export."
+        ) from e
+    return cutlass_fp4_scale_to_modelopt_fp4_scale(scale, weight.metadata["shape"][-2:])
+
+
 def _export_quantized_weight(
     sub_module: nn.Module,
     dtype: torch.dtype,
@@ -586,6 +612,27 @@ def _export_quantized_weight(
         sub_module, quantizer_attrs.output_quantizer, None
     )
 
+    # Already real-quantized weights (``mtq.compress`` / ``hf_ptq --low_memory_mode``) hold packed
+    # nibbles -- half the logical last dim -- so per-block scales cannot be recomputed from them.
+    # Use the scale the quantizer captured at compression time instead.
+    uses_compressed_nvfp4_scale = isinstance(weight, QTensorWrapper) and quantization_format in [
+        QUANTIZATION_NVFP4,
+        QUANTIZATION_NVFP4_AWQ,
+        QUANTIZATION_NVFP4_SVDQUANT,
+        QUANTIZATION_W4A16_NVFP4,
+    ]
+    compressed_weight_scale = (
+        _compressed_per_block_scale(weight_quantizer, weight)
+        if uses_compressed_nvfp4_scale
+        else None
+    )
+    compressed_weight_scale_2 = (
+        getattr(weight_quantizer, "_double_scale", None) if uses_compressed_nvfp4_scale else None
+    )
+    use_compressed_scale = (
+        compressed_weight_scale is not None and compressed_weight_scale_2 is not None
+    )
+
     if quantization_format == QUANTIZATION_FP8:
         # Convert amax to float32
         weight_quantizer._amax = weight_quantizer._amax.to(torch.float32)
@@ -635,7 +682,7 @@ def _export_quantized_weight(
             sub_module.register_buffer(quantizer_attrs.weight_scale, e8m0_scale)
             if hasattr(weight_quantizer, "_scale") and weight_quantizer._scale is not None:
                 del weight_quantizer._scale
-        else:
+        elif not use_compressed_scale:
             sub_module.register_buffer(
                 quantizer_attrs.weight_scale, get_weight_scaling_factor(sub_module, weight_name)
             )
@@ -694,7 +741,21 @@ def _export_quantized_weight(
             weight, is_bmm_expert_weight=is_bmm_expert_weight
         )
 
-        if NVFP4QTensor._is_static_quantizer(weight_quantizer):
+        if use_compressed_scale and weight_scale_2 is not None:
+            # Dequant is ``nibble * weight_scale * weight_scale_2``; the stored per-block scale is
+            # normalized against the compression-time global scale, so rescale to keep that product.
+            # The nibbles cannot be re-quantized here (the high-precision weight is gone), so once
+            # ``preprocess_linear_fusion`` unifies ``weight_scale_2`` over a fused group the ratio
+            # below is 1 only for the member owning the group max; the others take one extra E4M3
+            # rounding (<= half-ULP, 6.25%). Avoiding that needs a shared scale at compress time.
+            assert compressed_weight_scale is not None and compressed_weight_scale_2 is not None
+            device = compressed_weight_scale.device
+            weight_scale = _cast_per_block_scale_to_fp8(
+                compressed_weight_scale.float()
+                * compressed_weight_scale_2.float().to(device)
+                / weight_scale_2.float().to(device)
+            )
+        elif NVFP4QTensor._is_static_quantizer(weight_quantizer):
             weight_scale = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(
                 weight_quantizer,
                 weight,

modelopt/torch/quantization/nn/modules/quant_linear.py

@@ -192,7 +192,7 @@ def fold_weight(self, keep_attrs: bool = False):
 class RealQuantLinear(QuantModule):
     """Quantized version of nn.Linear with real quantization."""
 
-    list_of_scale_tensors = ["_scale", "double_scale", "_scale_zeros"]
+    list_of_scale_tensors = ["_scale", "_double_scale", "_scale_zeros"]
     allow_real_quant_gemm = True
 
     @property

tests/gpu/torch/export/test_export_weight_gpu.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
 import math
 
 import torch
@@ -22,6 +23,7 @@
 from torch.nn import init
 
 import modelopt.torch.quantization as mtq
+from modelopt.torch.export.quant_utils import postprocess_state_dict
 from modelopt.torch.export.unified_export_hf import _export_quantized_weight
 from modelopt.torch.quantization.nn.modules.quant_module import QuantModule, QuantModuleRegistry
 from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer
@@ -121,3 +123,45 @@ def test_export_per_block_quantized_weight():
     assert hasattr(model.linears[2], quantizer_attrs.output_quantizer)
     assert not getattr(model.linears[2], quantizer_attrs.output_quantizer).is_enabled
     assert not hasattr(model.linears[2], quantizer_attrs.output_scale)
+
+
+def test_export_compressed_nvfp4_weight():
+    """``mtq.compress`` (used by ``hf_ptq --low_memory_mode``) leaves the weight as packed NVFP4
+    nibbles, so per-block scales cannot be recomputed from it. The export must reuse the scales
+    stored on the quantizer and must not leak those internal buffers into the state_dict.
+    """
+    in_features, block_size = 256, 16
+    calib = lambda x: x(torch.randn(1, 4, in_features).cuda())  # noqa: E731
+
+    model = ToyModel(dims=[in_features, in_features, in_features, in_features]).cuda()
+    reference = mtq.quantize(copy.deepcopy(model), mtq.NVFP4_DEFAULT_CFG, calib)
+    compressed = mtq.quantize(copy.deepcopy(model), mtq.NVFP4_DEFAULT_CFG, calib)
+    mtq.compress(compressed)
+
+    quantizer_attrs = quantizer_attr_names("weight")
+    ref_module, compressed_module = reference.linears[2], compressed.linears[2]
+    _export_quantized_weight(ref_module, torch.float16, "weight")
+    _export_quantized_weight(compressed_module, torch.float16, "weight")
+
+    ref_scale = getattr(ref_module, quantizer_attrs.weight_scale)
+    compressed_scale = getattr(compressed_module, quantizer_attrs.weight_scale)
+
+    # Per-block scale covers the logical input dim, not the packed one.
+    assert compressed_scale.shape == ref_scale.shape
+    assert compressed_scale.shape[-1] == in_features // block_size
+
+    # weight_scale * weight_scale_2 is what dequantization consumes; it must match the
+    # uncompressed export rather than the compression-time normalization.
+    ref_2 = getattr(ref_module, quantizer_attrs.weight_scale_2)
+    compressed_2 = getattr(compressed_module, quantizer_attrs.weight_scale_2)
+    assert torch.allclose(
+        compressed_scale.float() * compressed_2.float(),
+        ref_scale.float() * ref_2.float(),
+        rtol=0.05,
+    )
+
+    # Internal compression buffers must be stripped by postprocess_state_dict, which keys off
+    # RealQuantLinear.list_of_scale_tensors -- a missing underscore there let _double_scale leak.
+    stripped = postprocess_state_dict(compressed.state_dict(), 1.0, None)
+    assert not any(key.endswith("weight_quantizer._double_scale") for key in stripped)
+    assert not any(key.endswith("weight_quantizer._scale") for key in stripped)

tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py

@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+import pytest
+import torch
+from _test_utils.torch.export.utils import ToyModel
+
+import modelopt.torch.quantization as mtq
+from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+from modelopt.torch.quantization.backends.utils import fp4_compatible
+from modelopt.torch.quantization.utils import quantizer_attr_names
+
+BLOCK_SIZE = 16
+
+
+@pytest.mark.skipif(not fp4_compatible(), reason="FP4 is not supported on this GPU")
+def test_export_compressed_nvfp4_weight_trtllm_scale():
+    """Export a compressed NVFP4 weight whose block scale is cutlass-swizzled.
+
+    With TensorRT-LLM importable on an FP4-capable device, ``TensorQuantizer._real_quantize``
+    routes through ``torch.ops.trtllm.fp4_quantize`` and stores a 1-D uint8 swizzled scale in
+    ``weight_quantizer._scale`` instead of the modelopt 2-D E4M3 layout. The export must
+    un-swizzle it; using it as-is would write a scale of raw byte values.
+    """
+    pytest.importorskip("tensorrt_llm")
+
+    in_features = 256
+    calib = lambda x: x(torch.randn(1, 4, in_features).cuda().half())  # noqa: E731
+
+    model = ToyModel(dims=[in_features] * 4).cuda().half()
+    reference = mtq.quantize(copy.deepcopy(model), mtq.NVFP4_DEFAULT_CFG, calib)
+    compressed = mtq.quantize(copy.deepcopy(model), mtq.NVFP4_DEFAULT_CFG, calib)
+    mtq.compress(compressed)
+
+    quantizer_attrs = quantizer_attr_names("weight")
+    ref_module, compressed_module = reference.linears[2], compressed.linears[2]
+
+    # Precondition: this environment really does produce the swizzled layout, otherwise the
+    # test would silently degrade into the dense-scale case already covered in tests/gpu.
+    stored_scale = getattr(compressed_module, quantizer_attrs.weight_quantizer)._scale
+    assert stored_scale.dtype == torch.uint8 and stored_scale.ndim == 1, (
+        f"expected a cutlass-swizzled scale, got {stored_scale.dtype} with ndim {stored_scale.ndim}"
+    )
+
+    _export_quantized_weight(ref_module, torch.float16, "weight")
+    _export_quantized_weight(compressed_module, torch.float16, "weight")
+
+    ref_scale = getattr(ref_module, quantizer_attrs.weight_scale)
+    compressed_scale = getattr(compressed_module, quantizer_attrs.weight_scale)
+
+    # Un-swizzled back to the checkpoint layout, not left as 1-D bytes.
+    assert compressed_scale.shape == ref_scale.shape
+    assert compressed_scale.shape[-1] == in_features // BLOCK_SIZE
+    assert compressed_scale.dtype == ref_scale.dtype
+
+    # weight_scale * weight_scale_2 is what dequantization consumes.
+    ref_2 = getattr(ref_module, quantizer_attrs.weight_scale_2)
+    compressed_2 = getattr(compressed_module, quantizer_attrs.weight_scale_2)
+    assert torch.allclose(
+        compressed_scale.float() * compressed_2.float(),
+        ref_scale.float() * ref_2.float(),
+        rtol=0.05,
+    )