From 8766bf21e6eedc73ef67460758362d3fd60dfa44 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Wed, 10 Jun 2026 17:59:14 +0000 Subject: [PATCH] Fix GPT-OSS MXFP4->NVFP4 PTQ load, export, and cast (nvbug 6295279, 6295242) The documented GPT-OSS MXFP4->NVFP4 command hf_ptq.py --pyt_ckpt_path openai/gpt-oss-20b --qformat nvfp4_mlp_only \ --cast_mxfp4_to_nvfp4 --export_path ... failed in three ways; all are fixed and the command now runs end-to-end, producing a bit-exact (100% lossless) NVFP4 checkpoint. 1. nvbug 6295242 - CUDA illegal memory access on load. GPT-OSS ships native MXFP4 weights that Transformers dequantizes to BF16; the threaded weight loader trips an illegal-memory access when device_map="auto" shards the dequant across multiple GPUs (the missing optional 'kernels' package only forces the dequant path, it is not the root cause). get_model now detects MXFP4 checkpoints and loads them with Mxfp4Config(dequantize=True) on a sequential device map so the dequant stays on a single device. 2. nvbug 6295279 #1 - unified HF export raised NotImplementedError for experts type 'Mxfp4GptOssExperts'. Forcing dequantize=True yields plain GptOssExperts (even when 'kernels' is installed), which ModelOpt wraps and exports normally. 3. nvbug 6295279 #2 - the --cast_mxfp4_to_nvfp4 step treated --pyt_ckpt_path as a local dir, so a HF Hub ID failed with FileNotFoundError. Resolve it to the cached local snapshot dir via _resolve_model_path before the cast. Also fixes a static-block NVFP4 regression (surfaced by the cast's force_weight_quantizers_static) introduced by #1560's unconditional weight_only_quantize: _QuantGptOssExperts / _QuantLlama4TextExperts quantize their expert weights transposed in the forward (_transposed_quantize) but the inherited iter_weights_for_calibration fed the non-transposed weight, locking a mismatched block-quant _original_shape and raising 'Input shape has changed'. Override iter_weights_for_calibration to calibrate on the transposed view, matching both the forward and the export's _amax orientation. Adds unit tests for get_original_hf_quant_method and the transposed expert-weight calibration. Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 2 + examples/llm_ptq/example_utils.py | 42 ++++++ examples/llm_ptq/hf_ptq.py | 8 +- .../torch/quantization/plugins/huggingface.py | 32 ++++- .../llm_ptq/test_cast_mxfp4_to_nvfp4.py | 27 ++++ tests/examples/llm_ptq/test_example_utils.py | 35 +++++ .../test_gpt_oss_mxfp4_nvfp4_cast_cuda.py | 124 ++++++++++++++++++ .../quantization/plugins/test_huggingface.py | 53 ++++++++ 8 files changed, 320 insertions(+), 3 deletions(-) create mode 100644 tests/gpu/torch/quantization/test_gpt_oss_mxfp4_nvfp4_cast_cuda.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 373994d8662..400227b4d81 100755 --- a/CHANGELOG.rst +++ b/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** diff --git a/examples/llm_ptq/example_utils.py b/examples/llm_ptq/example_utils.py index 9c692e5b7aa..dd2c049fee0 100755 --- a/examples/llm_ptq/example_utils.py +++ b/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] diff --git a/examples/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py index f3e7445e8b4..e424802e907 100755 --- a/examples/llm_ptq/hf_ptq.py +++ b/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, diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 631226dd090..97e13f419f9 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/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()`. diff --git a/tests/examples/llm_ptq/test_cast_mxfp4_to_nvfp4.py b/tests/examples/llm_ptq/test_cast_mxfp4_to_nvfp4.py index c6446d27b93..da145f8eec4 100644 --- a/tests/examples/llm_ptq/test_cast_mxfp4_to_nvfp4.py +++ b/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 diff --git a/tests/examples/llm_ptq/test_example_utils.py b/tests/examples/llm_ptq/test_example_utils.py index 0bbc31dcde0..0ed392340ad 100644 --- a/tests/examples/llm_ptq/test_example_utils.py +++ b/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 + ) diff --git a/tests/gpu/torch/quantization/test_gpt_oss_mxfp4_nvfp4_cast_cuda.py b/tests/gpu/torch/quantization/test_gpt_oss_mxfp4_nvfp4_cast_cuda.py new file mode 100644 index 00000000000..252f8390d69 --- /dev/null +++ b/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 diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py index 34cf6c7f95d..800a8159d46 100644 --- a/tests/unit/torch/quantization/plugins/test_huggingface.py +++ b/tests/unit/torch/quantization/plugins/test_huggingface.py @@ -33,6 +33,7 @@ import modelopt.torch.quantization as mtq from modelopt.torch.quantization.nn import QuantLinear, QuantModuleRegistry from modelopt.torch.quantization.plugins.huggingface import ( + _TransposedExpertsCalibMixin, get_homogeneous_hf_decoder_layers, is_homogeneous_hf_model, ) @@ -233,6 +234,58 @@ def test_is_homogeneous_hf_model_gpt_oss(): assert is_homogeneous_hf_model(model) +def test_gpt_oss_experts_iter_weights_for_calibration_transposed(): + """``_QuantGptOssExperts`` quantizes its expert weights *transposed* in the forward + (``_transposed_quantize`` puts the contraction ``in_dim`` last). Weight-only + calibration must yield the same transposed view; otherwise the unconditional + ``weight_only_quantize`` locks a non-transposed block-quant ``_original_shape`` and the + calibration forward then raises "Input shape has changed" for static-block NVFP4. + """ + # Use intermediate_size != hidden_size so both expert weights are non-square and the + # transpose is observable in the shape. + model = get_tiny_gpt_oss(num_hidden_layers=1, hidden_size=32, intermediate_size=48) + mtq.replace_quant_module(model) + experts = model.model.layers[0].mlp.experts + assert hasattr(experts, "gate_up_proj_weight_quantizer") + + yielded = {q: w for w, q in experts.iter_weights_for_calibration()} + # Stored weights are (num_experts, in_dim, out_dim); calibration must see (…, out_dim, in_dim). + assert ( + yielded[experts.gate_up_proj_weight_quantizer].shape + == experts.gate_up_proj.transpose(-1, -2).shape + ) + assert ( + yielded[experts.down_proj_weight_quantizer].shape + == experts.down_proj.transpose(-1, -2).shape + ) + + +def test_transposed_experts_calib_mixin_yields_transposed_views(): + """Unit-level guard for the shared ``_TransposedExpertsCalibMixin`` (no GPU / no model + conversion needed): it must yield the transposed ``(num_experts, out, in)`` weight view + paired with the matching weight quantizer, so weight-only calibration agrees with the + experts' transposed forward (regression for the static-block "Input shape has changed"). + """ + + class _FakeExperts(_TransposedExpertsCalibMixin): + def __init__(self): + # Non-square so the transpose is observable; (num_experts, in_dim, out_dim). + self.gate_up_proj = torch.randn(8, 64, 192) + self.down_proj = torch.randn(8, 96, 64) + self.gate_up_proj_weight_quantizer = nn.Identity() + self.down_proj_weight_quantizer = nn.Identity() + + experts = _FakeExperts() + pairs = list(experts.iter_weights_for_calibration()) + + assert len(pairs) == 2 + (gate_up_w, gate_up_q), (down_w, down_q) = pairs + assert torch.equal(gate_up_w, experts.gate_up_proj.transpose(-1, -2)) + assert gate_up_q is experts.gate_up_proj_weight_quantizer + assert torch.equal(down_w, experts.down_proj.transpose(-1, -2)) + assert down_q is experts.down_proj_weight_quantizer + + def test_hf_decoder_discoverer_registration_path(): model = get_tiny_llama() assert any(