From e7aae46ada90da9ac37e5b78e5d07fe3bad1f22c Mon Sep 17 00:00:00 2001
From: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Date: Tue, 4 Aug 2026 17:48:38 +0000
Subject: [PATCH 1/7] bug fix

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
---
 examples/llm_qat/export.py                 | 30 ++++++++++------
 modelopt/torch/export/quant_utils.py       |  3 ++
 modelopt/torch/opt/plugins/transformers.py | 21 ++++++-----
 tests/examples/llm_qat/test_llm_qat.py     | 41 +++++++++++++++++++++-
 tests/gpu/torch/export/test_export.py      | 29 +++++++++++++++
 5 files changed, 103 insertions(+), 21 deletions(-)

diff --git a/examples/llm_qat/export.py b/examples/llm_qat/export.py
index f48e85c3ee4..49201baad52 100644
--- a/examples/llm_qat/export.py
+++ b/examples/llm_qat/export.py
@@ -15,6 +15,7 @@
 
 import argparse
 import json
+import os
 import warnings
 from pathlib import Path
 
@@ -23,7 +24,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,17 +49,24 @@ 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"):
-        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")
-
-        # Restore modelopt quantizer state dict
+    # Restore modelopt state for LoRA models. For QAT/QAD models from_pretrained call handles this.
+    # For QLoRA the base model checkpoint is itself quantized, so `from_pretrained` has already
+    # restored the modelopt state (and the quantizer buffers) via `enable_huggingface_checkpointing`.
+    # Restoring a second time would raise `Model already has modelopt state!`, so only restore the
+    # state here if the loaded model does not have it yet.
+    modelopt_state_path = os.path.join(ckpt_path, "modelopt_state_train.pth")
+    if hasattr(model, "peft_config") and os.path.isfile(modelopt_state_path):
+        modelopt_state = mto.load_modelopt_state(modelopt_state_path)
         modelopt_weights = modelopt_state.pop("modelopt_state_weights", None)
-        if modelopt_weights is not None:
-            set_quantizer_state_dict(model, modelopt_weights)
-            print_rank_0("Restored modelopt quantizer state dict")
+
+        if not ModeloptStateManager.is_converted(model):
+            restore_from_modelopt_state(model, modelopt_state)
+            print_rank_0("Restored modelopt state")
+
+            # Restore modelopt quantizer state dict
+            if modelopt_weights is not None:
+                set_quantizer_state_dict(model, modelopt_weights)
+                print_rank_0("Restored modelopt quantizer state dict")
 
     return model
 
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index ab2ef0d9029..cb153ab2e5e 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -998,6 +998,9 @@ def postprocess_state_dict(
                 "base_layer.weight": "weight",
                 "base_layer.input_scale": "input_scale",
                 "base_layer.weight_scale": "weight_scale",
+                # NVFP4 double quantization keeps a per-tensor global scale; without it the
+                # exported base model cannot be dequantized.
+                "base_layer.weight_scale_2": "weight_scale_2",
             }
         )
         skip_keys.append("base_layer")
diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py
index f72715d410c..7e9ceca1505 100644
--- a/modelopt/torch/opt/plugins/transformers.py
+++ b/modelopt/torch/opt/plugins/transformers.py
@@ -111,15 +111,18 @@ def _restore_qtensor_wrappers(model, model_path):
         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"],
-                )
+            if not isinstance(module, RealQuantLinear) or isinstance(module.weight, QTensorWrapper):
+                continue
+            # QLoRA: adapters are loaded on top of the compressed base checkpoint, so PEFT has
+            # re-parented the quantized linear as `<name>.base_layer` while `q_tensor_state` is
+            # still keyed by the unwrapped name it was saved with.
+            key = name if name in q_tensor_state else 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"],
+            )
 
 
 def _new_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwargs):
diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py
index f87117501f5..77b7111ef7a 100644
--- a/tests/examples/llm_qat/test_llm_qat.py
+++ b/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"
@@ -219,14 +235,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)
diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py
index cac0a9a9aef..2dc2e9c664a 100644
--- a/tests/gpu/torch/export/test_export.py
+++ b/tests/gpu/torch/export/test_export.py
@@ -259,6 +259,35 @@ 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():
+    """QLoRA base weights live under `base_layer.*` and must all survive the rename.
+
+    `weight_scale_2` is the NVFP4 per-tensor global scale; dropping it makes the exported
+    base model impossible to dequantize.
+    """
+    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]),
+        # 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]),
+    }
+
+    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",
+    }
+    assert processed_state_dict["layer1.weight_scale_2"] == torch.tensor([0.5])
+
+
 @pytest.mark.parametrize(
     ("config", "expected"),
     [

From 3990f0c405eaefa84409f063afecbd1654f06d10 Mon Sep 17 00:00:00 2001
From: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Date: Tue, 4 Aug 2026 18:48:15 +0000
Subject: [PATCH 2/7] Shorten comments in QLoRA export fix

Trim the explanatory comments added in the previous commit to two lines each.
No functional change.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
---
 examples/llm_qat/export.py                 | 5 +----
 modelopt/torch/export/quant_utils.py       | 3 +--
 modelopt/torch/opt/plugins/transformers.py | 5 ++---
 tests/gpu/torch/export/test_export.py      | 3 +--
 4 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/examples/llm_qat/export.py b/examples/llm_qat/export.py
index 49201baad52..271a5ab2641 100644
--- a/examples/llm_qat/export.py
+++ b/examples/llm_qat/export.py
@@ -50,10 +50,7 @@ def get_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.
-    # For QLoRA the base model checkpoint is itself quantized, so `from_pretrained` has already
-    # restored the modelopt state (and the quantizer buffers) via `enable_huggingface_checkpointing`.
-    # Restoring a second time would raise `Model already has modelopt state!`, so only restore the
-    # state here if the loaded model does not have it yet.
+    # For QLoRA the base checkpoint is quantized, so from_pretrained already restored the state.
     modelopt_state_path = os.path.join(ckpt_path, "modelopt_state_train.pth")
     if hasattr(model, "peft_config") and os.path.isfile(modelopt_state_path):
         modelopt_state = mto.load_modelopt_state(modelopt_state_path)
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index cb153ab2e5e..9096665e68f 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -998,8 +998,7 @@ def postprocess_state_dict(
                 "base_layer.weight": "weight",
                 "base_layer.input_scale": "input_scale",
                 "base_layer.weight_scale": "weight_scale",
-                # NVFP4 double quantization keeps a per-tensor global scale; without it the
-                # exported base model cannot be dequantized.
+                # NVFP4 global scale; the exported model cannot be dequantized without it.
                 "base_layer.weight_scale_2": "weight_scale_2",
             }
         )
diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py
index 7e9ceca1505..121a183e257 100644
--- a/modelopt/torch/opt/plugins/transformers.py
+++ b/modelopt/torch/opt/plugins/transformers.py
@@ -113,9 +113,8 @@ def _restore_qtensor_wrappers(model, model_path):
         for name, module in model.named_modules():
             if not isinstance(module, RealQuantLinear) or isinstance(module.weight, QTensorWrapper):
                 continue
-            # QLoRA: adapters are loaded on top of the compressed base checkpoint, so PEFT has
-            # re-parented the quantized linear as `<name>.base_layer` while `q_tensor_state` is
-            # still keyed by the unwrapped name it was saved with.
+            # 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")
             if key not in q_tensor_state:
                 continue
diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py
index 2dc2e9c664a..2befa0b577b 100644
--- a/tests/gpu/torch/export/test_export.py
+++ b/tests/gpu/torch/export/test_export.py
@@ -262,8 +262,7 @@ def test_postprocess_state_dict(state_dict, quantization, maxbound, expected_sta
 def test_postprocess_state_dict_qlora_strips_base_layer():
     """QLoRA base weights live under `base_layer.*` and must all survive the rename.
 
-    `weight_scale_2` is the NVFP4 per-tensor global scale; dropping it makes the exported
-    base model impossible to dequantize.
+    Dropping `weight_scale_2` makes the exported base model impossible to dequantize.
     """
     state_dict = {
         "layer1.base_layer.weight": torch.ones(4, 2, dtype=torch.uint8),

From 24cf9f3ddb1e0834a94dd221297bc4b991e9633a Mon Sep 17 00:00:00 2001
From: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Date: Tue, 4 Aug 2026 19:37:10 +0000
Subject: [PATCH 3/7] Address PR review: generic base_layer strip, loud restore
 failure, unit test

- postprocess_state_dict: strip ".base_layer" generically instead of enumerating
  renames. The enumerated map silently dropped every unlisted key, which lost
  linear biases (Qwen2 q/k/v have them) and left "base_layer" in the exported AWQ
  pre_quant_scale key.
- export.py: drop the os.path.isfile guard so a LoRA checkpoint with no modelopt
  state still fails loudly instead of exporting unquantized, and pop the quantizer
  weights inside the branch that uses them.
- Add a CPU unit test for the .base_layer fallback in _restore_qtensor_wrappers,
  and cover bias / pre_quant_scale in the export rename test.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
---
 examples/llm_qat/export.py                    | 22 ++++-----
 modelopt/torch/export/quant_utils.py          | 20 +++------
 tests/gpu/torch/export/test_export.py         | 11 ++++-
 .../torch/opt/plugins/test_hf_patching.py     | 45 +++++++++++++++++++
 4 files changed, 69 insertions(+), 29 deletions(-)

diff --git a/examples/llm_qat/export.py b/examples/llm_qat/export.py
index 271a5ab2641..67f50f7b430 100644
--- a/examples/llm_qat/export.py
+++ b/examples/llm_qat/export.py
@@ -15,7 +15,6 @@
 
 import argparse
 import json
-import os
 import warnings
 from pathlib import Path
 
@@ -51,19 +50,16 @@ 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.
-    modelopt_state_path = os.path.join(ckpt_path, "modelopt_state_train.pth")
-    if hasattr(model, "peft_config") and os.path.isfile(modelopt_state_path):
-        modelopt_state = mto.load_modelopt_state(modelopt_state_path)
-        modelopt_weights = modelopt_state.pop("modelopt_state_weights", None)
-
-        if not ModeloptStateManager.is_converted(model):
-            restore_from_modelopt_state(model, modelopt_state)
-            print_rank_0("Restored modelopt 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")
 
-            # Restore modelopt quantizer state dict
-            if modelopt_weights is not None:
-                set_quantizer_state_dict(model, modelopt_weights)
-                print_rank_0("Restored modelopt quantizer state dict")
+        # Restore modelopt quantizer state dict
+        modelopt_weights = modelopt_state.pop("modelopt_state_weights", None)
+        if modelopt_weights is not None:
+            set_quantizer_state_dict(model, modelopt_weights)
+            print_rank_0("Restored modelopt quantizer state dict")
 
     return model
 
diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py
index 9096665e68f..8af87c817ec 100755
--- a/modelopt/torch/export/quant_utils.py
+++ b/modelopt/torch/export/quant_utils.py
@@ -991,18 +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",
-                # NVFP4 global scale; the exported model cannot be dequantized without it.
-                "base_layer.weight_scale_2": "weight_scale_2",
-            }
-        )
-        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 = {}
 
@@ -1014,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
@@ -1035,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
diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py
index 2befa0b577b..2b795817b49 100644
--- a/tests/gpu/torch/export/test_export.py
+++ b/tests/gpu/torch/export/test_export.py
@@ -260,18 +260,23 @@ def test_postprocess_state_dict(state_dict, quantization, maxbound, expected_sta
 
 
 def test_postprocess_state_dict_qlora_strips_base_layer():
-    """QLoRA base weights live under `base_layer.*` and must all survive the rename.
+    """Every QLoRA `base_layer.*` tensor needed for deployment must survive the rename.
 
-    Dropping `weight_scale_2` makes the exported base model impossible to dequantize.
+    `weight_scale_2` is the NVFP4 global scale and `bias` matters for architectures such as
+    Qwen2; dropping either 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(
@@ -283,6 +288,8 @@ def test_postprocess_state_dict_qlora_strips_base_layer():
         "layer1.weight_scale",
         "layer1.weight_scale_2",
         "layer1.input_scale",
+        "layer1.bias",
+        "layer1.pre_quant_scale",
     }
     assert processed_state_dict["layer1.weight_scale_2"] == torch.tensor([0.5])
 
diff --git a/tests/unit/torch/opt/plugins/test_hf_patching.py b/tests/unit/torch/opt/plugins/test_hf_patching.py
index 8a44ad23c76..328e8df1878 100644
--- a/tests/unit/torch/opt/plugins/test_hf_patching.py
+++ b/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,43 @@ 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):
+    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")
+
+    # transformers>=5 loads weights by assigning a plain Parameter holding the already-packed
+    # data, which leaves the module without its QTensorWrapper.
+    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])
+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)
+    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])

