Address second review round: symmetric base_layer lookup, loud miss, value oracle

- _restore_qtensor_wrappers: normalize the .base_layer suffix on both the saved
  q_tensor_state keys and the module names, so the lookup works whether the model
  was compressed before adapters were attached (quantize.py --compress) or after
  (QATTrainer._quantize_model). Warn when saved weights match no module at all,
  which is the silent failure this PR set out to fix.
- export.py: document why skipping the restore is safe (QATTrainer snapshots the
  state at trainer init from the base state from_pretrained already restored).
- LoRA-QAT example test: compare exported scales against a direct PTQ export
  rather than only asserting key presence.
- Unit tests: cover the reverse key direction and the no-match warning.
- CHANGELOG: add the 0.47 Bug Fixes entry.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>

Author: sugunav14

CHANGELOG.rst

@@ -21,6 +21,8 @@ Changelog
 
 **Bug Fixes**
 
+- 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.46 (2026-08-xx)
 ^^^^^^^^^^^^^^^^^
 

examples/llm_qat/export.py

@@ -50,6 +50,9 @@ def get_model(
 
     # 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 the restore below is only safe because QATTrainer writes modelopt_state_train.pth at
+    # trainer init from that same base state, so it carries no quantizer values from_pretrained did
+    # not already load. Revisit this if the trainer ever snapshots state later in training.
     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)

modelopt/torch/opt/plugins/transformers.py

@@ -110,18 +110,38 @@ 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 not isinstance(module, RealQuantLinear) or isinstance(module.weight, QTensorWrapper):
-                continue
-            # PEFT renames the quantized linear to `<name>.base_layer`, but `q_tensor_state` is
-            # keyed by the name it was saved with, so fall back to the stripped name.
-            key = name if name in q_tensor_state else name.removesuffix(".base_layer")
+        # PEFT nests the quantized linear as `<name>.base_layer`, and either side may carry that
+        # suffix: the state is saved without it when the base model was compressed before adapters
+        # were attached (`quantize.py --compress`) and with it when compressed after
+        # (`QATTrainer._quantize_model`). Normalize both so the lookup works in either direction.
+        # This assumes transformers injects adapters in place; moving the trainer to
+        # `get_peft_model` would prefix module names with `base_model.model.` and break it.
+        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 the module names were remapped by a wrapper we do not know about.
+        # Warn here rather than letting it surface as an opaque shape error during 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

@@ -190,17 +190,22 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path):
     with open(base_model_dir / "hf_quant_config.json") as f:
         assert json.load(f)["quantization"]["quant_algo"] == "NVFP4"
 
-    # Scales are derived from the calibrated amaxes; missing ones mean the export silently
-    # fell back to uncalibrated quantizers.
     base_weights = load_file(base_model_dir / "model.safetensors")
-    weight_scales = [k for k in base_weights if k.endswith(".weight_scale")]
-    assert weight_scales, "no NVFP4 weight scales in the exported base model"
-    for key in weight_scales:
-        prefix = key.removesuffix(".weight_scale")
-        assert f"{prefix}.weight_scale_2" in base_weights
-        assert f"{prefix}.input_scale" in base_weights
     assert not any("base_layer" in k or k.endswith("_amax") for k in base_weights)
 
+    # LoRA freezes the base model, so exporting the PTQ checkpoint directly is a trusted oracle
+    # for every calibrated value. Comparing against it catches scales that survive as keys but
+    # were silently reset to defaults, which key-presence assertions alone would miss.
+    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",

tests/gpu/torch/export/test_export.py

@@ -291,7 +291,8 @@ def test_postprocess_state_dict_qlora_strips_base_layer():
         "layer1.bias",
         "layer1.pre_quant_scale",
     }
-    assert processed_state_dict["layer1.weight_scale_2"] == torch.tensor([0.5])
+    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(

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

@@ -69,14 +69,22 @@ def __init__(self, base_layer):
         self.base_layer = base_layer
 
 
-def _compressed_model_and_state_dir(tmp_path):
+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)
 
-    torch.save(mto.modelopt_state(model), tmp_path / "modelopt_state.pth")
+    state = mto.modelopt_state(model)
+    if state_keyed_with_base_layer:
+        # Compressing after the adapters are attached (QATTrainer._quantize_model) saves the
+        # keys with the peft suffix already in them.
+        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 loads weights by assigning a plain Parameter holding the already-packed
     # data, which leaves the module without its QTensorWrapper.
@@ -88,9 +96,10 @@ def _compressed_model_and_state_dir(tmp_path):
 
 
 @pytest.mark.parametrize("wrap_in_lora", [False, True])
-def test_restore_qtensor_wrappers(tmp_path, wrap_in_lora):
-    """`q_tensor_state` is keyed by the pre-peft name, so `<name>.base_layer` must still match."""
-    model = _compressed_model_and_state_dir(tmp_path)
+@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)
 
@@ -99,3 +108,14 @@ def test_restore_qtensor_wrappers(tmp_path, wrap_in_lora):
     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)