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)