Fix GPT-OSS MXFP4->NVFP4 PTQ load, export, and cast (nvbug 6295279, 6295242)

CHANGELOG.rst

@@ -71,6 +71,8 @@ Changelog
 - Fix Megatron-Core HF importer to load fused ``TELayerNormColumnParallelLinear.layer_norm_weight`` from HF for GPT-family models (Qwen3 etc.) under ``--export-default-te-spec``. Importer now prefers per-context keys ``fused_input_layernorm`` / ``fused_pre_mlp_layernorm`` (fallback ``fused_norm`` for Nemotron-H backward compatibility); ``mcore_qwen.py`` provides the new rules. Without this fix, post-prune MMLU sat at chance.
 - Fix ONNX AutoCast ``keep_io_types=True`` sanity-check failure (``Unexpected type in I/O tensor ...``) when a network input/output is an empty tensor (a dimension of size 0). Such tensors were "fake-cast" (retyped in place) to the low precision type; because the value-info map aliases the ``graph.input``/``graph.output`` ``ValueInfoProto``, this silently changed the model's I/O type. AutoCast now inserts a real ``Cast`` for protected I/O tensors instead.
 - Fix INT8 entropy calibration of fp16 ONNX models raising ``ValueError: Too many bins for data range`` on numpy >= 2.0. ``_collect_value`` in ``modelopt.onnx.quantization.ort_patching`` now casts the histogram range endpoints to Python float so bin edges are computed in float64, instead of inheriting the fp16 dtype of an activation tensor with a small range (which collapsed the 128-bin linspace under NEP-50 promotion).
+- Fix the GPT-OSS MXFP4 → NVFP4 PTQ path in ``examples/llm_ptq/hf_ptq.py`` (used with ``--cast_mxfp4_to_nvfp4``). ``get_model`` now loads native MXFP4 checkpoints (``openai/gpt-oss-*``) dequantized to BF16 ``GptOssExperts`` via ``Mxfp4Config(dequantize=True)`` on a sequential device map. This fixes a CUDA illegal-memory access during the multi-GPU dequant load and the ``NotImplementedError`` for experts type ``Mxfp4GptOssExperts`` during unified HF export (the packed-kernel experts wrapper, used when the optional ``kernels`` package is installed, is unsupported by export); ``kernels`` is no longer required. The ``--cast_mxfp4_to_nvfp4`` step now also resolves a HF Hub ID ``--pyt_ckpt_path`` to its local snapshot directory instead of failing with ``FileNotFoundError``.
+- Fix ``_QuantGptOssExperts`` / ``_QuantLlama4TextExperts`` static-block NVFP4 weight calibration raising ``ValueError: Input shape has changed`` during the calibration forward. These experts quantize their weights transposed (``_transposed_quantize``); ``iter_weights_for_calibration`` now yields the same transposed view so weight-only calibration and the forward agree on the block-quant shape (and the export ``_amax`` orientation).
 
 **Deprecations**
 

examples/llm_ptq/example_utils.py

@@ -587,6 +587,25 @@ def _resolve_file(filename):
         module.__dict__.pop("weight", None)
 
 
+def get_original_hf_quant_method(config) -> str | None:
+    """Return the checkpoint's original ``quantization_config.quant_method``, if any.
+
+    Returns e.g. ``"mxfp4"`` for native MXFP4 checkpoints (OpenAI's gpt-oss family), or
+    ``None`` for unquantized models. Handles ``quantization_config`` stored as a dict or a
+    config object, and the nested ``text_config`` of multi-modal models.
+    """
+    for cfg in (config, getattr(config, "text_config", None)):
+        quant_cfg = getattr(cfg, "quantization_config", None)
+        method = (
+            quant_cfg.get("quant_method")
+            if isinstance(quant_cfg, dict)
+            else getattr(quant_cfg, "quant_method", None)
+        )
+        if method:
+            return str(method)
+    return None
+
+
 def get_model(
     ckpt_path,
     device="cuda",
@@ -690,6 +709,29 @@ def has_pack_quantized_config(config):
                     trust_remote_code=trust_remote_code,
                     dtype="auto",
                 )
