Fix folding transposed GPT-OSS and Llama 4 expert weights

CHANGELOG.rst

@@ -16,6 +16,7 @@ Changelog
 
 **Bug Fixes**
 
+- Fix ``fold_weight`` failing on quantized GPT-OSS and Llama 4 models with transposed expert weights.
 - 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.

modelopt/torch/quantization/model_quant.py

@@ -43,7 +43,7 @@
 from .config import QuantizeAlgoCfgType
 from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg
 from .nn import QuantModule, SequentialQuantizer, TensorQuantizer
-from .utils import is_quantized
+from .utils import enable_weight_access_and_writeback, is_quantized, module_name_maps
 
 __all__ = [
     "auto_quantize",
@@ -731,9 +731,11 @@ def fold_weight(model: nn.Module, keep_attrs: bool = False):
     Any weight-quantizer rotation is folded into the weights and disabled so subsequent
     forwards do not re-rotate the already-folded weights.
     """
+    names = module_name_maps(model)
     for name, module in model.named_modules():
         if isinstance(module, QuantModule):
-            module.fold_weight(keep_attrs)
+            with enable_weight_access_and_writeback(module, model, names):
+                module.fold_weight(keep_attrs)
 
 
 @contextmanager

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

@@ -145,6 +145,8 @@ def _fold_weight_quantizer(
         quantizer: TensorQuantizer,
         weights: Iterable[torch.Tensor],
         keep_attrs: bool = False,
+        *,
+        quantize_dtype: torch.dtype = torch.float32,
     ):
         """Fold ``quantizer`` into each weight view in place, then disable and clean it once.
 
@@ -156,7 +158,7 @@ def _fold_weight_quantizer(
             return
 
         for weight in weights:
-            weight.data.copy_(quantizer(weight.float().contiguous()).to(weight.dtype))
+            weight.data.copy_(quantizer(weight.to(quantize_dtype).contiguous()).to(weight.dtype))
         quantizer.disable()
         quantizer.disable_rotate()
         if keep_attrs and hasattr(quantizer, "_pre_quant_scale"):

modelopt/torch/quantization/plugins/accelerate.py

@@ -50,7 +50,8 @@ def _writeback_params_to_weights_map(module, align_hook):
             continue
         if isinstance(align_hook.weights_map, PrefixedDataset):
             key = align_hook.weights_map.prefix + name
-            w_map = align_hook.weights_map.dataset.state_dict
+            dataset = align_hook.weights_map.dataset
+            w_map = getattr(dataset, "state_dict", dataset)
         else:
             w_map = align_hook.weights_map
             key = name

modelopt/torch/quantization/plugins/huggingface.py

@@ -620,6 +620,14 @@ def iter_weights_for_calibration(self):
             weight = getattr(self, weight_name)
             yield weight.transpose(-1, -2), getattr(self, f"{weight_name}_weight_quantizer")
 
+    def fold_weight(self, keep_attrs: bool = False):
+        """Fold expert weights in the transposed orientation used by their forward."""
+        for weight, quantizer in self.iter_weights_for_calibration():
+            if isinstance(quantizer, TensorQuantizer):
+                QuantModule._fold_weight_quantizer(
+                    quantizer, (weight,), keep_attrs, quantize_dtype=weight.dtype
+                )
+
 
 class _QuantSparseSequentialMoe(QuantModule):
     """Quantization wrapper for HuggingFace sparse MoE blocks.

tests/unit/torch/quantization/plugins/test_huggingface.py

@@ -382,6 +382,84 @@ def __init__(self):
     assert down_q is experts.down_proj_weight_quantizer
 
 
+@pytest.mark.parametrize("model_type", ["gpt_oss", "llama4"])
+@pytest.mark.parametrize("keep_attrs", [False, True])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("offload", [None, "cpu", "disk"])
+def test_fold_transposed_expert_weights_preserves_forward(
+    model_type, keep_attrs, dtype, offload, tmp_path
+):
+    set_seed()
+    if model_type == "gpt_oss":
+        model = get_tiny_gpt_oss(num_hidden_layers=1, hidden_size=32, intermediate_size=48)
+    else:
+        model = AutoModelForCausalLM.from_config(
+            transformers.Llama4TextConfig(
+                num_hidden_layers=1,
+                num_local_experts=4,
+                hidden_size=32,
+                intermediate_size=48,
+                intermediate_size_mlp=48,
+                num_attention_heads=2,
+                num_key_value_heads=1,
+                head_dim=16,
+                vocab_size=32,
+                pad_token_id=0,
+                bos_token_id=1,
+                eos_token_id=2,
+                no_rope_layers=[0],
+            )
+        )
+    model = model.to(dtype=dtype).eval()
+    input_ids = torch.randint(0, 32, (2, 7))
+    config = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*.experts.*weight_quantizer",
+                "cfg": {"num_bits": 8, "axis": (0, 1)},
+            },
+        ],
+        "algorithm": "max",
+    }
+    mtq.quantize(model, config, lambda m: m(input_ids))
+    experts = next(
+        module for module in model.modules() if isinstance(module, _TransposedExpertsCalibMixin)
+    )
+    pairs = list(experts.iter_weights_for_calibration())
+    for weight, quantizer in pairs:
+        quantizer.pre_quant_scale = torch.linspace(0.75, 1.25, weight.shape[-1], dtype=dtype)
+    pointers = [weight.data_ptr() for weight, _ in pairs]
+
+    with torch.no_grad():
+        expected = model(input_ids).logits
+        if offload is not None:
+            # Accelerate is optional; only offload cases require it.
+            accelerate = pytest.importorskip("accelerate")
+            if offload == "cpu":
+                accelerate.cpu_offload(experts, execution_device=torch.device("cpu"))
+            else:
+                accelerate.disk_offload(
+                    experts, tmp_path / "offload", execution_device=torch.device("cpu")
+                )
+            assert experts.gate_up_proj.is_meta
+            torch.testing.assert_close(model(input_ids).logits, expected)
+        mtq.fold_weight(model, keep_attrs=keep_attrs)
+        torch.testing.assert_close(model(input_ids).logits, expected)
+        for (weight, quantizer), pointer in zip(pairs, pointers):
+            if offload is None:
+                assert weight.data_ptr() == pointer
+            assert not quantizer.is_enabled
+            assert quantizer.pre_quant_scale is None
+            assert hasattr(quantizer, "_amax") == keep_attrs
+            assert hasattr(quantizer, "_pre_quant_scale") == keep_attrs
+        mtq.fold_weight(model, keep_attrs=keep_attrs)
+        torch.testing.assert_close(model(input_ids).logits, expected)
+        if offload is not None:
+            assert experts.gate_up_proj.is_meta
+            assert experts.down_proj.is_meta
+
+
 def test_hf_decoder_discoverer_registration_path():
     model = get_tiny_llama()
     assert any(