diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index f287c409c71..d6be3951069 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -48,6 +48,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 optional Top-P (nucleus) truncation to the Megatron top-k logits KD loss (now ``TopLogitsKLLoss``) via ``logit_kl_top_p`` and ``logit_kl_top_p_min_k`` in ``DistillationConfig``: after the global Top-K selection, only the smallest prefix whose cumulative teacher probability reaches ``top_p`` (with a floor of ``min_k`` entries) contributes to the KL, mirroring ``--logits-save-top-p`` / ``--logits-save-top-p-min-k`` in Megatron-LM's logits saver.
- 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 `_ for details.
- Add ``--mlflow `` 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.
diff --git a/examples/megatron_bridge/README.md b/examples/megatron_bridge/README.md
index fd7c20a499c..fb4bc930f07 100644
--- a/examples/megatron_bridge/README.md
+++ b/examples/megatron_bridge/README.md
@@ -327,6 +327,39 @@ 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, 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.
+
### 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=` 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..ada87d14bcd
--- /dev/null
+++ b/modelopt/torch/quantization/lora.py
@@ -0,0 +1,217 @@
+# 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
+from collections import Counter
+
+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")
+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"})
+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))
+ nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+ self.lora_scale = config.alpha / config.rank
+
+ @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)
+ try:
+ return super().forward(*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."""
+ 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):
+ """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
+ 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)}"
+ )
+ 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
+
+
+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()
+ elif isinstance(module, _QuantLoRAInput):
+ module.export()
+ metadata = {}
+ _update_quant_lora(model, config, metadata)
+ return model, metadata
+
+
+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)
+
+
+@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/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
new file mode 100644
index 00000000000..b2affb365ce
--- /dev/null
+++ b/modelopt/torch/quantization/plugins/megatron_lora.py
@@ -0,0 +1,134 @@
+# 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.transformer_block import TransformerBlock
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
+from megatron.core.utils import WrappedTensor
+
+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",
+ QuantModuleRegistry[RowParallelLinear]: "RowParallelLinear",
+ }
+)
+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
+ 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):
+ """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:
+ 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):
+ """Include adapter factors with their tensor-parallel checkpoint axes."""
+ 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):
+ """Merge local factors and remove adapter-specific parallel state."""
+ 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 0360f55b5ef..b3df875f3ce 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,
@@ -127,6 +157,9 @@ def test_qad(tmp_path: Path, num_gpus, create_student):
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"
@@ -135,9 +168,45 @@ 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)
+ # 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.
- hf_export_path = tmp_path / "qad_fp8_hf"
+ hf_export_path = tmp_path / "qad_hf"
export_cmd = extend_cmd_parts(
[
"torchrun",
@@ -152,6 +221,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..86dc9c057c3
--- /dev/null
+++ b/tests/unit/torch/quantization/test_quant_lora.py
@@ -0,0 +1,248 @@
+# 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
+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
+
+
+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)
+ 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():
+ """Preserve training across restore and quantized outputs across adapter 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):
+ """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))
+ 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()}
+
+
+@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.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,
+ 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()}
+
+
+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)