+        elif get_original_hf_quant_method(hf_config) == "mxfp4":
+            # Native MXFP4 checkpoints (e.g. openai/gpt-oss-*) must be dequantized to
+            # plain BF16 experts (``GptOssExperts``) so ModelOpt can insert and export
+            # quantizers: the packed-kernel experts wrapper (``Mxfp4GptOssExperts``,
+            # used when the optional ``kernels`` package is present) is not supported by
+            # the unified HF export. Force dequantization regardless of whether
+            # ``kernels`` is installed.
+            # Local import: ``Mxfp4Config`` only exists in newer Transformers (gpt-oss support);
+            # importing it at module scope would break example_utils for users on older
+            # Transformers running unrelated (non-MXFP4) models.
+            from transformers import Mxfp4Config
+
+            # Load with a *sequential* device map (not "auto"): the MXFP4->BF16 dequant
+            # runs inside Transformers' threaded weight loader, and an "auto"/balanced
+            # split across multiple GPUs trips a CUDA illegal-memory access during dequant
+            # materialization. Sequential keeps each shard's dequant on a single device
+            # (the whole model lands on one GPU when it fits there).
+            model_kwargs["quantization_config"] = Mxfp4Config(dequantize=True)
+            model = AutoModelForCausalLM.from_pretrained(
+                ckpt_path,
+                device_map="cpu" if device == "cpu" else "sequential",
+                **model_kwargs,
+            )
         else:
             architecture = hf_config.architectures[0]
 

examples/llm_ptq/hf_ptq.py

