From af3dee7e16fad8f84d9bb6d3636ba10130214390 Mon Sep 17 00:00:00 2001
From: Chenhan Yu <chenhany@nvidia.com>
Date: Sun, 4 Oct 2026 19:35:44 +0000
Subject: [PATCH 1/3] Add mergeable LoRA QAD for Megatron Bridge

Signed-off-by: Chenhan Yu <chenhany@nvidia.com>
---
 CHANGELOG.rst                                 |   1 +
 examples/megatron_bridge/README.md            |  31 ++++
 .../export_quantized_megatron_to_hf.py        |   2 +
 examples/megatron_bridge/quantize.py          |  25 +++
 modelopt/torch/quantization/__init__.py       |   9 +
 modelopt/torch/quantization/lora.py           | 171 ++++++++++++++++++
 .../quantization/plugins/megatron_lora.py     | 113 ++++++++++++
 tests/examples/megatron_bridge/test_qad.py    |  68 ++++++-
 .../torch/quantization/test_quant_lora.py     |  68 +++++++
 .../torch/quantization/test_quant_lora.py     | 117 ++++++++++++
 10 files changed, 597 insertions(+), 8 deletions(-)
 create mode 100644 modelopt/torch/quantization/lora.py
 create mode 100644 modelopt/torch/quantization/plugins/megatron_lora.py
 create mode 100644 tests/gpu_megatron/torch/quantization/test_quant_lora.py
 create mode 100644 tests/unit/torch/quantization/test_quant_lora.py

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index adcaec96729..370f7ab8562 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -41,6 +41,7 @@ Changelog
 
 *Megatron Framework (M-LM / M-Bridge)*
 
+- Add mergeable LoRA QAD for fake-quantized dense Megatron/Transformer Engine linear layers. Enable it with ``--lora_rank`` in Megatron-Bridge PTQ, train through the normal QAD flow, and export a standalone quantized checkpoint without deployment adapters.
 - Add an end-to-end W4A4 NVFP4 PTQ and QAD tutorial for Qwen3.6-35B-A3B also covering evaluation and vLLM throughput benchmarking. See `examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/megatron_bridge/tutorials/Qwen3.6-35B-A3B/>`_ for details.
 - Add ``--mlflow <tracking-uri>`` to the ``examples/megatron_bridge`` scripts that write a checkpoint -- ``prune_minitron.py``, ``quantize.py``, ``distill.py``, ``export_quantized_megatron_to_hf.py`` and ``export_distilled_megatron_to_hf.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). Each run records the invocation, its arguments as searchable params and its log -- a distillation records its training metrics instead of the rank-0 log -- and writes ``.experiment.json`` into the checkpoint it produced, so a pruning, a quantization, the distillation that refines its checkpoint and the export that deploys it can be traced to one another; uploading the checkpoints themselves stays off unless ``--mlflow_log_checkpoints`` is passed.
 - Add ``--ep_size`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` so large MoE models with grouped-GEMM experts can be exported with their experts sharded across GPUs. Checkpoints built with ``--no_moe_grouped_gemm`` must still be exported at ``--ep_size 1``.
diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index fd7c20a499c..2463cb90029 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -327,6 +327,37 @@ torchrun --nproc_per_node 8 distill.py \
 
 The distilled checkpoint retains the ModelOpt quantization state, so it can be converted to a deployable HuggingFace checkpoint with [export_quantized_megatron_to_hf.py](export_quantized_megatron_to_hf.py) (point `--megatron_path` at `/output/qwen3_8b_nvfp4_qad/checkpoints`), exactly like the PTQ checkpoint in [step 2 above](#post-training-quantization).
 
+### LoRA QAD
+
+Add `--lora_rank` to PTQ to train low-rank weight updates with the usual QAD flow:
+
+```bash
+torchrun --nproc_per_node 2 quantize.py \
+    --hf_model_name_or_path Qwen/Qwen3-0.6B \
+    --recipe general/ptq/nvfp4_default-kv_fp8 \
+    --tp_size 2 --lora_rank 8 --lora_alpha 16 \
+    --export_megatron_path /output/qwen3_06b_nvfp4_lora
+```
+
+PTQ calibrates the quantizers before adding the adapters. The backbone is frozen; `A` starts
+random and `B` starts at zero. Each linear layer quantizes `W + (alpha / rank) B @ A` before
+its matrix multiplication. The default targets are `linear_qkv`, `linear_proj`, `linear_fc1`,
+and `linear_fc2`; select other module-name globs with `--lora_target_modules`.
+
+Run `distill.py` with this checkpoint as `--student_megatron_path`, just as for regular QAD.
+ModelOpt reconstructs the quantizers and factors before DDP and optimizer setup, then loads
+the distributed checkpoint tensors. Resuming the output directory also restores the optimizer.
+
+`export_quantized_megatron_to_hf.py` merges the factors into the backbone before exporting a
+standalone quantized checkpoint. Deployment does not require a LoRA adapter. The Python APIs
+are `modelopt.torch.quantization.enable_quant_lora` and `merge_quant_lora`.
+
+Supported targets are fake-quantized dense PyTorch and Megatron/Transformer Engine linear
+layers, including tensor-parallel row/column and fused layer-normalization linear layers.
+Grouped MoE experts and compressed quantized weights are unsupported. The backbone remains
+in floating point and the implementation forms dense weight updates, so memory savings come
+primarily from the smaller optimizer state rather than compressed backbone storage.
+
 ### Slurm Usage
 
 To run the distillation script on a Slurm cluster for multi-node training, you just need use `python` instead of `torchrun` and set the number of nodes using `#SBATCH --nodes=<num_nodes>` clause in your Slurm script.
diff --git a/examples/megatron_bridge/export_quantized_megatron_to_hf.py b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
index 7078b234b97..bb087e1c3ab 100644
--- a/examples/megatron_bridge/export_quantized_megatron_to_hf.py
+++ b/examples/megatron_bridge/export_quantized_megatron_to_hf.py
@@ -50,6 +50,7 @@
     resolve_mlflow_args,
 )
 
+import modelopt.torch.quantization as mtq
 import modelopt.torch.utils.distributed as dist
 from modelopt.torch.export import export_mcore_gpt_to_hf
 from modelopt.torch.utils import print_args, print_rank_0
@@ -178,6 +179,7 @@ def main(args: argparse.Namespace):
     print_rank_0(f"Loading quantized Megatron checkpoint from {args.megatron_path}...")
     load_modelopt_megatron_checkpoint(model, args.megatron_path)
     unwrapped_model = unwrap_model(model[0])
+    mtq.merge_quant_lora(unwrapped_model)
 
     # Extra modules (Medusa / EAGLE / MTP) only exist on the last pipeline stage. Use an all-reduce
     # MAX over all ranks (rather than a broadcast from a hard-coded source rank) so the decision is
diff --git a/examples/megatron_bridge/quantize.py b/examples/megatron_bridge/quantize.py
index 0cc0999f202..e12789cb3ef 100644
--- a/examples/megatron_bridge/quantize.py
+++ b/examples/megatron_bridge/quantize.py
@@ -203,6 +203,16 @@ def get_args() -> argparse.Namespace:
         action="store_true",
         help="Compress weights to a real low-bit representation (instead of fake quantization).",
     )