From e8d77d7a5f1a118cdb950f7796e9bf6519ff2685 Mon Sep 17 00:00:00 2001
From: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Date: Tue, 4 Aug 2026 20:33:17 +0000
Subject: [PATCH 4/7] Add export coverage to the LoRA-QAT example test

The is_converted guard also changes the fake-quant LoRA-QAT path, where the
calibrated amaxes rather than packed weights are what must survive the load.
Verified they do: the base PTQ safetensors carries all 32 _amax buffers, so
from_pretrained restores them and the export matches a direct PTQ export
byte-for-byte. Extend the test so that path stops being untested.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
---
 tests/examples/llm_qat/test_llm_qat.py | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py
index 77b7111ef7a..606b70c91cc 100644
--- a/tests/examples/llm_qat/test_llm_qat.py
+++ b/tests/examples/llm_qat/test_llm_qat.py
@@ -168,18 +168,39 @@ 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. Unlike QLoRA 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"
+
+    # 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)
+
 
 @pytest.mark.parametrize("backend", [
     "fsdp2",

From 5363203a7ea6ed91b26042d702d23daeb7deacba Mon Sep 17 00:00:00 2001
From: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Date: Tue, 4 Aug 2026 23:10:40 +0000
Subject: [PATCH 5/7] 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>
---
 CHANGELOG.rst                                 |  2 ++
 examples/llm_qat/export.py                    |  3 ++
 modelopt/torch/opt/plugins/transformers.py    | 32 +++++++++++++++----
 tests/examples/llm_qat/test_llm_qat.py        | 21 +++++++-----
 tests/gpu/torch/export/test_export.py         |  3 +-
 .../torch/opt/plugins/test_hf_patching.py     | 30 ++++++++++++++---
 6 files changed, 71 insertions(+), 20 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 7a188669cb8..bb551333520 100755
--- a/CHANGELOG.rst
+++ b/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)
 ^^^^^^^^^^^^^^^^^
 
diff --git a/examples/llm_qat/export.py b/examples/llm_qat/export.py
index 67f50f7b430..33f06d9fb4a 100644
--- a/examples/llm_qat/export.py
+++ b/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)
diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py
index 121a183e257..e0e956a3f35 100644
--- a/modelopt/torch/opt/plugins/transformers.py
+++ b/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):
diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py
index 606b70c91cc..f055c89c9c2 100644
--- a/tests/examples/llm_qat/test_llm_qat.py
+++ b/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",
diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py
index 2b795817b49..f2891b04b28 100644
--- a/tests/gpu/torch/export/test_export.py
+++ b/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(
diff --git a/tests/unit/torch/opt/plugins/test_hf_patching.py b/tests/unit/torch/opt/plugins/test_hf_patching.py
index 328e8df1878..7731ee9a862 100644
--- a/tests/unit/torch/opt/plugins/test_hf_patching.py
+++ b/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)

