From 2314cb7f027ee28381a35410ea25c31c9c462b99 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Fri, 31 Jul 2026 17:08:39 +0000 Subject: [PATCH] Fix unified HF export of compressed NVFP4 weights (--low_memory_mode) mtq.compress (used by hf_ptq --low_memory_mode) replaces the weight with packed NVFP4 nibbles and stores the per-block scale on the quantizer. The NVFP4 export branch had no handling for this and recomputed the scale from the weight, which is now uint8 packed data: QTensorWrapper reports the packed shape, so the block count came out as in_dim/32 instead of in_dim/16 and the amax was taken over nibble-pair bytes. On DeepSeek-R1-Distill- Llama-70B this wrote weight_scale [8192, 256] instead of [8192, 512] with meaningless values, which dequantizes to garbage. The internal _scale / _double_scale buffers also stayed registered and were serialized as *.weight_quantizer._double_scale entries (560 of them in the 70B checkpoint), so vLLM and the TensorRT-LLM PyTorch backend failed with KeyError before loading any weight. Reuse the compression-time per-block scale, rescaled into the exported weight_scale_2 convention so the weight_scale * weight_scale_2 product that dequantization consumes is preserved, and drop the internal buffers after use. Verified on DGX Spark: the exported checkpoint now matches a normal export (identical key set, weight_scale_2 and input_scale identical, all 154 quantized tensors at the same dequantization error) and serves correctly instead of failing to load. NVBug 5987078 Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 1 + modelopt/torch/export/unified_export_hf.py | 65 +++++++++++++++- .../quantization/nn/modules/quant_linear.py | 2 +- .../torch/export/test_export_weight_gpu.py | 44 +++++++++++ .../export/test_export_compressed_nvfp4.py | 77 +++++++++++++++++++ 5 files changed, 186 insertions(+), 3 deletions(-) create mode 100644 tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 45b7527a1a3..4313f3a17b3 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -57,6 +57,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. 0.45 (2026-07-02) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index cee64c22c05..619056ab3c0 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -56,6 +56,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 @@ -537,6 +539,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, @@ -584,6 +610,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) @@ -633,7 +680,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) ) @@ -692,7 +739,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, diff --git a/modelopt/torch/quantization/nn/modules/quant_linear.py b/modelopt/torch/quantization/nn/modules/quant_linear.py index bb65d59077c..da1b79a2f60 100644 --- a/modelopt/torch/quantization/nn/modules/quant_linear.py +++ b/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 diff --git a/tests/gpu/torch/export/test_export_weight_gpu.py b/tests/gpu/torch/export/test_export_weight_gpu.py index 2167f1e7936..9db2b51114b 100644 --- a/tests/gpu/torch/export/test_export_weight_gpu.py +++ b/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) diff --git a/tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py b/tests/gpu_trtllm/torch/export/test_export_compressed_nvfp4.py new file mode 100644 index 00000000000..08184b2cd18 --- /dev/null +++ b/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, + )