+    parser.add_argument(
+        "--lora_rank", type=int, default=0, help="Enable LoRA-QAT/QAD; 0 disables it."
+    )
+    parser.add_argument("--lora_alpha", type=float, default=16.0)
+    parser.add_argument(
+        "--lora_target_modules",
+        nargs="+",
+        default=["*.linear_qkv", "*.linear_proj", "*.linear_fc1", "*.linear_fc2"],
+        help="Glob patterns of fake-quantized linear layers to adapt.",
+    )
     # Calibration dataset arguments (matched to hf_ptq.py)
     parser.add_argument(
         "--calib_dataset_name",
@@ -251,6 +261,11 @@ def get_args() -> argparse.Namespace:
     args = parser.parse_args()
     resolve_mlflow_args(args, parser, QUANTIZE)
 
+    if args.lora_rank < 0:
+        parser.error("--lora_rank must be non-negative")
+    if args.lora_rank and args.compress:
+        parser.error("--lora_rank requires fake quantization; cannot be combined with --compress")
+
     print_args(masked_args(args))
 
     # Flipped by main() once the Megatron checkpoint is on disk. The MLflow provenance
@@ -436,6 +451,16 @@ def forward_loop(_model=None):
     else:
         mtq.quantize(unwrapped_model, mtq_config, forward_loop)
 
+    if args.lora_rank:
+        mtq.enable_quant_lora(
+            unwrapped_model,
+            {
+                "rank": args.lora_rank,
+                "alpha": args.lora_alpha,
+                "target_modules": args.lora_target_modules,
+            },
+        )
+
     # Free calibration/quantization memory before generate
     gc.collect()
     torch.cuda.empty_cache()
diff --git a/modelopt/torch/quantization/__init__.py b/modelopt/torch/quantization/__init__.py
index 29f42e8f72b..ec0acafd06b 100644
--- a/modelopt/torch/quantization/__init__.py
+++ b/modelopt/torch/quantization/__init__.py
@@ -17,6 +17,8 @@
 
 from importlib import import_module as _import_module
 
+from modelopt.torch.utils import import_plugin as _import_plugin
+
 # Initialize mode and plugins
 from . import mode, plugins, utils
 
@@ -24,10 +26,17 @@
 from .compress import *
 from .config import *
 from .conversion import *
+from .lora import *
 from .model_quant import *
 from .nn.modules.quant_module import QuantModuleRegistry
 from .utils import update_quant_cfg_with_kv_cache_quant
 
+# LoRA modes depend on the quantization mode registry, which framework plugin
+# initialization can reach before mode.py finishes importing.
+with _import_plugin("megatron LoRA"):
+    _import_module(".plugins.megatron_lora", __name__)
+del _import_plugin
+
 # Imported last to register the backend without cycling through quantization.qtensor.
 # A dynamic import prevents isort from hoisting it into the import block above.
 ggml = _import_module(".ggml", __name__)
diff --git a/modelopt/torch/quantization/lora.py b/modelopt/torch/quantization/lora.py
new file mode 100644
index 00000000000..beea775dfbe
--- /dev/null
+++ b/modelopt/torch/quantization/lora.py
@@ -0,0 +1,171 @@
+# 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.
+
+"""Mergeable low-rank adaptation of fake-quantized linear weights."""
+
+import fnmatch
+import math
+
+import torch
+from torch import nn
+
+from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField
+from modelopt.torch.opt.conversion import apply_mode
+from modelopt.torch.opt.dynamic import DynamicModule, _DMRegistryCls
+from modelopt.torch.opt.mode import ModeDescriptor
+
+from .config import QuantizeConfig
+from .conversion import restore_quantizer_state, update_quantize_metadata
+from .mode import QuantizeModeRegistry
+from .nn import QuantModule, QuantModuleRegistry
+from .qtensor import QTensorWrapper
+
+__all__ = ["QuantLoRAConfig", "enable_quant_lora", "merge_quant_lora"]
+
+
+class QuantLoRAConfig(ModeloptBaseConfig):
+    """Configuration for adapters applied to combined weights before fake quantization."""
+
+    rank: int = ModeloptField(8, gt=0)
+    alpha: float = ModeloptField(16.0, gt=0, allow_inf_nan=False)
+    target_modules: list[str] = ModeloptField(["*"], min_length=1)
+
+
+QuantLoRARegistry = _DMRegistryCls("quant_lora")
+
+
+@QuantLoRARegistry.register({QuantModuleRegistry[nn.Linear]: "nn.Linear"})
+class _QuantLoRALinear(DynamicModule):
+    def _setup(self, config: QuantLoRAConfig):
+        weight = self._parameters["weight"]
+        self.lora_A = nn.Parameter(weight.new_empty(config.rank, weight.shape[1]))
+        self.lora_B = nn.Parameter(weight.new_zeros(weight.shape[0], config.rank))
+        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+        self.lora_scale = config.alpha / config.rank
+
+    @staticmethod
+    def _effective_weight(module, weight):
+        return weight + (module.lora_scale * module.lora_B @ module.lora_A).to(weight.dtype)
+
+    def forward(self, *args, **kwargs):
+        # Keep the combined weight outside quantizers and TE custom autograd functions.
+        weight = self._parameters["weight"]
+        self._parameters["weight"] = self._effective_weight(self, weight)
+        try:
+            return super().forward(*args, **kwargs)
+        finally:
+            self._parameters["weight"] = weight
+
+    @torch.no_grad()
+    def merge_lora(self):
+        weight = self._parameters["weight"]
+        weight.copy_(self._effective_weight(self, weight))
+        self.export()
+        del self.lora_A, self.lora_B, self.lora_scale
+
+
+def enable_quant_lora(model: nn.Module, config: dict | QuantLoRAConfig | None = None) -> nn.Module:
+    """Freeze a fake-quantized backbone and add zero-initialized, mergeable adapters.
+
+    The forward quantizes ``W + (alpha / rank) * B @ A``. Save and restore using
+    ModelOpt checkpoint APIs to reconstruct adapters before loading tensor/optimizer state.
+    """
+    if any(isinstance(module, _QuantLoRALinear) for module in model.modules()):
+        raise ValueError("Quantization-aware LoRA is already enabled.")
+    return apply_mode(model, [("quant_lora", config or {})], registry=QuantizeModeRegistry)
+
+
+def merge_quant_lora(model: nn.Module) -> nn.Module:
+    """Merge adapters into floating-point weights, retaining quantizers for deployment export."""
+    if not any(isinstance(module, _QuantLoRALinear) for module in model.modules()):
+        return model
+    return apply_mode(model, "quant_lora_export", registry=QuantizeModeRegistry)
+
+
+def _convert_quant_lora(model, config):
+    targets = []
+    for name, module in model.named_modules():
+        if not isinstance(module, QuantModule) or not hasattr(module, "weight_quantizer"):
+            continue
+        if not any(fnmatch.fnmatchcase(name, pattern) for pattern in config.target_modules):
+            continue
+        weight = module._parameters.get("weight")
+        if isinstance(weight, QTensorWrapper):
+            raise ValueError(
+                "Quantization-aware LoRA requires uncompressed floating-point weights."
+            )
+        if weight is None or weight.ndim != 2 or type(module) not in QuantLoRARegistry:
+            raise ValueError(
+                f"Quantization-aware LoRA does not support layer {name}: {type(module)}"
+            )
+        targets.append(module)
+    if not targets:
+        raise ValueError("No supported fake-quantized linear layers match target_modules.")
+    model.requires_grad_(False)
+    for module in targets:
+        QuantLoRARegistry[type(module)].convert(module, config=config)
+    metadata = {}
+    _update_quant_lora(model, config, metadata)
+    return model, metadata
+
+
+def _restore_quant_lora(model, config, metadata):
+    restore_quantizer_state(model, QuantizeConfig(), metadata)
+    return _convert_quant_lora(model, config)[0]
+
+
+def _update_quant_lora(model, config, metadata):
+    update_quantize_metadata(model, QuantizeConfig(), metadata)
+
+
+def _merge_quant_lora(model, config):
+    for module in list(model.modules()):
+        if isinstance(module, _QuantLoRALinear):
+            module.merge_lora()
+    metadata = {}
+    _update_quant_lora(model, config, metadata)
+    return model, metadata
+
+
+def _restore_merged_quant_lora(model, config, metadata):
+    model = _merge_quant_lora(model, config)[0]
+    return restore_quantizer_state(model, QuantizeConfig(), metadata)
+
+
+@QuantizeModeRegistry.register_mode
+class QuantLoRAModeDescriptor(ModeDescriptor):
+    """Record adapters after quantization and reconstruct them before checkpoint loading."""
+
+    name = "quant_lora"
+    config_class = QuantLoRAConfig
+    export_mode = "quant_lora_export"
+    next_prohibited_modes = {"quant_lora", "quantize", "auto_quantize", "real_quantize"}
+    convert = staticmethod(_convert_quant_lora)
+    restore = staticmethod(_restore_quant_lora)
+    update_for_save = staticmethod(_update_quant_lora)
+    update_for_new_mode = staticmethod(_update_quant_lora)
+
+
+@QuantizeModeRegistry.register_mode
+class QuantLoRAExportModeDescriptor(ModeDescriptor):
+    """Record merged adapters so exported checkpoints restore without adapter parameters."""
+
+    name = "quant_lora_export"
+    config_class = ModeloptBaseConfig
+    is_export_mode = True
+    convert = staticmethod(_merge_quant_lora)
+    restore = staticmethod(_restore_merged_quant_lora)
+    update_for_save = staticmethod(_update_quant_lora)
+    update_for_new_mode = staticmethod(_update_quant_lora)
diff --git a/modelopt/torch/quantization/plugins/megatron_lora.py b/modelopt/torch/quantization/plugins/megatron_lora.py
new file mode 100644
index 00000000000..9a7271fd00b
--- /dev/null
+++ b/modelopt/torch/quantization/plugins/megatron_lora.py
@@ -0,0 +1,113 @@
+# 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.
+
+"""Tensor-parallel adapters and distributed checkpointing for Megatron LoRA-QAT/QAD."""
+
+import torch
+from megatron.core.extensions.transformer_engine import (
+    TEColumnParallelLinear,
+    TELayerNormColumnParallelLinear,
+    TERowParallelLinear,
+)
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.tensor_parallel.mappings import copy_to_tensor_model_parallel_region
+from megatron.core.tensor_parallel.random import get_cuda_rng_tracker
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
+
+from ..lora import QuantLoRARegistry, _QuantLoRALinear
+from ..nn import QuantModuleRegistry
+from .megatron import _QuantMegatronMLP
+
+__all__ = []
+
+
+@QuantLoRARegistry.register(
+    {
+        QuantModuleRegistry[ColumnParallelLinear]: "ColumnParallelLinear",
+        QuantModuleRegistry[RowParallelLinear]: "RowParallelLinear",
+    }
+)
+class _MegatronQuantLoRALinear(_QuantLoRALinear):
+    def _setup(self, config):
+        super()._setup(config)
+        base_weight = self._parameters["weight"]
+        self.lora_tp_group = self.parallel_state.tensor_parallel_group.group
+        for name, parameter in (("lora_A", self.lora_A), ("lora_B", self.lora_B)):
+            is_sharded = (name == "lora_A" and self._is_row_parallel) or (
+                name == "lora_B" and self._is_column_parallel
+            )
+            parameter.tensor_model_parallel = is_sharded
+            parameter.partition_dim = (1 if name == "lora_A" else 0) if is_sharded else -1
+            parameter.partition_stride = getattr(base_weight, "partition_stride", 1)
+            parameter.allreduce = getattr(base_weight, "allreduce", True)
+            parameter.sequence_parallel = False
+        if torch.distributed.get_world_size(self.lora_tp_group) > 1:
+            if self._is_row_parallel and self.lora_A.is_cuda:
+                with get_cuda_rng_tracker().fork():
+                    torch.nn.init.kaiming_uniform_(self.lora_A, a=5**0.5)
+            if not self._is_row_parallel:
+                torch.distributed.broadcast(
+                    self.lora_A.data,
+                    src=torch.distributed.get_global_rank(self.lora_tp_group, 0),
+                    group=self.lora_tp_group,
+                )
+
+    @staticmethod
+    def _effective_weight(module, weight):
+        a, b = module.lora_A, module.lora_B
+        if torch.distributed.get_world_size(module.lora_tp_group) > 1:
+            if module._is_column_parallel:
+                a = copy_to_tensor_model_parallel_region(a, group=module.lora_tp_group)
+            elif module._is_row_parallel:
+                b = copy_to_tensor_model_parallel_region(b, group=module.lora_tp_group)
+        return weight + (module.lora_scale * b @ a).to(weight.dtype)
+
+    def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
+        state = super().sharded_state_dict(prefix, sharded_offsets, metadata)
+        axes = {}
+        if self._is_column_parallel:
+            axes["lora_B"] = 0
+        elif self._is_row_parallel:
+            axes["lora_A"] = 1
+        state.update(
+            make_sharded_tensors_for_checkpoint(
+                {"lora_A": self.lora_A, "lora_B": self.lora_B},
+                prefix,
+                axes,
+                sharded_offsets,
+                tp_group=self.lora_tp_group,
+                dp_cp_group=self.parallel_state.data_parallel_group.group,
+            )
+        )
+        return state
+
+    def merge_lora(self):
+        super().merge_lora()
+        del self.lora_tp_group
+
+
+_QuantMegatronMLP._modelopt_state_keys = [
+    *_QuantMegatronMLP._modelopt_state_keys,
+    r".*linear_fc1\.lora_B$",
+]
+
+
+QuantLoRARegistry.register(
+    {
+        QuantModuleRegistry[TEColumnParallelLinear]: "TEColumnParallelLinear",
+        QuantModuleRegistry[TELayerNormColumnParallelLinear]: "TELayerNormColumnParallelLinear",
+        QuantModuleRegistry[TERowParallelLinear]: "TERowParallelLinear",
+    }
+)(_MegatronQuantLoRALinear)
diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py
index b54dffa16d1..98ed6ad0de2 100644
--- a/tests/examples/megatron_bridge/test_qad.py
+++ b/tests/examples/megatron_bridge/test_qad.py
@@ -19,6 +19,8 @@
 from pathlib import Path
 
 import pytest
+import torch
+import torch.distributed.checkpoint as dcp
 from _test_utils.examples.run_command import extend_cmd_parts, run_example_command
 from _test_utils.torch.export.unified_checkpoint import (
     assert_exported_checkpoint_matches,
@@ -32,6 +34,7 @@
     create_tiny_qwen3_5_moe_vl_dir,
     create_tiny_qwen3_dir,
 )
+from torch.distributed.checkpoint import FileSystemReader
 
 # The fix ships in the nemo:26.10 container.
 # TODO(Megatron-Bridge#6243): drop this probe once the minimum Megatron-Bridge carries it.
@@ -45,9 +48,18 @@
 
 @pytest.mark.timeout(720)  # Multiple steps in one test hence takes longer than the default timeout
 @pytest.mark.parametrize(
-    "create_student",
+    ("create_student", "lora_rank"),
     [
-        lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True),
+        (lambda tmp_path: create_tiny_qwen3_dir(tmp_path, with_tokenizer=True), 0),
+        (
+            lambda tmp_path: create_tiny_qwen3_dir(
+                tmp_path,
+                with_tokenizer=True,
+                hidden_size=128,
+                intermediate_size=256,
+            ),
+            4,
+        ),
         pytest.param(
             lambda tmp_path: create_tiny_qwen3_5_moe_vl_dir(
                 tmp_path,
@@ -56,11 +68,12 @@
                 num_hidden_layers=2,
                 layer_types=["linear_attention", "full_attention"],
             ),
+            0,
         ),
     ],
-    ids=["qwen3", "qwen3_5_moe_vl"],
+    ids=["qwen3", "qwen3_lora_nvfp4", "qwen3_5_moe_vl"],
 )