From 7a116e76388efaa670df49bceb5bc35b5909d91b Mon Sep 17 00:00:00 2001
From: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Date: Tue, 4 Aug 2026 23:18:30 +0000
Subject: [PATCH 6/7] Move the QLoRA export changelog entry to 0.46

The PR is labeled cherry-pick-0.46.0, so the entry belongs in the 0.46 section.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
---
 CHANGELOG.rst | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index bb551333520..d6d6f1328ae 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -21,8 +21,6 @@ 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)
 ^^^^^^^^^^^^^^^^^
 
@@ -104,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)
 ^^^^^^^^^^^^^^^^^

From 46596f07aaf9497ed8e4b9d57e5192743456c9b3 Mon Sep 17 00:00:00 2001
From: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Date: Tue, 4 Aug 2026 23:35:34 +0000
Subject: [PATCH 7/7] Shorten comments from the review-fix rounds

Trim the added comments and test docstrings to two lines each. No code change.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
---
 examples/llm_qat/export.py                       |  5 ++---
 modelopt/torch/opt/plugins/transformers.py       | 12 ++++--------
 tests/examples/llm_qat/test_llm_qat.py           |  9 ++++-----
 tests/gpu/torch/export/test_export.py            |  3 +--
 tests/unit/torch/opt/plugins/test_hf_patching.py |  6 ++----
 5 files changed, 13 insertions(+), 22 deletions(-)