@@ -29,6 +29,7 @@
 from example_utils import (
     _get_auto_quantize_cost_excluded_patterns,
     _get_auto_quantize_disabled_layers,
+    _resolve_model_path,
     build_quant_cfg,
     copy_custom_model_files,
     create_vlm_calibration_loop,
@@ -1163,7 +1164,12 @@ def _is_layerwise(obj):
     # to NVFP4StaticQuantizer with a data-derived ``_global_amax``); we just
     # override that scalar with the closed-form value before export.
     if args.cast_mxfp4_to_nvfp4:
-        apply_cast_mxfp4_to_nvfp4(language_model, args.pyt_ckpt_path)
+        # The cast reads the source MXFP4 ``*_scales``/``*_blocks`` tensors from a local
+        # checkpoint directory. ``--pyt_ckpt_path`` may be a HF Hub ID (e.g.
+        # ``openai/gpt-oss-20b``); resolve it to the local snapshot dir that load_model's
+        # ``from_pretrained`` already populated so the cast works with the documented command.
+        source_ckpt_dir = _resolve_model_path(args.pyt_ckpt_path, args.trust_remote_code)
+        apply_cast_mxfp4_to_nvfp4(language_model, source_ckpt_dir)
 
     post_quantize(
         args,

modelopt/torch/quantization/plugins/huggingface.py

@@ -471,6 +471,34 @@ def backward(ctx, grad_output):
 _transposed_quantize = _TransposedQuantization.apply
 
 
+class _TransposedExpertsCalibMixin:
+    """Weight-only calibration for BMM-style experts that quantize their weights transposed.
+
+    ``_QuantGptOssExperts`` / ``_QuantLlama4TextExperts`` hold 3-D expert weights of shape
+    ``(num_experts, in_dim, out_dim)`` and quantize them *transposed* in the forward (see
+    :class:`_TransposedQuantization`: per-channel / per-block quantization expects the
+    contraction ``in_dim`` as the last axis). Weight-only calibration
+    (``max_calibrate`` -> ``weight_only_quantize``) must feed the weight quantizer the same
+    transposed view; otherwise static-block NVFP4 locks ``_original_shape`` from the
+    non-transposed weight and the forward then raises "Input shape has changed".
+    Calibrating transposed also matches the orientation the unified HF export reads ``_amax``
+    in (it transposes BMM expert weights before deriving scales).
+
+    The transposed view is not made contiguous (unlike the forward's ``_transposed_quantize``,
+    which needs it for the matmul): calibration only reads the shape and reduces for ``_amax``,
+    both of which the quantizer handles on a non-contiguous view via ``reshape``.
+
+    For ``_QuantGptOssExperts`` ``gate_up_proj``/``down_proj`` are dynamic attributes; weight-only
+    calibration runs with weight quantization disabled, so they return the raw weight here.
+    """
+
+    def iter_weights_for_calibration(self):
+        """Yield ``(transposed_weight, weight_quantizer)`` for each expert projection."""
+        for weight_name in ("gate_up_proj", "down_proj"):
+            weight = getattr(self, weight_name)
+            yield weight.transpose(-1, -2), getattr(self, f"{weight_name}_weight_quantizer")
+
+
 class _QuantSparseSequentialMoe(QuantModule):
     """Quantization wrapper for HuggingFace sparse MoE blocks.
 
@@ -601,7 +629,7 @@ def layer_sync_moe_local_experts_amax(self, sync_weight_amax=False):
         sync_moe_expert_amax(self.experts, sync_weight_amax=sync_weight_amax)
 
 
-class _QuantLlama4TextExperts(QuantModule):
+class _QuantLlama4TextExperts(_TransposedExpertsCalibMixin, QuantModule):
     def _setup(self):
         self.gate_up_proj_input_quantizer = TensorQuantizer()
         self.gate_up_proj_weight_quantizer = TensorQuantizer()
@@ -1253,7 +1281,7 @@ def unpack_weight(self):
     pass
 
 
-class _QuantGptOssExperts(_QuantFunctionalMixin):
+class _QuantGptOssExperts(_TransposedExpertsCalibMixin, _QuantFunctionalMixin):
     """Quantized wrapper for `transformers.GptOssExperts`.
 
     Quantizes `gate_up_proj` and `down_proj` weights via dynamic attributes inside `quantize_weight()`.

tests/examples/llm_ptq/test_cast_mxfp4_to_nvfp4.py

@@ -254,3 +254,30 @@ def __init__(self):
 
     with pytest.raises(AssertionError, match="expected NVFP4StaticQuantizer"):
         cast.apply_to_model(_Wrong(), ckpt_dir)
+
+
+# ---------- force_weight_quantizers_static ----------------------------------
+
+
+def test_force_weight_quantizers_static():
+    """Only ``*weight_quantizer`` entries with a ``block_sizes`` dict flip to ``type='static'``;
+    input quantizers and entries without block_sizes are left untouched."""
+    quant_cfg = [
+        {"quantizer_name": "*", "enable": False},
+        {
+            "quantizer_name": "*weight_quantizer",
+            "cfg": {"num_bits": (2, 1), "block_sizes": {-1: 16, "type": "dynamic"}},
+        },
+        {
+            "quantizer_name": "*input_quantizer",
+            "cfg": {"num_bits": (2, 1), "block_sizes": {-1: 16, "type": "dynamic"}},
+        },
+        {"quantizer_name": "*router*", "enable": False},  # no cfg / block_sizes
+    ]
+
+    cast.force_weight_quantizers_static(quant_cfg)
+
+    assert quant_cfg[1]["cfg"]["block_sizes"]["type"] == "static"  # weight quantizer forced
+    assert quant_cfg[1]["cfg"]["block_sizes"][-1] == 16  # other block_sizes keys preserved
+    assert quant_cfg[2]["cfg"]["block_sizes"]["type"] == "dynamic"  # input quantizer untouched
+    assert "cfg" not in quant_cfg[3]  # entry without block_sizes untouched

tests/examples/llm_ptq/test_example_utils.py

@@ -159,3 +159,38 @@ def test_load_mtp_weights_no_mtp_returns_empty(tmp_path):
     prefixes, orphans = example_utils.load_mtp_weights(model, str(tmp_path))
     assert prefixes == []
     assert orphans == {}
+
+
+# ---------- get_original_hf_quant_method -------------------------------------
+# get_model uses this to detect native MXFP4 checkpoints (e.g. openai/gpt-oss-*) and load
+# them dequantized to BF16 GptOssExperts (so ModelOpt can quantize/export the experts).
+
+
+def test_get_original_hf_quant_method_mxfp4_dict():
+    # gpt-oss layout: quantization_config is a plain dict carrying quant_method.
+    cfg = SimpleNamespace(
+        quantization_config={"quant_method": "mxfp4", "modules_to_not_convert": []}
+    )
+    assert example_utils.get_original_hf_quant_method(cfg) == "mxfp4"
+
+
+def test_get_original_hf_quant_method_object():
+    # Some configs expose quantization_config as an object with a quant_method attribute.
+    cfg = SimpleNamespace(quantization_config=SimpleNamespace(quant_method="fp8"))
+    assert example_utils.get_original_hf_quant_method(cfg) == "fp8"
+
+
+def test_get_original_hf_quant_method_nested_text_config():
+    # Multi-modal models nest the quantization_config under text_config.
+    cfg = SimpleNamespace(
+        text_config=SimpleNamespace(quantization_config={"quant_method": "mxfp4"})
+    )
+    assert example_utils.get_original_hf_quant_method(cfg) == "mxfp4"
+
+
+def test_get_original_hf_quant_method_none_for_unquantized():
+    assert example_utils.get_original_hf_quant_method(SimpleNamespace()) is None
+    assert (
+        example_utils.get_original_hf_quant_method(SimpleNamespace(quantization_config=None))
+        is None
+    )

tests/gpu/torch/quantization/test_gpt_oss_mxfp4_nvfp4_cast_cuda.py

@@ -0,0 +1,124 @@
+# 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.
+
+"""End-to-end regression guard for the GPT-OSS MXFP4 -> NVFP4 cast path.
+
+This exercises the exact combination that silently regressed (nvbug 6295279 / 6295242):
+a transposed-quantize MoE (``GptOssExperts``) with **static-block** NVFP4 weight quantizers
+-- which is what ``examples/llm_ptq --cast_mxfp4_to_nvfp4`` produces via
+``force_weight_quantizers_static`` -- calibrated with a forward loop.
+
+The regression (the unconditional ``weight_only_quantize`` from #1560 feeding the
+*non-transposed* expert weight while the forward feeds the transposed one) raised
+``ValueError: Input shape has changed`` during ``mtq.quantize``. Without the
+``iter_weights_for_calibration`` fix this test fails at ``mtq.quantize`` below.
+
+Static-block NVFP4 fake quant uses a Triton kernel, so this is GPU-only.
+"""
+
+import copy
+import json
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+from _test_utils.torch.transformers_models import get_tiny_gpt_oss
+from safetensors.torch import save_file
+
+import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.config import NVFP4_DEFAULT_CFG
+from modelopt.torch.quantization.nn import NVFP4StaticQuantizer
+
+# The cast helpers live next to the example script, not in the ``modelopt`` package.
+_LLM_PTQ_DIR = Path(__file__).resolve().parents[4] / "examples" / "llm_ptq"
+if str(_LLM_PTQ_DIR) not in sys.path:
+    sys.path.insert(0, str(_LLM_PTQ_DIR))
+
+from cast_mxfp4_to_nvfp4 import apply_to_model, force_weight_quantizers_static
+
+_EXPERT_WEIGHT_QUANTIZERS = ("gate_up_proj_weight_quantizer", "down_proj_weight_quantizer")
+
+
+def _write_lossless_mxfp4_source(model, ckpt_dir: Path) -> None:
+    """Write a synthetic MXFP4 source whose ``*_scales``/``*_blocks`` match each expert
+    weight quantizer's per-block ``_amax`` count.
+
+    Every E8M0 scale is ``127`` (exponent ``k = 0``), so all blocks are "lossless" / in-range:
+    the cast derives a per-block amax of ``6 * 2^0 = 6`` and a per-tensor
+    ``global_amax = E2M1_MAX * E4M3_MAX * 2^(k_max - 8) = 6 * 448 * 2^-8 = 10.5``.
+    """
+    state: dict[str, torch.Tensor] = {}
+    for layer_idx in range(model.config.num_hidden_layers):
+        experts = model.model.layers[layer_idx].mlp.experts
+        for wname in ("gate_up_proj", "down_proj"):
+            quantizer = getattr(experts, f"{wname}_weight_quantizer")
+            # Two NVFP4 blocks (of 16) share one MXFP4 block (of 32).
+            n_mxfp4_blocks = quantizer._amax.numel() // 2
+            base = f"model.layers.{layer_idx}.mlp.experts.{wname}"
+            state[f"{base}_scales"] = torch.full((n_mxfp4_blocks,), 127, dtype=torch.uint8)
+            state[f"{base}_blocks"] = torch.zeros((n_mxfp4_blocks, 16), dtype=torch.uint8)
+    save_file(state, str(ckpt_dir / "model-00001-of-00001.safetensors"))
+    (ckpt_dir / "model.safetensors.index.json").write_text(
+        json.dumps(
+            {"metadata": {}, "weight_map": dict.fromkeys(state, "model-00001-of-00001.safetensors")}
+        )
+    )
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="static-block NVFP4 needs a CUDA Triton kernel"
+)
+def test_gpt_oss_mxfp4_to_nvfp4_cast(tmp_path):
+    # intermediate_size != hidden_size so the expert weights are non-square and the
+    # transposed-vs-non-transposed calibration orientations are actually distinguishable.
+    model = (
+        get_tiny_gpt_oss(num_hidden_layers=2, hidden_size=64, intermediate_size=96).cuda().eval()
+    )
+
+    # Mirror ``--cast_mxfp4_to_nvfp4``: force the NVFP4 weight quantizers to static block.
+    quant_cfg = copy.deepcopy(NVFP4_DEFAULT_CFG)
+    force_weight_quantizers_static(quant_cfg["quant_cfg"])
+
+    def forward_loop(m):
+        m(torch.randint(0, model.config.vocab_size, (2, 16), device="cuda"))
+
+    # Regression guard: pre-fix this raised "Input shape has changed" because weight-only
+    # calibration fed the non-transposed expert weight while the forward fed the transposed one.
+    mtq.quantize(model, quant_cfg, forward_loop)
+
+    # Calibration must have promoted every expert weight quantizer to a static NVFP4 quantizer
+    # with a per-block ``_amax`` (the cast's precondition), sized from the *transposed* weight.
+    for layer_idx in range(model.config.num_hidden_layers):
+        experts = model.model.layers[layer_idx].mlp.experts
+        for qname in _EXPERT_WEIGHT_QUANTIZERS:
+            quantizer = getattr(experts, qname)
+            weight = getattr(experts, qname[: -len("_weight_quantizer")])
+            assert isinstance(quantizer, NVFP4StaticQuantizer)
+            assert quantizer._amax is not None
+            assert quantizer._amax.numel() == weight.numel() // 16  # NVFP4 block_size = 16
+            assert quantizer._amax.numel() > 1  # per-block, not per-tensor
+
+    # Run the closed-form cast against a matching MXFP4 source and confirm it overrides every
+    # expert weight quantizer with the closed-form values (matches names + per-block numel).
+    _write_lossless_mxfp4_source(model, tmp_path)
+    apply_to_model(model, tmp_path)
+
+    for layer_idx in range(model.config.num_hidden_layers):
+        experts = model.model.layers[layer_idx].mlp.experts
+        for qname in _EXPERT_WEIGHT_QUANTIZERS:
+            quantizer = getattr(experts, qname)
+            assert torch.allclose(quantizer._amax, torch.full_like(quantizer._amax, 6.0))
+            assert abs(float(quantizer.global_amax) - 10.5) < 1e-3