Bug fix: 6542481

CHANGELOG.rst

@@ -102,6 +102,7 @@ Changelog
 - Fix ``examples/vllm_serve`` serving shared experts uncalibrated: their ``gate_proj``/``up_proj`` quantizer keys were not merged into ``gate_up_proj`` on reload, so they matched no module and were dropped.
 - Fix Qwen3-VL MoE PTQ failing on ``transformers>=5.12`` with ``AttributeError: 'QuantQwen3VLMoeTextExperts' object has no attribute 'hidden_size'`` (NVBug 6518551). transformers 5.12 moved ``Qwen3VLMoeTextExperts`` onto the standard ``@use_experts_implementation`` fused layout (``hidden_size``/``expert_dim`` renamed to ``hidden_dim``/``intermediate_dim``, ``gate_up_proj`` transposed to ``(num_experts, 2*intermediate_dim, hidden_dim)``, two ``F.linear`` calls per expert), but the legacy ``_QuantQwen3VLMoeTextExperts`` wrapper stayed statically registered and shadowed on-the-fly detection. The new layout is now left to ``register_fused_experts_on_the_fly``, which claims it with the generic ``_QuantFusedExperts``; the legacy wrapper is still registered on ``transformers<5.12``, whose ``torch.bmm``-based forward the generic wrapper cannot intercept.
 - Fix ``examples/hf_ptq`` multi-node FSDP2 export (``--use_fsdp2``) failing with ``RuntimeError: Cannot set version_counter for inference tensor``. ``export_quantized`` wrapped its whole body in ``torch.inference_mode()``, so the full params gathered by ``get_model_state_dict(full_state_dict=True)`` were inference tensors and the subsequent ``state_dict()`` -> ``param.detach()`` could not set their version counter. The export context is now ``torch.no_grad()``, which still disables autograd but keeps the gathered params as normal tensors.
+- Fix QLoRA export in ``examples/llm_qat/export.py`` failing with ``AssertionError: Model already has modelopt state!`` (NVBug 6542481). The QLoRA training output is an adapter-only checkpoint, so ``from_pretrained`` resolves the quantized base model from ``adapter_config.json`` and ``enable_huggingface_checkpointing`` already restores its ModelOpt state; the export then restored a second time. It now restores only when the loaded model is not already converted. Two further breakages on the same path are also fixed: ``_restore_qtensor_wrappers`` matched no modules because PEFT re-parents the quantized linear as ``<name>.base_layer`` while ``q_tensor_state`` is keyed by the name it was saved with (the packed NVFP4 weight then reached ``F.linear`` and raised a shape error), and ``postprocess_state_dict`` silently dropped every ``base_layer.*`` key missing from a hand-maintained rename map — losing the NVFP4 ``weight_scale_2`` global scale and any linear ``bias`` (Qwen2-style q/k/v biases), and leaving ``base_layer`` in the exported AWQ ``pre_quant_scale`` key. The rename is now a generic ``.base_layer.`` strip.
 
 0.45 (2026-07-02)
 ^^^^^^^^^^^^^^^^^

examples/llm_qat/export.py

@@ -23,7 +23,7 @@
 import modelopt.torch.opt as mto
 from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format
 from modelopt.torch.export.unified_export_hf import _export_transformers_checkpoint
-from modelopt.torch.opt.conversion import restore_from_modelopt_state
+from modelopt.torch.opt.conversion import ModeloptStateManager, restore_from_modelopt_state
 from modelopt.torch.quantization.utils import set_quantizer_state_dict
 from modelopt.torch.utils import print_rank_0
 