diff --git a/examples/llm_qat/export.py b/examples/llm_qat/export.py
index 33f06d9fb4a..afe2bd4d1cf 100644
--- a/examples/llm_qat/export.py
+++ b/examples/llm_qat/export.py
@@ -50,9 +50,8 @@ 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.
+    # 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)
diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py
index e0e956a3f35..0052525b6e9 100644
--- a/modelopt/torch/opt/plugins/transformers.py
+++ b/modelopt/torch/opt/plugins/transformers.py
@@ -110,12 +110,8 @@ def _restore_qtensor_wrappers(model, model_path):
         q_tensor_state = mode_config.get("metadata", {}).get("q_tensor_state", {})
         if not q_tensor_state:
             continue
-        # 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.
+        # 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 = [
@@ -134,8 +130,8 @@ def _restore_qtensor_wrappers(model, model_path):
             )
             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.
+        # 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 "
diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py
index f055c89c9c2..4cd67d7f905 100644
--- a/tests/examples/llm_qat/test_llm_qat.py
+++ b/tests/examples/llm_qat/test_llm_qat.py
@@ -181,8 +181,8 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path):
         cache_dir=cache_dir,
     )
 
-    # Step 3: Export. Unlike QLoRA this checkpoint is fake-quantized, so the calibrated amaxes
-    # rather than packed weights are what must survive the load.
+    # 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))
 
@@ -193,9 +193,8 @@ def test_qwen3_lora_qat_nvfp4(tiny_qwen3_path, tmp_path):
     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 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.
+    # 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")
diff --git a/tests/gpu/torch/export/test_export.py b/tests/gpu/torch/export/test_export.py
index f2891b04b28..55137a64639 100644
--- a/tests/gpu/torch/export/test_export.py
+++ b/tests/gpu/torch/export/test_export.py
@@ -262,8 +262,7 @@ def test_postprocess_state_dict(state_dict, quantization, maxbound, expected_sta
 def test_postprocess_state_dict_qlora_strips_base_layer():
     """Every QLoRA `base_layer.*` tensor needed for deployment must survive the rename.
 
-    `weight_scale_2` is the NVFP4 global scale and `bias` matters for architectures such as
-    Qwen2; dropping either yields an undeployable checkpoint.
+    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),
diff --git a/tests/unit/torch/opt/plugins/test_hf_patching.py b/tests/unit/torch/opt/plugins/test_hf_patching.py
index 7731ee9a862..0476f44c56c 100644
--- a/tests/unit/torch/opt/plugins/test_hf_patching.py
+++ b/tests/unit/torch/opt/plugins/test_hf_patching.py
@@ -78,16 +78,14 @@ def _compressed_model_and_state_dir(tmp_path, state_keyed_with_base_layer=False)
 
     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.
+        # 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 loads weights by assigning a plain Parameter holding the already-packed
-    # data, which leaves the module without its QTensorWrapper.
+    # 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)