-def test_qad(tmp_path: Path, num_gpus, create_student):
+def test_qad(tmp_path: Path, num_gpus, create_student, lora_rank):
     """Quantize a tiny model, run QAD from the quantized student, and export the result.
 
     Covers what only QAD exercises: that the ModelOpt state survives distillation. Per-architecture
@@ -83,16 +96,33 @@ def test_qad(tmp_path: Path, num_gpus, create_student):
     distill_output_dir = tmp_path / "qad_output"
     train_iters = 3
     early_exit_iter = 2
+    calib_dataset = "cnn_dailymail"
+    if lora_rank:
+        # Keep the LoRA smoke test independent of dataset downloads on the GPU node.
+        calib_path = tmp_path / "calibration.jsonl"
+        calib_path.write_text(
+            "\n".join(
+                json.dumps(
+                    {"text": f"Sample {i}: " + "The quick brown fox jumps over the lazy dog. " * 4}
+                )
+                for i in range(8)
+            ),
+            encoding="utf-8",
+        )
+        calib_dataset = str(calib_path)
 
-    # Step 1: PTQ the (language) model to FP8 and save a Megatron checkpoint carrying the ModelOpt state.
+    # Step 1: PTQ the model and save its quantizers and optional adapters.
     quantize_cmd = extend_cmd_parts(
         # QAD below must load this checkpoint at the same TP, so size the PTQ run to tp_size.
         ["torchrun", f"--nproc_per_node={tp_size}", "quantize.py", "--skip_generate"],
         hf_model_name_or_path=hf_model_path,
-        recipe="general/ptq/fp8_default-kv_fp8",
+        recipe="general/ptq/nvfp4_default-kv_fp8"
+        if lora_rank
+        else "general/ptq/fp8_default-kv_fp8",
+        lora_rank=lora_rank,
         tp_size=tp_size,
         pp_size=1,
-        calib_dataset_name="cnn_dailymail",  # text dataset -> (for VLMs) text-only LM calibration
+        calib_dataset_name=calib_dataset,  # text dataset -> (for VLMs) text-only LM calibration
         calib_num_samples=8,
         calib_batch_size=2,
         seq_length=16,
@@ -135,9 +165,25 @@ def test_qad(tmp_path: Path, num_gpus, create_student):
     assert (distilled_megatron_path / "iter_0000001").is_dir()
     assert_has_modelopt_state(distilled_megatron_path)
 
+    if lora_rank:
+        # Reuse the output directory to exercise optimizer and trained-adapter restoration.
+        run_example_command(distill_cmd, example_path="megatron_bridge", setup_free_port=True)
+        assert tracker.read_text(encoding="utf-8").strip() == str(train_iters)
+        checkpoint_dir = distilled_megatron_path / f"iter_{train_iters:07d}"
+        reader = FileSystemReader(str(checkpoint_dir))
+        metadata = reader.read_metadata().state_dict_metadata
+        adapters = {
+            name: torch.empty(meta.size, dtype=meta.properties.dtype)
+            for name, meta in metadata.items()
+            if name.endswith("lora_B")
+        }
+        assert adapters, "Trained checkpoint lost the LoRA factors"
+        dcp.load(adapters, storage_reader=reader)
+        assert any(value.count_nonzero() > 0 for value in adapters.values())
+
     # Step 3: export the distilled quantized checkpoint to a unified HF checkpoint. hf_quant_config.json
     # is only written for a quantized model, so its presence confirms the quantizers survived QAD.
-    hf_export_path = tmp_path / "qad_fp8_hf"
+    hf_export_path = tmp_path / "qad_hf"
     export_cmd = extend_cmd_parts(
         [
             "torchrun",
@@ -152,6 +198,12 @@ def test_qad(tmp_path: Path, num_gpus, create_student):
     run_example_command(export_cmd, example_path="megatron_bridge", setup_free_port=True)
     assert (hf_export_path / "config.json").exists()
     assert (hf_export_path / "hf_quant_config.json").exists()
+    if lora_rank:
+        assert not (hf_export_path / "adapter_config.json").exists()
+        index = json.loads((hf_export_path / "model.safetensors.index.json").read_text())
+        assert not any("lora_" in key for key in index["weight_map"])
+        quant_config = json.loads((hf_export_path / "hf_quant_config.json").read_text())
+        assert quant_config["quantization"]["quant_algo"] == "NVFP4"
     # A quantized export writes routed experts one per expert while the BF16 reference packs
     # them, so both sides of that expansion differ from the reference.
     text_config = json.loads((hf_model_path / "config.json").read_text())
diff --git a/tests/gpu_megatron/torch/quantization/test_quant_lora.py b/tests/gpu_megatron/torch/quantization/test_quant_lora.py
new file mode 100644
index 00000000000..8b2d3a6d2ea
--- /dev/null
+++ b/tests/gpu_megatron/torch/quantization/test_quant_lora.py
@@ -0,0 +1,68 @@
+# 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.
+
+"""Megatron tensor-parallel LoRA gradients and distributed checkpoint restoration."""
+
+from functools import partial
+
+import pytest
+import torch
+from _test_utils.torch.megatron.models import MegatronModel
+from _test_utils.torch.megatron.utils import initialize_for_megatron, sharded_state_dict_test_helper
+from torch.nn import functional as F
+
+import modelopt.torch.quantization as mtq
+
+
+def _test_quant_lora(checkpoint_dir, rank, size):
+    initialize_for_megatron(tensor_model_parallel_size=size)
+    teacher = MegatronModel(tp_size=size).cuda()
+    student = MegatronModel(tp_size=size).cuda()
+    student.load_state_dict(teacher.state_dict())
+    inputs = student.get_dummy_input(seed=123).cuda()
+    with torch.no_grad():
+        targets = teacher(inputs).softmax(dim=-1)
+    mtq.quantize(student, mtq.INT8_DEFAULT_CFG, lambda model: model(inputs))
+    baseline = student(inputs).detach()
+    weights = {name: p.detach().clone() for name, p in student.named_parameters()}
+    mtq.enable_quant_lora(student, {"rank": 4})
+    torch.testing.assert_close(student(inputs), baseline, rtol=0, atol=0)
+    optimizer = torch.optim.Adam([p for p in student.parameters() if p.requires_grad], lr=0.01)
+    for _ in range(2):
+        optimizer.zero_grad()
+        loss = F.kl_div(student(inputs).log_softmax(dim=-1), targets, reduction="batchmean")
+        assert torch.isfinite(loss)
+        loss.backward()
+        for parameter in (student.fc1.lora_A, student.fc2.lora_B):
+            replicas = [torch.empty_like(parameter.grad) for _ in range(size)]
+            torch.distributed.all_gather(replicas, parameter.grad)
+            for replica in replicas:
+                torch.testing.assert_close(replica, parameter.grad)
+        optimizer.step()
+    assert student.fc1.lora_A.grad.count_nonzero() > 0
+    for name, value in weights.items():
+        torch.testing.assert_close(dict(student.named_parameters())[name], value, rtol=0, atol=0)
+    restored = MegatronModel(tp_size=size).cuda()
+    sharded_state_dict_test_helper(checkpoint_dir, student, restored, lambda m: m(inputs))
+    assert all(p.requires_grad == ("lora_" in name) for name, p in restored.named_parameters())
+    expected = restored(inputs).detach()
+    mtq.merge_quant_lora(restored)
+    torch.testing.assert_close(restored(inputs), expected)
+
+
+@pytest.mark.parametrize("tp_size", [1, 2])
+def test_quant_lora_distributed_checkpoint(tmp_path, request, tp_size):
+    workers = request.getfixturevalue(f"dist_workers_size_{tp_size}")
+    workers.run(partial(_test_quant_lora, str(tmp_path)))
diff --git a/tests/unit/torch/quantization/test_quant_lora.py b/tests/unit/torch/quantization/test_quant_lora.py
new file mode 100644
index 00000000000..f331955fdd8
--- /dev/null
+++ b/tests/unit/torch/quantization/test_quant_lora.py
@@ -0,0 +1,117 @@
+# 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.
+
+"""Combined-weight quantization, adapter training, and ModelOpt checkpoint round trips."""
+
+import copy
+import io
+
+import pytest
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+import modelopt.torch.opt as mto
+import modelopt.torch.quantization as mtq
+
+
+def _model():
+    return nn.Sequential(nn.Linear(16, 32), nn.GELU(), nn.Linear(32, 16))
+
+
+def _optimizer(model):
+    return torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=0.01)
+
+
+def _step(model, teacher, optimizer, inputs):
+    optimizer.zero_grad()
+    with torch.no_grad():
+        target = teacher(inputs).softmax(dim=-1)
+    loss = F.kl_div(model(inputs).log_softmax(dim=-1), target, reduction="batchmean")
+    assert torch.isfinite(loss)
+    loss.backward()
+    optimizer.step()
+
+
+def test_quant_lora_training_restore_and_merge():
+    torch.manual_seed(42)
+    teacher = _model()
+    student = copy.deepcopy(teacher)
+    inputs = torch.randn(4, 16)
+    mtq.quantize(student, mtq.INT8_DEFAULT_CFG, lambda m: m(inputs))
+    baseline = student(inputs).detach()
+    backbone = {name: p.detach().clone() for name, p in student.named_parameters()}
+    mtq.enable_quant_lora(student, {"rank": 4, "alpha": 8})
+    torch.testing.assert_close(student(inputs), baseline, rtol=0, atol=0)
+    assert all(p.requires_grad == ("lora_" in name) for name, p in student.named_parameters())
+    optimizer = _optimizer(student)
+    for _ in range(2):
+        _step(student, teacher, optimizer, inputs)
+    assert student[0].lora_A.grad.count_nonzero() > 0
+    assert student[0].lora_B.count_nonzero() > 0
+    for name, value in backbone.items():
+        torch.testing.assert_close(dict(student.named_parameters())[name], value, rtol=0, atol=0)
+
+    restored = mto.restore_from_modelopt_state(_model(), mto.modelopt_state(student))
+    restored_optimizer = _optimizer(restored)
+    restored.load_state_dict(student.state_dict())
+    restored_optimizer.load_state_dict(copy.deepcopy(optimizer.state_dict()))
+    assert all(p.requires_grad == ("lora_" in name) for name, p in restored.named_parameters())
+    for model, opt in ((student, optimizer), (restored, restored_optimizer)):
+        _step(model, teacher, opt, inputs)
+    torch.testing.assert_close(restored(inputs), student(inputs), rtol=0, atol=0)
+
+    # A nonzero update must be quantized jointly, rather than added after quantization.
+    layer = restored[0]
+    with torch.no_grad():
+        layer.lora_B.normal_()
+    weight = layer._parameters["weight"]
+    combined = weight + layer.lora_scale * layer.lora_B @ layer.lora_A
+    expected = layer.output_quantizer(
+        F.linear(layer.input_quantizer(inputs), layer.weight_quantizer(combined), layer.bias)
+    )
+    torch.testing.assert_close(layer(inputs), expected, rtol=0, atol=0)
+    separate = F.linear(layer.input_quantizer(inputs), layer.weight_quantizer(weight), layer.bias)
+    separate += F.linear(layer.input_quantizer(inputs), combined - weight)
+    assert not torch.allclose(layer(inputs), separate)
+
+    output = restored(inputs).detach()
+    mtq.merge_quant_lora(restored)
+    assert not any("lora_" in name for name, _ in restored.named_parameters())
+    torch.testing.assert_close(restored(inputs), output, rtol=0, atol=0)
+    checkpoint = io.BytesIO()
+    mto.save(restored, checkpoint)
+    checkpoint.seek(0)
+    merged_restored = mto.restore(_model(), checkpoint)
+    torch.testing.assert_close(merged_restored(inputs), output, rtol=0, atol=0)
+    assert not any("lora_" in name for name, _ in merged_restored.named_parameters())
+
+
+@pytest.mark.parametrize("targets", [["0"], ["2"]])
+def test_quant_lora_targets(targets):
+    student = _model()
+    inputs = torch.randn(2, 16)
+    mtq.quantize(student, mtq.INT8_DEFAULT_CFG, lambda m: m(inputs))
+    mtq.enable_quant_lora(student, {"target_modules": targets})
+    assert [name for name, m in student.named_modules() if hasattr(m, "lora_A")] == targets
+    before = {name: p.requires_grad for name, p in student.named_parameters()}
+    with pytest.raises(ValueError, match="already enabled"):
+        mtq.enable_quant_lora(student, {"target_modules": targets})
+    assert before == {name: p.requires_grad for name, p in student.named_parameters()}
+
+
+def test_quant_lora_rejects_unquantized_model():
+    with pytest.raises(ValueError, match="No supported fake-quantized"):
+        mtq.enable_quant_lora(_model())

From 6c4628fff4032b7b13a9dc1f44f73f137d158116 Mon Sep 17 00:00:00 2001
From: Chenhan Yu <chenhany@nvidia.com>
Date: Tue, 6 Oct 2026 16:24:30 +0000
Subject: [PATCH 2/3] Strengthen LoRA QAD rejection and checkpoint resume
 coverage

Signed-off-by: Chenhan Yu <chenhany@nvidia.com>
---
 modelopt/torch/quantization/lora.py           |  9 ++++
 .../quantization/plugins/megatron_lora.py     |  6 +++
 tests/examples/megatron_bridge/test_qad.py    | 42 ++++++++++++-----
 .../torch/quantization/test_quant_lora.py     | 47 +++++++++++++++++--
 4 files changed, 90 insertions(+), 14 deletions(-)

diff --git a/modelopt/torch/quantization/lora.py b/modelopt/torch/quantization/lora.py
index beea775dfbe..5eaad2daaaa 100644
--- a/modelopt/torch/quantization/lora.py
+++ b/modelopt/torch/quantization/lora.py
@@ -49,6 +49,7 @@ class QuantLoRAConfig(ModeloptBaseConfig):
 @QuantLoRARegistry.register({QuantModuleRegistry[nn.Linear]: "nn.Linear"})
 class _QuantLoRALinear(DynamicModule):
     def _setup(self, config: QuantLoRAConfig):
+        """Initialize factors with a zero effective update."""
         weight = self._parameters["weight"]
         self.lora_A = nn.Parameter(weight.new_empty(config.rank, weight.shape[1]))
         self.lora_B = nn.Parameter(weight.new_zeros(weight.shape[0], config.rank))
@@ -57,9 +58,11 @@ def _setup(self, config: QuantLoRAConfig):
 
     @staticmethod
     def _effective_weight(module, weight):
+        """Combine the frozen weight and trainable low-rank update."""
         return weight + (module.lora_scale * module.lora_B @ module.lora_A).to(weight.dtype)
 
     def forward(self, *args, **kwargs):
+        """Run the quantized base layer on the combined weight."""
         # Keep the combined weight outside quantizers and TE custom autograd functions.
         weight = self._parameters["weight"]
         self._parameters["weight"] = self._effective_weight(self, weight)
@@ -70,6 +73,7 @@ def forward(self, *args, **kwargs):
 
     @torch.no_grad()
     def merge_lora(self):
+        """Fold the update into the base weight and remove the factors."""
         weight = self._parameters["weight"]
         weight.copy_(self._effective_weight(self, weight))
         self.export()
@@ -95,6 +99,7 @@ def merge_quant_lora(model: nn.Module) -> nn.Module:
 
 
 def _convert_quant_lora(model, config):
+    """Validate all targets before freezing the backbone and adding factors."""
     targets = []
     for name, module in model.named_modules():
         if not isinstance(module, QuantModule) or not hasattr(module, "weight_quantizer"):
@@ -122,15 +127,18 @@ def _convert_quant_lora(model, config):
 
 
 def _restore_quant_lora(model, config, metadata):
+    """Restore quantizers and recreate factors before loading checkpoint tensors."""
     restore_quantizer_state(model, QuantizeConfig(), metadata)
     return _convert_quant_lora(model, config)[0]
 
 
 def _update_quant_lora(model, config, metadata):
+    """Record quantizer metadata for adapter and merged checkpoints."""
     update_quantize_metadata(model, QuantizeConfig(), metadata)
 
 
 def _merge_quant_lora(model, config):
+    """Merge every adapter and record the remaining quantizers."""
     for module in list(model.modules()):
         if isinstance(module, _QuantLoRALinear):
             module.merge_lora()
@@ -140,6 +148,7 @@ def _merge_quant_lora(model, config):
 
 
 def _restore_merged_quant_lora(model, config, metadata):
+    """Remove reconstructed factors before restoring the merged checkpoint state."""
     model = _merge_quant_lora(model, config)[0]
     return restore_quantizer_state(model, QuantizeConfig(), metadata)
 
diff --git a/modelopt/torch/quantization/plugins/megatron_lora.py b/modelopt/torch/quantization/plugins/megatron_lora.py
index 9a7271fd00b..0648b4b917a 100644
--- a/modelopt/torch/quantization/plugins/megatron_lora.py
+++ b/modelopt/torch/quantization/plugins/megatron_lora.py
@@ -40,7 +40,10 @@
     }
 )
 class _MegatronQuantLoRALinear(_QuantLoRALinear):
+    """Adapt local weight shards with synchronized replicated-factor gradients."""
+
     def _setup(self, config):
+        """Initialize sharded factors and synchronize replicated factors."""
         super()._setup(config)
         base_weight = self._parameters["weight"]
         self.lora_tp_group = self.parallel_state.tensor_parallel_group.group
@@ -66,6 +69,7 @@ def _setup(self, config):
 
     @staticmethod
     def _effective_weight(module, weight):
+        """Combine local weight shards with TP-correct adapter gradients."""
         a, b = module.lora_A, module.lora_B
         if torch.distributed.get_world_size(module.lora_tp_group) > 1:
             if module._is_column_parallel:
@@ -75,6 +79,7 @@ def _effective_weight(module, weight):
         return weight + (module.lora_scale * b @ a).to(weight.dtype)
 
     def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
+        """Include adapter factors with their tensor-parallel checkpoint axes."""
         state = super().sharded_state_dict(prefix, sharded_offsets, metadata)
         axes = {}
         if self._is_column_parallel:
@@ -94,6 +99,7 @@ def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
         return state
 
     def merge_lora(self):
+        """Merge local factors and remove adapter-specific parallel state."""
         super().merge_lora()
         del self.lora_tp_group
 
diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py
index fa2bebdd4ec..4d87036bd76 100644
--- a/tests/examples/megatron_bridge/test_qad.py
+++ b/tests/examples/megatron_bridge/test_qad.py
@@ -169,17 +169,37 @@ def test_qad(tmp_path: Path, num_gpus, create_student, lora_rank):
         # Reuse the output directory to exercise optimizer and trained-adapter restoration.
         run_example_command(distill_cmd, example_path="megatron_bridge", setup_free_port=True)
         assert tracker.read_text(encoding="utf-8").strip() == str(train_iters)
-        checkpoint_dir = distilled_megatron_path / f"iter_{train_iters:07d}"
-        reader = FileSystemReader(str(checkpoint_dir))
-        metadata = reader.read_metadata().state_dict_metadata
-        adapters = {
-            name: torch.empty(meta.size, dtype=meta.properties.dtype)
-            for name, meta in metadata.items()
-            if name.endswith("lora_B")
-        }
-        assert adapters, "Trained checkpoint lost the LoRA factors"
-        dcp.load(adapters, storage_reader=reader)
-        assert any(value.count_nonzero() > 0 for value in adapters.values())
+        # Compare against the same seed, data, schedule, and PTQ checkpoint without a restart.
+        reference_dir = tmp_path / "qad_uninterrupted"
+        reference_cmd = distill_cmd.copy()
+        reference_cmd[reference_cmd.index("--output_dir") + 1] = str(reference_dir)
+        reference_cmd[reference_cmd.index("--exit_interval") + 1] = str(train_iters)
+        run_example_command(reference_cmd, example_path="megatron_bridge", setup_free_port=True)
+        reference_checkpoint = reference_dir / "checkpoints"
+        assert (
+            reference_checkpoint / "latest_checkpointed_iteration.txt"
+        ).read_text().strip() == str(train_iters)
+        adapter_states = []
+        for checkpoint_path in (distilled_megatron_path, reference_checkpoint):
+            reader = FileSystemReader(str(checkpoint_path / f"iter_{train_iters:07d}"))
+            metadata = reader.read_metadata().state_dict_metadata
+            adapters = {
+                name: torch.empty(meta.size, dtype=meta.properties.dtype)
+                for name, meta in metadata.items()
+                if name.endswith(("lora_A", "lora_B"))
+            }
+            assert adapters, "Trained checkpoint lost the LoRA factors"
+            dcp.load(adapters, storage_reader=reader)
+            assert any(
+                value.count_nonzero() > 0
+                for name, value in adapters.items()
+                if name.endswith("lora_B")
+            )
+            adapter_states.append(adapters)
+        resumed, uninterrupted = adapter_states
+        assert resumed.keys() == uninterrupted.keys()
+        for name, value in resumed.items():
+            torch.testing.assert_close(value, uninterrupted[name], rtol=0, atol=0)
 
     # Step 3: export the distilled quantized checkpoint to a unified HF checkpoint. hf_quant_config.json
     # is only written for a quantized model, so its presence confirms the quantizers survived QAD.
diff --git a/tests/unit/torch/quantization/test_quant_lora.py b/tests/unit/torch/quantization/test_quant_lora.py
index f331955fdd8..5e8d0b0ed74 100644
--- a/tests/unit/torch/quantization/test_quant_lora.py
+++ b/tests/unit/torch/quantization/test_quant_lora.py
@@ -25,17 +25,21 @@
 
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.qtensor import INT4QTensor, QTensorWrapper
 
 
 def _model():
+    """Build an offline dense student with two eligible layers."""
     return nn.Sequential(nn.Linear(16, 32), nn.GELU(), nn.Linear(32, 16))
 
 
 def _optimizer(model):
+    """Optimize only the trainable parameters."""
     return torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=0.01)
 
 
 def _step(model, teacher, optimizer, inputs):
+    """Take one distillation step against a frozen teacher."""
     optimizer.zero_grad()
     with torch.no_grad():
         target = teacher(inputs).softmax(dim=-1)
@@ -46,6 +50,7 @@ def _step(model, teacher, optimizer, inputs):
 
 
 def test_quant_lora_training_restore_and_merge():
+    """Preserve training across restore and quantized outputs across adapter merge."""
     torch.manual_seed(42)
     teacher = _model()
     student = copy.deepcopy(teacher)
@@ -101,6 +106,7 @@ def test_quant_lora_training_restore_and_merge():
 
 @pytest.mark.parametrize("targets", [["0"], ["2"]])
 def test_quant_lora_targets(targets):
+    """Adapt only matching layers and reject repeated conversion without mutation."""
     student = _model()
     inputs = torch.randn(2, 16)
     mtq.quantize(student, mtq.INT8_DEFAULT_CFG, lambda m: m(inputs))
@@ -112,6 +118,41 @@ def test_quant_lora_targets(targets):
     assert before == {name: p.requires_grad for name, p in student.named_parameters()}
 
 
-def test_quant_lora_rejects_unquantized_model():
-    with pytest.raises(ValueError, match="No supported fake-quantized"):
-        mtq.enable_quant_lora(_model())
+@pytest.mark.parametrize("case", ["unquantized", "compressed", "unsupported", "unmatched"])
+def test_quant_lora_rejection_preserves_model(case):
+    """Reject invalid targets before modifying any parameters or trainability."""
+    student = _model()
+    config = {}
+    error = "No supported fake-quantized"
+    if case == "unsupported":
+        student = nn.Sequential(nn.Linear(16, 32), nn.Conv2d(1, 2, 1))
+        mtq.quantize(
+            student,
+            mtq.INT8_DEFAULT_CFG,
+            lambda m: (m[0](torch.randn(2, 16)), m[1](torch.randn(1, 1, 4, 4))),
+        )
+        error = "does not support layer 1"
+    elif case != "unquantized":
+        mtq.quantize(student, mtq.INT8_DEFAULT_CFG, lambda m: m(torch.randn(2, 16)))
+        if case == "compressed":
+            packed, _ = INT4QTensor.quantize(student[2].weight.detach(), block_size=16)
+            student[2].weight = QTensorWrapper(packed)
+            error = "requires uncompressed"
+        else:
+            config = {"target_modules": ["missing*"]}
+    student[0].bias.requires_grad_(False)
+    parameters = dict(student.named_parameters())
+    values = {name: p.detach().clone() for name, p in parameters.items()}
+    trainability = {name: p.requires_grad for name, p in parameters.items()}
+    module_types = {name: type(m) for name, m in student.named_modules()}
+
+    with pytest.raises(ValueError, match=error):
+        mtq.enable_quant_lora(student, config)
+
+    after = dict(student.named_parameters())
+    assert after.keys() == parameters.keys()
+    for name, parameter in after.items():
+        assert parameter is parameters[name]
+        assert parameter.requires_grad == trainability[name]
+        torch.testing.assert_close(parameter, values[name], rtol=0, atol=0)
+    assert module_types == {name: type(m) for name, m in student.named_modules()}

From a840b96ec5b503d00f2d7891ef1701c84ce5cf20 Mon Sep 17 00:00:00 2001
From: Hung-Yueh Chiang <hungyuehc@nvidia.com>
Date: Thu, 8 Oct 2026 14:53:07 -0700
Subject: [PATCH 3/3] Fix LoRA QAD recomputation and unsafe weight
 transformations

Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>
---
 examples/megatron_bridge/README.md            |  4 +-
 modelopt/torch/quantization/lora.py           | 37 ++++++++
 modelopt/torch/quantization/model_quant.py    |  3 +
 .../quantization/nn/modules/quant_module.py   |  3 +
 .../quantization/plugins/megatron_lora.py     | 17 +++-
 tests/examples/megatron_bridge/test_qad.py    |  3 +
 .../torch/quantization/test_quant_lora.py     | 94 ++++++++++++++++++-
 7 files changed, 157 insertions(+), 4 deletions(-)

diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index 2463cb90029..fb4bc930f07 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -354,7 +354,9 @@ are `modelopt.torch.quantization.enable_quant_lora` and `merge_quant_lora`.
 
 Supported targets are fake-quantized dense PyTorch and Megatron/Transformer Engine linear
 layers, including tensor-parallel row/column and fused layer-normalization linear layers.
-Grouped MoE experts and compressed quantized weights are unsupported. The backbone remains
+Grouped MoE experts, shared target weights, and compressed quantized weights are unsupported.
+Merge adapters with `merge_quant_lora` before using weight-folding APIs. Full activation
+recomputation is supported without unfreezing the backbone. The backbone remains
 in floating point and the implementation forms dense weight updates, so memory savings come
 primarily from the smaller optimizer state rather than compressed backbone storage.
 
diff --git a/modelopt/torch/quantization/lora.py b/modelopt/torch/quantization/lora.py
index 5eaad2daaaa..ada87d14bcd 100644
--- a/modelopt/torch/quantization/lora.py
+++ b/modelopt/torch/quantization/lora.py
@@ -17,6 +17,7 @@
 
 import fnmatch
 import math
+from collections import Counter
 
 import torch
 from torch import nn
@@ -44,6 +45,23 @@ class QuantLoRAConfig(ModeloptBaseConfig):
 
 
 QuantLoRARegistry = _DMRegistryCls("quant_lora")
+QuantLoRAInputRegistry = _DMRegistryCls("quant_lora_input")
+
+
+class _QuantLoRAInput(DynamicModule):
+    def _setup(self):
+        """Preserve activation gradients without adding trainable backbone parameters."""
+
+
+@QuantLoRAInputRegistry.register(
+    {nn.Embedding: "nn.Embedding", QuantModuleRegistry[nn.Embedding]: "QuantEmbedding"}
+)
+class _QuantLoRAEmbedding(_QuantLoRAInput):
+    def forward(self, *args, **kwargs):
+        output = super().forward(*args, **kwargs)
+        if self.training and torch.is_grad_enabled():
+            output.requires_grad_(True)
+        return output
 
 
 @QuantLoRARegistry.register({QuantModuleRegistry[nn.Linear]: "nn.Linear"})
@@ -71,6 +89,12 @@ def forward(self, *args, **kwargs):
         finally:
             self._parameters["weight"] = weight
 
+    def _validate_fold_weight(self):
+        raise ValueError("Merge quantization-aware LoRA adapters before folding weights.")
+
+    def fold_weight(self, keep_attrs=False):
+        self._validate_fold_weight()
+
     @torch.no_grad()
     def merge_lora(self):
         """Fold the update into the base weight and remove the factors."""
@@ -101,6 +125,12 @@ def merge_quant_lora(model: nn.Module) -> nn.Module:
 def _convert_quant_lora(model, config):
     """Validate all targets before freezing the backbone and adding factors."""
     targets = []
+    parameter_uses = Counter(
+        id(parameter)
+        for module in model.modules()
+        for parameter in module._parameters.values()
+        if parameter is not None
+    )
     for name, module in model.named_modules():
         if not isinstance(module, QuantModule) or not hasattr(module, "weight_quantizer"):
             continue
@@ -115,12 +145,17 @@ def _convert_quant_lora(model, config):
             raise ValueError(
                 f"Quantization-aware LoRA does not support layer {name}: {type(module)}"
             )
+        if parameter_uses[id(weight)] > 1:
+            raise ValueError(f"Quantization-aware LoRA does not support shared weight at {name}.")
         targets.append(module)
     if not targets:
         raise ValueError("No supported fake-quantized linear layers match target_modules.")
     model.requires_grad_(False)
     for module in targets:
         QuantLoRARegistry[type(module)].convert(module, config=config)
+    for module in list(model.modules()):
+        if type(module) in QuantLoRAInputRegistry:
+            QuantLoRAInputRegistry[type(module)].convert(module)
     metadata = {}
     _update_quant_lora(model, config, metadata)
     return model, metadata
@@ -142,6 +177,8 @@ def _merge_quant_lora(model, config):
     for module in list(model.modules()):
         if isinstance(module, _QuantLoRALinear):
             module.merge_lora()
+        elif isinstance(module, _QuantLoRAInput):
+            module.export()
     metadata = {}
     _update_quant_lora(model, config, metadata)
     return model, metadata
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 5b53526aba6..f40df709a04 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -1078,6 +1078,9 @@ 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.
     """