@@ -48,8 +48,11 @@ def get_model(
     # Load model
     model = AutoModelForCausalLM.from_pretrained(ckpt_path, device_map=device_map)
 
-    # Restore modelopt state for LoRA models. For QAT/QAD models from_pretrained call handles this
-    if hasattr(model, "peft_config"):
+    # Restore modelopt state for LoRA models. For QAT/QAD models from_pretrained call handles this.
+    # For QLoRA the base checkpoint is quantized, so from_pretrained already restored the state.
+    # Skipping is safe only because QATTrainer writes modelopt_state_train.pth at trainer init,
+    # from that same base state.
+    if hasattr(model, "peft_config") and not ModeloptStateManager.is_converted(model):
         modelopt_state = mto.load_modelopt_state(f"{ckpt_path}/modelopt_state_train.pth")
         restore_from_modelopt_state(model, modelopt_state)
         print_rank_0("Restored modelopt state")

modelopt/torch/export/quant_utils.py

@@ -991,16 +991,10 @@ def postprocess_state_dict(
         "weight_shape",
     ]
 
-    # For modelopt-trained LoRA models, we need to remove the base_layer prefix from the keys for deployment
-    if is_modelopt_qlora:
-        replacements.update(
-            {
-                "base_layer.weight": "weight",
-                "base_layer.input_scale": "input_scale",
-                "base_layer.weight_scale": "weight_scale",
-            }
-        )
-        skip_keys.append("base_layer")
+    def _export_key(key: str) -> str:
+        # PEFT nests the quantized module under `base_layer`, which deployment does not expect.
+        # Strip it generically so new key types (bias, scales) do not need to be enumerated here.
+        return key.replace(".base_layer.", ".") if is_modelopt_qlora else key
 
     post_state_dict = {}
 
@@ -1012,7 +1006,7 @@ def postprocess_state_dict(
 
         # Skip keys not related to quantizers
         if all(skip_key not in key for skip_key in skip_keys):
-            post_state_dict[key] = value
+            post_state_dict[_export_key(key)] = value
             continue
 
         # Apply replacements if the key matches any suffix in the replacements dict
@@ -1033,7 +1027,7 @@ def postprocess_state_dict(
                         logger.warning(
                             "Large KV activations detected. Quantized KV cache may lead to higher accuracy drop."
                         )
-                post_state_dict[prefix + new_suffix] = value
+                post_state_dict[_export_key(prefix + new_suffix)] = value
                 break
 
     # Squeeze scales with a leading dimension of 1

modelopt/torch/opt/plugins/transformers.py

@@ -110,16 +110,34 @@ def _restore_qtensor_wrappers(model, model_path):
         q_tensor_state = mode_config.get("metadata", {}).get("q_tensor_state", {})
         if not q_tensor_state:
             continue
-        for name, module in model.named_modules():
-            if (
-                isinstance(module, RealQuantLinear)
-                and name in q_tensor_state
-                and not isinstance(module.weight, QTensorWrapper)
-            ):
-                module._parameters["weight"] = QTensorWrapper(
-                    qtensor=module.weight.data,
-                    metadata=q_tensor_state[name]["metadata"],
-                )
+        # PEFT nests the quantized linear as `<name>.base_layer`, and either the saved keys or the
+        # live names may carry that suffix. Normalize both so the lookup works in either direction.
+        q_tensor_state = {k.removesuffix(".base_layer"): v for k, v in q_tensor_state.items()}
+
+        pending = [
+            (name, module)
+            for name, module in model.named_modules()
+            if isinstance(module, RealQuantLinear) and not isinstance(module.weight, QTensorWrapper)
+        ]
+        matched = 0
+        for name, module in pending:
+            key = name.removesuffix(".base_layer")
+            if key not in q_tensor_state:
+                continue
+            module._parameters["weight"] = QTensorWrapper(
+                qtensor=module.weight.data,
+                metadata=q_tensor_state[key]["metadata"],
+            )
+            matched += 1
+
+        # A total miss means some wrapper renamed the modules. Warn instead of letting it surface
+        # as an opaque shape error at dequantization.
+        if pending and not matched:
+            warnings.warn(
+                f"Found {len(q_tensor_state)} compressed weight(s) in {modelopt_state_path} but "
+                f"re-wrapped none of the {len(pending)} candidate module(s); their names may have "
+                "been remapped. The model will likely fail when the packed weights are used."
+            )
 
 
 def _new_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwargs):

tests/examples/llm_qat/test_llm_qat.py

@@ -14,8 +14,12 @@
 # limitations under the License.
 
 
+import json
+
 import pytest
+import torch
 from _test_utils.examples.run_command import run_example_command
+from safetensors.torch import load_file
 
 # Mapping from backend name to accelerate config file
 BACKEND_CONFIGS = {
@@ -86,6 +90,18 @@ def _run_train(config: str, extra_cmd_args: list[str], backend: str = "fsdp2", c
         setup_free_port=True,
     )
 
+
+def _run_export(ckpt_dir: str, export_dir: str):
+    run_example_command(
+        [
+            "python", "export.py",
+            "--pyt_ckpt_path", ckpt_dir,
+            "--export_path", export_dir,
+        ],
+        "llm_qat",
+    )
+
+
 def test_dataset_utils_pretokenize(tiny_qwen3_path, tmp_path):
     """Test dataset_utils.py standalone CLI pre-tokenization."""
     cache_dir = tmp_path / "dataset_cache"
@@ -152,18 +168,43 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path):
     )
 
     # Step 2: LoRA QAT
+    lora_qat_output_dir = tmp_path / "lora_qat"
     _run_train(
         "configs/train/qat_nvfp4.yaml",
         [
             "--model_name_or_path", str(ptq_output_dir),
             "--do_train", "True",
             "--lora", "True",
-            "--output_dir", str(tmp_path / "lora_qat"),
+            "--output_dir", str(lora_qat_output_dir),
         ],
         backend="fsdp2",
         cache_dir=cache_dir,
     )
 
+    # Step 3: Export. This checkpoint is fake-quantized, so the calibrated amaxes rather than
+    # packed weights are what must survive the load.
+    export_dir = tmp_path / "lora_qat_export"
+    _run_export(str(lora_qat_output_dir), str(export_dir))
+
+    base_model_dir = export_dir / "base_model"
+    with open(base_model_dir / "hf_quant_config.json") as f:
+        assert json.load(f)["quantization"]["quant_algo"] == "NVFP4"
+
+    base_weights = load_file(base_model_dir / "model.safetensors")
+    assert not any("base_layer" in k or k.endswith("_amax") for k in base_weights)
+
+    # LoRA freezes the base model, so a direct PTQ export is a trusted oracle for every calibrated
+    # value. This catches scales that keep their key but were reset to defaults.
+    ptq_export_dir = tmp_path / "ptq_export"
+    _run_export(str(ptq_output_dir), str(ptq_export_dir))
+    reference = load_file(ptq_export_dir / "model.safetensors")
+
+    scales = [k for k in reference if k.endswith(("_scale", "_scale_2"))]
+    assert scales, "no NVFP4 scales in the reference PTQ export"
+    for key in scales:
+        assert key in base_weights, f"{key} missing from the LoRA-QAT export"
+        assert torch.equal(base_weights[key], reference[key]), f"{key} does not match PTQ export"
+
 
 @pytest.mark.parametrize("backend", [
     "fsdp2",
@@ -219,14 +260,37 @@ def test_qwen3_qlora_nvfp4(tiny_qwen3_path, tmp_path):
     )
 
     # Step 2: QLoRA training
+    qlora_output_dir = tmp_path / "qlora"
     _run_train(
         "configs/train/qlora_nvfp4.yaml",
         [
             "--model_name_or_path", str(ptq_output_dir),
             "--do_train", "True",
             "--lora", "True",
-            "--output_dir", str(tmp_path / "qlora"),
+            "--output_dir", str(qlora_output_dir),
         ],
         backend="ddp",
         cache_dir=cache_dir,
     )
+
+    # Step 3: Export the QLoRA checkpoint for deployment
+    export_dir = tmp_path / "qlora_export"
+    _run_export(str(qlora_output_dir), str(export_dir))
+
+    # The base model is exported compressed; the adapters stay at the top level.
+    base_model_dir = export_dir / "base_model"
+    assert (export_dir / "adapter_model.safetensors").is_file()
+    assert (base_model_dir / "hf_quant_config.json").is_file()
+
+    with open(base_model_dir / "hf_quant_config.json") as f:
+        assert json.load(f)["quantization"]["quant_algo"] == "NVFP4"
+
+    # NVFP4 needs the packed weight and *both* scales to be dequantizable downstream.
+    base_weights = load_file(base_model_dir / "model.safetensors")
+    packed_weights = [k for k, v in base_weights.items() if k.endswith(".weight") and v.dtype == torch.uint8]
+    assert packed_weights, "no NVFP4-packed weights found in the exported base model"
+    for key in packed_weights:
+        prefix = key.removesuffix(".weight")
+        assert f"{prefix}.weight_scale" in base_weights
+        assert f"{prefix}.weight_scale_2" in base_weights
+    assert not any("base_layer" in k or "lora" in k for k in base_weights)

tests/gpu/torch/export/test_export.py

@@ -259,6 +259,41 @@ def test_postprocess_state_dict(state_dict, quantization, maxbound, expected_sta
     assert processed_state_dict == expected_state_dict
 
 
+def test_postprocess_state_dict_qlora_strips_base_layer():
+    """Every QLoRA `base_layer.*` tensor needed for deployment must survive the rename.
+
+    Dropping the NVFP4 global scale or a bias yields an undeployable checkpoint.
+    """
+    state_dict = {
+        "layer1.base_layer.weight": torch.ones(4, 2, dtype=torch.uint8),
+        "layer1.base_layer.weight_scale": torch.ones(4, 1),
+        "layer1.base_layer.weight_scale_2": torch.tensor([0.5]),
+        "layer1.base_layer.input_scale": torch.tensor([0.25]),
+        "layer1.base_layer.bias": torch.arange(4.0),
+        "layer1.base_layer.input_quantizer._pre_quant_scale": torch.ones(2),
+        # Quantizer internals must still be dropped.
+        "layer1.base_layer.weight_quantizer._amax": torch.tensor([1.0]),
+        "layer1.base_layer.input_quantizer._amax": torch.tensor([1.0]),
+        "layer1.base_layer.weight_quantizer._scale": torch.ones(4, 1),
+        "layer1.base_layer.weight_quantizer._double_scale": torch.tensor([0.5]),
+    }
+
+    processed_state_dict = postprocess_state_dict(
+        state_dict, 448.0, QUANTIZATION_NONE, is_modelopt_qlora=True
+    )
+
+    assert set(processed_state_dict) == {
+        "layer1.weight",
+        "layer1.weight_scale",
+        "layer1.weight_scale_2",
+        "layer1.input_scale",
+        "layer1.bias",
+        "layer1.pre_quant_scale",
+    }
+    assert torch.equal(processed_state_dict["layer1.weight_scale_2"], torch.tensor([0.5]))
+    assert torch.equal(processed_state_dict["layer1.bias"], torch.arange(4.0))
+
+
 @pytest.mark.parametrize(
     ("config", "expected"),
     [

tests/unit/torch/opt/plugins/test_hf_patching.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 
 import pytest
+import torch
+import torch.nn as nn
 from _test_utils.torch.transformers_models import (
     create_tiny_llama_dir,
     get_tiny_qwen3,
@@ -23,6 +25,9 @@
 
 import modelopt.torch.distill as mtd
 import modelopt.torch.opt as mto
+import modelopt.torch.quantization as mtq
+from modelopt.torch.opt.plugins.transformers import _restore_qtensor_wrappers
+from modelopt.torch.quantization.qtensor import QTensorWrapper
 
 
 @pytest.mark.parametrize(
@@ -54,3 +59,61 @@ def test_nested_model_save_restore(tmp_path, model_cls, teacher_model_type):
     tf_output_tester(model, model_test)
     # KD state is not saved and it should be empty
     assert not mto.ModeloptStateManager(model_test).has_state
+
+
+class _LoraLike(nn.Module):
+    """Stand-in for peft's `lora.Linear`, which nests the original module under `base_layer`."""
+
+    def __init__(self, base_layer):
+        super().__init__()
+        self.base_layer = base_layer
+
+
+def _compressed_model_and_state_dir(tmp_path, state_keyed_with_base_layer=False):
+    model = nn.Sequential()
+    model.fc = nn.Linear(64, 32)
+    mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(2, 64)))
+    mtq.compress(model)
+    assert isinstance(model.fc.weight, QTensorWrapper)
+
+    state = mto.modelopt_state(model)
+    if state_keyed_with_base_layer:
+        # Compressing after the adapters are attached saves the keys with the peft suffix.
+        for _, mode_config in state["modelopt_state_dict"]:
+            q_tensor_state = mode_config.get("metadata", {}).get("q_tensor_state", {})
+            for key in list(q_tensor_state):
+                q_tensor_state[f"{key}.base_layer"] = q_tensor_state.pop(key)
+    torch.save(state, tmp_path / "modelopt_state.pth")
+
+    # transformers>=5 assigns a plain Parameter holding the packed data, dropping the wrapper.
+    packed = model.fc.weight.data.clone()
+    del model.fc._parameters["weight"]
+    model.fc._parameters["weight"] = nn.Parameter(packed, requires_grad=False)
+    assert not isinstance(model.fc.weight, QTensorWrapper)
+    return model
+
+
+@pytest.mark.parametrize("wrap_in_lora", [False, True])
+@pytest.mark.parametrize("state_keyed_with_base_layer", [False, True])
+def test_restore_qtensor_wrappers(tmp_path, wrap_in_lora, state_keyed_with_base_layer):
+    """Either side may carry the `.base_layer` suffix, so the lookup must work in both directions."""
+    model = _compressed_model_and_state_dir(tmp_path, state_keyed_with_base_layer)
+    if wrap_in_lora:
+        model.fc = _LoraLike(model.fc)
+
+    _restore_qtensor_wrappers(model, str(tmp_path))
+
+    linear = model.fc.base_layer if wrap_in_lora else model.fc
+    assert isinstance(linear.weight, QTensorWrapper)
+    assert linear.weight.metadata["shape"] == torch.Size([32, 64])
+
+
+def test_restore_qtensor_wrappers_warns_when_nothing_matches(tmp_path):
+    """A total miss must be loud -- it otherwise surfaces as an opaque shape error at dequant."""
+    model = _compressed_model_and_state_dir(tmp_path)
+    model.fc = _LoraLike(_LoraLike(model.fc))  # a nesting the lookup does not know about
+
+    with pytest.warns(UserWarning, match="re-wrapped none"):
+        _restore_qtensor_wrappers(model, str(tmp_path))
+
+    assert not isinstance(model.fc.base_layer.base_layer.weight, QTensorWrapper)