From 1678e3629bcfc900692d9bef355e3ed9ed6dc202 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Wed, 5 Aug 2026 00:20:09 -0700 Subject: [PATCH 1/3] Add FastGen quantization-aware distillation Signed-off-by: Jingyu Xin --- examples/diffusers/fastgen/qad/README.md | 174 +++++ examples/diffusers/fastgen/qad/__init__.py | 16 + examples/diffusers/fastgen/qad/artifacts.py | 428 +++++++++++++ .../fastgen/qad/configs/qwen_image_nvfp4.yaml | 108 ++++ .../configs/qwen_image_svdquant_nvfp4.yaml | 112 ++++ examples/diffusers/fastgen/qad/finetune.py | 43 ++ examples/diffusers/fastgen/qad/modeling.py | 259 ++++++++ examples/diffusers/fastgen/qad/pipeline.py | 74 +++ examples/diffusers/fastgen/qad/recipe.py | 595 ++++++++++++++++++ 9 files changed, 1809 insertions(+) create mode 100644 examples/diffusers/fastgen/qad/README.md create mode 100644 examples/diffusers/fastgen/qad/__init__.py create mode 100644 examples/diffusers/fastgen/qad/artifacts.py create mode 100644 examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml create mode 100644 examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml create mode 100644 examples/diffusers/fastgen/qad/finetune.py create mode 100644 examples/diffusers/fastgen/qad/modeling.py create mode 100644 examples/diffusers/fastgen/qad/pipeline.py create mode 100644 examples/diffusers/fastgen/qad/recipe.py diff --git a/examples/diffusers/fastgen/qad/README.md b/examples/diffusers/fastgen/qad/README.md new file mode 100644 index 00000000000..175f59a2e02 --- /dev/null +++ b/examples/diffusers/fastgen/qad/README.md @@ -0,0 +1,174 @@ +# FastGen Quantization-Aware Distillation + +This example trains a quantized diffusion student against a frozen, BF16 +Diffusers teacher with ModelOpt's distillation API. It is a standalone FastGen +recipe: it does not use DMD2, reduce the sampling schedule, create a fake-score +model, or add a GAN/EMA training phase. + +The initial Qwen-Image recipe uses the official `Qwen/Qwen-Image` Diffusers +checkpoint as the teacher. Set `qad.teacher_model_name_or_path` to +`nvidia/Qwen-Image-Flash` when the four-step DMD2-trained Qwen-Image checkpoint +should be the teacher instead. Both follow the standard Diffusers checkpoint +interface. QAD intentionally does not interpret FastGen/DMD2 intermediate +checkpoint sidecars or standalone transformer safetensors as teacher inputs. + +Every training micro-batch samples one noisy latent and one timestep, then sends +the same latent, timestep, prompt conditioning, and guidance inputs to the +teacher and student. + +## Supported students + +The `qad.student.mode` field selects one of two restore contracts. + +### Regular NVFP4 + +Set `qad.student.mode=nvfp4` and point +`model.pretrained_model_name_or_path` at the unquantized Diffusers model. The recipe +loads its weights first and then restores the weight-free ModelOpt NVFP4 state +from `qad.student.quant_state_path`. The quantizer state must have been calibrated +against exactly the same student weights. This mode trains all student +parameters, so its only valid `train_scope` is `all`. + +Use [`configs/qwen_image_nvfp4.yaml`](configs/qwen_image_nvfp4.yaml) as the +starting configuration. + +### NVFP4 SVDQuant with Hugging Face PEFT + +Set `qad.student.mode=nvfp4_svdquant` and point +`model.pretrained_model_name_or_path` at a user-prepared, ModelOpt-enabled Diffusers +training bundle. The bundle must contain the complete SVDQuant student: + +- a DiffusionPipeline root with `model_index.json` (not only a standalone + transformer `save_pretrained` directory); +- the ModelOpt topology and quantizer state; +- the residual weights produced by SVDQuant calibration; and +- the Hugging Face PEFT A/B factors for the SVDQuant low-rank branch. + +For the standard Diffusers layout, the transformer files and ModelOpt sidecar +are under `transformer/`, including `transformer/modelopt_state.pth`. The path +given to QAD is the parent DiffusionPipeline directory. + +A weight-free NVFP4 quantizer-state file is not a valid SVDQuant bundle. +SVDQuant subtracts the low-rank branch from the original weight, so both the +resulting residual weight and the PEFT factors are required. A unified +deployment export is also not a training bundle and must not be used here. + +The SVDQuant topology is restored before FSDP and before optimizer construction. +`qad.student.train_scope=all` is the default and trains both the residual/base +parameters and the PEFT factors. Set it to `lora_only` to freeze every student +parameter except the SVDQuant PEFT A/B factors. In both scopes, +`pre_quant_scale` remains ModelOpt buffer state and is never placed in the +optimizer. + +Use [`configs/qwen_image_svdquant_nvfp4.yaml`](configs/qwen_image_svdquant_nvfp4.yaml) +as the starting configuration. + +QAD is restore-only in both modes. It does not calibrate a student during +distributed training. + +## Distillation losses + +Output distillation is always MSE. The canonical setting is: + +```yaml +qad: + output_loss: + type: mse + weight: 1.0 + task_loss: + weight: 0.0 +``` + +At `weight: 1.0`, the optimized objective is pure teacher-output MSE and the +ordinary flow-matching target has weight zero because `task_loss.weight` defaults +to `0.0`. All coefficients are independent and additive. For example, setting +both output and task weights to `0.5` produces an equal output-MSE/flow-matching +mixture; adding layerwise terms does not silently renormalize either coefficient. + +Optional layerwise MSE can be added without changing the output loss: + +```yaml +qad: + layerwise: + enabled: true + pairs: + - student_layer: transformer_blocks.29 + teacher_layer: transformer_blocks.29 + selector: hidden_states + weight: 0.05 +``` + +Each pair is an exact module name relative to the student or teacher +transformer. Its weight is additive to the output/task objective. Start with +output-only training: layer hooks retain activations and therefore increase +memory use, especially when activation checkpointing is enabled. + +The recipe logs the flow-matching loss, output MSE, every configured layerwise +MSE, and the final combined loss separately. + +## Configuration and launch + +The entry point is `examples/diffusers/fastgen/qad/finetune.py`. It uses the +same YAML plus dotted-command-line override convention as the other FastGen +recipes: + +```bash +torchrun --nproc-per-node=4 \ + examples/diffusers/fastgen/qad/finetune.py \ + --config examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml \ + --fsdp.dp_size=4 \ + --model.pretrained_model_name_or_path=/path/to/qwen-image-nvfp4-svdquant-training-bundle \ + --data.dataloader.cache_dir=/path/to/qwen_image_1024p \ + --checkpoint.checkpoint_dir=/path/to/qad/checkpoints +``` + +Cluster launchers can keep the established `CONFIG`, `RUN_ID`, and +`EXTRA_ARGS` interface. For example: + +```bash +EXTRA_ARGS="--step_scheduler.max_steps=50000 \ +--step_scheduler.ckpt_every_steps=1000 \ +--step_scheduler.num_epochs=200 \ +--step_scheduler.global_batch_size=64 \ +--optim.learning_rate=2e-6 \ +--lr_scheduler.min_lr=2e-6 \ +--fsdp.dp_size=64 \ +--qad.teacher_model_name_or_path=Qwen/Qwen-Image \ +--qad.output_loss.weight=1.0 \ +--qad.task_loss.weight=0.0 \ +--qad.student.mode=nvfp4_svdquant \ +--model.pretrained_model_name_or_path=/path/to/qwen-image-nvfp4-svdquant-training-bundle \ +--qad.student.train_scope=all \ +--data.dataloader.cache_dir=/path/to/qwen_image_1024p" \ +CONFIG=examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml \ +RUN_ID=qad_qwen_image_svdquant_nvfp4_16n \ +NODES=16 \ +GPUS_PER_NODE=4 \ +TIME=05:00:00 \ +PARTITION=batch \ +bash /path/to/experiments/qad_qwen_image/launch.sh +``` + +The launcher must invoke `examples/diffusers/fastgen/qad/finetune.py`. +Pointing the existing DMD2 launcher at a QAD YAML is not sufficient when that +launcher still hard-codes `dmd2_finetune.py`. + +The launch environment contains no Attention Grill settings. It also contains +no DMD2 timestep, fake-score, discriminator, negative-prompt, GAN, or EMA +settings. + +## Restore and checkpoint invariants + +On a fresh run the recipe restores the complete student first, constructs its +final ModelOpt/PEFT topology, applies FSDP, builds the optimizer from the selected +training scope, and only then creates the frozen teacher and distillation +controller. On resume, the same immutable student source reconstructs the +topology before the QAD checkpoint is loaded. + +The teacher and the transient ModelOpt distillation controller are not training +checkpoint payloads. Checkpoints contain the student state required by the +selected training scope together with optimizer, scheduler, dataloader, RNG, and +global-step state. Resolved dotted CLI overrides are materialized into the saved +`config.yaml`. Resume validates the student bundle, quantizer state, mode, train +scope, teacher, and loss configuration before loading optimizer shards; do not +change them while resuming an existing run. diff --git a/examples/diffusers/fastgen/qad/__init__.py b/examples/diffusers/fastgen/qad/__init__.py new file mode 100644 index 00000000000..d44f8e3a87e --- /dev/null +++ b/examples/diffusers/fastgen/qad/__init__.py @@ -0,0 +1,16 @@ +# 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. + +"""Quantization-aware distillation example for Diffusers models.""" diff --git a/examples/diffusers/fastgen/qad/artifacts.py b/examples/diffusers/fastgen/qad/artifacts.py new file mode 100644 index 00000000000..2721fb93eeb --- /dev/null +++ b/examples/diffusers/fastgen/qad/artifacts.py @@ -0,0 +1,428 @@ +# 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. + +"""Student artifact restore and train-scope handling for the QAD example. + +The generic AutoModel diffusion builder intentionally owns FSDP and optimizer +construction. QAD only needs two narrowly-scoped hooks around that builder: + +* restore a standalone, weight-free ModelOpt state before FSDP for regular NVFP4; +* after FSDP, optionally freeze everything except ModelOpt SVDQuant's HF PEFT A/B + parameters and rebuild AdamW from the live sharded parameters. + +SVDQuant itself is never calibrated here. Its complete topology and weights must +already be present in a ModelOpt-enabled Diffusers training bundle. +""" + +from __future__ import annotations + +import contextlib +import dataclasses +import inspect +import logging +import os +import re +from typing import TYPE_CHECKING, Any + +import modelopt.torch.opt as mto +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.utils.core_utils import set_quantizer_state_dict + +if TYPE_CHECKING: + from collections.abc import Iterator + + import torch + from torch import nn + +_SVDQUANT_PARAMETER_RE = re.compile(r"(?:^|\.)lora_[AB]\.modelopt_svdquant\.weight$") +_SUPPORTED_STUDENT_MODES = frozenset({"nvfp4", "nvfp4_svdquant"}) +_SUPPORTED_TRAIN_SCOPES = frozenset({"all", "lora_only"}) + + +@dataclasses.dataclass(frozen=True) +class StudentSettings: + """Resolved ``qad.student`` configuration.""" + + mode: str + model_name_or_path: str + train_scope: str = "all" + quant_state_path: str | None = None + + def validate(self) -> None: + if self.mode not in _SUPPORTED_STUDENT_MODES: + raise ValueError( + f"qad.student.mode must be one of {sorted(_SUPPORTED_STUDENT_MODES)}, " + f"got {self.mode!r}." + ) + if not self.model_name_or_path: + raise ValueError("model.pretrained_model_name_or_path is required for the student.") + if self.train_scope not in _SUPPORTED_TRAIN_SCOPES: + raise ValueError( + f"qad.student.train_scope must be 'all' or 'lora_only', got {self.train_scope!r}." + ) + + if self.mode == "nvfp4": + if not self.quant_state_path: + raise ValueError( + "qad.student.quant_state_path is required for a regular NVFP4 student." + ) + if self.train_scope != "all": + raise ValueError("Regular NVFP4 supports only qad.student.train_scope=all.") + elif self.quant_state_path: + raise ValueError( + "A SVDQuant student must be restored from a complete ModelOpt-enabled " + "Diffusers training bundle; do not set qad.student.quant_state_path." + ) + + +@dataclasses.dataclass +class StudentBuildState: + """Information captured while AutoModel builds the student.""" + + parallel_scheme: dict[str, dict[str, Any]] | None = None + quantizer_count: int = 0 + svdquant_parameter_names: tuple[str, ...] = () + + +def _restore_regular_quant_state(model: nn.Module, path: str, device: torch.device) -> int: + """Restore a regular weight-free ModelOpt quantizer state before FSDP.""" + if not os.path.isfile(path): + raise FileNotFoundError(f"ModelOpt quantizer state does not exist: {path}") + if mto.ModeloptStateManager.is_converted(model): + raise RuntimeError( + "The regular NVFP4 base model already contains ModelOpt topology while " + "qad.student.quant_state_path was also supplied. Use exactly one restore source." + ) + + state = mto.load_modelopt_state(path) + mode_states = state.get("modelopt_state_dict", ()) + if any( + mode_name == "svdquant_calibrate" or mode_state.get("metadata", {}).get("svdquant_peft") + for mode_name, mode_state in mode_states + ): + raise RuntimeError( + "qad.student.mode=nvfp4 cannot restore an SVDQuant state. Use " + "mode=nvfp4_svdquant with a complete ModelOpt-enabled Diffusers bundle." + ) + quantizer_state = state.pop("modelopt_state_weights", None) + if quantizer_state is None: + raise RuntimeError( + f"{path} has no modelopt_state_weights payload. Expected the weight-free " + "quantizer state produced by the Diffusers quantization example." + ) + + mto.restore_from_modelopt_state(model, state) + set_quantizer_state_dict(model, quantizer_state) + + quantizers = [module for module in model.modules() if isinstance(module, TensorQuantizer)] + if not quantizers: + raise RuntimeError(f"No TensorQuantizer modules were restored from {path}.") + for quantizer in quantizers: + quantizer.to(device) + _validate_nvfp4_quantizers(model, artifact_name=path) + + logging.info( + "[QAD] restored regular ModelOpt quantizer state before FSDP: %s (%d quantizers)", + path, + len(quantizers), + ) + return len(quantizers) + + +def _is_block16_nvfp4(quantizer: TensorQuantizer) -> bool: + block_sizes = quantizer.block_sizes or {} + return bool( + (quantizer.is_nvfp4_dynamic or quantizer.is_nvfp4_static) and block_sizes.get(-1) == 16 + ) + + +def _enabled_quantizer_leaves(module: Any) -> tuple[TensorQuantizer, ...]: + if module is None or not hasattr(module, "modules"): + return () + return tuple( + child + for child in module.modules() + if isinstance(child, TensorQuantizer) and child.is_enabled + ) + + +def _validate_nvfp4_quantizers( + model: nn.Module, + *, + artifact_name: str, + required_targets: tuple[str, ...] = (), +) -> None: + """Reject non-NVFP4 artifacts before FSDP obscures their module topology.""" + enabled_by_slot: dict[str, list[tuple[str, TensorQuantizer]]] = { + "weight": [], + "input": [], + } + for name, module in model.named_modules(): + if not isinstance(module, TensorQuantizer) or not module.is_enabled: + continue + path_parts = name.split(".") + for slot in enabled_by_slot: + if f"{slot}_quantizer" in path_parts: + enabled_by_slot[slot].append((name, module)) + + missing_slots = [slot for slot, entries in enabled_by_slot.items() if not entries] + if missing_slots: + raise RuntimeError( + f"{artifact_name} is not an NVFP4 W4A4 training artifact: no enabled " + + "/".join(missing_slots) + + " quantizers were found." + ) + + incompatible = [ + name + for entries in enabled_by_slot.values() + for name, quantizer in entries + if not _is_block16_nvfp4(quantizer) + ] + if incompatible: + raise RuntimeError( + f"{artifact_name} contains enabled GEMM quantizers that are not block-16 NVFP4 " + "(E2M1 values with E4M3 scales): " + ", ".join(incompatible[:5]) + ) + + for target_name in required_targets: + target = model.get_submodule(target_name) + get_base_layer = getattr(target, "get_base_layer", None) + base_layer = get_base_layer() if callable(get_base_layer) else target + for slot in ("weight", "input"): + leaves = _enabled_quantizer_leaves(getattr(base_layer, f"{slot}_quantizer", None)) + if not leaves or any(not _is_block16_nvfp4(quantizer) for quantizer in leaves): + raise RuntimeError( + f"SVDQuant target {target_name!r} does not have an enabled block-16 " + f"NVFP4 {slot}_quantizer." + ) + + logging.info( + "[QAD] validated block-16 NVFP4 W4A4 quantizers before FSDP: %d weight, %d input", + len(enabled_by_slot["weight"]), + len(enabled_by_slot["input"]), + ) + + +def _svdquant_metadata(model: nn.Module) -> dict[str, Any] | None: + if not mto.ModeloptStateManager.is_converted(model): + return None + for mode_name, mode_state in mto.modelopt_state(model)["modelopt_state_dict"]: + if mode_name == "svdquant_calibrate": + return mode_state.get("metadata", {}).get("svdquant_peft") + return None + + +def _validate_svdquant_bundle(model: nn.Module) -> tuple[str, ...]: + metadata = _svdquant_metadata(model) + if not metadata: + raise RuntimeError( + "qad.student.mode=nvfp4_svdquant requires a ModelOpt-enabled Diffusers " + "training bundle with svdquant_peft metadata. The topology must be restored " + "during Diffusers from_pretrained(), before FSDP." + ) + + expected_targets = tuple(metadata.get("target_modules", ())) + names = tuple( + name for name, _ in model.named_parameters() if _SVDQUANT_PARAMETER_RE.search(name) + ) + expected_names = { + f"{target_name}.lora_{factor}.modelopt_svdquant.weight" + for target_name in expected_targets + for factor in ("A", "B") + } + if not expected_targets or set(names) != expected_names: + raise RuntimeError( + "The SVDQuant bundle did not restore a complete pair of " + "lora_A/lora_B.modelopt_svdquant weights for every target module. " + "A weight-free quantizer state or a deployment export is not a valid " + "QAD training bundle." + ) + _validate_nvfp4_quantizers( + model, + artifact_name="The SVDQuant student bundle", + required_targets=expected_targets, + ) + missing_pre_quant_scale_buffers: list[str] = [] + for target_name in expected_targets: + target = model.get_submodule(target_name) + get_base_layer = getattr(target, "get_base_layer", None) + base_layer = get_base_layer() if callable(get_base_layer) else target + input_quantizer = getattr(base_layer, "input_quantizer", None) + pre_quant_scale = getattr(input_quantizer, "_pre_quant_scale", None) + if ( + pre_quant_scale is None + or getattr(input_quantizer, "_buffers", {}).get("_pre_quant_scale") + is not pre_quant_scale + ): + missing_pre_quant_scale_buffers.append(target_name) + if missing_pre_quant_scale_buffers: + raise RuntimeError( + "SVDQuant pre_quant_scale must be restored as frozen TensorQuantizer buffer " + "state for every target; missing or non-buffer targets: " + + ", ".join(missing_pre_quant_scale_buffers[:5]) + ) + logging.info( + "[QAD] validated SVDQuant training bundle before FSDP: %d targets, %d A/B tensors", + len(expected_targets), + len(names), + ) + return names + + +def _apply_train_scope(model: nn.Module, scope: str) -> list[nn.Parameter]: + if scope == "lora_only": + for name, parameter in model.named_parameters(): + parameter.requires_grad_(_SVDQUANT_PARAMETER_RE.search(name) is not None) + + trainable = [parameter for parameter in model.parameters() if parameter.requires_grad] + if not trainable: + raise RuntimeError(f"qad.student.train_scope={scope!r} selected no parameters.") + + if scope == "lora_only": + live_names = tuple( + name for name, parameter in model.named_parameters() if parameter.requires_grad + ) + invalid = [name for name in live_names if not _SVDQUANT_PARAMETER_RE.search(name)] + if invalid: + raise RuntimeError( + "lora_only left non-SVDQuant parameters trainable: " + ", ".join(invalid[:5]) + ) + + parameter_pre_scales = [ + name for name, _ in model.named_parameters() if "pre_quant_scale" in name + ] + if parameter_pre_scales: + raise RuntimeError( + "pre_quant_scale must remain a buffer and must never enter the optimizer: " + + ", ".join(parameter_pre_scales[:5]) + ) + return trainable + + +def _rebuild_optimizer_from_live_parameters( + optimizer: torch.optim.Optimizer, + parameters: list[nn.Parameter], +) -> torch.optim.Optimizer: + """Recreate the just-built optimizer without carrying stale parameter refs.""" + if optimizer.state: + raise RuntimeError("QAD expected a newly-created optimizer with no state.") + if len(optimizer.param_groups) != 1: + raise RuntimeError( + "QAD lora_only currently expects AutoModel to create one optimizer parameter group." + ) + return type(optimizer)(parameters, **dict(optimizer.defaults)) + + +def _validate_optimizer_membership( + model: nn.Module, + optimizer: torch.optim.Optimizer, +) -> None: + expected = {id(parameter) for parameter in model.parameters() if parameter.requires_grad} + actual_list = [parameter for group in optimizer.param_groups for parameter in group["params"]] + actual = {id(parameter) for parameter in actual_list} + if len(actual) != len(actual_list): + raise RuntimeError("The student optimizer contains duplicate parameter references.") + if actual != expected: + raise RuntimeError( + "Student optimizer membership does not exactly match the live post-FSDP " + f"trainable parameters (missing={len(expected - actual)}, extra={len(actual - expected)})." + ) + + +def _guard_automodel_hooks(diffusion_train: Any, auto_pipeline: Any) -> None: + builder_parameters = inspect.signature(diffusion_train.build_model_and_optimizer).parameters + required_builder_parameters = { + "model_id", + "learning_rate", + "device", + "dtype", + "optimizer_cfg", + } + missing = required_builder_parameters - set(builder_parameters) + if missing: + raise RuntimeError( + "Unsupported nemo_automodel diffusion builder; missing parameters: " + + ", ".join(sorted(missing)) + ) + if not hasattr(auto_pipeline, "_apply_parallelization"): + raise RuntimeError( + "Unsupported nemo_automodel: auto_diffusion_pipeline._apply_parallelization is missing." + ) + + +@contextlib.contextmanager +def patch_student_build( + settings: StudentSettings, +) -> Iterator[StudentBuildState]: + """Patch the two example-local seams needed during the parent ``setup`` call. + + Both module globals are restored in ``finally``. The patch is active only while + the one student is being constructed; teacher construction happens afterwards. + """ + from nemo_automodel._diffusers import auto_diffusion_pipeline as auto_pipeline + from nemo_automodel.recipes.diffusion import train as diffusion_train + + _guard_automodel_hooks(diffusion_train, auto_pipeline) + original_apply_parallelization = auto_pipeline._apply_parallelization + original_build_model_and_optimizer = diffusion_train.build_model_and_optimizer + state = StudentBuildState() + apply_calls = 0 + + def apply_parallelization(pipe, parallel_scheme): + nonlocal apply_calls + apply_calls += 1 + if apply_calls != 1: + raise RuntimeError( + "QAD's guarded student build expected exactly one parallelized component load." + ) + state.parallel_scheme = parallel_scheme + transformer = pipe.transformer + if settings.mode == "nvfp4": + quant_state_path = settings.quant_state_path + assert quant_state_path is not None + state.quantizer_count = _restore_regular_quant_state( + transformer, + quant_state_path, + next(transformer.parameters()).device, + ) + else: + state.svdquant_parameter_names = _validate_svdquant_bundle(transformer) + return original_apply_parallelization(pipe, parallel_scheme) + + def build_model_and_optimizer(**kwargs): + pipe, optimizer, device_mesh = original_build_model_and_optimizer(**kwargs) + trainable = _apply_train_scope(pipe.transformer, settings.train_scope) + if settings.train_scope == "lora_only": + optimizer = _rebuild_optimizer_from_live_parameters(optimizer, trainable) + logging.info( + "[QAD] rebuilt optimizer after FSDP for lora_only: %d live A/B tensors", + len(trainable), + ) + _validate_optimizer_membership(pipe.transformer, optimizer) + return pipe, optimizer, device_mesh + + auto_pipeline._apply_parallelization = apply_parallelization + diffusion_train.build_model_and_optimizer = build_model_and_optimizer + try: + yield state + finally: + diffusion_train.build_model_and_optimizer = original_build_model_and_optimizer + auto_pipeline._apply_parallelization = original_apply_parallelization + + if apply_calls != 1 or state.parallel_scheme is None: + raise RuntimeError( + "QAD did not observe the expected pre-FSDP student parallelization point." + ) diff --git a/examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml b/examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml new file mode 100644 index 00000000000..70bc99f8d09 --- /dev/null +++ b/examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml @@ -0,0 +1,108 @@ +# Qwen-Image QAD with a regular ModelOpt NVFP4 student. +# +# The student is loaded from its Diffusers checkpoint, then the weight-free +# quantizer state is restored. Paths are placeholders and should be supplied by +# dotted CLI overrides in production launches. + +seed: 42 + +wandb: + project: fastgen-qad-qwen-image + mode: online + name: qwen_image_qad_nvfp4 + +dist_env: + backend: nccl + timeout_minutes: 60 + +model: + # Canonical student source; AutoModel records this field in checkpoints. + pretrained_model_name_or_path: Qwen/Qwen-Image + mode: finetune + +step_scheduler: + global_batch_size: 64 + local_batch_size: 1 + ckpt_every_steps: 1000 + num_epochs: 200 + log_every: 1 + max_steps: 50000 + +qad: + # The teacher is always an unquantized, frozen Diffusers model. + teacher_model_name_or_path: Qwen/Qwen-Image + + # weight=1.0 is pure teacher-output MSE; the flow-matching task weight is 0. + output_loss: + type: mse + weight: 1.0 + + # Independent additive coefficient for the ordinary flow-matching target. + task_loss: + weight: 0.0 + + # Layer pairs use exact module names relative to each transformer. + layerwise: + enabled: false + pairs: [] + + student: + mode: nvfp4 + quant_state_path: /path/to/qwen-image/quant/transformer.nvfp4.pt + train_scope: all + +optim: + learning_rate: 2.0e-6 + optimizer: + weight_decay: 0.01 + betas: [0.9, 0.999] + +lr_scheduler: + lr_decay_style: constant + lr_warmup_steps: 0 + min_lr: 2.0e-6 + +fsdp: + tp_size: 1 + cp_size: 1 + pp_size: 1 + dp_replicate_size: 1 + dp_size: 1 + activation_checkpointing: true + +flow_matching: + adapter_type: qwen_image + timestep_sampling: logit_normal + logit_mean: 0.0 + logit_std: 1.0 + flow_shift: 3.0 + mix_uniform_ratio: 0.1 + use_sigma_noise: true + sigma_min: 0.0 + sigma_max: 1.0 + num_train_timesteps: 1000 + cfg_dropout_prob: 0.0 + use_loss_weighting: true + loss_weighting_scheme: linear + adapter_kwargs: + guidance_scale: 3.5 + use_guidance_embeds: false + +data: + dataloader: + _target_: fastgen_data.build_text_to_image_multiresolution_dataloader + cache_dir: /path/to/preprocessed/qwen_image_1024p + base_resolution: [1024, 1024] + batch_size: 1 + drop_last: false + shuffle: true + num_workers: 0 + +checkpoint: + enabled: true + checkpoint_dir: /path/to/output/qwen_image_qad_nvfp4/checkpoints + model_save_format: safetensors + # Save full sharded student DCP. Deployment export is a separate operation. + save_consolidated: false + diffusers_compatible: false + restore_from: LATEST diff --git a/examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml b/examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml new file mode 100644 index 00000000000..99873aeae43 --- /dev/null +++ b/examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml @@ -0,0 +1,112 @@ +# Qwen-Image QAD with a ModelOpt NVFP4 SVDQuant + Hugging Face PEFT student. +# +# model.pretrained_model_name_or_path must be a complete, user-prepared ModelOpt-enabled +# Diffusers training bundle. A weight-free quantizer-state file or a unified +# deployment export is not sufficient. + +seed: 42 + +wandb: + project: fastgen-qad-qwen-image + mode: online + name: qwen_image_qad_svdquant_nvfp4 + +dist_env: + backend: nccl + timeout_minutes: 60 + +model: + # Complete user-prepared ModelOpt-enabled Diffusers training bundle. + pretrained_model_name_or_path: /path/to/qwen-image-nvfp4-svdquant-training-bundle + mode: finetune + +step_scheduler: + global_batch_size: 64 + local_batch_size: 1 + ckpt_every_steps: 1000 + num_epochs: 200 + log_every: 1 + max_steps: 50000 + +qad: + # Keep the teacher independent from the quantized student bundle. + teacher_model_name_or_path: Qwen/Qwen-Image + + # weight=1.0 is pure teacher-output MSE; the flow-matching task weight is 0. + output_loss: + type: mse + weight: 1.0 + + # Independent additive coefficient for the ordinary flow-matching target. + task_loss: + weight: 0.0 + + # Example when enabled: + # pairs: + # - student_layer: transformer_blocks.29 + # teacher_layer: transformer_blocks.29 + # weight: 0.05 + layerwise: + enabled: false + pairs: [] + + student: + mode: nvfp4_svdquant + # all is canonical; lora_only trains only the SVDQuant HF PEFT A/B factors. + train_scope: all + +optim: + learning_rate: 2.0e-6 + optimizer: + weight_decay: 0.01 + betas: [0.9, 0.999] + +lr_scheduler: + lr_decay_style: constant + lr_warmup_steps: 0 + min_lr: 2.0e-6 + +fsdp: + tp_size: 1 + cp_size: 1 + pp_size: 1 + dp_replicate_size: 1 + dp_size: 1 + activation_checkpointing: true + +flow_matching: + adapter_type: qwen_image + timestep_sampling: logit_normal + logit_mean: 0.0 + logit_std: 1.0 + flow_shift: 3.0 + mix_uniform_ratio: 0.1 + use_sigma_noise: true + sigma_min: 0.0 + sigma_max: 1.0 + num_train_timesteps: 1000 + cfg_dropout_prob: 0.0 + use_loss_weighting: true + loss_weighting_scheme: linear + adapter_kwargs: + guidance_scale: 3.5 + use_guidance_embeds: false + +data: + dataloader: + _target_: fastgen_data.build_text_to_image_multiresolution_dataloader + cache_dir: /path/to/preprocessed/qwen_image_1024p + base_resolution: [1024, 1024] + batch_size: 1 + drop_last: false + shuffle: true + num_workers: 0 + +checkpoint: + enabled: true + checkpoint_dir: /path/to/output/qwen_image_qad_svdquant_nvfp4/checkpoints + model_save_format: safetensors + # Save full sharded student DCP. Deployment export is a separate operation. + save_consolidated: false + diffusers_compatible: false + restore_from: LATEST diff --git a/examples/diffusers/fastgen/qad/finetune.py b/examples/diffusers/fastgen/qad/finetune.py new file mode 100644 index 00000000000..495c2bef484 --- /dev/null +++ b/examples/diffusers/fastgen/qad/finetune.py @@ -0,0 +1,43 @@ +# 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. + +"""Entrypoint for FastGen quantization-aware distillation.""" + +from __future__ import annotations + +import os +import sys + +_QAD_DIR = os.path.dirname(os.path.abspath(__file__)) +_FASTGEN_DIR = os.path.dirname(_QAD_DIR) +if _FASTGEN_DIR not in sys.path: + sys.path.insert(0, _FASTGEN_DIR) + +from nemo_automodel.components.config._arg_parser import parse_args_and_load_config # noqa: E402 + +from qad.recipe import QADDiffusionRecipe # noqa: E402 + + +def main( + default_config_path: str = ("examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml"), +) -> None: + cfg = parse_args_and_load_config(default_config_path) + recipe = QADDiffusionRecipe(cfg) + recipe.setup() + recipe.run_train_validation_loop() + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/qad/modeling.py b/examples/diffusers/fastgen/qad/modeling.py new file mode 100644 index 00000000000..759302795bb --- /dev/null +++ b/examples/diffusers/fastgen/qad/modeling.py @@ -0,0 +1,259 @@ +# 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. + +"""Model and ModelOpt-distillation helpers for QAD.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn.functional as F +from torch import nn + +import modelopt.torch.distill as mtd + +if TYPE_CHECKING: + from collections.abc import Sequence + +try: + from nemo_automodel.components.distributed.parallelizer import ( + PARALLELIZATION_STRATEGIES, + DefaultParallelizationStrategy, + register_parallel_strategy, + ) +except ImportError as exc: + raise ImportError( + "The FastGen QAD example requires nemo_automodel. Install " + "examples/diffusers/fastgen/requirements.txt." + ) from exc + + +class _QwenImageParallelizationStrategy(DefaultParallelizationStrategy): + """Checkpoint complete Qwen transformer blocks before AutoModel applies FSDP.""" + + 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( + "[QAD] 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: + """Register the Qwen strategy unless AutoModel already ships a native strategy.""" + 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() + + +def _extract_tensor(output: Any, selector: str) -> torch.Tensor: + """Select a tensor from a Diffusers root or Qwen dual-stream block output.""" + normalized = selector.lower() + if torch.is_tensor(output): + return output + if hasattr(output, "sample") and normalized in {"sample", "output", "tensor"}: + return output.sample + if isinstance(output, dict): + if selector not in output: + raise KeyError(f"Selector {selector!r} is not present in layer output keys.") + selected = output[selector] + if not torch.is_tensor(selected): + raise TypeError(f"Layer output {selector!r} is not a tensor.") + return selected + if isinstance(output, tuple | list): + index_by_name = { + "sample": 0, + "output": 0, + "tensor": 0, + "first": 0, + "encoder_hidden_states": 0, + "text": 0, + "hidden_states": 1, + "image": 1, + "last": -1, + } + if normalized not in index_by_name: + raise ValueError( + f"Unsupported tuple selector {selector!r}; use hidden_states/image, " + "encoder_hidden_states/text, first, last, or sample." + ) + selected = output[index_by_name[normalized]] + if not torch.is_tensor(selected): + raise TypeError(f"Selected {selector!r} output is not a tensor.") + return selected + raise TypeError(f"Cannot select {selector!r} from output type {type(output).__name__}.") + + +class TensorOutputDelegate(nn.Module): + """Forward to a model while exposing only its final tensor to ModelOpt KD. + + The wrapped model is deliberately stored outside ``nn.Module._modules``. This + keeps the controller parameter-free and prevents its state_dict from duplicating + either the FSDP student or teacher. ``get_submodule`` still routes layerwise + criterion paths to the live wrapped transformer. + """ + + def __init__(self, target: nn.Module): + super().__init__() + self.__dict__["_qad_target"] = target + + @property + def target(self) -> nn.Module: + return self.__dict__["_qad_target"] + + def forward(self, *args, **kwargs) -> torch.Tensor: + return _extract_tensor(self.target(*args, **kwargs), "sample") + + def get_submodule(self, target: str) -> nn.Module: + if target == "": + return self + return self.target.get_submodule(target) + + +class SelectedMSELoss(nn.modules.loss._Loss): + """FP32 MSE after selecting a stream from a captured layer output.""" + + def __init__(self, selector: str = "sample"): + super().__init__(reduction="mean") + self.selector = selector + + def forward(self, student_output: Any, teacher_output: Any) -> torch.Tensor: + student = _extract_tensor(student_output, self.selector) + teacher = _extract_tensor(teacher_output, self.selector) + return F.mse_loss(student.float(), teacher.float(), reduction="mean") + + +class AdditiveLossBalancer(mtd.DistillationLossBalancer): + """Apply independent, additive weights to task and KD loss terms.""" + + def __init__(self, *, task_weight: float, kd_weights: Sequence[float]): + super().__init__() + self.task_weight = float(task_weight) + self.kd_weights = tuple(float(weight) for weight in kd_weights) + + def forward(self, losses: dict[str, torch.Tensor]) -> torch.Tensor: + losses = dict(losses) + student_loss = losses.pop("student_loss", None) + if not losses: + raise RuntimeError("QAD received no KD loss terms.") + total = None + if self.task_weight != 0.0: + if student_loss is None: + raise RuntimeError("A nonzero QAD task weight requires student_loss.") + total = student_loss * self.task_weight + + if len(losses) != len(self.kd_weights): + raise RuntimeError( + "ModelOpt returned an unexpected number of KD losses: " + f"expected {len(self.kd_weights)}, got {len(losses)}." + ) + for loss, weight in zip(losses.values(), self.kd_weights): + # Multiplying a disabled NaN/Inf term by zero would still poison the + # objective. Skip disabled terms completely while retaining their + # detached diagnostics in the pipeline. + if weight == 0.0: + continue + weighted_loss = loss * weight + total = weighted_loss if total is None else total + weighted_loss + if total is None: + raise RuntimeError("QAD has no nonzero loss coefficient.") + return total + + +def build_distillation_controller( + *, + student: nn.Module, + teacher: nn.Module, + output_weight: float, + task_weight: float, + layer_pairs: Sequence[dict[str, Any]], +) -> tuple[nn.Module, tuple[str, ...]]: + """Create a parameter-free ModelOpt KD controller around live FSDP models.""" + criterion: dict[tuple[str, str], nn.modules.loss._Loss] = {("", ""): SelectedMSELoss("sample")} + names = ["output_mse"] + weights = [float(output_weight)] + seen_pairs = {("", "")} + + for index, pair in enumerate(layer_pairs): + student_layer = str(pair["student_layer"]) + teacher_layer = str(pair.get("teacher_layer", student_layer)) + selector = str(pair.get("selector", "hidden_states")) + weight = float(pair.get("weight", 1.0)) + key = (student_layer, teacher_layer) + if key in seen_pairs: + raise ValueError(f"Duplicate QAD layer pair: {key!r}") + seen_pairs.add(key) + criterion[key] = SelectedMSELoss(selector) + names.append(f"layer_{index}_{student_layer}_{selector}_mse") + weights.append(weight) + + controller = mtd.convert( + TensorOutputDelegate(student), + mode=[ + ( + "kd_loss", + { + "teacher_model": TensorOutputDelegate(teacher), + "criterion": criterion, + "loss_balancer": AdditiveLossBalancer( + task_weight=task_weight, + kd_weights=weights, + ), + "expose_minimal_state_dict": True, + }, + ) + ], + ) + return controller, tuple(names) + + +def clear_captured_outputs(controller: nn.Module) -> None: + """Release activation references before forwards and after checkpoint recompute.""" + for student_layer, teacher_layer in controller._layers_to_loss: + student_layer._intermediate_output = None + teacher_layer._intermediate_output = None diff --git a/examples/diffusers/fastgen/qad/pipeline.py b/examples/diffusers/fastgen/qad/pipeline.py new file mode 100644 index 00000000000..734a904fd31 --- /dev/null +++ b/examples/diffusers/fastgen/qad/pipeline.py @@ -0,0 +1,74 @@ +# 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. + +"""QAD loss pipeline layered on AutoModel's flow-matching input preparation.""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + +from .modeling import clear_captured_outputs + + +class QADPipeline: + """Run teacher/student on identical inputs and aggregate ModelOpt KD losses.""" + + def __init__(self, flow_matching_pipeline, controller: nn.Module, loss_names: tuple[str, ...]): + self.flow_matching_pipeline = flow_matching_pipeline + self.controller = controller + self.loss_names = loss_names + + def step( + self, + *, + batch: dict[str, Any], + device: torch.device, + dtype: torch.dtype, + global_step: int, + check_loss: bool, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + clear_captured_outputs(self.controller) + _, task_loss, _, _ = self.flow_matching_pipeline.step( + model=self.controller, + batch=batch, + device=device, + dtype=dtype, + global_step=global_step, + collect_metrics=False, + # The flow target is optional in QAD. Validate the actual combined loss below. + check_loss=False, + ) + losses = self.controller.compute_kd_loss( + student_loss=task_loss, + skip_balancer=True, + ) + total = self.controller.loss_balancer(losses) + if check_loss and not bool(torch.isfinite(total.detach()).all()): + raise FloatingPointError(f"Non-finite QAD loss at step {global_step}.") + + kd_values = [value for key, value in losses.items() if key != "student_loss"] + if len(kd_values) != len(self.loss_names): + raise RuntimeError( + "QAD loss-name mapping is out of sync with ModelOpt's returned losses." + ) + metrics = {"task_loss": task_loss.detach(), "total_loss": total.detach()} + metrics.update({name: value.detach() for name, value in zip(self.loss_names, kd_values)}) + return total, metrics + + def clear(self) -> None: + clear_captured_outputs(self.controller) diff --git a/examples/diffusers/fastgen/qad/recipe.py b/examples/diffusers/fastgen/qad/recipe.py new file mode 100644 index 00000000000..276abac32ac --- /dev/null +++ b/examples/diffusers/fastgen/qad/recipe.py @@ -0,0 +1,595 @@ +# 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. + +"""FastGen quantization-aware distillation recipe. + +QAD is deliberately separate from DMD2: one frozen Diffusers teacher and one +quantized student see the same noisy latent, timestep, and conditioning, and +ModelOpt's standard ``kd_loss`` API supplies output and optional representation +MSE losses. +""" + +from __future__ import annotations + +import logging +import math +import os +from typing import Any + +import torch +import wandb +import yaml +from torch import nn +from torchdata.stateful_dataloader import StatefulDataLoader + +import modelopt.torch.distill as mtd +import modelopt.torch.opt as mto +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.nn import TensorQuantizer + +try: + from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline + from nemo_automodel.components.training.utils import ( + clip_grad_norm, + prepare_after_first_microbatch, + prepare_for_final_backward, + prepare_for_grad_accumulation, + ) + from nemo_automodel.recipes.base_recipe import ( + _find_latest_checkpoint, + _resolve_restore_from_to_ckpt_dir, + ) + from nemo_automodel.recipes.diffusion.train import TrainDiffusionRecipe, is_main_process +except ImportError as exc: + raise ImportError( + "The FastGen QAD example requires nemo_automodel. Install dependencies with:\n" + " pip install -r examples/diffusers/fastgen/requirements.txt" + ) from exc + +from fastgen_checkpoint import make_optimizer_partial_load_tolerant + +from .artifacts import StudentSettings, patch_student_build +from .modeling import build_distillation_controller, clear_captured_outputs +from .pipeline import QADPipeline + + +def _as_dict(value: Any) -> dict[str, Any]: + if value is None: + return {} + if hasattr(value, "to_dict"): + return value.to_dict() + return dict(value) + + +class QADDiffusionRecipe(TrainDiffusionRecipe): + """AutoModel diffusion recipe with a ModelOpt KD controller.""" + + def __init__(self, cfg) -> None: + # AutoModel's dotted CLI setter updates live ConfigNodes but not the + # raw_config later written beside checkpoints. Materialize the resolved + # runtime values so QAD paths, scope, teacher, and losses are reproducible. + if hasattr(cfg, "to_yaml_dict"): + cfg.__dict__["_raw_config"] = cfg.to_yaml_dict( + resolve_env=False, + redact_sensitive=True, + use_orig_values=False, + ) + super().__init__(cfg) + + def setup(self) -> None: + settings, loss_config = self._resolve_qad_config() + self.__dict__["_qad_resume_signature"] = self._resume_signature( + settings, + loss_config, + ) + + if self.cfg.get("peft", None) is not None: + raise ValueError( + "Do not set AutoModel's top-level peft block for QAD. SVDQuant's " + "modelopt_svdquant HF PEFT topology must come from the student bundle." + ) + if str(self.cfg.get("model.mode", "finetune")).lower() != "finetune": + raise ValueError("QAD supports model.mode=finetune only.") + if self.cfg.get("ddp", None) is not None: + raise ValueError("QAD currently supports AutoModel FSDP2, not DDP.") + + # Diffusers' ModelMixin must be patched before from_pretrained so a + # SVDQuant bundle rebuilds quant/PEFT topology before loading its weights. + mto.enable_huggingface_checkpointing() + + with patch_student_build(settings) as build_state: + super().setup() + + # Diffusers loads ModelMixin objects in eval mode. QAD owns the student + # train/eval boundary because the controller delegate intentionally does + # not register the live FSDP module as a child. + self.model.train() + + # Parent checkpoint restore established the exact next-step RNG state. + # Teacher construction/sharding is transient setup and must not perturb + # the first fresh or resumed training sample. + training_rng_state = self.rng.state_dict() + try: + parallel_scheme = build_state.parallel_scheme + if parallel_scheme is None: + raise RuntimeError("QAD failed to capture the student's parallel scheme.") + teacher = self._load_frozen_teacher( + loss_config["teacher_model_name_or_path"], + parallel_scheme, + ) + controller, loss_names = build_distillation_controller( + student=self.model, + teacher=teacher, + output_weight=loss_config["output_weight"], + task_weight=loss_config["task_weight"], + layer_pairs=loss_config["layer_pairs"], + ) + if any(True for _ in controller.parameters()): + raise RuntimeError( + "The QAD controller must remain parameter-free; optimizer/checkpoint " + "ownership belongs exclusively to self.model." + ) + + # BaseRecipe tracks nn.Module assignments. Bypass it for the frozen teacher + # and transient controller so checkpoint selection cannot mistake either for + # the student. + object.__setattr__(self, "_qad_teacher", teacher) + object.__setattr__(self, "_qad_controller", controller) + object.__setattr__( + self, + "_qad_pipeline", + QADPipeline(self.flow_matching_pipeline, controller, loss_names), + ) + object.__setattr__(self, "_qad_student_settings", settings) + object.__setattr__(self, "_qad_loss_config", loss_config) + + tracked = self.__dict__.get("__state_tracked", set()) + forbidden = {"_qad_teacher", "_qad_controller", "_qad_pipeline"} & set(tracked) + if forbidden: + raise RuntimeError( + f"QAD transient objects were accidentally state-tracked: {forbidden}" + ) + self._validate_state_ownership() + finally: + self.rng.load_state_dict(training_rng_state) + + if is_main_process(): + logging.info( + "[QAD] initialized: teacher=%s student=%s mode=%s train_scope=%s " + "task_weight=%g output_weight=%g layer_pairs=%d", + loss_config["teacher_model_name_or_path"], + settings.model_name_or_path, + settings.mode, + settings.train_scope, + loss_config["task_weight"], + loss_config["output_weight"], + len(loss_config["layer_pairs"]), + ) + logging.info("[QAD] student quantizer summary:") + mtq.print_quant_summary(self.model) + + def _resolve_qad_config(self) -> tuple[StudentSettings, dict[str, Any]]: + qad = _as_dict(self.cfg.get("qad", None)) + if not qad: + raise ValueError("Missing required qad configuration block.") + + student_cfg = _as_dict(qad.get("student")) + model_name_or_path = self.cfg.get("model.pretrained_model_name_or_path", None) + if not model_name_or_path: + raise ValueError( + "model.pretrained_model_name_or_path is required and is the canonical " + "student source recorded in checkpoints." + ) + duplicate_student_path = student_cfg.get("model_name_or_path") + if duplicate_student_path is not None and str(duplicate_student_path) != str( + model_name_or_path + ): + raise ValueError( + "qad.student.model_name_or_path conflicts with the canonical " + "model.pretrained_model_name_or_path. Remove the duplicate QAD field." + ) + mode = str(student_cfg.get("mode", "nvfp4")).lower() + # Accept the early design spelling while emitting one canonical name. + if mode == "svdquant_nvfp4": + mode = "nvfp4_svdquant" + settings = StudentSettings( + mode=mode, + model_name_or_path=str(model_name_or_path), + quant_state_path=student_cfg.get("quant_state_path"), + train_scope=str(student_cfg.get("train_scope", "all")).lower(), + ) + settings.validate() + + teacher_model_name_or_path = qad.get("teacher_model_name_or_path") + if not teacher_model_name_or_path: + raise ValueError("qad.teacher_model_name_or_path is required.") + + output_cfg = _as_dict(qad.get("output_loss")) + if str(output_cfg.get("type", "mse")).lower() != "mse": + raise ValueError("QAD currently supports only output_loss.type=mse.") + output_weight = float(output_cfg.get("weight", 1.0)) + + task_cfg = _as_dict(qad.get("task_loss")) + task_weight = float(task_cfg.get("weight", 0.0)) + + layerwise_cfg = _as_dict(qad.get("layerwise")) + layer_pairs = layerwise_cfg.get("pairs", []) if layerwise_cfg.get("enabled", False) else [] + layer_pairs = [_as_dict(pair) for pair in layer_pairs] + + all_weights = [output_weight, task_weight] + [ + float(pair.get("weight", 1.0)) for pair in layer_pairs + ] + if any(not math.isfinite(weight) or weight < 0.0 for weight in all_weights): + raise ValueError("QAD loss weights must be finite and non-negative.") + if not any(weight > 0.0 for weight in all_weights): + raise ValueError("At least one QAD loss weight must be positive.") + for index, pair in enumerate(layer_pairs): + if not pair.get("student_layer"): + raise ValueError(f"qad.layerwise.pairs[{index}].student_layer is required.") + + return settings, { + "teacher_model_name_or_path": str(teacher_model_name_or_path), + "output_weight": output_weight, + "task_weight": task_weight, + "layer_pairs": layer_pairs, + } + + @staticmethod + def _resume_signature( + settings: StudentSettings, + loss_config: dict[str, Any], + ) -> dict[str, Any]: + return { + "student_source": settings.model_name_or_path, + "student_mode": settings.mode, + "quant_state_path": settings.quant_state_path, + "train_scope": settings.train_scope, + "teacher_source": loss_config["teacher_model_name_or_path"], + "output_weight": float(loss_config["output_weight"]), + "task_weight": float(loss_config["task_weight"]), + "layer_pairs": tuple( + ( + str(pair["student_layer"]), + str(pair.get("teacher_layer", pair["student_layer"])), + str(pair.get("selector", "hidden_states")), + float(pair.get("weight", 1.0)), + ) + for pair in loss_config["layer_pairs"] + ), + } + + @classmethod + def _resume_signature_from_saved_config(cls, config: dict[str, Any]) -> dict[str, Any]: + model_cfg = _as_dict(config.get("model")) + qad_cfg = _as_dict(config.get("qad")) + student_cfg = _as_dict(qad_cfg.get("student")) + output_cfg = _as_dict(qad_cfg.get("output_loss")) + task_cfg = _as_dict(qad_cfg.get("task_loss")) + layerwise_cfg = _as_dict(qad_cfg.get("layerwise")) + + mode = str(student_cfg.get("mode", "nvfp4")).lower() + if mode == "svdquant_nvfp4": + mode = "nvfp4_svdquant" + raw_pairs = layerwise_cfg.get("pairs", []) if layerwise_cfg.get("enabled", False) else [] + loss_config = { + "teacher_model_name_or_path": str(qad_cfg.get("teacher_model_name_or_path", "")), + "output_weight": float(output_cfg.get("weight", 1.0)), + "task_weight": float(task_cfg.get("weight", 0.0)), + "layer_pairs": [_as_dict(pair) for pair in raw_pairs], + } + settings = StudentSettings( + mode=mode, + model_name_or_path=str(model_cfg.get("pretrained_model_name_or_path", "")), + quant_state_path=student_cfg.get("quant_state_path"), + train_scope=str(student_cfg.get("train_scope", "all")).lower(), + ) + return cls._resume_signature(settings, loss_config) + + def _resolved_checkpoint_dir(self, restore_from: str | None) -> str | None: + if not self.checkpointer.config.enabled: + return None + if restore_from: + resolved = _resolve_restore_from_to_ckpt_dir( + self.checkpointer.config.checkpoint_dir, + restore_from, + ) + else: + resolved = _find_latest_checkpoint(self.checkpointer.config.checkpoint_dir) + if resolved is None: + return None + return os.fspath(resolved) + + def _validate_qad_checkpoint_signature(self, checkpoint_dir: str) -> None: + config_path = os.path.join(checkpoint_dir, "config.yaml") + if not os.path.isfile(config_path): + raise RuntimeError( + "QAD cannot safely restore optimizer shards from a checkpoint without " + f"config.yaml: {checkpoint_dir}" + ) + with open(config_path) as config_file: + saved_config = yaml.safe_load(config_file) or {} + saved_signature = self._resume_signature_from_saved_config(saved_config) + current_signature = self.__dict__["_qad_resume_signature"] + if saved_signature != current_signature: + changed = [ + key + for key in current_signature + if saved_signature.get(key) != current_signature[key] + ] + raise RuntimeError( + "QAD resume contract changed for " + + ", ".join(changed) + + ". Use the same student artifact, quantization mode/state, train scope, " + "teacher, and loss configuration as the saved run." + ) + + def load_checkpoint(self, restore_from: str | None = None) -> None: + """Validate QAD topology before enabling FSDP2 partial-shard optimizer load.""" + checkpoint_dir = self._resolved_checkpoint_dir(restore_from) + if checkpoint_dir is not None and os.path.isdir(checkpoint_dir): + self._validate_qad_checkpoint_signature(checkpoint_dir) + make_optimizer_partial_load_tolerant(self.checkpointer) + super().load_checkpoint(restore_from) + + def _rebuild_dataloader_for_resume(self, global_step: int) -> None: + """Rebuild the loader and deterministically skip to the restored data position.""" + epoch_len = int(getattr(self.step_scheduler, "epoch_len", 0) or 0) + grad_acc = int(getattr(self.step_scheduler, "grad_acc_steps", 1) or 1) + if epoch_len <= 0 or self.sampler is None or global_step <= 0: + return + + current_epoch = global_step // epoch_len + skip_batches = (global_step % epoch_len) * grad_acc + old_dataloader = self.dataloader + dataloader_kwargs = { + "collate_fn": getattr(old_dataloader, "collate_fn", None), + "num_workers": int(getattr(old_dataloader, "num_workers", 0) or 0), + "pin_memory": bool(getattr(old_dataloader, "pin_memory", False)), + } + if dataloader_kwargs["num_workers"] > 0: + dataloader_kwargs["prefetch_factor"] = getattr( + old_dataloader, + "prefetch_factor", + 2, + ) + dataloader_kwargs["persistent_workers"] = bool( + getattr(old_dataloader, "persistent_workers", False) + ) + + # Keep the parent's existing tracked state key while replacing the + # StatefulDataLoader object whose restored cursor is known to stick. + self.__dict__["dataloader"] = StatefulDataLoader( + old_dataloader.dataset, + batch_sampler=self.sampler, + **dataloader_kwargs, + ) + self.step_scheduler.epoch = current_epoch + self.sampler.set_epoch(current_epoch) + self.sampler._batches_to_skip = skip_batches + if is_main_process(): + logging.info( + "[QAD][resume] rebuilt dataloader at epoch=%d skip_batches=%d " + "(global_step=%d epoch_len=%d grad_acc=%d)", + current_epoch, + skip_batches, + global_step, + epoch_len, + grad_acc, + ) + + def _load_frozen_teacher( + self, + model_name_or_path: str, + parallel_scheme: dict[str, dict[str, Any]], + ) -> nn.Module: + pipe, _ = NeMoAutoDiffusionPipeline.from_pretrained( + model_name_or_path, + torch_dtype=self.bf16, + device=self.device, + parallel_scheme=parallel_scheme, + components_to_load=["transformer"], + load_for_training=False, + low_cpu_mem_usage=True, + ) + teacher = pipe.transformer + if mto.ModeloptStateManager.is_converted(teacher): + raise RuntimeError( + "QAD teacher must be a plain BF16 Diffusers checkpoint without ModelOpt modes." + ) + if any(isinstance(module, TensorQuantizer) for module in teacher.modules()): + raise RuntimeError("QAD teacher must be an unquantized BF16 Diffusers checkpoint.") + teacher.eval() + teacher.requires_grad_(False) + return teacher + + def _validate_state_ownership(self) -> None: + optimizer_parameters = { + id(parameter) for group in self.optimizer.param_groups for parameter in group["params"] + } + student_parameters = { + id(parameter) for parameter in self.model.parameters() if parameter.requires_grad + } + teacher_parameters = {id(parameter) for parameter in self._qad_teacher.parameters()} + if optimizer_parameters != student_parameters: + raise RuntimeError("QAD optimizer does not exactly own the trainable student state.") + if optimizer_parameters & teacher_parameters: + raise RuntimeError("Frozen teacher parameters leaked into the student optimizer.") + if any(parameter.requires_grad for parameter in self._qad_teacher.parameters()): + raise RuntimeError("QAD teacher must be completely frozen.") + + def run_train_validation_loop(self) -> None: + """Run a conventional optimizer loop using the QAD objective.""" + self.model.train() + logging.info( + "[QAD] starting training: global_batch_size=%s local_batch_size=%s dp_size=%s", + self.global_batch_size, + self.local_batch_size, + self.dp_size, + ) + global_step = int(self.step_scheduler.step) + self._rebuild_dataloader_for_resume(global_step) + + try: + for epoch in self.step_scheduler.epochs: + if self.sampler is not None and hasattr(self.sampler, "set_epoch"): + self.sampler.set_epoch(epoch) + + tqdm_initial = int(getattr(self.sampler, "_batches_to_skip", 0) or 0) + if is_main_process(): + from tqdm import tqdm + + self.step_scheduler.dataloader = tqdm( + self.dataloader, + desc=f"Epoch {epoch + 1}/{self.num_epochs} (global step {global_step})", + initial=tqdm_initial, + ) + else: + self.step_scheduler.dataloader = self.dataloader + + epoch_loss = 0.0 + num_steps = 0 + for batch_group in self.step_scheduler: + # StepScheduler increments only after control returns to its + # generator, so refresh at the top of every yielded group. + global_step = int(self.step_scheduler.step) + self.optimizer.zero_grad(set_to_none=True) + prepare_for_grad_accumulation([self.model], pp_enabled=False) + num_microbatches = len(batch_group) + micro_metrics: list[dict[str, torch.Tensor]] = [] + + for microbatch_index, micro_batch in enumerate(batch_group): + if microbatch_index == num_microbatches - 1: + prepare_for_final_backward([self.model], pp_enabled=False) + try: + total_loss, metrics = self._qad_pipeline.step( + batch=micro_batch, + device=self.device, + dtype=self.bf16, + global_step=global_step, + check_loss=self.check_loss, + ) + (total_loss / num_microbatches).backward() + micro_metrics.append(metrics) + finally: + # Full-block NO_REENTRANT checkpoint wrappers avoid hook + # repopulation during recompute; this final cleanup is also + # safe when activation checkpointing is disabled. + self._qad_pipeline.clear() + + if microbatch_index == 0: + prepare_after_first_microbatch() + + self._validate_first_step_gradients(global_step) + grad_norm = clip_grad_norm( + self.clip_grad_max_norm, + [self.model], + foreach=self.grad_clip_foreach, + ) + grad_norm = float(grad_norm) if torch.is_tensor(grad_norm) else grad_norm + self.optimizer.step() + if self.lr_scheduler is not None: + self.lr_scheduler[0].step(1) + + reduced_metrics = { + name: float( + torch.stack([metrics[name] for metrics in micro_metrics]).mean().item() + ) + for name in micro_metrics[0] + } + group_loss = reduced_metrics["total_loss"] + epoch_loss += group_loss + num_steps += 1 + + if self.log_every and global_step % self.log_every == 0 and is_main_process(): + log_dict = { + "train_loss": group_loss, + "train_avg_loss": epoch_loss / num_steps, + "lr": self.optimizer.param_groups[0]["lr"], + "grad_norm": grad_norm, + "epoch": epoch, + "global_step": global_step, + **{f"qad/{name}": value for name, value in reduced_metrics.items()}, + } + if wandb.run is not None: + wandb.log(log_dict, step=global_step) + component_text = " ".join( + f"{name}={value:.6f}" for name, value in reduced_metrics.items() + ) + logging.info( + "[QAD][TRAIN] step=%d epoch=%d %s lr=%.3e grad_norm=%.3f", + global_step, + epoch, + component_text, + self.optimizer.param_groups[0]["lr"], + grad_norm, + ) + if hasattr(self.step_scheduler.dataloader, "set_postfix"): + self.step_scheduler.dataloader.set_postfix( + { + "loss": f"{group_loss:.4f}", + "lr": f"{self.optimizer.param_groups[0]['lr']:.2e}", + "gn": f"{grad_norm:.2f}", + } + ) + + if self.step_scheduler.is_ckpt_step: + self.save_checkpoint(epoch, global_step, epoch_loss / num_steps) + + if num_steps == 0: + logging.info( + "[QAD] epoch %d skipped (already completed in previous run)", epoch + 1 + ) + continue + logging.info( + "[QAD] epoch %d complete: avg_loss=%.6f", + epoch + 1, + epoch_loss / num_steps, + ) + + if is_main_process() and wandb.run is not None: + wandb.finish() + logging.info("[QAD] training complete at step %d", global_step) + finally: + self._release_distillation_controller() + + def _validate_first_step_gradients(self, global_step: int) -> None: + if global_step != 0: + return + trainable = [parameter for parameter in self.model.parameters() if parameter.requires_grad] + if not any(parameter.grad is not None for parameter in trainable): + raise RuntimeError("QAD produced no gradients for any trainable student parameter.") + if self._qad_student_settings.train_scope == "lora_only": + missing = [ + name + for name, parameter in self.model.named_parameters() + if parameter.requires_grad and parameter.grad is None + ] + if missing: + raise RuntimeError( + "SVDQuant lora_only parameters missing gradients on the first step: " + + ", ".join(missing[:5]) + ) + + def _release_distillation_controller(self) -> None: + controller = getattr(self, "_qad_controller", None) + if controller is None or not hasattr(controller, "_layers_to_loss"): + return + layer_pairs = tuple(controller._layers_to_loss) + clear_captured_outputs(controller) + mtd.export(controller) + for student_layer, teacher_layer in layer_pairs: + for layer in (student_layer, teacher_layer): + if hasattr(layer, "_intermediate_output"): + delattr(layer, "_intermediate_output") From bc47080df22e7b424eced5a6daa0e8d74cb1a1a1 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Wed, 5 Aug 2026 13:30:14 -0700 Subject: [PATCH 2/3] Save quantized Diffusers training bundles Signed-off-by: Jingyu Xin --- examples/diffusers/README.md | 18 ++- examples/diffusers/fastgen/qad/README.md | 73 ++++++++-- examples/diffusers/fastgen/qad/artifacts.py | 136 ++++++++---------- .../fastgen/qad/configs/qwen_image_nvfp4.yaml | 10 +- .../configs/qwen_image_svdquant_nvfp4.yaml | 7 +- examples/diffusers/fastgen/qad/recipe.py | 35 ++++- .../quantization/ONNX-TRT-Deployment.md | 14 +- .../quantization/build_sdxl_8bit_engine.sh | 4 +- .../diffusers/quantization/calibration.py | 36 +++++ .../diffusers/quantization/diffusion_trt.py | 17 ++- examples/diffusers/quantization/quantize.py | 122 +++++++--------- .../diffusers/quantization/quantize_config.py | 27 +++- .../fastgen/test_quant_state_roundtrip.py | 6 +- tests/examples/diffusers/test_diffusers.py | 16 +-- 14 files changed, 308 insertions(+), 213 deletions(-) diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index a9efb5fc3a3..77a93d2f304 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -101,7 +101,7 @@ python quantize.py \ --format int8 --batch-size 2 \ --calib-size 32 --alpha 0.8 --n-steps 20 \ --model-dtype {Half/BFloat16} \ - --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt \ + --output-bundle ./{MODEL_NAME}-training-bundle \ --hf-ckpt-dir ./hf_ckpt ``` @@ -112,7 +112,7 @@ python quantize.py \ --model {flux-dev|flux-schnell|sdxl-1.0|sdxl-turbo|sd3-medium|sd3.5-medium|ltx-video-dev|wan2.2-t2v-14b|wan2.2-t2v-5b} \ --model-dtype {Half|BFloat16} \ --format {fp8|fp4} --batch-size 2 --calib-size {128|256} --quantize-mha \ - --n-steps 20 --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --collect-method default \ + --n-steps 20 --output-bundle ./{MODEL_NAME}-training-bundle --collect-method default \ --hf-ckpt-dir ./hf_ckpt ``` @@ -127,7 +127,7 @@ python quantize.py \ --format fp4 --quant-algo max --collect-method default \ --model-dtype BFloat16 --trt-high-precision-dtype BFloat16 \ --batch-size 1 --calib-size 32 --n-steps 30 \ - --quantized-torch-ckpt-save-path ./wan22_vae_fp4.pt + --output-bundle ./wan22-vae-fp4-training-bundle ``` #### [LTX-2](https://github.com/Lightricks/LTX-2) FP4 @@ -146,7 +146,11 @@ python quantize.py \ > (including quantized or distilled checkpoints) remain subject to the LTX Community License > Agreement and are **not** covered by Apache 2.0. -This example produces three outputs: a PyTorch checkpoint (`--quantized-torch-ckpt-save-path`), a Hugging Face checkpoint (`--hf-ckpt-dir`), and a ComfyUI-compatible merged safetensor (`--extra-param merged_base_safetensor_path`). +LTX-2 uses a third-party pipeline that does not implement native Diffusers +`save_pretrained`/`from_pretrained`, so `--output-bundle` is not available for +this model. The example produces a Hugging Face deployment checkpoint +(`--hf-ckpt-dir`) and a ComfyUI-compatible merged safetensor +(`--extra-param merged_base_safetensor_path`). ```sh python quantize.py \ @@ -156,7 +160,6 @@ python quantize.py \ --extra-param spatial_upsampler_path=./ltx-2-spatial-upscaler-x2-1.0.safetensors \ --extra-param gemma_root=./gemma-3-12b-it-qat-q4_0-unquantized \ --extra-param fp8transformer=true \ - --quantized-torch-ckpt-save-path ./ltx-2-transformer.pt \ --hf-ckpt-dir ./LTX2-NVFP4/ \ --extra-param merged_base_safetensor_path=./ltx-2-19b-dev-fp8.safetensors ``` @@ -175,7 +178,10 @@ To additionally apply NVFP4 scale swizzle and padding , add: - `calib-size`: For SDXL INT8, we recommend 32 or 64, for SDXL FP8, 128 is recommended. - `n_steps`: Recommendation: SD/SDXL 20 or 30, SDXL-Turbo 4. -**You can use the generated checkpoint directly in PyTorch, export a Hugging Face checkpoint (`--hf-ckpt-dir`) to deploy the model on SGLang/vLLM/TRTLLM, or follow the ONNX/TensorRT workflow in [`quantization/ONNX-TRT-Deployment.md`](./quantization/ONNX-TRT-Deployment.md).** +**You can restore the generated Diffusers training bundle directly in PyTorch, +export a Hugging Face deployment checkpoint (`--hf-ckpt-dir`) for +SGLang/vLLM/TRTLLM, or follow the ONNX/TensorRT workflow in +[`quantization/ONNX-TRT-Deployment.md`](./quantization/ONNX-TRT-Deployment.md).** ## Quantization Aware Training (QAT) diff --git a/examples/diffusers/fastgen/qad/README.md b/examples/diffusers/fastgen/qad/README.md index 175f59a2e02..2ae9b305439 100644 --- a/examples/diffusers/fastgen/qad/README.md +++ b/examples/diffusers/fastgen/qad/README.md @@ -18,16 +18,20 @@ teacher and student. ## Supported students -The `qad.student.mode` field selects one of two restore contracts. +The `qad.student.mode` field selects one of two bundle validation contracts. In +both cases, `model.pretrained_model_name_or_path` is the only student artifact +path: it points to a complete, calibrated Diffusers pipeline written by +`quantize.py --output-bundle`. QAD restores the pipeline's weights and +component-local ModelOpt state together before FSDP; it does not accept a second +quantizer-state or transformer-checkpoint path. ### Regular NVFP4 Set `qad.student.mode=nvfp4` and point -`model.pretrained_model_name_or_path` at the unquantized Diffusers model. The recipe -loads its weights first and then restores the weight-free ModelOpt NVFP4 state -from `qad.student.quant_state_path`. The quantizer state must have been calibrated -against exactly the same student weights. This mode trains all student -parameters, so its only valid `train_scope` is `all`. +`model.pretrained_model_name_or_path` at a regular NVFP4 training bundle. The +bundle includes the calibrated weights and ModelOpt quantizer topology/state. +This mode trains all student parameters, so its only valid `train_scope` is +`all`. Use [`configs/qwen_image_nvfp4.yaml`](configs/qwen_image_nvfp4.yaml) as the starting configuration. @@ -48,10 +52,11 @@ For the standard Diffusers layout, the transformer files and ModelOpt sidecar are under `transformer/`, including `transformer/modelopt_state.pth`. The path given to QAD is the parent DiffusionPipeline directory. -A weight-free NVFP4 quantizer-state file is not a valid SVDQuant bundle. -SVDQuant subtracts the low-rank branch from the original weight, so both the -resulting residual weight and the PEFT factors are required. A unified -deployment export is also not a training bundle and must not be used here. +A standalone weight-free NVFP4 quantizer-state file is not a QAD student bundle. +This is especially important for SVDQuant: calibration subtracts the low-rank +branch from the original weight, so both the resulting residual weight and the +PEFT factors are required. A unified deployment export is also not a training +bundle and must not be used here. The SVDQuant topology is restored before FSDP and before optimizer construction. `qad.student.train_scope=all` is the default and trains both the residual/base @@ -66,6 +71,52 @@ as the starting configuration. QAD is restore-only in both modes. It does not calibrate a student during distributed training. +## Prepare a student bundle + +Patch Diffusers ModelMixin support and save the complete pipeline through the +quantization entry point. `quantize.py` does this automatically before model +load and calls `pipe.save_pretrained(output_bundle)` after calibration. For +example: + +```bash +# Regular NVFP4 +python examples/diffusers/quantization/quantize.py \ + --model qwen-image \ + --override-model-path /path/to/Qwen-Image \ + --model-dtype BFloat16 \ + --format fp4 \ + --quant-algo max \ + --block-size 16 \ + --batch-size 1 \ + --calib-size 32 \ + --n-steps 50 \ + --extra-param true_cfg_scale=4.0 \ + --extra-param "negative_prompt= " \ + --output-bundle /path/to/Qwen-Image-NVFP4-Calib32 + +# NVFP4 SVDQuant, rank 32 +python examples/diffusers/quantization/quantize.py \ + --model qwen-image \ + --override-model-path /path/to/Qwen-Image \ + --model-dtype BFloat16 \ + --format fp4 \ + --quant-algo svdquant \ + --lowrank 32 \ + --block-size 16 \ + --batch-size 1 \ + --calib-size 32 \ + --n-steps 50 \ + --extra-param true_cfg_scale=4.0 \ + --extra-param "negative_prompt= " \ + --output-bundle /path/to/Qwen-Image-NVFP4-SVDQuant-Calib32 +``` + +The saved root must contain `model_index.json`; the converted transformer must +contain `transformer/modelopt_state.pth`. For Qwen-Image-Flash, use `--n-steps 4` +and `--extra-param true_cfg_scale=1.0`; omit `negative_prompt`. Standard output +includes ModelOpt's full quantizer summary; capture it with `tee` and retain that +log with the bundle. + ## Distillation losses Output distillation is always MSE. The canonical setting is: @@ -169,6 +220,6 @@ The teacher and the transient ModelOpt distillation controller are not training checkpoint payloads. Checkpoints contain the student state required by the selected training scope together with optimizer, scheduler, dataloader, RNG, and global-step state. Resolved dotted CLI overrides are materialized into the saved -`config.yaml`. Resume validates the student bundle, quantizer state, mode, train +`config.yaml`. Resume validates the student bundle, quantization mode, train scope, teacher, and loss configuration before loading optimizer shards; do not change them while resuming an existing run. diff --git a/examples/diffusers/fastgen/qad/artifacts.py b/examples/diffusers/fastgen/qad/artifacts.py index 2721fb93eeb..92c977a6c50 100644 --- a/examples/diffusers/fastgen/qad/artifacts.py +++ b/examples/diffusers/fastgen/qad/artifacts.py @@ -18,12 +18,13 @@ The generic AutoModel diffusion builder intentionally owns FSDP and optimizer construction. QAD only needs two narrowly-scoped hooks around that builder: -* restore a standalone, weight-free ModelOpt state before FSDP for regular NVFP4; +* validate the ModelOpt topology restored by a native Diffusers training bundle + before FSDP; * after FSDP, optionally freeze everything except ModelOpt SVDQuant's HF PEFT A/B parameters and rebuild AdamW from the live sharded parameters. -SVDQuant itself is never calibrated here. Its complete topology and weights must -already be present in a ModelOpt-enabled Diffusers training bundle. +Quantization itself is never calibrated here. The complete topology, weights, +and quantizer buffers must already be present in the student bundle. """ from __future__ import annotations @@ -32,13 +33,11 @@ import dataclasses import inspect import logging -import os import re from typing import TYPE_CHECKING, Any import modelopt.torch.opt as mto from modelopt.torch.quantization.nn import TensorQuantizer -from modelopt.torch.quantization.utils.core_utils import set_quantizer_state_dict if TYPE_CHECKING: from collections.abc import Iterator @@ -58,7 +57,6 @@ class StudentSettings: mode: str model_name_or_path: str train_scope: str = "all" - quant_state_path: str | None = None def validate(self) -> None: if self.mode not in _SUPPORTED_STUDENT_MODES: @@ -73,18 +71,8 @@ def validate(self) -> None: f"qad.student.train_scope must be 'all' or 'lora_only', got {self.train_scope!r}." ) - if self.mode == "nvfp4": - if not self.quant_state_path: - raise ValueError( - "qad.student.quant_state_path is required for a regular NVFP4 student." - ) - if self.train_scope != "all": - raise ValueError("Regular NVFP4 supports only qad.student.train_scope=all.") - elif self.quant_state_path: - raise ValueError( - "A SVDQuant student must be restored from a complete ModelOpt-enabled " - "Diffusers training bundle; do not set qad.student.quant_state_path." - ) + if self.mode == "nvfp4" and self.train_scope != "all": + raise ValueError("Regular NVFP4 supports only qad.student.train_scope=all.") @dataclasses.dataclass @@ -96,51 +84,6 @@ class StudentBuildState: svdquant_parameter_names: tuple[str, ...] = () -def _restore_regular_quant_state(model: nn.Module, path: str, device: torch.device) -> int: - """Restore a regular weight-free ModelOpt quantizer state before FSDP.""" - if not os.path.isfile(path): - raise FileNotFoundError(f"ModelOpt quantizer state does not exist: {path}") - if mto.ModeloptStateManager.is_converted(model): - raise RuntimeError( - "The regular NVFP4 base model already contains ModelOpt topology while " - "qad.student.quant_state_path was also supplied. Use exactly one restore source." - ) - - state = mto.load_modelopt_state(path) - mode_states = state.get("modelopt_state_dict", ()) - if any( - mode_name == "svdquant_calibrate" or mode_state.get("metadata", {}).get("svdquant_peft") - for mode_name, mode_state in mode_states - ): - raise RuntimeError( - "qad.student.mode=nvfp4 cannot restore an SVDQuant state. Use " - "mode=nvfp4_svdquant with a complete ModelOpt-enabled Diffusers bundle." - ) - quantizer_state = state.pop("modelopt_state_weights", None) - if quantizer_state is None: - raise RuntimeError( - f"{path} has no modelopt_state_weights payload. Expected the weight-free " - "quantizer state produced by the Diffusers quantization example." - ) - - mto.restore_from_modelopt_state(model, state) - set_quantizer_state_dict(model, quantizer_state) - - quantizers = [module for module in model.modules() if isinstance(module, TensorQuantizer)] - if not quantizers: - raise RuntimeError(f"No TensorQuantizer modules were restored from {path}.") - for quantizer in quantizers: - quantizer.to(device) - _validate_nvfp4_quantizers(model, artifact_name=path) - - logging.info( - "[QAD] restored regular ModelOpt quantizer state before FSDP: %s (%d quantizers)", - path, - len(quantizers), - ) - return len(quantizers) - - def _is_block16_nvfp4(quantizer: TensorQuantizer) -> bool: block_sizes = quantizer.block_sizes or {} return bool( @@ -216,22 +159,58 @@ def _validate_nvfp4_quantizers( ) -def _svdquant_metadata(model: nn.Module) -> dict[str, Any] | None: +def _modelopt_mode_states(model: nn.Module) -> dict[str, dict[str, Any]]: if not mto.ModeloptStateManager.is_converted(model): - return None - for mode_name, mode_state in mto.modelopt_state(model)["modelopt_state_dict"]: - if mode_name == "svdquant_calibrate": - return mode_state.get("metadata", {}).get("svdquant_peft") - return None + return {} + return dict(mto.modelopt_state(model)["modelopt_state_dict"]) + + +def _reject_non_training_modes(mode_states: dict[str, dict[str, Any]]) -> None: + if "real_quantize" in mode_states: + raise RuntimeError( + "QAD cannot train a compressed real-quantized bundle. Recalibrate without " + "quantize.py --compress and provide the resulting fake-quantized training bundle." + ) + + +def _validate_regular_bundle(model: nn.Module) -> int: + mode_states = _modelopt_mode_states(model) + if not mode_states: + raise RuntimeError( + "qad.student.mode=nvfp4 requires a ModelOpt-aware Diffusers training bundle. " + "Calibrate it with quantize.py --output-bundle before starting QAD." + ) + _reject_non_training_modes(mode_states) + if "svdquant_calibrate" in mode_states: + raise RuntimeError( + "qad.student.mode=nvfp4 received an SVDQuant bundle; use mode=nvfp4_svdquant." + ) + quantizers = [module for module in model.modules() if isinstance(module, TensorQuantizer)] + if not quantizers: + raise RuntimeError("The regular NVFP4 student bundle restored no TensorQuantizers.") + _validate_nvfp4_quantizers(model, artifact_name="The regular NVFP4 student bundle") + logging.info( + "[QAD] validated regular ModelOpt NVFP4 bundle before FSDP: %d quantizers", + len(quantizers), + ) + return len(quantizers) def _validate_svdquant_bundle(model: nn.Module) -> tuple[str, ...]: - metadata = _svdquant_metadata(model) + mode_states = _modelopt_mode_states(model) + _reject_non_training_modes(mode_states) + mode_state = mode_states.get("svdquant_calibrate") + if mode_state is None: + raise RuntimeError( + "qad.student.mode=nvfp4_svdquant requires a bundle containing the " + "svdquant_calibrate ModelOpt mode." + ) + metadata = mode_state.get("metadata", {}).get("svdquant_peft") if not metadata: raise RuntimeError( - "qad.student.mode=nvfp4_svdquant requires a ModelOpt-enabled Diffusers " - "training bundle with svdquant_peft metadata. The topology must be restored " - "during Diffusers from_pretrained(), before FSDP." + "The SVDQuant bundle is malformed or predates the HF PEFT contract: its " + "svdquant_calibrate mode has no svdquant_peft metadata. Recalibrate it " + "with quantize.py --output-bundle." ) expected_targets = tuple(metadata.get("target_modules", ())) @@ -391,15 +370,12 @@ def apply_parallelization(pipe, parallel_scheme): state.parallel_scheme = parallel_scheme transformer = pipe.transformer if settings.mode == "nvfp4": - quant_state_path = settings.quant_state_path - assert quant_state_path is not None - state.quantizer_count = _restore_regular_quant_state( - transformer, - quant_state_path, - next(transformer.parameters()).device, - ) + state.quantizer_count = _validate_regular_bundle(transformer) else: state.svdquant_parameter_names = _validate_svdquant_bundle(transformer) + state.quantizer_count = sum( + isinstance(module, TensorQuantizer) for module in transformer.modules() + ) return original_apply_parallelization(pipe, parallel_scheme) def build_model_and_optimizer(**kwargs): diff --git a/examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml b/examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml index 70bc99f8d09..2cc67afd74c 100644 --- a/examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml +++ b/examples/diffusers/fastgen/qad/configs/qwen_image_nvfp4.yaml @@ -1,8 +1,7 @@ # Qwen-Image QAD with a regular ModelOpt NVFP4 student. # -# The student is loaded from its Diffusers checkpoint, then the weight-free -# quantizer state is restored. Paths are placeholders and should be supplied by -# dotted CLI overrides in production launches. +# model.pretrained_model_name_or_path must be a complete ModelOpt-aware +# Diffusers training bundle produced by quantize.py --output-bundle. seed: 42 @@ -16,8 +15,8 @@ dist_env: timeout_minutes: 60 model: - # Canonical student source; AutoModel records this field in checkpoints. - pretrained_model_name_or_path: Qwen/Qwen-Image + # Complete calibrated NVFP4 bundle; AutoModel records this field in checkpoints. + pretrained_model_name_or_path: /path/to/qwen-image-nvfp4-training-bundle mode: finetune step_scheduler: @@ -48,7 +47,6 @@ qad: student: mode: nvfp4 - quant_state_path: /path/to/qwen-image/quant/transformer.nvfp4.pt train_scope: all optim: diff --git a/examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml b/examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml index 99873aeae43..12ecc1a3535 100644 --- a/examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml +++ b/examples/diffusers/fastgen/qad/configs/qwen_image_svdquant_nvfp4.yaml @@ -1,8 +1,7 @@ # Qwen-Image QAD with a ModelOpt NVFP4 SVDQuant + Hugging Face PEFT student. # -# model.pretrained_model_name_or_path must be a complete, user-prepared ModelOpt-enabled -# Diffusers training bundle. A weight-free quantizer-state file or a unified -# deployment export is not sufficient. +# model.pretrained_model_name_or_path must be a complete ModelOpt-aware +# Diffusers training bundle produced by quantize.py --output-bundle. seed: 42 @@ -16,7 +15,7 @@ dist_env: timeout_minutes: 60 model: - # Complete user-prepared ModelOpt-enabled Diffusers training bundle. + # Complete calibrated NVFP4 SVDQuant bundle. pretrained_model_name_or_path: /path/to/qwen-image-nvfp4-svdquant-training-bundle mode: finetune diff --git a/examples/diffusers/fastgen/qad/recipe.py b/examples/diffusers/fastgen/qad/recipe.py index 276abac32ac..05136cd3d7a 100644 --- a/examples/diffusers/fastgen/qad/recipe.py +++ b/examples/diffusers/fastgen/qad/recipe.py @@ -98,15 +98,16 @@ def setup(self) -> None: if self.cfg.get("peft", None) is not None: raise ValueError( "Do not set AutoModel's top-level peft block for QAD. SVDQuant's " - "modelopt_svdquant HF PEFT topology must come from the student bundle." + "modelopt_svdquant HF PEFT topology comes from the student bundle." ) if str(self.cfg.get("model.mode", "finetune")).lower() != "finetune": raise ValueError("QAD supports model.mode=finetune only.") if self.cfg.get("ddp", None) is not None: raise ValueError("QAD currently supports AutoModel FSDP2, not DDP.") - # Diffusers' ModelMixin must be patched before from_pretrained so a - # SVDQuant bundle rebuilds quant/PEFT topology before loading its weights. + # Diffusers' ModelMixin must be patched before from_pretrained so both + # regular NVFP4 and SVDQuant bundles rebuild their ModelOpt topology and + # load component-local modelopt_state.pth before AutoModel applies FSDP. mto.enable_huggingface_checkpointing() with patch_student_build(settings) as build_state: @@ -186,6 +187,18 @@ def _resolve_qad_config(self) -> tuple[StudentSettings, dict[str, Any]]: raise ValueError("Missing required qad configuration block.") student_cfg = _as_dict(qad.get("student")) + secondary_artifact_fields = sorted( + field + for field in ("quant_state_path", "modelopt_state_path") + if student_cfg.get(field) is not None + ) + if secondary_artifact_fields: + raise ValueError( + "QAD accepts one complete Diffusers student bundle through " + "model.pretrained_model_name_or_path; remove unsupported secondary " + "artifact field(s): " + + ", ".join(f"qad.student.{field}" for field in secondary_artifact_fields) + ) model_name_or_path = self.cfg.get("model.pretrained_model_name_or_path", None) if not model_name_or_path: raise ValueError( @@ -207,7 +220,6 @@ def _resolve_qad_config(self) -> tuple[StudentSettings, dict[str, Any]]: settings = StudentSettings( mode=mode, model_name_or_path=str(model_name_or_path), - quant_state_path=student_cfg.get("quant_state_path"), train_scope=str(student_cfg.get("train_scope", "all")).lower(), ) settings.validate() @@ -254,7 +266,6 @@ def _resume_signature( return { "student_source": settings.model_name_or_path, "student_mode": settings.mode, - "quant_state_path": settings.quant_state_path, "train_scope": settings.train_scope, "teacher_source": loss_config["teacher_model_name_or_path"], "output_weight": float(loss_config["output_weight"]), @@ -275,6 +286,17 @@ def _resume_signature_from_saved_config(cls, config: dict[str, Any]) -> dict[str model_cfg = _as_dict(config.get("model")) qad_cfg = _as_dict(config.get("qad")) student_cfg = _as_dict(qad_cfg.get("student")) + secondary_artifact_fields = sorted( + field + for field in ("quant_state_path", "modelopt_state_path") + if student_cfg.get(field) is not None + ) + if secondary_artifact_fields: + raise RuntimeError( + "The saved QAD checkpoint uses unsupported secondary student artifact " + "field(s): " + + ", ".join(f"qad.student.{field}" for field in secondary_artifact_fields) + ) output_cfg = _as_dict(qad_cfg.get("output_loss")) task_cfg = _as_dict(qad_cfg.get("task_loss")) layerwise_cfg = _as_dict(qad_cfg.get("layerwise")) @@ -292,7 +314,6 @@ def _resume_signature_from_saved_config(cls, config: dict[str, Any]) -> dict[str settings = StudentSettings( mode=mode, model_name_or_path=str(model_cfg.get("pretrained_model_name_or_path", "")), - quant_state_path=student_cfg.get("quant_state_path"), train_scope=str(student_cfg.get("train_scope", "all")).lower(), ) return cls._resume_signature(settings, loss_config) @@ -331,7 +352,7 @@ def _validate_qad_checkpoint_signature(self, checkpoint_dir: str) -> None: raise RuntimeError( "QAD resume contract changed for " + ", ".join(changed) - + ". Use the same student artifact, quantization mode/state, train scope, " + + ". Use the same student bundle, quantization mode, train scope, " "teacher, and loss configuration as the saved run." ) diff --git a/examples/diffusers/quantization/ONNX-TRT-Deployment.md b/examples/diffusers/quantization/ONNX-TRT-Deployment.md index 57448b8a38e..fb32f91702a 100644 --- a/examples/diffusers/quantization/ONNX-TRT-Deployment.md +++ b/examples/diffusers/quantization/ONNX-TRT-Deployment.md @@ -23,7 +23,7 @@ python quantize.py \ --format int8 --batch-size 2 \ --calib-size 32 --alpha 0.8 --n-steps 20 \ --model-dtype {Half/BFloat16} --trt-high-precision-dtype {Half|BFloat16} \ - --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --onnx-dir {ONNX_DIR} + --output-bundle ./{MODEL_NAME}-training-bundle --onnx-dir {ONNX_DIR} ``` #### FLUX-Dev|SDXL|SDXL-Turbo|LTX-Video FP8/FP4 [Script](./quantize.py) @@ -34,7 +34,7 @@ python quantize.py \ python quantize.py \ --model {flux-dev|sdxl-1.0|sdxl-turbo|ltx-video-dev} --model-dtype {Half|BFloat16} --trt-high-precision-dtype {Half|BFloat16} \ --format {fp8|fp4} --batch-size 2 --calib-size {128|256} --quantize-mha \ - --n-steps 20 --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --collect-method default \ + --n-steps 20 --output-bundle ./{MODEL_NAME}-training-bundle --collect-method default \ --onnx-dir {ONNX_DIR} ``` @@ -106,7 +106,8 @@ Note, the engines must be built on the same GPU, and ensure that the INT8 engine DeviceModel is an interface designed to run TensorRT engines like torch models. It takes torch inputs and returns torch outputs. Under the hood, DeviceModel exports a torch checkpoint to ONNX and then generates a TensorRT engine from it. This allows you to swap the backbone of the diffusion pipeline with DeviceModel and execute the pipeline for your desired prompt. -Generate a quantized torch checkpoint using the [Script](./quantize.py) shown below: +Generate a native quantized Diffusers training bundle using the +[Script](./quantize.py) shown below: ```bash python quantize.py \ @@ -115,18 +116,19 @@ python quantize.py \ --batch-size {1|2} \ --calib-size 128 \ --n-steps 20 \ - --quantized-torch-ckpt-save-path ./{MODEL}_fp8.pt \ + --output-bundle ./{MODEL}-fp8-training-bundle \ --collect-method default ``` -Generate images for the quantized checkpoint with the following [Script](./diffusion_trt.py): +Generate images from the quantized bundle with the following +[Script](./diffusion_trt.py): ```bash python diffusion_trt.py \ --model {sdxl-1.0|sdxl-turbo|sd3-medium|flux-dev} \ --prompt "A cat holding a sign that says hello world" \ [--override-model-path /path/to/model] \ - [--restore-from ./{MODEL}_fp8.pt] \ + [--restore-from ./{MODEL}-fp8-training-bundle] \ [--onnx-load-path {ONNX_DIR}] \ [--trt-engine-load-path {ENGINE_DIR}] \ [--dq-only] \ diff --git a/examples/diffusers/quantization/build_sdxl_8bit_engine.sh b/examples/diffusers/quantization/build_sdxl_8bit_engine.sh index 0f09bbf2b41..43a2b60d8e4 100755 --- a/examples/diffusers/quantization/build_sdxl_8bit_engine.sh +++ b/examples/diffusers/quantization/build_sdxl_8bit_engine.sh @@ -53,9 +53,9 @@ cleaned_m="${model//\//-}" curt_exp="${cleaned_m}_${format}" echo "=====>Processing $curt_exp" if [ "$format" == "fp8" ]; then - python quantize.py --model "$model" --format "$format" --batch-size 2 --calib-size 128 --n-steps 20 --quantized-torch-ckpt-save-path "$curt_exp".pt --collect-method default --onnx-dir "$curt_exp".onnx + python quantize.py --model "$model" --format "$format" --batch-size 2 --calib-size 128 --n-steps 20 --output-bundle "$curt_exp"-training-bundle --collect-method default --onnx-dir "$curt_exp".onnx else - python quantize.py --model "$model" --format "$format" --batch-size 2 --calib-size 32 --collect-method "min-mean" --percentile 1.0 --alpha 0.8 --n-steps 20 --quantized-torch-ckpt-save-path "$curt_exp".pt --onnx-dir "$curt_exp".onnx + python quantize.py --model "$model" --format "$format" --batch-size 2 --calib-size 32 --collect-method "min-mean" --percentile 1.0 --alpha 0.8 --n-steps 20 --output-bundle "$curt_exp"-training-bundle --onnx-dir "$curt_exp".onnx fi echo "=====>Exported to ONNX model." diff --git a/examples/diffusers/quantization/calibration.py b/examples/diffusers/quantization/calibration.py index bebc61970a3..4c9b8fd1f9a 100644 --- a/examples/diffusers/quantization/calibration.py +++ b/examples/diffusers/quantization/calibration.py @@ -82,6 +82,15 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: """ self.logger.info(f"Starting calibration with {self.config.num_batches} batches") extra_args = MODEL_DEFAULTS.get(self.model_type, {}).get("inference_extra_args", {}) + if self.model_type == ModelType.QWEN_IMAGE: + extra_params = self.pipeline_manager.config.extra_params + self.logger.info( + "Qwen-Image calibration path: steps=%d true_cfg_scale=%s " + "negative_prompt=%s output_type=latent", + self.config.n_steps, + extra_params.get("true_cfg_scale", "pipeline default"), + "provided" if "negative_prompt" in extra_params else "omitted", + ) with tqdm(total=self.config.num_batches, desc="Calibration", unit="batch") as pbar: for i, prompt_batch in enumerate(batched_prompts): @@ -99,6 +108,8 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: 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) + elif self.model_type == ModelType.QWEN_IMAGE: + self._run_qwen_image_calibration(prompt_batch, extra_args) else: common_args = { "prompt": prompt_batch, @@ -109,6 +120,31 @@ 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_calibration( + self, prompt_batch: list[str], extra_args: dict[str, Any] + ) -> None: + """Run Qwen-Image's standard denoising loop without the unused VAE decode. + + ``true_cfg_scale`` alone does not enable true CFG in QwenImagePipeline; a + negative prompt is also required. Keep both values explicit launch-time + parameters so the 50-step base model and four-step Flash model can share + this model type while calibrating against their actual inference paths. + """ + extra_params = self.pipeline_manager.config.extra_params + kwargs = { + "height": extra_params.get("height", extra_args.get("height", 1024)), + "width": extra_params.get("width", extra_args.get("width", 1024)), + "num_inference_steps": self.config.n_steps, + "output_type": "latent", + } + for name in ("negative_prompt", "true_cfg_scale", "guidance_scale"): + if name in extra_params: + value = extra_params[name] + if name == "negative_prompt" and isinstance(value, str): + value = [value] * len(prompt_batch) + kwargs[name] = value + self.pipe(prompt=prompt_batch, **kwargs).images + def _run_qwen_image_dmd2_calibration(self, prompt_batch: list[str]) -> None: """Calibrate a DMD2 Qwen-Image student via its few-step sampler. diff --git a/examples/diffusers/quantization/diffusion_trt.py b/examples/diffusers/quantization/diffusion_trt.py index eb62f2b0937..3897220c97d 100644 --- a/examples/diffusers/quantization/diffusion_trt.py +++ b/examples/diffusers/quantization/diffusion_trt.py @@ -15,6 +15,7 @@ import argparse from contextlib import nullcontext +from pathlib import Path import numpy as np import torch @@ -151,7 +152,10 @@ def main(): help="Path to the model if not using default paths in MODEL_ID mapping.", ) parser.add_argument( - "--restore-from", type=str, default=None, help="Path to the modelopt quantized checkpoint" + "--restore-from", + type=str, + default=None, + help="Native Diffusers training bundle or legacy ModelOpt backbone checkpoint", ) parser.add_argument( "--prompt", @@ -196,11 +200,18 @@ def main(): image_name = args.save_image_as or f"{args.model}.png" model_dtype = DTYPE_MAP[args.model] + restore_bundle = bool( + args.restore_from + and Path(args.restore_from).is_dir() + and (Path(args.restore_from) / "model_index.json").is_file() + ) + if restore_bundle: + mto.enable_huggingface_checkpointing() pipe = PipelineManager.create_pipeline_from( MODEL_ID[args.model], torch_dtype=model_dtype, - override_model_path=args.override_model_path, + override_model_path=args.restore_from if restore_bundle else args.override_model_path, ) if args.torch_compile: @@ -218,7 +229,7 @@ def main(): else: raise ValueError("Pipeline does not have a transformer or unet backbone") - if args.restore_from: + if args.restore_from and not restore_bundle: mto.restore(backbone, args.restore_from) if args.torch: diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 26aac8ae07f..dace4e9bafc 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -50,8 +50,9 @@ QuantFormat, QuantizationConfig, ) -from utils import check_conv_and_mha, check_lora, restore_quantizer_state, save_quantizer_state +from utils import check_conv_and_mha, check_lora +import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.export import export_hf_checkpoint @@ -281,40 +282,34 @@ def _has_conv_layers(self, model: torch.nn.Module) -> bool: True if model contains Conv layers, False otherwise """ for module in model.modules(): - if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)) and ( + if isinstance(module, torch.nn.Conv1d | torch.nn.Conv2d | torch.nn.Conv3d) and ( module.input_quantizer.is_enabled or module.weight_quantizer.is_enabled ): return True return False - def save_checkpoint( - self, - backbone: torch.nn.Module, - backbone_name: str | None = None, - ) -> None: - """ - Save quantized model checkpoint. - - Args: - backbone: The quantized backbone module to save (must be the same instance - that was passed to mtq.quantize, as it carries the _modelopt_state). - backbone_name: Optional name for the backbone file (defaults to "backbone"). - """ - if not self.config.quantized_torch_ckpt_path: + def save_training_bundle(self, pipe: DiffusionPipeline) -> None: + """Save the calibrated pipeline in native, ModelOpt-aware Diffusers format.""" + if not self.config.output_bundle: return - ckpt_path = self.config.quantized_torch_ckpt_path - ckpt_path.mkdir(parents=True, exist_ok=True) - filename = f"{backbone_name}.pt" if backbone_name else "backbone.pt" - target_path = ckpt_path / filename - - # 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)) + output_bundle = self.config.output_bundle + self.logger.info("Saving ModelOpt training bundle to %s", output_bundle) + pipe.save_pretrained(output_bundle) - self.logger.info("Checkpoint saved successfully") + model_index_path = output_bundle / "model_index.json" + if not model_index_path.is_file(): + raise RuntimeError(f"Training bundle is missing {model_index_path}.") + if self.pipeline_manager is None: + raise RuntimeError("Pipeline manager is required to validate the training bundle.") + for backbone_name, backbone in self.pipeline_manager.iter_backbones(): + if mto.ModeloptStateManager.is_converted(backbone): + state_path = output_bundle / backbone_name / "modelopt_state.pth" + if not state_path.is_file(): + raise RuntimeError( + f"ModelOpt state was not saved for {backbone_name}: {state_path}" + ) + self.logger.info("ModelOpt training bundle saved successfully") def export_onnx( self, @@ -360,34 +355,6 @@ def export_onnx( self.logger.info("ONNX export completed successfully") - def restore_checkpoint(self) -> None: - """ - Restore a previously quantized model. - - """ - if not self.config.restore_from: - return - - restore_path = self.config.restore_from - if self.pipeline_manager is None: - raise RuntimeError("Pipeline manager is required for per-backbone checkpoints.") - - if not restore_path.exists() or not restore_path.is_dir(): - raise FileNotFoundError(f"Checkpoint directory not found: {restore_path}") - - for backbone_name, backbone in self.pipeline_manager.iter_backbones(): - source_path = restore_path / f"{backbone_name}.pt" - if not source_path.exists(): - raise FileNotFoundError( - f"Checkpoint not found for '{backbone_name}' in {restore_path}" - ) - self.logger.info(f"Restoring {backbone_name} from {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") - # TODO: should not do the any data type def export_hf_ckpt(self, pipe: Any, model_config: ModelConfig | None = None) -> None: """ @@ -460,8 +427,8 @@ def create_argument_parser() -> argparse.ArgumentParser: # Faster LTX-Video quantization (skip upsampler) %(prog)s --model ltx-video-dev --format fp8 --batch-size 1 --calib-size 32 --ltx-skip-upsampler - # Restore and export a previously quantized model - %(prog)s --model flux-schnell --restore-from checkpoint.pt --onnx-dir ./exports/ + # Restore and export a previously quantized native training bundle + %(prog)s --model flux-schnell --restore-from ./flux-schnell-int8 --onnx-dir ./exports/ """, ) model_group = parser.add_argument_group("Model Configuration") @@ -579,9 +546,14 @@ def create_argument_parser() -> argparse.ArgumentParser: export_group = parser.add_argument_group("Export Configuration") export_group.add_argument( + "--output-bundle", "--quantized-torch-ckpt-save-path", + dest="output_bundle", type=str, - help="Path to save quantized PyTorch checkpoint", + help=( + "Directory for the native ModelOpt-aware Diffusers training bundle. " + "The legacy --quantized-torch-ckpt-save-path spelling is accepted as an alias." + ), ) export_group.add_argument("--onnx-dir", type=str, help="Directory for ONNX export") export_group.add_argument( @@ -590,7 +562,7 @@ def create_argument_parser() -> argparse.ArgumentParser: help="Directory for HuggingFace checkpoint export", ) export_group.add_argument( - "--restore-from", type=str, help="Path to restore from previous checkpoint" + "--restore-from", type=str, help="Native Diffusers training bundle to restore" ) export_group.add_argument( "--trt-high-precision-dtype", @@ -613,7 +585,17 @@ def main() -> None: parser = create_argument_parser() args, unknown_args = parser.parse_known_args() + # Patch Diffusers ModelMixin before any pipeline is loaded or saved. Converted + # components then restore/save their topology in /modelopt_state.pth. + mto.enable_huggingface_checkpointing() + model_type = ModelType(args.model) + if model_type == ModelType.LTX2 and (args.output_bundle or args.restore_from): + parser.error( + "LTX-2 uses a third-party TI2VidTwoStagesPipeline without native Diffusers " + "save_pretrained/from_pretrained support, so --output-bundle and " + "--restore-from are unavailable. Use --hf-ckpt-dir for its deployment export." + ) if args.backbone is None: args.backbone = [MODEL_DEFAULTS[model_type]["backbone"]] s = time.time() @@ -633,9 +615,13 @@ def main() -> None: model_dtype=model_dtype, backbone=args.backbone, trt_high_precision_dtype=DataType(args.trt_high_precision_dtype), - override_model_path=Path(args.override_model_path) - if args.override_model_path - else None, + override_model_path=( + Path(args.restore_from) + if args.restore_from + else Path(args.override_model_path) + if args.override_model_path + else None + ), cpu_offloading=args.cpu_offloading, ltx_skip_upsampler=args.ltx_skip_upsampler, extra_params=extra_params, @@ -669,9 +655,7 @@ def main() -> None: ) export_config = ExportConfig( - quantized_torch_ckpt_path=Path(args.quantized_torch_ckpt_save_path) - if args.quantized_torch_ckpt_save_path - else None, + output_bundle=Path(args.output_bundle) if args.output_bundle else None, onnx_dir=Path(args.onnx_dir) if args.onnx_dir else None, hf_ckpt_dir=Path(args.hf_ckpt_dir) if args.hf_ckpt_dir else None, restore_from=Path(args.restore_from) if args.restore_from else None, @@ -689,10 +673,7 @@ def main() -> None: export_manager = ExportManager(export_config, logger, pipeline_manager) - if export_config.restore_from and export_config.restore_from.exists(): - export_manager.restore_checkpoint() - - else: + if not export_config.restore_from: logger.info("Initializing calibration...") calibrator = Calibrator(pipeline_manager, calib_config, model_config.model_type, logger) batched_prompts = calibrator.load_and_batch_prompts() @@ -727,9 +708,8 @@ def forward_loop(mod): backbone, quant_config.format == QuantFormat.FP4, quant_config.quantize_mha ) - export_manager.save_checkpoint(backbone, backbone_name) - pipeline_manager.print_quant_summary() + export_manager.save_training_bundle(pipe) for backbone_name, backbone in pipeline_manager.iter_backbones(): export_manager.export_onnx( diff --git a/examples/diffusers/quantization/quantize_config.py b/examples/diffusers/quantization/quantize_config.py index a92dd4e8147..173ec420f1c 100644 --- a/examples/diffusers/quantization/quantize_config.py +++ b/examples/diffusers/quantization/quantize_config.py @@ -140,20 +140,35 @@ def model_path(self) -> str: class ExportConfig: """Configuration for model export.""" - quantized_torch_ckpt_path: Path | None = None + output_bundle: Path | None = None onnx_dir: Path | None = None hf_ckpt_dir: Path | None = None restore_from: Path | None = None def validate(self) -> None: """Validate export configuration.""" - if self.restore_from and not self.restore_from.exists(): - raise FileNotFoundError(f"Restore checkpoint not found: {self.restore_from}") - - if self.quantized_torch_ckpt_path: - parent_dir = self.quantized_torch_ckpt_path.parent + if self.restore_from: + if not self.restore_from.is_dir(): + raise FileNotFoundError( + f"Diffusers training bundle directory not found: {self.restore_from}" + ) + if not (self.restore_from / "model_index.json").is_file(): + raise FileNotFoundError( + f"Diffusers training bundle is missing model_index.json: {self.restore_from}" + ) + + if self.output_bundle: + parent_dir = self.output_bundle.parent if not parent_dir.exists(): parent_dir.mkdir(parents=True, exist_ok=True) + if self.output_bundle.exists() and not self.output_bundle.is_dir(): + raise FileExistsError( + f"Output training bundle path is not a directory: {self.output_bundle}" + ) + if self.output_bundle.exists() and any(self.output_bundle.iterdir()): + raise FileExistsError( + f"Output training bundle directory is not empty: {self.output_bundle}" + ) if self.onnx_dir and not self.onnx_dir.exists(): self.onnx_dir.mkdir(parents=True, exist_ok=True) diff --git a/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py b/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py index 2747012ec59..e4d7a724b80 100644 --- a/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py +++ b/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py @@ -23,9 +23,9 @@ 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. +The on-disk state is built here with ModelOpt's own weight-free state idiom. This pins +compatibility with the historical DMD2 QAT artifacts independently from the current +quantization example's complete Diffusers training-bundle output. Dependency-guarded with ``importorskip`` so it skips where torch / modelopt are absent. """ diff --git a/tests/examples/diffusers/test_diffusers.py b/tests/examples/diffusers/test_diffusers.py index 15c5eb44934..d6840f59c31 100644 --- a/tests/examples/diffusers/test_diffusers.py +++ b/tests/examples/diffusers/test_diffusers.py @@ -68,8 +68,8 @@ def quantize(self, tmp_path: Path) -> None: *self._format_args(), "--trt-high-precision-dtype", self.dtype, - "--quantized-torch-ckpt-save-path", - str(tmp_path / f"{self.name}_{self.format_type}.pt"), + "--output-bundle", + str(tmp_path / f"{self.name}_{self.format_type}_bundle"), "--onnx-dir", str(tmp_path / f"{self.name}_{self.format_type}_onnx"), ) @@ -81,7 +81,7 @@ def restore(self, tmp_path: Path) -> None: "--trt-high-precision-dtype", self.dtype, "--restore-from", - str(tmp_path / f"{self.name}_{self.format_type}.pt"), + str(tmp_path / f"{self.name}_{self.format_type}_bundle"), "--onnx-dir", str(tmp_path / f"{self.name}_{self.format_type}_onnx"), ) @@ -160,10 +160,10 @@ class Wan22Model(NamedTuple): quant_algo: str collect_method: str - def _ckpt_path(self, tmp_path: Path) -> str: + def _bundle_path(self, tmp_path: Path) -> str: stem = self.model.replace("wan2.2-t2v-", "") parts = [stem, *([self.backbone] if self.backbone else []), self.format_type] - return str(tmp_path / f"wan22_{'_'.join(parts)}.pt") + return str(tmp_path / f"wan22_{'_'.join(parts)}_bundle") def _common_args(self, tiny_wan22_path: str) -> list[str]: cmd_args = [ @@ -205,15 +205,15 @@ def quantize(self, tiny_wan22_path: str, tmp_path: Path) -> None: run_example_command( [ *self._common_args(tiny_wan22_path), - "--quantized-torch-ckpt-save-path", - self._ckpt_path(tmp_path), + "--output-bundle", + self._bundle_path(tmp_path), ], "diffusers/quantization", ) def restore(self, tiny_wan22_path: str, tmp_path: Path) -> None: run_example_command( - [*self._common_args(tiny_wan22_path), "--restore-from", self._ckpt_path(tmp_path)], + [*self._common_args(tiny_wan22_path), "--restore-from", self._bundle_path(tmp_path)], "diffusers/quantization", ) From 8e21cd43ab78c258b13f349c46ad64cd812bb057 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Wed, 5 Aug 2026 13:58:13 -0700 Subject: [PATCH 3/3] Limit QAD changes to training flow Signed-off-by: Jingyu Xin --- examples/diffusers/README.md | 18 +-- examples/diffusers/fastgen/qad/README.md | 4 +- .../quantization/ONNX-TRT-Deployment.md | 14 +-- .../quantization/build_sdxl_8bit_engine.sh | 4 +- .../diffusers/quantization/diffusion_trt.py | 17 +-- examples/diffusers/quantization/quantize.py | 113 +++++++++++++----- .../diffusers/quantization/quantize_config.py | 17 ++- .../fastgen/test_quant_state_roundtrip.py | 6 +- tests/examples/diffusers/test_diffusers.py | 16 +-- 9 files changed, 122 insertions(+), 87 deletions(-) diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 77a93d2f304..a9efb5fc3a3 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -101,7 +101,7 @@ python quantize.py \ --format int8 --batch-size 2 \ --calib-size 32 --alpha 0.8 --n-steps 20 \ --model-dtype {Half/BFloat16} \ - --output-bundle ./{MODEL_NAME}-training-bundle \ + --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt \ --hf-ckpt-dir ./hf_ckpt ``` @@ -112,7 +112,7 @@ python quantize.py \ --model {flux-dev|flux-schnell|sdxl-1.0|sdxl-turbo|sd3-medium|sd3.5-medium|ltx-video-dev|wan2.2-t2v-14b|wan2.2-t2v-5b} \ --model-dtype {Half|BFloat16} \ --format {fp8|fp4} --batch-size 2 --calib-size {128|256} --quantize-mha \ - --n-steps 20 --output-bundle ./{MODEL_NAME}-training-bundle --collect-method default \ + --n-steps 20 --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --collect-method default \ --hf-ckpt-dir ./hf_ckpt ``` @@ -127,7 +127,7 @@ python quantize.py \ --format fp4 --quant-algo max --collect-method default \ --model-dtype BFloat16 --trt-high-precision-dtype BFloat16 \ --batch-size 1 --calib-size 32 --n-steps 30 \ - --output-bundle ./wan22-vae-fp4-training-bundle + --quantized-torch-ckpt-save-path ./wan22_vae_fp4.pt ``` #### [LTX-2](https://github.com/Lightricks/LTX-2) FP4 @@ -146,11 +146,7 @@ python quantize.py \ > (including quantized or distilled checkpoints) remain subject to the LTX Community License > Agreement and are **not** covered by Apache 2.0. -LTX-2 uses a third-party pipeline that does not implement native Diffusers -`save_pretrained`/`from_pretrained`, so `--output-bundle` is not available for -this model. The example produces a Hugging Face deployment checkpoint -(`--hf-ckpt-dir`) and a ComfyUI-compatible merged safetensor -(`--extra-param merged_base_safetensor_path`). +This example produces three outputs: a PyTorch checkpoint (`--quantized-torch-ckpt-save-path`), a Hugging Face checkpoint (`--hf-ckpt-dir`), and a ComfyUI-compatible merged safetensor (`--extra-param merged_base_safetensor_path`). ```sh python quantize.py \ @@ -160,6 +156,7 @@ python quantize.py \ --extra-param spatial_upsampler_path=./ltx-2-spatial-upscaler-x2-1.0.safetensors \ --extra-param gemma_root=./gemma-3-12b-it-qat-q4_0-unquantized \ --extra-param fp8transformer=true \ + --quantized-torch-ckpt-save-path ./ltx-2-transformer.pt \ --hf-ckpt-dir ./LTX2-NVFP4/ \ --extra-param merged_base_safetensor_path=./ltx-2-19b-dev-fp8.safetensors ``` @@ -178,10 +175,7 @@ To additionally apply NVFP4 scale swizzle and padding , add: - `calib-size`: For SDXL INT8, we recommend 32 or 64, for SDXL FP8, 128 is recommended. - `n_steps`: Recommendation: SD/SDXL 20 or 30, SDXL-Turbo 4. -**You can restore the generated Diffusers training bundle directly in PyTorch, -export a Hugging Face deployment checkpoint (`--hf-ckpt-dir`) for -SGLang/vLLM/TRTLLM, or follow the ONNX/TensorRT workflow in -[`quantization/ONNX-TRT-Deployment.md`](./quantization/ONNX-TRT-Deployment.md).** +**You can use the generated checkpoint directly in PyTorch, export a Hugging Face checkpoint (`--hf-ckpt-dir`) to deploy the model on SGLang/vLLM/TRTLLM, or follow the ONNX/TensorRT workflow in [`quantization/ONNX-TRT-Deployment.md`](./quantization/ONNX-TRT-Deployment.md).** ## Quantization Aware Training (QAT) diff --git a/examples/diffusers/fastgen/qad/README.md b/examples/diffusers/fastgen/qad/README.md index 2ae9b305439..d887d637444 100644 --- a/examples/diffusers/fastgen/qad/README.md +++ b/examples/diffusers/fastgen/qad/README.md @@ -55,8 +55,8 @@ given to QAD is the parent DiffusionPipeline directory. A standalone weight-free NVFP4 quantizer-state file is not a QAD student bundle. This is especially important for SVDQuant: calibration subtracts the low-rank branch from the original weight, so both the resulting residual weight and the -PEFT factors are required. A unified deployment export is also not a training -bundle and must not be used here. +PEFT factors are required. Deployment artifacts are not training bundles and +must not be used here. The SVDQuant topology is restored before FSDP and before optimizer construction. `qad.student.train_scope=all` is the default and trains both the residual/base diff --git a/examples/diffusers/quantization/ONNX-TRT-Deployment.md b/examples/diffusers/quantization/ONNX-TRT-Deployment.md index fb32f91702a..57448b8a38e 100644 --- a/examples/diffusers/quantization/ONNX-TRT-Deployment.md +++ b/examples/diffusers/quantization/ONNX-TRT-Deployment.md @@ -23,7 +23,7 @@ python quantize.py \ --format int8 --batch-size 2 \ --calib-size 32 --alpha 0.8 --n-steps 20 \ --model-dtype {Half/BFloat16} --trt-high-precision-dtype {Half|BFloat16} \ - --output-bundle ./{MODEL_NAME}-training-bundle --onnx-dir {ONNX_DIR} + --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --onnx-dir {ONNX_DIR} ``` #### FLUX-Dev|SDXL|SDXL-Turbo|LTX-Video FP8/FP4 [Script](./quantize.py) @@ -34,7 +34,7 @@ python quantize.py \ python quantize.py \ --model {flux-dev|sdxl-1.0|sdxl-turbo|ltx-video-dev} --model-dtype {Half|BFloat16} --trt-high-precision-dtype {Half|BFloat16} \ --format {fp8|fp4} --batch-size 2 --calib-size {128|256} --quantize-mha \ - --n-steps 20 --output-bundle ./{MODEL_NAME}-training-bundle --collect-method default \ + --n-steps 20 --quantized-torch-ckpt-save-path ./{MODEL_NAME}.pt --collect-method default \ --onnx-dir {ONNX_DIR} ``` @@ -106,8 +106,7 @@ Note, the engines must be built on the same GPU, and ensure that the INT8 engine DeviceModel is an interface designed to run TensorRT engines like torch models. It takes torch inputs and returns torch outputs. Under the hood, DeviceModel exports a torch checkpoint to ONNX and then generates a TensorRT engine from it. This allows you to swap the backbone of the diffusion pipeline with DeviceModel and execute the pipeline for your desired prompt. -Generate a native quantized Diffusers training bundle using the -[Script](./quantize.py) shown below: +Generate a quantized torch checkpoint using the [Script](./quantize.py) shown below: ```bash python quantize.py \ @@ -116,19 +115,18 @@ python quantize.py \ --batch-size {1|2} \ --calib-size 128 \ --n-steps 20 \ - --output-bundle ./{MODEL}-fp8-training-bundle \ + --quantized-torch-ckpt-save-path ./{MODEL}_fp8.pt \ --collect-method default ``` -Generate images from the quantized bundle with the following -[Script](./diffusion_trt.py): +Generate images for the quantized checkpoint with the following [Script](./diffusion_trt.py): ```bash python diffusion_trt.py \ --model {sdxl-1.0|sdxl-turbo|sd3-medium|flux-dev} \ --prompt "A cat holding a sign that says hello world" \ [--override-model-path /path/to/model] \ - [--restore-from ./{MODEL}-fp8-training-bundle] \ + [--restore-from ./{MODEL}_fp8.pt] \ [--onnx-load-path {ONNX_DIR}] \ [--trt-engine-load-path {ENGINE_DIR}] \ [--dq-only] \ diff --git a/examples/diffusers/quantization/build_sdxl_8bit_engine.sh b/examples/diffusers/quantization/build_sdxl_8bit_engine.sh index 43a2b60d8e4..0f09bbf2b41 100755 --- a/examples/diffusers/quantization/build_sdxl_8bit_engine.sh +++ b/examples/diffusers/quantization/build_sdxl_8bit_engine.sh @@ -53,9 +53,9 @@ cleaned_m="${model//\//-}" curt_exp="${cleaned_m}_${format}" echo "=====>Processing $curt_exp" if [ "$format" == "fp8" ]; then - python quantize.py --model "$model" --format "$format" --batch-size 2 --calib-size 128 --n-steps 20 --output-bundle "$curt_exp"-training-bundle --collect-method default --onnx-dir "$curt_exp".onnx + python quantize.py --model "$model" --format "$format" --batch-size 2 --calib-size 128 --n-steps 20 --quantized-torch-ckpt-save-path "$curt_exp".pt --collect-method default --onnx-dir "$curt_exp".onnx else - python quantize.py --model "$model" --format "$format" --batch-size 2 --calib-size 32 --collect-method "min-mean" --percentile 1.0 --alpha 0.8 --n-steps 20 --output-bundle "$curt_exp"-training-bundle --onnx-dir "$curt_exp".onnx + python quantize.py --model "$model" --format "$format" --batch-size 2 --calib-size 32 --collect-method "min-mean" --percentile 1.0 --alpha 0.8 --n-steps 20 --quantized-torch-ckpt-save-path "$curt_exp".pt --onnx-dir "$curt_exp".onnx fi echo "=====>Exported to ONNX model." diff --git a/examples/diffusers/quantization/diffusion_trt.py b/examples/diffusers/quantization/diffusion_trt.py index 3897220c97d..eb62f2b0937 100644 --- a/examples/diffusers/quantization/diffusion_trt.py +++ b/examples/diffusers/quantization/diffusion_trt.py @@ -15,7 +15,6 @@ import argparse from contextlib import nullcontext -from pathlib import Path import numpy as np import torch @@ -152,10 +151,7 @@ def main(): help="Path to the model if not using default paths in MODEL_ID mapping.", ) parser.add_argument( - "--restore-from", - type=str, - default=None, - help="Native Diffusers training bundle or legacy ModelOpt backbone checkpoint", + "--restore-from", type=str, default=None, help="Path to the modelopt quantized checkpoint" ) parser.add_argument( "--prompt", @@ -200,18 +196,11 @@ def main(): image_name = args.save_image_as or f"{args.model}.png" model_dtype = DTYPE_MAP[args.model] - restore_bundle = bool( - args.restore_from - and Path(args.restore_from).is_dir() - and (Path(args.restore_from) / "model_index.json").is_file() - ) - if restore_bundle: - mto.enable_huggingface_checkpointing() pipe = PipelineManager.create_pipeline_from( MODEL_ID[args.model], torch_dtype=model_dtype, - override_model_path=args.restore_from if restore_bundle else args.override_model_path, + override_model_path=args.override_model_path, ) if args.torch_compile: @@ -229,7 +218,7 @@ def main(): else: raise ValueError("Pipeline does not have a transformer or unet backbone") - if args.restore_from and not restore_bundle: + if args.restore_from: mto.restore(backbone, args.restore_from) if args.torch: diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index dace4e9bafc..5fbaa7c9ccb 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -50,7 +50,7 @@ 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 @@ -282,14 +282,43 @@ def _has_conv_layers(self, model: torch.nn.Module) -> bool: True if model contains Conv layers, False otherwise """ for module in model.modules(): - if isinstance(module, torch.nn.Conv1d | torch.nn.Conv2d | torch.nn.Conv3d) and ( + if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)) and ( module.input_quantizer.is_enabled or module.weight_quantizer.is_enabled ): return True return False + def save_checkpoint( + self, + backbone: torch.nn.Module, + backbone_name: str | None = None, + ) -> None: + """ + Save quantized model checkpoint. + + Args: + backbone: The quantized backbone module to save (must be the same instance + that was passed to mtq.quantize, as it carries the _modelopt_state). + backbone_name: Optional name for the backbone file (defaults to "backbone"). + """ + if not self.config.quantized_torch_ckpt_path: + return + + ckpt_path = self.config.quantized_torch_ckpt_path + ckpt_path.mkdir(parents=True, exist_ok=True) + filename = f"{backbone_name}.pt" if backbone_name else "backbone.pt" + target_path = ckpt_path / filename + + # 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") + def save_training_bundle(self, pipe: DiffusionPipeline) -> None: - """Save the calibrated pipeline in native, ModelOpt-aware Diffusers format.""" + """Save a calibrated pipeline for QAD through native Diffusers checkpointing.""" if not self.config.output_bundle: return @@ -355,6 +384,34 @@ def export_onnx( self.logger.info("ONNX export completed successfully") + def restore_checkpoint(self) -> None: + """ + Restore a previously quantized model. + + """ + if not self.config.restore_from: + return + + restore_path = self.config.restore_from + if self.pipeline_manager is None: + raise RuntimeError("Pipeline manager is required for per-backbone checkpoints.") + + if not restore_path.exists() or not restore_path.is_dir(): + raise FileNotFoundError(f"Checkpoint directory not found: {restore_path}") + + for backbone_name, backbone in self.pipeline_manager.iter_backbones(): + source_path = restore_path / f"{backbone_name}.pt" + if not source_path.exists(): + raise FileNotFoundError( + f"Checkpoint not found for '{backbone_name}' in {restore_path}" + ) + self.logger.info(f"Restoring {backbone_name} from {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") + # TODO: should not do the any data type def export_hf_ckpt(self, pipe: Any, model_config: ModelConfig | None = None) -> None: """ @@ -427,8 +484,8 @@ def create_argument_parser() -> argparse.ArgumentParser: # Faster LTX-Video quantization (skip upsampler) %(prog)s --model ltx-video-dev --format fp8 --batch-size 1 --calib-size 32 --ltx-skip-upsampler - # Restore and export a previously quantized native training bundle - %(prog)s --model flux-schnell --restore-from ./flux-schnell-int8 --onnx-dir ./exports/ + # Restore and export a previously quantized model + %(prog)s --model flux-schnell --restore-from checkpoint.pt --onnx-dir ./exports/ """, ) model_group = parser.add_argument_group("Model Configuration") @@ -546,14 +603,14 @@ def create_argument_parser() -> argparse.ArgumentParser: export_group = parser.add_argument_group("Export Configuration") export_group.add_argument( - "--output-bundle", "--quantized-torch-ckpt-save-path", - dest="output_bundle", type=str, - help=( - "Directory for the native ModelOpt-aware Diffusers training bundle. " - "The legacy --quantized-torch-ckpt-save-path spelling is accepted as an alias." - ), + help="Path to save quantized PyTorch checkpoint", + ) + export_group.add_argument( + "--output-bundle", + type=str, + help="Directory for a native ModelOpt-aware Diffusers training bundle used by QAD", ) export_group.add_argument("--onnx-dir", type=str, help="Directory for ONNX export") export_group.add_argument( @@ -562,7 +619,7 @@ def create_argument_parser() -> argparse.ArgumentParser: help="Directory for HuggingFace checkpoint export", ) export_group.add_argument( - "--restore-from", type=str, help="Native Diffusers training bundle to restore" + "--restore-from", type=str, help="Path to restore from previous checkpoint" ) export_group.add_argument( "--trt-high-precision-dtype", @@ -585,17 +642,11 @@ def main() -> None: parser = create_argument_parser() args, unknown_args = parser.parse_known_args() - # Patch Diffusers ModelMixin before any pipeline is loaded or saved. Converted - # components then restore/save their topology in /modelopt_state.pth. - mto.enable_huggingface_checkpointing() + if args.output_bundle: + # Install ModelOpt's Diffusers save/load hooks before pipeline construction. + mto.enable_huggingface_checkpointing() model_type = ModelType(args.model) - if model_type == ModelType.LTX2 and (args.output_bundle or args.restore_from): - parser.error( - "LTX-2 uses a third-party TI2VidTwoStagesPipeline without native Diffusers " - "save_pretrained/from_pretrained support, so --output-bundle and " - "--restore-from are unavailable. Use --hf-ckpt-dir for its deployment export." - ) if args.backbone is None: args.backbone = [MODEL_DEFAULTS[model_type]["backbone"]] s = time.time() @@ -615,13 +666,9 @@ def main() -> None: model_dtype=model_dtype, backbone=args.backbone, trt_high_precision_dtype=DataType(args.trt_high_precision_dtype), - override_model_path=( - Path(args.restore_from) - if args.restore_from - else Path(args.override_model_path) - if args.override_model_path - else None - ), + override_model_path=Path(args.override_model_path) + if args.override_model_path + else None, cpu_offloading=args.cpu_offloading, ltx_skip_upsampler=args.ltx_skip_upsampler, extra_params=extra_params, @@ -655,6 +702,9 @@ def main() -> None: ) export_config = ExportConfig( + quantized_torch_ckpt_path=Path(args.quantized_torch_ckpt_save_path) + if args.quantized_torch_ckpt_save_path + else None, output_bundle=Path(args.output_bundle) if args.output_bundle else None, onnx_dir=Path(args.onnx_dir) if args.onnx_dir else None, hf_ckpt_dir=Path(args.hf_ckpt_dir) if args.hf_ckpt_dir else None, @@ -673,7 +723,10 @@ def main() -> None: export_manager = ExportManager(export_config, logger, pipeline_manager) - if not export_config.restore_from: + if export_config.restore_from and export_config.restore_from.exists(): + export_manager.restore_checkpoint() + + else: logger.info("Initializing calibration...") calibrator = Calibrator(pipeline_manager, calib_config, model_config.model_type, logger) batched_prompts = calibrator.load_and_batch_prompts() @@ -708,6 +761,8 @@ def forward_loop(mod): backbone, quant_config.format == QuantFormat.FP4, quant_config.quantize_mha ) + export_manager.save_checkpoint(backbone, backbone_name) + pipeline_manager.print_quant_summary() export_manager.save_training_bundle(pipe) diff --git a/examples/diffusers/quantization/quantize_config.py b/examples/diffusers/quantization/quantize_config.py index 173ec420f1c..7bbb724ee59 100644 --- a/examples/diffusers/quantization/quantize_config.py +++ b/examples/diffusers/quantization/quantize_config.py @@ -140,6 +140,7 @@ def model_path(self) -> str: class ExportConfig: """Configuration for model export.""" + quantized_torch_ckpt_path: Path | None = None output_bundle: Path | None = None onnx_dir: Path | None = None hf_ckpt_dir: Path | None = None @@ -147,15 +148,13 @@ class ExportConfig: def validate(self) -> None: """Validate export configuration.""" - if self.restore_from: - if not self.restore_from.is_dir(): - raise FileNotFoundError( - f"Diffusers training bundle directory not found: {self.restore_from}" - ) - if not (self.restore_from / "model_index.json").is_file(): - raise FileNotFoundError( - f"Diffusers training bundle is missing model_index.json: {self.restore_from}" - ) + if self.restore_from and not self.restore_from.exists(): + raise FileNotFoundError(f"Restore checkpoint not found: {self.restore_from}") + + if self.quantized_torch_ckpt_path: + parent_dir = self.quantized_torch_ckpt_path.parent + if not parent_dir.exists(): + parent_dir.mkdir(parents=True, exist_ok=True) if self.output_bundle: parent_dir = self.output_bundle.parent diff --git a/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py b/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py index e4d7a724b80..2747012ec59 100644 --- a/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py +++ b/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py @@ -23,9 +23,9 @@ 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 weight-free state idiom. This pins -compatibility with the historical DMD2 QAT artifacts independently from the current -quantization example's complete Diffusers training-bundle output. +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. """ diff --git a/tests/examples/diffusers/test_diffusers.py b/tests/examples/diffusers/test_diffusers.py index d6840f59c31..15c5eb44934 100644 --- a/tests/examples/diffusers/test_diffusers.py +++ b/tests/examples/diffusers/test_diffusers.py @@ -68,8 +68,8 @@ def quantize(self, tmp_path: Path) -> None: *self._format_args(), "--trt-high-precision-dtype", self.dtype, - "--output-bundle", - str(tmp_path / f"{self.name}_{self.format_type}_bundle"), + "--quantized-torch-ckpt-save-path", + str(tmp_path / f"{self.name}_{self.format_type}.pt"), "--onnx-dir", str(tmp_path / f"{self.name}_{self.format_type}_onnx"), ) @@ -81,7 +81,7 @@ def restore(self, tmp_path: Path) -> None: "--trt-high-precision-dtype", self.dtype, "--restore-from", - str(tmp_path / f"{self.name}_{self.format_type}_bundle"), + str(tmp_path / f"{self.name}_{self.format_type}.pt"), "--onnx-dir", str(tmp_path / f"{self.name}_{self.format_type}_onnx"), ) @@ -160,10 +160,10 @@ class Wan22Model(NamedTuple): quant_algo: str collect_method: str - def _bundle_path(self, tmp_path: Path) -> str: + def _ckpt_path(self, tmp_path: Path) -> str: stem = self.model.replace("wan2.2-t2v-", "") parts = [stem, *([self.backbone] if self.backbone else []), self.format_type] - return str(tmp_path / f"wan22_{'_'.join(parts)}_bundle") + return str(tmp_path / f"wan22_{'_'.join(parts)}.pt") def _common_args(self, tiny_wan22_path: str) -> list[str]: cmd_args = [ @@ -205,15 +205,15 @@ def quantize(self, tiny_wan22_path: str, tmp_path: Path) -> None: run_example_command( [ *self._common_args(tiny_wan22_path), - "--output-bundle", - self._bundle_path(tmp_path), + "--quantized-torch-ckpt-save-path", + self._ckpt_path(tmp_path), ], "diffusers/quantization", ) def restore(self, tiny_wan22_path: str, tmp_path: Path) -> None: run_example_command( - [*self._common_args(tiny_wan22_path), "--restore-from", self._bundle_path(tmp_path)], + [*self._common_args(tiny_wan22_path), "--restore-from", self._ckpt_path(tmp_path)], "diffusers/quantization", )