+    for module in model.modules():
+        if isinstance(module, QuantModule):
+            module._validate_fold_weight()
     for name, module in model.named_modules():
         if isinstance(module, QuantModule):
             module.fold_weight(keep_attrs)
diff --git a/modelopt/torch/quantization/nn/modules/quant_module.py b/modelopt/torch/quantization/nn/modules/quant_module.py
index 522f2d0db09..8b1d8a4d349 100644
--- a/modelopt/torch/quantization/nn/modules/quant_module.py
+++ b/modelopt/torch/quantization/nn/modules/quant_module.py
@@ -168,6 +168,9 @@ def _fold_weight_quantizer(
                 if hasattr(quantizer, attr_name):
                     delattr(quantizer, attr_name)
 
+    def _validate_fold_weight(self):
+        """Validate that weight folding is supported before any module is modified."""
+
     def fold_weight(self, keep_attrs: bool = False):
         """Bake each fake-quant weight quantizer into its weight for faster eval.
 
diff --git a/modelopt/torch/quantization/plugins/megatron_lora.py b/modelopt/torch/quantization/plugins/megatron_lora.py
index 0648b4b917a..b2affb365ce 100644
--- a/modelopt/torch/quantization/plugins/megatron_lora.py
+++ b/modelopt/torch/quantization/plugins/megatron_lora.py
@@ -24,15 +24,30 @@
 from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
 from megatron.core.tensor_parallel.mappings import copy_to_tensor_model_parallel_region
 from megatron.core.tensor_parallel.random import get_cuda_rng_tracker
+from megatron.core.transformer.transformer_block import TransformerBlock
 from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
+from megatron.core.utils import WrappedTensor
 
-from ..lora import QuantLoRARegistry, _QuantLoRALinear
+from ..lora import QuantLoRAInputRegistry, QuantLoRARegistry, _QuantLoRAInput, _QuantLoRALinear
 from ..nn import QuantModuleRegistry
 from .megatron import _QuantMegatronMLP
 
 __all__ = []
 
 
+@QuantLoRAInputRegistry.register({TransformerBlock: "TransformerBlock"})
+class _MegatronQuantLoRATransformerBlock(_QuantLoRAInput):
+    def forward(self, hidden_states, *args, **kwargs):
+        # Reentrant checkpoints need a differentiable input even with frozen embeddings.
+        if self.training and torch.is_grad_enabled():
+            activation = hidden_states if self.pre_process else self.input_tensor
+            if isinstance(activation, WrappedTensor):
+                activation = activation.unwrap()
+            if isinstance(activation, torch.Tensor):
+                activation.requires_grad_(True)
+        return super().forward(hidden_states, *args, **kwargs)
+
+
 @QuantLoRARegistry.register(
     {
         QuantModuleRegistry[ColumnParallelLinear]: "ColumnParallelLinear",
diff --git a/tests/examples/megatron_bridge/test_qad.py b/tests/examples/megatron_bridge/test_qad.py
index 4d87036bd76..b3df875f3ce 100644
--- a/tests/examples/megatron_bridge/test_qad.py
+++ b/tests/examples/megatron_bridge/test_qad.py
@@ -157,6 +157,9 @@ def test_qad(tmp_path: Path, num_gpus, create_student, lora_rank):
         log_interval=1,
         exit_interval=early_exit_iter,
         exit_duration_in_mins=10,
+        recompute_granularity="full" if lora_rank else None,
+        recompute_method="uniform" if lora_rank else None,
+        recompute_num_layers=1 if lora_rank else None,
     )
     run_example_command(distill_cmd, example_path="megatron_bridge", setup_free_port=True)
     distilled_megatron_path = distill_output_dir / "checkpoints"
diff --git a/tests/unit/torch/quantization/test_quant_lora.py b/tests/unit/torch/quantization/test_quant_lora.py
index 5e8d0b0ed74..86dc9c057c3 100644
--- a/tests/unit/torch/quantization/test_quant_lora.py
+++ b/tests/unit/torch/quantization/test_quant_lora.py
@@ -17,14 +17,17 @@
 
 import copy
 import io
+from contextlib import nullcontext
 
 import pytest
 import torch
 from torch import nn
 from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
 
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.nn import TensorQuantizer
 from modelopt.torch.quantization.qtensor import INT4QTensor, QTensorWrapper
 
 
@@ -118,13 +121,22 @@ def test_quant_lora_targets(targets):
     assert before == {name: p.requires_grad for name, p in student.named_parameters()}
 
 
-@pytest.mark.parametrize("case", ["unquantized", "compressed", "unsupported", "unmatched"])
+@pytest.mark.parametrize(
+    "case", ["unquantized", "compressed", "unsupported", "unmatched", "shared_both", "shared_one"]
+)
 def test_quant_lora_rejection_preserves_model(case):
     """Reject invalid targets before modifying any parameters or trainability."""
     student = _model()
     config = {}
     error = "No supported fake-quantized"
-    if case == "unsupported":
+    if case.startswith("shared"):
+        student = nn.Sequential(*(nn.Linear(16, 16) for _ in range(3)))
+        student[2].weight = student[1].weight
+        mtq.quantize(student, mtq.INT8_DEFAULT_CFG, lambda m: m(torch.randn(2, 16)))
+        if case == "shared_one":
+            config = {"target_modules": ["0", "1"]}
+        error = "does not support shared weight"
+    elif case == "unsupported":
         student = nn.Sequential(nn.Linear(16, 32), nn.Conv2d(1, 2, 1))
         mtq.quantize(
             student,
@@ -156,3 +168,81 @@ def test_quant_lora_rejection_preserves_model(case):
         assert parameter.requires_grad == trainability[name]
         torch.testing.assert_close(parameter, values[name], rtol=0, atol=0)
     assert module_types == {name: type(m) for name, m in student.named_modules()}
+
+
+class _CheckpointModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embedding = nn.Embedding(16, 16)
+        self.layers = _model()
+
+    def forward(self, tokens):
+        return checkpoint(self.layers, self.embedding(tokens), use_reentrant=True)
+
+
+def test_quant_lora_reentrant_checkpoint_restore():
+    """Train adapters with frozen embeddings before and after checkpoint restoration."""
+    tokens = torch.tensor([0, 1, 2, 3])
+    student = _CheckpointModel()
+    mtq.quantize(student, mtq.INT8_DEFAULT_CFG, lambda m: m(tokens))
+    embedding_type = type(student.embedding)
+    mtq.enable_quant_lora(student, {"rank": 4, "target_modules": ["layers.*"]})
+    restored = mto.restore_from_modelopt_state(_CheckpointModel(), mto.modelopt_state(student))
+    restored.load_state_dict(student.state_dict())
+    for model in (student, restored):
+        backbone = {
+            name: p.detach().clone() for name, p in model.named_parameters() if "lora_" not in name
+        }
+        optimizer = _optimizer(model)
+        for _ in range(2):
+            optimizer.zero_grad()
+            model(tokens).square().mean().backward()
+            optimizer.step()
+        for layer in (model.layers[0], model.layers[2]):
+            assert layer.lora_A.grad.count_nonzero() > 0
+            assert layer.lora_B.grad.count_nonzero() > 0
+        for name, value in backbone.items():
+            parameter = dict(model.named_parameters())[name]
+            assert not parameter.requires_grad
+            torch.testing.assert_close(parameter, value, rtol=0, atol=0)
+        model.eval()
+        assert not model.embedding(tokens).requires_grad
+        model.train()
+        mtq.merge_quant_lora(model)
+        assert type(model.embedding) is embedding_type
+
+
+@pytest.mark.parametrize("temporary", [False, True])
+def test_quant_lora_folding_requires_merge(temporary):
+    """Reject folding atomically with active adapters, and retain parity after explicit merge."""
+    student = _model()
+    inputs = torch.randn(4, 16)
+    mtq.quantize(student, mtq.INT8_DEFAULT_CFG, lambda m: m(inputs))
+    # The earlier unadapted layer must also be untouched when folding is rejected.
+    mtq.enable_quant_lora(student, {"target_modules": ["2"]})
+    with torch.no_grad():
+        student[2].lora_B.normal_()
+    output = student(inputs).detach()
+    before = copy.deepcopy(student.state_dict())
+    parameters = dict(student.named_parameters())
+    enabled = {
+        name: m.is_enabled for name, m in student.named_modules() if isinstance(m, TensorQuantizer)
+    }
+    with pytest.raises(ValueError, match=r"Merge.*LoRA.*before folding"):
+        if temporary:
+            with mtq.temporarily_fold_weights(student):
+                pytest.fail("Active adapters must not enter the folding context")
+        else:
+            mtq.fold_weight(student)
+    for name, value in student.state_dict().items():
+        torch.testing.assert_close(value, before[name], rtol=0, atol=0)
+    assert all(dict(student.named_parameters())[name] is p for name, p in parameters.items())
+    assert enabled == {
+        name: m.is_enabled for name, m in student.named_modules() if isinstance(m, TensorQuantizer)
+    }
+    torch.testing.assert_close(student(inputs), output, rtol=0, atol=0)
+    mtq.merge_quant_lora(student)
+    with mtq.temporarily_fold_weights(student) if temporary else nullcontext():
+        if not temporary:
+            mtq.fold_weight(student)
+        torch.testing.assert_close(student(inputs), output, rtol=0, atol=0)