From 1583d68e35f7e1d74bd05353afc0030514b6d88e Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Tue, 4 Aug 2026 19:13:20 -0700 Subject: [PATCH] Add Qwen-Image DMD2 QAT and PEFT-backed SVDQuant Add Qwen-Image support to the Diffusers quantization and FastGen DMD2/QAT flows. Move calibrated SVDQuant low-rank factors into trainable Hugging Face PEFT adapters for NVFP4, INT4, and INT8 recipes, and preserve their topology and values through ModelOpt save/restore. Keep the final DMD2 example scoped to Qwen-Image by removing the intermediate Qwen-Image-Edit additions. Signed-off-by: Jingyu Xin --- CHANGELOG.rst | 3 + docs/source/guides/_pytorch_quantization.rst | 21 ++ examples/diffusers/fastgen/README.md | 67 ++++- examples/diffusers/fastgen/dmd2_recipe.py | 278 +++++++++++++++++- .../fastgen/fastgen_data/collate_fns.py | 6 + .../fastgen/inference_dmd2_qwen_image.py | 35 ++- .../diffusers/quantization/calibration.py | 20 ++ .../diffusers/quantization/models_utils.py | 19 ++ .../quantization/pipeline_manager.py | 115 ++++++++ examples/diffusers/quantization/quantize.py | 14 +- .../quantization/qwen_image_dmd2_sampler.py | 260 ++++++++++++++++ .../quantization/sanity_check_dmd2.py | 175 +++++++++++ examples/diffusers/quantization/utils.py | 38 +++ modelopt/torch/fastgen/plugins/qwen_image.py | 36 ++- modelopt/torch/quantization/conversion.py | 7 +- modelopt/torch/quantization/mode.py | 29 ++ .../quantization/nn/modules/quant_linear.py | 30 +- .../plugins/diffusion/diffusers.py | 27 +- .../quantization/plugins/svdquant_peft.py | 124 ++++++++ .../plugins/transformers_trainer.py | 10 + .../torch/quantization/quantize_common.py | 9 +- .../fastgen/test_quant_state_roundtrip.py | 95 ++++++ .../torch/fastgen/test_qwen_image_plugin.py | 29 +- .../plugins/test_svdquant_peft_modelopt.py | 208 +++++++++++++ tests/unit/torch/quantization/test_calib.py | 21 +- 25 files changed, 1611 insertions(+), 65 deletions(-) create mode 100644 examples/diffusers/quantization/qwen_image_dmd2_sampler.py create mode 100644 examples/diffusers/quantization/sanity_check_dmd2.py create mode 100644 modelopt/torch/quantization/plugins/svdquant_peft.py create mode 100644 tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py create mode 100644 tests/unit/torch/quantization/plugins/test_svdquant_peft_modelopt.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7a188669cb8..037adc835e9 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog *Quantization* +- Store SVDQuant low-rank residual factors as trainable Hugging Face PEFT adapter parameters on exactly the layers calibrated by SVDQuant. The forward remains ``Q(W_residual)x + B(Ax)``. Full ``mto.save`` / ``mto.restore`` checkpoints and the split ``modelopt_state`` + complete ``state_dict`` flow preserve the adapter topology and A/B values. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Plain ``max`` sets a tensor's global scale from the largest per-block amax seen during calibration, leaving no room above it, so any activation larger than the calibration max saturates. ``nvfp4_act_headroom`` instead anchors the global scale to a low percentile of the per-block amax distribution, ``amax = max(rho * anchor, upper)``, placing the calibrated blocks in the lower part of the FP8 block-scale range and leaving the rest as headroom. ``upper`` is the top of the range the scale commits to representing and defaults to the 99.99th percentile rather than the literal maximum: chasing a lone freak block would drag the global scale up until every other block's FP8 block scale falls below subnormal and flushes to zero, so the rarest blocks are clipped instead. Set ``upper_percentile=100`` to use the literal observed max, which guarantees no calibration data is clipped. The calibrator warns when the per-block range is too wide for ``rho`` to clear any headroom. Tunable via ``anchor_percentile`` (default 1), ``upper_percentile`` (default 99.99) and ``rho`` (default 16384). Applies only to NVFP4 dynamic-block input quantizers. Weight scales are an orthogonal axis, selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian`` with that algorithm's own options), so one recipe can combine a weight calibration with this activation policy in a single pass. ``SequentialQuantizer`` activation quantizers are not supported and raise. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` (dynamic NVFP4 W4A4 plus FP8 KV-cache cast) with only the calibration algorithm swapped, so it exports a standard NVFP4 checkpoint. *Misc* @@ -17,6 +18,8 @@ Changelog **Backward Breaking Changes** +- New SVDQuant checkpoints no longer expose low-rank factors through ``weight_quantizer.svdquant_lora_a/b``; Hugging Face PEFT owns them as trainable ``modelopt_svdquant`` adapter parameters. Code that directly inspected or optimized the quantizer buffers must use the PEFT parameters instead. + **Deprecations** **Bug Fixes** diff --git a/docs/source/guides/_pytorch_quantization.rst b/docs/source/guides/_pytorch_quantization.rst index f8f12b068ba..b2e495f595a 100644 --- a/docs/source/guides/_pytorch_quantization.rst +++ b/docs/source/guides/_pytorch_quantization.rst @@ -114,6 +114,22 @@ Here is an example of performing QAT: # Adjust learning rate and training duration train(model, train_loader, optimizer, scheduler, ...) +For SVDQuant, the calibrated low-rank residual is represented by a structural +Hugging Face PEFT adapter on only the layers to which SVDQuant was applied. The +adapter A/B factors are trainable immediately after calibration or restoration, +so they are included by a normally constructed optimizer. Construct distributed +wrappers and the optimizer after quantization or checkpoint restoration: + +.. code-block:: python + + model = mtq.quantize(model, quant_config, forward_loop) + model = DistributedDataParallel(model) + optimizer = create_optimizer(parameter for parameter in model.parameters() if parameter.requires_grad) + +The adapter implements the SVDQuant decomposition +``Q(W_residual)x + B(Ax)`` and requires the ``nvidia-modelopt[hf]`` optional +dependencies. + .. tip:: We recommend QAT for 10% of the original training epochs. For LLMs, we find that QAT fine-tuning for even @@ -126,6 +142,11 @@ The model weights and quantizer states need to saved for future use or to resume Please see :ref:`saving and restoring of ModelOpt-modified models ` to learn how to save and restore the quantized model. +For PEFT-backed SVDQuant, ``mto.save`` / ``mto.restore`` preserve the adapter +topology and factor values. The split checkpoint flow is also supported: restore +``mto.modelopt_state(model)`` first, then load the complete +``model.state_dict()`` before constructing the distributed wrapper and optimizer. + Optimal Partial Quantization using ``auto_quantize`` =================================================================== diff --git a/examples/diffusers/fastgen/README.md b/examples/diffusers/fastgen/README.md index 9c9373807a9..8cd794f1ba8 100644 --- a/examples/diffusers/fastgen/README.md +++ b/examples/diffusers/fastgen/README.md @@ -119,10 +119,73 @@ Any `DMDConfig` field can be overridden on the CLI (e.g. `--dmd2.guidance_scale= Checkpoints land under `checkpoint.checkpoint_dir`. Alongside the student, the recipe saves the DMD2 sidecars needed to resume exactly: the fake-score model + optimizer, the -student EMA (`ema_shadow.pt`), and the DMD iteration counter (`dmd_state.pt`). With +DMD iteration counter (`dmd_state.pt`), and, when EMA is enabled, the student EMA +(`ema_shadow.pt`). With `restore_from: LATEST` a re-launch auto-resumes from the newest checkpoint; pin a specific one with `--checkpoint.restore_from=epoch_0_step_500`. +## Quantization-aware training (QAT) + +Continue a full-precision DMD2 run with the **student quantized**, so the few-step model +stays accurate at FP8/NVFP4. QAT here is **restore-only**: the trainer loads a ModelOpt +quantizer state (recipe + frozen `amax`) from disk and **never calibrates**. Only the +student is quantized; the frozen teacher and trainable fake-score stay full precision so +the distribution-matching gradient is exact, and `amax` stays frozen for the whole run. + +QAT is driven by a `dmd2.quant` block — there's no dedicated config file. The cleanest +way to launch is to **reuse the exact config + overrides of the full-precision run you're +continuing** and add only the three `dmd2.quant.*` keys (plus a reduced LR), so the QAT +run is provably identical to the FP run except for quantization and learning rate. The +CLI parser creates the `dmd2.quant` subtree even when it's absent from the YAML. + +| Key | Role | +| --- | --- | +| `dmd2.quant.enabled` | Turn QAT on (restore-only student quantization). | +| `dmd2.quant.quant_state_path` | The `transformer.pt` from step 1 below (recipe + frozen `amax`). | +| `dmd2.quant.init_weights_from` | FP DMD2 checkpoint to warm-start student / fake-score / optimizers from on the first launch (the run `amax` was calibrated against). | + +1. **Calibrate once** with the quantization example to produce the quantizer state + (`amax`, no weights) for a trained student checkpoint: + + ```bash + python examples/diffusers/quantization/quantize.py \ + --model qwen-image-dmd2 --format fp8 \ + --extra-param student_path=<.../epoch_4_step_15999/model/consolidated> \ + --quantized-torch-ckpt-save-path <.../epoch_4_step_15999/quant> + # -> writes <.../epoch_4_step_15999/quant/transformer.pt> + ``` + +2. **Launch QAT** by re-running the FP run's command with a new output dir, a reduced + student LR, and the three quant keys appended: + + ```bash + torchrun --nproc-per-node= \ + examples/diffusers/fastgen/dmd2_finetune.py \ + --config examples/diffusers/fastgen/configs/.yaml \ + --checkpoint.checkpoint_dir= \ + <... the FP run's other overrides, unchanged ...> \ + --optim.learning_rate= --lr_scheduler.min_lr= \ + --dmd2.quant.enabled=true \ + --dmd2.quant.quant_state_path=<.../epoch_4_step_15999/quant/transformer.pt> \ + --dmd2.quant.init_weights_from=<.../epoch_4_step_15999> + ``` + +On the first launch (empty `checkpoint_dir`) the student / fake-score / discriminator / +optimizers warm-start from `init_weights_from`, then the student is quantized from +`quant_state_path`. `restore_from: LATEST` auto-resumes the new `checkpoint_dir` +thereafter. Because QAT is restore-only — amax never recalibrates — the recipe re-applies +`quant_state_path` on every resume rather than persisting a per-checkpoint copy, so keep +that file accessible for the whole run (it's the only quantization dependency). The saved +student weights are clean full precision (`model/consolidated` is a normal +`QwenImageTransformer2DModel`); re-apply `quant_state_path` to deploy or evaluate the +quantized QAT student via the quantization example. + +> The `quant_state_path` `amax` must have been calibrated against the student in +> `init_weights_from`, with the same few-step schedule (`dmd2.sample_t_cfg.t_list`) the +> student trains/infers with. Pass `dmd2.quant.enabled=true` on every resume too (it is +> what tells the recipe to quantize). Reduce only the student LR by keeping +> `--dmd2.fake_score_lr` / `--dmd2.discriminator_lr` at the FP value. + ## Inference After training, sample from the distilled student. The pipeline loads your consolidated @@ -170,7 +233,7 @@ student). | `model` | `mode` | `finetune` — loads the pretrained weights. | | `step_scheduler` | `global_batch_size`, `local_batch_size`, `max_steps`, `ckpt_every_steps`, `log_every` | Standard AutoModel scheduling knobs. | | `dmd2` | `recipe_path` | Built-in fastgen recipe to hydrate `DMDConfig` from (`general/distillation/dmd2_qwen_image`). | -| `dmd2` | `pipeline_plugin` | `qwen_image` — selects `QwenImageDMDPipeline` (2×2 patch packing / img_shapes). | +| `dmd2` | `pipeline_plugin` | `qwen_image` — selects `QwenImageDMDPipeline` (2×2 patch packing / `img_shapes`). | | `dmd2` | `student_sample_steps` | Number of student sampling steps (e.g. 4). | | `dmd2` | `guidance_scale` | CFG strength on the teacher (`null` disables CFG; requires a negative-prompt embedding when set). | | `dmd2` | `gan_loss_weight_gen`, `gan_r1_reg_weight`, `gan_feature_indices`, … | GAN branch (set `gan_loss_weight_gen: 0` to disable). | diff --git a/examples/diffusers/fastgen/dmd2_recipe.py b/examples/diffusers/fastgen/dmd2_recipe.py index 7934a07cf13..a1e64b73cc7 100644 --- a/examples/diffusers/fastgen/dmd2_recipe.py +++ b/examples/diffusers/fastgen/dmd2_recipe.py @@ -37,6 +37,8 @@ from __future__ import annotations +import contextlib +import dataclasses import json import logging import os @@ -54,6 +56,11 @@ # and surfaced as a downstream ``TypeError: takes no arguments``. try: from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline + from nemo_automodel.components.distributed.parallelizer import ( + PARALLELIZATION_STRATEGIES, + DefaultParallelizationStrategy, + register_parallel_strategy, + ) from nemo_automodel.recipes.diffusion.train import TrainDiffusionRecipe, is_main_process except ImportError as exc: raise ImportError( @@ -68,10 +75,67 @@ from torch import nn import modelopt.torch.fastgen as mtf +import modelopt.torch.opt as mto from modelopt.torch.fastgen.config import DMDConfig from modelopt.torch.fastgen.discriminators import Discriminator_ImageDiT from modelopt.torch.fastgen.methods.dmd import DMDPipeline from modelopt.torch.fastgen.plugins import qwen_image as qwen_image_plugin +from modelopt.torch.quantization.utils.core_utils import set_quantizer_state_dict + + +class _QwenImageParallelizationStrategy(DefaultParallelizationStrategy): + """Add full-block activation checkpointing to AutoModel's native FSDP flow. + + AutoModel's default decoder-layer checkpointing recognizes ``self_attn`` / ``mlp`` + attributes. Diffusers' Qwen image blocks instead contain joint ``attn``, ``img_mlp``, + and ``txt_mlp`` paths, so that generic logic wraps nothing. Diffusers checkpoints each + complete block; mirror that boundary, then delegate TP/FSDP behavior unchanged. + """ + + def parallelize( + self, + model, + device_mesh, + activation_checkpointing: bool = False, + **kwargs, + ): + if activation_checkpointing: + from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + CheckpointImpl, + checkpoint_wrapper, + ) + + blocks = getattr(model, "transformer_blocks", None) + if blocks is None: + raise AttributeError( + "QwenImageTransformer2DModel does not expose `transformer_blocks`" + ) + for index, block in enumerate(blocks): + blocks[index] = checkpoint_wrapper( + block, + checkpoint_impl=CheckpointImpl.NO_REENTRANT, + ) + logging.info( + "[DMD2] Qwen-Image activation checkpointing enabled for %d full blocks", + len(blocks), + ) + + return super().parallelize( + model, + device_mesh, + activation_checkpointing=False, + **kwargs, + ) + + +def _register_qwen_image_parallelization_strategy() -> None: + """Install the Qwen strategy unless AutoModel already provides a native one.""" + model_class_name = "QwenImageTransformer2DModel" + if model_class_name not in PARALLELIZATION_STRATEGIES: + register_parallel_strategy(name=model_class_name)(_QwenImageParallelizationStrategy) + + +_register_qwen_image_parallelization_strategy() # Keys under the ``dmd2:`` YAML block that shadow fields on :class:`DMDConfig`. The # recipe deep-merges these on top of the loaded built-in recipe so users can tweak DMD2 @@ -98,6 +162,31 @@ def _deep_merge_dicts(base: dict, override: dict) -> dict: return merged +def restore_quantizer_state(model: nn.Module, path: str) -> nn.Module: + """Re-insert the quantizer modules + load the frozen amax onto ``model`` from ``path``. + + ``path`` is the weight-free ModelOpt quantizer state written by the + ``examples/diffusers/quantization`` calibration example + (``--quantized-torch-ckpt-save-path`` → ``transformer.pt``): ``mto.modelopt_state`` + (which layers are quantized, FP8/NVFP4, axes) bundled with the per-quantizer buffers + (``amax``, ...) under ``modelopt_state_weights``. It carries NO model weights; ``model`` + must already hold the weights the amax was calibrated against (here: the DMD2 student + warm-started from the FP checkpoint). + + This re-applies the quantization recipe (module conversion -- ``nn.Linear`` -> + ``QuantLinear``, in place, preserving the existing ``weight``/``bias`` Parameter objects + so any pre-built FSDP2 optimizer's references stay valid) and loads the saved amax. It + is RESTORE-ONLY: no calibration forward pass is run, so amax stays exactly as it was on + disk and remains frozen for the whole training run. + """ + modelopt_state = mto.load_modelopt_state(str(path)) + quantizer_state = modelopt_state.pop("modelopt_state_weights", None) + mto.restore_from_modelopt_state(model, modelopt_state) + if quantizer_state is not None: + set_quantizer_state_dict(model, quantizer_state) + return model + + # Auto-detect substrings (matched case-insensitively against ``model_id``) that map to # DMDPipeline plugin subclasses. Keep this list small — adding a new entry is only the # right move when the model has a non-diffusers transformer signature that requires a @@ -111,6 +200,29 @@ def _deep_merge_dicts(base: dict, override: dict) -> dict: _DMD_COMPLETE_MARKER = "dmd2_complete.marker" +@dataclasses.dataclass(frozen=True) +class _QuantSettings: + """Resolved ``dmd2.quant`` block — restore-only QAT of the student (no calibration). + + Attributes: + enabled: When ``True`` the student is quantized by RESTORING a ModelOpt quantizer + state from disk (recipe + frozen amax). The trainer never calibrates. + quant_state_path: Path to the quantizer-state file produced by the + ``examples/diffusers/quantization`` calibration example + (``--quantized-torch-ckpt-save-path`` → ``transformer.pt``). Used on the FIRST + launch (warm-start) to quantize the FP student. Required when ``enabled``. + init_weights_from: Optional path to the full-precision DMD2 checkpoint to + warm-start the student / fake_score / discriminator / EMA / optimizers from + when the run's own ``checkpoint_dir`` has no QAT checkpoint yet. The amax in + ``quant_state_path`` must have been calibrated against this checkpoint's + student weights. + """ + + enabled: bool + quant_state_path: str | None + init_weights_from: str | None + + class DMD2DiffusionRecipe(TrainDiffusionRecipe): """DMD2 recipe that reuses ``TrainDiffusionRecipe`` for the student path. @@ -379,7 +491,6 @@ def run_train_validation_loop(self) -> None: neg_text_embeds, neg_text_mask, ) = self._prepare_micro_batch(micro_batch) - if is_student_phase: # ``compute_student_loss`` reads ``guidance_scale`` from the # DMDConfig when this kwarg is None. We pass the negative @@ -538,10 +649,33 @@ def load_checkpoint(self, restore_from: str | None = None): # ``nemo_automodel`` can be used unmodified. make_optimizer_partial_load_tolerant(self.checkpointer) + quant = self._resolve_quant_settings() + resolved = self._resolve_complete_dmd_checkpoint(restore_from) + + # QAT first launch: the run's own checkpoint_dir has no QAT checkpoint yet, so + # warm-start the FP student / fake_score / EMA / optimizers from the full-precision + # checkpoint named by ``dmd2.quant.init_weights_from`` (the FP run the amax was + # calibrated against). On later resumes ``resolved`` already points at a QAT + # checkpoint in checkpoint_dir, so this branch is skipped. + if quant.enabled and resolved is None and quant.init_weights_from: + resolved = self._resolve_complete_dmd_checkpoint(quant.init_weights_from) + if is_main_process(): + logging.info( + "[DMD2][qat] no QAT checkpoint in checkpoint_dir; warm-starting FP " + "state from dmd2.quant.init_weights_from=%s", + quant.init_weights_from, + ) + self.__dict__["_dmd2_resolved_restore_from"] = resolved if resolved is None: + if quant.enabled and is_main_process(): + logging.warning( + "[DMD2][qat] QAT enabled but no checkpoint resolved (checkpoint_dir empty " + "and dmd2.quant.init_weights_from unset). Quantizing a fresh base model — " + "its weights will NOT match the calibrated amax." + ) if ( restore_from is not None and str(restore_from).upper() == "LATEST" @@ -552,10 +686,137 @@ def load_checkpoint(self, restore_from: str | None = None): "Starting fresh.", self.checkpointer.config.checkpoint_dir, ) + # Even with no weights to restore, honor QAT by quantizing from the recipe file. + if quant.enabled: + self._quantize_student(quant.quant_state_path) return + # Single restore path. The student-weight checkpoints are always clean FP (the QAT + # save hides the quantizer buffers — see ``_student_quant_buffers_hidden``), so the + # strict DCP load matches an unquantized student whether this is a warm-start from + # the FP checkpoint or a resume from a QAT checkpoint. Quantization always happens + # AFTER the weights are in place. super().load_checkpoint(resolved) + if quant.enabled: + # Restore-only QAT never recalibrates, so the recipe + amax are invariant for the + # whole run — re-apply the same configured quantizer-state file on every (re)start + # rather than persisting an unchanging copy per checkpoint. + self._quantize_student(quant.quant_state_path) + + def _resolve_quant_settings(self) -> _QuantSettings: + """Parse and cache the ``dmd2.quant`` block (restore-only student QAT).""" + cached = self.__dict__.get("_quant_settings") + if cached is not None: + return cached + + node = self.cfg.get("dmd2.quant", None) + if node is None: + settings = _QuantSettings(enabled=False, quant_state_path=None, init_weights_from=None) + self.__dict__["_quant_settings"] = settings + return settings + + d = node.to_dict() if hasattr(node, "to_dict") else dict(node) + enabled = bool(d.get("enabled", False)) + quant_state_path = d.get("quant_state_path") + init_weights_from = d.get("init_weights_from") + if enabled and not quant_state_path: + raise ValueError( + "dmd2.quant.enabled is true but dmd2.quant.quant_state_path is not set. Point it " + "at the quantizer-state file produced by examples/diffusers/quantization " + "(its --quantized-torch-ckpt-save-path, e.g. .../quant/transformer.pt)." + ) + settings = _QuantSettings( + enabled=enabled, + quant_state_path=quant_state_path, + init_weights_from=init_weights_from, + ) + self.__dict__["_quant_settings"] = settings + return settings + + def _quantize_student(self, quant_state_path: str | None) -> None: + """Quantize the student by RESTORING a ModelOpt quantizer state from disk. + + Restore-only: re-inserts the quantizer modules from the saved recipe and loads the + precomputed, frozen amax. No ``mtq.quantize`` / calibration forward pass is ever + run, so amax stays exactly as it was on disk for the whole training run. Only the + student (``self.model``) is quantized — the frozen teacher and the trainable + fake_score stay full precision so the distribution-matching gradient is exact. + + Safe after FSDP2 wrapping: the conversion is an in-place ``__class__`` swap that + preserves the existing ``weight``/``bias`` Parameter objects, so the already-built + student optimizer's references stay valid (verified by ModelOpt's + ``tests/gpu/torch/quantization/test_fsdp2.py``). + """ + if not quant_state_path: + raise ValueError( + "[DMD2][qat] _quantize_student called without a quant_state_path. Set " + "dmd2.quant.quant_state_path." + ) + if not os.path.isfile(quant_state_path): + raise FileNotFoundError( + f"[DMD2][qat] quantizer-state file not found: {quant_state_path}. QAT here is " + "restore-only (no on-the-fly calibration); produce it with " + "examples/diffusers/quantization first." + ) + + if is_main_process(): + logging.info( + "[DMD2][qat] restoring student quantizer state (recipe + frozen amax) <- %s", + quant_state_path, + ) + restore_quantizer_state(self.model, quant_state_path) + + # amax (and any other quantizer buffers) come off disk on CPU. Move just the + # TensorQuantizer buffers onto the student device — these modules carry no + # parameters, so this never touches the FSDP2 DTensor weights. + from modelopt.torch.quantization.nn import TensorQuantizer + + for module in self.model.modules(): + if isinstance(module, TensorQuantizer): + module.to(self.device) + + if is_main_process(): + import modelopt.torch.quantization as mtq + + logging.info("[DMD2][qat] student quantized (restore-only). Quantizer summary:") + mtq.print_quant_summary(self.model) + + @contextlib.contextmanager + def _student_quant_buffers_hidden(self): + """Temporarily mark the student's quantizer buffers non-persistent. + + Wraps the parent ``save_checkpoint`` so the student's DCP shards AND the + consolidated/diffusers export stay clean full-precision (``ModelState.state_dict()`` + — and hence the consolidated index, which re-adds every state_dict key + (checkpointing.py:941-948) — excludes the ``amax`` buffers). The frozen amax is not + persisted per checkpoint at all: it is re-applied from ``dmd2.quant.quant_state_path`` + on every (re)start. Keeping the saved student weights amax-free both yields a clean + ``model/consolidated`` (a normal ``QwenImageTransformer2DModel``, re-quantizable for + deploy/eval) and makes every restore a clean ``load FP weights -> quantize`` path + (the strict DCP load matches an unquantized student). Mirrors ModelOpt's own trick in + ``quantization/plugins/transformers_trainer.py`` (``_modelopt_prepare``). + + No-op when QAT is disabled. + """ + if not self._resolve_quant_settings().enabled: + yield + return + + from modelopt.torch.quantization.nn import TensorQuantizer + + saved: list[tuple[TensorQuantizer, set[str]]] = [] + for module in self.model.modules(): + if isinstance(module, TensorQuantizer): + saved.append((module, set(module._non_persistent_buffers_set))) + module._non_persistent_buffers_set.update(module._buffers.keys()) + try: + yield + finally: + for module, original in saved: + module._non_persistent_buffers_set.clear() + module._non_persistent_buffers_set.update(original) + def save_checkpoint( self, epoch: int, @@ -588,7 +849,11 @@ def save_checkpoint( self.checkpointer.config.checkpoint_dir ) - super().save_checkpoint(epoch, step, train_loss, val_loss, best_metric_key) + # Hide the student's quantizer buffers during the parent save so the DCP shards and + # consolidated/diffusers export stay clean FP. Frozen amax remains external in + # ``dmd2.quant.quant_state_path`` and is re-applied on restore. No-op when QAT is disabled. + with self._student_quant_buffers_hidden(): + super().save_checkpoint(epoch, step, train_loss, val_loss, best_metric_key) if not self.checkpointer.config.enabled: return @@ -861,6 +1126,9 @@ def _is_dmd_checkpoint_complete(self, path: str) -> bool: if not complete: return False + # QAT adds no per-checkpoint artifact (the quantizer recipe + frozen amax are + # re-applied from dmd2.quant.quant_state_path on every (re)start), so QAT + # checkpoints use the same completeness criteria as full-precision ones. if self._cfg_gan_enabled(): return os.path.isfile(os.path.join(path, "discriminator.pt")) and os.path.isfile( os.path.join(path, "discriminator_optimizer.pt") @@ -976,10 +1244,10 @@ def _build_discriminator_optimizer(self) -> torch.optim.Optimizer | None: def _attach_gan_feature_capture(self) -> None: """Install Qwen-Image feature-capture hooks on the teacher when GAN is enabled. - Reads the latent resolution from the dataloader so the hook can reshape + Reads an initial latent resolution from the dataloader so the hook can reshape ``[B, num_image_patches, 3072]`` into ``[B, 3072, H_lat//2, W_lat//2]``. - Mock dataloader → spatial_h/spatial_w from the YAML. Real dataloader → - base_resolution / vae_scale. + The Qwen plugin refreshes that shape before every teacher forward, so real + multiresolution batches are captured using their actual target dimensions. """ feature_indices = list(self.cfg.get("dmd2.gan_feature_indices", [30])) diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index d669d2a7c4a..ed7221dd93b 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -240,3 +240,9 @@ def build_text_to_image_multiresolution_dataloader( dp_world_size, ) return dataloader, sampler + + +__all__ = [ + "build_text_to_image_multiresolution_dataloader", + "collate_fn_text_to_image", +] diff --git a/examples/diffusers/fastgen/inference_dmd2_qwen_image.py b/examples/diffusers/fastgen/inference_dmd2_qwen_image.py index 5907d0f1b86..695be95abb9 100644 --- a/examples/diffusers/fastgen/inference_dmd2_qwen_image.py +++ b/examples/diffusers/fastgen/inference_dmd2_qwen_image.py @@ -69,17 +69,26 @@ import logging import os from dataclasses import dataclass +from importlib.metadata import PackageNotFoundError, version from typing import TYPE_CHECKING import torch from diffusers import QwenImagePipeline, QwenImageTransformer2DModel from diffusers.utils.torch_utils import randn_tensor +from packaging.version import Version if TYPE_CHECKING: from pathlib import Path logger = logging.getLogger(__name__) +try: + # Diffusers 0.35/0.36 needs explicit sequence lengths; starting in 0.37 the + # attention mask is authoritative and ``txt_seq_lens`` was removed. + _DIFFUSERS_NEEDS_TXT_SEQ_LENS = Version(version("diffusers")) < Version("0.37.0") +except PackageNotFoundError: # pragma: no cover - diffusers is imported above + _DIFFUSERS_NEEDS_TXT_SEQ_LENS = False + @dataclass class QwenImageDMDOutput: @@ -294,15 +303,19 @@ def __call__( num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, ) - txt_seq_lens = ( - prompt_embeds_mask.sum(dim=1).int().tolist() if prompt_embeds_mask is not None else None - ) - neg_txt_seq_lens = ( - neg_prompt_embeds_mask.sum(dim=1).int().tolist() - if neg_prompt_embeds_mask is not None - else None - ) - + positive_transformer_kwargs = {} + negative_transformer_kwargs = {} + if _DIFFUSERS_NEEDS_TXT_SEQ_LENS: + positive_transformer_kwargs["txt_seq_lens"] = ( + prompt_embeds_mask.sum(dim=1).int().tolist() + if prompt_embeds_mask is not None + else None + ) + negative_transformer_kwargs["txt_seq_lens"] = ( + neg_prompt_embeds_mask.sum(dim=1).int().tolist() + if neg_prompt_embeds_mask is not None + else None + ) # ---- 3. Build initial noisy latents at t = schedule[0] --------------- if isinstance(prompt, str): batch_size = 1 @@ -334,9 +347,9 @@ def __call__( encoder_hidden_states_mask=prompt_embeds_mask, timestep=timestep, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, guidance=None, return_dict=False, + **positive_transformer_kwargs, )[0] if do_cfg: # CFG two-pass: ``v_cfg = v_neg + s*(v_pos - v_neg)``. Equivalent @@ -352,9 +365,9 @@ def __call__( encoder_hidden_states_mask=neg_prompt_embeds_mask, timestep=timestep, img_shapes=img_shapes, - txt_seq_lens=neg_txt_seq_lens, guidance=None, return_dict=False, + **negative_transformer_kwargs, )[0] flow_packed = ( neg_flow_packed.to(torch.float64) diff --git a/examples/diffusers/quantization/calibration.py b/examples/diffusers/quantization/calibration.py index 27b1ec22436..bebc61970a3 100644 --- a/examples/diffusers/quantization/calibration.py +++ b/examples/diffusers/quantization/calibration.py @@ -21,6 +21,7 @@ from models_utils import MODEL_DEFAULTS, ModelType from pipeline_manager import PipelineManager from quantize_config import CalibrationConfig +from qwen_image_dmd2_sampler import dmd2_sample from tqdm import tqdm from utils import load_calib_prompts @@ -95,6 +96,9 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: elif self.model_type in [ModelType.WAN22_T2V_14b, ModelType.WAN22_T2V_5b]: # Special handling for WAN video models self._run_wan_video_calibration(prompt_batch, extra_args) + elif self.model_type == ModelType.QWEN_IMAGE_DMD2: + # DMD2 students use a custom few-step sampler, not the standard loop. + self._run_qwen_image_dmd2_calibration(prompt_batch) else: common_args = { "prompt": prompt_batch, @@ -105,6 +109,22 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: self.logger.debug(f"Completed calibration batch {i + 1}/{self.config.num_batches}") self.logger.info("Calibration completed successfully") + def _run_qwen_image_dmd2_calibration(self, prompt_batch: list[str]) -> None: + """Calibrate a DMD2 Qwen-Image student via its few-step sampler. + + Drives the same few-step DMD unroll the student was trained/served with + (NOT the standard denoising loop) so the collected activation statistics + are representative of inference. The VAE decode is skipped — calibration + only needs the transformer forwards. + """ + cfg = self.pipeline_manager.dmd_sampler_cfg + if cfg is None: + raise RuntimeError( + "DMD2 sampler config is not set; the qwen-image-dmd2 pipeline must be created " + "via PipelineManager.create_pipeline() before calibration." + ) + dmd2_sample(self.pipe, prompt_batch, decode=False, **cfg) + def _run_wan_video_calibration( self, prompt_batch: list[str], extra_args: dict[str, Any] ) -> None: diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index 4d1bd803305..3dc41138797 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -64,6 +64,10 @@ class ModelType(str, Enum): WAN22_T2V_14b = "wan2.2-t2v-14b" WAN22_T2V_5b = "wan2.2-t2v-5b" QWEN_IMAGE = "qwen-image" + # DMD2-distilled few-step Qwen-Image student (from examples/diffusers/fastgen). + # Same architecture as QWEN_IMAGE, but loaded from a consolidated student dir + # and calibrated with the few-step DMD sampler instead of the standard loop. + QWEN_IMAGE_DMD2 = "qwen-image-dmd2" _FILTER_FUNC_MAP: dict[ModelType, Callable[[str], bool]] = { @@ -74,6 +78,7 @@ class ModelType(str, Enum): ModelType.WAN22_T2V_14b: filter_func_wan_video, ModelType.WAN22_T2V_5b: filter_func_wan_video, ModelType.QWEN_IMAGE: filter_func_qwen_image, + ModelType.QWEN_IMAGE_DMD2: filter_func_qwen_image, } _VAE_FILTER_FUNC_MAP: dict[tuple[ModelType, str], Callable[[str], bool]] = { @@ -107,6 +112,11 @@ def get_model_filter_func( ModelType.WAN22_T2V_14b: "Wan-AI/Wan2.2-T2V-A14B-Diffusers", ModelType.WAN22_T2V_5b: "Wan-AI/Wan2.2-TI2V-5B-Diffusers", ModelType.QWEN_IMAGE: "Qwen/Qwen-Image", + # Base pipeline (VAE / text-encoder / tokenizer / scheduler) for DMD2 students; + # the trained transformer is loaded separately from a consolidated dir via the + # ``student_path`` extra-param. Override with ``--override-model-path`` or + # ``--extra-param base_pipeline_path=...``. + ModelType.QWEN_IMAGE_DMD2: "Qwen/Qwen-Image", } MODEL_PIPELINE: dict[ModelType, type[DiffusionPipeline] | None] = { @@ -122,6 +132,7 @@ def get_model_filter_func( ModelType.WAN22_T2V_14b: WanPipeline, ModelType.WAN22_T2V_5b: WanPipeline, ModelType.QWEN_IMAGE: QwenImagePipeline, + ModelType.QWEN_IMAGE_DMD2: QwenImagePipeline, } # Shared dataset configurations @@ -273,6 +284,14 @@ def get_model_filter_func( }, } +# DMD2 students share Qwen-Image's architecture, so they reuse the same block-range +# recipe, high-precision filter, base pipeline, and calibration dataset. They differ +# only in (a) loading -- a consolidated student dir swapped into the base pipeline +# (PipelineManager._create_qwen_image_dmd2_pipeline) -- and (b) calibration, which +# drives the few-step DMD sampler instead of the standard denoising loop +# (Calibrator._run_qwen_image_dmd2_calibration). Inherit so the recipe stays in sync. +MODEL_DEFAULTS[ModelType.QWEN_IMAGE_DMD2] = {**MODEL_DEFAULTS[ModelType.QWEN_IMAGE]} + def _coerce_extra_param_value(value: str) -> Any: lowered = value.lower() diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index af89ed568ff..4e76ee0d661 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -43,6 +43,9 @@ def __init__(self, config: ModelConfig, logger: logging.Logger): self.pipe_upsample: LTXLatentUpsamplePipeline | None = None # For LTX-Video upsampling self._transformer: torch.nn.Module | None = None self._video_decoder: torch.nn.Module | None = None + # Few-step sampler config for DMD2 students (populated when loading a + # qwen-image-dmd2 pipeline); consumed by the calibrator / sanity check. + self.dmd_sampler_cfg: dict[str, Any] | None = None @staticmethod def create_pipeline_from( @@ -100,6 +103,11 @@ def create_pipeline(self) -> Any: self.logger.info("LTX-2 pipeline created successfully") return self.pipe + if self.config.model_type == ModelType.QWEN_IMAGE_DMD2: + self.pipe = self._create_qwen_image_dmd2_pipeline() + self.logger.info("Qwen-Image DMD2 pipeline created successfully") + return self.pipe + pipeline_cls = MODEL_PIPELINE[self.config.model_type] if pipeline_cls is None: raise ValueError( @@ -266,6 +274,113 @@ def _create_ltx2_pipeline(self) -> Any: pipeline_kwargs.update(params) return TI2VidTwoStagesPipeline(**pipeline_kwargs) + def _create_qwen_image_dmd2_pipeline(self) -> Any: + """Build a QwenImagePipeline whose transformer is a DMD2-trained student. + + Loads the consolidated student transformer (the ``model/consolidated`` dir + produced by ``examples/diffusers/fastgen`` training), optionally overlays an + EMA shadow, and swaps it into the base Qwen-Image pipeline so the VAE / + text-encoder / tokenizer / scheduler come from the base checkpoint. + + Reads from ``extra_params``: + student_path (required): consolidated student dir. + base_pipeline_path: base Qwen-Image dir/HF id (defaults to the + registry id or ``--override-model-path``). + ema_path: optional ``ema_shadow.pt`` to overlay onto the student. + sample_steps / t_list / sample_type / guidance_scale / max_t: + few-step sampler schedule (defaults match the canonical 4-step + shift=3 student); stashed in ``self.dmd_sampler_cfg``. + """ + from qwen_image_dmd2_sampler import DEFAULT_MAX_T, resolve_schedule + + try: + from diffusers import QwenImagePipeline, QwenImageTransformer2DModel + except ImportError as e: + raise ImportError( + "qwen-image-dmd2 requires a diffusers version providing QwenImagePipeline " + "and QwenImageTransformer2DModel; upgrade diffusers." + ) from e + + params = dict(self.config.extra_params) + student_path = params.get("student_path") + if not student_path: + raise ValueError( + "Missing required extra_param: student_path (the consolidated DMD2 student " + "dir, e.g. .../epoch_4_step_17999/model/consolidated)." + ) + base_pipeline_path = params.get("base_pipeline_path") or self.config.model_path + ema_path = params.get("ema_path") + + default_dtype = self.config.model_dtype["default"] + transformer_dtype = self.config.model_dtype.get("transformer", default_dtype) + if torch.float16 in (default_dtype, transformer_dtype): + self.logger.warning( + "Qwen-Image is trained/served in bfloat16; float16 (Half) can overflow the " + "VAE and produce NaNs. Consider --model-dtype BFloat16." + ) + + self.logger.info("Loading DMD2 student transformer from %s", student_path) + transformer = QwenImageTransformer2DModel.from_pretrained( + student_path, torch_dtype=transformer_dtype + ) + + if ema_path: + self.logger.info("Overlaying EMA shadow from %s", ema_path) + ema_state = torch.load(str(ema_path), map_location="cpu") + shadow = ( + ema_state.get("shadow", ema_state) if isinstance(ema_state, dict) else ema_state + ) + if not isinstance(shadow, dict): + raise ValueError( + f"ema_path content has unexpected type {type(shadow).__name__}; " + "expected dict[str, Tensor]." + ) + missing, unexpected = transformer.load_state_dict(shadow, strict=False) + if unexpected: + self.logger.warning("EMA overlay had %d unexpected key(s)", len(unexpected)) + if missing: + self.logger.warning("EMA overlay missed %d student key(s)", len(missing)) + + transformer.eval() + + self.logger.info( + "Loading base Qwen-Image pipeline from %s (transformer replaced by student)", + base_pipeline_path, + ) + pipe = QwenImagePipeline.from_pretrained( + base_pipeline_path, transformer=transformer, torch_dtype=default_dtype + ) + pipe.set_progress_bar_config(disable=True) + + # Resolve and stash the few-step sampler config. Defaults match the + # canonical 4-step shift=3 student; the schedule MUST match training. + sample_steps = params.get("sample_steps") + sample_steps = int(sample_steps) if sample_steps is not None else 4 + t_list = params.get("t_list") + if isinstance(t_list, str): + t_list = [float(x) for x in t_list.split(",") if x.strip()] + max_t = float(params.get("max_t", DEFAULT_MAX_T)) + schedule = resolve_schedule(t_list, sample_steps, max_t) + defaults = MODEL_DEFAULTS[self.config.model_type].get("inference_extra_args", {}) + self.dmd_sampler_cfg = { + "schedule": schedule, + "sample_type": str(params.get("sample_type", "ode")), + "guidance_scale": float(params.get("guidance_scale", 1.0)), + "negative_prompt": params.get("negative_prompt"), + "height": int(params.get("height", defaults.get("height", 1024))), + "width": int(params.get("width", defaults.get("width", 1024))), + "max_sequence_length": int(params.get("max_sequence_length", 512)), + } + self.logger.info( + "DMD2 few-step sampler: steps=%d schedule=%s sample_type=%s guidance_scale=%s " + "(schedule must match the student's training t_list)", + len(schedule) - 1, + schedule, + self.dmd_sampler_cfg["sample_type"], + self.dmd_sampler_cfg["guidance_scale"], + ) + return pipe + def print_quant_summary(self): for name, backbone in self.iter_backbones(): self.logger.info(f"{name} quantization info:") diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 1d71c088652..26aac8ae07f 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -50,9 +50,8 @@ QuantFormat, QuantizationConfig, ) -from utils import check_conv_and_mha, check_lora +from utils import check_conv_and_mha, check_lora, restore_quantizer_state, save_quantizer_state -import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.export import export_hf_checkpoint @@ -309,8 +308,11 @@ def save_checkpoint( filename = f"{backbone_name}.pt" if backbone_name else "backbone.pt" target_path = ckpt_path / filename - self.logger.info(f"Saving backbone to {target_path}") - mto.save(backbone, str(target_path)) + # Save ONLY the quantization state (recipe + quantizer buffers incl. amax), + # not the model weights. The weights live in the base HF/diffusers checkpoint + # and are reloaded there on restore; this keeps the artifact tiny. + self.logger.info(f"Saving quantizer state (amax + recipe, no weights) to {target_path}") + save_quantizer_state(backbone, str(target_path)) self.logger.info("Checkpoint saved successfully") @@ -380,7 +382,9 @@ def restore_checkpoint(self) -> None: f"Checkpoint not found for '{backbone_name}' in {restore_path}" ) self.logger.info(f"Restoring {backbone_name} from {source_path}") - mto.restore(backbone, str(source_path)) + # The pipeline was just created with the base (unquantized) weights, so + # this re-applies the quantization recipe + amax on top of them. + restore_quantizer_state(backbone, str(source_path)) self.logger.info("Checkpoints restored successfully") diff --git a/examples/diffusers/quantization/qwen_image_dmd2_sampler.py b/examples/diffusers/quantization/qwen_image_dmd2_sampler.py new file mode 100644 index 00000000000..518bc6b47d5 --- /dev/null +++ b/examples/diffusers/quantization/qwen_image_dmd2_sampler.py @@ -0,0 +1,260 @@ +# 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. + +"""Compact DMD2 few-step sampler for Qwen-Image students. + +This is a vendored, calibration-friendly version of the few-step unroll in +``examples/diffusers/fastgen/inference_dmd2_qwen_image.py``. It is kept here so +the quantization example is self-contained (no cross-example ``sys.path`` +imports) and so calibration can run the **same forward logic the student was +trained/served with** — which is what makes the collected ``amax`` statistics +representative. + +The single :func:`dmd2_sample` entry point serves two callers: + +* **Calibration** (``decode=False``): runs only the transformer forwards of the + DMD unroll and returns ``None``. The VAE / image post-processing is skipped + because quantization only needs the transformer's activation statistics, and + skipping the VAE saves substantial time and memory on the 60-layer student. +* **Sanity inference** (``decode=True``): additionally runs the VAE decode and + returns a list of images, used to confirm a restored (quantized) student + still produces a finite image. + +The math is bit-aligned with the training-time ``_build_student_input`` in +``modelopt/torch/fastgen/methods/dmd.py`` and with the inference reference: + + for (t_cur, t_next) in pairwise(t_list): + v = student(x, t=t_cur, text_emb) # flow at t_cur + x_0 = x - t_cur * v # RF identity -> x_0 estimate + if t_next > 0: + eps = (x - (1 - t_cur) * x_0) / t_cur # ODE: invert RF forward + x = (1 - t_next) * x_0 + t_next * eps # re-noise to t_next + else: + x = x_0 # final step + +``t_list`` MUST match the student's training schedule (e.g. the LightX2V +"shift=3" 4-step shape ``[1.0, 0.9, 0.75, 0.5, 0.0]``); a mismatch produces a +train/inference gap and therefore misleading calibration statistics. +""" + +from __future__ import annotations + +import itertools + +import torch +from diffusers.utils.torch_utils import randn_tensor + +# Canonical 4-step "shift=3" student schedule (LightX2V-Qwen-Image-Lightning +# shape). t_list has student_sample_steps + 1 entries: the first N are the +# timesteps the student is evaluated at, the trailing 0.0 is the terminal the +# final Euler step lands on (NOT an extra evaluation). +DEFAULT_T_LIST: tuple[float, ...] = (1.0, 0.9, 0.75, 0.5, 0.0) +DEFAULT_MAX_T: float = 0.999 + + +def resolve_schedule( + t_list: list[float] | tuple[float, ...] | None, + sample_steps: int | None, + max_t: float = DEFAULT_MAX_T, +) -> list[float]: + """Resolve the sampling schedule (timesteps + terminal 0.0). + + Priority: + 1. An explicit ``t_list`` (must end at 0.0 and have ``sample_steps + 1`` + entries when ``sample_steps`` is given). + 2. ``sample_steps == 1`` -> ``[max_t, 0.0]`` (canonical single-step). + 3. ``sample_steps == 4`` (or None) with no ``t_list`` -> ``DEFAULT_T_LIST``. + 4. Otherwise a linear ``linspace(max_t, 0, sample_steps + 1)`` fallback. + """ + if t_list is not None: + schedule = [float(t) for t in t_list] + if abs(schedule[-1]) > 1e-6: + raise ValueError( + f"t_list must end at 0.0 (got {schedule[-1]}); the final step lands on x_0." + ) + if sample_steps is not None and len(schedule) != sample_steps + 1: + raise ValueError( + f"t_list must have sample_steps+1 entries " + f"(got {len(schedule)} for sample_steps={sample_steps})." + ) + return schedule + + if sample_steps == 1: + return [float(max_t), 0.0] + if sample_steps in (None, 4): + return list(DEFAULT_T_LIST) + return torch.linspace(float(max_t), 0.0, sample_steps + 1).tolist() + + +@torch.no_grad() +def dmd2_sample( + pipe, + prompt: str | list[str], + *, + schedule: list[float], + sample_type: str = "ode", + guidance_scale: float = 1.0, + negative_prompt: str | list[str] | None = None, + height: int = 1024, + width: int = 1024, + num_images_per_prompt: int = 1, + generator: torch.Generator | None = None, + max_sequence_length: int = 512, + decode: bool = False, + output_type: str = "pil", +) -> list | None: + """Run the DMD few-step unroll on ``pipe.transformer``. + + Args: + pipe: A ``QwenImagePipeline`` whose ``transformer`` is the DMD2 student. + prompt: A prompt or list of prompts (one calibration batch). + schedule: Full timestep schedule incl. trailing 0.0 (see + :func:`resolve_schedule`). + sample_type: ``"ode"`` (deterministic, recover eps via RF identity) or + ``"sde"`` (fresh Gaussian noise between steps). Must match training. + guidance_scale: Inference-time CFG. Leave at ``1.0`` for students trained + with an internalised (non-null) ``dmd2.guidance_scale`` — passing + ``> 1.0`` there would double-apply CFG. + negative_prompt: Negative prompt for CFG; defaults to ``""`` when CFG is + engaged and none is given. + height/width: Output spatial size (must be VAE-compatible). + num_images_per_prompt: Images per prompt. + generator: Optional RNG for reproducible noise. + max_sequence_length: Text-encoder max sequence length. + decode: If ``True`` run VAE decode + post-process and return images. If + ``False`` (calibration) skip the VAE and return ``None``. + output_type: Passed to the image processor when ``decode=True``. + + Returns: + A list of images when ``decode=True``, else ``None``. + """ + if sample_type not in ("ode", "sde"): + raise ValueError(f"sample_type must be 'ode' or 'sde', got {sample_type!r}") + + do_cfg = guidance_scale != 1.0 + if do_cfg and negative_prompt is None: + negative_prompt = "" + + device = pipe.transformer.device + dtype = next(pipe.transformer.parameters()).dtype + + # ---- Encode prompt(s) ------------------------------------------------ + prompt_embeds, prompt_embeds_mask = pipe.encode_prompt( + prompt=prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + neg_prompt_embeds = neg_prompt_embeds_mask = None + if do_cfg: + neg_prompt_embeds, neg_prompt_embeds_mask = pipe.encode_prompt( + prompt=negative_prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + txt_seq_lens = ( + prompt_embeds_mask.sum(dim=1).int().tolist() if prompt_embeds_mask is not None else None + ) + neg_txt_seq_lens = ( + neg_prompt_embeds_mask.sum(dim=1).int().tolist() + if neg_prompt_embeds_mask is not None + else None + ) + + # ---- Build initial noisy latents at t = schedule[0] ------------------ + batch_size = (1 if isinstance(prompt, str) else len(prompt)) * num_images_per_prompt + num_channels_latents = pipe.transformer.config.in_channels // 4 # 64 // 4 = 16 + h_lat = 2 * (height // (pipe.vae_scale_factor * 2)) + w_lat = 2 * (width // (pipe.vae_scale_factor * 2)) + latent_shape = (batch_size, 1, num_channels_latents, h_lat, w_lat) + + noise = randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype) + latents_5d = noise * schedule[0] # RF: sigma(t0) = t0 + x_packed = pipe._pack_latents(latents_5d, batch_size, num_channels_latents, h_lat, w_lat) + img_shapes = [[(1, h_lat // 2, w_lat // 2)]] * batch_size + + # ---- DMD few-step unroll (transformer forwards) ---------------------- + for t_cur, t_next in itertools.pairwise(schedule): + timestep = torch.tensor([t_cur], device=device, dtype=dtype).expand(batch_size) + flow_packed = pipe.transformer( + hidden_states=x_packed, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_embeds_mask, + timestep=timestep, + img_shapes=img_shapes, + txt_seq_lens=txt_seq_lens, + guidance=None, + return_dict=False, + )[0] + if do_cfg: + neg_flow_packed = pipe.transformer( + hidden_states=x_packed, + encoder_hidden_states=neg_prompt_embeds, + encoder_hidden_states_mask=neg_prompt_embeds_mask, + timestep=timestep, + img_shapes=img_shapes, + txt_seq_lens=neg_txt_seq_lens, + guidance=None, + return_dict=False, + )[0] + flow_packed = ( + neg_flow_packed.to(torch.float64) + + float(guidance_scale) + * (flow_packed.to(torch.float64) - neg_flow_packed.to(torch.float64)) + ).to(dtype) + + # RF identity: x_0 = x_t - t_cur * v (fp64 for stability). + x0_packed = (x_packed.to(torch.float64) - float(t_cur) * flow_packed.to(torch.float64)).to( + dtype + ) + + if t_next > 1e-6: + if sample_type == "ode": + alpha_cur = 1.0 - float(t_cur) + eps_packed = ( + (x_packed.to(torch.float64) - alpha_cur * x0_packed.to(torch.float64)) + / max(float(t_cur), 1e-6) + ).to(dtype) + else: + eps_packed = torch.randn( + x_packed.shape, generator=generator, device=device, dtype=dtype + ) + alpha_next = 1.0 - float(t_next) + x_packed = ( + alpha_next * x0_packed.to(torch.float64) + + float(t_next) * eps_packed.to(torch.float64) + ).to(dtype) + else: + x_packed = x0_packed + + if not decode: + # Calibration path: transformer forwards already ran; nothing to decode. + return None + + # ---- VAE decode (sanity-inference path only) ------------------------- + x0_5d = pipe._unpack_latents(x_packed, height, width, pipe.vae_scale_factor) + latents_mean = ( + torch.tensor(pipe.vae.config.latents_mean) + .view(1, pipe.vae.config.z_dim, 1, 1, 1) + .to(device=device, dtype=dtype) + ) + latents_std = 1.0 / torch.tensor(pipe.vae.config.latents_std).view( + 1, pipe.vae.config.z_dim, 1, 1, 1 + ).to(device=device, dtype=dtype) + x0_scaled = x0_5d / latents_std + latents_mean + image_5d = pipe.vae.decode(x0_scaled, return_dict=False)[0] + image_4d = image_5d[:, :, 0] # Qwen-Image treats images as 1-frame videos + return pipe.image_processor.postprocess(image_4d, output_type=output_type) diff --git a/examples/diffusers/quantization/sanity_check_dmd2.py b/examples/diffusers/quantization/sanity_check_dmd2.py new file mode 100644 index 00000000000..9f78901f20c --- /dev/null +++ b/examples/diffusers/quantization/sanity_check_dmd2.py @@ -0,0 +1,175 @@ +# 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. + +"""Restore a quantized DMD2 Qwen-Image student and run one few-step inference. + +Confirms the round trip of the new ``qwen-image-dmd2`` quantization flow: + + 1. Load the base Qwen-Image pipeline with the consolidated student swapped in + (via the same :class:`PipelineManager` path quantize.py uses) -- this brings + the original (unquantized) weights. + 2. Reapply the weight-free quantization checkpoint saved by ``quantize.py`` + (``save_quantizer_state`` -> ``transformer.pt``) via + ``restore_quantizer_state``, which re-applies the quantizer recipe **and the + calibrated amax** buffers on top of the loaded weights. + 3. Run a single few-step DMD inference (with VAE decode) and assert the image + is finite and non-constant. + +This deliberately reuses :class:`PipelineManager` and +:func:`qwen_image_dmd2_sampler.dmd2_sample` so the inference path is identical to +calibration's (minus the VAE decode, which is enabled here). + +Usage:: + + python sanity_check_dmd2.py \\ + --quantized-ckpt ./qwen_dmd2_fp8/transformer.pt \\ + --student-path /.../epoch_4_step_17999/model/consolidated \\ + --base-pipeline-path /.../models/Qwen-Image \\ + --output-png ./qwen_dmd2_fp8/sanity.png +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys + +import torch +from models_utils import ModelType +from pipeline_manager import PipelineManager +from quantize_config import ModelConfig +from qwen_image_dmd2_sampler import dmd2_sample +from utils import restore_quantizer_state + +import modelopt.torch.quantization as mtq + +logger = logging.getLogger("sanity_check_dmd2") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--quantized-ckpt", + required=True, + help="Path to the quantized checkpoint saved by quantize.py (e.g. .../transformer.pt).", + ) + parser.add_argument( + "--student-path", + required=True, + help="Consolidated DMD2 student dir (provides architecture + base weights to restore into).", + ) + parser.add_argument( + "--base-pipeline-path", + default="Qwen/Qwen-Image", + help="Base Qwen-Image dir/HF id for the VAE / text-encoder / tokenizer / scheduler.", + ) + parser.add_argument("--ema-path", default=None, help="Optional EMA shadow overlaid on load.") + parser.add_argument("--output-png", default="./qwen_dmd2_sanity.png") + parser.add_argument("--prompt", default="a small red cube on a white table") + parser.add_argument("--height", type=int, default=1024) + parser.add_argument("--width", type=int, default=1024) + parser.add_argument("--seed", type=int, default=42) + # Few-step sampler knobs (defaults match the canonical 4-step shift=3 student). + parser.add_argument("--sample-steps", type=int, default=4) + parser.add_argument( + "--t-list", + default=None, + help="Comma-separated schedule incl. trailing 0.0, e.g. '1.0,0.9,0.75,0.5,0.0'.", + ) + parser.add_argument("--sample-type", default="ode", choices=["ode", "sde"]) + parser.add_argument("--guidance-scale", type=float, default=1.0) + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s" + ) + + # 1. Build the base pipeline with the student swapped in (unquantized). + extra_params: dict[str, object] = { + "student_path": args.student_path, + "base_pipeline_path": args.base_pipeline_path, + "sample_steps": args.sample_steps, + "sample_type": args.sample_type, + "guidance_scale": args.guidance_scale, + "height": args.height, + "width": args.width, + } + if args.ema_path: + extra_params["ema_path"] = args.ema_path + if args.t_list: + extra_params["t_list"] = args.t_list + + model_config = ModelConfig( + model_type=ModelType.QWEN_IMAGE_DMD2, + model_dtype={"default": torch.bfloat16}, + backbone=["transformer"], + extra_params=extra_params, + ) + pm = PipelineManager(model_config, logger) + pipe = pm.create_pipeline() + + # 2. Restore the quantized architecture + calibrated amax into the student. + logger.info( + "Restoring quantizer state (amax + recipe) from %s onto the loaded student", + args.quantized_ckpt, + ) + restore_quantizer_state(pipe.transformer, args.quantized_ckpt) + mtq.print_quant_summary(pipe.transformer) + pm.setup_device() + + # 3. One few-step inference (with VAE decode). + gen = torch.Generator(device=pipe.transformer.device).manual_seed(args.seed) + images = dmd2_sample(pipe, [args.prompt], decode=True, generator=gen, **pm.dmd_sampler_cfg) + image = images[0] + + import numpy as np + + arr = np.asarray(image) + stats = { + "prompt": args.prompt, + "quantized_ckpt": args.quantized_ckpt, + "schedule": pm.dmd_sampler_cfg["schedule"], + "image_shape": list(arr.shape), + "image_dtype": str(arr.dtype), + "image_min": float(arr.min()), + "image_max": float(arr.max()), + "image_mean": float(arr.mean()), + "image_std": float(arr.std()), + "is_finite": bool(np.isfinite(arr).all()), + "is_not_constant": bool(arr.std() > 0), + } + + os.makedirs(os.path.dirname(os.path.abspath(args.output_png)), exist_ok=True) + image.save(args.output_png) + with open(args.output_png.replace(".png", "_stats.json"), "w") as f: + json.dump(stats, f, indent=2) + print(json.dumps(stats, indent=2)) + + if not stats["is_finite"]: + logger.error("Sanity check FAILED: image contains non-finite values.") + sys.exit(1) + if not stats["is_not_constant"]: + logger.error("Sanity check FAILED: image is constant (std == 0).") + sys.exit(1) + logger.info( + "Sanity check PASSED: restored quantized student produced a finite image -> %s", + args.output_png, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index c3cfdcd5cdd..9a6b841a52e 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -24,8 +24,13 @@ from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear from diffusers.utils import load_image +import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.plugins.diffusion.diffusers import AttentionModuleMixin +from modelopt.torch.quantization.utils.core_utils import ( + get_quantizer_state_dict, + set_quantizer_state_dict, +) USE_PEFT = True try: @@ -193,3 +198,36 @@ def mha_filter_func(name): if hasattr(F, "scaled_dot_product_attention"): mtq.disable_quantizer(backbone, mha_filter_func) + + +def save_quantizer_state(model: torch.nn.Module, path: str) -> None: + """Save ONLY ModelOpt's quantization state -- the recipe plus the quantizer + buffers (amax, pre_quant_scale, ...) -- and NOT the model weights. + + This is the same idiom ModelOpt uses internally (see + ``modelopt.torch.quantization.plugins.transformers_trainer``): the + ``modelopt_state`` (architecture/recipe from :func:`mto.modelopt_state`) is + bundled with the per-quantizer state from + :func:`get_quantizer_state_dict` under the ``modelopt_state_weights`` key. + The resulting checkpoint is tiny (KBs-MBs) and is reloaded on top of the + original (unquantized) model via :func:`restore_quantizer_state`. + """ + modelopt_state = mto.modelopt_state(model) + modelopt_state["modelopt_state_weights"] = get_quantizer_state_dict(model) + torch.save(modelopt_state, str(path)) + + +def restore_quantizer_state(model: torch.nn.Module, path: str) -> torch.nn.Module: + """Reload a checkpoint written by :func:`save_quantizer_state` onto ``model``. + + ``model`` must already hold its original (unquantized) weights (e.g. freshly + loaded from the base HF/diffusers checkpoint); this re-applies the + quantization recipe and loads the calibrated amax/quantizer buffers on top. + Mirrors ModelOpt's ``_restore_modelopt_state_with_weights``. + """ + modelopt_state = mto.load_modelopt_state(str(path)) + quantizer_state = modelopt_state.pop("modelopt_state_weights", None) + mto.restore_from_modelopt_state(model, modelopt_state) + if quantizer_state is not None: + set_quantizer_state_dict(model, quantizer_state) + return model diff --git a/modelopt/torch/fastgen/plugins/qwen_image.py b/modelopt/torch/fastgen/plugins/qwen_image.py index 08a32b09301..3ae90b611ae 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image.py +++ b/modelopt/torch/fastgen/plugins/qwen_image.py @@ -48,9 +48,11 @@ from __future__ import annotations import contextlib +from importlib.metadata import PackageNotFoundError, version from typing import TYPE_CHECKING, Any import torch +from packaging.version import Version from torch import nn from ..methods.dmd import DMDPipeline @@ -58,6 +60,16 @@ if TYPE_CHECKING: from ..config import DMDConfig + +try: + # Diffusers 0.35/0.36 requires explicit Python sequence lengths in Qwen's + # positional-embedding path. Starting in 0.37 the mask is authoritative and + # passing txt_seq_lens is deprecated. Keep the plugin compatible with the + # broader diffusers versions supported by ModelOpt. + _DIFFUSERS_NEEDS_TXT_SEQ_LENS = Version(version("diffusers")) < Version("0.37.0") +except PackageNotFoundError: # pragma: no cover - optional plugin import guard + _DIFFUSERS_NEEDS_TXT_SEQ_LENS = False + __all__ = [ "QwenImageDMDPipeline", "attach_feature_capture", @@ -65,6 +77,7 @@ "pack_latents", "remove_feature_capture", "unpack_latents", + "update_feature_capture_shape", ] @@ -206,6 +219,7 @@ def _call_model( packed = pack_latents(hidden_states) img_shapes = build_img_shapes(b, h, w) + update_feature_capture_shape(model, h, w) call_kwargs: dict[str, Any] = dict(model_kwargs) call_kwargs.pop("hidden_states", None) @@ -214,8 +228,10 @@ def _call_model( call_kwargs.pop("guidance", None) call_kwargs.pop("return_dict", None) txt_seq_lens = call_kwargs.pop("txt_seq_lens", None) - if txt_seq_lens is None and encoder_hidden_states_mask is not None: - txt_seq_lens = encoder_hidden_states_mask.sum(dim=1).int().tolist() + if _DIFFUSERS_NEEDS_TXT_SEQ_LENS: + if txt_seq_lens is None and encoder_hidden_states_mask is not None: + txt_seq_lens = encoder_hidden_states_mask.sum(dim=1).int().tolist() + call_kwargs["txt_seq_lens"] = txt_seq_lens guidance = None if self._guidance_value is not None: @@ -232,7 +248,6 @@ def _call_model( encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_mask=encoder_hidden_states_mask, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, guidance=guidance, return_dict=False, **call_kwargs, @@ -266,6 +281,16 @@ def _call_model( _SHAPE_ATTR = "_fastgen_capture_shape" +def update_feature_capture_shape(model: nn.Module, h_lat: int, w_lat: int) -> None: + """Refresh a hooked teacher's target shape for the current multiresolution batch.""" + if h_lat % 2 or w_lat % 2: + raise ValueError( + f"feature capture requires even latent dims, got h_lat={h_lat}, w_lat={w_lat}." + ) + if hasattr(model, _HANDLES_ATTR): + setattr(model, _SHAPE_ATTR, (h_lat // 2, w_lat // 2)) + + def attach_feature_capture( teacher: nn.Module, feature_indices: list[int], @@ -336,8 +361,6 @@ def attach_feature_capture( setattr(teacher, _SHAPE_ATTR, (h_lat // 2, w_lat // 2)) handles: list[Any] = [] - h_half = h_lat // 2 - w_half = w_lat // 2 for idx in sorted_indices: block = blocks[idx] @@ -354,12 +377,13 @@ def _hook(_module: nn.Module, _inputs: Any, output: Any) -> None: ) # hidden: [B, num_image_patches, C] -> [B, C, H_half, W_half]. b, s, c = hidden.shape + h_half, w_half = getattr(teacher, _SHAPE_ATTR) expected_s = h_half * w_half if s != expected_s: raise RuntimeError( f"QwenImage feature-capture got hidden_states seq_len={s} but expected " f"{expected_s} = (h_lat // 2) * (w_lat // 2). Did the input resolution " - f"drift from the attach_feature_capture-time setting?" + "drift from the current batch setting?" ) feat = hidden.permute(0, 2, 1).reshape(b, c, h_half, w_half) captured.append(feat) diff --git a/modelopt/torch/quantization/conversion.py b/modelopt/torch/quantization/conversion.py index 00187d291c0..6eb26cb85cd 100644 --- a/modelopt/torch/quantization/conversion.py +++ b/modelopt/torch/quantization/conversion.py @@ -184,8 +184,13 @@ def create_and_replace_svdquant_linear_on_the_fly(model): def restore_svdquant_model(model: nn.Module, config: QuantizeConfig, metadata: MetadataDict): - """Restore the svdquant states from the given state dict.""" + """Restore SVDQuant and rebuild its PEFT adapter topology when present.""" create_and_replace_svdquant_linear_on_the_fly(model) + peft_metadata = metadata.get("svdquant_peft") + if peft_metadata is not None: + from .plugins.svdquant_peft import _restore_svdquant_peft + + _restore_svdquant_peft(model, peft_metadata) restore_quantizer_state(model, config, metadata) return model diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index c096aaeb00e..57b3587c445 100644 --- a/modelopt/torch/quantization/mode.py +++ b/modelopt/torch/quantization/mode.py @@ -538,6 +538,35 @@ def config_class(self) -> type[QuantizeAlgorithmConfig]: # root model, which is not present when layerwise_calibrate dispatches per decoder layer. _supports_layerwise = False + @property + def convert(self) -> ConvertEntrypoint: + """Calibrate SVDQuant and move its low-rank branch into HF PEFT. + + This applies to every ``algorithm.method=svdquant`` recipe, independent + of the quantization format. + """ + + def wrapped_func(model, config, forward_loop=None): + from .plugins.svdquant_peft import _externalize_svdquant_lora + + rank = config.lowrank + model, _ = wrapped_calib_func( + model, + config, + forward_loop, + func=self.__class__._calib_func, + supports_layerwise=self.__class__._supports_layerwise, + ) + + peft_metadata = _externalize_svdquant_lora(model, rank) + metadata = {} + update_quantize_metadata(model, config, metadata) + if peft_metadata is not None: + metadata["svdquant_peft"] = peft_metadata + return model, metadata + + return wrapped_func + @property def restore(self) -> RestoreEntrypoint: """The mode's entrypoint for restoring a model.""" diff --git a/modelopt/torch/quantization/nn/modules/quant_linear.py b/modelopt/torch/quantization/nn/modules/quant_linear.py index da1b79a2f60..79a464faa39 100644 --- a/modelopt/torch/quantization/nn/modules/quant_linear.py +++ b/modelopt/torch/quantization/nn/modules/quant_linear.py @@ -136,22 +136,23 @@ def _apply_pre_quant_scale(self, input: torch.Tensor): def _compute_lora_residual(self, input: torch.Tensor): """Compute the LoRA residual if present, otherwise return None.""" + lora_a_weight = getattr(self.weight_quantizer, "svdquant_lora_a", None) + lora_b_weight = getattr(self.weight_quantizer, "svdquant_lora_b", None) if ( self._not_sequential_quantizers() - and self.weight_quantizer.svdquant_lora_a is not None - and self.weight_quantizer.svdquant_lora_b is not None + and lora_a_weight is not None + and lora_b_weight is not None ): - lora_a = F.linear(input, weight=self.weight_quantizer.svdquant_lora_a) - lora_b = F.linear(lora_a, weight=self.weight_quantizer.svdquant_lora_b) + lora_a = F.linear(input, weight=lora_a_weight) + lora_b = F.linear(lora_a, weight=lora_b_weight) return lora_b return None def forward(self, input, *args, **kwargs): """SVDQuant layer forward function.""" - has_svdquant_lora = ( - self._not_sequential_quantizers() - and self.weight_quantizer.svdquant_lora_a is not None - and self.weight_quantizer.svdquant_lora_b is not None + has_svdquant_lora = self._not_sequential_quantizers() and all( + getattr(self.weight_quantizer, factor_name, None) is not None + for factor_name in ("svdquant_lora_a", "svdquant_lora_b") ) if has_svdquant_lora: input = self._apply_pre_quant_scale(input) @@ -170,15 +171,10 @@ def fold_weight(self, keep_attrs: bool = False): and hasattr(self, "weight") and self.weight_quantizer.fake_quant ): - if ( - self._not_sequential_quantizers() - and self.weight_quantizer.svdquant_lora_a is not None - and self.weight_quantizer.svdquant_lora_b is not None - ): - self.weight.data.copy_( - self.weight - + self.weight_quantizer.svdquant_lora_b @ self.weight_quantizer.svdquant_lora_a - ) + lora_a = getattr(self.weight_quantizer, "svdquant_lora_a", None) + lora_b = getattr(self.weight_quantizer, "svdquant_lora_b", None) + if self._not_sequential_quantizers() and lora_a is not None and lora_b is not None: + self.weight.data.copy_(self.weight + lora_b @ lora_a) if not keep_attrs: _attrs = [ "_svdquant_lora_a", diff --git a/modelopt/torch/quantization/plugins/diffusion/diffusers.py b/modelopt/torch/quantization/plugins/diffusion/diffusers.py index f2f6a702479..fdb5e3443c7 100644 --- a/modelopt/torch/quantization/plugins/diffusion/diffusers.py +++ b/modelopt/torch/quantization/plugins/diffusion/diffusers.py @@ -142,9 +142,16 @@ def _quantized_sdpa(self, *args, **kwargs): k_quantized_scale = self.k_bmm_quantizer._get_amax(key) v_quantized_scale = self.v_bmm_quantizer._get_amax(value) - # We don't need to calibrate the output of softmax - return self.bmm2_output_quantizer( - fp8_sdpa( + # We don't need to calibrate the output of softmax. + # ``FP8SDPA`` is an export-only autograd Function: it exists solely to attach the ONNX + # ``symbolic`` (export_fp8_mha), and its forward is just + # ``original_scaled_dot_product_attention``. It implements no ``backward``, so routing + # through it at runtime makes quantized attention non-differentiable and breaks training + # (QAT) -- ``loss.backward()`` raises "must implement either the backward or vjp method". + # Use it only during ONNX export; at runtime call SDPA directly (identical forward math, + # with q/k/v already fake-quantized above) so autograd works. + if torch.onnx.is_in_onnx_export(): + attn_output = fp8_sdpa( query, key, value, @@ -157,7 +164,19 @@ def _quantized_sdpa(self, *args, **kwargs): else "Half", self._disable_fp8_mha if hasattr(self, "_disable_fp8_mha") else True, ) - ) + else: + # Pass attn_mask/dropout_p/is_causal/scale as keywords (``scale`` is keyword-only in + # recent torch), mirroring FP8SDPA.forward's own call to SDPA. + attn_output = original_scaled_dot_product_attention( + query, + key, + value, + attn_mask=param_dict["attn_mask"], + dropout_p=param_dict["dropout_p"], + is_causal=param_dict["is_causal"], + scale=param_dict["scale"], + ) + return self.bmm2_output_quantizer(attn_output) class _QuantAttention(_QuantFunctionalMixin): diff --git a/modelopt/torch/quantization/plugins/svdquant_peft.py b/modelopt/torch/quantization/plugins/svdquant_peft.py new file mode 100644 index 00000000000..de2d1833251 --- /dev/null +++ b/modelopt/torch/quantization/plugins/svdquant_peft.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hugging Face PEFT ownership for the trainable SVDQuant residual.""" + +import re +from typing import Any + +import torch +import torch.nn as nn +from peft import LoraConfig, inject_adapter_in_model +from peft.tuners.lora.layer import Linear as LoraLinear + +from ..nn import SVDQuantLinear + +__all__ = [] + +_SVDQUANT_ADAPTER_NAME = "modelopt_svdquant" + + +class _SVDQuantPeftLinear(LoraLinear): + """PEFT layer implementing ``Q(W_residual)x + B(Ax)`` for SVDQuant.""" + + def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + """Apply the AWQ pre-scale exactly once to both SVDQuant branches.""" + base_layer = self.get_base_layer() + scaled_x = base_layer._apply_pre_quant_scale(x) + with base_layer.input_quantizer.disable_pre_quant_scale(): + return super().forward(scaled_x, *args, **kwargs) + + +def _svdquant_peft_config(target_names: list[str], rank: int) -> LoraConfig: + target_regex = "^(?:" + "|".join(re.escape(name) for name in target_names) + ")$" + config = LoraConfig( + task_type=None, + r=rank, + lora_alpha=rank, + lora_dropout=0.0, + target_modules=target_regex, + bias="none", + lora_bias=False, + modules_to_save=None, + fan_in_fan_out=False, + use_rslora=False, + use_dora=False, + init_lora_weights=False, + inference_mode=False, + ) + config._register_custom_module({SVDQuantLinear: _SVDQuantPeftLinear}) + return config + + +def _delete_quantizer_svdquant_factors(weight_quantizer: nn.Module) -> None: + """Remove calibration factors stored as buffers or plain quantizer attributes.""" + for public_name in ("svdquant_lora_a", "svdquant_lora_b"): + for storage_name in (f"_{public_name}", public_name): + if ( + storage_name in weight_quantizer.__dict__ + or storage_name in weight_quantizer._buffers + ): + delattr(weight_quantizer, storage_name) + break + + +def _inject_svdquant_peft( + model: nn.Module, + target_names: list[str], + rank: int, + factors: dict[str, tuple[torch.Tensor, torch.Tensor]] | None, +) -> dict[str, Any]: + target_names = sorted(target_names) + base_trainability = [(parameter, parameter.requires_grad) for parameter in model.parameters()] + config = _svdquant_peft_config(target_names, rank) + inject_adapter_in_model( + config, + model, + adapter_name=_SVDQUANT_ADAPTER_NAME, + low_cpu_mem_usage=False, + ) + for parameter, requires_grad in base_trainability: + parameter.requires_grad_(requires_grad) + + for name in target_names: + module = model.get_submodule(name) + lora_a = module.lora_A[_SVDQUANT_ADAPTER_NAME].weight + lora_b = module.lora_B[_SVDQUANT_ADAPTER_NAME].weight + if factors is not None: + source_a, source_b = factors[name] + with torch.no_grad(): + lora_a.copy_(source_a.to(device=lora_a.device, dtype=lora_a.dtype)) + lora_b.copy_(source_b.to(device=lora_b.device, dtype=lora_b.dtype)) + _delete_quantizer_svdquant_factors(module.get_base_layer().weight_quantizer) + lora_a.requires_grad_(True) + lora_b.requires_grad_(True) + + return { + "rank": rank, + "target_modules": target_names, + } + + +def _externalize_svdquant_lora(model: nn.Module, rank: int) -> dict[str, Any] | None: + """Move calibrated SVDQuant factors from weight quantizers into HF PEFT.""" + factors = {} + for name, module in model.named_modules(): + if not isinstance(module, SVDQuantLinear): + continue + lora_a = getattr(module.weight_quantizer, "svdquant_lora_a", None) + lora_b = getattr(module.weight_quantizer, "svdquant_lora_b", None) + if lora_a is not None and lora_b is not None: + factors[name] = (lora_a.detach(), lora_b.detach()) + if not factors: + return None + return _inject_svdquant_peft(model, list(factors), rank, factors) + + +def _restore_svdquant_peft(model: nn.Module, metadata: dict[str, Any]) -> None: + """Rebuild the PEFT topology before the complete model state is loaded.""" + _inject_svdquant_peft( + model, + metadata["target_modules"], + metadata["rank"], + factors=None, + ) diff --git a/modelopt/torch/quantization/plugins/transformers_trainer.py b/modelopt/torch/quantization/plugins/transformers_trainer.py index 981f3d990d4..64e56e34228 100644 --- a/modelopt/torch/quantization/plugins/transformers_trainer.py +++ b/modelopt/torch/quantization/plugins/transformers_trainer.py @@ -303,13 +303,23 @@ def prediction_step(self, *args, **kwargs): def evaluate(self, *args, **kwargs): """Evaluate the model.""" + if self.quant_cfg is not None and not is_quantized(self.model): + self._quantize_model() if self.args.do_eval and not self.args.do_train and self.accelerator.is_fsdp2: # [Not related to ModelOpt] HF does not support eval only for FSDP2. self.model = self._prepare_model(self.model) return super().evaluate(*args, **kwargs) + def predict(self, *args, **kwargs): + """Run prediction.""" + if self.quant_cfg is not None and not is_quantized(self.model): + self._quantize_model() + return super().predict(*args, **kwargs) + def train(self, *args, **kwargs): """Train the model.""" + if self.quant_cfg is not None and not is_quantized(self.model): + self._quantize_model() outputs = super().train(*args, **kwargs) print_rank_0( "Training completed. Please save the final model using `Trainer.save_model()` to preserve ModelOpt states." diff --git a/tests/_test_utils/torch/quantization/quantize_common.py b/tests/_test_utils/torch/quantization/quantize_common.py index 46259203b24..b79e54f5f3e 100644 --- a/tests/_test_utils/torch/quantization/quantize_common.py +++ b/tests/_test_utils/torch/quantization/quantize_common.py @@ -72,8 +72,13 @@ def forward_loop(model, run_backward=False): if compress: mtq.compress(model) - for module in model.modules(): - assert not isinstance(module, torch.nn.Linear) or is_quantized_linear(module) + for name, module in model.named_modules(): + is_peft_adapter = ".lora_A." in name or ".lora_B." in name + assert ( + not isinstance(module, torch.nn.Linear) + or is_quantized_linear(module) + or is_peft_adapter + ) model.train() diff --git a/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py b/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py new file mode 100644 index 00000000000..2747012ec59 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 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. + +"""Regression test for the DMD2 QAT (restore-only) quantizer-state restore. + +The fastgen QAT path NEVER calibrates during training -- it RESTORES a ModelOpt quantizer +state (recipe + frozen amax) saved by the calibration example and re-applies it on every +(re)start. This test pins the guarantee the recipe depends on, on a tiny CPU model (no +GPU, milliseconds): ``dmd2_recipe.restore_quantizer_state`` onto a *fresh* model with +DIFFERENT weights reproduces the amax bit-identically and leaves that model's weights +untouched -- i.e. amax stays exactly as calibrated and the warm-started student weights +are preserved. + +The on-disk state is built here with ModelOpt's own idiom (the same one the calibration +example's ``--quantized-torch-ckpt-save-path`` uses), so the test also pins format +compatibility with that file. + +Dependency-guarded with ``importorskip`` so it skips where torch / modelopt are absent. +""" + +from __future__ import annotations + +import pathlib +import sys + +import pytest + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) + +torch = pytest.importorskip("torch") +mtq = pytest.importorskip("modelopt.torch.quantization") +mto = pytest.importorskip("modelopt.torch.opt") +dmd2_recipe = pytest.importorskip("dmd2_recipe") + +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.utils.core_utils import get_quantizer_state_dict + + +def _tiny_model(seed: int) -> torch.nn.Module: + torch.manual_seed(seed) + return torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.ReLU(), torch.nn.Linear(8, 4)) + + +def _amax_by_name(model: torch.nn.Module) -> dict[str, torch.Tensor]: + return { + name: module.amax.detach().clone() + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and module.amax is not None + } + + +def test_restore_quantizer_state_is_bit_identical_and_weight_free(tmp_path): + # Calibrate a tiny model (this is the ONLY place quantize/calibration happens -- the + # calibration example; the trainer never does this) and save its quantizer state in the + # weight-free format the calibration example writes (mto.modelopt_state + amax). + model = _tiny_model(seed=0) + calib = torch.randn(16, 8) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(calib)) + src_amax = _amax_by_name(model) + assert src_amax, "expected at least one calibrated TensorQuantizer amax" + + state = mto.modelopt_state(model) + state["modelopt_state_weights"] = get_quantizer_state_dict(model) + path = tmp_path / "transformer.pt" + torch.save(state, str(path)) + + # Restore onto a FRESH model with DIFFERENT weights; amax must come back + # bit-identically and the fresh model's weights must be untouched. + fresh = _tiny_model(seed=999) + before = {n: p.detach().clone() for n, p in fresh.named_parameters()} + dmd2_recipe.restore_quantizer_state(fresh, str(path)) + + restored_amax = _amax_by_name(fresh) + assert set(restored_amax) == set(src_amax) + for name, amax in src_amax.items(): + assert torch.equal(restored_amax[name], amax), f"amax mismatch at {name}" + + for n, p in fresh.named_parameters(): + if n in before: + assert torch.equal(p.detach(), before[n]), f"restore changed weight {n}" diff --git a/tests/unit/torch/fastgen/test_qwen_image_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_plugin.py index 498b6ce5f9f..07bd75fdce4 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_plugin.py @@ -33,6 +33,7 @@ from torch import nn from modelopt.torch.fastgen import DMDConfig +from modelopt.torch.fastgen.plugins import qwen_image as qwen_image_plugin from modelopt.torch.fastgen.plugins.qwen_image import ( QwenImageDMDPipeline, build_img_shapes, @@ -162,12 +163,13 @@ def _make_pipeline(student: nn.Module) -> QwenImageDMDPipeline: ) -def test_call_model_forwards_qwen_kwargs(): +def test_call_model_forwards_qwen_kwargs(monkeypatch): """``_call_model`` must forward the exact Qwen signature (hidden_states packed to ``[B, num_patches, 64]``, encoder_hidden_states verbatim, - encoder_hidden_states_mask verbatim, txt_seq_lens derived from the mask, + encoder_hidden_states_mask verbatim (Diffusers derives sequence lengths from it), img_shapes as ``[[(1, h//2, w//2)]] * B``, guidance=None, return_dict=False, timestep verbatim with no /1000 rescale).""" + monkeypatch.setattr(qwen_image_plugin, "_DIFFUSERS_NEEDS_TXT_SEQ_LENS", False) b, c, h, w = 2, 16, 32, 32 student = _CapturingModel(out_shape=(b, (h // 2) * (w // 2), c * 4), style="tensor") pipe = _make_pipeline(student) @@ -191,7 +193,7 @@ def test_call_model_forwards_qwen_kwargs(): assert tuple(kw["hidden_states"].shape) == (b, (h // 2) * (w // 2), c * 4) assert tuple(kw["encoder_hidden_states"].shape) == (b, 512, 3584) assert torch.equal(kw["encoder_hidden_states_mask"], mask) - assert kw["txt_seq_lens"] == [37, 42] + assert "txt_seq_lens" not in kw assert kw["img_shapes"] == [[(1, h // 2, w // 2)]] * b assert kw["guidance"] is None assert kw["return_dict"] is False @@ -199,6 +201,27 @@ def test_call_model_forwards_qwen_kwargs(): assert tuple(out.shape) == (b, c, h, w) +def test_call_model_forwards_legacy_txt_seq_lens(monkeypatch): + """Diffusers 0.35/0.36 still needs lengths derived from the attention mask.""" + monkeypatch.setattr(qwen_image_plugin, "_DIFFUSERS_NEEDS_TXT_SEQ_LENS", True) + b, c, h, w = 2, 16, 8, 8 + student = _CapturingModel(out_shape=(b, (h // 2) * (w // 2), c * 4)) + pipe = _make_pipeline(student) + mask = torch.zeros(b, 9, dtype=torch.long) + mask[0, :4] = 1 + mask[1, :7] = 1 + + pipe._call_model( + student, + torch.randn(b, c, h, w), + torch.tensor([0.25, 0.5]), + encoder_hidden_states=torch.randn(b, 9, 32), + encoder_hidden_states_mask=mask, + ) + + assert student.last_kwargs["txt_seq_lens"] == [4, 7] + + @pytest.mark.parametrize("style", ["tensor", "tuple", "sample"]) def test_call_model_unpacks_return_styles(style): """``_call_model`` must unpack ``tensor`` / ``tuple`` / ``.sample`` return diff --git a/tests/unit/torch/quantization/plugins/test_svdquant_peft_modelopt.py b/tests/unit/torch/quantization/plugins/test_svdquant_peft_modelopt.py new file mode 100644 index 00000000000..a0697b8e9cc --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_svdquant_peft_modelopt.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the ModelOpt checkpoint contract of PEFT-backed SVDQuant.""" + +import copy +import io + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +peft = pytest.importorskip("peft", minversion="0.17.0") + +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins.svdquant_peft import ( + _SVDQUANT_ADAPTER_NAME, + _SVDQuantPeftLinear, +) + + +class _TinyMLP(nn.Module): + def __init__(self): + super().__init__() + self.svdquant = nn.Linear(8, 8, bias=False) + self.skipped = nn.Linear(8, 8, bias=False) + self.disabled = nn.Linear(8, 8, bias=False) + + def forward(self, x): + x = F.silu(self.svdquant(x)) + x = F.silu(self.skipped(x)) + return self.disabled(x) + + +def _svdquant_config(*, select_one_target=False): + config = copy.deepcopy(mtq.INT8_SMOOTHQUANT_CFG) + config["algorithm"] = { + "method": "svdquant", + "lowrank": 4, + "skip_layers": ["skipped"] if select_one_target else None, + } + if select_one_target: + # A disabled weight quantizer must not cause PEFT injection either. + config["quant_cfg"].append({"quantizer_name": "disabled.weight_quantizer", "enable": False}) + return config + + +def _quantize(model, *, select_one_target=False): + reference = next(model.parameters()) + calibration_input = torch.randn(4, 8, device=reference.device, dtype=reference.dtype) + return mtq.quantize( + model, + _svdquant_config(select_one_target=select_one_target), + forward_loop=lambda current: current(calibration_input), + ) + + +def _factor_state(model): + return { + name: parameter.detach().clone() + for name, parameter in model.named_parameters() + if ".lora_A." in name or ".lora_B." in name + } + + +def _svdquant_metadata(model): + matches = [ + mode_state["metadata"]["svdquant_peft"] + for mode_name, mode_state in mto.modelopt_state(model)["modelopt_state_dict"] + if mode_name == "svdquant_calibrate" + ] + assert len(matches) == 1 + return matches[0] + + +def test_svdquant_peft_uses_exact_targets_and_preserves_unquantized_forward(): + torch.manual_seed(17) + original = _TinyMLP() + model = copy.deepcopy(original) + model = _quantize(model, select_one_target=True) + + assert isinstance(model.svdquant, _SVDQuantPeftLinear) + assert not isinstance(model.skipped, _SVDQuantPeftLinear) + assert not isinstance(model.disabled, _SVDQuantPeftLinear) + assert _svdquant_metadata(model)["target_modules"] == ["svdquant"] + + adapter = model.svdquant + base_layer = adapter.get_base_layer() + assert base_layer.weight_quantizer.svdquant_lora_a is None + assert base_layer.weight_quantizer.svdquant_lora_b is None + assert adapter.scaling[_SVDQUANT_ADAPTER_NAME] == 1.0 + assert list(adapter.active_adapters) == [_SVDQUANT_ADAPTER_NAME] + assert adapter.lora_A[_SVDQUANT_ADAPTER_NAME].weight.requires_grad + assert adapter.lora_B[_SVDQUANT_ADAPTER_NAME].weight.requires_grad + assert not model.disabled.weight_quantizer.is_enabled + + # Both branches see the same AWQ-scaled input, while output quantization applies only + # to the residual-weight branch, matching the historical SVDQuant contract. + probe = torch.randn(3, 8) + scaled_probe = base_layer._apply_pre_quant_scale(probe) + with base_layer.input_quantizer.disable_pre_quant_scale(): + raw_base_output = base_layer(scaled_probe) + lora_output = F.linear( + F.linear(scaled_probe, adapter.lora_A[_SVDQUANT_ADAPTER_NAME].weight), + adapter.lora_B[_SVDQUANT_ADAPTER_NAME].weight, + ) + base_layer.output_quantizer.amax = raw_base_output.detach().abs().amax().mul(0.5) + base_layer.output_quantizer.enable() + expected = base_layer.output_quantizer(raw_base_output) + lora_output + combined_quantized = base_layer.output_quantizer(raw_base_output + lora_output) + actual = adapter(probe) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert not torch.equal(actual, combined_quantized) + + # With Q/DQ disabled, W_residual + BA must reproduce the original BF16/FP32 model. + for module in model.modules(): + if isinstance(module, TensorQuantizer): + module.disable() + torch.testing.assert_close(model(probe), original(probe), rtol=2e-5, atol=2e-5) + + bf16_model = _quantize(_TinyMLP().to(torch.bfloat16), select_one_target=True) + bf16_adapter = bf16_model.svdquant + assert bf16_adapter.lora_A[_SVDQUANT_ADAPTER_NAME].weight.dtype == torch.bfloat16 + assert bf16_adapter.lora_B[_SVDQUANT_ADAPTER_NAME].weight.dtype == torch.bfloat16 + + +def test_mto_save_restore_preserves_qat_mutated_factors(): + torch.manual_seed(19) + model = _quantize(_TinyMLP()) + initial_factors = _factor_state(model) + assert initial_factors + + factor_parameters = [ + parameter + for name, parameter in model.named_parameters() + if ".lora_A." in name or ".lora_B." in name + ] + assert factor_parameters and all(parameter.requires_grad for parameter in factor_parameters) + optimizer = torch.optim.SGD(factor_parameters, lr=5e-2) + optimizer.zero_grad(set_to_none=True) + model(torch.randn(5, 8)).square().mean().backward() + assert all(parameter.grad is not None for parameter in factor_parameters) + optimizer.step() + + trained_factors = _factor_state(model) + assert any( + not torch.equal(initial_factors[name], trained_factors[name]) + for name in trained_factors + if ".lora_A." in name + ) + assert any( + not torch.equal(initial_factors[name], trained_factors[name]) + for name in trained_factors + if ".lora_B." in name + ) + + probe = torch.randn(2, 8) + expected = model(probe).detach().clone() + checkpoint = io.BytesIO() + mto.save(model, checkpoint) + checkpoint.seek(0) + restored = mto.restore(_TinyMLP(), checkpoint) + + restored_factors = _factor_state(restored) + assert trained_factors.keys() == restored_factors.keys() + for name in trained_factors: + assert torch.equal(trained_factors[name], restored_factors[name]), name + assert all( + parameter.requires_grad + for name, parameter in restored.named_parameters() + if ".lora_A." in name or ".lora_B." in name + ) + assert restored.peft_config[_SVDQUANT_ADAPTER_NAME].inference_mode is False + for module in restored.modules(): + if isinstance(module, _SVDQuantPeftLinear): + assert list(module.active_adapters) == [_SVDQUANT_ADAPTER_NAME] + torch.testing.assert_close(restored(probe), expected, rtol=0, atol=0) + + # Optimizers are application-owned, but restored PEFT parameters must be discoverable + # before a resume optimizer is constructed. + resume_optimizer = torch.optim.SGD( + (parameter for parameter in restored.parameters() if parameter.requires_grad), lr=1e-3 + ) + optimizer_ids = { + id(parameter) for group in resume_optimizer.param_groups for parameter in group["params"] + } + assert all( + id(parameter) in optimizer_ids + for name, parameter in restored.named_parameters() + if ".lora_A." in name or ".lora_B." in name + ) + + # The split ModelOpt API follows the same topology-then-state lifecycle. + modelopt_state = copy.deepcopy(mto.modelopt_state(model)) + model_state = copy.deepcopy(model.state_dict()) + manually_restored = mto.restore_from_modelopt_state(_TinyMLP(), modelopt_state) + manually_restored.load_state_dict(model_state) + for name, tensor in trained_factors.items(): + assert torch.equal(_factor_state(manually_restored)[name], tensor), name + assert all( + parameter.requires_grad + for name, parameter in manually_restored.named_parameters() + if ".lora_A." in name or ".lora_B." in name + ) + torch.testing.assert_close(manually_restored(probe), expected, rtol=0, atol=0) diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index b609761f12a..713c83f8b05 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -403,6 +403,7 @@ def test_postprocess_amax(): def test_svdquant_lora_weights(): + pytest.importorskip("peft", minversion="0.17.0") model = _SimpleMLP(64, 64, 64, 64) quant_config = mtq.INT8_SMOOTHQUANT_CFG.copy() @@ -410,15 +411,17 @@ def test_svdquant_lora_weights(): mtq.quantize(model, quant_config, partial(forward_loop, dataloader=[torch.randn(2, 64, 64)])) - for module in model.modules(): - if isinstance(module, torch.nn.Linear): - assert module.weight_quantizer.svdquant_lora_a is not None - assert module.weight_quantizer.svdquant_lora_b is not None - - lora_residual = ( - module.weight_quantizer.svdquant_lora_b @ module.weight_quantizer.svdquant_lora_a - ) - assert lora_residual.shape == module.weight.shape + adapters = [ + module for module in model.modules() if "modelopt_svdquant" in getattr(module, "lora_A", {}) + ] + assert adapters + for adapter in adapters: + lora_a = adapter.lora_A["modelopt_svdquant"].weight + lora_b = adapter.lora_B["modelopt_svdquant"].weight + lora_residual = lora_b @ lora_a + assert lora_residual.shape == adapter.get_base_layer().weight.shape + assert adapter.get_base_layer().weight_quantizer.svdquant_lora_a is None + assert adapter.get_base_layer().weight_quantizer.svdquant_lora_b is None def test_layerwise_calibrate_support_gate():