From 99a04445072973b27f011af14092aca463b53dc0 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 16:50:03 -0700 Subject: [PATCH 01/22] Add Qwen-Image registration to diffusers quantization example Register Qwen/Qwen-Image as a supported model in the diffusers quantization example: - ModelType.QWEN_IMAGE and lazy-imported QwenImagePipeline (so the example still imports on older diffusers). - MODEL_REGISTRY / MODEL_PIPELINE / MODEL_DEFAULTS entries (backbone="transformer", text-to-image calibration dataset). - An actionable ImportError when the installed diffusers lacks Qwen classes, instead of an opaque failure. - filter_func_qwen_image: quantize only transformer_blocks, keeping the first two and last two of the 60 blocks (and everything outside transformer_blocks) in original precision. Enables the plain FP8/NVFP4 export path for Qwen-Image. Core SVDQuant code is unchanged. (Qwen-Image SVDQuant checkpoint work, RLCR round 0 / M1.) Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- .../diffusers/quantization/models_utils.py | 20 ++++++++++++++++ .../quantization/pipeline_manager.py | 6 +++++ examples/diffusers/quantization/utils.py | 24 +++++++++++++++++++ 3 files changed, 50 insertions(+) diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index b59744282f6..ecacc38d407 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -30,11 +30,19 @@ from diffusers import Flux2Pipeline except ImportError: Flux2Pipeline = None + +# Qwen-Image classes were added in a recent diffusers release; import lazily so +# this example still imports on older diffusers versions. +try: + from diffusers import QwenImagePipeline +except ImportError: + QwenImagePipeline = None from utils import ( filter_func_default, filter_func_flux_dev, filter_func_ltx2_vae, filter_func_ltx_video, + filter_func_qwen_image, filter_func_wan_vae, filter_func_wan_video, ) @@ -54,6 +62,7 @@ class ModelType(str, Enum): LTX2 = "ltx-2" WAN22_T2V_14b = "wan2.2-t2v-14b" WAN22_T2V_5b = "wan2.2-t2v-5b" + QWEN_IMAGE = "qwen-image" _FILTER_FUNC_MAP: dict[ModelType, Callable[[str], bool]] = { @@ -63,6 +72,7 @@ class ModelType(str, Enum): ModelType.LTX2: filter_func_ltx_video, ModelType.WAN22_T2V_14b: filter_func_wan_video, ModelType.WAN22_T2V_5b: filter_func_wan_video, + ModelType.QWEN_IMAGE: filter_func_qwen_image, } _VAE_FILTER_FUNC_MAP: dict[tuple[ModelType, str], Callable[[str], bool]] = { @@ -95,6 +105,7 @@ def get_model_filter_func( ModelType.LTX2: "Lightricks/LTX-2", ModelType.WAN22_T2V_14b: "Wan-AI/Wan2.2-T2V-A14B-Diffusers", ModelType.WAN22_T2V_5b: "Wan-AI/Wan2.2-TI2V-5B-Diffusers", + ModelType.QWEN_IMAGE: "Qwen/Qwen-Image", } MODEL_PIPELINE: dict[ModelType, type[DiffusionPipeline] | None] = { @@ -109,6 +120,7 @@ def get_model_filter_func( ModelType.LTX2: None, ModelType.WAN22_T2V_14b: WanPipeline, ModelType.WAN22_T2V_5b: WanPipeline, + ModelType.QWEN_IMAGE: QwenImagePipeline, } # Shared dataset configurations @@ -226,6 +238,14 @@ def get_model_filter_func( ), }, }, + ModelType.QWEN_IMAGE: { + "backbone": "transformer", + "dataset": _SD_PROMPTS_DATASET, + "inference_extra_args": { + "height": 1024, + "width": 1024, + }, + }, } diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index 85e335ba787..f3878c24f49 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -99,6 +99,12 @@ def create_pipeline(self) -> Any: pipeline_cls = MODEL_PIPELINE[self.config.model_type] if pipeline_cls is None: + if self.config.model_type == ModelType.QWEN_IMAGE: + raise ImportError( + "Qwen-Image requires a diffusers version that provides " + "QwenImagePipeline. Please upgrade diffusers (e.g. " + "`pip install -U diffusers`) to a release that includes Qwen-Image." + ) raise ValueError( f"Model type {self.config.model_type.value} does not use diffusers pipelines." ) diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index d102e83e068..be3f6276db9 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -111,6 +111,30 @@ def filter_func_wan_video(name: str) -> bool: return pattern.match(name) is not None +# Qwen-Image's transformer has 60 ``transformer_blocks``. The recipe quantizes +# only those blocks while keeping the first two and last two -- and everything +# outside ``transformer_blocks`` -- in original precision. The model-agnostic, +# pre-calibration form of this recipe (deriving the block count from the model) +# lives in quantize.py; this name-only filter covers the plain FP8/NVFP4 path +# for the full 60-block Qwen-Image transformer. +QWEN_IMAGE_NUM_TRANSFORMER_BLOCKS = 60 +_QWEN_IMAGE_BLOCK_RE = re.compile(r"(?:^|\.)transformer_blocks\.(\d+)(?:\.|$)") + + +def filter_func_qwen_image(name: str) -> bool: + """Filter function specifically for Qwen-Image models. + + Returns ``True`` for modules to keep in original precision (quantization + disabled): everything outside ``transformer_blocks``, plus the first two and + last two transformer blocks. + """ + match = _QWEN_IMAGE_BLOCK_RE.search(name) + if match is None: + return True + block_idx = int(match.group(1)) + return block_idx < 2 or block_idx >= QWEN_IMAGE_NUM_TRANSFORMER_BLOCKS - 2 + + def load_calib_prompts( batch_size, calib_data_path: str | Path = "Gustavosta/Stable-Diffusion-Prompts", From e0c7910766cf9460fe618800dc7f00dbf066b143 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 17:16:44 -0700 Subject: [PATCH 02/22] Qwen-Image SVDQuant: block-range recipe, AWQ-style diffusers export, harness Implements the Qwen-Image NVFP4/FP8/SVDQuant diffusers quantization feature (RLCR round 0 / M2-M5), keeping core SVDQuant code unchanged: M2 (recipe): build_block_range_quant_cfg() emits ordered quant_cfg rules (disable-all -> enable *.transformer_blocks.* -> disable first/last-N), applied pre-calibration in Quantizer.get_quant_config so SVDQuant never mutates the excluded blocks. Driven by a MODEL_DEFAULTS["block_range"] entry for Qwen-Image (exclude first 2 / last 2; n derived from the model; n>=first+last+1 enforced). M3 (export): _export_diffusers_checkpoint now promotes quantizer-owned tensors to clean module-level safetensors keys before hide_quantizers_from_state_dict (diffusers path only; the transformers path keeps its postprocess_state_dict rename): input_quantizer._pre_quant_scale -> .pre_quant_scale (AWQ key), weight_quantizer.svdquant_lora_a/b -> .svdquant_lora_a/b. Adds an NVFP4_SVD branch to convert_hf_config (modeled on nvfp4_awq: pre_quant_scale + lora_rank), and process_layer_quant_config now flags SVDQuant with pre_quant_scale=True. This also resolves the diffusers pre_quant_scale TODO for AWQ-style exports. M4 (tests): unit tests for the block-range recipe (first/last-2 exclusion, n>=6 validation) and the NVFP4_SVD HF config conversion. M5 (harness): quantize.py --sanity-image-path (in-memory quantized-inference image, pre-export) + examples/diffusers/quantization/qwen_image_svdquant/ {run_qwen_image_quantization.sh, README.md} (parameterized container/model/ export flow for FP8/NVFP4/SVDQuant). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- .../diffusers/quantization/models_utils.py | 71 ++++++++++++ examples/diffusers/quantization/quantize.py | 59 +++++++++- .../qwen_image_svdquant/README.md | 108 ++++++++++++++++++ .../run_qwen_image_quantization.sh | 94 +++++++++++++++ modelopt/torch/export/convert_hf_config.py | 35 ++++++ modelopt/torch/export/quant_utils.py | 4 + modelopt/torch/export/unified_export_hf.py | 65 +++++++++++ .../diffusers/test_qwen_block_range_recipe.py | 91 +++++++++++++++ .../export/test_convert_hf_config_svdquant.py | 85 ++++++++++++++ 9 files changed, 611 insertions(+), 1 deletion(-) create mode 100644 examples/diffusers/quantization/qwen_image_svdquant/README.md create mode 100755 examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh create mode 100644 tests/examples/diffusers/test_qwen_block_range_recipe.py create mode 100644 tests/unit/torch/export/test_convert_hf_config_svdquant.py diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index ecacc38d407..05f59b232c9 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -18,6 +18,7 @@ from enum import Enum from typing import Any +import torch from diffusers import ( DiffusionPipeline, FluxPipeline, @@ -245,6 +246,15 @@ def get_model_filter_func( "height": 1024, "width": 1024, }, + # Quantize only ``transformer_blocks``; keep the first 2 and last 2 blocks + # (and everything outside ``transformer_blocks``) in original precision. + # Applied pre-calibration via ``build_block_range_quant_cfg`` so SVDQuant + # never mutates the excluded blocks' weights. + "block_range": { + "exclude_first_n": 2, + "exclude_last_n": 2, + "block_module": "transformer_blocks", + }, }, } @@ -292,3 +302,64 @@ def parse_extra_params( i += 1 return extra_params + + +def build_block_range_quant_cfg( + backbone: torch.nn.Module, + exclude_first_n: int, + exclude_last_n: int, + block_module: str = "transformer_blocks", +) -> list[dict[str, Any]]: + """Build ordered ``quant_cfg`` rules for a transformer-block-only recipe. + + The rules quantize only the linears under ``block_module`` while keeping the + first ``exclude_first_n`` and last ``exclude_last_n`` blocks -- and everything + outside ``block_module`` -- in original precision. + + The rules are meant to be appended to the ``quant_cfg`` list consumed by + ``mtq.quantize`` so the selection is applied BEFORE calibration. This is + required for SVDQuant, whose calibration subtracts a low-rank residual from + the weights of every *enabled* linear: disabling the excluded blocks only + after calibration would leave their weights mutated instead of bit-identical + to the original precision. + + Rules are applied in order with later rules overriding earlier ones: + 1. disable every linear weight/input quantizer, + 2. re-enable only those under ``block_module``, + 3. disable the first/last ``n`` blocks. + + Raises: + ValueError: if the backbone has no ``block_module`` list, or it has fewer + than ``exclude_first_n + exclude_last_n + 1`` blocks. + """ + blocks = getattr(backbone, block_module, None) + if blocks is None or not hasattr(blocks, "__len__"): + raise ValueError( + f"Backbone {type(backbone).__name__} has no '{block_module}' module list; " + "cannot build the transformer-block-range recipe." + ) + num_blocks = len(blocks) + min_blocks = exclude_first_n + exclude_last_n + 1 + if num_blocks < min_blocks: + raise ValueError( + f"'{block_module}' has only {num_blocks} block(s); excluding the first " + f"{exclude_first_n} and last {exclude_last_n} requires at least {min_blocks} blocks." + ) + + excluded = sorted( + set(range(exclude_first_n)) | set(range(num_blocks - exclude_last_n, num_blocks)) + ) + rules: list[dict[str, Any]] = [ + {"quantizer_name": "*weight_quantizer", "cfg": {"enable": False}}, + {"quantizer_name": "*input_quantizer", "cfg": {"enable": False}}, + {"quantizer_name": f"*{block_module}.*weight_quantizer", "cfg": {"enable": True}}, + {"quantizer_name": f"*{block_module}.*input_quantizer", "cfg": {"enable": True}}, + ] + for idx in excluded: + rules.append( + {"quantizer_name": f"*{block_module}.{idx}.*weight_quantizer", "cfg": {"enable": False}} + ) + rules.append( + {"quantizer_name": f"*{block_module}.{idx}.*input_quantizer", "cfg": {"enable": False}} + ) + return rules diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 299a101172a..886eb775eb3 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -32,7 +32,13 @@ set_quant_config_attr, ) from diffusers import DiffusionPipeline -from models_utils import MODEL_DEFAULTS, ModelType, get_model_filter_func, parse_extra_params +from models_utils import ( + MODEL_DEFAULTS, + ModelType, + build_block_range_quant_cfg, + get_model_filter_func, + parse_extra_params, +) from onnx_utils.export import generate_fp8_scales, modelopt_export_sd from pipeline_manager import PipelineManager from quantize_config import ( @@ -163,6 +169,28 @@ def get_quant_config(self, n_steps: int, backbone: torch.nn.Module) -> Any: } ) + # Apply the transformer-block-range recipe (e.g. Qwen-Image) BEFORE + # calibration. This restricts quantization to `transformer_blocks` and + # excludes the first/last N blocks. It must run pre-calibration so that + # SVDQuant does not mutate the weights of the excluded blocks. The recipe + # is format-agnostic (applies to FP8/NVFP4/SVDQuant alike). + block_range = MODEL_DEFAULTS.get(self.model_config.model_type, {}).get("block_range") + if block_range is not None: + recipe_rules = build_block_range_quant_cfg( + backbone, + exclude_first_n=block_range.get("exclude_first_n", 2), + exclude_last_n=block_range.get("exclude_last_n", 2), + block_module=block_range.get("block_module", "transformer_blocks"), + ) + self.logger.info( + f"Applying block-range recipe ({len(recipe_rules)} rules) for " + f"{self.model_config.model_type.value}: quantize only " + f"'{block_range.get('block_module', 'transformer_blocks')}' excluding " + f"first {block_range.get('exclude_first_n', 2)} / last " + f"{block_range.get('exclude_last_n', 2)} blocks." + ) + quant_cfg_list.extend(recipe_rules) + quant_config = {**base_cfg, "quant_cfg": quant_cfg_list} set_quant_config_attr( quant_config, @@ -542,6 +570,15 @@ def create_argument_parser() -> argparse.ArgumentParser: export_group.add_argument( "--restore-from", type=str, help="Path to restore from previous checkpoint" ) + export_group.add_argument( + "--sanity-image-path", + type=str, + default=None, + help="If set, generate one image from the in-memory quantized pipeline (after " + "quantization, before the weights are packed for export) and save it here. This is " + "a quick functional sanity check of quantized inference; it does NOT reload the " + "exported checkpoint.", + ) export_group.add_argument( "--trt-high-precision-dtype", type=str, @@ -681,6 +718,26 @@ def forward_loop(mod): pipeline_manager.print_quant_summary() + # Optional functional sanity check: generate one image from the in-memory + # quantized pipeline. This runs BEFORE export (while weights are still + # fake-quantized and runnable, not yet packed) and does not reload the + # exported checkpoint. + if args.sanity_image_path: + try: + logger.info(f"Generating sanity image to {args.sanity_image_path}") + inference_args = MODEL_DEFAULTS.get(model_type, {}).get("inference_extra_args", {}) + result = pipe( + prompt="A high-quality photo of a cat wearing sunglasses", + num_inference_steps=calib_config.n_steps, + **inference_args, + ) + sanity_path = Path(args.sanity_image_path) + sanity_path.parent.mkdir(parents=True, exist_ok=True) + result.images[0].save(str(sanity_path)) + logger.info("Sanity image saved successfully") + except Exception as sanity_error: # noqa: BLE001 + logger.warning(f"Sanity image generation failed (non-fatal): {sanity_error}") + for backbone_name, backbone in pipeline_manager.iter_backbones(): export_manager.export_onnx( pipe, diff --git a/examples/diffusers/quantization/qwen_image_svdquant/README.md b/examples/diffusers/quantization/qwen_image_svdquant/README.md new file mode 100644 index 00000000000..02daf02ac8c --- /dev/null +++ b/examples/diffusers/quantization/qwen_image_svdquant/README.md @@ -0,0 +1,108 @@ +# Qwen-Image Quantization (FP8 / NVFP4 / NVFP4-SVDQuant) + +A reproducible harness for quantizing [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) +with the diffusers quantization example and exporting HuggingFace checkpoints. + +## What it does + +- Registers Qwen-Image in the diffusers quantization example (`--model qwen-image`). +- **Recipe**: quantizes only the linears under `transformer_blocks`, keeping the + **first 2 and last 2** of the 60 blocks (and everything outside + `transformer_blocks`: text encoder, VAE, embedders, norms, `proj_out`, …) in + original precision. The exclusion is applied **before calibration** so that for + SVDQuant the excluded blocks' weights stay bit-identical to the original. +- Produces three checkpoints: **FP8**, **NVFP4** (max), and **NVFP4 + SVDQuant**. +- Exports a HuggingFace unified checkpoint per component (safetensors + `config.json`). + +### SVDQuant checkpoint format (AWQ-aligned) + +For the SVDQuant export, the quantizer-owned tensors are promoted to clean, +module-level safetensors keys (mirroring how AWQ exports `pre_quant_scale`): + +| Tensor | Safetensors key | +|--------|-----------------| +| AWQ smoothing scale (`input_quantizer._pre_quant_scale`) | `.pre_quant_scale` | +| Low-rank factor A (`weight_quantizer.svdquant_lora_a`) | `.svdquant_lora_a` | +| Low-rank factor B (`weight_quantizer.svdquant_lora_b`) | `.svdquant_lora_b` | + +They are embedded in the component's main safetensors (no sidecar). The +`config.json`'s `quantization_config` follows the `nvfp4_awq` shape with +`"pre_quant_scale": true` plus the SVDQuant `lora_rank`, so a consumer can +reconstruct `y = NVFP4_GEMM(x) + (x @ lora_a^T) @ lora_b^T`. (No in-repo runtime +applies this residual yet; the checkpoint is a documented on-disk artifact.) + +## Layout (kernel-dev defaults) + +| Env var | Default | Purpose | +|---------|---------|---------| +| `KERNEL_DEV_ROOT` | `/lustre/fsw/coreai_dlalgo_modelopt/users/jingyux/kernel-dev` | Root for container/models/output | +| `MODEL_DIR` | `${KERNEL_DEV_ROOT}/models/Qwen-Image` | Local model cache | +| `OUTPUT_DIR` | `${KERNEL_DEV_ROOT}/qwen_image_ckpts` | Exported checkpoints | +| `HF_TOKEN_FILE` | `${KERNEL_DEV_ROOT}/HF_TOKEN.txt` | Hugging Face token file | +| `FORMATS` | `fp8 nvfp4 svdquant` | Formats to run | +| `CALIB_SIZE` / `BATCH_SIZE` / `N_STEPS` / `LOWRANK` | `64 / 2 / 20 / 32` | Calibration knobs | + +## 1. Build the container (once) + +The diffusers example needs a recent `diffusers` (with `QwenImagePipeline`) and +modelopt installed from source. From a base NGC PyTorch image: + +```bash +CONTAINER_DIR=/lustre/fsw/coreai_dlalgo_modelopt/users/jingyux/kernel-dev/container +mkdir -p "${CONTAINER_DIR}" + +# Import a base image to an enroot squashfs (adjust the tag as needed). +enroot import -o "${CONTAINER_DIR}/modelopt-diffusers.sqsh" \ + docker://nvcr.io#nvidia/pytorch:25.04-py3 + +# Install modelopt (from source) + example deps into the container, then re-save. +srun --container-image="${CONTAINER_DIR}/modelopt-diffusers.sqsh" \ + --container-mounts=/lustre:/lustre --container-save="${CONTAINER_DIR}/modelopt-diffusers.sqsh" \ + bash -lc ' + cd /lustre/fsw/coreai_dlalgo_modelopt/users/jingyux/kernel-dev/source/Model-Optimizer && + pip install -e ".[dev]" && + pip install -U "diffusers>=0.35" "transformers>=4.52" accelerate datasets && + python -c "from diffusers import QwenImagePipeline; print(\"QwenImagePipeline OK\")" + ' +``` + +## 2. Run quantization + +Inside the container (or via `srun`), run the harness: + +```bash +srun --gpus=1 \ + --container-image=/lustre/fsw/coreai_dlalgo_modelopt/users/jingyux/kernel-dev/container/modelopt-diffusers.sqsh \ + --container-mounts=/lustre:/lustre \ + bash examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh +``` + +This downloads `Qwen/Qwen-Image` to `MODEL_DIR` (idempotent), then for each +format writes `${OUTPUT_DIR}/qwen-image-/` (HF checkpoint + `sanity.png`). + +Run a single format, or preview the commands without executing: + +```bash +FORMATS=svdquant LOWRANK=32 bash .../run_qwen_image_quantization.sh +DRY_RUN=1 bash .../run_qwen_image_quantization.sh # print planned commands only +``` + +The equivalent direct `quantize.py` invocation for SVDQuant: + +```bash +python examples/diffusers/quantization/quantize.py \ + --model qwen-image --override-model-path "${MODEL_DIR}" --model-dtype BFloat16 \ + --format fp4 --quant-algo svdquant --lowrank 32 \ + --calib-size 64 --batch-size 2 --n-steps 20 \ + --hf-ckpt-dir "${OUTPUT_DIR}/qwen-image-svdquant" \ + --sanity-image-path "${OUTPUT_DIR}/qwen-image-svdquant/sanity.png" +``` + +## Notes + +- `Qwen/Qwen-Image` loads without `trust_remote_code`. +- The transformer is ~20B params; calibration needs a GPU with enough memory + (use `--cpu-offloading` if VRAM-limited). +- The `--sanity-image-path` image is generated from the **in-memory** quantized + pipeline before the weights are packed for export (a functional check of + quantized inference; it does not reload the exported checkpoint). diff --git a/examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh b/examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh new file mode 100755 index 00000000000..1ba6e440e69 --- /dev/null +++ b/examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh @@ -0,0 +1,94 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Reproducible Qwen-Image quantization (FP8 / NVFP4 / NVFP4-SVDQuant) using the +# diffusers quantization example. This script is meant to run INSIDE a container +# that already has NVIDIA Model Optimizer installed from source and a +# Qwen-capable diffusers (see README.md for building the container and the +# Slurm/srun wrapper). +# +# It downloads Qwen/Qwen-Image (idempotently), then for each requested format +# runs `quantize.py` to calibrate the transformer (only `transformer_blocks`, +# excluding the first 2 / last 2 blocks), generate a quantized-inference sanity +# image, and export a HuggingFace checkpoint. +# +# All paths are parameterized via environment variables; the defaults match the +# kernel-dev experiment layout described in README.md. +set -euo pipefail + +# --- Configuration (override via environment) -------------------------------- +KERNEL_DEV_ROOT="${KERNEL_DEV_ROOT:-/lustre/fsw/coreai_dlalgo_modelopt/users/jingyux/kernel-dev}" +MODEL_ID="${MODEL_ID:-Qwen/Qwen-Image}" +MODEL_DIR="${MODEL_DIR:-${KERNEL_DEV_ROOT}/models/Qwen-Image}" +OUTPUT_DIR="${OUTPUT_DIR:-${KERNEL_DEV_ROOT}/qwen_image_ckpts}" +HF_TOKEN_FILE="${HF_TOKEN_FILE:-${KERNEL_DEV_ROOT}/HF_TOKEN.txt}" +# Path to the diffusers quantization example (this script lives one level below it). +QUANT_DIR="${QUANT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}" + +# Formats to run: any of {fp8, nvfp4, svdquant}. +FORMATS="${FORMATS:-fp8 nvfp4 svdquant}" + +# Calibration knobs (small defaults for a quick run; raise CALIB_SIZE for quality). +CALIB_SIZE="${CALIB_SIZE:-64}" +BATCH_SIZE="${BATCH_SIZE:-2}" +N_STEPS="${N_STEPS:-20}" +LOWRANK="${LOWRANK:-32}" +MODEL_DTYPE="${MODEL_DTYPE:-BFloat16}" + +# Set DRY_RUN=1 to print the planned commands without executing them. +DRY_RUN="${DRY_RUN:-0}" + +log() { echo "[qwen-image-quant] $*"; } +run() { + log "+ $*" + if [[ "${DRY_RUN}" != "1" ]]; then + "$@" + fi +} + +# --- Hugging Face token ------------------------------------------------------ +if [[ ! -r "${HF_TOKEN_FILE}" ]]; then + echo "ERROR: HF token file not found or not readable: ${HF_TOKEN_FILE}" >&2 + echo " Set HF_TOKEN_FILE to a readable file containing your Hugging Face token." >&2 + exit 1 +fi +HF_TOKEN="$(tr -d '[:space:]' < "${HF_TOKEN_FILE}")" +if [[ -z "${HF_TOKEN}" ]]; then + echo "ERROR: HF token file is empty: ${HF_TOKEN_FILE}" >&2 + exit 1 +fi +export HF_TOKEN +export HUGGING_FACE_HUB_TOKEN="${HF_TOKEN}" + +# --- Download the model (idempotent) ---------------------------------------- +log "Downloading ${MODEL_ID} -> ${MODEL_DIR} (skipped if already present)" +run mkdir -p "${MODEL_DIR}" +run huggingface-cli download "${MODEL_ID}" --local-dir "${MODEL_DIR}" --exclude "*.onnx" + +# --- Quantize + export for each format -------------------------------------- +mkdir -p "${OUTPUT_DIR}" +for fmt in ${FORMATS}; do + case "${fmt}" in + fp8) quant_args=(--format fp8 --quant-algo max) ;; + nvfp4) quant_args=(--format fp4 --quant-algo max) ;; + svdquant) quant_args=(--format fp4 --quant-algo svdquant --lowrank "${LOWRANK}") ;; + *) echo "ERROR: unknown format '${fmt}' (expected fp8|nvfp4|svdquant)" >&2; exit 1 ;; + esac + + out="${OUTPUT_DIR}/qwen-image-${fmt}" + log "=== Quantizing Qwen-Image (${fmt}) -> ${out} ===" + run python "${QUANT_DIR}/quantize.py" \ + --model qwen-image \ + --override-model-path "${MODEL_DIR}" \ + --model-dtype "${MODEL_DTYPE}" \ + "${quant_args[@]}" \ + --calib-size "${CALIB_SIZE}" \ + --batch-size "${BATCH_SIZE}" \ + --n-steps "${N_STEPS}" \ + --hf-ckpt-dir "${out}" \ + --sanity-image-path "${out}/sanity.png" + log "Done: ${fmt}. Checkpoint at ${out}, sanity image at ${out}/sanity.png" +done + +log "All requested formats complete. Checkpoints under ${OUTPUT_DIR}" diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index 06e5923a30f..70016a54bca 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -62,6 +62,18 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) return { "weights": {"dynamic": False, "num_bits": 4, "type": "float", "group_size": gs}, } + elif quant_algo == "NVFP4_SVD": + gs = group_size or 16 + return { + "input_activations": { + "dynamic": False, + "num_bits": 4, + "type": "float", + "group_size": gs, + }, + "weights": {"dynamic": False, "num_bits": 4, "type": "float", "group_size": gs}, + "pre_quant_scale": True, + } elif quant_algo in ("NVFP4_AWQ", "W4A8_AWQ"): gs = group_size or 128 return { @@ -196,6 +208,29 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An "targets": ["Linear"], } new_config["config_groups"] = {"group_0": config_group_details} + elif quant_algo_value == "NVFP4_SVD": + # NVFP4 + SVDQuant: NVFP4 weights/activations plus an AWQ-style + # pre_quant_scale and a low-rank residual (svdquant_lora_a/b) stored as + # .pre_quant_scale / .svdquant_lora_{a,b} in the + # safetensors. The config mirrors NVFP4 with a pre_quant_scale flag and + # the LoRA rank so consumers can reconstruct + # ``y = NVFP4_GEMM(x) + (x @ lora_a^T) @ lora_b^T``. + group_size = original_quantization_details.get("group_size", 16) + config_group_details = { + "input_activations": { + "dynamic": False, + "num_bits": 4, + "type": "float", + "group_size": group_size, + }, + "weights": {"dynamic": False, "num_bits": 4, "type": "float", "group_size": group_size}, + "pre_quant_scale": True, + "targets": ["Linear"], + } + lora_rank = original_quantization_details.get("lora_rank") + if lora_rank is not None: + config_group_details["lora_rank"] = lora_rank + new_config["config_groups"] = {"group_0": config_group_details} elif quant_algo_value == "MIXED_PRECISION": quantized_layers = original_quantization_details.get("quantized_layers", {}) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index d8ddf442924..e35ead6253b 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -749,9 +749,13 @@ def process_layer_quant_config(layer_config_dict): "group_size": block_size_value, } elif v == "nvfp4_svdquant": + # SVDQuant builds on the AWQ-style pre_quant_scale smoothing, so its + # config mirrors nvfp4_awq (group_size + pre_quant_scale flag). layer_config = { "quant_algo": "NVFP4_SVD", "group_size": block_size_value, + "has_zero_point": False, + "pre_quant_scale": True, } elif v == "mxfp8": layer_config = { diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index ef5757aa0cb..0a02f751855 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1013,6 +1013,58 @@ def _fuse_qkv_linears_diffusion( ) +def _detect_svdquant_rank(component: nn.Module) -> int | None: + """Return the SVDQuant low-rank dimension from the first SVDQuant linear, if any. + + ``svdquant_lora_a`` has shape ``(rank, in_features)``, so its first dimension + is the low-rank size. + """ + for _, sub_module in component.named_modules(): + weight_quantizer = getattr(sub_module, "weight_quantizer", None) + lora_a = getattr(weight_quantizer, "svdquant_lora_a", None) + if lora_a is not None: + return int(lora_a.shape[0]) + return None + + +def _promote_quantizer_tensors_to_module(component: nn.Module) -> None: + """Promote quantizer-owned export tensors onto their parent linear module. + + The diffusers export path saves via ``save_pretrained`` inside + :func:`hide_quantizers_from_state_dict` (which deletes the ``weight_quantizer`` + / ``input_quantizer`` submodules) and -- unlike the transformers path -- does + NOT run :func:`postprocess_state_dict`. Without this step the AWQ smoothing + scale and the SVDQuant low-rank factors would be dropped from the exported + checkpoint. We register them as module buffers under clean, AWQ-aligned keys + so they are embedded in the component's main safetensors: + + - ``input_quantizer._pre_quant_scale`` -> ``.pre_quant_scale`` + (the same key the transformers/AWQ path produces via postprocess_state_dict) + - ``weight_quantizer.svdquant_lora_a`` -> ``.svdquant_lora_a`` + - ``weight_quantizer.svdquant_lora_b`` -> ``.svdquant_lora_b`` + + This runs after :func:`_process_quantized_modules` (which leaves these + quantizer buffers in place) and before ``save_pretrained``. + """ + for _, sub_module in component.named_modules(): + if not is_quantlinear(sub_module): + continue + + input_quantizer = getattr(sub_module, "input_quantizer", None) + pre_quant_scale = getattr(input_quantizer, "_pre_quant_scale", None) + if pre_quant_scale is not None and not hasattr(sub_module, "pre_quant_scale"): + sub_module.register_buffer("pre_quant_scale", pre_quant_scale.detach().clone()) + + weight_quantizer = getattr(sub_module, "weight_quantizer", None) + lora_a = getattr(weight_quantizer, "svdquant_lora_a", None) + lora_b = getattr(weight_quantizer, "svdquant_lora_b", None) + if lora_a is not None and lora_b is not None: + if not hasattr(sub_module, "svdquant_lora_a"): + sub_module.register_buffer("svdquant_lora_a", lora_a.detach().clone()) + if not hasattr(sub_module, "svdquant_lora_b"): + sub_module.register_buffer("svdquant_lora_b", lora_b.detach().clone()) + + def _export_diffusers_checkpoint( pipe: Any, dtype: torch.dtype | None, @@ -1091,8 +1143,21 @@ def _export_diffusers_checkpoint( # Step 4: Process quantized modules (convert weights, register scales) _process_quantized_modules(component, component_dtype, is_modelopt_qlora=False) + # Step 4.5: Promote quantizer-owned tensors (AWQ pre_quant_scale and + # SVDQuant LoRA factors) onto the module so they survive + # hide_quantizers_from_state_dict and are embedded in the component's + # main safetensors under clean, AWQ-aligned keys. + _promote_quantizer_tensors_to_module(component) + # Step 5: Build quantization config quant_config = get_quant_config(component, is_modelopt_qlora=False) + if quant_config: + quantization_details = quant_config.get("quantization", {}) + # Record the SVDQuant low-rank size so consumers know the LoRA shape. + if quantization_details.get("quant_algo") == "NVFP4_SVD": + svdquant_rank = _detect_svdquant_rank(component) + if svdquant_rank is not None: + quantization_details["lora_rank"] = svdquant_rank hf_quant_config = convert_hf_quant_config_format(quant_config) if quant_config else None # Step 6: Save the component diff --git a/tests/examples/diffusers/test_qwen_block_range_recipe.py b/tests/examples/diffusers/test_qwen_block_range_recipe.py new file mode 100644 index 00000000000..906dd1fb60e --- /dev/null +++ b/tests/examples/diffusers/test_qwen_block_range_recipe.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Unit tests for the transformer-block-range quantization recipe (e.g. Qwen-Image). + +The recipe must quantize only the linears under ``transformer_blocks`` while +excluding the first/last N blocks, and it must be expressible as pre-calibration +``quant_cfg`` rules (so SVDQuant never mutates the excluded blocks' weights). +""" + +import re +import sys +from pathlib import Path + +import pytest + +# Importing the example module pulls in diffusers/torch/datasets/modelopt. +pytest.importorskip("diffusers") +pytest.importorskip("torch") + +# Make the diffusers quantization example importable. +_EXAMPLE_DIR = Path(__file__).parents[3] / "examples" / "diffusers" / "quantization" +if str(_EXAMPLE_DIR) not in sys.path: + sys.path.insert(0, str(_EXAMPLE_DIR)) + +from models_utils import build_block_range_quant_cfg # noqa: E402 + +_BLOCK_RULE_RE = re.compile(r"\*transformer_blocks\.(\d+)\.\*(?:weight|input)_quantizer") + + +class _StubBackbone: + """Minimal stand-in exposing a ``transformer_blocks`` sequence of length n.""" + + def __init__(self, num_blocks: int): + self.transformer_blocks = list(range(num_blocks)) + + +def _disabled_block_indices(rules): + """Indices of transformer blocks explicitly disabled by per-block rules.""" + indices = set() + for rule in rules: + if rule["cfg"].get("enable") is False: + match = _BLOCK_RULE_RE.fullmatch(rule["quantizer_name"]) + if match: + indices.add(int(match.group(1))) + return indices + + +def test_recipe_excludes_first_and_last_two_blocks(): + rules = build_block_range_quant_cfg(_StubBackbone(6), exclude_first_n=2, exclude_last_n=2) + + # 1. disable-all rules come first (weight + input). + assert rules[0] == {"quantizer_name": "*weight_quantizer", "cfg": {"enable": False}} + assert rules[1] == {"quantizer_name": "*input_quantizer", "cfg": {"enable": False}} + # 2. then enable only the transformer_blocks. + assert {"quantizer_name": "*transformer_blocks.*weight_quantizer", "cfg": {"enable": True}} in rules + assert {"quantizer_name": "*transformer_blocks.*input_quantizer", "cfg": {"enable": True}} in rules + # 3. then disable the first 2 and last 2 of the 6 blocks -> {0, 1, 4, 5}; quantize {2, 3}. + assert _disabled_block_indices(rules) == {0, 1, 4, 5} + + +def test_recipe_block_count_scales_with_model(): + # For a 60-block model (Qwen-Image), exclude {0, 1, 58, 59}; quantize 2..57. + rules = build_block_range_quant_cfg(_StubBackbone(60), exclude_first_n=2, exclude_last_n=2) + assert _disabled_block_indices(rules) == {0, 1, 58, 59} + + +def test_recipe_rejects_too_few_blocks(): + # 2 + 2 exclusion needs at least 5 blocks; 4 blocks must raise a clear error. + with pytest.raises(ValueError, match="at least"): + build_block_range_quant_cfg(_StubBackbone(4), exclude_first_n=2, exclude_last_n=2) + + +def test_recipe_missing_block_module_raises(): + class _NoBlocks: + pass + + with pytest.raises(ValueError, match="transformer_blocks"): + build_block_range_quant_cfg(_NoBlocks(), exclude_first_n=2, exclude_last_n=2) diff --git a/tests/unit/torch/export/test_convert_hf_config_svdquant.py b/tests/unit/torch/export/test_convert_hf_config_svdquant.py new file mode 100644 index 00000000000..d00365703b1 --- /dev/null +++ b/tests/unit/torch/export/test_convert_hf_config_svdquant.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Unit tests for the NVFP4_SVD (SVDQuant) HF quantization-config conversion.""" + +from modelopt.torch.export.convert_hf_config import ( + _quant_algo_to_group_config, + convert_hf_quant_config_format, +) + + +def test_nvfp4_svd_group_config_mirrors_awq_with_pre_quant_scale(): + """The NVFP4_SVD config group is NVFP4 weights/activations + a pre_quant_scale flag.""" + group = _quant_algo_to_group_config("NVFP4_SVD", group_size=16) + assert group["pre_quant_scale"] is True + assert group["weights"] == { + "dynamic": False, + "num_bits": 4, + "type": "float", + "group_size": 16, + } + assert group["input_activations"]["num_bits"] == 4 + assert group["input_activations"]["type"] == "float" + assert group["input_activations"]["group_size"] == 16 + + +def test_convert_hf_quant_config_format_nvfp4_svd(): + """A full NVFP4_SVD quantization dict converts to a complete compressed-tensors config.""" + input_config = { + "producer": {"name": "modelopt", "version": "0.0.0"}, + "quantization": { + "quant_algo": "NVFP4_SVD", + "group_size": 16, + "has_zero_point": False, + "pre_quant_scale": True, + "lora_rank": 32, + "exclude_modules": ["transformer_blocks.0.*", "proj_out"], + "kv_cache_quant_algo": None, + }, + } + + out = convert_hf_quant_config_format(input_config) + + # A real config group is emitted (not a bare {"quant_algo": ...} fallback). + assert "config_groups" in out + group = out["config_groups"]["group_0"] + assert group["pre_quant_scale"] is True + assert group["lora_rank"] == 32 + assert group["weights"]["num_bits"] == 4 + assert group["weights"]["type"] == "float" + assert group["weights"]["group_size"] == 16 + assert group["input_activations"]["num_bits"] == 4 + assert group["targets"] == ["Linear"] + + # Top-level metadata is preserved. + assert out["quant_algo"] == "NVFP4_SVD" + assert out["ignore"] == ["transformer_blocks.0.*", "proj_out"] + assert out["quant_method"] == "modelopt" + + +def test_convert_hf_quant_config_format_nvfp4_svd_without_rank(): + """lora_rank is optional; omitting it must not break the conversion.""" + input_config = { + "quantization": { + "quant_algo": "NVFP4_SVD", + "group_size": 16, + "pre_quant_scale": True, + }, + } + out = convert_hf_quant_config_format(input_config) + group = out["config_groups"]["group_0"] + assert "lora_rank" not in group + assert group["pre_quant_scale"] is True From 6718e722ca58897a169157ffc9a963afc524704a Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 17:42:27 -0700 Subject: [PATCH 03/22] Qwen-Image SVDQuant: fix export/recipe defects, add Qwen QKV fusion + tests Addresses the round-0 Codex review (RLCR round 1): Blocking fixes: - convert_hf_config: NVFP4_SVD config groups now keep `has_zero_point: False` (both convert_hf_quant_config_format and _quant_algo_to_group_config); asserted in the unit test. - build_block_range_quant_cfg: minimum is now first+last+2 (>=2 quantized middle blocks; n>=6 for the 2+2 Qwen recipe); recipe test rejects 5/4/3-block models. - quantize.py --sanity-image-path failures are now fatal (re-raise -> non-zero exit) so the harness cannot report success without the image; the harness also verifies sanity.png + safetensors + config.json exist per format. Qwen export enablement: - diffusers_utils.generate_diffusion_dummy_inputs: add a QwenImageTransformer2DModel branch (packed latents [B,(H//2)(W//2),C], encoder_hidden_states_mask, img_shapes, txt_seq_lens, optional guidance, continuous timestep). - unified_export_hf._fuse_qkv_linears_diffusion gains strict=; Qwen QKV fusion now fails hard instead of silently skipping. Promotion buffers now overwrite on re-export. create_pipeline_from gives the same actionable Qwen import error. Tests: - New tests/unit/torch/quantization/test_svdquant_forward_fold.py: LoRA stays on weight_quantizer, forward includes a nonzero residual, fold_weight folds it and drops the buffers (existing test_svdquant_lora_weights left unmodified). Deferred to Round 2 / cluster: tiny Qwen2_5_VL fixture + full diffusers e2e export test (needs a Qwen-capable diffusers + GPU); the actual AC-7 checkpoint run. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- .../diffusers/quantization/models_utils.py | 8 +- .../quantization/pipeline_manager.py | 6 + examples/diffusers/quantization/quantize.py | 7 +- .../run_qwen_image_quantization.sh | 7 ++ modelopt/torch/export/convert_hf_config.py | 2 + modelopt/torch/export/diffusers_utils.py | 40 ++++++ modelopt/torch/export/unified_export_hf.py | 25 ++-- .../diffusers/test_qwen_block_range_recipe.py | 10 +- .../export/test_convert_hf_config_svdquant.py | 2 + .../test_svdquant_forward_fold.py | 116 ++++++++++++++++++ 10 files changed, 209 insertions(+), 14 deletions(-) create mode 100644 tests/unit/torch/quantization/test_svdquant_forward_fold.py diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index 05f59b232c9..e7faf7589c1 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -339,11 +339,15 @@ def build_block_range_quant_cfg( "cannot build the transformer-block-range recipe." ) num_blocks = len(blocks) - min_blocks = exclude_first_n + exclude_last_n + 1 + # Require at least two quantized middle blocks so the recipe actually + # quantizes something (excluding first/last alone could otherwise leave 0-1 + # quantized blocks). For the default 2+2 recipe this means n >= 6. + min_blocks = exclude_first_n + exclude_last_n + 2 if num_blocks < min_blocks: raise ValueError( f"'{block_module}' has only {num_blocks} block(s); excluding the first " - f"{exclude_first_n} and last {exclude_last_n} requires at least {min_blocks} blocks." + f"{exclude_first_n} and last {exclude_last_n} requires at least {min_blocks} blocks " + f"(at least 2 quantized middle blocks)." ) excluded = sorted( diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index f3878c24f49..5496f131ba2 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -61,6 +61,12 @@ def create_pipeline_from( """ pipeline_cls = MODEL_PIPELINE[model_type] if pipeline_cls is None: + if model_type == ModelType.QWEN_IMAGE: + raise ImportError( + "Qwen-Image requires a diffusers version that provides " + "QwenImagePipeline. Please upgrade diffusers (e.g. " + "`pip install -U diffusers`) to a release that includes Qwen-Image." + ) raise ValueError(f"Model type {model_type.value} does not use diffusers pipelines.") model_id = ( MODEL_REGISTRY[model_type] if override_model_path is None else override_model_path diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 886eb775eb3..57b9db589db 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -735,8 +735,11 @@ def forward_loop(mod): sanity_path.parent.mkdir(parents=True, exist_ok=True) result.images[0].save(str(sanity_path)) logger.info("Sanity image saved successfully") - except Exception as sanity_error: # noqa: BLE001 - logger.warning(f"Sanity image generation failed (non-fatal): {sanity_error}") + except Exception as sanity_error: + # A requested sanity image is a positive success criterion: if it + # cannot be produced, fail loudly rather than reporting success. + logger.error(f"Sanity image generation failed: {sanity_error}", exc_info=True) + raise for backbone_name, backbone in pipeline_manager.iter_backbones(): export_manager.export_onnx( diff --git a/examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh b/examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh index 1ba6e440e69..5e0569c0e40 100755 --- a/examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh +++ b/examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh @@ -88,6 +88,13 @@ for fmt in ${FORMATS}; do --n-steps "${N_STEPS}" \ --hf-ckpt-dir "${out}" \ --sanity-image-path "${out}/sanity.png" + + # Verify the expected artifacts were produced (a missing artifact is a failure). + if [[ "${DRY_RUN}" != "1" ]]; then + [[ -f "${out}/sanity.png" ]] || { echo "ERROR: missing sanity image ${out}/sanity.png" >&2; exit 1; } + find "${out}" -name '*.safetensors' | grep -q . || { echo "ERROR: no safetensors under ${out}" >&2; exit 1; } + find "${out}" -name 'config.json' | grep -q . || { echo "ERROR: no config.json under ${out}" >&2; exit 1; } + fi log "Done: ${fmt}. Checkpoint at ${out}, sanity image at ${out}/sanity.png" done diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index 70016a54bca..6f7dedb97c8 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -72,6 +72,7 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None) "group_size": gs, }, "weights": {"dynamic": False, "num_bits": 4, "type": "float", "group_size": gs}, + "has_zero_point": False, "pre_quant_scale": True, } elif quant_algo in ("NVFP4_AWQ", "W4A8_AWQ"): @@ -224,6 +225,7 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An "group_size": group_size, }, "weights": {"dynamic": False, "num_bits": 4, "type": "float", "group_size": group_size}, + "has_zero_point": False, "pre_quant_scale": True, "targets": ["Linear"], } diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index 9620c97c10e..823f26c8e5d 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -142,6 +142,11 @@ def _is_model_type(module_path: str, class_name: str, fallback: bool) -> bool: "UNet2DConditionModel", "unet" in model_class_name.lower(), ) + is_qwen = _is_model_type( + "diffusers.models.transformers", + "QwenImageTransformer2DModel", + "qwen" in model_class_name.lower(), + ) cfg = getattr(model, "config", None) @@ -321,6 +326,40 @@ def _wan_inputs() -> dict[str, torch.Tensor]: "return_dict": False, } + def _qwen_inputs() -> dict[str, Any]: + # QwenImageTransformer2DModel does NOT take the standard + # (hidden_states[B,C,H,W], timestep, encoder_hidden_states) triple. It expects + # *packed* latents [B, (H//2)*(W//2), in_channels] plus encoder_hidden_states, + # encoder_hidden_states_mask, img_shapes, txt_seq_lens, and optional guidance. + # Timesteps are continuous in [0, 1] (not the diffusers [0, 1000] scale). + in_channels = getattr(cfg, "in_channels", 64) + joint_attention_dim = getattr(cfg, "joint_attention_dim", 3584) + guidance_embeds = getattr(cfg, "guidance_embeds", False) + + # Small packed spatial grid (already divided by the 2x2 patch size). + packed_h = packed_w = 4 + img_seq_len = packed_h * packed_w + text_seq_len = 8 + + dummy_inputs: dict[str, Any] = { + "hidden_states": torch.randn( + batch_size, img_seq_len, in_channels, device=device, dtype=dtype + ), + "encoder_hidden_states": torch.randn( + batch_size, text_seq_len, joint_attention_dim, device=device, dtype=dtype + ), + "encoder_hidden_states_mask": torch.ones( + batch_size, text_seq_len, device=device, dtype=torch.int64 + ), + "timestep": torch.tensor([0.5], device=device, dtype=dtype).expand(batch_size), + "img_shapes": [[(1, packed_h, packed_w)]] * batch_size, + "txt_seq_lens": [text_seq_len] * batch_size, + "return_dict": False, + } + if guidance_embeds: + dummy_inputs["guidance"] = torch.tensor([4.0], device=device, dtype=torch.float32) + return dummy_inputs + def _generic_transformer_inputs() -> dict[str, torch.Tensor] | None: # Try generic transformer handling for other model types # Check if model has common transformer attributes @@ -366,6 +405,7 @@ def _generic_transformer_inputs() -> dict[str, torch.Tensor] | None: ("dit", is_dit, _dit_inputs), ("wan", is_wan, _wan_inputs), ("unet", is_unet, _unet_inputs), + ("qwen", is_qwen, _qwen_inputs), ] for _, matches, build_inputs in model_input_builders: diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 0a02f751855..6d870a726bd 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -960,7 +960,9 @@ def _export_transformers_checkpoint( def _fuse_qkv_linears_diffusion( - model: nn.Module, dummy_forward_fn: Callable[[], None] | None = None + model: nn.Module, + dummy_forward_fn: Callable[[], None] | None = None, + strict: bool = False, ) -> None: """Fuse QKV linear layers that share the same input for diffusion models. @@ -994,6 +996,11 @@ def _fuse_qkv_linears_diffusion( model, dummy_forward_fn, collect_layernorms=False ) except Exception as e: + if strict: + raise RuntimeError( + f"QKV fusion dummy forward failed for {type(model).__name__}; a working " + f"dummy forward is required to export this model correctly. Original error: {e}" + ) from e print(f"Warning: Failed to run dummy forward for QKV fusion: {e}") print("Skipping QKV fusion. Quantization may still work but amax values won't be unified.") return @@ -1050,19 +1057,19 @@ def _promote_quantizer_tensors_to_module(component: nn.Module) -> None: if not is_quantlinear(sub_module): continue + # register_buffer overwrites an existing buffer of the same name, so a + # repeated export refreshes (rather than keeps stale) promoted tensors. input_quantizer = getattr(sub_module, "input_quantizer", None) pre_quant_scale = getattr(input_quantizer, "_pre_quant_scale", None) - if pre_quant_scale is not None and not hasattr(sub_module, "pre_quant_scale"): + if pre_quant_scale is not None: sub_module.register_buffer("pre_quant_scale", pre_quant_scale.detach().clone()) weight_quantizer = getattr(sub_module, "weight_quantizer", None) lora_a = getattr(weight_quantizer, "svdquant_lora_a", None) lora_b = getattr(weight_quantizer, "svdquant_lora_b", None) if lora_a is not None and lora_b is not None: - if not hasattr(sub_module, "svdquant_lora_a"): - sub_module.register_buffer("svdquant_lora_a", lora_a.detach().clone()) - if not hasattr(sub_module, "svdquant_lora_b"): - sub_module.register_buffer("svdquant_lora_b", lora_b.detach().clone()) + sub_module.register_buffer("svdquant_lora_a", lora_a.detach().clone()) + sub_module.register_buffer("svdquant_lora_b", lora_b.detach().clone()) def _export_diffusers_checkpoint( @@ -1138,7 +1145,11 @@ def _export_diffusers_checkpoint( # This is similar to requantize_resmooth_fused_llm_layers but simplified for diffusion # TODO: Add pre_quant_scale handling and FFN fusion for AWQ-style quantization print(f" Running QKV fusion for {component_name}...") - _fuse_qkv_linears_diffusion(component) + # Qwen-Image's packed-latent forward signature is non-standard; if the + # dummy forward fails for it, fail loudly rather than silently skipping + # fusion (which would export un-unified amax values). + is_qwen_component = "qwen" in type(component).__name__.lower() + _fuse_qkv_linears_diffusion(component, strict=is_qwen_component) # Step 4: Process quantized modules (convert weights, register scales) _process_quantized_modules(component, component_dtype, is_modelopt_qlora=False) diff --git a/tests/examples/diffusers/test_qwen_block_range_recipe.py b/tests/examples/diffusers/test_qwen_block_range_recipe.py index 906dd1fb60e..82690ae252c 100644 --- a/tests/examples/diffusers/test_qwen_block_range_recipe.py +++ b/tests/examples/diffusers/test_qwen_block_range_recipe.py @@ -77,10 +77,14 @@ def test_recipe_block_count_scales_with_model(): assert _disabled_block_indices(rules) == {0, 1, 58, 59} -def test_recipe_rejects_too_few_blocks(): - # 2 + 2 exclusion needs at least 5 blocks; 4 blocks must raise a clear error. +@pytest.mark.parametrize("num_blocks", [5, 4, 3]) +def test_recipe_rejects_too_few_blocks(num_blocks): + # A 2 + 2 exclusion needs at least 6 blocks (>= 2 quantized middle blocks). + # A 5-block model leaves only 1 middle block and must be rejected too. with pytest.raises(ValueError, match="at least"): - build_block_range_quant_cfg(_StubBackbone(4), exclude_first_n=2, exclude_last_n=2) + build_block_range_quant_cfg( + _StubBackbone(num_blocks), exclude_first_n=2, exclude_last_n=2 + ) def test_recipe_missing_block_module_raises(): diff --git a/tests/unit/torch/export/test_convert_hf_config_svdquant.py b/tests/unit/torch/export/test_convert_hf_config_svdquant.py index d00365703b1..13a22ab22d5 100644 --- a/tests/unit/torch/export/test_convert_hf_config_svdquant.py +++ b/tests/unit/torch/export/test_convert_hf_config_svdquant.py @@ -25,6 +25,7 @@ def test_nvfp4_svd_group_config_mirrors_awq_with_pre_quant_scale(): """The NVFP4_SVD config group is NVFP4 weights/activations + a pre_quant_scale flag.""" group = _quant_algo_to_group_config("NVFP4_SVD", group_size=16) assert group["pre_quant_scale"] is True + assert group["has_zero_point"] is False assert group["weights"] == { "dynamic": False, "num_bits": 4, @@ -57,6 +58,7 @@ def test_convert_hf_quant_config_format_nvfp4_svd(): assert "config_groups" in out group = out["config_groups"]["group_0"] assert group["pre_quant_scale"] is True + assert group["has_zero_point"] is False assert group["lora_rank"] == 32 assert group["weights"]["num_bits"] == 4 assert group["weights"]["type"] == "float" diff --git a/tests/unit/torch/quantization/test_svdquant_forward_fold.py b/tests/unit/torch/quantization/test_svdquant_forward_fold.py new file mode 100644 index 00000000000..f6a8a83f4b4 --- /dev/null +++ b/tests/unit/torch/quantization/test_svdquant_forward_fold.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""SVDQuant forward / fold coverage. + +These tests protect the invariants the diffusers SVDQuant export relies on +(DEC-5: the LoRA factors stay on the ``weight_quantizer`` in the live model; the +export layer promotes them). They complement (and do not modify) the existing +``test_calib.py::test_svdquant_lora_weights``. +""" + +from functools import partial + +import torch +import torch.nn as nn + +import modelopt.torch.quantization as mtq + + +class _SVDMLP(nn.Module): + def __init__(self, dim: int = 64): + super().__init__() + self.fc1 = nn.Linear(dim, dim) + self.fc2 = nn.Linear(dim, dim) + + def forward(self, x): + return self.fc2(torch.relu(self.fc1(x))) + + +def _forward_loop(model, dataloader): + for batch in dataloader: + model(batch) + + +def _quantize_svdquant(dim: int = 64) -> nn.Module: + model = _SVDMLP(dim) + quant_config = mtq.INT8_SMOOTHQUANT_CFG.copy() + quant_config["algorithm"] = {"method": "svdquant", "lowrank": 8} + data = [torch.randn(2, dim) for _ in range(2)] + mtq.quantize(model, quant_config, partial(_forward_loop, dataloader=data)) + return model + + +def _quantized_linears(model: nn.Module): + return [m for m in model.modules() if isinstance(m, torch.nn.Linear)] + + +def test_svdquant_lora_stays_on_weight_quantizer(): + """DEC-5: LoRA lives on the quantizer, not the module (export promotes it).""" + model = _quantize_svdquant() + linears = _quantized_linears(model) + assert linears + for module in linears: + wq = module.weight_quantizer + assert wq.svdquant_lora_a is not None + assert wq.svdquant_lora_b is not None + # Not refactored onto the module. + assert not hasattr(module, "svdquant_lora_a") + assert not hasattr(module, "svdquant_lora_b") + + +def test_svdquant_forward_includes_nonzero_residual(): + """The forward output includes a nonzero low-rank residual term.""" + model = _quantize_svdquant() + for module in _quantized_linears(model): + x = torch.randn(2, module.in_features) + + residual = module._compute_lora_residual(x) + assert residual is not None + assert torch.count_nonzero(residual) > 0 + + full = module(x) + + # Temporarily drop the LoRA buffers to get the base (no-residual) output. + wq = module.weight_quantizer + lora_a = wq._svdquant_lora_a + lora_b = wq._svdquant_lora_b + delattr(wq, "_svdquant_lora_a") + delattr(wq, "_svdquant_lora_b") + try: + base = module(x) + finally: + wq.register_buffer("_svdquant_lora_a", lora_a) + wq.register_buffer("_svdquant_lora_b", lora_b) + + # The residual measurably changes the forward output. + assert not torch.allclose(full, base) + + +def test_svdquant_fold_weight_removes_buffers_and_changes_weight(): + """fold_weight() folds the residual into the weight and drops the buffers.""" + model = _quantize_svdquant() + for module in _quantized_linears(model): + wq = module.weight_quantizer + assert hasattr(wq, "_svdquant_lora_a") + assert hasattr(wq, "_svdquant_lora_b") + + weight_before = module.weight.detach().clone() + module.fold_weight() + + assert not hasattr(wq, "_svdquant_lora_a") + assert not hasattr(wq, "_svdquant_lora_b") + # Folding (quantized weight + low-rank residual) changes the stored weight. + assert not torch.allclose(module.weight, weight_before) From 8e3b3ed12d48f1a9f667b0fa55d982cf4fc93d13 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 17:58:44 -0700 Subject: [PATCH 04/22] Qwen-Image SVDQuant: add export/fusion/promotion tests; drop plan terminology Round 2 (addresses round-1 Codex review: the round-1 code had no direct test coverage). Adds tests/unit/torch/export/test_diffusers_qwen_export.py: - Qwen dummy inputs: generate_diffusion_dummy_inputs builds the expected keys for a real tiny QwenImageTransformer2DModel, and the generated dummy forward runs on it (this is what catches any wrong shape/kwarg in the dummy-input builder). - Strict fusion: _fuse_qkv_linears_diffusion(strict=True) re-raises on a failing dummy forward; strict=False does not. - Structural export: _promote_quantizer_tensors_to_module promotes SVDQuant LoRA + pre_quant_scale to clean module keys that survive hide_quantizers_from_state_dict (promoted .svdquant_lora_a/b + .pre_quant_scale present; weight_quantizer / input_quantizer keys absent), on a calibrated tiny SVDQuant MLP. Also removes plan/workflow terminology (DEC-5, "pre-calibration") from source and test comments per the plan code-style note. Still pending (Round 3 / cluster): the full tiny Qwen pipeline fixture + e2e subprocess export test (needs diffusers' tokenizer/text-encoder construction and a GPU) and the AC-7 cluster run. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- .../diffusers/quantization/models_utils.py | 2 +- examples/diffusers/quantization/quantize.py | 2 +- examples/diffusers/quantization/utils.py | 2 +- .../diffusers/test_qwen_block_range_recipe.py | 4 +- .../export/test_diffusers_qwen_export.py | 134 ++++++++++++++++++ .../test_svdquant_forward_fold.py | 8 +- 6 files changed, 143 insertions(+), 9 deletions(-) create mode 100644 tests/unit/torch/export/test_diffusers_qwen_export.py diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index e7faf7589c1..eb2ff1f6039 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -248,7 +248,7 @@ def get_model_filter_func( }, # Quantize only ``transformer_blocks``; keep the first 2 and last 2 blocks # (and everything outside ``transformer_blocks``) in original precision. - # Applied pre-calibration via ``build_block_range_quant_cfg`` so SVDQuant + # Applied before calibration via ``build_block_range_quant_cfg`` so SVDQuant # never mutates the excluded blocks' weights. "block_range": { "exclude_first_n": 2, diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 57b9db589db..e81dc51a66b 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -171,7 +171,7 @@ def get_quant_config(self, n_steps: int, backbone: torch.nn.Module) -> Any: # Apply the transformer-block-range recipe (e.g. Qwen-Image) BEFORE # calibration. This restricts quantization to `transformer_blocks` and - # excludes the first/last N blocks. It must run pre-calibration so that + # excludes the first/last N blocks. It must run before calibration so that # SVDQuant does not mutate the weights of the excluded blocks. The recipe # is format-agnostic (applies to FP8/NVFP4/SVDQuant alike). block_range = MODEL_DEFAULTS.get(self.model_config.model_type, {}).get("block_range") diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index be3f6276db9..fd57c328378 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -114,7 +114,7 @@ def filter_func_wan_video(name: str) -> bool: # Qwen-Image's transformer has 60 ``transformer_blocks``. The recipe quantizes # only those blocks while keeping the first two and last two -- and everything # outside ``transformer_blocks`` -- in original precision. The model-agnostic, -# pre-calibration form of this recipe (deriving the block count from the model) +# before-calibration form of this recipe (deriving the block count from the model) # lives in quantize.py; this name-only filter covers the plain FP8/NVFP4 path # for the full 60-block Qwen-Image transformer. QWEN_IMAGE_NUM_TRANSFORMER_BLOCKS = 60 diff --git a/tests/examples/diffusers/test_qwen_block_range_recipe.py b/tests/examples/diffusers/test_qwen_block_range_recipe.py index 82690ae252c..8214ec1c124 100644 --- a/tests/examples/diffusers/test_qwen_block_range_recipe.py +++ b/tests/examples/diffusers/test_qwen_block_range_recipe.py @@ -16,8 +16,8 @@ """Unit tests for the transformer-block-range quantization recipe (e.g. Qwen-Image). The recipe must quantize only the linears under ``transformer_blocks`` while -excluding the first/last N blocks, and it must be expressible as pre-calibration -``quant_cfg`` rules (so SVDQuant never mutates the excluded blocks' weights). +excluding the first/last N blocks, and it must be expressible as ``quant_cfg`` +rules applied before calibration (so SVDQuant never mutates the excluded blocks). """ import re diff --git a/tests/unit/torch/export/test_diffusers_qwen_export.py b/tests/unit/torch/export/test_diffusers_qwen_export.py new file mode 100644 index 00000000000..6577a8800f0 --- /dev/null +++ b/tests/unit/torch/export/test_diffusers_qwen_export.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Tests for the Qwen-Image SVDQuant diffusers export path. + +Covers the three pieces added for Qwen support: +- the Qwen branch of ``generate_diffusion_dummy_inputs`` (validated by running the + dummy forward on a real tiny ``QwenImageTransformer2DModel``), +- the strict-failure mode of ``_fuse_qkv_linears_diffusion``, +- promotion of quantizer-owned SVDQuant tensors to clean module-level keys that + survive ``hide_quantizers_from_state_dict``. +""" + +from functools import partial + +import pytest +import torch +import torch.nn as nn + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.diffusers_utils import ( + generate_diffusion_dummy_forward_fn, + generate_diffusion_dummy_inputs, + hide_quantizers_from_state_dict, +) +from modelopt.torch.export.unified_export_hf import ( + _fuse_qkv_linears_diffusion, + _promote_quantizer_tensors_to_module, +) + + +class _MLP(nn.Module): + def __init__(self, dim: int = 64): + super().__init__() + self.fc1 = nn.Linear(dim, dim) + self.fc2 = nn.Linear(dim, dim) + + def forward(self, x): + return self.fc2(torch.relu(self.fc1(x))) + + +def _forward_loop(model, data): + for batch in data: + model(batch) + + +def _quantize(model: nn.Module, algorithm=None, dim: int = 64) -> nn.Module: + cfg = mtq.INT8_SMOOTHQUANT_CFG.copy() + if algorithm is not None: + cfg["algorithm"] = algorithm + data = [torch.randn(2, dim) for _ in range(2)] + mtq.quantize(model, cfg, partial(_forward_loop, data=data)) + return model + + +def test_qwen_dummy_inputs_drive_real_transformer_forward(): + """The Qwen dummy inputs must actually drive a real tiny Qwen transformer.""" + pytest.importorskip("diffusers") + from _test_utils.torch.diffusers_models import get_tiny_qwen_image_transformer + + transformer = get_tiny_qwen_image_transformer().to("cpu", torch.float32).eval() + + inputs = generate_diffusion_dummy_inputs(transformer, torch.device("cpu"), torch.float32) + assert inputs is not None + for key in ( + "hidden_states", + "encoder_hidden_states", + "encoder_hidden_states_mask", + "img_shapes", + "txt_seq_lens", + ): + assert key in inputs, f"missing Qwen dummy input '{key}'" + assert inputs["hidden_states"].shape[-1] == transformer.config.in_channels + assert inputs["encoder_hidden_states"].shape[-1] == transformer.config.joint_attention_dim + + # Strongest check: the generated dummy inputs run through the real model. + with torch.no_grad(): + generate_diffusion_dummy_forward_fn(transformer)() + + +def test_qwen_qkv_fusion_strict_raises_on_failed_dummy_forward(): + """strict=True turns a dummy-forward failure into a hard error; strict=False does not.""" + model = _quantize(_MLP()) + + def _boom(): + raise RuntimeError("dummy forward failed") + + with pytest.raises(RuntimeError): + _fuse_qkv_linears_diffusion(model, dummy_forward_fn=_boom, strict=True) + + # Non-strict path warns and returns without raising. + _fuse_qkv_linears_diffusion(model, dummy_forward_fn=_boom, strict=False) + + +def test_svdquant_promotion_survives_hide_quantizers(): + """Promoted LoRA + pre_quant_scale land on the module under clean keys and + survive ``hide_quantizers_from_state_dict`` (which strips the quantizers).""" + model = _quantize(_MLP(), algorithm={"method": "svdquant", "lowrank": 8}) + + _promote_quantizer_tensors_to_module(model) + + linears = [m for m in model.modules() if isinstance(m, torch.nn.Linear)] + assert linears + for module in linears: + assert hasattr(module, "svdquant_lora_a") + assert hasattr(module, "svdquant_lora_b") + # INT8_SMOOTHQUANT produces a pre_quant_scale that is promoted too. + assert hasattr(module, "pre_quant_scale") + # Rank-consistent shapes: lora_a [rank, in], lora_b [out, rank]. + assert module.svdquant_lora_a.shape[1] == module.in_features + assert module.svdquant_lora_b.shape[0] == module.out_features + assert module.svdquant_lora_a.shape[0] == module.svdquant_lora_b.shape[1] + + with hide_quantizers_from_state_dict(model): + keys = list(model.state_dict().keys()) + + assert any(k.endswith(".svdquant_lora_a") for k in keys) + assert any(k.endswith(".svdquant_lora_b") for k in keys) + assert any(k.endswith(".pre_quant_scale") for k in keys) + # Clean keys only: no quantizer-prefixed keys remain once quantizers are hidden. + assert not any("weight_quantizer" in k for k in keys) + assert not any("input_quantizer" in k for k in keys) diff --git a/tests/unit/torch/quantization/test_svdquant_forward_fold.py b/tests/unit/torch/quantization/test_svdquant_forward_fold.py index f6a8a83f4b4..085414f93a3 100644 --- a/tests/unit/torch/quantization/test_svdquant_forward_fold.py +++ b/tests/unit/torch/quantization/test_svdquant_forward_fold.py @@ -15,9 +15,9 @@ """SVDQuant forward / fold coverage. -These tests protect the invariants the diffusers SVDQuant export relies on -(DEC-5: the LoRA factors stay on the ``weight_quantizer`` in the live model; the -export layer promotes them). They complement (and do not modify) the existing +These tests protect the invariants the diffusers SVDQuant export relies on: the +LoRA factors stay on the ``weight_quantizer`` in the live model and the export +layer promotes them. They complement (and do not modify) the existing ``test_calib.py::test_svdquant_lora_weights``. """ @@ -58,7 +58,7 @@ def _quantized_linears(model: nn.Module): def test_svdquant_lora_stays_on_weight_quantizer(): - """DEC-5: LoRA lives on the quantizer, not the module (export promotes it).""" + """LoRA lives on the quantizer, not the module (the export layer promotes it).""" model = _quantize_svdquant() linears = _quantized_linears(model) assert linears From 027a5e200a362f554fe1d5fcb90348c50b7de2b6 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 18:15:05 -0700 Subject: [PATCH 05/22] Qwen-Image SVDQuant: offline tiny Qwen fixture + e2e export test Round 3 (addresses round-2 Codex review): - Fix the tiny Qwen-Image pipeline fixture (tests/_test_utils/torch/diffusers_models.py): build the Qwen2.5-VL text encoder inline from a tiny Qwen2_5_VLConfig (no Hub model load; the previous hf-internal-testing/...Qwen2_5_VL id does not exist), load the tokenizer from the tiny ...Qwen2VL id diffusers' own fast test uses, build the transformer with num_layers=6 (so the corrected first-2/last-2 block-range recipe, which needs >=6 blocks, is valid) and joint_attention_dim=16 matching the text encoder hidden_size, and a z_dim=4 VAE. Mirrors diffusers' QwenImagePipelineFastTests.get_dummy_components. - Add Qwen FP8 / NVFP4 / NVFP4-SVDQuant cases to test_export_diffusers_hf_ckpt.py using the tiny fixture. The test opens transformer/config.json and the exported safetensors and asserts: quant_method=modelopt; no weight_quantizer / input_quantizer._amax keys; for SVDQuant, promoted .svdquant_lora_a/b + .pre_quant_scale keys, config group pre_quant_scale/has_zero_point/ lora_rank, and non-empty ignore (excluded blocks); for plain formats, weight_scale. GPU/diffusers skip-guarded. - Drop remaining workflow terminology (Step 4.5, before-calibration) from the comments I introduced. Still cluster-only (no GPU here): executing these tests and the AC-7 harness run. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- examples/diffusers/quantization/utils.py | 2 +- modelopt/torch/export/unified_export_hf.py | 4 +- tests/_test_utils/torch/diffusers_models.py | 87 +++++++++++------ .../test_export_diffusers_hf_ckpt.py | 93 +++++++++++++++++++ 4 files changed, 153 insertions(+), 33 deletions(-) diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index fd57c328378..c3cfdcd5cdd 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -114,7 +114,7 @@ def filter_func_wan_video(name: str) -> bool: # Qwen-Image's transformer has 60 ``transformer_blocks``. The recipe quantizes # only those blocks while keeping the first two and last two -- and everything # outside ``transformer_blocks`` -- in original precision. The model-agnostic, -# before-calibration form of this recipe (deriving the block count from the model) +# config-driven form of this recipe (deriving the block count from the model) # lives in quantize.py; this name-only filter covers the plain FP8/NVFP4 path # for the full 60-block Qwen-Image transformer. QWEN_IMAGE_NUM_TRANSFORMER_BLOCKS = 60 diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 6d870a726bd..6f08b653fa4 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1154,8 +1154,8 @@ def _export_diffusers_checkpoint( # Step 4: Process quantized modules (convert weights, register scales) _process_quantized_modules(component, component_dtype, is_modelopt_qlora=False) - # Step 4.5: Promote quantizer-owned tensors (AWQ pre_quant_scale and - # SVDQuant LoRA factors) onto the module so they survive + # Promote quantizer-owned tensors (AWQ pre_quant_scale and SVDQuant + # LoRA factors) onto the module so they survive # hide_quantizers_from_state_dict and are embedded in the component's # main safetensors under clean, AWQ-aligned keys. _promote_quantizer_tensors_to_module(component) diff --git a/tests/_test_utils/torch/diffusers_models.py b/tests/_test_utils/torch/diffusers_models.py index a42ebdd8ffb..3a5e4adc982 100644 --- a/tests/_test_utils/torch/diffusers_models.py +++ b/tests/_test_utils/torch/diffusers_models.py @@ -312,20 +312,16 @@ def get_tiny_qwen_image_vae(**config_kwargs): def create_tiny_qwen_image_pipeline_dir(tmp_path: Path) -> Path: - """Create and save a tiny Qwen-Image pipeline to a directory (SKETCH). - - Mirrors ``create_tiny_wan22_pipeline_dir``. Needs in-container validation; the - fragile piece is the Qwen2.5-VL text encoder. This prefers a tiny-random HF model - (as Wan uses ``hf-internal-testing/tiny-random-t5``); if that id drifts or the - config schema differs across transformers versions, copy the text-encoder - construction from diffusers' own QwenImage fast test - (``tests/pipelines/qwenimage/test_qwenimage.py``). - - For the DMD2 mock-data training path the transformer consumes the dataloader's - embeddings rather than the text encoder, so the bundled tiny text encoder only - needs to load; its hidden size is intentionally decoupled from the transformer's - ``joint_attention_dim`` (set the dataloader's ``text_embed_dim`` to match instead). - The saved dir loads with ``QwenImagePipeline.from_pretrained(path)``. + """Create and save a tiny, (mostly) offline Qwen-Image pipeline to a directory. + + Mirrors diffusers' ``QwenImagePipelineFastTests.get_dummy_components``: the + Qwen2.5-VL text encoder is built inline from a tiny ``Qwen2_5_VLConfig`` (no Hub + model load); only the tokenizer is fetched from the tiny ``Qwen2VL`` test repo + (building a Qwen tokenizer fully offline is impractical). The transformer uses + ``num_layers=6`` so the first-2/last-2 block-range recipe is valid, and its + ``joint_attention_dim`` matches the text encoder ``hidden_size`` (16) so the + pipeline runs end-to-end during quantization calibration. The saved dir loads + with ``QwenImagePipeline.from_pretrained(path)``. """ if QwenImageTransformer2DModel is None or AutoencoderKLQwenImage is None: pytest.skip("QwenImage diffusers classes not available in this diffusers version.") @@ -333,27 +329,58 @@ def create_tiny_qwen_image_pipeline_dir(tmp_path: Path) -> Path: transformers = pytest.importorskip("transformers") - # Tiny Qwen2.5-VL text encoder + matching Qwen2 tokenizer (loaded, but bypassed - # during DMD2 mock-data training). - # NOTE (validated 2026-06-06): the hf-internal-testing id below does NOT exist on the - # Hub, so this fixture currently skips. To make the recipe e2e runnable in CI, - # construct the encoder inline from a tiny ``Qwen2_5_VLConfig`` (nested text + vision - # config) — mirror diffusers' ``QwenImagePipelineFastTests.get_dummy_components`` in - # ``tests/pipelines/qwenimage/test_qwenimage.py``. - tiny_id = "hf-internal-testing/tiny-random-Qwen2_5_VLForConditionalGeneration" + # Tiny Qwen2.5-VL text encoder, built offline from a tiny config (no Hub model + # load), mirroring diffusers' QwenImagePipelineFastTests.get_dummy_components. + qwen_vl_config = transformers.Qwen2_5_VLConfig( + text_config={ + "hidden_size": 16, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "rope_scaling": { + "mrope_section": [1, 1, 2], + "rope_type": "default", + "type": "default", + }, + "rope_theta": 1000000.0, + }, + vision_config={ + "depth": 2, + "hidden_size": 16, + "intermediate_size": 16, + "num_heads": 2, + "out_hidden_size": 16, + }, + hidden_size=16, + vocab_size=152064, + vision_end_token_id=151653, + vision_start_token_id=151652, + vision_token_id=151654, + ) + text_encoder = transformers.Qwen2_5_VLForConditionalGeneration(qwen_vl_config).eval() + + # The Qwen tokenizer cannot be built fully offline; load the tiny one diffusers' + # own fast test uses (this id exists, unlike the Qwen2.5-VL one previously tried). try: - text_encoder = transformers.Qwen2_5_VLForConditionalGeneration.from_pretrained(tiny_id) - tokenizer = transformers.Qwen2Tokenizer.from_pretrained(tiny_id) - except Exception as exc: # pragma: no cover - depends on hub availability / version - pytest.skip( - f"tiny Qwen2.5-VL text encoder unavailable ({exc}); " - "copy the fixture from diffusers' QwenImage fast test" + tokenizer = transformers.Qwen2Tokenizer.from_pretrained( + "hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration" ) + except Exception as exc: # pragma: no cover - depends on hub availability + pytest.skip(f"tiny Qwen tokenizer unavailable ({exc})") torch.manual_seed(0) - transformer = get_tiny_qwen_image_transformer() + # num_layers=6 so the first-2/last-2 block-range recipe (which needs >=6 blocks) + # is valid; joint_attention_dim must match the text encoder hidden_size (16). + transformer = get_tiny_qwen_image_transformer( + num_layers=6, + in_channels=16, + out_channels=4, + joint_attention_dim=16, + num_attention_heads=3, + ) torch.manual_seed(0) - vae = get_tiny_qwen_image_vae() + vae = get_tiny_qwen_image_vae(z_dim=4, latents_mean=[0.0] * 4, latents_std=[1.0] * 4) scheduler = FlowMatchEulerDiscreteScheduler( base_image_seq_len=256, diff --git a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py index 88821bbf8f7..6ea8d22b51e 100644 --- a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py +++ b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json from pathlib import Path from typing import NamedTuple @@ -130,6 +131,98 @@ def test_diffusers_hf_ckpt_export(model: DiffuserHfExportModel, tmp_path: Path) assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}" +class QwenHfExportModel(NamedTuple): + format_type: str + quant_algo: str + is_svdquant: bool + + def quantize_and_export_hf(self, tiny_qwen_image_path: str, tmp_path: Path) -> Path: + hf_ckpt_dir = tmp_path / f"qwen_{self.format_type}_{self.quant_algo}_hf_ckpt" + cmd_args = [ + "python", + "quantize.py", + "--model", + "qwen-image", + "--override-model-path", + str(tiny_qwen_image_path), + "--format", + self.format_type, + "--quant-algo", + self.quant_algo, + "--collect-method", + "default", + "--model-dtype", + "BFloat16", + "--trt-high-precision-dtype", + "BFloat16", + "--calib-size", + "2", + "--batch-size", + "1", + "--n-steps", + "2", + "--hf-ckpt-dir", + str(hf_ckpt_dir), + ] + if self.is_svdquant: + cmd_args.extend(["--lowrank", "8"]) + run_example_command(cmd_args, "diffusers/quantization") + return hf_ckpt_dir + + +@pytest.mark.parametrize( + "qwen_model", + [ + pytest.param(QwenHfExportModel("fp8", "max", False), marks=minimum_sm(89)), + pytest.param(QwenHfExportModel("fp4", "max", False), marks=minimum_sm(89)), + pytest.param(QwenHfExportModel("fp4", "svdquant", True), marks=minimum_sm(89)), + ], + ids=["qwen_fp8_max", "qwen_nvfp4_max", "qwen_nvfp4_svdquant"], +) +def test_qwen_image_hf_ckpt_export( + qwen_model: QwenHfExportModel, tiny_qwen_image_path: str, tmp_path: Path +) -> None: + from safetensors import safe_open + + hf_ckpt_dir = qwen_model.quantize_and_export_hf(tiny_qwen_image_path, tmp_path) + assert hf_ckpt_dir.exists(), f"HF checkpoint directory was not created: {hf_ckpt_dir}" + + # The transformer is the quantized component. + transformer_dir = hf_ckpt_dir / "transformer" + config_path = transformer_dir / "config.json" + assert config_path.exists(), f"no transformer/config.json in {hf_ckpt_dir}" + quant_config = json.loads(config_path.read_text()).get("quantization_config") + assert quant_config is not None, "missing quantization_config" + assert quant_config.get("quant_method") == "modelopt" + + keys: set[str] = set() + safetensors_files = sorted(transformer_dir.rglob("*.safetensors")) + assert safetensors_files, f"no safetensors in {transformer_dir}" + for path in safetensors_files: + with safe_open(str(path), framework="pt") as handle: + keys.update(handle.keys()) + + # No live quantizer state should leak into the exported checkpoint. + assert not any("weight_quantizer" in k for k in keys), "quantizer keys leaked into export" + assert not any("input_quantizer._amax" in k for k in keys) + + if qwen_model.is_svdquant: + # Promoted SVDQuant tensors under clean module-level keys. + assert any(k.endswith(".svdquant_lora_a") for k in keys) + assert any(k.endswith(".svdquant_lora_b") for k in keys) + assert any(k.endswith(".pre_quant_scale") for k in keys) + # config schema modeled on nvfp4_awq. + group = next(iter(quant_config.get("config_groups", {}).values()), {}) + assert group.get("pre_quant_scale") is True + assert group.get("has_zero_point") is False + assert "lora_rank" in group + # Excluded first-2/last-2 blocks + non-transformer_blocks modules are ignored. + assert quant_config.get("ignore"), "expected excluded modules in 'ignore'" + else: + # Plain FP8/NVFP4 expose weight scales. + assert any(k.endswith(".weight_scale") for k in keys), "no weight_scale in export" + + class Wan22HfExportModel(NamedTuple): model: str backbone: str | None From a51105814724655bfa645e8a0f95dab96ef2c5c3 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 18:28:49 -0700 Subject: [PATCH 06/22] Qwen-Image SVDQuant: offline tokenizer, stronger export test, drop Step comments Round 4 (addresses round-3 Codex review): - Offline tiny Qwen tokenizer: _build_local_qwen2_tokenizer builds a deterministic byte-level Qwen2 tokenizer locally (GPT-2 byte->unicode vocab + Qwen specials, empty merges) instead of a Hub load; removes the tokenizer-unavailable skip path. - Strengthen test_qwen_image_hf_ckpt_export: assert equal module-prefix sets for .svdquant_lora_a/.svdquant_lora_b/.pre_quant_scale; promoted linears are a subset of weight-scaled linears; only the middle blocks {2,3} of 6 are quantized (first-2/ last-2 excluded); lora_a=[rank,in]/lora_b=[out,rank] with rank == --lowrank (8); NVFP4 weight_scale_2 present; exact config (quant_algo=NVFP4_SVD, lora_rank=8, pre_quant_scale=True, has_zero_point=False, non-empty ignore). - Remove the remaining "Step N:" workflow comments from unified_export_hf.py (the round-3 "grep clean" claim was wrong; verified clean across the whole file). Still cluster-only (no GPU/torch/diffusers here): executing these tests and the AC-7 harness run. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- modelopt/torch/export/unified_export_hf.py | 22 +++--- tests/_test_utils/torch/diffusers_models.py | 39 ++++++++--- .../test_export_diffusers_hf_ckpt.py | 69 ++++++++++++++++--- 3 files changed, 101 insertions(+), 29 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 6f08b653fa4..164f2a0e1fe 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1100,7 +1100,7 @@ def _export_diffusers_checkpoint( """ export_dir = Path(export_dir) - # Step 1: Get all pipeline components (nn.Module, tokenizers, schedulers, etc.) + # Get all pipeline components (nn.Module, tokenizers, schedulers, etc.) all_components = get_diffusion_components(pipe, components) if not all_components: @@ -1122,7 +1122,7 @@ def _export_diffusers_checkpoint( except Exception: is_diffusers_pipe = False - # Step 3: Export each nn.Module component with quantization handling + # Export each nn.Module component with quantization handling for component_name, component in module_components.items(): is_quantized = has_quantized_modules(component) status = "quantized" if is_quantized else "non-quantized" @@ -1141,7 +1141,7 @@ def _export_diffusers_checkpoint( component_dtype = dtype if dtype is not None else infer_dtype_from_model(component) if is_quantized: - # Step 3.5: Fuse QKV linears that share the same input (unify amax values) + # Fuse QKV linears that share the same input (unify amax values) # This is similar to requantize_resmooth_fused_llm_layers but simplified for diffusion # TODO: Add pre_quant_scale handling and FFN fusion for AWQ-style quantization print(f" Running QKV fusion for {component_name}...") @@ -1151,7 +1151,7 @@ def _export_diffusers_checkpoint( is_qwen_component = "qwen" in type(component).__name__.lower() _fuse_qkv_linears_diffusion(component, strict=is_qwen_component) - # Step 4: Process quantized modules (convert weights, register scales) + # Process quantized modules (convert weights, register scales) _process_quantized_modules(component, component_dtype, is_modelopt_qlora=False) # Promote quantizer-owned tensors (AWQ pre_quant_scale and SVDQuant @@ -1160,7 +1160,7 @@ def _export_diffusers_checkpoint( # main safetensors under clean, AWQ-aligned keys. _promote_quantizer_tensors_to_module(component) - # Step 5: Build quantization config + # Build quantization config quant_config = get_quant_config(component, is_modelopt_qlora=False) if quant_config: quantization_details = quant_config.get("quantization", {}) @@ -1171,7 +1171,7 @@ def _export_diffusers_checkpoint( quantization_details["lora_rank"] = svdquant_rank hf_quant_config = convert_hf_quant_config_format(quant_config) if quant_config else None - # Step 6: Save the component + # Save the component # - diffusers ModelMixin.save_pretrained does NOT accept state_dict parameter # - for non-diffusers modules (e.g., LTX-2 transformer), fall back to torch.save if hasattr(component, "save_pretrained"): @@ -1181,7 +1181,7 @@ def _export_diffusers_checkpoint( with hide_quantizers_from_state_dict(component): _save_component_state_dict_safetensors(component, component_export_dir) - # Step 7: Post-process — merge, metadata, padding, swizzle + # Post-process — merge, metadata, padding, swizzle _postprocess_safetensors( component_export_dir, pipe, @@ -1189,7 +1189,7 @@ def _export_diffusers_checkpoint( **kwargs, ) - # Step 8: Update config.json with quantization info + # Update config.json with quantization info if hf_quant_config is not None: config_path = component_export_dir / "config.json" if config_path.exists(): @@ -1204,7 +1204,7 @@ def _export_diffusers_checkpoint( else: _save_component_state_dict_safetensors(component, component_export_dir) - # Step 9: Update config.json with sparse attention info (both quantized and non-quantized) + # Update config.json with sparse attention info (both quantized and non-quantized) if export_sparse_attention_config is not None: sparse_attn_config = export_sparse_attention_config(component) if sparse_attn_config is not None: @@ -1219,7 +1219,7 @@ def _export_diffusers_checkpoint( print(f" Saved to: {component_export_dir}") - # Step 4: Export non-nn.Module components (tokenizers, schedulers, feature extractors, etc.) + # Export non-nn.Module components (tokenizers, schedulers, feature extractors, etc.) if is_diffusers_pipe: for component_name, component in all_components.items(): # Skip nn.Module components (already handled above) @@ -1247,7 +1247,7 @@ def _export_diffusers_checkpoint( print(f" Saved to: {component_export_dir}") - # Step 5: For pipelines, also save model_index.json + # For pipelines, also save model_index.json if is_diffusers_pipe: model_index_path = export_dir / "model_index.json" is_partial_export = components is not None diff --git a/tests/_test_utils/torch/diffusers_models.py b/tests/_test_utils/torch/diffusers_models.py index 3a5e4adc982..71e0e19cf78 100644 --- a/tests/_test_utils/torch/diffusers_models.py +++ b/tests/_test_utils/torch/diffusers_models.py @@ -311,6 +311,35 @@ def get_tiny_qwen_image_vae(**config_kwargs): return AutoencoderKLQwenImage(**kwargs) +def _build_local_qwen2_tokenizer(out_dir: Path): + """Build a tiny, fully offline byte-level Qwen2 tokenizer (no Hub access). + + Uses the GPT-2/Qwen byte->unicode mapping for the 256 single-byte tokens plus + Qwen's core special tokens, with an empty merge table (pure byte-level + fallback). This is enough to tokenize calibration prompts so the pipeline runs + end-to-end; it is not meant for high-quality text. + """ + import json + + import transformers + from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode + + out_dir.mkdir(parents=True, exist_ok=True) + vocab = {token: idx for idx, token in enumerate(bytes_to_unicode().values())} + for special in ("<|endoftext|>", "<|im_start|>", "<|im_end|>"): + vocab.setdefault(special, len(vocab)) + (out_dir / "vocab.json").write_text(json.dumps(vocab)) + (out_dir / "merges.txt").write_text("#version: 0.2\n") + + return transformers.Qwen2Tokenizer( + vocab_file=str(out_dir / "vocab.json"), + merges_file=str(out_dir / "merges.txt"), + unk_token="<|endoftext|>", + eos_token="<|endoftext|>", + pad_token="<|endoftext|>", + ) + + def create_tiny_qwen_image_pipeline_dir(tmp_path: Path) -> Path: """Create and save a tiny, (mostly) offline Qwen-Image pipeline to a directory. @@ -360,14 +389,8 @@ def create_tiny_qwen_image_pipeline_dir(tmp_path: Path) -> Path: ) text_encoder = transformers.Qwen2_5_VLForConditionalGeneration(qwen_vl_config).eval() - # The Qwen tokenizer cannot be built fully offline; load the tiny one diffusers' - # own fast test uses (this id exists, unlike the Qwen2.5-VL one previously tried). - try: - tokenizer = transformers.Qwen2Tokenizer.from_pretrained( - "hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration" - ) - except Exception as exc: # pragma: no cover - depends on hub availability - pytest.skip(f"tiny Qwen tokenizer unavailable ({exc})") + # Deterministic local byte-level Qwen2 tokenizer (built offline; no Hub, no skip). + tokenizer = _build_local_qwen2_tokenizer(tmp_path / "qwen_tokenizer") torch.manual_seed(0) # num_layers=6 so the first-2/last-2 block-range recipe (which needs >=6 blocks) diff --git a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py index 6ea8d22b51e..58fe587cf0e 100644 --- a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py +++ b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py @@ -179,6 +179,29 @@ def quantize_and_export_hf(self, tiny_qwen_image_path: str, tmp_path: Path) -> P ], ids=["qwen_fp8_max", "qwen_nvfp4_max", "qwen_nvfp4_svdquant"], ) +def _module_prefixes(keys: set[str], suffix: str) -> set[str]: + """Module paths (key minus suffix) for every key ending in ``suffix``.""" + return {k[: -len(suffix)] for k in keys if k.endswith(suffix)} + + +def _block_indices(prefixes: set[str]) -> set[int]: + """transformer_blocks indices referenced by a set of module prefixes.""" + import re + + indices = set() + for prefix in prefixes: + match = re.search(r"transformer_blocks\.(\d+)\.", prefix) + if match: + indices.add(int(match.group(1))) + return indices + + +# Tiny Qwen fixture has 6 transformer blocks; the recipe excludes the first 2 and +# last 2, so only blocks 2 and 3 are quantized. +_QWEN_QUANTIZED_BLOCKS = {2, 3} +_QWEN_LORA_RANK = 8 + + def test_qwen_image_hf_ckpt_export( qwen_model: QwenHfExportModel, tiny_qwen_image_path: str, tmp_path: Path ) -> None: @@ -196,31 +219,57 @@ def test_qwen_image_hf_ckpt_export( assert quant_config.get("quant_method") == "modelopt" keys: set[str] = set() + lora_tensors: dict[str, "object"] = {} safetensors_files = sorted(transformer_dir.rglob("*.safetensors")) assert safetensors_files, f"no safetensors in {transformer_dir}" for path in safetensors_files: with safe_open(str(path), framework="pt") as handle: - keys.update(handle.keys()) + for key in handle.keys(): + keys.add(key) + if key.endswith(".svdquant_lora_a") or key.endswith(".svdquant_lora_b"): + lora_tensors[key] = handle.get_tensor(key) # No live quantizer state should leak into the exported checkpoint. assert not any("weight_quantizer" in k for k in keys), "quantizer keys leaked into export" assert not any("input_quantizer._amax" in k for k in keys) + # Recipe: only the middle blocks are quantized (first-2/last-2 excluded). + weight_scale_prefixes = _module_prefixes(keys, ".weight_scale") + assert _block_indices(weight_scale_prefixes) == _QWEN_QUANTIZED_BLOCKS, ( + f"expected only blocks {_QWEN_QUANTIZED_BLOCKS} quantized" + ) + if qwen_model.is_svdquant: - # Promoted SVDQuant tensors under clean module-level keys. - assert any(k.endswith(".svdquant_lora_a") for k in keys) - assert any(k.endswith(".svdquant_lora_b") for k in keys) - assert any(k.endswith(".pre_quant_scale") for k in keys) - # config schema modeled on nvfp4_awq. + a_prefixes = _module_prefixes(keys, ".svdquant_lora_a") + b_prefixes = _module_prefixes(keys, ".svdquant_lora_b") + pqs_prefixes = _module_prefixes(keys, ".pre_quant_scale") + assert a_prefixes, "no promoted svdquant_lora_a keys" + # Every promoted linear carries lora_a, lora_b, and pre_quant_scale. + assert a_prefixes == b_prefixes == pqs_prefixes + # ...and each is a quantized linear, only in the middle blocks. + assert a_prefixes <= weight_scale_prefixes + assert _block_indices(a_prefixes) == _QWEN_QUANTIZED_BLOCKS + # Rank-consistent shapes; lora_a=[rank, in], lora_b=[out, rank], rank == --lowrank. + for key, tensor in lora_tensors.items(): + if key.endswith(".svdquant_lora_a"): + assert tensor.shape[0] == _QWEN_LORA_RANK + else: + assert tensor.shape[1] == _QWEN_LORA_RANK + # NVFP4 secondary scales are present. + assert any(k.endswith(".weight_scale_2") for k in keys) + # config schema (modeled on nvfp4_awq). + assert quant_config.get("quant_algo") == "NVFP4_SVD" group = next(iter(quant_config.get("config_groups", {}).values()), {}) + assert group.get("lora_rank") == _QWEN_LORA_RANK assert group.get("pre_quant_scale") is True assert group.get("has_zero_point") is False - assert "lora_rank" in group - # Excluded first-2/last-2 blocks + non-transformer_blocks modules are ignored. assert quant_config.get("ignore"), "expected excluded modules in 'ignore'" else: - # Plain FP8/NVFP4 expose weight scales. - assert any(k.endswith(".weight_scale") for k in keys), "no weight_scale in export" + # Plain FP8/NVFP4: weight scales present, no SVDQuant tensors. + assert weight_scale_prefixes, "no weight_scale in export" + assert not any(k.endswith(".svdquant_lora_a") for k in keys) + if qwen_model.format_type == "fp4": + assert any(k.endswith(".weight_scale_2") for k in keys) class Wan22HfExportModel(NamedTuple): From a6a3d590bcf0b7e4a8581b1443a449f8dc06f5e6 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 18:39:10 -0700 Subject: [PATCH 07/22] Qwen-Image SVDQuant: fix misplaced parametrize decorator + tighten export test Round 5 (addresses round-4 Codex review, which found a regression I introduced): - The round-4 edit inserted the _module_prefixes/_block_indices helpers between @pytest.mark.parametrize("qwen_model", ...) and test_qwen_image_hf_ckpt_export, so the decorator was attached to the helper and the test would request an undefined qwen_model fixture. Moved the helpers/constants above the decorator so it directly decorates the test (verified via ast: the test now carries the qwen_model parametrization and the helper is undecorated). - Tightened SVDQuant assertions: require a_prefixes == b_prefixes == pqs_prefixes == weight_scale_prefixes (every quantized linear is promoted, no gaps), and assert every quantized prefix is under transformer_blocks (nothing outside is quantized), in addition to the {2,3}-only block check. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- .../test_export_diffusers_hf_ckpt.py | 32 +++++++++++-------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py index 58fe587cf0e..9e2e30a160c 100644 --- a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py +++ b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py @@ -170,15 +170,6 @@ def quantize_and_export_hf(self, tiny_qwen_image_path: str, tmp_path: Path) -> P return hf_ckpt_dir -@pytest.mark.parametrize( - "qwen_model", - [ - pytest.param(QwenHfExportModel("fp8", "max", False), marks=minimum_sm(89)), - pytest.param(QwenHfExportModel("fp4", "max", False), marks=minimum_sm(89)), - pytest.param(QwenHfExportModel("fp4", "svdquant", True), marks=minimum_sm(89)), - ], - ids=["qwen_fp8_max", "qwen_nvfp4_max", "qwen_nvfp4_svdquant"], -) def _module_prefixes(keys: set[str], suffix: str) -> set[str]: """Module paths (key minus suffix) for every key ending in ``suffix``.""" return {k[: -len(suffix)] for k in keys if k.endswith(suffix)} @@ -202,6 +193,15 @@ def _block_indices(prefixes: set[str]) -> set[int]: _QWEN_LORA_RANK = 8 +@pytest.mark.parametrize( + "qwen_model", + [ + pytest.param(QwenHfExportModel("fp8", "max", False), marks=minimum_sm(89)), + pytest.param(QwenHfExportModel("fp4", "max", False), marks=minimum_sm(89)), + pytest.param(QwenHfExportModel("fp4", "svdquant", True), marks=minimum_sm(89)), + ], + ids=["qwen_fp8_max", "qwen_nvfp4_max", "qwen_nvfp4_svdquant"], +) def test_qwen_image_hf_ckpt_export( qwen_model: QwenHfExportModel, tiny_qwen_image_path: str, tmp_path: Path ) -> None: @@ -233,8 +233,13 @@ def test_qwen_image_hf_ckpt_export( assert not any("weight_quantizer" in k for k in keys), "quantizer keys leaked into export" assert not any("input_quantizer._amax" in k for k in keys) - # Recipe: only the middle blocks are quantized (first-2/last-2 excluded). + # Recipe: only the middle transformer blocks are quantized — first-2/last-2 of + # transformer_blocks are excluded, and nothing outside transformer_blocks. weight_scale_prefixes = _module_prefixes(keys, ".weight_scale") + assert weight_scale_prefixes, "no quantized linears found in export" + assert all("transformer_blocks." in p for p in weight_scale_prefixes), ( + f"a non-transformer_blocks module was quantized: {weight_scale_prefixes}" + ) assert _block_indices(weight_scale_prefixes) == _QWEN_QUANTIZED_BLOCKS, ( f"expected only blocks {_QWEN_QUANTIZED_BLOCKS} quantized" ) @@ -244,10 +249,9 @@ def test_qwen_image_hf_ckpt_export( b_prefixes = _module_prefixes(keys, ".svdquant_lora_b") pqs_prefixes = _module_prefixes(keys, ".pre_quant_scale") assert a_prefixes, "no promoted svdquant_lora_a keys" - # Every promoted linear carries lora_a, lora_b, and pre_quant_scale. - assert a_prefixes == b_prefixes == pqs_prefixes - # ...and each is a quantized linear, only in the middle blocks. - assert a_prefixes <= weight_scale_prefixes + # Every promoted linear carries lora_a, lora_b, and pre_quant_scale, and + # every quantized linear is promoted (the sets are identical). + assert a_prefixes == b_prefixes == pqs_prefixes == weight_scale_prefixes assert _block_indices(a_prefixes) == _QWEN_QUANTIZED_BLOCKS # Rank-consistent shapes; lora_a=[rank, in], lora_b=[out, rank], rank == --lowrank. for key, tensor in lora_tensors.items(): From 1cfe0b36c9590e3b28f478e364402d3fb7f3a6c4 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 18:48:09 -0700 Subject: [PATCH 08/22] Qwen-Image SVDQuant: fix stale tiny-fixture tokenizer docstring Round 6 (round-5 review found no code blocker; only the queued docstring nit): the create_tiny_qwen_image_pipeline_dir docstring still said the tokenizer was fetched from the Hub, but Round 4 switched it to a local offline build (_build_local_qwen2_tokenizer). Updated the wording to "fully offline". Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- tests/_test_utils/torch/diffusers_models.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/_test_utils/torch/diffusers_models.py b/tests/_test_utils/torch/diffusers_models.py index 71e0e19cf78..6e08518bdce 100644 --- a/tests/_test_utils/torch/diffusers_models.py +++ b/tests/_test_utils/torch/diffusers_models.py @@ -341,16 +341,16 @@ def _build_local_qwen2_tokenizer(out_dir: Path): def create_tiny_qwen_image_pipeline_dir(tmp_path: Path) -> Path: - """Create and save a tiny, (mostly) offline Qwen-Image pipeline to a directory. - - Mirrors diffusers' ``QwenImagePipelineFastTests.get_dummy_components``: the - Qwen2.5-VL text encoder is built inline from a tiny ``Qwen2_5_VLConfig`` (no Hub - model load); only the tokenizer is fetched from the tiny ``Qwen2VL`` test repo - (building a Qwen tokenizer fully offline is impractical). The transformer uses - ``num_layers=6`` so the first-2/last-2 block-range recipe is valid, and its - ``joint_attention_dim`` matches the text encoder ``hidden_size`` (16) so the - pipeline runs end-to-end during quantization calibration. The saved dir loads - with ``QwenImagePipeline.from_pretrained(path)``. + """Create and save a tiny, fully offline Qwen-Image pipeline to a directory. + + Mirrors diffusers' ``QwenImagePipelineFastTests.get_dummy_components`` but with + no Hub access: the Qwen2.5-VL text encoder is built inline from a tiny + ``Qwen2_5_VLConfig``, and the tokenizer is built locally by + ``_build_local_qwen2_tokenizer`` (byte-level vocab written to a temp dir). The + transformer uses ``num_layers=6`` so the first-2/last-2 block-range recipe is + valid, and its ``joint_attention_dim`` matches the text encoder ``hidden_size`` + (16) so the pipeline runs end-to-end during quantization calibration. The saved + dir loads with ``QwenImagePipeline.from_pretrained(path)``. """ if QwenImageTransformer2DModel is None or AutoencoderKLQwenImage is None: pytest.skip("QwenImage diffusers classes not available in this diffusers version.") From 789f4ef0a002d7a08ed2372cb3dd59037d2d8966 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 19:01:23 -0700 Subject: [PATCH 09/22] Qwen-Image SVDQuant: add immutability + negative-loading tests; fix stale docs Round 7 (addresses round-6 Codex review's two missing-coverage items): - AC-2.2 SVDQuant immutability test (test_qwen_block_range_recipe.py): builds a 6-block backbone, snapshots the excluded first/last block linear weights, runs SVDQuant via build_block_range_quant_cfg, and asserts the excluded blocks' weights are bit-identical (never calibrated) with no LoRA, while the middle blocks {2,3} receive LoRA and have their weights modified. - AC-1 negative-loading tests (new test_qwen_pipeline_loading.py): monkeypatch MODEL_PIPELINE[QWEN_IMAGE]=None and assert the actionable ImportError; a fake pipeline asserts create_pipeline does not pass trust_remote_code. Stale-doc cleanups: the resolved pre_quant_scale TODO wording in unified_export_hf.py; the build_block_range_quant_cfg docstring (first+last+1 -> +2); the conftest "SKETCH" wording (the fixture is now a working offline build). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- .../diffusers/quantization/models_utils.py | 3 +- modelopt/torch/export/unified_export_hf.py | 6 +- tests/examples/diffusers/conftest.py | 8 +- .../diffusers/test_qwen_block_range_recipe.py | 64 ++++++++++++++++ .../diffusers/test_qwen_pipeline_loading.py | 74 +++++++++++++++++++ 5 files changed, 149 insertions(+), 6 deletions(-) create mode 100644 tests/examples/diffusers/test_qwen_pipeline_loading.py diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index eb2ff1f6039..86fa0750c42 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -330,7 +330,8 @@ def build_block_range_quant_cfg( Raises: ValueError: if the backbone has no ``block_module`` list, or it has fewer - than ``exclude_first_n + exclude_last_n + 1`` blocks. + than ``exclude_first_n + exclude_last_n + 2`` blocks (it requires at + least two quantized middle blocks). """ blocks = getattr(backbone, block_module, None) if blocks is None or not hasattr(blocks, "__len__"): diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 164f2a0e1fe..c47b12a3fd6 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -973,7 +973,8 @@ def _fuse_qkv_linears_diffusion( Note: This is a simplified version for diffusion models that: - Handles QKV fusion (shared input detection) - Filters to only fuse actual QKV projection layers (not AdaLN, FFN, etc.) - - Skips pre_quant_scale handling (TODO for future) + - Skips pre_quant_scale *fusion* (the export path promotes pre_quant_scale to + module-level keys separately; see _promote_quantizer_tensors_to_module) - Skips FFN fusion with layernorm (TODO for future) Args: @@ -1143,7 +1144,8 @@ def _export_diffusers_checkpoint( if is_quantized: # Fuse QKV linears that share the same input (unify amax values) # This is similar to requantize_resmooth_fused_llm_layers but simplified for diffusion - # TODO: Add pre_quant_scale handling and FFN fusion for AWQ-style quantization + # TODO: Add FFN fusion for AWQ-style quantization (pre_quant_scale is + # promoted to module keys at export by _promote_quantizer_tensors_to_module below) print(f" Running QKV fusion for {component_name}...") # Qwen-Image's packed-latent forward signature is non-standard; if the # dummy forward fails for it, fail loudly rather than silently skipping diff --git a/tests/examples/diffusers/conftest.py b/tests/examples/diffusers/conftest.py index e704f6d5879..625f9bc3415 100644 --- a/tests/examples/diffusers/conftest.py +++ b/tests/examples/diffusers/conftest.py @@ -35,9 +35,11 @@ def tiny_wan22_path(tmp_path_factory): def tiny_qwen_image_path(tmp_path_factory): """Create a tiny Qwen-Image pipeline and return its path (built once per session). - SKETCH fixture for the recipe-level DMD2 e2e (``test_fastgen_recipe_e2e.py``). - See ``create_tiny_qwen_image_pipeline_dir`` for caveats — notably the tiny - Qwen2.5-VL text encoder, which needs in-container validation. + Used by the diffusers Qwen export tests and the recipe-level DMD2 e2e + (``test_fastgen_recipe_e2e.py``). The pipeline is built fully offline by + ``create_tiny_qwen_image_pipeline_dir`` (inline tiny Qwen2.5-VL text encoder + + local byte-level tokenizer); it skips only when the diffusers Qwen classes are + unavailable. """ try: from _test_utils.torch.diffusers_models import create_tiny_qwen_image_pipeline_dir diff --git a/tests/examples/diffusers/test_qwen_block_range_recipe.py b/tests/examples/diffusers/test_qwen_block_range_recipe.py index 8214ec1c124..7e0110afacc 100644 --- a/tests/examples/diffusers/test_qwen_block_range_recipe.py +++ b/tests/examples/diffusers/test_qwen_block_range_recipe.py @@ -93,3 +93,67 @@ class _NoBlocks: with pytest.raises(ValueError, match="transformer_blocks"): build_block_range_quant_cfg(_NoBlocks(), exclude_first_n=2, exclude_last_n=2) + + +def test_svdquant_recipe_leaves_excluded_blocks_bit_identical(): + """AC-2.2: the pre-calibration recipe must keep the excluded first/last blocks + bit-identical through SVDQuant (whose calibration subtracts a residual from + every *enabled* linear), while the middle blocks receive LoRA.""" + import torch + import torch.nn as nn + + import modelopt.torch.quantization as mtq + + class _Block(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.proj = nn.Linear(dim, dim) + + def forward(self, x): + return self.proj(x) + + class _Backbone(nn.Module): + def __init__(self, num_blocks: int = 6, dim: int = 32): + super().__init__() + self.transformer_blocks = nn.ModuleList(_Block(dim) for _ in range(num_blocks)) + + def forward(self, x): + for block in self.transformer_blocks: + x = block(x) + return x + + torch.manual_seed(0) + model = _Backbone(num_blocks=6, dim=32) + weights_before = { + i: model.transformer_blocks[i].proj.weight.detach().clone() for i in range(6) + } + + # Base rules quantize every linear weight/input quantizer; the recipe then + # disables all and re-enables only the middle transformer blocks (2, 3). + quant_cfg = { + "quant_cfg": [ + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + *build_block_range_quant_cfg(model, exclude_first_n=2, exclude_last_n=2), + ], + "algorithm": {"method": "svdquant", "lowrank": 4}, + } + calib_data = [torch.randn(2, 32) for _ in range(2)] + mtq.quantize(model, quant_cfg, lambda m: [m(batch) for batch in calib_data]) + + excluded = {0, 1, 4, 5} + for idx in range(6): + proj = model.transformer_blocks[idx].proj + lora_a = getattr(getattr(proj, "weight_quantizer", None), "svdquant_lora_a", None) + if idx in excluded: + # Never calibrated -> weight bit-identical, no LoRA residual. + assert torch.equal(proj.weight, weights_before[idx]), ( + f"excluded block {idx} weight was modified" + ) + assert lora_a is None, f"excluded block {idx} unexpectedly has SVDQuant LoRA" + else: + # Calibrated -> LoRA present and the residual was subtracted from the weight. + assert lora_a is not None, f"middle block {idx} is missing SVDQuant LoRA" + assert not torch.equal(proj.weight, weights_before[idx]), ( + f"middle block {idx} weight was not modified by SVDQuant" + ) diff --git a/tests/examples/diffusers/test_qwen_pipeline_loading.py b/tests/examples/diffusers/test_qwen_pipeline_loading.py new file mode 100644 index 00000000000..612251633b5 --- /dev/null +++ b/tests/examples/diffusers/test_qwen_pipeline_loading.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Negative-path loading tests for Qwen-Image in the diffusers quantization example. + +These cover the AC-1 negative criteria without a GPU or a real model: +- selecting Qwen-Image when diffusers lacks the Qwen classes raises a clear, + actionable error (not an opaque failure); +- Qwen loading does not pass ``trust_remote_code``. +""" + +import logging +import sys +from pathlib import Path + +import pytest + +pytest.importorskip("diffusers") +pytest.importorskip("torch") + +_EXAMPLE_DIR = Path(__file__).parents[3] / "examples" / "diffusers" / "quantization" +if str(_EXAMPLE_DIR) not in sys.path: + sys.path.insert(0, str(_EXAMPLE_DIR)) + +import models_utils # noqa: E402 +import pipeline_manager # noqa: E402 +from models_utils import ModelType # noqa: E402 +from quantize_config import ModelConfig # noqa: E402 + + +def _qwen_pipeline_manager() -> "pipeline_manager.PipelineManager": + config = ModelConfig(model_type=ModelType.QWEN_IMAGE, backbone=["transformer"]) + return pipeline_manager.PipelineManager(config, logging.getLogger("qwen-loading-test")) + + +def test_missing_qwen_pipeline_raises_actionable_error(monkeypatch): + # Simulate a diffusers version without QwenImagePipeline. + monkeypatch.setitem(models_utils.MODEL_PIPELINE, ModelType.QWEN_IMAGE, None) + manager = _qwen_pipeline_manager() + with pytest.raises(ImportError, match="Qwen-Image requires"): + manager.create_pipeline() + + +def test_qwen_loading_does_not_pass_trust_remote_code(monkeypatch): + captured_kwargs: dict = {} + + class _FakeQwenPipeline: + @classmethod + def from_pretrained(cls, model_id, **kwargs): + captured_kwargs.update(kwargs) + return cls() + + def set_progress_bar_config(self, **kwargs): + pass + + monkeypatch.setitem(models_utils.MODEL_PIPELINE, ModelType.QWEN_IMAGE, _FakeQwenPipeline) + manager = _qwen_pipeline_manager() + manager.create_pipeline() + + # Qwen-Image must load without trust_remote_code. + assert captured_kwargs.get("trust_remote_code") is not True + assert "trust_remote_code" not in captured_kwargs From 521cda09585a540ec5c3659c0c882acb91a13a77 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 19:13:12 -0700 Subject: [PATCH 10/22] Qwen-Image SVDQuant: drop invalid txt_seq_lens dummy kwarg + signature-gate Round 8 (addresses round-7 Codex review, which verified against the diffusers source that QwenImageTransformer2DModel.forward has no txt_seq_lens parameter): - _qwen_inputs no longer passes txt_seq_lens (the real forward signature is hidden_states, encoder_hidden_states, encoder_hidden_states_mask, timestep, img_shapes, guidance, return_dict). Passing txt_seq_lens would have raised an unexpected-keyword error and, because Qwen export uses strict QKV fusion, hard-failed the export. - Signature-gate the dummy inputs: filter to the kwargs the installed model's forward actually accepts (via inspect.signature), so diffusers-version drift cannot hard-fail strict fusion either. - Update test_diffusers_qwen_export.py: no longer require txt_seq_lens. - Remove AC- plan terminology from two test docstrings (code-style note). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- modelopt/torch/export/diffusers_utils.py | 14 ++++++++++++-- .../diffusers/test_qwen_block_range_recipe.py | 6 +++--- .../diffusers/test_qwen_pipeline_loading.py | 2 +- .../torch/export/test_diffusers_qwen_export.py | 1 - 4 files changed, 16 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index 823f26c8e5d..fedfe98e723 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -330,7 +330,7 @@ def _qwen_inputs() -> dict[str, Any]: # QwenImageTransformer2DModel does NOT take the standard # (hidden_states[B,C,H,W], timestep, encoder_hidden_states) triple. It expects # *packed* latents [B, (H//2)*(W//2), in_channels] plus encoder_hidden_states, - # encoder_hidden_states_mask, img_shapes, txt_seq_lens, and optional guidance. + # encoder_hidden_states_mask, img_shapes, and optional guidance. # Timesteps are continuous in [0, 1] (not the diffusers [0, 1000] scale). in_channels = getattr(cfg, "in_channels", 64) joint_attention_dim = getattr(cfg, "joint_attention_dim", 3584) @@ -353,11 +353,21 @@ def _qwen_inputs() -> dict[str, Any]: ), "timestep": torch.tensor([0.5], device=device, dtype=dtype).expand(batch_size), "img_shapes": [[(1, packed_h, packed_w)]] * batch_size, - "txt_seq_lens": [text_seq_len] * batch_size, "return_dict": False, } if guidance_embeds: dummy_inputs["guidance"] = torch.tensor([4.0], device=device, dtype=torch.float32) + + # Only pass kwargs the installed QwenImageTransformer2DModel.forward accepts + # (signatures vary across diffusers versions); prevents the strict QKV-fusion + # dummy forward from failing on an unexpected keyword argument. + import inspect + + try: + accepted = set(inspect.signature(model.forward).parameters) + dummy_inputs = {k: v for k, v in dummy_inputs.items() if k in accepted} + except (TypeError, ValueError): + pass return dummy_inputs def _generic_transformer_inputs() -> dict[str, torch.Tensor] | None: diff --git a/tests/examples/diffusers/test_qwen_block_range_recipe.py b/tests/examples/diffusers/test_qwen_block_range_recipe.py index 7e0110afacc..d1fdf3b7c65 100644 --- a/tests/examples/diffusers/test_qwen_block_range_recipe.py +++ b/tests/examples/diffusers/test_qwen_block_range_recipe.py @@ -96,9 +96,9 @@ class _NoBlocks: def test_svdquant_recipe_leaves_excluded_blocks_bit_identical(): - """AC-2.2: the pre-calibration recipe must keep the excluded first/last blocks - bit-identical through SVDQuant (whose calibration subtracts a residual from - every *enabled* linear), while the middle blocks receive LoRA.""" + """The block-range recipe must keep the excluded first/last blocks bit-identical + through SVDQuant (whose calibration subtracts a residual from every *enabled* + linear), while the middle blocks receive LoRA.""" import torch import torch.nn as nn diff --git a/tests/examples/diffusers/test_qwen_pipeline_loading.py b/tests/examples/diffusers/test_qwen_pipeline_loading.py index 612251633b5..cbc49aa5ea3 100644 --- a/tests/examples/diffusers/test_qwen_pipeline_loading.py +++ b/tests/examples/diffusers/test_qwen_pipeline_loading.py @@ -15,7 +15,7 @@ """Negative-path loading tests for Qwen-Image in the diffusers quantization example. -These cover the AC-1 negative criteria without a GPU or a real model: +These cover the negative loading paths without a GPU or a real model: - selecting Qwen-Image when diffusers lacks the Qwen classes raises a clear, actionable error (not an opaque failure); - Qwen loading does not pass ``trust_remote_code``. diff --git a/tests/unit/torch/export/test_diffusers_qwen_export.py b/tests/unit/torch/export/test_diffusers_qwen_export.py index 6577a8800f0..ad4b62b048f 100644 --- a/tests/unit/torch/export/test_diffusers_qwen_export.py +++ b/tests/unit/torch/export/test_diffusers_qwen_export.py @@ -79,7 +79,6 @@ def test_qwen_dummy_inputs_drive_real_transformer_forward(): "encoder_hidden_states", "encoder_hidden_states_mask", "img_shapes", - "txt_seq_lens", ): assert key in inputs, f"missing Qwen dummy input '{key}'" assert inputs["hidden_states"].shape[-1] == transformer.config.in_channels From d776ffb5ec079aa89298283350c0ea8d1e324891 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 19:27:54 -0700 Subject: [PATCH 11/22] Qwen-Image SVDQuant: remove promoted export buffers from live module after export Round 9 (clears the last queued code item from Codex; no code blockers remain): _promote_quantizer_tensors_to_module left the temporary .svdquant_lora_a/b + .pre_quant_scale buffers on the live module after export. Add _remove_promoted_quantizer_tensors and call it after each quantized diffusers component is saved, so the live module is unchanged post-export (repeated export / module reuse stay correct). The quantizer-owned tensors are untouched. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- modelopt/torch/export/unified_export_hf.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index c47b12a3fd6..abbea9fc37e 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1073,6 +1073,21 @@ def _promote_quantizer_tensors_to_module(component: nn.Module) -> None: sub_module.register_buffer("svdquant_lora_b", lora_b.detach().clone()) +def _remove_promoted_quantizer_tensors(component: nn.Module) -> None: + """Undo :func:`_promote_quantizer_tensors_to_module`. + + Removes the temporary module-level export buffers (``svdquant_lora_a/b`` and + ``pre_quant_scale``) so the live module is unchanged after export, keeping + repeated export / post-export module reuse correct. The quantizer-owned tensors + (``weight_quantizer.svdquant_lora_a/b``, ``input_quantizer._pre_quant_scale``) + are left untouched. + """ + for _, sub_module in component.named_modules(): + for buffer_name in ("svdquant_lora_a", "svdquant_lora_b", "pre_quant_scale"): + if buffer_name in getattr(sub_module, "_buffers", {}): + del sub_module._buffers[buffer_name] + + def _export_diffusers_checkpoint( pipe: Any, dtype: torch.dtype | None, @@ -1200,6 +1215,10 @@ def _export_diffusers_checkpoint( config_data["quantization_config"] = hf_quant_config with open(config_path, "w") as file: json.dump(config_data, file, indent=4) + + # Drop the temporary promoted export buffers so the live module is + # unchanged after export (supports repeated export / module reuse). + _remove_promoted_quantizer_tensors(component) # Non-quantized component: just save as-is elif hasattr(component, "save_pretrained"): component.save_pretrained(component_export_dir, max_shard_size=max_shard_size) From c2250cb5b5b0fa89f56d353cbf229ee8030177fc Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 11 Jun 2026 23:01:12 -0700 Subject: [PATCH 12/22] Qwen-Image diffusers PTQ: fix 5 execution-surfaced bugs (fp8/nvfp4/svdquant) Validated end-to-end on GB200 against the real Qwen/Qwen-Image: all three formats export correct HF checkpoints (only transformer_blocks 2..57; nothing outside), no quantizer-state leak, and the focused tests pass. - models_utils: build_block_range_quant_cfg now uses the top-level enable QuantizerCfgEntry field (a None cfg retains the base preset's params) instead of nesting cfg.enable, which the QuantizerAttributeConfig validator rejects/mis-applies (the old form left every block quantized). - quantize.py: import onnx_utils.export lazily (only needed for --onnx-dir; avoids a hard onnx_graphsurgeon dependency), and pass max_shard_size so the ~20B transformer saves as a single safetensors -- the unified export's layerwise-metadata post-processing does not support sharded files. - diffusers_utils: hide_quantizers_from_state_dict strips quantizer submodules from all modules, not only is_quantlinear, so enabled input quantizers on norm layers no longer leak input_quantizer._amax into the checkpoint. - tests: the tiny QwenImageTransformer2DModel fixture signature-gates its kwargs (diffusers 0.38 removed pooled_projection_dim from the constructor); the recipe test asserts the corrected top-level enable schema. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- .../diffusers/quantization/models_utils.py | 23 ++++++++++--------- examples/diffusers/quantization/quantize.py | 10 +++++++- modelopt/torch/export/diffusers_utils.py | 23 +++++++++++-------- tests/_test_utils/torch/diffusers_models.py | 8 +++++++ .../diffusers/test_qwen_block_range_recipe.py | 13 ++++++----- 5 files changed, 50 insertions(+), 27 deletions(-) diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index 86fa0750c42..1126a421390 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -325,7 +325,8 @@ def build_block_range_quant_cfg( Rules are applied in order with later rules overriding earlier ones: 1. disable every linear weight/input quantizer, - 2. re-enable only those under ``block_module``, + 2. re-enable only those under ``block_module`` (``enable`` is a top-level + QuantizerCfgEntry toggle; a ``None`` cfg keeps the base preset's quant params), 3. disable the first/last ``n`` blocks. Raises: @@ -354,17 +355,17 @@ def build_block_range_quant_cfg( excluded = sorted( set(range(exclude_first_n)) | set(range(num_blocks - exclude_last_n, num_blocks)) ) + # `enable` is a top-level QuantizerCfgEntry field (independent of `cfg`); a `None` + # cfg leaves the base preset's quant params untouched, so disabling then + # re-enabling restores the original (FP8/NVFP4/...) attributes. Putting `enable` + # under `cfg` is rejected by the QuantizerAttributeConfig validator. rules: list[dict[str, Any]] = [ - {"quantizer_name": "*weight_quantizer", "cfg": {"enable": False}}, - {"quantizer_name": "*input_quantizer", "cfg": {"enable": False}}, - {"quantizer_name": f"*{block_module}.*weight_quantizer", "cfg": {"enable": True}}, - {"quantizer_name": f"*{block_module}.*input_quantizer", "cfg": {"enable": True}}, + {"quantizer_name": "*weight_quantizer", "enable": False}, + {"quantizer_name": "*input_quantizer", "enable": False}, + {"quantizer_name": f"*{block_module}.*weight_quantizer", "enable": True}, + {"quantizer_name": f"*{block_module}.*input_quantizer", "enable": True}, ] for idx in excluded: - rules.append( - {"quantizer_name": f"*{block_module}.{idx}.*weight_quantizer", "cfg": {"enable": False}} - ) - rules.append( - {"quantizer_name": f"*{block_module}.{idx}.*input_quantizer", "cfg": {"enable": False}} - ) + rules.append({"quantizer_name": f"*{block_module}.{idx}.*weight_quantizer", "enable": False}) + rules.append({"quantizer_name": f"*{block_module}.{idx}.*input_quantizer", "enable": False}) return rules diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index e81dc51a66b..325d39e5bd7 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -39,7 +39,6 @@ get_model_filter_func, parse_extra_params, ) -from onnx_utils.export import generate_fp8_scales, modelopt_export_sd from pipeline_manager import PipelineManager from quantize_config import ( CalibrationConfig, @@ -319,6 +318,11 @@ def export_onnx( if not self.config.onnx_dir: return + # onnx_graphsurgeon (pulled in by onnx_utils.export) is an optional dependency + # only needed for the ONNX export path; import lazily so the HF-checkpoint + # export runs without it installed. + from onnx_utils.export import generate_fp8_scales, modelopt_export_sd + self.logger.info(f"Starting ONNX export to {self.config.onnx_dir}") if quant_format == QuantFormat.FP8 and self._has_conv_layers(backbone): @@ -407,6 +411,10 @@ def export_hf_ckpt(self, pipe: Any, model_config: ModelConfig | None = None) -> f"Invalid padding_strategy: {padding!r}. Expected 'row' or 'row_col'." ) kwargs["padding_strategy"] = padding + # The diffusion transformer is large (~20B params); the unified export's + # layerwise-metadata post-processing does not support sharded safetensors, so + # save each component as a single file (no *.safetensors.index.json). + kwargs.setdefault("max_shard_size", "200GB") export_hf_checkpoint(pipe, export_dir=self.config.hf_ckpt_dir, **kwargs) self.logger.info("HuggingFace checkpoint export completed successfully") diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index fedfe98e723..075f25a9101 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -735,15 +735,20 @@ def hide_quantizers_from_state_dict(model: nn.Module): # Store references to quantizers that we'll temporarily remove quantizer_backup: dict[str, dict[str, nn.Module]] = {} - for name, module in model.named_modules(): - if is_quantlinear(module): - backup = {} - for attr in ["weight_quantizer", "input_quantizer", "output_quantizer"]: - if hasattr(module, attr): - backup[attr] = getattr(module, attr) - delattr(module, attr) - if backup: - quantizer_backup[name] = backup + # Remove every quantizer submodule from *all* modules, not only recognized + # quant-linears: enabled input quantizers can also live on non-linear modules + # (e.g. norm layers whose activations were calibrated), and their ``_amax`` + # buffers must not leak into the saved checkpoint. Snapshot the module list + # first since we mutate the module tree while iterating. + for name, module in list(model.named_modules()): + backup = {} + for attr in ["weight_quantizer", "input_quantizer", "output_quantizer"]: + child = getattr(module, attr, None) + if isinstance(child, nn.Module): + backup[attr] = child + delattr(module, attr) + if backup: + quantizer_backup[name] = backup try: yield diff --git a/tests/_test_utils/torch/diffusers_models.py b/tests/_test_utils/torch/diffusers_models.py index 6e08518bdce..caf4966399f 100644 --- a/tests/_test_utils/torch/diffusers_models.py +++ b/tests/_test_utils/torch/diffusers_models.py @@ -289,6 +289,14 @@ def get_tiny_qwen_image_transformer(**config_kwargs): "axes_dims_rope": (8, 4, 4), # sums to attention_head_dim (16) } kwargs.update(**config_kwargs) + # Drop kwargs the installed diffusers QwenImageTransformer2DModel doesn't accept. + # `pooled_projection_dim` is present in the published config.json but was removed + # from the constructor in newer diffusers: from_pretrained tolerates the extra + # config key, but a direct constructor call raises TypeError. + import inspect + + accepted = set(inspect.signature(QwenImageTransformer2DModel.__init__).parameters) + kwargs = {k: v for k, v in kwargs.items() if k in accepted} return QwenImageTransformer2DModel(**kwargs) diff --git a/tests/examples/diffusers/test_qwen_block_range_recipe.py b/tests/examples/diffusers/test_qwen_block_range_recipe.py index d1fdf3b7c65..57cd5619591 100644 --- a/tests/examples/diffusers/test_qwen_block_range_recipe.py +++ b/tests/examples/diffusers/test_qwen_block_range_recipe.py @@ -51,7 +51,7 @@ def _disabled_block_indices(rules): """Indices of transformer blocks explicitly disabled by per-block rules.""" indices = set() for rule in rules: - if rule["cfg"].get("enable") is False: + if rule.get("enable") is False: match = _BLOCK_RULE_RE.fullmatch(rule["quantizer_name"]) if match: indices.add(int(match.group(1))) @@ -62,11 +62,12 @@ def test_recipe_excludes_first_and_last_two_blocks(): rules = build_block_range_quant_cfg(_StubBackbone(6), exclude_first_n=2, exclude_last_n=2) # 1. disable-all rules come first (weight + input). - assert rules[0] == {"quantizer_name": "*weight_quantizer", "cfg": {"enable": False}} - assert rules[1] == {"quantizer_name": "*input_quantizer", "cfg": {"enable": False}} - # 2. then enable only the transformer_blocks. - assert {"quantizer_name": "*transformer_blocks.*weight_quantizer", "cfg": {"enable": True}} in rules - assert {"quantizer_name": "*transformer_blocks.*input_quantizer", "cfg": {"enable": True}} in rules + assert rules[0] == {"quantizer_name": "*weight_quantizer", "enable": False} + assert rules[1] == {"quantizer_name": "*input_quantizer", "enable": False} + # 2. then re-enable only the transformer_blocks (top-level `enable`; a `None` cfg + # keeps the base preset's quant params). + assert {"quantizer_name": "*transformer_blocks.*weight_quantizer", "enable": True} in rules + assert {"quantizer_name": "*transformer_blocks.*input_quantizer", "enable": True} in rules # 3. then disable the first 2 and last 2 of the 6 blocks -> {0, 1, 4, 5}; quantize {2, 3}. assert _disabled_block_indices(rules) == {0, 1, 4, 5} From fb23155a7b1078802b17599d539f1db61c0b46f9 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Fri, 12 Jun 2026 16:21:44 -0700 Subject: [PATCH 13/22] Qwen-Image: drop operator-specific quantization harness from the example run_qwen_image_quantization.sh and its README are cluster-specific experiment/operator scripts (hard-coded /lustre paths) that do not belong in the upstream diffusers example. The feature itself (model registration, block-range recipe, FP8/NVFP4/SVDQuant export) is covered by the committed tests. The scripts are kept locally outside the repo. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- .../qwen_image_svdquant/README.md | 108 ------------------ .../run_qwen_image_quantization.sh | 101 ---------------- 2 files changed, 209 deletions(-) delete mode 100644 examples/diffusers/quantization/qwen_image_svdquant/README.md delete mode 100755 examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh diff --git a/examples/diffusers/quantization/qwen_image_svdquant/README.md b/examples/diffusers/quantization/qwen_image_svdquant/README.md deleted file mode 100644 index 02daf02ac8c..00000000000 --- a/examples/diffusers/quantization/qwen_image_svdquant/README.md +++ /dev/null @@ -1,108 +0,0 @@ -# Qwen-Image Quantization (FP8 / NVFP4 / NVFP4-SVDQuant) - -A reproducible harness for quantizing [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) -with the diffusers quantization example and exporting HuggingFace checkpoints. - -## What it does - -- Registers Qwen-Image in the diffusers quantization example (`--model qwen-image`). -- **Recipe**: quantizes only the linears under `transformer_blocks`, keeping the - **first 2 and last 2** of the 60 blocks (and everything outside - `transformer_blocks`: text encoder, VAE, embedders, norms, `proj_out`, …) in - original precision. The exclusion is applied **before calibration** so that for - SVDQuant the excluded blocks' weights stay bit-identical to the original. -- Produces three checkpoints: **FP8**, **NVFP4** (max), and **NVFP4 + SVDQuant**. -- Exports a HuggingFace unified checkpoint per component (safetensors + `config.json`). - -### SVDQuant checkpoint format (AWQ-aligned) - -For the SVDQuant export, the quantizer-owned tensors are promoted to clean, -module-level safetensors keys (mirroring how AWQ exports `pre_quant_scale`): - -| Tensor | Safetensors key | -|--------|-----------------| -| AWQ smoothing scale (`input_quantizer._pre_quant_scale`) | `.pre_quant_scale` | -| Low-rank factor A (`weight_quantizer.svdquant_lora_a`) | `.svdquant_lora_a` | -| Low-rank factor B (`weight_quantizer.svdquant_lora_b`) | `.svdquant_lora_b` | - -They are embedded in the component's main safetensors (no sidecar). The -`config.json`'s `quantization_config` follows the `nvfp4_awq` shape with -`"pre_quant_scale": true` plus the SVDQuant `lora_rank`, so a consumer can -reconstruct `y = NVFP4_GEMM(x) + (x @ lora_a^T) @ lora_b^T`. (No in-repo runtime -applies this residual yet; the checkpoint is a documented on-disk artifact.) - -## Layout (kernel-dev defaults) - -| Env var | Default | Purpose | -|---------|---------|---------| -| `KERNEL_DEV_ROOT` | `/lustre/fsw/coreai_dlalgo_modelopt/users/jingyux/kernel-dev` | Root for container/models/output | -| `MODEL_DIR` | `${KERNEL_DEV_ROOT}/models/Qwen-Image` | Local model cache | -| `OUTPUT_DIR` | `${KERNEL_DEV_ROOT}/qwen_image_ckpts` | Exported checkpoints | -| `HF_TOKEN_FILE` | `${KERNEL_DEV_ROOT}/HF_TOKEN.txt` | Hugging Face token file | -| `FORMATS` | `fp8 nvfp4 svdquant` | Formats to run | -| `CALIB_SIZE` / `BATCH_SIZE` / `N_STEPS` / `LOWRANK` | `64 / 2 / 20 / 32` | Calibration knobs | - -## 1. Build the container (once) - -The diffusers example needs a recent `diffusers` (with `QwenImagePipeline`) and -modelopt installed from source. From a base NGC PyTorch image: - -```bash -CONTAINER_DIR=/lustre/fsw/coreai_dlalgo_modelopt/users/jingyux/kernel-dev/container -mkdir -p "${CONTAINER_DIR}" - -# Import a base image to an enroot squashfs (adjust the tag as needed). -enroot import -o "${CONTAINER_DIR}/modelopt-diffusers.sqsh" \ - docker://nvcr.io#nvidia/pytorch:25.04-py3 - -# Install modelopt (from source) + example deps into the container, then re-save. -srun --container-image="${CONTAINER_DIR}/modelopt-diffusers.sqsh" \ - --container-mounts=/lustre:/lustre --container-save="${CONTAINER_DIR}/modelopt-diffusers.sqsh" \ - bash -lc ' - cd /lustre/fsw/coreai_dlalgo_modelopt/users/jingyux/kernel-dev/source/Model-Optimizer && - pip install -e ".[dev]" && - pip install -U "diffusers>=0.35" "transformers>=4.52" accelerate datasets && - python -c "from diffusers import QwenImagePipeline; print(\"QwenImagePipeline OK\")" - ' -``` - -## 2. Run quantization - -Inside the container (or via `srun`), run the harness: - -```bash -srun --gpus=1 \ - --container-image=/lustre/fsw/coreai_dlalgo_modelopt/users/jingyux/kernel-dev/container/modelopt-diffusers.sqsh \ - --container-mounts=/lustre:/lustre \ - bash examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh -``` - -This downloads `Qwen/Qwen-Image` to `MODEL_DIR` (idempotent), then for each -format writes `${OUTPUT_DIR}/qwen-image-/` (HF checkpoint + `sanity.png`). - -Run a single format, or preview the commands without executing: - -```bash -FORMATS=svdquant LOWRANK=32 bash .../run_qwen_image_quantization.sh -DRY_RUN=1 bash .../run_qwen_image_quantization.sh # print planned commands only -``` - -The equivalent direct `quantize.py` invocation for SVDQuant: - -```bash -python examples/diffusers/quantization/quantize.py \ - --model qwen-image --override-model-path "${MODEL_DIR}" --model-dtype BFloat16 \ - --format fp4 --quant-algo svdquant --lowrank 32 \ - --calib-size 64 --batch-size 2 --n-steps 20 \ - --hf-ckpt-dir "${OUTPUT_DIR}/qwen-image-svdquant" \ - --sanity-image-path "${OUTPUT_DIR}/qwen-image-svdquant/sanity.png" -``` - -## Notes - -- `Qwen/Qwen-Image` loads without `trust_remote_code`. -- The transformer is ~20B params; calibration needs a GPU with enough memory - (use `--cpu-offloading` if VRAM-limited). -- The `--sanity-image-path` image is generated from the **in-memory** quantized - pipeline before the weights are packed for export (a functional check of - quantized inference; it does not reload the exported checkpoint). diff --git a/examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh b/examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh deleted file mode 100755 index 5e0569c0e40..00000000000 --- a/examples/diffusers/quantization/qwen_image_svdquant/run_qwen_image_quantization.sh +++ /dev/null @@ -1,101 +0,0 @@ -#!/usr/bin/env bash -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Reproducible Qwen-Image quantization (FP8 / NVFP4 / NVFP4-SVDQuant) using the -# diffusers quantization example. This script is meant to run INSIDE a container -# that already has NVIDIA Model Optimizer installed from source and a -# Qwen-capable diffusers (see README.md for building the container and the -# Slurm/srun wrapper). -# -# It downloads Qwen/Qwen-Image (idempotently), then for each requested format -# runs `quantize.py` to calibrate the transformer (only `transformer_blocks`, -# excluding the first 2 / last 2 blocks), generate a quantized-inference sanity -# image, and export a HuggingFace checkpoint. -# -# All paths are parameterized via environment variables; the defaults match the -# kernel-dev experiment layout described in README.md. -set -euo pipefail - -# --- Configuration (override via environment) -------------------------------- -KERNEL_DEV_ROOT="${KERNEL_DEV_ROOT:-/lustre/fsw/coreai_dlalgo_modelopt/users/jingyux/kernel-dev}" -MODEL_ID="${MODEL_ID:-Qwen/Qwen-Image}" -MODEL_DIR="${MODEL_DIR:-${KERNEL_DEV_ROOT}/models/Qwen-Image}" -OUTPUT_DIR="${OUTPUT_DIR:-${KERNEL_DEV_ROOT}/qwen_image_ckpts}" -HF_TOKEN_FILE="${HF_TOKEN_FILE:-${KERNEL_DEV_ROOT}/HF_TOKEN.txt}" -# Path to the diffusers quantization example (this script lives one level below it). -QUANT_DIR="${QUANT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}" - -# Formats to run: any of {fp8, nvfp4, svdquant}. -FORMATS="${FORMATS:-fp8 nvfp4 svdquant}" - -# Calibration knobs (small defaults for a quick run; raise CALIB_SIZE for quality). -CALIB_SIZE="${CALIB_SIZE:-64}" -BATCH_SIZE="${BATCH_SIZE:-2}" -N_STEPS="${N_STEPS:-20}" -LOWRANK="${LOWRANK:-32}" -MODEL_DTYPE="${MODEL_DTYPE:-BFloat16}" - -# Set DRY_RUN=1 to print the planned commands without executing them. -DRY_RUN="${DRY_RUN:-0}" - -log() { echo "[qwen-image-quant] $*"; } -run() { - log "+ $*" - if [[ "${DRY_RUN}" != "1" ]]; then - "$@" - fi -} - -# --- Hugging Face token ------------------------------------------------------ -if [[ ! -r "${HF_TOKEN_FILE}" ]]; then - echo "ERROR: HF token file not found or not readable: ${HF_TOKEN_FILE}" >&2 - echo " Set HF_TOKEN_FILE to a readable file containing your Hugging Face token." >&2 - exit 1 -fi -HF_TOKEN="$(tr -d '[:space:]' < "${HF_TOKEN_FILE}")" -if [[ -z "${HF_TOKEN}" ]]; then - echo "ERROR: HF token file is empty: ${HF_TOKEN_FILE}" >&2 - exit 1 -fi -export HF_TOKEN -export HUGGING_FACE_HUB_TOKEN="${HF_TOKEN}" - -# --- Download the model (idempotent) ---------------------------------------- -log "Downloading ${MODEL_ID} -> ${MODEL_DIR} (skipped if already present)" -run mkdir -p "${MODEL_DIR}" -run huggingface-cli download "${MODEL_ID}" --local-dir "${MODEL_DIR}" --exclude "*.onnx" - -# --- Quantize + export for each format -------------------------------------- -mkdir -p "${OUTPUT_DIR}" -for fmt in ${FORMATS}; do - case "${fmt}" in - fp8) quant_args=(--format fp8 --quant-algo max) ;; - nvfp4) quant_args=(--format fp4 --quant-algo max) ;; - svdquant) quant_args=(--format fp4 --quant-algo svdquant --lowrank "${LOWRANK}") ;; - *) echo "ERROR: unknown format '${fmt}' (expected fp8|nvfp4|svdquant)" >&2; exit 1 ;; - esac - - out="${OUTPUT_DIR}/qwen-image-${fmt}" - log "=== Quantizing Qwen-Image (${fmt}) -> ${out} ===" - run python "${QUANT_DIR}/quantize.py" \ - --model qwen-image \ - --override-model-path "${MODEL_DIR}" \ - --model-dtype "${MODEL_DTYPE}" \ - "${quant_args[@]}" \ - --calib-size "${CALIB_SIZE}" \ - --batch-size "${BATCH_SIZE}" \ - --n-steps "${N_STEPS}" \ - --hf-ckpt-dir "${out}" \ - --sanity-image-path "${out}/sanity.png" - - # Verify the expected artifacts were produced (a missing artifact is a failure). - if [[ "${DRY_RUN}" != "1" ]]; then - [[ -f "${out}/sanity.png" ]] || { echo "ERROR: missing sanity image ${out}/sanity.png" >&2; exit 1; } - find "${out}" -name '*.safetensors' | grep -q . || { echo "ERROR: no safetensors under ${out}" >&2; exit 1; } - find "${out}" -name 'config.json' | grep -q . || { echo "ERROR: no config.json under ${out}" >&2; exit 1; } - fi - log "Done: ${fmt}. Checkpoint at ${out}, sanity image at ${out}/sanity.png" -done - -log "All requested formats complete. Checkpoints under ${OUTPUT_DIR}" From 80e22443a308a0d10956615d9c80d1661331dab0 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Fri, 12 Jun 2026 16:32:29 -0700 Subject: [PATCH 14/22] Qwen-Image: consolidate tests into the shared diffusers export suite Remove the standalone Qwen test files. The fp8/nvfp4/svdquant cases in test_export_diffusers_hf_ckpt.py already cover the block-range recipe (only transformer_blocks 2..57 quantized), the promoted SVDQuant keys + pre_quant_scale, the NVFP4_SVD quantization_config, and the no-leak check -- matching how SDXL/Flux/Wan are tested in the same file. Core SVDQuant forward/fold is unchanged and remains covered by existing upstream tests. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- .../diffusers/test_qwen_block_range_recipe.py | 160 ------------------ .../diffusers/test_qwen_pipeline_loading.py | 74 -------- .../export/test_convert_hf_config_svdquant.py | 87 ---------- .../export/test_diffusers_qwen_export.py | 133 --------------- .../test_svdquant_forward_fold.py | 116 ------------- 5 files changed, 570 deletions(-) delete mode 100644 tests/examples/diffusers/test_qwen_block_range_recipe.py delete mode 100644 tests/examples/diffusers/test_qwen_pipeline_loading.py delete mode 100644 tests/unit/torch/export/test_convert_hf_config_svdquant.py delete mode 100644 tests/unit/torch/export/test_diffusers_qwen_export.py delete mode 100644 tests/unit/torch/quantization/test_svdquant_forward_fold.py diff --git a/tests/examples/diffusers/test_qwen_block_range_recipe.py b/tests/examples/diffusers/test_qwen_block_range_recipe.py deleted file mode 100644 index 57cd5619591..00000000000 --- a/tests/examples/diffusers/test_qwen_block_range_recipe.py +++ /dev/null @@ -1,160 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - -"""Unit tests for the transformer-block-range quantization recipe (e.g. Qwen-Image). - -The recipe must quantize only the linears under ``transformer_blocks`` while -excluding the first/last N blocks, and it must be expressible as ``quant_cfg`` -rules applied before calibration (so SVDQuant never mutates the excluded blocks). -""" - -import re -import sys -from pathlib import Path - -import pytest - -# Importing the example module pulls in diffusers/torch/datasets/modelopt. -pytest.importorskip("diffusers") -pytest.importorskip("torch") - -# Make the diffusers quantization example importable. -_EXAMPLE_DIR = Path(__file__).parents[3] / "examples" / "diffusers" / "quantization" -if str(_EXAMPLE_DIR) not in sys.path: - sys.path.insert(0, str(_EXAMPLE_DIR)) - -from models_utils import build_block_range_quant_cfg # noqa: E402 - -_BLOCK_RULE_RE = re.compile(r"\*transformer_blocks\.(\d+)\.\*(?:weight|input)_quantizer") - - -class _StubBackbone: - """Minimal stand-in exposing a ``transformer_blocks`` sequence of length n.""" - - def __init__(self, num_blocks: int): - self.transformer_blocks = list(range(num_blocks)) - - -def _disabled_block_indices(rules): - """Indices of transformer blocks explicitly disabled by per-block rules.""" - indices = set() - for rule in rules: - if rule.get("enable") is False: - match = _BLOCK_RULE_RE.fullmatch(rule["quantizer_name"]) - if match: - indices.add(int(match.group(1))) - return indices - - -def test_recipe_excludes_first_and_last_two_blocks(): - rules = build_block_range_quant_cfg(_StubBackbone(6), exclude_first_n=2, exclude_last_n=2) - - # 1. disable-all rules come first (weight + input). - assert rules[0] == {"quantizer_name": "*weight_quantizer", "enable": False} - assert rules[1] == {"quantizer_name": "*input_quantizer", "enable": False} - # 2. then re-enable only the transformer_blocks (top-level `enable`; a `None` cfg - # keeps the base preset's quant params). - assert {"quantizer_name": "*transformer_blocks.*weight_quantizer", "enable": True} in rules - assert {"quantizer_name": "*transformer_blocks.*input_quantizer", "enable": True} in rules - # 3. then disable the first 2 and last 2 of the 6 blocks -> {0, 1, 4, 5}; quantize {2, 3}. - assert _disabled_block_indices(rules) == {0, 1, 4, 5} - - -def test_recipe_block_count_scales_with_model(): - # For a 60-block model (Qwen-Image), exclude {0, 1, 58, 59}; quantize 2..57. - rules = build_block_range_quant_cfg(_StubBackbone(60), exclude_first_n=2, exclude_last_n=2) - assert _disabled_block_indices(rules) == {0, 1, 58, 59} - - -@pytest.mark.parametrize("num_blocks", [5, 4, 3]) -def test_recipe_rejects_too_few_blocks(num_blocks): - # A 2 + 2 exclusion needs at least 6 blocks (>= 2 quantized middle blocks). - # A 5-block model leaves only 1 middle block and must be rejected too. - with pytest.raises(ValueError, match="at least"): - build_block_range_quant_cfg( - _StubBackbone(num_blocks), exclude_first_n=2, exclude_last_n=2 - ) - - -def test_recipe_missing_block_module_raises(): - class _NoBlocks: - pass - - with pytest.raises(ValueError, match="transformer_blocks"): - build_block_range_quant_cfg(_NoBlocks(), exclude_first_n=2, exclude_last_n=2) - - -def test_svdquant_recipe_leaves_excluded_blocks_bit_identical(): - """The block-range recipe must keep the excluded first/last blocks bit-identical - through SVDQuant (whose calibration subtracts a residual from every *enabled* - linear), while the middle blocks receive LoRA.""" - import torch - import torch.nn as nn - - import modelopt.torch.quantization as mtq - - class _Block(nn.Module): - def __init__(self, dim: int): - super().__init__() - self.proj = nn.Linear(dim, dim) - - def forward(self, x): - return self.proj(x) - - class _Backbone(nn.Module): - def __init__(self, num_blocks: int = 6, dim: int = 32): - super().__init__() - self.transformer_blocks = nn.ModuleList(_Block(dim) for _ in range(num_blocks)) - - def forward(self, x): - for block in self.transformer_blocks: - x = block(x) - return x - - torch.manual_seed(0) - model = _Backbone(num_blocks=6, dim=32) - weights_before = { - i: model.transformer_blocks[i].proj.weight.detach().clone() for i in range(6) - } - - # Base rules quantize every linear weight/input quantizer; the recipe then - # disables all and re-enables only the middle transformer blocks (2, 3). - quant_cfg = { - "quant_cfg": [ - {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, - {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, - *build_block_range_quant_cfg(model, exclude_first_n=2, exclude_last_n=2), - ], - "algorithm": {"method": "svdquant", "lowrank": 4}, - } - calib_data = [torch.randn(2, 32) for _ in range(2)] - mtq.quantize(model, quant_cfg, lambda m: [m(batch) for batch in calib_data]) - - excluded = {0, 1, 4, 5} - for idx in range(6): - proj = model.transformer_blocks[idx].proj - lora_a = getattr(getattr(proj, "weight_quantizer", None), "svdquant_lora_a", None) - if idx in excluded: - # Never calibrated -> weight bit-identical, no LoRA residual. - assert torch.equal(proj.weight, weights_before[idx]), ( - f"excluded block {idx} weight was modified" - ) - assert lora_a is None, f"excluded block {idx} unexpectedly has SVDQuant LoRA" - else: - # Calibrated -> LoRA present and the residual was subtracted from the weight. - assert lora_a is not None, f"middle block {idx} is missing SVDQuant LoRA" - assert not torch.equal(proj.weight, weights_before[idx]), ( - f"middle block {idx} weight was not modified by SVDQuant" - ) diff --git a/tests/examples/diffusers/test_qwen_pipeline_loading.py b/tests/examples/diffusers/test_qwen_pipeline_loading.py deleted file mode 100644 index cbc49aa5ea3..00000000000 --- a/tests/examples/diffusers/test_qwen_pipeline_loading.py +++ /dev/null @@ -1,74 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - -"""Negative-path loading tests for Qwen-Image in the diffusers quantization example. - -These cover the negative loading paths without a GPU or a real model: -- selecting Qwen-Image when diffusers lacks the Qwen classes raises a clear, - actionable error (not an opaque failure); -- Qwen loading does not pass ``trust_remote_code``. -""" - -import logging -import sys -from pathlib import Path - -import pytest - -pytest.importorskip("diffusers") -pytest.importorskip("torch") - -_EXAMPLE_DIR = Path(__file__).parents[3] / "examples" / "diffusers" / "quantization" -if str(_EXAMPLE_DIR) not in sys.path: - sys.path.insert(0, str(_EXAMPLE_DIR)) - -import models_utils # noqa: E402 -import pipeline_manager # noqa: E402 -from models_utils import ModelType # noqa: E402 -from quantize_config import ModelConfig # noqa: E402 - - -def _qwen_pipeline_manager() -> "pipeline_manager.PipelineManager": - config = ModelConfig(model_type=ModelType.QWEN_IMAGE, backbone=["transformer"]) - return pipeline_manager.PipelineManager(config, logging.getLogger("qwen-loading-test")) - - -def test_missing_qwen_pipeline_raises_actionable_error(monkeypatch): - # Simulate a diffusers version without QwenImagePipeline. - monkeypatch.setitem(models_utils.MODEL_PIPELINE, ModelType.QWEN_IMAGE, None) - manager = _qwen_pipeline_manager() - with pytest.raises(ImportError, match="Qwen-Image requires"): - manager.create_pipeline() - - -def test_qwen_loading_does_not_pass_trust_remote_code(monkeypatch): - captured_kwargs: dict = {} - - class _FakeQwenPipeline: - @classmethod - def from_pretrained(cls, model_id, **kwargs): - captured_kwargs.update(kwargs) - return cls() - - def set_progress_bar_config(self, **kwargs): - pass - - monkeypatch.setitem(models_utils.MODEL_PIPELINE, ModelType.QWEN_IMAGE, _FakeQwenPipeline) - manager = _qwen_pipeline_manager() - manager.create_pipeline() - - # Qwen-Image must load without trust_remote_code. - assert captured_kwargs.get("trust_remote_code") is not True - assert "trust_remote_code" not in captured_kwargs diff --git a/tests/unit/torch/export/test_convert_hf_config_svdquant.py b/tests/unit/torch/export/test_convert_hf_config_svdquant.py deleted file mode 100644 index 13a22ab22d5..00000000000 --- a/tests/unit/torch/export/test_convert_hf_config_svdquant.py +++ /dev/null @@ -1,87 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - -"""Unit tests for the NVFP4_SVD (SVDQuant) HF quantization-config conversion.""" - -from modelopt.torch.export.convert_hf_config import ( - _quant_algo_to_group_config, - convert_hf_quant_config_format, -) - - -def test_nvfp4_svd_group_config_mirrors_awq_with_pre_quant_scale(): - """The NVFP4_SVD config group is NVFP4 weights/activations + a pre_quant_scale flag.""" - group = _quant_algo_to_group_config("NVFP4_SVD", group_size=16) - assert group["pre_quant_scale"] is True - assert group["has_zero_point"] is False - assert group["weights"] == { - "dynamic": False, - "num_bits": 4, - "type": "float", - "group_size": 16, - } - assert group["input_activations"]["num_bits"] == 4 - assert group["input_activations"]["type"] == "float" - assert group["input_activations"]["group_size"] == 16 - - -def test_convert_hf_quant_config_format_nvfp4_svd(): - """A full NVFP4_SVD quantization dict converts to a complete compressed-tensors config.""" - input_config = { - "producer": {"name": "modelopt", "version": "0.0.0"}, - "quantization": { - "quant_algo": "NVFP4_SVD", - "group_size": 16, - "has_zero_point": False, - "pre_quant_scale": True, - "lora_rank": 32, - "exclude_modules": ["transformer_blocks.0.*", "proj_out"], - "kv_cache_quant_algo": None, - }, - } - - out = convert_hf_quant_config_format(input_config) - - # A real config group is emitted (not a bare {"quant_algo": ...} fallback). - assert "config_groups" in out - group = out["config_groups"]["group_0"] - assert group["pre_quant_scale"] is True - assert group["has_zero_point"] is False - assert group["lora_rank"] == 32 - assert group["weights"]["num_bits"] == 4 - assert group["weights"]["type"] == "float" - assert group["weights"]["group_size"] == 16 - assert group["input_activations"]["num_bits"] == 4 - assert group["targets"] == ["Linear"] - - # Top-level metadata is preserved. - assert out["quant_algo"] == "NVFP4_SVD" - assert out["ignore"] == ["transformer_blocks.0.*", "proj_out"] - assert out["quant_method"] == "modelopt" - - -def test_convert_hf_quant_config_format_nvfp4_svd_without_rank(): - """lora_rank is optional; omitting it must not break the conversion.""" - input_config = { - "quantization": { - "quant_algo": "NVFP4_SVD", - "group_size": 16, - "pre_quant_scale": True, - }, - } - out = convert_hf_quant_config_format(input_config) - group = out["config_groups"]["group_0"] - assert "lora_rank" not in group - assert group["pre_quant_scale"] is True diff --git a/tests/unit/torch/export/test_diffusers_qwen_export.py b/tests/unit/torch/export/test_diffusers_qwen_export.py deleted file mode 100644 index ad4b62b048f..00000000000 --- a/tests/unit/torch/export/test_diffusers_qwen_export.py +++ /dev/null @@ -1,133 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - -"""Tests for the Qwen-Image SVDQuant diffusers export path. - -Covers the three pieces added for Qwen support: -- the Qwen branch of ``generate_diffusion_dummy_inputs`` (validated by running the - dummy forward on a real tiny ``QwenImageTransformer2DModel``), -- the strict-failure mode of ``_fuse_qkv_linears_diffusion``, -- promotion of quantizer-owned SVDQuant tensors to clean module-level keys that - survive ``hide_quantizers_from_state_dict``. -""" - -from functools import partial - -import pytest -import torch -import torch.nn as nn - -import modelopt.torch.quantization as mtq -from modelopt.torch.export.diffusers_utils import ( - generate_diffusion_dummy_forward_fn, - generate_diffusion_dummy_inputs, - hide_quantizers_from_state_dict, -) -from modelopt.torch.export.unified_export_hf import ( - _fuse_qkv_linears_diffusion, - _promote_quantizer_tensors_to_module, -) - - -class _MLP(nn.Module): - def __init__(self, dim: int = 64): - super().__init__() - self.fc1 = nn.Linear(dim, dim) - self.fc2 = nn.Linear(dim, dim) - - def forward(self, x): - return self.fc2(torch.relu(self.fc1(x))) - - -def _forward_loop(model, data): - for batch in data: - model(batch) - - -def _quantize(model: nn.Module, algorithm=None, dim: int = 64) -> nn.Module: - cfg = mtq.INT8_SMOOTHQUANT_CFG.copy() - if algorithm is not None: - cfg["algorithm"] = algorithm - data = [torch.randn(2, dim) for _ in range(2)] - mtq.quantize(model, cfg, partial(_forward_loop, data=data)) - return model - - -def test_qwen_dummy_inputs_drive_real_transformer_forward(): - """The Qwen dummy inputs must actually drive a real tiny Qwen transformer.""" - pytest.importorskip("diffusers") - from _test_utils.torch.diffusers_models import get_tiny_qwen_image_transformer - - transformer = get_tiny_qwen_image_transformer().to("cpu", torch.float32).eval() - - inputs = generate_diffusion_dummy_inputs(transformer, torch.device("cpu"), torch.float32) - assert inputs is not None - for key in ( - "hidden_states", - "encoder_hidden_states", - "encoder_hidden_states_mask", - "img_shapes", - ): - assert key in inputs, f"missing Qwen dummy input '{key}'" - assert inputs["hidden_states"].shape[-1] == transformer.config.in_channels - assert inputs["encoder_hidden_states"].shape[-1] == transformer.config.joint_attention_dim - - # Strongest check: the generated dummy inputs run through the real model. - with torch.no_grad(): - generate_diffusion_dummy_forward_fn(transformer)() - - -def test_qwen_qkv_fusion_strict_raises_on_failed_dummy_forward(): - """strict=True turns a dummy-forward failure into a hard error; strict=False does not.""" - model = _quantize(_MLP()) - - def _boom(): - raise RuntimeError("dummy forward failed") - - with pytest.raises(RuntimeError): - _fuse_qkv_linears_diffusion(model, dummy_forward_fn=_boom, strict=True) - - # Non-strict path warns and returns without raising. - _fuse_qkv_linears_diffusion(model, dummy_forward_fn=_boom, strict=False) - - -def test_svdquant_promotion_survives_hide_quantizers(): - """Promoted LoRA + pre_quant_scale land on the module under clean keys and - survive ``hide_quantizers_from_state_dict`` (which strips the quantizers).""" - model = _quantize(_MLP(), algorithm={"method": "svdquant", "lowrank": 8}) - - _promote_quantizer_tensors_to_module(model) - - linears = [m for m in model.modules() if isinstance(m, torch.nn.Linear)] - assert linears - for module in linears: - assert hasattr(module, "svdquant_lora_a") - assert hasattr(module, "svdquant_lora_b") - # INT8_SMOOTHQUANT produces a pre_quant_scale that is promoted too. - assert hasattr(module, "pre_quant_scale") - # Rank-consistent shapes: lora_a [rank, in], lora_b [out, rank]. - assert module.svdquant_lora_a.shape[1] == module.in_features - assert module.svdquant_lora_b.shape[0] == module.out_features - assert module.svdquant_lora_a.shape[0] == module.svdquant_lora_b.shape[1] - - with hide_quantizers_from_state_dict(model): - keys = list(model.state_dict().keys()) - - assert any(k.endswith(".svdquant_lora_a") for k in keys) - assert any(k.endswith(".svdquant_lora_b") for k in keys) - assert any(k.endswith(".pre_quant_scale") for k in keys) - # Clean keys only: no quantizer-prefixed keys remain once quantizers are hidden. - assert not any("weight_quantizer" in k for k in keys) - assert not any("input_quantizer" in k for k in keys) diff --git a/tests/unit/torch/quantization/test_svdquant_forward_fold.py b/tests/unit/torch/quantization/test_svdquant_forward_fold.py deleted file mode 100644 index 085414f93a3..00000000000 --- a/tests/unit/torch/quantization/test_svdquant_forward_fold.py +++ /dev/null @@ -1,116 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 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. - -"""SVDQuant forward / fold coverage. - -These tests protect the invariants the diffusers SVDQuant export relies on: the -LoRA factors stay on the ``weight_quantizer`` in the live model and the export -layer promotes them. They complement (and do not modify) the existing -``test_calib.py::test_svdquant_lora_weights``. -""" - -from functools import partial - -import torch -import torch.nn as nn - -import modelopt.torch.quantization as mtq - - -class _SVDMLP(nn.Module): - def __init__(self, dim: int = 64): - super().__init__() - self.fc1 = nn.Linear(dim, dim) - self.fc2 = nn.Linear(dim, dim) - - def forward(self, x): - return self.fc2(torch.relu(self.fc1(x))) - - -def _forward_loop(model, dataloader): - for batch in dataloader: - model(batch) - - -def _quantize_svdquant(dim: int = 64) -> nn.Module: - model = _SVDMLP(dim) - quant_config = mtq.INT8_SMOOTHQUANT_CFG.copy() - quant_config["algorithm"] = {"method": "svdquant", "lowrank": 8} - data = [torch.randn(2, dim) for _ in range(2)] - mtq.quantize(model, quant_config, partial(_forward_loop, dataloader=data)) - return model - - -def _quantized_linears(model: nn.Module): - return [m for m in model.modules() if isinstance(m, torch.nn.Linear)] - - -def test_svdquant_lora_stays_on_weight_quantizer(): - """LoRA lives on the quantizer, not the module (the export layer promotes it).""" - model = _quantize_svdquant() - linears = _quantized_linears(model) - assert linears - for module in linears: - wq = module.weight_quantizer - assert wq.svdquant_lora_a is not None - assert wq.svdquant_lora_b is not None - # Not refactored onto the module. - assert not hasattr(module, "svdquant_lora_a") - assert not hasattr(module, "svdquant_lora_b") - - -def test_svdquant_forward_includes_nonzero_residual(): - """The forward output includes a nonzero low-rank residual term.""" - model = _quantize_svdquant() - for module in _quantized_linears(model): - x = torch.randn(2, module.in_features) - - residual = module._compute_lora_residual(x) - assert residual is not None - assert torch.count_nonzero(residual) > 0 - - full = module(x) - - # Temporarily drop the LoRA buffers to get the base (no-residual) output. - wq = module.weight_quantizer - lora_a = wq._svdquant_lora_a - lora_b = wq._svdquant_lora_b - delattr(wq, "_svdquant_lora_a") - delattr(wq, "_svdquant_lora_b") - try: - base = module(x) - finally: - wq.register_buffer("_svdquant_lora_a", lora_a) - wq.register_buffer("_svdquant_lora_b", lora_b) - - # The residual measurably changes the forward output. - assert not torch.allclose(full, base) - - -def test_svdquant_fold_weight_removes_buffers_and_changes_weight(): - """fold_weight() folds the residual into the weight and drops the buffers.""" - model = _quantize_svdquant() - for module in _quantized_linears(model): - wq = module.weight_quantizer - assert hasattr(wq, "_svdquant_lora_a") - assert hasattr(wq, "_svdquant_lora_b") - - weight_before = module.weight.detach().clone() - module.fold_weight() - - assert not hasattr(wq, "_svdquant_lora_a") - assert not hasattr(wq, "_svdquant_lora_b") - # Folding (quantized weight + low-rank residual) changes the stored weight. - assert not torch.allclose(module.weight, weight_before) From 9b472b234f530cc5e040532d44ee94c15129b43d Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Fri, 12 Jun 2026 16:40:20 -0700 Subject: [PATCH 15/22] Qwen-Image: add a fast CPU test for the diffusers SVDQuant export promotion Covers svdquant calibration -> _promote_quantizer_tensors_to_module -> clean module-level keys (svdquant_lora_a/b, pre_quant_scale) with the quantizers hidden, plus the post-export cleanup. Runs on CPU in <1s (INT8_SMOOTHQUANT + svdquant on a tiny linear stack). The full NVFP4 end-to-end check remains test_qwen_image_hf_ckpt_export[qwen_nvfp4_svdquant]; svdquant calibration is already covered by test_calib.py. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- .../torch/export/test_export_diffusers.py | 51 +++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 1a7a3495158..2f1e5085a3e 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -117,3 +117,54 @@ def test_flux2_dummy_inputs_shape(): # guidance_embeds defaults to True for Flux2 assert "guidance" in inputs + + +def test_svdquant_diffusers_export_promotes_clean_keys(): + """Fast CPU check of the diffusers SVDQuant export promotion. + + SVDQuant calibration stores the low-rank factors on ``weight_quantizer`` and the + smoothing scale on ``input_quantizer``; the diffusers export promotes both to + clean module-level keys (``svdquant_lora_a/b``, ``pre_quant_scale``) and hides the + quantizers, so the saved state dict carries no live quantizer tensors. The full + NVFP4 end-to-end coverage lives in the GPU test + ``tests/examples/diffusers/test_export_diffusers_hf_ckpt.py`` (``qwen_nvfp4_svdquant``). + """ + import copy + + import torch.nn as nn + + import modelopt.torch.quantization as mtq + from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict + + torch.manual_seed(0) + model = nn.Sequential(nn.Linear(64, 64), nn.Linear(64, 64)) + + quant_config = copy.deepcopy(mtq.INT8_SMOOTHQUANT_CFG) + quant_config["algorithm"] = {"method": "svdquant", "lowrank": 8} + mtq.quantize(model, quant_config, lambda m: m(torch.randn(8, 64))) + + # Calibration populated the quantizer-owned SVDQuant tensors. + linear = model[0] + assert linear.weight_quantizer.svdquant_lora_a is not None + assert linear.weight_quantizer.svdquant_lora_b is not None + assert getattr(linear.input_quantizer, "_pre_quant_scale", None) is not None + + # Export promotes them to clean module-level keys and hides the quantizers. + unified_export_hf._promote_quantizer_tensors_to_module(model) + with hide_quantizers_from_state_dict(model): + keys = set(model.state_dict().keys()) + + assert any(k.endswith(".svdquant_lora_a") for k in keys) + assert any(k.endswith(".svdquant_lora_b") for k in keys) + assert any(k.endswith(".pre_quant_scale") for k in keys) + assert not any("weight_quantizer" in k or "input_quantizer" in k for k in keys), ( + "live quantizer state leaked into the exported state dict" + ) + + # The promotion is undone after export, leaving the live module unchanged. + unified_export_hf._remove_promoted_quantizer_tensors(model) + keys_after = set(model.state_dict().keys()) + assert not any( + k.endswith((".svdquant_lora_a", ".svdquant_lora_b", ".pre_quant_scale")) + for k in keys_after + ) From 4a92a696b257e574e25b707d88d73f8194be989d Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Tue, 23 Jun 2026 11:53:13 -0700 Subject: [PATCH 16/22] Qwen-Image diffusers PTQ: address review feedback (cleanup safety, rank check, imports) - unified_export_hf: wrap promote -> save -> post-process -> config-update in try/finally so the temporary promoted export buffers are always removed even if an exception occurs, keeping the live module reusable for a repeated export. - unified_export_hf: _detect_svdquant_rank collects the unique SVDQuant ranks across modules and raises on a mismatch instead of silently recording the first module's rank as the single config-level lora_rank. - Hoist deferred imports to module level (inspect in diffusers_utils and the tiny Qwen fixture; copy/torch.nn/mtq/hide_quantizers in the export-promotion test) per the test coding guidelines. Verified on GB200: ruff clean; the fast CPU export-promotion test and the three Qwen e2e export cases (fp8/nvfp4/svdquant) pass. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- modelopt/torch/export/diffusers_utils.py | 5 +- modelopt/torch/export/unified_export_hf.py | 105 ++++++++++-------- tests/_test_utils/torch/diffusers_models.py | 3 +- .../torch/export/test_export_diffusers.py | 18 ++- 4 files changed, 70 insertions(+), 61 deletions(-) diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index 075f25a9101..0309c83e1f6 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -15,6 +15,7 @@ """Code that export quantized Hugging Face models for deployment.""" +import inspect import json import warnings from collections.abc import Callable @@ -26,8 +27,6 @@ import torch.nn as nn from safetensors.torch import load_file, safe_open -from .layer_utils import is_quantlinear - DiffusionPipeline: type[Any] | None ModelMixin: type[Any] | None try: # diffusers is optional for LTX-2 export paths @@ -361,8 +360,6 @@ def _qwen_inputs() -> dict[str, Any]: # Only pass kwargs the installed QwenImageTransformer2DModel.forward accepts # (signatures vary across diffusers versions); prevents the strict QKV-fusion # dummy forward from failing on an unexpected keyword argument. - import inspect - try: accepted = set(inspect.signature(model.forward).parameters) dummy_inputs = {k: v for k, v in dummy_inputs.items() if k in accepted} diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index abbea9fc37e..8605addfed2 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1022,17 +1022,25 @@ def _fuse_qkv_linears_diffusion( def _detect_svdquant_rank(component: nn.Module) -> int | None: - """Return the SVDQuant low-rank dimension from the first SVDQuant linear, if any. + """Return the single SVDQuant low-rank dimension shared by the SVDQuant linears. - ``svdquant_lora_a`` has shape ``(rank, in_features)``, so its first dimension - is the low-rank size. + ``svdquant_lora_a`` has shape ``(rank, in_features)``, so its first dimension is + the low-rank size. A single global ``lora_rank`` is written to the checkpoint + config, so all SVDQuant linears are expected to share one rank; an inconsistency + is raised rather than silently recording one module's rank for all. Returns + ``None`` when no SVDQuant LoRA factors are present. """ + ranks: set[int] = set() for _, sub_module in component.named_modules(): weight_quantizer = getattr(sub_module, "weight_quantizer", None) lora_a = getattr(weight_quantizer, "svdquant_lora_a", None) if lora_a is not None: - return int(lora_a.shape[0]) - return None + ranks.add(int(lora_a.shape[0])) + if not ranks: + return None + if len(ranks) > 1: + raise ValueError(f"Inconsistent SVDQuant ranks across modules: {sorted(ranks)}") + return next(iter(ranks)) def _promote_quantizer_tensors_to_module(component: nn.Module) -> None: @@ -1177,48 +1185,55 @@ def _export_diffusers_checkpoint( # main safetensors under clean, AWQ-aligned keys. _promote_quantizer_tensors_to_module(component) - # Build quantization config - quant_config = get_quant_config(component, is_modelopt_qlora=False) - if quant_config: - quantization_details = quant_config.get("quantization", {}) - # Record the SVDQuant low-rank size so consumers know the LoRA shape. - if quantization_details.get("quant_algo") == "NVFP4_SVD": - svdquant_rank = _detect_svdquant_rank(component) - if svdquant_rank is not None: - quantization_details["lora_rank"] = svdquant_rank - hf_quant_config = convert_hf_quant_config_format(quant_config) if quant_config else None - - # Save the component - # - diffusers ModelMixin.save_pretrained does NOT accept state_dict parameter - # - for non-diffusers modules (e.g., LTX-2 transformer), fall back to torch.save - if hasattr(component, "save_pretrained"): - with hide_quantizers_from_state_dict(component): - component.save_pretrained(component_export_dir, max_shard_size=max_shard_size) - else: - with hide_quantizers_from_state_dict(component): - _save_component_state_dict_safetensors(component, component_export_dir) - - # Post-process — merge, metadata, padding, swizzle - _postprocess_safetensors( - component_export_dir, - pipe, - hf_quant_config=hf_quant_config, - **kwargs, - ) + # Build the quantization config + save inside try/finally so the temporary + # promoted buffers are always removed, even if save / post-process / config + # update raises (keeps the live module reusable for a repeated export). + try: + quant_config = get_quant_config(component, is_modelopt_qlora=False) + if quant_config: + quantization_details = quant_config.get("quantization", {}) + # Record the SVDQuant low-rank size so consumers know the LoRA shape. + if quantization_details.get("quant_algo") == "NVFP4_SVD": + svdquant_rank = _detect_svdquant_rank(component) + if svdquant_rank is not None: + quantization_details["lora_rank"] = svdquant_rank + hf_quant_config = ( + convert_hf_quant_config_format(quant_config) if quant_config else None + ) - # Update config.json with quantization info - if hf_quant_config is not None: - config_path = component_export_dir / "config.json" - if config_path.exists(): - with open(config_path) as file: - config_data = json.load(file) - config_data["quantization_config"] = hf_quant_config - with open(config_path, "w") as file: - json.dump(config_data, file, indent=4) + # Save the component + # - diffusers ModelMixin.save_pretrained does NOT accept state_dict parameter + # - for non-diffusers modules (e.g., LTX-2 transformer), fall back to torch.save + if hasattr(component, "save_pretrained"): + with hide_quantizers_from_state_dict(component): + component.save_pretrained( + component_export_dir, max_shard_size=max_shard_size + ) + else: + with hide_quantizers_from_state_dict(component): + _save_component_state_dict_safetensors(component, component_export_dir) + + # Post-process — merge, metadata, padding, swizzle + _postprocess_safetensors( + component_export_dir, + pipe, + hf_quant_config=hf_quant_config, + **kwargs, + ) - # Drop the temporary promoted export buffers so the live module is - # unchanged after export (supports repeated export / module reuse). - _remove_promoted_quantizer_tensors(component) + # Update config.json with quantization info + if hf_quant_config is not None: + config_path = component_export_dir / "config.json" + if config_path.exists(): + with open(config_path) as file: + config_data = json.load(file) + config_data["quantization_config"] = hf_quant_config + with open(config_path, "w") as file: + json.dump(config_data, file, indent=4) + finally: + # Drop the temporary promoted export buffers so the live module is + # unchanged after export (supports repeated export / module reuse). + _remove_promoted_quantizer_tensors(component) # Non-quantized component: just save as-is elif hasattr(component, "save_pretrained"): component.save_pretrained(component_export_dir, max_shard_size=max_shard_size) diff --git a/tests/_test_utils/torch/diffusers_models.py b/tests/_test_utils/torch/diffusers_models.py index caf4966399f..4c88288360c 100644 --- a/tests/_test_utils/torch/diffusers_models.py +++ b/tests/_test_utils/torch/diffusers_models.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import inspect from pathlib import Path import pytest @@ -293,8 +294,6 @@ def get_tiny_qwen_image_transformer(**config_kwargs): # `pooled_projection_dim` is present in the published config.json but was removed # from the constructor in newer diffusers: from_pretrained tolerates the extra # config key, but a direct constructor call raises TypeError. - import inspect - accepted = set(inspect.signature(QwenImageTransformer2DModel.__init__).parameters) kwargs = {k: v for k, v in kwargs.items() if k in accepted} return QwenImageTransformer2DModel(**kwargs) diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 2f1e5085a3e..62388677e51 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -13,10 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import json import pytest import torch +import torch.nn as nn from _test_utils.torch.diffusers_models import ( get_tiny_dit, get_tiny_flux, @@ -27,8 +29,12 @@ pytest.importorskip("diffusers") import modelopt.torch.export.unified_export_hf as unified_export_hf +import modelopt.torch.quantization as mtq from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format -from modelopt.torch.export.diffusers_utils import generate_diffusion_dummy_inputs +from modelopt.torch.export.diffusers_utils import ( + generate_diffusion_dummy_inputs, + hide_quantizers_from_state_dict, +) from modelopt.torch.export.unified_export_hf import export_hf_checkpoint @@ -129,13 +135,6 @@ def test_svdquant_diffusers_export_promotes_clean_keys(): NVFP4 end-to-end coverage lives in the GPU test ``tests/examples/diffusers/test_export_diffusers_hf_ckpt.py`` (``qwen_nvfp4_svdquant``). """ - import copy - - import torch.nn as nn - - import modelopt.torch.quantization as mtq - from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict - torch.manual_seed(0) model = nn.Sequential(nn.Linear(64, 64), nn.Linear(64, 64)) @@ -165,6 +164,5 @@ def test_svdquant_diffusers_export_promotes_clean_keys(): unified_export_hf._remove_promoted_quantizer_tensors(model) keys_after = set(model.state_dict().keys()) assert not any( - k.endswith((".svdquant_lora_a", ".svdquant_lora_b", ".pre_quant_scale")) - for k in keys_after + k.endswith((".svdquant_lora_a", ".svdquant_lora_b", ".pre_quant_scale")) for k in keys_after ) From 2ad6a8f4980cbaac79d8cab7d474ea901f67cee6 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Tue, 23 Jun 2026 16:53:53 -0700 Subject: [PATCH 17/22] Qwen-Image: simplify unavailable-pipeline error in PipelineManager This PR had added a Qwen-Image-specific ImportError in both PipelineManager.create_pipeline_from and .create_pipeline, raised when MODEL_PIPELINE[model_type] is None. That over-specialized the message: the same None condition also covers Flux2 (whose pipeline import is version-gated) and any future version-gated pipeline, which fell through to the misleading generic "does not use diffusers pipelines." Replace both blocks with a single ValueError stating the model type is not supported by the installed diffusers version, keeping the actionable "upgrade diffusers" hint. This matches the existing "Raises: ValueError" docstrings; no caller or test depends on the previous ImportError type (create_pipeline re-raises through a generic except Exception). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- .../quantization/pipeline_manager.py | 21 +++++++------------ 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index 5496f131ba2..af89ed568ff 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -61,13 +61,10 @@ def create_pipeline_from( """ pipeline_cls = MODEL_PIPELINE[model_type] if pipeline_cls is None: - if model_type == ModelType.QWEN_IMAGE: - raise ImportError( - "Qwen-Image requires a diffusers version that provides " - "QwenImagePipeline. Please upgrade diffusers (e.g. " - "`pip install -U diffusers`) to a release that includes Qwen-Image." - ) - raise ValueError(f"Model type {model_type.value} does not use diffusers pipelines.") + raise ValueError( + f"Model type {model_type.value} is not supported by the installed diffusers " + "version; upgrade diffusers to a release that provides its pipeline." + ) model_id = ( MODEL_REGISTRY[model_type] if override_model_path is None else override_model_path ) @@ -105,14 +102,10 @@ def create_pipeline(self) -> Any: pipeline_cls = MODEL_PIPELINE[self.config.model_type] if pipeline_cls is None: - if self.config.model_type == ModelType.QWEN_IMAGE: - raise ImportError( - "Qwen-Image requires a diffusers version that provides " - "QwenImagePipeline. Please upgrade diffusers (e.g. " - "`pip install -U diffusers`) to a release that includes Qwen-Image." - ) raise ValueError( - f"Model type {self.config.model_type.value} does not use diffusers pipelines." + f"Model type {self.config.model_type.value} is not supported by the " + "installed diffusers version; upgrade diffusers to a release that " + "provides its pipeline." ) self.pipe = pipeline_cls.from_pretrained( self.config.model_path, From 4bfe484945285c09a2060a660f8ff3aae632be9a Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Tue, 23 Jun 2026 17:00:41 -0700 Subject: [PATCH 18/22] Qwen-Image: drop the optional --sanity-image-path inference check Remove the --sanity-image-path flag and the in-memory sanity-image generation block from the diffusers quantization example. It was an optional developer convenience (render one image from the fake-quantized pipeline before export) that no test exercises and that the regular quantize/export flow does not need. Dropping it simplifies the example's main path. No functional impact on export: pipe, MODEL_DEFAULTS, and Path remain used elsewhere, and the Qwen e2e never invoked the flag. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- examples/diffusers/quantization/quantize.py | 32 --------------------- 1 file changed, 32 deletions(-) diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 0e0dbca0035..50632212e6b 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -574,15 +574,6 @@ def create_argument_parser() -> argparse.ArgumentParser: export_group.add_argument( "--restore-from", type=str, help="Path to restore from previous checkpoint" ) - export_group.add_argument( - "--sanity-image-path", - type=str, - default=None, - help="If set, generate one image from the in-memory quantized pipeline (after " - "quantization, before the weights are packed for export) and save it here. This is " - "a quick functional sanity check of quantized inference; it does NOT reload the " - "exported checkpoint.", - ) export_group.add_argument( "--trt-high-precision-dtype", type=str, @@ -722,29 +713,6 @@ def forward_loop(mod): pipeline_manager.print_quant_summary() - # Optional functional sanity check: generate one image from the in-memory - # quantized pipeline. This runs BEFORE export (while weights are still - # fake-quantized and runnable, not yet packed) and does not reload the - # exported checkpoint. - if args.sanity_image_path: - try: - logger.info(f"Generating sanity image to {args.sanity_image_path}") - inference_args = MODEL_DEFAULTS.get(model_type, {}).get("inference_extra_args", {}) - result = pipe( - prompt="A high-quality photo of a cat wearing sunglasses", - num_inference_steps=calib_config.n_steps, - **inference_args, - ) - sanity_path = Path(args.sanity_image_path) - sanity_path.parent.mkdir(parents=True, exist_ok=True) - result.images[0].save(str(sanity_path)) - logger.info("Sanity image saved successfully") - except Exception as sanity_error: - # A requested sanity image is a positive success criterion: if it - # cannot be produced, fail loudly rather than reporting success. - logger.error(f"Sanity image generation failed: {sanity_error}", exc_info=True) - raise - for backbone_name, backbone in pipeline_manager.iter_backbones(): export_manager.export_onnx( pipe, From 2269e96abe32a03f71a97b27e4ad43ad213c9ab2 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Tue, 23 Jun 2026 17:28:32 -0700 Subject: [PATCH 19/22] Qwen-Image: move onnx_utils import back to module top in quantize.py Restore the top-level `from onnx_utils.export import generate_fp8_scales, modelopt_export_sd` instead of importing it lazily inside export_onnx, matching main's placement and keeping the import style consistent with the rest of the example rather than special-casing this one import. This restores main's contract that the diffusers example depends on the modelopt `[onnx]` extra (onnx-graphsurgeon, onnxruntime, ...) to import quantize.py. The HF-export logic is unchanged; the module is ruff/isort clean and compiles. The top-level import path is the same one main's SDXL/Flux/Wan export tests already exercise in CI. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- examples/diffusers/quantization/quantize.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 50632212e6b..41a90089129 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -39,6 +39,7 @@ get_model_filter_func, parse_extra_params, ) +from onnx_utils.export import generate_fp8_scales, modelopt_export_sd from pipeline_manager import PipelineManager from quantize_config import ( CalibrationConfig, @@ -318,11 +319,6 @@ def export_onnx( if not self.config.onnx_dir: return - # onnx_graphsurgeon (pulled in by onnx_utils.export) is an optional dependency - # only needed for the ONNX export path; import lazily so the HF-checkpoint - # export runs without it installed. - from onnx_utils.export import generate_fp8_scales, modelopt_export_sd - self.logger.info(f"Starting ONNX export to {self.config.onnx_dir}") if quant_format == QuantFormat.FP8 and self._has_conv_layers(backbone): From daf6d0f31f46211226870d4e0118a294f6280bfb Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Tue, 23 Jun 2026 18:05:10 -0700 Subject: [PATCH 20/22] Qwen-Image: fix code-quality CI (mypy dict-item, ruff format/lint) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The code-quality job (pre-commit) failed on three hooks; all addressed: - mypy: convert_hf_config.py — annotate the shared `config_group_details` (first/FP8 branch) as `dict[str, Any]`. The NVFP4_SVD branch adds bool flags (`has_zero_point`, `pre_quant_scale`) that broke mypy's inferred `dict[str, Collection[str]]` (dict-item errors on the two bool entries). - ruff format: models_utils.py — wrap the over-length `rules.append(...)` line in build_block_range_quant_cfg. - ruff check: test_export_diffusers_hf_ckpt.py — UP037 (drop quotes on the `object` annotation), PIE810 (merge the two `endswith` calls into one tuple call), and `# noqa: SIM118` on the `safe_open(...).keys()` loop (safe_open is not directly iterable; matches the convention used elsewhere in the repo). No runtime behavior changes. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Jingyu Xin --- examples/diffusers/quantization/models_utils.py | 4 +++- modelopt/torch/export/convert_hf_config.py | 2 +- tests/examples/diffusers/test_export_diffusers_hf_ckpt.py | 6 +++--- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index 1126a421390..9d366dc7402 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -366,6 +366,8 @@ def build_block_range_quant_cfg( {"quantizer_name": f"*{block_module}.*input_quantizer", "enable": True}, ] for idx in excluded: - rules.append({"quantizer_name": f"*{block_module}.{idx}.*weight_quantizer", "enable": False}) + rules.append( + {"quantizer_name": f"*{block_module}.{idx}.*weight_quantizer", "enable": False} + ) rules.append({"quantizer_name": f"*{block_module}.{idx}.*input_quantizer", "enable": False}) return rules diff --git a/modelopt/torch/export/convert_hf_config.py b/modelopt/torch/export/convert_hf_config.py index 6f7dedb97c8..45fa0c30f3b 100644 --- a/modelopt/torch/export/convert_hf_config.py +++ b/modelopt/torch/export/convert_hf_config.py @@ -182,7 +182,7 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An # This structure is derived based on the example for "FP8" and "NVFP4" # TODO: Handle other quantization algorithms if quant_algo_value == "FP8": - config_group_details = { + config_group_details: dict[str, Any] = { "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"}, "weights": {"dynamic": False, "num_bits": 8, "type": "float"}, "targets": ["Linear"], diff --git a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py index 9e2e30a160c..618804f30c6 100644 --- a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py +++ b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py @@ -219,14 +219,14 @@ def test_qwen_image_hf_ckpt_export( assert quant_config.get("quant_method") == "modelopt" keys: set[str] = set() - lora_tensors: dict[str, "object"] = {} + lora_tensors: dict[str, object] = {} safetensors_files = sorted(transformer_dir.rglob("*.safetensors")) assert safetensors_files, f"no safetensors in {transformer_dir}" for path in safetensors_files: with safe_open(str(path), framework="pt") as handle: - for key in handle.keys(): + for key in handle.keys(): # noqa: SIM118 - safe_open is not iterable keys.add(key) - if key.endswith(".svdquant_lora_a") or key.endswith(".svdquant_lora_b"): + if key.endswith((".svdquant_lora_a", ".svdquant_lora_b")): lora_tensors[key] = handle.get_tensor(key) # No live quantizer state should leak into the exported checkpoint. From a689f3100640f44770d4e9ba5b3b6cfe891b0210 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Thu, 25 Jun 2026 11:21:44 -0700 Subject: [PATCH 21/22] Add Qwen-Image DMD2 PTQ support; save quantizer state (amax) without weights Two related changes to the diffusers quantization example. 1) Weight-free quantizer-state checkpoint (all models) The torch checkpoint written by `--quantized-torch-ckpt-save-path` now stores ONLY ModelOpt's quantization state -- the recipe plus the quantizer buffers (amax, pre_quant_scale, ...) -- and NOT the model weights. The weights live in the base HF/diffusers checkpoint and are reloaded there on restore. This uses ModelOpt's own idiom (mirrors plugins/transformers_trainer.py): save: modelopt_state = mto.modelopt_state(model) modelopt_state["modelopt_state_weights"] = get_quantizer_state_dict(model) torch.save(modelopt_state, path) restore: modelopt_state = mto.load_modelopt_state(path) weights = modelopt_state.pop("modelopt_state_weights", None) mto.restore_from_modelopt_state(model, modelopt_state) set_quantizer_state_dict(model, weights) Wrapped as utils.save_quantizer_state / restore_quantizer_state and wired into ExportManager.save_checkpoint / restore_checkpoint. Restore auto-applies on top of the freshly-loaded base weights (the pipeline is created before restore). Effect: the artifact drops from a full-model checkpoint to KBs-MBs (a 60-layer Qwen-Image student: 40.8 GB -> 2.0 MB) while amax round-trips bit-identically. 2) qwen-image-dmd2 model type (DMD2 few-step Qwen-Image students) For students distilled by examples/diffusers/fastgen. Reuses the existing Qwen-Image quantization stack from the base branch (filter_func_qwen_image, the block-range recipe, QwenImagePipeline registration) and adds only what differs: - pipeline_manager: load the consolidated student transformer (+ optional EMA) and swap it into the base QwenImagePipeline; stash the few-step sampler config (defaults to the canonical 4-step shift=3 [1.0, 0.9, 0.75, 0.5, 0.0] ODE schedule, guidance_scale=1.0). - calibration: drive the few-step DMD sampler instead of the standard denoising loop, so collected amax matches how the student is actually run. - qwen_image_dmd2_sampler.py (new): vendored compact DMD unroll, bit-aligned with fastgen/inference_dmd2_qwen_image.py; transformer forwards only for calibration (decode=False), optional VAE decode for sanity inference (decode=True). - sanity_check_dmd2.py (new): restore the quantizer-state checkpoint and run one few-step inference to validate the round trip. - quantize.py: import the ONNX export tooling lazily so the calibration + save path runs without onnx_graphsurgeon (e.g. the diffusers/fastgen container). Validated end-to-end (FP8, single GB200): calibrate -> save (2.0 MB) -> restore_quantizer_state reproduces amax bit-identically (e.g. block 30 attn.to_q input amax 1.26e+03) and the restored quantized student renders a finite, non-constant 1024x1024 image. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jingyu Xin --- .../diffusers/quantization/calibration.py | 20 ++ .../diffusers/quantization/models_utils.py | 19 ++ .../quantization/pipeline_manager.py | 115 ++++++++ examples/diffusers/quantization/quantize.py | 21 +- .../quantization/qwen_image_dmd2_sampler.py | 260 ++++++++++++++++++ .../quantization/sanity_check_dmd2.py | 175 ++++++++++++ examples/diffusers/quantization/utils.py | 38 +++ 7 files changed, 642 insertions(+), 6 deletions(-) create mode 100644 examples/diffusers/quantization/qwen_image_dmd2_sampler.py create mode 100644 examples/diffusers/quantization/sanity_check_dmd2.py diff --git a/examples/diffusers/quantization/calibration.py b/examples/diffusers/quantization/calibration.py index 27b1ec22436..bebc61970a3 100644 --- a/examples/diffusers/quantization/calibration.py +++ b/examples/diffusers/quantization/calibration.py @@ -21,6 +21,7 @@ from models_utils import MODEL_DEFAULTS, ModelType from pipeline_manager import PipelineManager from quantize_config import CalibrationConfig +from qwen_image_dmd2_sampler import dmd2_sample from tqdm import tqdm from utils import load_calib_prompts @@ -95,6 +96,9 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: elif self.model_type in [ModelType.WAN22_T2V_14b, ModelType.WAN22_T2V_5b]: # Special handling for WAN video models self._run_wan_video_calibration(prompt_batch, extra_args) + elif self.model_type == ModelType.QWEN_IMAGE_DMD2: + # DMD2 students use a custom few-step sampler, not the standard loop. + self._run_qwen_image_dmd2_calibration(prompt_batch) else: common_args = { "prompt": prompt_batch, @@ -105,6 +109,22 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: self.logger.debug(f"Completed calibration batch {i + 1}/{self.config.num_batches}") self.logger.info("Calibration completed successfully") + def _run_qwen_image_dmd2_calibration(self, prompt_batch: list[str]) -> None: + """Calibrate a DMD2 Qwen-Image student via its few-step sampler. + + Drives the same few-step DMD unroll the student was trained/served with + (NOT the standard denoising loop) so the collected activation statistics + are representative of inference. The VAE decode is skipped — calibration + only needs the transformer forwards. + """ + cfg = self.pipeline_manager.dmd_sampler_cfg + if cfg is None: + raise RuntimeError( + "DMD2 sampler config is not set; the qwen-image-dmd2 pipeline must be created " + "via PipelineManager.create_pipeline() before calibration." + ) + dmd2_sample(self.pipe, prompt_batch, decode=False, **cfg) + def _run_wan_video_calibration( self, prompt_batch: list[str], extra_args: dict[str, Any] ) -> None: diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index 9d366dc7402..99e8b5e6c12 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -64,6 +64,10 @@ class ModelType(str, Enum): WAN22_T2V_14b = "wan2.2-t2v-14b" WAN22_T2V_5b = "wan2.2-t2v-5b" QWEN_IMAGE = "qwen-image" + # DMD2-distilled few-step Qwen-Image student (from examples/diffusers/fastgen). + # Same architecture as QWEN_IMAGE, but loaded from a consolidated student dir + # and calibrated with the few-step DMD sampler instead of the standard loop. + QWEN_IMAGE_DMD2 = "qwen-image-dmd2" _FILTER_FUNC_MAP: dict[ModelType, Callable[[str], bool]] = { @@ -74,6 +78,7 @@ class ModelType(str, Enum): ModelType.WAN22_T2V_14b: filter_func_wan_video, ModelType.WAN22_T2V_5b: filter_func_wan_video, ModelType.QWEN_IMAGE: filter_func_qwen_image, + ModelType.QWEN_IMAGE_DMD2: filter_func_qwen_image, } _VAE_FILTER_FUNC_MAP: dict[tuple[ModelType, str], Callable[[str], bool]] = { @@ -107,6 +112,11 @@ def get_model_filter_func( ModelType.WAN22_T2V_14b: "Wan-AI/Wan2.2-T2V-A14B-Diffusers", ModelType.WAN22_T2V_5b: "Wan-AI/Wan2.2-TI2V-5B-Diffusers", ModelType.QWEN_IMAGE: "Qwen/Qwen-Image", + # Base pipeline (VAE / text-encoder / tokenizer / scheduler) for DMD2 students; + # the trained transformer is loaded separately from a consolidated dir via the + # ``student_path`` extra-param. Override with ``--override-model-path`` or + # ``--extra-param base_pipeline_path=...``. + ModelType.QWEN_IMAGE_DMD2: "Qwen/Qwen-Image", } MODEL_PIPELINE: dict[ModelType, type[DiffusionPipeline] | None] = { @@ -122,6 +132,7 @@ def get_model_filter_func( ModelType.WAN22_T2V_14b: WanPipeline, ModelType.WAN22_T2V_5b: WanPipeline, ModelType.QWEN_IMAGE: QwenImagePipeline, + ModelType.QWEN_IMAGE_DMD2: QwenImagePipeline, } # Shared dataset configurations @@ -258,6 +269,14 @@ def get_model_filter_func( }, } +# DMD2 students share Qwen-Image's architecture, so they reuse the same block-range +# recipe, high-precision filter, base pipeline, and calibration dataset. They differ +# only in (a) loading -- a consolidated student dir swapped into the base pipeline +# (PipelineManager._create_qwen_image_dmd2_pipeline) -- and (b) calibration, which +# drives the few-step DMD sampler instead of the standard denoising loop +# (Calibrator._run_qwen_image_dmd2_calibration). Inherit so the recipe stays in sync. +MODEL_DEFAULTS[ModelType.QWEN_IMAGE_DMD2] = {**MODEL_DEFAULTS[ModelType.QWEN_IMAGE]} + def _coerce_extra_param_value(value: str) -> Any: lowered = value.lower() diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index af89ed568ff..4e76ee0d661 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -43,6 +43,9 @@ def __init__(self, config: ModelConfig, logger: logging.Logger): self.pipe_upsample: LTXLatentUpsamplePipeline | None = None # For LTX-Video upsampling self._transformer: torch.nn.Module | None = None self._video_decoder: torch.nn.Module | None = None + # Few-step sampler config for DMD2 students (populated when loading a + # qwen-image-dmd2 pipeline); consumed by the calibrator / sanity check. + self.dmd_sampler_cfg: dict[str, Any] | None = None @staticmethod def create_pipeline_from( @@ -100,6 +103,11 @@ def create_pipeline(self) -> Any: self.logger.info("LTX-2 pipeline created successfully") return self.pipe + if self.config.model_type == ModelType.QWEN_IMAGE_DMD2: + self.pipe = self._create_qwen_image_dmd2_pipeline() + self.logger.info("Qwen-Image DMD2 pipeline created successfully") + return self.pipe + pipeline_cls = MODEL_PIPELINE[self.config.model_type] if pipeline_cls is None: raise ValueError( @@ -266,6 +274,113 @@ def _create_ltx2_pipeline(self) -> Any: pipeline_kwargs.update(params) return TI2VidTwoStagesPipeline(**pipeline_kwargs) + def _create_qwen_image_dmd2_pipeline(self) -> Any: + """Build a QwenImagePipeline whose transformer is a DMD2-trained student. + + Loads the consolidated student transformer (the ``model/consolidated`` dir + produced by ``examples/diffusers/fastgen`` training), optionally overlays an + EMA shadow, and swaps it into the base Qwen-Image pipeline so the VAE / + text-encoder / tokenizer / scheduler come from the base checkpoint. + + Reads from ``extra_params``: + student_path (required): consolidated student dir. + base_pipeline_path: base Qwen-Image dir/HF id (defaults to the + registry id or ``--override-model-path``). + ema_path: optional ``ema_shadow.pt`` to overlay onto the student. + sample_steps / t_list / sample_type / guidance_scale / max_t: + few-step sampler schedule (defaults match the canonical 4-step + shift=3 student); stashed in ``self.dmd_sampler_cfg``. + """ + from qwen_image_dmd2_sampler import DEFAULT_MAX_T, resolve_schedule + + try: + from diffusers import QwenImagePipeline, QwenImageTransformer2DModel + except ImportError as e: + raise ImportError( + "qwen-image-dmd2 requires a diffusers version providing QwenImagePipeline " + "and QwenImageTransformer2DModel; upgrade diffusers." + ) from e + + params = dict(self.config.extra_params) + student_path = params.get("student_path") + if not student_path: + raise ValueError( + "Missing required extra_param: student_path (the consolidated DMD2 student " + "dir, e.g. .../epoch_4_step_17999/model/consolidated)." + ) + base_pipeline_path = params.get("base_pipeline_path") or self.config.model_path + ema_path = params.get("ema_path") + + default_dtype = self.config.model_dtype["default"] + transformer_dtype = self.config.model_dtype.get("transformer", default_dtype) + if torch.float16 in (default_dtype, transformer_dtype): + self.logger.warning( + "Qwen-Image is trained/served in bfloat16; float16 (Half) can overflow the " + "VAE and produce NaNs. Consider --model-dtype BFloat16." + ) + + self.logger.info("Loading DMD2 student transformer from %s", student_path) + transformer = QwenImageTransformer2DModel.from_pretrained( + student_path, torch_dtype=transformer_dtype + ) + + if ema_path: + self.logger.info("Overlaying EMA shadow from %s", ema_path) + ema_state = torch.load(str(ema_path), map_location="cpu") + shadow = ( + ema_state.get("shadow", ema_state) if isinstance(ema_state, dict) else ema_state + ) + if not isinstance(shadow, dict): + raise ValueError( + f"ema_path content has unexpected type {type(shadow).__name__}; " + "expected dict[str, Tensor]." + ) + missing, unexpected = transformer.load_state_dict(shadow, strict=False) + if unexpected: + self.logger.warning("EMA overlay had %d unexpected key(s)", len(unexpected)) + if missing: + self.logger.warning("EMA overlay missed %d student key(s)", len(missing)) + + transformer.eval() + + self.logger.info( + "Loading base Qwen-Image pipeline from %s (transformer replaced by student)", + base_pipeline_path, + ) + pipe = QwenImagePipeline.from_pretrained( + base_pipeline_path, transformer=transformer, torch_dtype=default_dtype + ) + pipe.set_progress_bar_config(disable=True) + + # Resolve and stash the few-step sampler config. Defaults match the + # canonical 4-step shift=3 student; the schedule MUST match training. + sample_steps = params.get("sample_steps") + sample_steps = int(sample_steps) if sample_steps is not None else 4 + t_list = params.get("t_list") + if isinstance(t_list, str): + t_list = [float(x) for x in t_list.split(",") if x.strip()] + max_t = float(params.get("max_t", DEFAULT_MAX_T)) + schedule = resolve_schedule(t_list, sample_steps, max_t) + defaults = MODEL_DEFAULTS[self.config.model_type].get("inference_extra_args", {}) + self.dmd_sampler_cfg = { + "schedule": schedule, + "sample_type": str(params.get("sample_type", "ode")), + "guidance_scale": float(params.get("guidance_scale", 1.0)), + "negative_prompt": params.get("negative_prompt"), + "height": int(params.get("height", defaults.get("height", 1024))), + "width": int(params.get("width", defaults.get("width", 1024))), + "max_sequence_length": int(params.get("max_sequence_length", 512)), + } + self.logger.info( + "DMD2 few-step sampler: steps=%d schedule=%s sample_type=%s guidance_scale=%s " + "(schedule must match the student's training t_list)", + len(schedule) - 1, + schedule, + self.dmd_sampler_cfg["sample_type"], + self.dmd_sampler_cfg["guidance_scale"], + ) + return pipe + def print_quant_summary(self): for name, backbone in self.iter_backbones(): self.logger.info(f"{name} quantization info:") diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 41a90089129..fa1cc729f92 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -39,7 +39,6 @@ get_model_filter_func, parse_extra_params, ) -from onnx_utils.export import generate_fp8_scales, modelopt_export_sd from pipeline_manager import PipelineManager from quantize_config import ( CalibrationConfig, @@ -51,9 +50,8 @@ QuantFormat, QuantizationConfig, ) -from utils import check_conv_and_mha, check_lora +from utils import check_conv_and_mha, check_lora, restore_quantizer_state, save_quantizer_state -import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.export import export_hf_checkpoint @@ -295,8 +293,11 @@ def save_checkpoint( filename = f"{backbone_name}.pt" if backbone_name else "backbone.pt" target_path = ckpt_path / filename - self.logger.info(f"Saving backbone to {target_path}") - mto.save(backbone, str(target_path)) + # Save ONLY the quantization state (recipe + quantizer buffers incl. amax), + # not the model weights. The weights live in the base HF/diffusers checkpoint + # and are reloaded there on restore; this keeps the artifact tiny. + self.logger.info(f"Saving quantizer state (amax + recipe, no weights) to {target_path}") + save_quantizer_state(backbone, str(target_path)) self.logger.info("Checkpoint saved successfully") @@ -319,6 +320,12 @@ def export_onnx( if not self.config.onnx_dir: return + # Imported lazily: the ONNX export tooling (onnx_graphsurgeon, etc.) is only + # needed when --onnx-dir is set, so the calibration + amax-save path still + # works in environments without the ONNX deps (e.g. the diffusers/fastgen + # container used for Qwen-Image DMD2 students). + from onnx_utils.export import generate_fp8_scales, modelopt_export_sd + self.logger.info(f"Starting ONNX export to {self.config.onnx_dir}") if quant_format == QuantFormat.FP8 and self._has_conv_layers(backbone): @@ -362,7 +369,9 @@ def restore_checkpoint(self) -> None: f"Checkpoint not found for '{backbone_name}' in {restore_path}" ) self.logger.info(f"Restoring {backbone_name} from {source_path}") - mto.restore(backbone, str(source_path)) + # The pipeline was just created with the base (unquantized) weights, so + # this re-applies the quantization recipe + amax on top of them. + restore_quantizer_state(backbone, str(source_path)) self.logger.info("Checkpoints restored successfully") diff --git a/examples/diffusers/quantization/qwen_image_dmd2_sampler.py b/examples/diffusers/quantization/qwen_image_dmd2_sampler.py new file mode 100644 index 00000000000..518bc6b47d5 --- /dev/null +++ b/examples/diffusers/quantization/qwen_image_dmd2_sampler.py @@ -0,0 +1,260 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compact DMD2 few-step sampler for Qwen-Image students. + +This is a vendored, calibration-friendly version of the few-step unroll in +``examples/diffusers/fastgen/inference_dmd2_qwen_image.py``. It is kept here so +the quantization example is self-contained (no cross-example ``sys.path`` +imports) and so calibration can run the **same forward logic the student was +trained/served with** — which is what makes the collected ``amax`` statistics +representative. + +The single :func:`dmd2_sample` entry point serves two callers: + +* **Calibration** (``decode=False``): runs only the transformer forwards of the + DMD unroll and returns ``None``. The VAE / image post-processing is skipped + because quantization only needs the transformer's activation statistics, and + skipping the VAE saves substantial time and memory on the 60-layer student. +* **Sanity inference** (``decode=True``): additionally runs the VAE decode and + returns a list of images, used to confirm a restored (quantized) student + still produces a finite image. + +The math is bit-aligned with the training-time ``_build_student_input`` in +``modelopt/torch/fastgen/methods/dmd.py`` and with the inference reference: + + for (t_cur, t_next) in pairwise(t_list): + v = student(x, t=t_cur, text_emb) # flow at t_cur + x_0 = x - t_cur * v # RF identity -> x_0 estimate + if t_next > 0: + eps = (x - (1 - t_cur) * x_0) / t_cur # ODE: invert RF forward + x = (1 - t_next) * x_0 + t_next * eps # re-noise to t_next + else: + x = x_0 # final step + +``t_list`` MUST match the student's training schedule (e.g. the LightX2V +"shift=3" 4-step shape ``[1.0, 0.9, 0.75, 0.5, 0.0]``); a mismatch produces a +train/inference gap and therefore misleading calibration statistics. +""" + +from __future__ import annotations + +import itertools + +import torch +from diffusers.utils.torch_utils import randn_tensor + +# Canonical 4-step "shift=3" student schedule (LightX2V-Qwen-Image-Lightning +# shape). t_list has student_sample_steps + 1 entries: the first N are the +# timesteps the student is evaluated at, the trailing 0.0 is the terminal the +# final Euler step lands on (NOT an extra evaluation). +DEFAULT_T_LIST: tuple[float, ...] = (1.0, 0.9, 0.75, 0.5, 0.0) +DEFAULT_MAX_T: float = 0.999 + + +def resolve_schedule( + t_list: list[float] | tuple[float, ...] | None, + sample_steps: int | None, + max_t: float = DEFAULT_MAX_T, +) -> list[float]: + """Resolve the sampling schedule (timesteps + terminal 0.0). + + Priority: + 1. An explicit ``t_list`` (must end at 0.0 and have ``sample_steps + 1`` + entries when ``sample_steps`` is given). + 2. ``sample_steps == 1`` -> ``[max_t, 0.0]`` (canonical single-step). + 3. ``sample_steps == 4`` (or None) with no ``t_list`` -> ``DEFAULT_T_LIST``. + 4. Otherwise a linear ``linspace(max_t, 0, sample_steps + 1)`` fallback. + """ + if t_list is not None: + schedule = [float(t) for t in t_list] + if abs(schedule[-1]) > 1e-6: + raise ValueError( + f"t_list must end at 0.0 (got {schedule[-1]}); the final step lands on x_0." + ) + if sample_steps is not None and len(schedule) != sample_steps + 1: + raise ValueError( + f"t_list must have sample_steps+1 entries " + f"(got {len(schedule)} for sample_steps={sample_steps})." + ) + return schedule + + if sample_steps == 1: + return [float(max_t), 0.0] + if sample_steps in (None, 4): + return list(DEFAULT_T_LIST) + return torch.linspace(float(max_t), 0.0, sample_steps + 1).tolist() + + +@torch.no_grad() +def dmd2_sample( + pipe, + prompt: str | list[str], + *, + schedule: list[float], + sample_type: str = "ode", + guidance_scale: float = 1.0, + negative_prompt: str | list[str] | None = None, + height: int = 1024, + width: int = 1024, + num_images_per_prompt: int = 1, + generator: torch.Generator | None = None, + max_sequence_length: int = 512, + decode: bool = False, + output_type: str = "pil", +) -> list | None: + """Run the DMD few-step unroll on ``pipe.transformer``. + + Args: + pipe: A ``QwenImagePipeline`` whose ``transformer`` is the DMD2 student. + prompt: A prompt or list of prompts (one calibration batch). + schedule: Full timestep schedule incl. trailing 0.0 (see + :func:`resolve_schedule`). + sample_type: ``"ode"`` (deterministic, recover eps via RF identity) or + ``"sde"`` (fresh Gaussian noise between steps). Must match training. + guidance_scale: Inference-time CFG. Leave at ``1.0`` for students trained + with an internalised (non-null) ``dmd2.guidance_scale`` — passing + ``> 1.0`` there would double-apply CFG. + negative_prompt: Negative prompt for CFG; defaults to ``""`` when CFG is + engaged and none is given. + height/width: Output spatial size (must be VAE-compatible). + num_images_per_prompt: Images per prompt. + generator: Optional RNG for reproducible noise. + max_sequence_length: Text-encoder max sequence length. + decode: If ``True`` run VAE decode + post-process and return images. If + ``False`` (calibration) skip the VAE and return ``None``. + output_type: Passed to the image processor when ``decode=True``. + + Returns: + A list of images when ``decode=True``, else ``None``. + """ + if sample_type not in ("ode", "sde"): + raise ValueError(f"sample_type must be 'ode' or 'sde', got {sample_type!r}") + + do_cfg = guidance_scale != 1.0 + if do_cfg and negative_prompt is None: + negative_prompt = "" + + device = pipe.transformer.device + dtype = next(pipe.transformer.parameters()).dtype + + # ---- Encode prompt(s) ------------------------------------------------ + prompt_embeds, prompt_embeds_mask = pipe.encode_prompt( + prompt=prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + neg_prompt_embeds = neg_prompt_embeds_mask = None + if do_cfg: + neg_prompt_embeds, neg_prompt_embeds_mask = pipe.encode_prompt( + prompt=negative_prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + txt_seq_lens = ( + prompt_embeds_mask.sum(dim=1).int().tolist() if prompt_embeds_mask is not None else None + ) + neg_txt_seq_lens = ( + neg_prompt_embeds_mask.sum(dim=1).int().tolist() + if neg_prompt_embeds_mask is not None + else None + ) + + # ---- Build initial noisy latents at t = schedule[0] ------------------ + batch_size = (1 if isinstance(prompt, str) else len(prompt)) * num_images_per_prompt + num_channels_latents = pipe.transformer.config.in_channels // 4 # 64 // 4 = 16 + h_lat = 2 * (height // (pipe.vae_scale_factor * 2)) + w_lat = 2 * (width // (pipe.vae_scale_factor * 2)) + latent_shape = (batch_size, 1, num_channels_latents, h_lat, w_lat) + + noise = randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype) + latents_5d = noise * schedule[0] # RF: sigma(t0) = t0 + x_packed = pipe._pack_latents(latents_5d, batch_size, num_channels_latents, h_lat, w_lat) + img_shapes = [[(1, h_lat // 2, w_lat // 2)]] * batch_size + + # ---- DMD few-step unroll (transformer forwards) ---------------------- + for t_cur, t_next in itertools.pairwise(schedule): + timestep = torch.tensor([t_cur], device=device, dtype=dtype).expand(batch_size) + flow_packed = pipe.transformer( + hidden_states=x_packed, + encoder_hidden_states=prompt_embeds, + encoder_hidden_states_mask=prompt_embeds_mask, + timestep=timestep, + img_shapes=img_shapes, + txt_seq_lens=txt_seq_lens, + guidance=None, + return_dict=False, + )[0] + if do_cfg: + neg_flow_packed = pipe.transformer( + hidden_states=x_packed, + encoder_hidden_states=neg_prompt_embeds, + encoder_hidden_states_mask=neg_prompt_embeds_mask, + timestep=timestep, + img_shapes=img_shapes, + txt_seq_lens=neg_txt_seq_lens, + guidance=None, + return_dict=False, + )[0] + flow_packed = ( + neg_flow_packed.to(torch.float64) + + float(guidance_scale) + * (flow_packed.to(torch.float64) - neg_flow_packed.to(torch.float64)) + ).to(dtype) + + # RF identity: x_0 = x_t - t_cur * v (fp64 for stability). + x0_packed = (x_packed.to(torch.float64) - float(t_cur) * flow_packed.to(torch.float64)).to( + dtype + ) + + if t_next > 1e-6: + if sample_type == "ode": + alpha_cur = 1.0 - float(t_cur) + eps_packed = ( + (x_packed.to(torch.float64) - alpha_cur * x0_packed.to(torch.float64)) + / max(float(t_cur), 1e-6) + ).to(dtype) + else: + eps_packed = torch.randn( + x_packed.shape, generator=generator, device=device, dtype=dtype + ) + alpha_next = 1.0 - float(t_next) + x_packed = ( + alpha_next * x0_packed.to(torch.float64) + + float(t_next) * eps_packed.to(torch.float64) + ).to(dtype) + else: + x_packed = x0_packed + + if not decode: + # Calibration path: transformer forwards already ran; nothing to decode. + return None + + # ---- VAE decode (sanity-inference path only) ------------------------- + x0_5d = pipe._unpack_latents(x_packed, height, width, pipe.vae_scale_factor) + latents_mean = ( + torch.tensor(pipe.vae.config.latents_mean) + .view(1, pipe.vae.config.z_dim, 1, 1, 1) + .to(device=device, dtype=dtype) + ) + latents_std = 1.0 / torch.tensor(pipe.vae.config.latents_std).view( + 1, pipe.vae.config.z_dim, 1, 1, 1 + ).to(device=device, dtype=dtype) + x0_scaled = x0_5d / latents_std + latents_mean + image_5d = pipe.vae.decode(x0_scaled, return_dict=False)[0] + image_4d = image_5d[:, :, 0] # Qwen-Image treats images as 1-frame videos + return pipe.image_processor.postprocess(image_4d, output_type=output_type) diff --git a/examples/diffusers/quantization/sanity_check_dmd2.py b/examples/diffusers/quantization/sanity_check_dmd2.py new file mode 100644 index 00000000000..9f78901f20c --- /dev/null +++ b/examples/diffusers/quantization/sanity_check_dmd2.py @@ -0,0 +1,175 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Restore a quantized DMD2 Qwen-Image student and run one few-step inference. + +Confirms the round trip of the new ``qwen-image-dmd2`` quantization flow: + + 1. Load the base Qwen-Image pipeline with the consolidated student swapped in + (via the same :class:`PipelineManager` path quantize.py uses) -- this brings + the original (unquantized) weights. + 2. Reapply the weight-free quantization checkpoint saved by ``quantize.py`` + (``save_quantizer_state`` -> ``transformer.pt``) via + ``restore_quantizer_state``, which re-applies the quantizer recipe **and the + calibrated amax** buffers on top of the loaded weights. + 3. Run a single few-step DMD inference (with VAE decode) and assert the image + is finite and non-constant. + +This deliberately reuses :class:`PipelineManager` and +:func:`qwen_image_dmd2_sampler.dmd2_sample` so the inference path is identical to +calibration's (minus the VAE decode, which is enabled here). + +Usage:: + + python sanity_check_dmd2.py \\ + --quantized-ckpt ./qwen_dmd2_fp8/transformer.pt \\ + --student-path /.../epoch_4_step_17999/model/consolidated \\ + --base-pipeline-path /.../models/Qwen-Image \\ + --output-png ./qwen_dmd2_fp8/sanity.png +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import sys + +import torch +from models_utils import ModelType +from pipeline_manager import PipelineManager +from quantize_config import ModelConfig +from qwen_image_dmd2_sampler import dmd2_sample +from utils import restore_quantizer_state + +import modelopt.torch.quantization as mtq + +logger = logging.getLogger("sanity_check_dmd2") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--quantized-ckpt", + required=True, + help="Path to the quantized checkpoint saved by quantize.py (e.g. .../transformer.pt).", + ) + parser.add_argument( + "--student-path", + required=True, + help="Consolidated DMD2 student dir (provides architecture + base weights to restore into).", + ) + parser.add_argument( + "--base-pipeline-path", + default="Qwen/Qwen-Image", + help="Base Qwen-Image dir/HF id for the VAE / text-encoder / tokenizer / scheduler.", + ) + parser.add_argument("--ema-path", default=None, help="Optional EMA shadow overlaid on load.") + parser.add_argument("--output-png", default="./qwen_dmd2_sanity.png") + parser.add_argument("--prompt", default="a small red cube on a white table") + parser.add_argument("--height", type=int, default=1024) + parser.add_argument("--width", type=int, default=1024) + parser.add_argument("--seed", type=int, default=42) + # Few-step sampler knobs (defaults match the canonical 4-step shift=3 student). + parser.add_argument("--sample-steps", type=int, default=4) + parser.add_argument( + "--t-list", + default=None, + help="Comma-separated schedule incl. trailing 0.0, e.g. '1.0,0.9,0.75,0.5,0.0'.", + ) + parser.add_argument("--sample-type", default="ode", choices=["ode", "sde"]) + parser.add_argument("--guidance-scale", type=float, default=1.0) + args = parser.parse_args() + + logging.basicConfig( + level=logging.INFO, format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s" + ) + + # 1. Build the base pipeline with the student swapped in (unquantized). + extra_params: dict[str, object] = { + "student_path": args.student_path, + "base_pipeline_path": args.base_pipeline_path, + "sample_steps": args.sample_steps, + "sample_type": args.sample_type, + "guidance_scale": args.guidance_scale, + "height": args.height, + "width": args.width, + } + if args.ema_path: + extra_params["ema_path"] = args.ema_path + if args.t_list: + extra_params["t_list"] = args.t_list + + model_config = ModelConfig( + model_type=ModelType.QWEN_IMAGE_DMD2, + model_dtype={"default": torch.bfloat16}, + backbone=["transformer"], + extra_params=extra_params, + ) + pm = PipelineManager(model_config, logger) + pipe = pm.create_pipeline() + + # 2. Restore the quantized architecture + calibrated amax into the student. + logger.info( + "Restoring quantizer state (amax + recipe) from %s onto the loaded student", + args.quantized_ckpt, + ) + restore_quantizer_state(pipe.transformer, args.quantized_ckpt) + mtq.print_quant_summary(pipe.transformer) + pm.setup_device() + + # 3. One few-step inference (with VAE decode). + gen = torch.Generator(device=pipe.transformer.device).manual_seed(args.seed) + images = dmd2_sample(pipe, [args.prompt], decode=True, generator=gen, **pm.dmd_sampler_cfg) + image = images[0] + + import numpy as np + + arr = np.asarray(image) + stats = { + "prompt": args.prompt, + "quantized_ckpt": args.quantized_ckpt, + "schedule": pm.dmd_sampler_cfg["schedule"], + "image_shape": list(arr.shape), + "image_dtype": str(arr.dtype), + "image_min": float(arr.min()), + "image_max": float(arr.max()), + "image_mean": float(arr.mean()), + "image_std": float(arr.std()), + "is_finite": bool(np.isfinite(arr).all()), + "is_not_constant": bool(arr.std() > 0), + } + + os.makedirs(os.path.dirname(os.path.abspath(args.output_png)), exist_ok=True) + image.save(args.output_png) + with open(args.output_png.replace(".png", "_stats.json"), "w") as f: + json.dump(stats, f, indent=2) + print(json.dumps(stats, indent=2)) + + if not stats["is_finite"]: + logger.error("Sanity check FAILED: image contains non-finite values.") + sys.exit(1) + if not stats["is_not_constant"]: + logger.error("Sanity check FAILED: image is constant (std == 0).") + sys.exit(1) + logger.info( + "Sanity check PASSED: restored quantized student produced a finite image -> %s", + args.output_png, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index c3cfdcd5cdd..9a6b841a52e 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -24,8 +24,13 @@ from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear from diffusers.utils import load_image +import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq from modelopt.torch.quantization.plugins.diffusion.diffusers import AttentionModuleMixin +from modelopt.torch.quantization.utils.core_utils import ( + get_quantizer_state_dict, + set_quantizer_state_dict, +) USE_PEFT = True try: @@ -193,3 +198,36 @@ def mha_filter_func(name): if hasattr(F, "scaled_dot_product_attention"): mtq.disable_quantizer(backbone, mha_filter_func) + + +def save_quantizer_state(model: torch.nn.Module, path: str) -> None: + """Save ONLY ModelOpt's quantization state -- the recipe plus the quantizer + buffers (amax, pre_quant_scale, ...) -- and NOT the model weights. + + This is the same idiom ModelOpt uses internally (see + ``modelopt.torch.quantization.plugins.transformers_trainer``): the + ``modelopt_state`` (architecture/recipe from :func:`mto.modelopt_state`) is + bundled with the per-quantizer state from + :func:`get_quantizer_state_dict` under the ``modelopt_state_weights`` key. + The resulting checkpoint is tiny (KBs-MBs) and is reloaded on top of the + original (unquantized) model via :func:`restore_quantizer_state`. + """ + modelopt_state = mto.modelopt_state(model) + modelopt_state["modelopt_state_weights"] = get_quantizer_state_dict(model) + torch.save(modelopt_state, str(path)) + + +def restore_quantizer_state(model: torch.nn.Module, path: str) -> torch.nn.Module: + """Reload a checkpoint written by :func:`save_quantizer_state` onto ``model``. + + ``model`` must already hold its original (unquantized) weights (e.g. freshly + loaded from the base HF/diffusers checkpoint); this re-applies the + quantization recipe and loads the calibrated amax/quantizer buffers on top. + Mirrors ModelOpt's ``_restore_modelopt_state_with_weights``. + """ + modelopt_state = mto.load_modelopt_state(str(path)) + quantizer_state = modelopt_state.pop("modelopt_state_weights", None) + mto.restore_from_modelopt_state(model, modelopt_state) + if quantizer_state is not None: + set_quantizer_state_dict(model, quantizer_state) + return model From f7a9140d8998895adcf1097cf3127c5f9607f0a2 Mon Sep 17 00:00:00 2001 From: Jingyu Xin Date: Tue, 30 Jun 2026 23:19:36 -0700 Subject: [PATCH 22/22] Add Qwen-Image-Edit DMD2 training and QAT support --- examples/diffusers/fastgen/README.md | 150 +++- .../configs/dmd2_qwen_image_edit_2511.yaml | 124 +++ examples/diffusers/fastgen/dmd2_finetune.py | 2 +- examples/diffusers/fastgen/dmd2_recipe.py | 360 +++++++- .../fastgen/fastgen_data/__init__.py | 6 + .../fastgen/fastgen_data/collate_fns.py | 171 ++++ .../fastgen_data/image_to_image_dataset.py | 144 ++++ .../fastgen/inference_dmd2_qwen_image.py | 11 - .../fastgen/inference_dmd2_qwen_image_edit.py | 330 ++++++++ .../fastgen/preprocess/processors/__init__.py | 2 + .../preprocess/processors/qwen_image_edit.py | 288 +++++++ .../fastgen/preprocess_qwen_image_edit.py | 792 ++++++++++++++++++ examples/diffusers/fastgen/requirements.txt | 4 + modelopt/torch/fastgen/plugins/__init__.py | 1 + modelopt/torch/fastgen/plugins/qwen_image.py | 52 +- .../torch/fastgen/plugins/qwen_image_edit.py | 249 ++++++ .../plugins/diffusion/diffusers.py | 27 +- .../fastgen/test_quant_state_roundtrip.py | 95 +++ .../fastgen/test_qwen_image_edit_data.py | 180 ++++ .../fastgen/test_qwen_image_edit_plugin.py | 312 +++++++ .../torch/fastgen/test_qwen_image_plugin.py | 29 +- 21 files changed, 3288 insertions(+), 41 deletions(-) create mode 100644 examples/diffusers/fastgen/configs/dmd2_qwen_image_edit_2511.yaml create mode 100644 examples/diffusers/fastgen/fastgen_data/image_to_image_dataset.py create mode 100644 examples/diffusers/fastgen/inference_dmd2_qwen_image_edit.py create mode 100644 examples/diffusers/fastgen/preprocess/processors/qwen_image_edit.py create mode 100644 examples/diffusers/fastgen/preprocess_qwen_image_edit.py create mode 100644 modelopt/torch/fastgen/plugins/qwen_image_edit.py create mode 100644 tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py create mode 100644 tests/examples/diffusers/fastgen/test_qwen_image_edit_data.py create mode 100644 tests/unit/torch/fastgen/test_qwen_image_edit_plugin.py diff --git a/examples/diffusers/fastgen/README.md b/examples/diffusers/fastgen/README.md index 9c9373807a9..977a7a033b8 100644 --- a/examples/diffusers/fastgen/README.md +++ b/examples/diffusers/fastgen/README.md @@ -1,4 +1,4 @@ -# DMD2 distillation for Qwen-Image +# DMD2 distillation for Qwen-Image and Qwen-Image-Edit Distill [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) into a **few-step generator** with DMD2 (Distribution Matching Distillation). The distilled student @@ -6,6 +6,10 @@ produces images in as few as **1–4 sampling steps** while matching the base mo output distribution. Built on `modelopt.torch.fastgen` and NeMo AutoModel's [`TrainDiffusionRecipe`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/recipes/diffusion/train.py). +The paired image-edit path supports +[`Qwen/Qwen-Image-Edit-2511`](https://huggingface.co/Qwen/Qwen-Image-Edit-2511), +including its one-or-more reference-image input contract. + > [!NOTE] > Qwen-Image is a third-party model with its own license terms. Review the > [Qwen-Image model card](https://huggingface.co/Qwen/Qwen-Image) before downloading or @@ -115,14 +119,135 @@ torchrun --nproc-per-node=8 \ Any `DMDConfig` field can be overridden on the CLI (e.g. `--dmd2.guidance_scale=3.5`). +## Qwen-Image-Edit-2511 training + +Image editing uses `configs/dmd2_qwen_image_edit_2511.yaml`. It is not a text-to-image +cache with an extra tensor: Edit-2511 conditions every transformer call with packed +reference-image latents, and its Qwen2.5-VL prompt embedding jointly encodes the edit +instruction and those same ordered references. Stable `diffusers>=0.37.0` is required so +the transformer's `zero_cond_t` path applies `t=0` modulation to the reference tokens. + +Preprocess native SpatialEdit WebDataset shards directly, without extracting the image +corpus: + +```bash +python examples/diffusers/fastgen/preprocess_qwen_image_edit.py \ + --input-dir /path/to/SpatialEdit-500K \ + --output-dir /path/to/qwen_image_edit_2511_cache \ + --model-name /path/to/Qwen-Image-Edit-2511 \ + --gpu-id 0 +``` + +The preprocessor also accepts `--manifest pairs.jsonl`. A row supplies a target, an edit +instruction, and one or more ordered references; image values may be local paths or +`{"archive": "/path/shard.tar", "member": "sample.0.jpg"}` descriptors: + +```json +{"id":"sample-1","target":"target.png","conditioning":["source.png"],"prompt":"Move the red cube left."} +``` + +Each cached sample contains the target latent, a list of deterministic reference latents, +and image-aware positive **and negative** embeddings. Consequently the edit dataloader +does not take `negative_prompt_embedding_path`; one global text-only negative embedding +would omit the reference-image visual tokens. + +Keep the edit dataloader at `batch_size: 1` with the current sampler. It buckets target +resolution only; batching multiple samples also requires matching reference count and every +reference-slot shape. The collate rejects incompatible batches instead of padding image tokens. + +```bash +torchrun --nproc-per-node=8 \ + examples/diffusers/fastgen/dmd2_finetune.py \ + --config examples/diffusers/fastgen/configs/dmd2_qwen_image_edit_2511.yaml \ + --model.pretrained_model_name_or_path=/path/to/Qwen-Image-Edit-2511 \ + --data.dataloader.cache_dir=/path/to/qwen_image_edit_2511_cache \ + --fsdp.dp_size=8 --step_scheduler.global_batch_size=8 +``` + +The `qwen_image_edit` plugin packs the noisy target first, appends every clean reference, +builds `img_shapes=[target, reference_1, ...]`, and slices the transformer prediction back +to the target-token prefix. The same references are forwarded through student, +teacher, fake-score, CFG, backward-simulation, and GAN paths. + ### Checkpoints & resuming Checkpoints land under `checkpoint.checkpoint_dir`. Alongside the student, the recipe saves the DMD2 sidecars needed to resume exactly: the fake-score model + optimizer, the -student EMA (`ema_shadow.pt`), and the DMD iteration counter (`dmd_state.pt`). With +DMD iteration counter (`dmd_state.pt`), and, when EMA is enabled, the student EMA +(`ema_shadow.pt`). With `restore_from: LATEST` a re-launch auto-resumes from the newest checkpoint; pin a specific one with `--checkpoint.restore_from=epoch_0_step_500`. +## Quantization-aware training (QAT) + +Continue a full-precision DMD2 run with the **student quantized**, so the few-step model +stays accurate at FP8/NVFP4. QAT here is **restore-only**: the trainer loads a ModelOpt +quantizer state (recipe + frozen `amax`) from disk and **never calibrates**. Only the +student is quantized; the frozen teacher and trainable fake-score stay full precision so +the distribution-matching gradient is exact, and `amax` stays frozen for the whole run. + +QAT is driven by a `dmd2.quant` block — there's no dedicated config file. The cleanest +way to launch is to **reuse the exact config + overrides of the full-precision run you're +continuing** and add only the three `dmd2.quant.*` keys (plus a reduced LR), so the QAT +run is provably identical to the FP run except for quantization and learning rate. The +CLI parser creates the `dmd2.quant` subtree even when it's absent from the YAML. + +| Key | Role | +| --- | --- | +| `dmd2.quant.enabled` | Turn QAT on (restore-only student quantization). | +| `dmd2.quant.quant_state_path` | The `transformer.pt` from step 1 below (recipe + frozen `amax`). | +| `dmd2.quant.init_weights_from` | FP DMD2 checkpoint to warm-start student / fake-score / optimizers from on the first launch (the run `amax` was calibrated against). | + +1. **Calibrate once** with the quantization example to produce the quantizer state + (`amax`, no weights) for a trained student checkpoint: + + ```bash + python examples/diffusers/quantization/quantize.py \ + --model qwen-image-dmd2 --format fp8 \ + --extra-param student_path=<.../epoch_4_step_15999/model/consolidated> \ + --quantized-torch-ckpt-save-path <.../epoch_4_step_15999/quant> + # -> writes <.../epoch_4_step_15999/quant/transformer.pt> + ``` + +2. **Launch QAT** by re-running the FP run's command with a new output dir, a reduced + student LR, and the three quant keys appended: + + ```bash + torchrun --nproc-per-node= \ + examples/diffusers/fastgen/dmd2_finetune.py \ + --config examples/diffusers/fastgen/configs/.yaml \ + --checkpoint.checkpoint_dir= \ + <... the FP run's other overrides, unchanged ...> \ + --optim.learning_rate= --lr_scheduler.min_lr= \ + --dmd2.quant.enabled=true \ + --dmd2.quant.quant_state_path=<.../epoch_4_step_15999/quant/transformer.pt> \ + --dmd2.quant.init_weights_from=<.../epoch_4_step_15999> + ``` + +On the first launch (empty `checkpoint_dir`) the student / fake-score / discriminator / +optimizers warm-start from `init_weights_from`, then the student is quantized from +`quant_state_path`. `restore_from: LATEST` auto-resumes the new `checkpoint_dir` +thereafter. Because QAT is restore-only — amax never recalibrates — the recipe re-applies +`quant_state_path` on every resume rather than persisting a per-checkpoint copy, so keep +that file accessible for the whole run (it's the only quantization dependency). The saved +student weights are clean full precision (`model/consolidated` is a normal +`QwenImageTransformer2DModel`); re-apply `quant_state_path` to deploy or evaluate the +quantized QAT student via the quantization example. + +> The `quant_state_path` `amax` must have been calibrated against the student in +> `init_weights_from`, with the same few-step schedule (`dmd2.sample_t_cfg.t_list`) the +> student trains/infers with. Pass `dmd2.quant.enabled=true` on every resume too (it is +> what tells the recipe to quantize). Reduce only the student LR by keeping +> `--dmd2.fake_score_lr` / `--dmd2.discriminator_lr` at the FP value. + +For Qwen-Image-Edit, the restore-only QAT path itself is unchanged because the student is +still a `QwenImageTransformer2DModel`. The quantizer state must, however, be calibrated on +the exact edit student using target + reference tokens, multimodal image/instruction +embeddings, and the same `t_list`. The existing `--model qwen-image-dmd2` text-only +calibrator does **not** provide representative activation ranges for Edit-2511. Also +calibrate the non-EMA weights restored by `dmd2.quant.init_weights_from`; do not calibrate +an EMA overlay for that warm start. + ## Inference After training, sample from the distilled student. The pipeline loads your consolidated @@ -162,6 +287,23 @@ Set `num_inference_steps` to the number of steps the student was trained for (`dmd2.student_sample_steps` — e.g. 4 for the canonical config, or 1 for a single-step student). +For an edit student, use the companion pipeline and pass one or more ordered references: + +```python +from inference_dmd2_qwen_image_edit import QwenImageEditDMDInferencePipeline +from diffusers.utils import load_image + +pipe = QwenImageEditDMDInferencePipeline.from_pretrained( + student_path="/path/to/checkpoint/model/consolidated", + base_pipeline_path="Qwen/Qwen-Image-Edit-2511", +).to("cuda") +image = pipe( + [load_image("source.png")], + "Move the red cube left.", + num_inference_steps=4, +).images[0] +``` + ## Config reference | Section | Key | Role | @@ -170,7 +312,7 @@ student). | `model` | `mode` | `finetune` — loads the pretrained weights. | | `step_scheduler` | `global_batch_size`, `local_batch_size`, `max_steps`, `ckpt_every_steps`, `log_every` | Standard AutoModel scheduling knobs. | | `dmd2` | `recipe_path` | Built-in fastgen recipe to hydrate `DMDConfig` from (`general/distillation/dmd2_qwen_image`). | -| `dmd2` | `pipeline_plugin` | `qwen_image` — selects `QwenImageDMDPipeline` (2×2 patch packing / img_shapes). | +| `dmd2` | `pipeline_plugin` | `qwen_image` for T2I or `qwen_image_edit` for target + reference token packing. | | `dmd2` | `student_sample_steps` | Number of student sampling steps (e.g. 4). | | `dmd2` | `guidance_scale` | CFG strength on the teacher (`null` disables CFG; requires a negative-prompt embedding when set). | | `dmd2` | `gan_loss_weight_gen`, `gan_r1_reg_weight`, `gan_feature_indices`, … | GAN branch (set `gan_loss_weight_gen: 0` to disable). | @@ -178,7 +320,7 @@ student). | `dmd2` | `sample_t_cfg`, `ema` | Timestep sampling + student EMA settings. | | `optim` | `learning_rate`, `optimizer.*` | Student AdamW knobs. | | `fsdp` | `dp_size`, `tp_size`, `activation_checkpointing`, … | FSDP2 parallelism (set `dp_size` to your GPU count). | -| `data` | `dataloader._target_`, `cache_dir`, `negative_prompt_embedding_path` | Latent cache dir + optional CFG negative-prompt embedding. | +| `data` | `dataloader._target_`, `cache_dir`, `negative_prompt_embedding_path` | Latent cache. The static negative path is T2I-only; edit negatives are cached per sample. | | `checkpoint` | `checkpoint_dir`, `model_save_format`, `restore_from` | Output dir, save format, resume behavior. | ## Troubleshooting diff --git a/examples/diffusers/fastgen/configs/dmd2_qwen_image_edit_2511.yaml b/examples/diffusers/fastgen/configs/dmd2_qwen_image_edit_2511.yaml new file mode 100644 index 00000000000..e515fa2e911 --- /dev/null +++ b/examples/diffusers/fastgen/configs/dmd2_qwen_image_edit_2511.yaml @@ -0,0 +1,124 @@ +# Qwen-Image-Edit-2511 DMD2 — paired image-edit training. +# +# Preprocess a paired dataset first (SpatialEdit-500K is supported directly): +# +# python examples/diffusers/fastgen/preprocess_qwen_image_edit.py \ +# --input-dir /path/to/SpatialEdit-500K \ +# --output-dir /path/to/qwen_image_edit_2511_cache \ +# --model-name Qwen/Qwen-Image-Edit-2511 +# +# Then launch with torchrun and override the cache/checkpoint paths as needed. +# A local model snapshot can be selected with: +# --model.pretrained_model_name_or_path=/path/to/Qwen-Image-Edit-2511 + +seed: 42 + +wandb: + project: fastgen-dmd2-qwen-image-edit + mode: online + name: qwen_image_edit_2511_dmd2 + +dist_env: + backend: nccl + timeout_minutes: 60 + +model: + pretrained_model_name_or_path: Qwen/Qwen-Image-Edit-2511 + mode: finetune + +step_scheduler: + global_batch_size: 128 + local_batch_size: 1 + ckpt_every_steps: 500 + num_epochs: 4 + log_every: 1 + max_steps: 5000 + +dmd2: + recipe_path: general/distillation/dmd2_qwen_image + # Unlike text-to-image, this plugin appends the cached reference-image tokens + # to the noisy target tokens and trains on the target prefix only. + pipeline_plugin: qwen_image_edit + qwen_image_guidance: + + pred_type: flow + num_train_timesteps: + guidance_scale: 4.0 + student_sample_steps: 4 + student_sample_type: ode + backward_simulation: false + student_update_freq: 5 + fake_score_pred_type: x0 + + gan_loss_weight_gen: 0.03 + gan_use_same_t_noise: true + gan_r1_reg_weight: 0.1 + gan_r1_reg_alpha: 0.1 + + fake_score_lr: 2.0e-6 + discriminator_lr: 2.0e-6 + gan_feature_indices: [30] + gan_num_blocks: 60 + gan_inner_dim: 3072 + + sample_t_cfg: + time_dist_type: uniform + min_t: 0.001 + max_t: 0.999 + p_mean: 0.0 + p_std: 1.0 + t_list: [0.999, 0.74925, 0.4995, 0.24975, 0.0] + + # Full-tensor EMA materializes a complete FP32 copy of this 20B-class student on + # every rank, which is not viable on 80-GiB workers alongside the three DMD2 models. + # Keep it disabled until the EMA checkpoint path supports sharded shadows end to end. + ema: null + + # Restore-only student QAT uses the same recipe. Enable after calibrating the + # exact FP edit student with reference-token + multimodal conditioning and this + # t_list; the text-only qwen-image-dmd2 calibration path is not representative. + quant: + enabled: false + quant_state_path: + init_weights_from: + +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: 128 + activation_checkpointing: true + +# Each cache item contains: +# target latent, 1..N reference latents, multimodal positive prompt embedding, +# and a per-sample multimodal negative embedding built from the same references. +data: + dataloader: + _target_: fastgen_data.build_image_to_image_multiresolution_dataloader + cache_dir: /path/to/preprocessed/qwen_image_edit_2511 + 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_edit_2511_dmd2/checkpoints + model_save_format: safetensors + save_consolidated: true + v4_compatible: true + diffusers_compatible: true + restore_from: LATEST diff --git a/examples/diffusers/fastgen/dmd2_finetune.py b/examples/diffusers/fastgen/dmd2_finetune.py index 6d91db94acd..22bddfb4aa1 100644 --- a/examples/diffusers/fastgen/dmd2_finetune.py +++ b/examples/diffusers/fastgen/dmd2_finetune.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Entrypoint for the DMD2 Qwen-Image AutoModel example. +"""Entrypoint for the DMD2 Qwen-Image / Qwen-Image-Edit AutoModel examples. Parses the YAML config + CLI overrides with AutoModel's argument parser, then hands control to :class:`DMD2DiffusionRecipe`. diff --git a/examples/diffusers/fastgen/dmd2_recipe.py b/examples/diffusers/fastgen/dmd2_recipe.py index 7934a07cf13..3ab18771755 100644 --- a/examples/diffusers/fastgen/dmd2_recipe.py +++ b/examples/diffusers/fastgen/dmd2_recipe.py @@ -20,10 +20,10 @@ drives ``modelopt.torch.fastgen.DMDPipeline`` (or a plugin subclass) through the three-phase DMD2 alternation (student update / fake-score update / EMA step). -Backbone: **Qwen-Image** (``Qwen/Qwen-Image``) — 4D ``image_latents``, -:class:`QwenImageDMDPipeline` handles 2x2 patch packing / img_shapes / -unpacking. Config: ``configs/dmd2_qwen_image.yaml`` — the canonical -real-data run (4-step + CFG + GAN). +Backbones: **Qwen-Image** and **Qwen-Image-Edit-2511** — both train 4D target +``image_latents``. :class:`QwenImageDMDPipeline` handles T2I patch packing; +:class:`QwenImageEditDMDPipeline` additionally appends clean reference-image tokens and +crops predictions back to the target prefix. Canonical configs live under ``configs/``. Launch:: @@ -37,6 +37,8 @@ from __future__ import annotations +import contextlib +import dataclasses import json import logging import os @@ -54,6 +56,11 @@ # and surfaced as a downstream ``TypeError: takes no arguments``. try: from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline + from nemo_automodel.components.distributed.parallelizer import ( + PARALLELIZATION_STRATEGIES, + DefaultParallelizationStrategy, + register_parallel_strategy, + ) from nemo_automodel.recipes.diffusion.train import TrainDiffusionRecipe, is_main_process except ImportError as exc: raise ImportError( @@ -68,10 +75,67 @@ from torch import nn import modelopt.torch.fastgen as mtf +import modelopt.torch.opt as mto from modelopt.torch.fastgen.config import DMDConfig from modelopt.torch.fastgen.discriminators import Discriminator_ImageDiT from modelopt.torch.fastgen.methods.dmd import DMDPipeline from modelopt.torch.fastgen.plugins import qwen_image as qwen_image_plugin +from modelopt.torch.quantization.utils.core_utils import set_quantizer_state_dict + + +class _QwenImageParallelizationStrategy(DefaultParallelizationStrategy): + """Add full-block activation checkpointing to AutoModel's native FSDP flow. + + AutoModel's default decoder-layer checkpointing recognizes ``self_attn`` / ``mlp`` + attributes. Diffusers' Qwen image blocks instead contain joint ``attn``, ``img_mlp``, + and ``txt_mlp`` paths, so that generic logic wraps nothing. Diffusers checkpoints each + complete block; mirror that boundary, then delegate TP/FSDP behavior unchanged. + """ + + def parallelize( + self, + model, + device_mesh, + activation_checkpointing: bool = False, + **kwargs, + ): + if activation_checkpointing: + from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import ( + CheckpointImpl, + checkpoint_wrapper, + ) + + blocks = getattr(model, "transformer_blocks", None) + if blocks is None: + raise AttributeError( + "QwenImageTransformer2DModel does not expose `transformer_blocks`" + ) + for index, block in enumerate(blocks): + blocks[index] = checkpoint_wrapper( + block, + checkpoint_impl=CheckpointImpl.NO_REENTRANT, + ) + logging.info( + "[DMD2] Qwen-Image activation checkpointing enabled for %d full blocks", + len(blocks), + ) + + return super().parallelize( + model, + device_mesh, + activation_checkpointing=False, + **kwargs, + ) + + +def _register_qwen_image_parallelization_strategy() -> None: + """Install the Qwen strategy unless AutoModel already provides a native one.""" + model_class_name = "QwenImageTransformer2DModel" + if model_class_name not in PARALLELIZATION_STRATEGIES: + register_parallel_strategy(name=model_class_name)(_QwenImageParallelizationStrategy) + + +_register_qwen_image_parallelization_strategy() # Keys under the ``dmd2:`` YAML block that shadow fields on :class:`DMDConfig`. The # recipe deep-merges these on top of the loaded built-in recipe so users can tweak DMD2 @@ -98,12 +162,41 @@ def _deep_merge_dicts(base: dict, override: dict) -> dict: return merged +def restore_quantizer_state(model: nn.Module, path: str) -> nn.Module: + """Re-insert the quantizer modules + load the frozen amax onto ``model`` from ``path``. + + ``path`` is the weight-free ModelOpt quantizer state written by the + ``examples/diffusers/quantization`` calibration example + (``--quantized-torch-ckpt-save-path`` → ``transformer.pt``): ``mto.modelopt_state`` + (which layers are quantized, FP8/NVFP4, axes) bundled with the per-quantizer buffers + (``amax``, ...) under ``modelopt_state_weights``. It carries NO model weights; ``model`` + must already hold the weights the amax was calibrated against (here: the DMD2 student + warm-started from the FP checkpoint). + + This re-applies the quantization recipe (module conversion -- ``nn.Linear`` -> + ``QuantLinear``, in place, preserving the existing ``weight``/``bias`` Parameter objects + so any pre-built FSDP2 optimizer's references stay valid) and loads the saved amax. It + is RESTORE-ONLY: no calibration forward pass is run, so amax stays exactly as it was on + disk and remains frozen for the whole training run. + """ + modelopt_state = mto.load_modelopt_state(str(path)) + quantizer_state = modelopt_state.pop("modelopt_state_weights", None) + mto.restore_from_modelopt_state(model, modelopt_state) + if quantizer_state is not None: + set_quantizer_state_dict(model, quantizer_state) + return model + + # Auto-detect substrings (matched case-insensitively against ``model_id``) that map to # DMDPipeline plugin subclasses. Keep this list small — adding a new entry is only the # right move when the model has a non-diffusers transformer signature that requires a # pack/unpack wrapper. Models with the standard ``(hidden_states, timestep, # encoder_hidden_states)`` signature work with the base :class:`DMDPipeline`. _PIPELINE_PLUGIN_BY_MODEL_SUBSTR = ( + # Edit must precede the generic Qwen-Image match: the edit transformer consumes + # target tokens followed by one or more reference-image token sequences. + ("qwen-image-edit", "qwen_image_edit"), + ("qwen_image_edit", "qwen_image_edit"), ("qwen-image", "qwen_image"), ("qwen_image", "qwen_image"), ) @@ -111,6 +204,29 @@ def _deep_merge_dicts(base: dict, override: dict) -> dict: _DMD_COMPLETE_MARKER = "dmd2_complete.marker" +@dataclasses.dataclass(frozen=True) +class _QuantSettings: + """Resolved ``dmd2.quant`` block — restore-only QAT of the student (no calibration). + + Attributes: + enabled: When ``True`` the student is quantized by RESTORING a ModelOpt quantizer + state from disk (recipe + frozen amax). The trainer never calibrates. + quant_state_path: Path to the quantizer-state file produced by the + ``examples/diffusers/quantization`` calibration example + (``--quantized-torch-ckpt-save-path`` → ``transformer.pt``). Used on the FIRST + launch (warm-start) to quantize the FP student. Required when ``enabled``. + init_weights_from: Optional path to the full-precision DMD2 checkpoint to + warm-start the student / fake_score / discriminator / EMA / optimizers from + when the run's own ``checkpoint_dir`` has no QAT checkpoint yet. The amax in + ``quant_state_path`` must have been calibrated against this checkpoint's + student weights. + """ + + enabled: bool + quant_state_path: str | None + init_weights_from: str | None + + class DMD2DiffusionRecipe(TrainDiffusionRecipe): """DMD2 recipe that reuses ``TrainDiffusionRecipe`` for the student path. @@ -379,6 +495,7 @@ def run_train_validation_loop(self) -> None: neg_text_embeds, neg_text_mask, ) = self._prepare_micro_batch(micro_batch) + model_kwargs = self._prepare_model_kwargs(micro_batch) if is_student_phase: # ``compute_student_loss`` reads ``guidance_scale`` from the @@ -394,6 +511,7 @@ def run_train_validation_loop(self) -> None: negative_encoder_hidden_states=neg_text_embeds, negative_encoder_hidden_states_mask=neg_text_mask, guidance_scale=None, + **model_kwargs, ) micro_vsd_losses.append(float(losses["vsd"].item())) else: @@ -402,6 +520,7 @@ def run_train_validation_loop(self) -> None: noise, encoder_hidden_states=text_embeds, encoder_hidden_states_mask=text_mask, + **model_kwargs, ) (losses["total"] / len(batch_group)).backward() @@ -421,6 +540,7 @@ def run_train_validation_loop(self) -> None: noise, encoder_hidden_states=text_embeds, encoder_hidden_states_mask=text_mask, + **model_kwargs, ) (disc_losses["total"] / len(batch_group)).backward() # Manual gradient all-reduce across DP ranks (the @@ -538,10 +658,33 @@ def load_checkpoint(self, restore_from: str | None = None): # ``nemo_automodel`` can be used unmodified. make_optimizer_partial_load_tolerant(self.checkpointer) + quant = self._resolve_quant_settings() + resolved = self._resolve_complete_dmd_checkpoint(restore_from) + + # QAT first launch: the run's own checkpoint_dir has no QAT checkpoint yet, so + # warm-start the FP student / fake_score / EMA / optimizers from the full-precision + # checkpoint named by ``dmd2.quant.init_weights_from`` (the FP run the amax was + # calibrated against). On later resumes ``resolved`` already points at a QAT + # checkpoint in checkpoint_dir, so this branch is skipped. + if quant.enabled and resolved is None and quant.init_weights_from: + resolved = self._resolve_complete_dmd_checkpoint(quant.init_weights_from) + if is_main_process(): + logging.info( + "[DMD2][qat] no QAT checkpoint in checkpoint_dir; warm-starting FP " + "state from dmd2.quant.init_weights_from=%s", + quant.init_weights_from, + ) + self.__dict__["_dmd2_resolved_restore_from"] = resolved if resolved is None: + if quant.enabled and is_main_process(): + logging.warning( + "[DMD2][qat] QAT enabled but no checkpoint resolved (checkpoint_dir empty " + "and dmd2.quant.init_weights_from unset). Quantizing a fresh base model — " + "its weights will NOT match the calibrated amax." + ) if ( restore_from is not None and str(restore_from).upper() == "LATEST" @@ -552,10 +695,137 @@ def load_checkpoint(self, restore_from: str | None = None): "Starting fresh.", self.checkpointer.config.checkpoint_dir, ) + # Even with no weights to restore, honor QAT by quantizing from the recipe file. + if quant.enabled: + self._quantize_student(quant.quant_state_path) return + # Single restore path. The student-weight checkpoints are always clean FP (the QAT + # save hides the quantizer buffers — see ``_student_quant_buffers_hidden``), so the + # strict DCP load matches an unquantized student whether this is a warm-start from + # the FP checkpoint or a resume from a QAT checkpoint. Quantization always happens + # AFTER the weights are in place. super().load_checkpoint(resolved) + if quant.enabled: + # Restore-only QAT never recalibrates, so the recipe + amax are invariant for the + # whole run — re-apply the same configured quantizer-state file on every (re)start + # rather than persisting an unchanging copy per checkpoint. + self._quantize_student(quant.quant_state_path) + + def _resolve_quant_settings(self) -> _QuantSettings: + """Parse and cache the ``dmd2.quant`` block (restore-only student QAT).""" + cached = self.__dict__.get("_quant_settings") + if cached is not None: + return cached + + node = self.cfg.get("dmd2.quant", None) + if node is None: + settings = _QuantSettings(enabled=False, quant_state_path=None, init_weights_from=None) + self.__dict__["_quant_settings"] = settings + return settings + + d = node.to_dict() if hasattr(node, "to_dict") else dict(node) + enabled = bool(d.get("enabled", False)) + quant_state_path = d.get("quant_state_path") + init_weights_from = d.get("init_weights_from") + if enabled and not quant_state_path: + raise ValueError( + "dmd2.quant.enabled is true but dmd2.quant.quant_state_path is not set. Point it " + "at the quantizer-state file produced by examples/diffusers/quantization " + "(its --quantized-torch-ckpt-save-path, e.g. .../quant/transformer.pt)." + ) + settings = _QuantSettings( + enabled=enabled, + quant_state_path=quant_state_path, + init_weights_from=init_weights_from, + ) + self.__dict__["_quant_settings"] = settings + return settings + + def _quantize_student(self, quant_state_path: str | None) -> None: + """Quantize the student by RESTORING a ModelOpt quantizer state from disk. + + Restore-only: re-inserts the quantizer modules from the saved recipe and loads the + precomputed, frozen amax. No ``mtq.quantize`` / calibration forward pass is ever + run, so amax stays exactly as it was on disk for the whole training run. Only the + student (``self.model``) is quantized — the frozen teacher and the trainable + fake_score stay full precision so the distribution-matching gradient is exact. + + Safe after FSDP2 wrapping: the conversion is an in-place ``__class__`` swap that + preserves the existing ``weight``/``bias`` Parameter objects, so the already-built + student optimizer's references stay valid (verified by ModelOpt's + ``tests/gpu/torch/quantization/test_fsdp2.py``). + """ + if not quant_state_path: + raise ValueError( + "[DMD2][qat] _quantize_student called without a quant_state_path. Set " + "dmd2.quant.quant_state_path." + ) + if not os.path.isfile(quant_state_path): + raise FileNotFoundError( + f"[DMD2][qat] quantizer-state file not found: {quant_state_path}. QAT here is " + "restore-only (no on-the-fly calibration); produce it with " + "examples/diffusers/quantization first." + ) + + if is_main_process(): + logging.info( + "[DMD2][qat] restoring student quantizer state (recipe + frozen amax) <- %s", + quant_state_path, + ) + restore_quantizer_state(self.model, quant_state_path) + + # amax (and any other quantizer buffers) come off disk on CPU. Move just the + # TensorQuantizer buffers onto the student device — these modules carry no + # parameters, so this never touches the FSDP2 DTensor weights. + from modelopt.torch.quantization.nn import TensorQuantizer + + for module in self.model.modules(): + if isinstance(module, TensorQuantizer): + module.to(self.device) + + if is_main_process(): + import modelopt.torch.quantization as mtq + + logging.info("[DMD2][qat] student quantized (restore-only). Quantizer summary:") + mtq.print_quant_summary(self.model) + + @contextlib.contextmanager + def _student_quant_buffers_hidden(self): + """Temporarily mark the student's quantizer buffers non-persistent. + + Wraps the parent ``save_checkpoint`` so the student's DCP shards AND the + consolidated/diffusers export stay clean full-precision (``ModelState.state_dict()`` + — and hence the consolidated index, which re-adds every state_dict key + (checkpointing.py:941-948) — excludes the ``amax`` buffers). The frozen amax is not + persisted per checkpoint at all: it is re-applied from ``dmd2.quant.quant_state_path`` + on every (re)start. Keeping the saved student weights amax-free both yields a clean + ``model/consolidated`` (a normal ``QwenImageTransformer2DModel``, re-quantizable for + deploy/eval) and makes every restore a clean ``load FP weights -> quantize`` path + (the strict DCP load matches an unquantized student). Mirrors ModelOpt's own trick in + ``quantization/plugins/transformers_trainer.py`` (``_modelopt_prepare``). + + No-op when QAT is disabled. + """ + if not self._resolve_quant_settings().enabled: + yield + return + + from modelopt.torch.quantization.nn import TensorQuantizer + + saved: list[tuple[TensorQuantizer, set[str]]] = [] + for module in self.model.modules(): + if isinstance(module, TensorQuantizer): + saved.append((module, set(module._non_persistent_buffers_set))) + module._non_persistent_buffers_set.update(module._buffers.keys()) + try: + yield + finally: + for module, original in saved: + module._non_persistent_buffers_set.clear() + module._non_persistent_buffers_set.update(original) + def save_checkpoint( self, epoch: int, @@ -588,7 +858,11 @@ def save_checkpoint( self.checkpointer.config.checkpoint_dir ) - super().save_checkpoint(epoch, step, train_loss, val_loss, best_metric_key) + # Hide the student's quantizer buffers during the parent save so the DCP shards and + # consolidated/diffusers export stay clean FP. Frozen amax remains external in + # ``dmd2.quant.quant_state_path`` and is re-applied on restore. No-op when QAT is disabled. + with self._student_quant_buffers_hidden(): + super().save_checkpoint(epoch, step, train_loss, val_loss, best_metric_key) if not self.checkpointer.config.enabled: return @@ -861,6 +1135,9 @@ def _is_dmd_checkpoint_complete(self, path: str) -> bool: if not complete: return False + # QAT adds no per-checkpoint artifact (the quantizer recipe + frozen amax are + # re-applied from dmd2.quant.quant_state_path on every (re)start), so QAT + # checkpoints use the same completeness criteria as full-precision ones. if self._cfg_gan_enabled(): return os.path.isfile(os.path.join(path, "discriminator.pt")) and os.path.isfile( os.path.join(path, "discriminator_optimizer.pt") @@ -976,10 +1253,10 @@ def _build_discriminator_optimizer(self) -> torch.optim.Optimizer | None: def _attach_gan_feature_capture(self) -> None: """Install Qwen-Image feature-capture hooks on the teacher when GAN is enabled. - Reads the latent resolution from the dataloader so the hook can reshape + Reads an initial latent resolution from the dataloader so the hook can reshape ``[B, num_image_patches, 3072]`` into ``[B, 3072, H_lat//2, W_lat//2]``. - Mock dataloader → spatial_h/spatial_w from the YAML. Real dataloader → - base_resolution / vae_scale. + The Qwen plugin refreshes that shape before every teacher forward, so real + multiresolution batches are captured using their actual target dimensions. """ feature_indices = list(self.cfg.get("dmd2.gan_feature_indices", [30])) @@ -1020,6 +1297,11 @@ def _attach_gan_feature_capture(self) -> None: feature_indices=feature_indices, h_lat=h_lat, w_lat=w_lat, + # Edit models append reference-image tokens after the target tokens. The GAN + # discriminator is defined on the generated target only, so capture its prefix. + target_prefix_only=( + self._resolve_pipeline_cls().__name__ == "QwenImageEditDMDPipeline" + ), ) if is_main_process(): logging.info( @@ -1123,13 +1405,18 @@ def _resolve_pipeline_cls(self) -> type[DMDPipeline]: from modelopt.torch.fastgen.plugins.qwen_image import QwenImageDMDPipeline return QwenImageDMDPipeline + if explicit == "qwen_image_edit": + from modelopt.torch.fastgen.plugins.qwen_image_edit import QwenImageEditDMDPipeline + + return QwenImageEditDMDPipeline raise ValueError( - f"Unknown dmd2.pipeline_plugin={explicit!r}. Supported: null/'base', 'qwen_image'." + f"Unknown dmd2.pipeline_plugin={explicit!r}. Supported: null/'base', " + "'qwen_image', 'qwen_image_edit'." ) def _resolve_pipeline_kwargs(self, pipeline_cls: type[DMDPipeline]) -> dict[str, Any]: """Extra kwargs to forward to the pipeline subclass constructor (plugin-specific).""" - if pipeline_cls.__name__ == "QwenImageDMDPipeline": + if pipeline_cls.__name__ in {"QwenImageDMDPipeline", "QwenImageEditDMDPipeline"}: # Optional ``guidance`` value passed to the transformer's guidance kwarg every # call. Independent of DMDConfig.guidance_scale (which drives the negative- # prompt CFG path on the teacher). Leave ``None`` to skip the embedding when @@ -1289,6 +1576,59 @@ def _prepare_micro_batch( noise = torch.randn_like(latents) return latents, noise, text_embeds, text_mask, negative_text_embeds, negative_text_mask + def _prepare_model_kwargs(self, micro_batch: dict[str, Any]) -> dict[str, Any]: + """Move optional model-specific conditioning to the training device. + + Qwen-Image-Edit caches one or more VAE-encoded reference images separately from + the clean target latent. The edit plugin packs these tensors and appends them to + every student / teacher / fake-score forward. Keeping the data as a list allows + references with different aspect ratios while retaining a regular batch dimension + for each reference slot. + """ + conditioning = micro_batch.get("conditioning_latents") + if conditioning is None: + return {} + + if torch.is_tensor(conditioning): + if conditioning.ndim == 4: + conditioning = [conditioning] + elif conditioning.ndim == 5: + # Collates may represent refs as [B, N, C, H, W]. Convert to the + # plugin's list-of-[B,C,H,W] contract. + conditioning = list(conditioning.unbind(dim=1)) + else: + raise ValueError( + "conditioning_latents tensor must be [B,C,H,W] or [B,N,C,H,W], " + f"got shape {tuple(conditioning.shape)}." + ) + elif isinstance(conditioning, (list, tuple)): + conditioning = list(conditioning) + else: + raise TypeError( + "conditioning_latents must be a tensor or a list/tuple of tensors; " + f"got {type(conditioning).__name__}." + ) + + if not conditioning: + raise ValueError("conditioning_latents must contain at least one reference image.") + + expected_batch = int(micro_batch["image_latents"].shape[0]) + prepared: list[torch.Tensor] = [] + for ref_index, ref in enumerate(conditioning): + if not torch.is_tensor(ref) or ref.ndim != 4: + shape = tuple(ref.shape) if torch.is_tensor(ref) else None + raise ValueError( + f"conditioning_latents[{ref_index}] must be [B,C,H,W], got {shape}." + ) + if int(ref.shape[0]) != expected_batch: + raise ValueError( + f"conditioning_latents[{ref_index}] batch={ref.shape[0]} does not match " + f"target batch={expected_batch}." + ) + prepared.append(ref.to(self.device, dtype=self.bf16, non_blocking=True)) + + return {"conditioning_latents": tuple(prepared)} + def _log_step( self, *, diff --git a/examples/diffusers/fastgen/fastgen_data/__init__.py b/examples/diffusers/fastgen/fastgen_data/__init__.py index 771b93b1c0b..52d7cb6675a 100644 --- a/examples/diffusers/fastgen/fastgen_data/__init__.py +++ b/examples/diffusers/fastgen/fastgen_data/__init__.py @@ -38,9 +38,12 @@ # Convert a missing-helper ImportError into an actionable message naming the supported range. try: from .collate_fns import ( + build_image_to_image_multiresolution_dataloader, build_text_to_image_multiresolution_dataloader, + collate_fn_image_to_image, collate_fn_text_to_image, ) + from .image_to_image_dataset import ImageToImageDataset from .text_to_image_dataset import TextToImageDataset except ImportError as exc: # pragma: no cover - environment guard raise ImportError( @@ -53,8 +56,11 @@ ) from exc __all__ = [ + "ImageToImageDataset", "TextToImageDataset", + "build_image_to_image_multiresolution_dataloader", "build_text_to_image_multiresolution_dataloader", + "collate_fn_image_to_image", "collate_fn_text_to_image", ] diff --git a/examples/diffusers/fastgen/fastgen_data/collate_fns.py b/examples/diffusers/fastgen/fastgen_data/collate_fns.py index d669d2a7c4a..4049b646ded 100644 --- a/examples/diffusers/fastgen/fastgen_data/collate_fns.py +++ b/examples/diffusers/fastgen/fastgen_data/collate_fns.py @@ -38,11 +38,116 @@ from nemo_automodel.components.datasets.diffusion.sampler import SequentialBucketSampler from torchdata.stateful_dataloader import StatefulDataLoader +from .image_to_image_dataset import ImageToImageDataset from .text_to_image_dataset import TextToImageDataset logger = logging.getLogger(__name__) +def _pad_text_conditioning( + batch: list[dict], + embedding_key: str, + mask_key: str, +) -> tuple[torch.Tensor, torch.Tensor]: + """Right-pad variable-length embeddings and masks from cached edit samples.""" + + embeddings = [item[embedding_key] for item in batch] + masks = [item[mask_key] for item in batch] + if any(not torch.is_tensor(value) or value.ndim != 2 for value in embeddings): + shapes = [getattr(value, "shape", None) for value in embeddings] + raise ValueError(f"{embedding_key} values must have shape [seq,dim], got {shapes}") + hidden_dims = {value.shape[1] for value in embeddings} + dtypes = {value.dtype for value in embeddings} + if len(hidden_dims) != 1 or len(dtypes) != 1: + raise ValueError( + f"{embedding_key} hidden dimensions/dtypes must match across a batch: " + f"dims={hidden_dims}, dtypes={dtypes}" + ) + + max_length = max(value.shape[0] for value in embeddings) + hidden_dim = embeddings[0].shape[1] + padded = embeddings[0].new_zeros((len(batch), max_length, hidden_dim)) + padded_mask = torch.zeros((len(batch), max_length), dtype=torch.long) + for index, (embedding, mask) in enumerate(zip(embeddings, masks)): + if not torch.is_tensor(mask) or mask.ndim != 1 or mask.shape[0] != embedding.shape[0]: + raise ValueError( + f"{mask_key} for sample {index} must have shape [{embedding.shape[0]}], " + f"got {getattr(mask, 'shape', None)}" + ) + length = embedding.shape[0] + padded[index, :length] = embedding + padded_mask[index, :length] = mask.to(dtype=torch.long) + return padded, padded_mask + + +def collate_fn_image_to_image(batch: list[dict]) -> dict: + """Build a Qwen-Image-Edit batch and validate reference compatibility. + + ``conditioning_latents`` remains a list, one entry per reference image. Each entry is a + stacked ``[B,C,H,W]`` tensor, which preserves support for references with different aspect + ratios while ensuring a given reference slot is stackable across the batch. + """ + + if not batch: + raise ValueError("Cannot collate an empty image-to-image batch") + resolutions = {tuple(item["crop_resolution"].tolist()) for item in batch} + if len(resolutions) != 1: + raise ValueError(f"Mixed target resolutions in batch: {resolutions}") + + reference_counts = {len(item["conditioning_latents"]) for item in batch} + if len(reference_counts) != 1: + raise ValueError(f"Mixed conditioning-image counts in batch: {reference_counts}") + reference_count = reference_counts.pop() + if reference_count < 1: + raise ValueError("Every image-to-image sample must contain at least one reference") + + conditioning_latents = [] + for reference_index in range(reference_count): + shapes = {tuple(item["conditioning_latents"][reference_index].shape) for item in batch} + if len(shapes) != 1: + raise ValueError( + f"Reference {reference_index} has mixed latent shapes in batch: {shapes}" + ) + conditioning_latents.append( + torch.stack([item["conditioning_latents"][reference_index] for item in batch]) + ) + + text_embeddings, text_mask = _pad_text_conditioning( + batch, "prompt_embeds", "prompt_embeds_mask" + ) + negative_embeddings, negative_mask = _pad_text_conditioning( + batch, "negative_prompt_embeds", "negative_prompt_embeds_mask" + ) + + image_batch = { + "image_latents": torch.stack([item["latent"] for item in batch]), + "conditioning_latents": conditioning_latents, + "data_type": "image_edit", + "text_embeddings": text_embeddings, + "text_embeddings_mask": text_mask, + "negative_text_embeddings": negative_embeddings, + "negative_text_embeddings_mask": negative_mask, + "metadata": { + "sample_ids": [item["sample_id"] for item in batch], + "prompts": [item["prompt"] for item in batch], + "negative_prompts": [item["negative_prompt"] for item in batch], + "image_paths": [item["image_path"] for item in batch], + "conditioning_image_paths": [item["conditioning_image_paths"] for item in batch], + "conditioning_resolutions": [item["conditioning_resolutions"] for item in batch], + "target_latent_shapes": [item["target_latent_shape"] for item in batch], + "conditioning_latent_shapes": [item["conditioning_latent_shapes"] for item in batch], + "bucket_ids": [item["bucket_id"] for item in batch], + "aspect_ratios": [item["aspect_ratio"] for item in batch], + "crop_resolution": torch.stack([item["crop_resolution"] for item in batch]), + "original_resolution": torch.stack([item["original_resolution"] for item in batch]), + "crop_offset": torch.stack([item["crop_offset"] for item in batch]), + }, + } + if "source_metadata" in batch[0]: + image_batch["metadata"]["source_metadata"] = [item.get("source_metadata") for item in batch] + return image_batch + + def collate_fn_text_to_image( batch: list[dict], negative_text_embeddings: torch.Tensor | None = None, @@ -240,3 +345,69 @@ def build_text_to_image_multiresolution_dataloader( dp_world_size, ) return dataloader, sampler + + +def build_image_to_image_multiresolution_dataloader( + *, + cache_dir: str, + train_text_encoder: bool = False, + batch_size: int = 1, + dp_rank: int = 0, + dp_world_size: int = 1, + base_resolution: tuple[int, int] = (256, 256), + drop_last: bool = True, + shuffle: bool = True, + dynamic_batch_size: bool = False, + num_workers: int = 4, + pin_memory: bool = True, + prefetch_factor: int = 2, +) -> tuple[StatefulDataLoader, SequentialBucketSampler]: + """Build the cached Qwen-Image-Edit multiresolution dataloader. + + Positive and negative embeddings are both sample-specific because each contains visual + tokens from that sample's references, so this builder intentionally has no static negative + prompt embedding argument. + """ + + dataset = ImageToImageDataset( + cache_dir=cache_dir, + train_text_encoder=train_text_encoder, + ) + sampler = SequentialBucketSampler( + dataset, + base_batch_size=batch_size, + base_resolution=base_resolution, + drop_last=drop_last, + shuffle_buckets=shuffle, + shuffle_within_bucket=shuffle, + dynamic_batch_size=dynamic_batch_size, + num_replicas=dp_world_size, + rank=dp_rank, + ) + dataloader = StatefulDataLoader( + dataset, + batch_sampler=sampler, + collate_fn=collate_fn_image_to_image, + num_workers=num_workers, + pin_memory=pin_memory, + prefetch_factor=prefetch_factor if num_workers > 0 else None, + persistent_workers=num_workers > 0, + ) + logger.info( + "image-to-image dataloader | cache_dir=%s size=%d batches/epoch=%d batch_size=%d dp=%d/%d", + cache_dir, + len(dataset), + len(sampler), + batch_size, + dp_rank, + dp_world_size, + ) + return dataloader, sampler + + +__all__ = [ + "build_image_to_image_multiresolution_dataloader", + "build_text_to_image_multiresolution_dataloader", + "collate_fn_image_to_image", + "collate_fn_text_to_image", +] diff --git a/examples/diffusers/fastgen/fastgen_data/image_to_image_dataset.py b/examples/diffusers/fastgen/fastgen_data/image_to_image_dataset.py new file mode 100644 index 00000000000..e12d7d7289c --- /dev/null +++ b/examples/diffusers/fastgen/fastgen_data/image_to_image_dataset.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Cached image-to-image dataset for Qwen-Image-Edit DMD2 training.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import torch +from nemo_automodel.components.datasets.diffusion.base_dataset import BaseMultiresolutionDataset + + +def _remove_cache_batch_dim( + tensor: torch.Tensor, + name: str, + unbatched_ndim: int, +) -> torch.Tensor: + """Remove the singleton encoder batch dimension while rejecting malformed caches.""" + + if not torch.is_tensor(tensor): + raise TypeError(f"Cached {name!r} must be a tensor, got {type(tensor).__name__}") + if tensor.ndim == unbatched_ndim + 1 and tensor.shape[0] == 1: + tensor = tensor.squeeze(0) + if tensor.ndim != unbatched_ndim: + raise ValueError( + f"Cached {name!r} must be {unbatched_ndim}D after removing an optional " + f"singleton batch dimension, got shape {tuple(tensor.shape)}" + ) + return tensor + + +class ImageToImageDataset(BaseMultiresolutionDataset): + """Read target/reference latents and per-sample multimodal prompt embeddings.""" + + def __init__(self, cache_dir: str, train_text_encoder: bool = False): + if train_text_encoder: + raise NotImplementedError( + "Qwen-Image-Edit requires cached multimodal embeddings; on-the-fly text encoder " + "training is not supported by ImageToImageDataset." + ) + self.train_text_encoder = False + super().__init__(cache_dir, quantization=64) + + def _validated_cache_file(self, item: dict[str, Any]) -> Path: + cache_file = Path(item["cache_file"]).resolve() + cache_dir = Path(self.cache_dir).resolve() + try: + cache_file.relative_to(cache_dir) + except ValueError as exc: + raise ValueError( + f"Cache file {cache_file} is outside cache directory {cache_dir}" + ) from exc + return cache_file + + def __getitem__(self, idx: int) -> dict[str, Any]: + item = self.metadata[idx] + data = torch.load( + self._validated_cache_file(item), + map_location="cpu", + weights_only=True, + ) + target_latent = data.get("latent") + if not torch.is_tensor(target_latent) or target_latent.ndim != 3: + raise ValueError(f"Cache item {idx} target latent must have shape [C,H,W]") + + conditioning_latents = data.get("conditioning_latents") + if not isinstance(conditioning_latents, list) or not conditioning_latents: + raise ValueError( + f"Cache item {idx} must contain a non-empty `conditioning_latents` list" + ) + if not all(torch.is_tensor(latent) and latent.ndim == 3 for latent in conditioning_latents): + raise ValueError(f"Cache item {idx} conditioning latents must all have shape [C,H,W]") + + resolution_key = "bucket_resolution" if "bucket_resolution" in item else "crop_resolution" + prompt_embeds = _remove_cache_batch_dim( + data["prompt_embeds"], "prompt_embeds", unbatched_ndim=2 + ) + negative_prompt_embeds = _remove_cache_batch_dim( + data["negative_prompt_embeds"], "negative_prompt_embeds", unbatched_ndim=2 + ) + prompt_mask = data.get("prompt_embeds_mask") + if prompt_mask is None: + prompt_mask = torch.ones(prompt_embeds.shape[0], dtype=torch.long) + else: + prompt_mask = _remove_cache_batch_dim( + prompt_mask, "prompt_embeds_mask", unbatched_ndim=1 + ).long() + negative_mask = data.get("negative_prompt_embeds_mask") + if negative_mask is None: + negative_mask = torch.ones(negative_prompt_embeds.shape[0], dtype=torch.long) + else: + negative_mask = _remove_cache_batch_dim( + negative_mask, + "negative_prompt_embeds_mask", + unbatched_ndim=1, + ).long() + + output = { + "latent": target_latent, + "conditioning_latents": conditioning_latents, + "prompt_embeds": prompt_embeds, + "prompt_embeds_mask": prompt_mask, + "negative_prompt_embeds": negative_prompt_embeds, + "negative_prompt_embeds_mask": negative_mask, + "crop_resolution": torch.tensor(item[resolution_key]), + "original_resolution": torch.tensor(item["original_resolution"]), + "crop_offset": torch.tensor(data["crop_offset"]), + "prompt": data["prompt"], + "negative_prompt": data.get("negative_prompt", " "), + "image_path": data["image_path"], + "conditioning_image_paths": data["conditioning_image_paths"], + "conditioning_resolutions": data.get( + "conditioning_resolutions", + [None] * len(conditioning_latents), + ), + "target_latent_shape": data.get("target_latent_shape", tuple(target_latent.shape)), + "conditioning_latent_shapes": data.get( + "conditioning_latent_shapes", + [tuple(value.shape) for value in conditioning_latents], + ), + "sample_id": data.get("sample_id", str(idx)), + "bucket_id": item["bucket_id"], + "aspect_ratio": item.get("aspect_ratio", 1.0), + } + if "source_metadata" in data: + output["source_metadata"] = data["source_metadata"] + return output + + +__all__ = ["ImageToImageDataset"] diff --git a/examples/diffusers/fastgen/inference_dmd2_qwen_image.py b/examples/diffusers/fastgen/inference_dmd2_qwen_image.py index 5907d0f1b86..7297679db8d 100644 --- a/examples/diffusers/fastgen/inference_dmd2_qwen_image.py +++ b/examples/diffusers/fastgen/inference_dmd2_qwen_image.py @@ -294,15 +294,6 @@ def __call__( num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, ) - txt_seq_lens = ( - prompt_embeds_mask.sum(dim=1).int().tolist() if prompt_embeds_mask is not None else None - ) - neg_txt_seq_lens = ( - neg_prompt_embeds_mask.sum(dim=1).int().tolist() - if neg_prompt_embeds_mask is not None - else None - ) - # ---- 3. Build initial noisy latents at t = schedule[0] --------------- if isinstance(prompt, str): batch_size = 1 @@ -334,7 +325,6 @@ def __call__( encoder_hidden_states_mask=prompt_embeds_mask, timestep=timestep, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, guidance=None, return_dict=False, )[0] @@ -352,7 +342,6 @@ def __call__( encoder_hidden_states_mask=neg_prompt_embeds_mask, timestep=timestep, img_shapes=img_shapes, - txt_seq_lens=neg_txt_seq_lens, guidance=None, return_dict=False, )[0] diff --git a/examples/diffusers/fastgen/inference_dmd2_qwen_image_edit.py b/examples/diffusers/fastgen/inference_dmd2_qwen_image_edit.py new file mode 100644 index 00000000000..25bdf687142 --- /dev/null +++ b/examples/diffusers/fastgen/inference_dmd2_qwen_image_edit.py @@ -0,0 +1,330 @@ +# 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. + +"""Few-step inference for a DMD2-trained Qwen-Image-Edit-2511 student. + +The stock EditPlus pipeline is reused for multimodal prompt encoding, reference-image +preprocessing, VAE encode/decode, and output postprocessing. Only its denoising loop is +replaced with the exact rectified-flow schedule used by DMD2 training. Target tokens are +followed by the fixed reference-image tokens on every transformer call; only the target +prediction prefix is stepped. +""" + +from __future__ import annotations + +import argparse +import itertools +import logging +import math +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from diffusers import QwenImageEditPlusPipeline, QwenImageTransformer2DModel +from diffusers.utils import load_image +from diffusers.utils.torch_utils import randn_tensor + +logger = logging.getLogger(__name__) + +_CONDITION_IMAGE_AREA = 384 * 384 +_VAE_IMAGE_AREA = 1024 * 1024 + + +def _calculate_dimensions(target_area: int, ratio: float) -> tuple[int, int]: + """Match Diffusers EditPlus' area-preserving, 32-pixel-quantized resize.""" + raw_width = math.sqrt(target_area * ratio) + raw_height = raw_width / ratio + return round(raw_width / 32) * 32, round(raw_height / 32) * 32 + + +def _overlay_ema(student: torch.nn.Module, ema_path: str | os.PathLike[str]) -> None: + payload = torch.load(str(ema_path), map_location="cpu", weights_only=True) + shadow = payload.get("shadow", payload) if isinstance(payload, dict) else payload + if not isinstance(shadow, dict): + raise TypeError(f"EMA payload must be a state dict, got {type(shadow).__name__}.") + missing, unexpected = student.load_state_dict(shadow, strict=False) + if missing or unexpected: + logger.warning("EMA overlay: %d missing, %d unexpected keys", len(missing), len(unexpected)) + + +@dataclass +class QwenImageEditDMDOutput: + images: list[Any] + + +class QwenImageEditDMDInferencePipeline: + """DMD sampler around a stock :class:`QwenImageEditPlusPipeline`.""" + + def __init__(self, pipeline: QwenImageEditPlusPipeline, max_t: float = 0.999) -> None: + self._pipe = pipeline + self.max_t = float(max_t) + + @classmethod + def from_pretrained( + cls, + student_path: str | os.PathLike[str], + base_pipeline_path: str | os.PathLike[str] = "Qwen/Qwen-Image-Edit-2511", + *, + ema_path: str | os.PathLike[str] | None = None, + torch_dtype: torch.dtype = torch.bfloat16, + max_t: float = 0.999, + ) -> QwenImageEditDMDInferencePipeline: + student_path = str(student_path) + if not os.path.isdir(student_path): + raise FileNotFoundError(f"student_path is not a directory: {student_path}") + student = QwenImageTransformer2DModel.from_pretrained(student_path, torch_dtype=torch_dtype) + if ema_path is not None: + _overlay_ema(student, ema_path) + student.eval() + pipeline = QwenImageEditPlusPipeline.from_pretrained( + str(base_pipeline_path), transformer=student, torch_dtype=torch_dtype + ) + return cls(pipeline, max_t=max_t) + + def to(self, device: str | torch.device) -> QwenImageEditDMDInferencePipeline: + self._pipe.to(device) + return self + + @property + def device(self) -> torch.device: + return self._pipe.transformer.device + + @property + def dtype(self) -> torch.dtype: + return next(self._pipe.transformer.parameters()).dtype + + @staticmethod + def _resolve_schedule( + num_inference_steps: int, + max_t: float, + t_list: list[float] | None, + ) -> list[float]: + if num_inference_steps < 1: + raise ValueError("num_inference_steps must be >= 1.") + if t_list is None: + return torch.linspace(max_t, 0.0, num_inference_steps + 1).tolist() + if len(t_list) != num_inference_steps + 1: + raise ValueError("t_list must contain num_inference_steps + 1 entries.") + schedule = [float(value) for value in t_list] + if abs(schedule[-1]) > 1e-6: + raise ValueError("t_list must end at 0.0.") + if any(left <= right for left, right in itertools.pairwise(schedule)): + raise ValueError("t_list must be strictly decreasing.") + return schedule + + @torch.no_grad() + def __call__( + self, + image: Any | list[Any], + prompt: str, + *, + negative_prompt: str | None = None, + num_inference_steps: int = 4, + guidance_scale: float = 1.0, + height: int | None = None, + width: int | None = None, + generator: torch.Generator | None = None, + max_t: float | None = None, + t_list: list[float] | None = None, + sample_type: str = "ode", + output_type: str = "pil", + ) -> QwenImageEditDMDOutput: + """Edit one image from one or more ordered references. + + A CFG-trained DMD2 student has already internalized teacher guidance, so the + default ``guidance_scale=1`` performs a single transformer call per step. + """ + if sample_type not in {"ode", "sde"}: + raise ValueError("sample_type must be 'ode' or 'sde'.") + references = list(image) if isinstance(image, (list, tuple)) else [image] + if not references: + raise ValueError("At least one reference image is required.") + + pipe = self._pipe + device, dtype = self.device, self.dtype + max_t = self.max_t if max_t is None else float(max_t) + schedule = self._resolve_schedule(num_inference_steps, max_t, t_list) + + # Match QwenImageEditPlusPipeline.__call__: the last reference determines the + # default target aspect ratio; each reference gets separate vision/VAE resolutions. + last_width, last_height = references[-1].size + default_width, default_height = _calculate_dimensions( + _VAE_IMAGE_AREA, last_width / last_height + ) + width = int(width or default_width) + height = int(height or default_height) + multiple = pipe.vae_scale_factor * 2 + width, height = width // multiple * multiple, height // multiple * multiple + + condition_images: list[Any] = [] + vae_images: list[torch.Tensor] = [] + vae_sizes: list[tuple[int, int]] = [] + for reference in references: + ref_width, ref_height = reference.size + ratio = ref_width / ref_height + cond_width, cond_height = _calculate_dimensions(_CONDITION_IMAGE_AREA, ratio) + vae_width, vae_height = _calculate_dimensions(_VAE_IMAGE_AREA, ratio) + condition_images.append(pipe.image_processor.resize(reference, cond_height, cond_width)) + vae_images.append( + pipe.image_processor.preprocess(reference, vae_height, vae_width).unsqueeze(2) + ) + vae_sizes.append((vae_width, vae_height)) + + prompt_embeds, prompt_mask = pipe.encode_prompt( + image=condition_images, prompt=prompt, device=device, num_images_per_prompt=1 + ) + do_cfg = guidance_scale != 1.0 + negative_embeds = negative_mask = None + if do_cfg: + negative_prompt = " " if negative_prompt is None else negative_prompt + negative_embeds, negative_mask = pipe.encode_prompt( + image=condition_images, + prompt=negative_prompt, + device=device, + num_images_per_prompt=1, + ) + + channels = pipe.transformer.config.in_channels // 4 + h_lat = 2 * (height // (pipe.vae_scale_factor * 2)) + w_lat = 2 * (width // (pipe.vae_scale_factor * 2)) + noise = randn_tensor( + (1, 1, channels, h_lat, w_lat), generator=generator, device=device, dtype=dtype + ) + target = pipe._pack_latents(noise * schedule[0], 1, channels, h_lat, w_lat) + target_tokens = target.shape[1] + + packed_references: list[torch.Tensor] = [] + for vae_image in vae_images: + ref_latent = pipe._encode_vae_image( + vae_image.to(device=device, dtype=dtype), generator=generator + ) + ref_h, ref_w = ref_latent.shape[3:] + packed_references.append(pipe._pack_latents(ref_latent, 1, channels, ref_h, ref_w)) + img_shapes = [ + [ + (1, h_lat // 2, w_lat // 2), + *[ + ( + 1, + vae_height // pipe.vae_scale_factor // 2, + vae_width // pipe.vae_scale_factor // 2, + ) + for vae_width, vae_height in vae_sizes + ], + ] + ] + + x = target + fixed_references = torch.cat(packed_references, dim=1) + for t_cur, t_next in itertools.pairwise(schedule): + model_input = torch.cat([x, fixed_references], dim=1) + timestep = torch.full((1,), float(t_cur), device=device, dtype=dtype) + flow = pipe.transformer( + hidden_states=model_input, + timestep=timestep, + guidance=None, + encoder_hidden_states_mask=prompt_mask, + encoder_hidden_states=prompt_embeds, + img_shapes=img_shapes, + return_dict=False, + )[0][:, :target_tokens] + if do_cfg: + negative_flow = pipe.transformer( + hidden_states=model_input, + timestep=timestep, + guidance=None, + encoder_hidden_states_mask=negative_mask, + encoder_hidden_states=negative_embeds, + img_shapes=img_shapes, + return_dict=False, + )[0][:, :target_tokens] + flow = ( + negative_flow.double() + + float(guidance_scale) * (flow.double() - negative_flow.double()) + ).to(dtype) + + x0 = (x.double() - float(t_cur) * flow.double()).to(dtype) + if t_next <= 1e-6: + x = x0 + continue + if sample_type == "ode": + eps = ( + (x.double() - (1.0 - float(t_cur)) * x0.double()) / max(float(t_cur), 1e-6) + ).to(dtype) + else: + eps = torch.randn(x.shape, generator=generator, device=device, dtype=dtype) + x = ((1.0 - float(t_next)) * x0.double() + float(t_next) * eps.double()).to(dtype) + + decoded_latents = pipe._unpack_latents(x, height, width, pipe.vae_scale_factor) + latents_mean = ( + torch.tensor(pipe.vae.config.latents_mean) + .view(1, pipe.vae.config.z_dim, 1, 1, 1) + .to(device, dtype) + ) + latents_std = ( + torch.tensor(pipe.vae.config.latents_std) + .view(1, pipe.vae.config.z_dim, 1, 1, 1) + .to(device, dtype) + ) + decoded_latents = decoded_latents * latents_std + latents_mean + decoded = pipe.vae.decode(decoded_latents, return_dict=False)[0][:, :, 0] + return QwenImageEditDMDOutput( + images=pipe.image_processor.postprocess(decoded, output_type=output_type) + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--student-path", required=True) + parser.add_argument("--base-pipeline-path", default="Qwen/Qwen-Image-Edit-2511") + parser.add_argument("--image", nargs="+", required=True, help="Ordered reference image(s).") + parser.add_argument("--prompt", required=True) + parser.add_argument("--negative-prompt") + parser.add_argument("--num-inference-steps", type=int, default=4) + parser.add_argument("--guidance-scale", type=float, default=1.0) + parser.add_argument("--height", type=int) + parser.add_argument("--width", type=int) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--ema-path") + parser.add_argument("--output", default="qwen_image_edit_dmd2.png") + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(args.seed) + pipeline = QwenImageEditDMDInferencePipeline.from_pretrained( + args.student_path, + args.base_pipeline_path, + ema_path=args.ema_path, + ).to(device) + references = [load_image(path).convert("RGB") for path in args.image] + output = pipeline( + references, + args.prompt, + negative_prompt=args.negative_prompt, + num_inference_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + height=args.height, + width=args.width, + generator=generator, + ) + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + output.images[0].save(args.output) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/preprocess/processors/__init__.py b/examples/diffusers/fastgen/preprocess/processors/__init__.py index 2660d27b105..a9d3e67448e 100644 --- a/examples/diffusers/fastgen/preprocess/processors/__init__.py +++ b/examples/diffusers/fastgen/preprocess/processors/__init__.py @@ -23,6 +23,7 @@ get_caption_loader, ) from .qwen_image import QwenImageProcessor +from .qwen_image_edit import QwenImageEditProcessor from .registry import ProcessorRegistry __all__ = [ @@ -33,6 +34,7 @@ "JSONSidecarCaptionLoader", "MetaJSONCaptionLoader", "ProcessorRegistry", + "QwenImageEditProcessor", "QwenImageProcessor", "get_caption_loader", ] diff --git a/examples/diffusers/fastgen/preprocess/processors/qwen_image_edit.py b/examples/diffusers/fastgen/preprocess/processors/qwen_image_edit.py new file mode 100644 index 00000000000..f1af7f8c5d4 --- /dev/null +++ b/examples/diffusers/fastgen/preprocess/processors/qwen_image_edit.py @@ -0,0 +1,288 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Qwen-Image-Edit preprocessing support. + +Qwen-Image-Edit-2511 conditions the denoiser through two independent paths: + +* every reference image is encoded by the Qwen2.5-VL prompt encoder together with the edit + instruction; and +* every reference image is encoded by the Qwen-Image VAE and appended to the noisy target + tokens by the denoiser adapter. + +This processor intentionally keeps those two representations separate in the cache. Target +latents use the same sampled-VAE convention as :class:`QwenImageProcessor`, while reference +latents use the deterministic posterior mode used by ``QwenImageEditPlusPipeline`` at inference. +""" + +from __future__ import annotations + +import logging +import math +from typing import TYPE_CHECKING, Any + +import torch + +from .qwen_image import QwenImageProcessor +from .registry import ProcessorRegistry + +if TYPE_CHECKING: + from PIL import Image + +logger = logging.getLogger(__name__) + +_CONDITION_IMAGE_AREA = 384 * 384 +_VAE_IMAGE_AREA = 1024 * 1024 + + +def _dimensions_for_area(image: Image.Image, area: int) -> tuple[int, int]: + """Return ``(width, height)`` preserving aspect ratio, rounded to multiples of 32.""" + + width, height = image.size + if width <= 0 or height <= 0: + raise ValueError(f"Invalid image dimensions: {image.size}") + ratio = width / height + resized_width = max(32, round(math.sqrt(area * ratio) / 32) * 32) + resized_height = max(32, round(math.sqrt(area / ratio) / 32) * 32) + return resized_width, resized_height + + +def _posterior_mode(encoder_output: Any) -> torch.Tensor: + """Extract the deterministic VAE posterior mode across diffusers return variants.""" + + if hasattr(encoder_output, "latent_dist"): + return encoder_output.latent_dist.mode() + if hasattr(encoder_output, "latents"): + return encoder_output.latents + raise AttributeError("Could not access VAE latents from encoder output") + + +@ProcessorRegistry.register("qwen_image_edit") +class QwenImageEditProcessor(QwenImageProcessor): + """Precompute the full Qwen-Image-Edit-2511 conditioning contract.""" + + @property + def model_type(self) -> str: + return "qwen_image_edit" + + @property + def default_model_name(self) -> str: + return "Qwen/Qwen-Image-Edit-2511" + + def load_models(self, model_name: str, device: str) -> dict[str, Any]: + """Load only the VAE and multimodal prompt encoder needed for caching.""" + + try: + from diffusers import QwenImageEditPlusPipeline + except ImportError as exc: # pragma: no cover - depends on the runtime environment + raise ImportError( + "Qwen-Image-Edit-2511 preprocessing requires a diffusers release that provides " + "QwenImageEditPlusPipeline (introduced in diffusers 0.36.0). Full 2511 " + "denoising also requires `zero_cond_t` support (stable diffusers>=0.37)." + ) from exc + + logger.info("[Qwen-Image-Edit] Loading preprocessing models from %s", model_name) + pipeline = QwenImageEditPlusPipeline.from_pretrained( + model_name, + transformer=None, + torch_dtype=torch.bfloat16, + ) + pipeline.vae.to(device=device, dtype=torch.bfloat16).eval() + pipeline.text_encoder.to(device).eval() + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return {"vae": pipeline.vae, "pipeline": pipeline} + + def _vision_images(self, images: list[Image.Image], pipeline: Any) -> list[Image.Image]: + """Resize references exactly as the Edit Plus prompt-encoding path does.""" + + prepared = [] + for image in images: + width, height = _dimensions_for_area(image, _CONDITION_IMAGE_AREA) + prepared.append(pipeline.image_processor.resize(image, height, width)) + return prepared + + def encode_conditioning_images( + self, + images: list[Image.Image], + models: dict[str, Any], + device: str, + *, + max_pixels: int = _VAE_IMAGE_AREA, + ) -> list[torch.Tensor]: + """Encode one or more references with deterministic VAE posterior modes. + + Each returned tensor is ``[C, H/8, W/8]``. References are kept as a list because the + Edit Plus model permits different aspect ratios for different references. + """ + + if not images: + raise ValueError("Qwen-Image-Edit requires at least one conditioning image") + vae = models["vae"] + pipeline = models["pipeline"] + latents = [] + for image in images: + width, height = _dimensions_for_area(image, max_pixels) + image_tensor = pipeline.image_processor.preprocess(image, height, width).unsqueeze(2) + image_tensor = image_tensor.to(device=device, dtype=torch.bfloat16) + with torch.no_grad(): + latent = _posterior_mode(vae.encode(image_tensor)) + + latents_mean = ( + torch.tensor(vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(latent.device, latent.dtype) + ) + latents_std = ( + torch.tensor(vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(latent.device, latent.dtype) + ) + latent = (latent - latents_mean) / latents_std + latents.append(latent.detach().cpu().to(torch.float16).squeeze(2).squeeze(0)) + return latents + + def encode_multimodal_text( + self, + prompt: str, + images: list[Image.Image], + models: dict[str, Any], + device: str, + ) -> dict[str, torch.Tensor]: + """Encode an instruction and its reference images with Qwen2.5-VL.""" + + if not images: + raise ValueError("Qwen-Image-Edit prompt encoding requires at least one image") + pipeline = models["pipeline"] + vision_images = self._vision_images(images, pipeline) + with torch.no_grad(): + prompt_embeds, prompt_embeds_mask = pipeline.encode_prompt( + prompt=prompt, + image=vision_images, + device=device, + ) + result = {"prompt_embeds": prompt_embeds.detach().cpu().to(torch.bfloat16)} + if prompt_embeds_mask is not None: + result["prompt_embeds_mask"] = prompt_embeds_mask.detach().cpu().to(torch.long) + return result + + def encode_edit_prompts( + self, + prompt: str, + negative_prompt: str, + images: list[Image.Image], + models: dict[str, Any], + device: str, + ) -> dict[str, torch.Tensor]: + """Encode positive and per-sample negative multimodal conditioning.""" + + positive = self.encode_multimodal_text(prompt, images, models, device) + negative = self.encode_multimodal_text(negative_prompt, images, models, device) + result = dict(positive) + result["negative_prompt_embeds"] = negative["prompt_embeds"] + if "prompt_embeds_mask" in negative: + result["negative_prompt_embeds_mask"] = negative["prompt_embeds_mask"] + return result + + def encode_text( + self, + prompt: str, + models: dict[str, Any], + device: str, + ) -> dict[str, torch.Tensor]: + """Reject text-only use, which would silently omit image tokens from the cache.""" + + raise ValueError( + "QwenImageEditProcessor.encode_text cannot encode a text-only prompt. Use " + "encode_multimodal_text/encode_edit_prompts with the conditioning images." + ) + + def verify_latent( + self, + latent: torch.Tensor, + models: dict[str, Any], + device: str, + ) -> bool: + """Decode a target latent with the VAE's actual dtype and validate finiteness.""" + + try: + vae = models["vae"] + vae_dtype = next(vae.parameters()).dtype + value = latent.unsqueeze(0).unsqueeze(2).to(device=device, dtype=vae_dtype) + latents_mean = ( + torch.tensor(vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(device=device, dtype=vae_dtype) + ) + latents_std = ( + torch.tensor(vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(device=device, dtype=vae_dtype) + ) + with torch.no_grad(): + decoded = vae.decode(value * latents_std + latents_mean).sample[:, :, 0] + return ( + decoded.ndim == 4 and decoded.shape[1] == 3 and torch.isfinite(decoded).all().item() + ) + except Exception as exc: + logger.warning("[Qwen-Image-Edit] Latent verification failed: %s", exc) + return False + + def get_cache_data( + self, + latent: torch.Tensor, + text_encodings: dict[str, torch.Tensor], + metadata: dict[str, Any], + ) -> dict[str, Any]: + """Construct an image-edit cache record consumed by ``ImageToImageDataset``.""" + + conditioning_latents = metadata.get("conditioning_latents") + if not isinstance(conditioning_latents, list) or not conditioning_latents: + raise ValueError("metadata['conditioning_latents'] must be a non-empty tensor list") + required_text = ("prompt_embeds", "negative_prompt_embeds") + missing = [key for key in required_text if key not in text_encodings] + if missing: + raise KeyError(f"Missing edit text encodings: {missing}") + + cache = { + "latent": latent, + "conditioning_latents": conditioning_latents, + "prompt_embeds": text_encodings["prompt_embeds"], + "negative_prompt_embeds": text_encodings["negative_prompt_embeds"], + "original_resolution": metadata["original_resolution"], + "bucket_resolution": metadata["bucket_resolution"], + "crop_offset": metadata["crop_offset"], + "prompt": metadata["prompt"], + "negative_prompt": metadata["negative_prompt"], + "image_path": metadata["image_path"], + "conditioning_image_paths": metadata["conditioning_image_paths"], + "conditioning_resolutions": metadata["conditioning_resolutions"], + "target_latent_shape": tuple(latent.shape), + "conditioning_latent_shapes": [tuple(value.shape) for value in conditioning_latents], + "bucket_id": metadata["bucket_id"], + "aspect_ratio": metadata["aspect_ratio"], + "sample_id": metadata["sample_id"], + "model_type": self.model_type, + } + for key in ("prompt_embeds_mask", "negative_prompt_embeds_mask"): + if key in text_encodings: + cache[key] = text_encodings[key] + if metadata.get("source_metadata") is not None: + cache["source_metadata"] = metadata["source_metadata"] + return cache + + +__all__ = ["QwenImageEditProcessor"] diff --git a/examples/diffusers/fastgen/preprocess_qwen_image_edit.py b/examples/diffusers/fastgen/preprocess_qwen_image_edit.py new file mode 100644 index 00000000000..e844e27b0e1 --- /dev/null +++ b/examples/diffusers/fastgen/preprocess_qwen_image_edit.py @@ -0,0 +1,792 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Precompute Qwen-Image-Edit-2511 DMD2 training caches. + +Two local input layouts are supported: + +* SpatialEdit-style WebDataset roots containing ``*.tar`` shards. A sample contains + ``.json`` and indexed images such as ``.0.jpg`` / ``.1.jpg``. Images are + ordered by index; the last image is the target and every preceding image is a reference. +* A JSONL manifest with ``target``, ``conditioning`` (a path or path list), and ``prompt``. + The data-tooling aliases ``generated_image`` / ``reference_image`` (targets) and + ``conditioning_images`` (sources) are also accepted, including ``{archive, member}`` + descriptors. Relative paths are resolved against the manifest directory. ``id``, + ``negative_prompt``, and ``metadata`` are optional. + +The output follows ``BaseMultiresolutionDataset``'s sharded ``metadata.json`` layout. Each +``.pt`` record contains a sampled target ``latent``, a list of deterministic +``conditioning_latents``, and positive/negative multimodal embeddings and masks. +""" + +from __future__ import annotations + +import argparse +import hashlib +import io +import json +import logging +import os +import re +import sys +import tarfile +import traceback +from collections import OrderedDict +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from PIL import Image + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + +_THIS_DIR = Path(__file__).resolve().parent +if str(_THIS_DIR) not in sys.path: + sys.path.insert(0, str(_THIS_DIR)) + +logger = logging.getLogger(__name__) + +_IMAGE_MEMBER_RE = re.compile( + r"^(?P.+)\.(?P\d+)\.(?:jpe?g|png|webp)$", + flags=re.IGNORECASE, +) +_IMAGE_MARKER_RE = re.compile(r"(?:\s*)+", flags=re.IGNORECASE) + + +@dataclass +class EditSample: + """One loaded edit pair. Only one instance is retained while streaming input shards.""" + + sample_id: str + target_image: Image.Image + conditioning_images: list[Image.Image] + prompt: str + negative_prompt: str + target_path: str + conditioning_paths: list[str] + source_metadata: dict[str, Any] | None = None + + +def _load_rgb(path: Path) -> Image.Image: + with Image.open(path) as image: + return image.convert("RGB") + + +def _load_tar_rgb(archive: tarfile.TarFile, member: tarfile.TarInfo) -> Image.Image: + fileobj = archive.extractfile(member) + if fileobj is None: + raise OSError(f"Could not read {member.name!r} from {archive.name!r}") + with Image.open(io.BytesIO(fileobj.read())) as image: + return image.convert("RGB") + + +def _strip_image_markers(text: str) -> str: + """Remove dataset placeholders because EditPlus inserts its own visual tokens.""" + + return _IMAGE_MARKER_RE.sub("", text).strip() + + +def _conversation_text(payload: dict[str, Any], key: str) -> str | None: + conversation = payload.get(key) + if not isinstance(conversation, list): + return None + for message in conversation: + if isinstance(message, dict) and message.get("from") in {"human", "user"}: + value = message.get("value") or message.get("text") + if isinstance(value, str) and value.strip(): + return value + return None + + +def extract_spatialedit_prompt(payload: dict[str, Any], variant: str = "human") -> str: + """Extract a usable instruction across SpatialEdit's three shard families.""" + + metadata = payload.get("metadata") + metadata = metadata if isinstance(metadata, dict) else {} + if variant == "human": + candidates = ( + metadata.get("instruction_human"), + _conversation_text(payload, "conversations_human"), + metadata.get("instruction"), + _conversation_text(payload, "conversations"), + ) + elif variant == "raw": + candidates = ( + metadata.get("instruction"), + _conversation_text(payload, "conversations"), + metadata.get("instruction_human"), + _conversation_text(payload, "conversations_human"), + ) + else: + raise ValueError(f"Unknown prompt variant: {variant!r}") + for candidate in candidates: + if isinstance(candidate, str) and candidate.strip(): + prompt = _strip_image_markers(candidate) + if prompt: + return prompt + raise ValueError("SpatialEdit sample has no non-empty edit instruction") + + +def _spatialedit_sample_id(payload: dict[str, Any], fallback: str) -> str: + metadata = payload.get("metadata") + candidates = ( + payload.get("SAMPLE_ID"), + payload.get("id"), + metadata.get("id") if isinstance(metadata, dict) else None, + fallback, + ) + return str(next(value for value in candidates if value is not None and str(value))) + + +def _spatialedit_metadata( + payload: dict[str, Any], + tar_path: Path, + sample_key: str, +) -> dict[str, Any]: + result: dict[str, Any] = { + "source": "SpatialEdit-500K", + "tar_path": str(tar_path.resolve()), + "webdataset_key": sample_key, + } + for key in ("metadata", "meta", "data_type", "multi_image", "only_text"): + if key in payload: + result[key] = payload[key] + return result + + +def iter_spatialedit_samples( + root: Path, + *, + negative_prompt: str = " ", + prompt_variant: str = "human", + shard_rank: int = 0, + shard_world: int = 1, +) -> Iterator[EditSample]: + """Stream native SpatialEdit WebDataset pairs without extracting shards to disk.""" + + tar_paths = sorted(path for path in root.rglob("*.tar") if path.is_file()) + if not tar_paths: + raise FileNotFoundError(f"No .tar shards found under {root}") + selected = tar_paths[shard_rank::shard_world] + logger.info( + "SpatialEdit input: %d/%d tar shards assigned to rank %d", + len(selected), + len(tar_paths), + shard_rank, + ) + + for tar_path in selected: + try: + with tarfile.open(tar_path, mode="r:*") as archive: + grouped: dict[str, dict[str, Any]] = {} + for member in archive.getmembers(): + if not member.isfile(): + continue + if member.name.lower().endswith(".json"): + key = member.name[: -len(".json")] + grouped.setdefault(key, {})["json"] = member + continue + match = _IMAGE_MEMBER_RE.match(member.name) + if match: + group = grouped.setdefault(match.group("key"), {}) + group.setdefault("images", {})[int(match.group("index"))] = member + + for sample_key in sorted(grouped): + members = grouped[sample_key] + image_members = members.get("images", {}) + if "json" not in members or len(image_members) < 2: + logger.warning( + "Skipping incomplete WebDataset sample %s::%s (json=%s, images=%d)", + tar_path, + sample_key, + "json" in members, + len(image_members), + ) + continue + json_file = archive.extractfile(members["json"]) + if json_file is None: + raise OSError(f"Could not read metadata for {tar_path}::{sample_key}") + payload = json.loads(json_file.read()) + if not isinstance(payload, dict): + raise ValueError(f"Metadata for {tar_path}::{sample_key} is not an object") + + ordered = sorted(image_members.items()) + conditioning = [_load_tar_rgb(archive, member) for _, member in ordered[:-1]] + _, target_member = ordered[-1] + target = _load_tar_rgb(archive, target_member) + display_prefix = f"{tar_path.resolve()}::" + yield EditSample( + sample_id=_spatialedit_sample_id(payload, sample_key), + target_image=target, + conditioning_images=conditioning, + prompt=extract_spatialedit_prompt(payload, prompt_variant), + negative_prompt=negative_prompt, + target_path=f"{display_prefix}{target_member.name}", + conditioning_paths=[ + f"{display_prefix}{member.name}" for _, member in ordered[:-1] + ], + source_metadata=_spatialedit_metadata(payload, tar_path, sample_key), + ) + except Exception: + logger.error("Failed while reading WebDataset shard %s", tar_path) + logger.debug(traceback.format_exc()) + raise + + +def _manifest_value(record: dict[str, Any], names: tuple[str, ...]) -> Any: + for name in names: + if name in record and record[name] is not None: + return record[name] + return None + + +def _resolve_manifest_path(value: Any, base_dir: Path, field: str) -> Path: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Manifest field {field!r} must be a non-empty local path") + path = Path(value).expanduser() + if not path.is_absolute(): + path = base_dir / path + path = path.resolve() + if not path.is_file(): + raise FileNotFoundError(f"Manifest {field} image does not exist: {path}") + return path + + +class _ArchiveImageLoader: + """Small LRU of open tar files for archive/member JSONL descriptors.""" + + def __init__(self, max_open_archives: int = 8) -> None: + self.max_open_archives = max_open_archives + self._archives: OrderedDict[Path, tarfile.TarFile] = OrderedDict() + + def _archive(self, path: Path) -> tarfile.TarFile: + archive = self._archives.pop(path, None) + if archive is None: + archive = tarfile.open(path, mode="r:*") + self._archives[path] = archive + while len(self._archives) > self.max_open_archives: + _, stale = self._archives.popitem(last=False) + stale.close() + return archive + + def load(self, archive_path: Path, member_name: str) -> Image.Image: + archive = self._archive(archive_path) + try: + member = archive.getmember(member_name) + except KeyError as exc: + raise FileNotFoundError( + f"Archive member does not exist: {archive_path}::{member_name}" + ) from exc + return _load_tar_rgb(archive, member) + + def close(self) -> None: + for archive in self._archives.values(): + archive.close() + self._archives.clear() + + +def _load_manifest_image( + value: Any, + base_dir: Path, + field: str, + archive_loader: _ArchiveImageLoader, +) -> tuple[Image.Image, str]: + """Load a local path or ``{archive, member}`` image descriptor.""" + + if isinstance(value, str): + path = _resolve_manifest_path(value, base_dir, field) + return _load_rgb(path), str(path) + if not isinstance(value, dict): + raise ValueError( + f"Manifest field {field!r} must be a path or an {{archive, member}} object" + ) + if "path" in value: + path = _resolve_manifest_path(value["path"], base_dir, field) + return _load_rgb(path), str(path) + + archive_path = _resolve_manifest_path(value.get("archive"), base_dir, f"{field}.archive") + member = value.get("member") + if not isinstance(member, str) or not member: + raise ValueError(f"Manifest field {field!r}.member must be a non-empty string") + return archive_loader.load(archive_path, member), f"{archive_path}::{member}" + + +def iter_jsonl_samples( + manifest: Path, + *, + negative_prompt: str = " ", + shard_rank: int = 0, + shard_world: int = 1, +) -> Iterator[EditSample]: + """Stream generic local edit records from a JSONL manifest.""" + + base_dir = manifest.resolve().parent + archive_loader = _ArchiveImageLoader() + try: + with manifest.open("r", encoding="utf-8") as handle: + for line_index, line in enumerate(handle): + if not line.strip() or line.lstrip().startswith("#"): + continue + record = json.loads(line) + if not isinstance(record, dict): + raise ValueError(f"Manifest line {line_index + 1} is not a JSON object") + shard_value = record.get("archive_index") + if shard_value is None: + shard_value = line_index + try: + assigned_rank = int(shard_value) % shard_world + except (TypeError, ValueError) as exc: + raise ValueError( + f"Manifest line {line_index + 1} archive_index must be an integer" + ) from exc + if assigned_rank != shard_rank: + continue + + target_value = _manifest_value( + record, + ( + "target", + "target_image", + "generated_image", + "output_image", + "output", + "reference_image", + ), + ) + conditioning_value = _manifest_value( + record, + ( + "conditioning", + "conditioning_images", + "reference_images", + "source_images", + "source", + "input", + ), + ) + prompt_value = _manifest_value( + record, ("prompt", "instruction", "edit_instruction") + ) + if isinstance(conditioning_value, (str, dict)): + conditioning_value = [conditioning_value] + if not isinstance(conditioning_value, list) or not conditioning_value: + raise ValueError( + f"Manifest line {line_index + 1} must provide one or more " + "conditioning images" + ) + if not isinstance(prompt_value, str) or not prompt_value.strip(): + raise ValueError(f"Manifest line {line_index + 1} has no edit prompt") + + target_image, target_path = _load_manifest_image( + target_value, + base_dir, + "target", + archive_loader, + ) + loaded_conditioning = [ + _load_manifest_image(value, base_dir, "conditioning", archive_loader) + for value in conditioning_value + ] + per_sample_negative = record.get("negative_prompt", negative_prompt) + if not isinstance(per_sample_negative, str): + raise ValueError( + f"Manifest line {line_index + 1} negative_prompt must be a string" + ) + sample_id = next( + str(value) + for value in ( + record.get("id"), + record.get("sample_id"), + record.get("source_id"), + record.get("key"), + line_index, + ) + if value is not None and str(value) + ) + yield EditSample( + sample_id=sample_id, + target_image=target_image, + conditioning_images=[image for image, _ in loaded_conditioning], + prompt=_strip_image_markers(prompt_value), + negative_prompt=per_sample_negative, + target_path=target_path, + conditioning_paths=[path for _, path in loaded_conditioning], + source_metadata=record.get("metadata"), + ) + finally: + archive_loader.close() + + +def _write_json_atomic(path: Path, payload: Any) -> None: + temporary = path.with_name(f".{path.name}.tmp-{os.getpid()}") + with temporary.open("w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, ensure_ascii=False) + os.replace(temporary, path) + + +class MetadataShardWriter: + """Incrementally write metadata shards so 500K samples need not stay in memory.""" + + def __init__( + self, + output_dir: Path, + shard_size: int, + shard_rank: int, + shard_world: int, + ) -> None: + self.output_dir = output_dir + self.shard_size = shard_size + self.shard_rank = shard_rank + self.shard_world = shard_world + self.buffer: list[dict[str, Any]] = [] + self.shards: list[str] = [] + self.total_items = 0 + + def add(self, item: dict[str, Any]) -> None: + self.buffer.append(item) + self.total_items += 1 + if len(self.buffer) >= self.shard_size: + self.flush() + + def flush(self) -> None: + if not self.buffer: + return + rank_prefix = f"r{self.shard_rank:02d}_" if self.shard_world > 1 else "" + filename = f"metadata_shard_{rank_prefix}s{len(self.shards):04d}.json" + _write_json_atomic(self.output_dir / filename, self.buffer) + self.shards.append(filename) + self.buffer = [] + + def finish(self, **config: Any) -> Path: + self.flush() + if not self.shards: + raise RuntimeError("No valid samples were preprocessed; metadata was not written") + index_name = ( + f"metadata_r{self.shard_rank:02d}.json" if self.shard_world > 1 else "metadata.json" + ) + payload = { + "processor": "qwen_image_edit", + "model_type": "qwen_image_edit", + "total_items": self.total_items, + "num_shards": len(self.shards), + "shard_size": self.shard_size, + "shards": self.shards, + **config, + } + if self.shard_world > 1: + payload.update(shard_rank=self.shard_rank, shard_world=self.shard_world) + index_path = self.output_dir / index_name + _write_json_atomic(index_path, payload) + return index_path + + +def _cache_identity( + sample: EditSample, + model_name: str, + resolution: tuple[int, int], + conditioning_max_pixels: int, +) -> str: + fields = ( + model_name, + sample.sample_id, + sample.target_path, + *sample.conditioning_paths, + sample.prompt, + sample.negative_prompt, + f"{resolution[0]}x{resolution[1]}", + str(conditioning_max_pixels), + ) + return hashlib.sha256("\0".join(fields).encode("utf-8")).hexdigest() + + +def _stable_sample_seed(sample_id: str, seed: int) -> int: + digest = hashlib.sha256(f"{seed}\0{sample_id}".encode()).digest() + return int.from_bytes(digest[:8], byteorder="big") % (2**31) + + +def preprocess_samples( + samples: Iterable[EditSample], + *, + output_dir: Path, + model_name: str, + device: str, + max_pixels: int, + conditioning_max_pixels: int, + metadata_shard_size: int, + shard_rank: int, + shard_world: int, + verify: bool, + overwrite: bool, + fail_fast: bool, + limit: int | None, + seed: int, + log_every: int, +) -> Path: + """Encode a stream of edit samples and return the generated metadata index path.""" + + import torch + from nemo_automodel.components.datasets.diffusion.multi_tier_bucketing import ( + MultiTierBucketCalculator, + ) + from preprocess.processors import QwenImageEditProcessor + + output_dir.mkdir(parents=True, exist_ok=True) + processor = QwenImageEditProcessor() + models = processor.load_models(model_name, device) + calculator = MultiTierBucketCalculator(quantization=64, max_pixels=max_pixels) + writer = MetadataShardWriter( + output_dir, + metadata_shard_size, + shard_rank, + shard_world, + ) + + attempted = 0 + failures = 0 + for sample in samples: + if limit is not None and attempted >= limit: + break + attempted += 1 + try: + original_width, original_height = sample.target_image.size + bucket = calculator.get_bucket_for_image(original_width, original_height) + target_width, target_height = bucket["resolution"] + resolution = (target_width, target_height) + cache_hash = _cache_identity( + sample, + model_name, + resolution, + conditioning_max_pixels, + ) + cache_subdir = output_dir / f"{target_width}x{target_height}" + cache_subdir.mkdir(parents=True, exist_ok=True) + cache_file = cache_subdir / f"{cache_hash}.pt" + + if overwrite or not cache_file.is_file(): + resized_target, crop_offset = calculator.resize_and_crop( + sample.target_image, + target_width, + target_height, + crop_mode="center", + ) + sample_seed = _stable_sample_seed(sample.sample_id, seed) + torch.manual_seed(sample_seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(sample_seed) + target_tensor = processor.preprocess_image(resized_target) + latent = processor.encode_image(target_tensor, models, device) + if verify and not processor.verify_latent(latent, models, device): + raise ValueError("target latent verification failed") + + conditioning_latents = processor.encode_conditioning_images( + sample.conditioning_images, + models, + device, + max_pixels=conditioning_max_pixels, + ) + text_encodings = processor.encode_edit_prompts( + sample.prompt, + sample.negative_prompt, + sample.conditioning_images, + models, + device, + ) + cache_metadata = { + "conditioning_latents": conditioning_latents, + "original_resolution": (original_width, original_height), + "bucket_resolution": resolution, + "crop_offset": crop_offset, + "prompt": sample.prompt, + "negative_prompt": sample.negative_prompt, + "image_path": sample.target_path, + "conditioning_image_paths": sample.conditioning_paths, + "conditioning_resolutions": [ + tuple(image.size) for image in sample.conditioning_images + ], + "bucket_id": bucket["id"], + "aspect_ratio": bucket["aspect_ratio"], + "sample_id": sample.sample_id, + "source_metadata": sample.source_metadata, + } + cache = processor.get_cache_data(latent, text_encodings, cache_metadata) + temporary = cache_file.with_name(f".{cache_file.name}.tmp-{os.getpid()}") + torch.save(cache, temporary) + os.replace(temporary, cache_file) + else: + crop_offset = (0, 0) + + writer.add( + { + "cache_file": str(cache_file.resolve()), + "image_path": sample.target_path, + "conditioning_image_paths": sample.conditioning_paths, + "conditioning_resolutions": [ + list(image.size) for image in sample.conditioning_images + ], + "sample_id": sample.sample_id, + "bucket_resolution": [target_width, target_height], + "original_resolution": [original_width, original_height], + "prompt": sample.prompt, + "bucket_id": bucket["id"], + "aspect_ratio": bucket["aspect_ratio"], + "pixels": target_width * target_height, + "model_type": processor.model_type, + } + ) + if log_every > 0 and attempted % log_every == 0: + logger.info( + "Processed %d samples (%d failures, %d cached records)", + attempted, + failures, + writer.total_items, + ) + except Exception as exc: + failures += 1 + logger.error("Failed sample %s: %s", sample.sample_id, exc) + logger.debug(traceback.format_exc()) + if fail_fast: + raise + + index_path = writer.finish( + model_name=model_name, + max_pixels=max_pixels, + conditioning_max_pixels=conditioning_max_pixels, + attempted_items=attempted, + failed_items=failures, + negative_prompt_is_per_sample=True, + ) + logger.info( + "Finished preprocessing: %d records, %d failures; metadata=%s", + writer.total_items, + failures, + index_path, + ) + return index_path + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be positive") + return parsed + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument( + "--webdataset-root", + "--input-dir", + dest="webdataset_root", + type=Path, + help="SpatialEdit-style root recursively containing WebDataset .tar shards", + ) + source.add_argument( + "--manifest", + type=Path, + help="Local JSONL with target, conditioning, and prompt fields", + ) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--model-name", default="Qwen/Qwen-Image-Edit-2511") + parser.add_argument( + "--device", + default=None, + help="Torch device. Defaults to cuda or cuda: when --gpu-id is provided.", + ) + parser.add_argument("--gpu-id", type=int, help="GPU index used when --device is omitted") + parser.add_argument("--max-pixels", type=_positive_int, default=1024 * 1024) + parser.add_argument( + "--conditioning-max-pixels", + type=_positive_int, + default=1024 * 1024, + help="Per-reference VAE pixel budget (the official EditPlus default is 1024^2)", + ) + parser.add_argument("--negative-prompt", default=" ") + parser.add_argument( + "--prompt-variant", + choices=("human", "raw"), + default="human", + help="SpatialEdit instruction variant; ignored for JSONL manifests", + ) + parser.add_argument("--metadata-shard-size", type=_positive_int, default=10_000) + parser.add_argument("--shard-rank", "--shard-idx", dest="shard_rank", type=int, default=0) + parser.add_argument( + "--shard-world", + "--shard-count", + dest="shard_world", + type=_positive_int, + default=1, + ) + parser.add_argument("--limit", "--max-samples", dest="limit", type=_positive_int) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--log-every", type=int, default=100) + parser.add_argument("--verify", action="store_true") + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--fail-fast", action="store_true") + return parser + + +def main(argv: list[str] | None = None) -> Path: + args = build_parser().parse_args(argv) + if not 0 <= args.shard_rank < args.shard_world: + raise ValueError( + f"shard_rank must satisfy 0 <= rank < world; got {args.shard_rank}/{args.shard_world}" + ) + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s | %(levelname)s | %(name)s | %(message)s", + ) + device = args.device + if device is None: + device = f"cuda:{args.gpu_id}" if args.gpu_id is not None else "cuda" + + if args.webdataset_root is not None: + samples = iter_spatialedit_samples( + args.webdataset_root, + negative_prompt=args.negative_prompt, + prompt_variant=args.prompt_variant, + shard_rank=args.shard_rank, + shard_world=args.shard_world, + ) + else: + samples = iter_jsonl_samples( + args.manifest, + negative_prompt=args.negative_prompt, + shard_rank=args.shard_rank, + shard_world=args.shard_world, + ) + + return preprocess_samples( + samples, + output_dir=args.output_dir, + model_name=args.model_name, + device=device, + max_pixels=args.max_pixels, + conditioning_max_pixels=args.conditioning_max_pixels, + metadata_shard_size=args.metadata_shard_size, + shard_rank=args.shard_rank, + shard_world=args.shard_world, + verify=args.verify, + overwrite=args.overwrite, + fail_fast=args.fail_fast, + limit=args.limit, + seed=args.seed, + log_every=args.log_every, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/diffusers/fastgen/requirements.txt b/examples/diffusers/fastgen/requirements.txt index 5e5e79f0d12..6c543061e42 100644 --- a/examples/diffusers/fastgen/requirements.txt +++ b/examples/diffusers/fastgen/requirements.txt @@ -9,5 +9,9 @@ # fastgen_data/__init__.py adds a runtime guard with an actionable message if the helpers move. nemo_automodel[diffusion]>=0.4.0,<1.0 +# Qwen-Image-Edit-2511 needs EditPlus plus the transformer's ``zero_cond_t`` handling, +# which landed after the stable 0.36 release (the checkpoint reports 0.36.0.dev0). +diffusers>=0.37.0 + # Optional but recommended for the smoke logs. wandb diff --git a/modelopt/torch/fastgen/plugins/__init__.py b/modelopt/torch/fastgen/plugins/__init__.py index 8810470b26f..c33f29ec9d1 100644 --- a/modelopt/torch/fastgen/plugins/__init__.py +++ b/modelopt/torch/fastgen/plugins/__init__.py @@ -25,3 +25,4 @@ with import_plugin("qwen_image"): from .qwen_image import * + from .qwen_image_edit import * diff --git a/modelopt/torch/fastgen/plugins/qwen_image.py b/modelopt/torch/fastgen/plugins/qwen_image.py index 08a32b09301..0a4d3f8fc64 100644 --- a/modelopt/torch/fastgen/plugins/qwen_image.py +++ b/modelopt/torch/fastgen/plugins/qwen_image.py @@ -48,9 +48,11 @@ from __future__ import annotations import contextlib +from importlib.metadata import PackageNotFoundError, version from typing import TYPE_CHECKING, Any import torch +from packaging.version import Version from torch import nn from ..methods.dmd import DMDPipeline @@ -58,6 +60,16 @@ if TYPE_CHECKING: from ..config import DMDConfig + +try: + # Diffusers 0.35/0.36 requires explicit Python sequence lengths in Qwen's + # positional-embedding path. Starting in 0.37 the mask is authoritative and + # passing txt_seq_lens is deprecated. Keep the shared T2I plugin compatible + # with ModelOpt's broader diffusers extra while Edit-2511 pins the newer API. + _DIFFUSERS_NEEDS_TXT_SEQ_LENS = Version(version("diffusers")) < Version("0.37.0") +except PackageNotFoundError: # pragma: no cover - optional plugin import guard + _DIFFUSERS_NEEDS_TXT_SEQ_LENS = False + __all__ = [ "QwenImageDMDPipeline", "attach_feature_capture", @@ -65,6 +77,7 @@ "pack_latents", "remove_feature_capture", "unpack_latents", + "update_feature_capture_shape", ] @@ -206,6 +219,7 @@ def _call_model( packed = pack_latents(hidden_states) img_shapes = build_img_shapes(b, h, w) + update_feature_capture_shape(model, h, w) call_kwargs: dict[str, Any] = dict(model_kwargs) call_kwargs.pop("hidden_states", None) @@ -214,8 +228,10 @@ def _call_model( call_kwargs.pop("guidance", None) call_kwargs.pop("return_dict", None) txt_seq_lens = call_kwargs.pop("txt_seq_lens", None) - if txt_seq_lens is None and encoder_hidden_states_mask is not None: - txt_seq_lens = encoder_hidden_states_mask.sum(dim=1).int().tolist() + if _DIFFUSERS_NEEDS_TXT_SEQ_LENS: + if txt_seq_lens is None and encoder_hidden_states_mask is not None: + txt_seq_lens = encoder_hidden_states_mask.sum(dim=1).int().tolist() + call_kwargs["txt_seq_lens"] = txt_seq_lens guidance = None if self._guidance_value is not None: @@ -232,7 +248,6 @@ def _call_model( encoder_hidden_states=encoder_hidden_states, encoder_hidden_states_mask=encoder_hidden_states_mask, img_shapes=img_shapes, - txt_seq_lens=txt_seq_lens, guidance=guidance, return_dict=False, **call_kwargs, @@ -266,6 +281,16 @@ def _call_model( _SHAPE_ATTR = "_fastgen_capture_shape" +def update_feature_capture_shape(model: nn.Module, h_lat: int, w_lat: int) -> None: + """Refresh a hooked teacher's target shape for the current multiresolution batch.""" + if h_lat % 2 or w_lat % 2: + raise ValueError( + f"feature capture requires even latent dims, got h_lat={h_lat}, w_lat={w_lat}." + ) + if hasattr(model, _HANDLES_ATTR): + setattr(model, _SHAPE_ATTR, (h_lat // 2, w_lat // 2)) + + def attach_feature_capture( teacher: nn.Module, feature_indices: list[int], @@ -273,6 +298,7 @@ def attach_feature_capture( w_lat: int, *, blocks_attr: str = "transformer_blocks", + target_prefix_only: bool = False, ) -> None: """Install forward hooks on ``teacher.transformer_blocks[i]`` for each ``i`` in ``feature_indices``. @@ -297,6 +323,11 @@ def attach_feature_capture( blocks_attr: Attribute under which the teacher exposes its block stack. Default ``"transformer_blocks"`` matches diffusers' ``QwenImageTransformer2DModel``. + target_prefix_only: When ``True``, allow the captured image-token sequence to + contain extra tokens after the target image and retain only the leading + ``(h_lat // 2) * (w_lat // 2)`` target tokens. Qwen-Image-Edit concatenates + packed reference-image tokens after the noisy target tokens. The default is + ``False`` so the text-to-image path keeps its strict sequence-length check. Raises: AttributeError: ``teacher`` does not expose ``blocks_attr``. @@ -336,8 +367,6 @@ def attach_feature_capture( setattr(teacher, _SHAPE_ATTR, (h_lat // 2, w_lat // 2)) handles: list[Any] = [] - h_half = h_lat // 2 - w_half = w_lat // 2 for idx in sorted_indices: block = blocks[idx] @@ -354,13 +383,20 @@ def _hook(_module: nn.Module, _inputs: Any, output: Any) -> None: ) # hidden: [B, num_image_patches, C] -> [B, C, H_half, W_half]. b, s, c = hidden.shape + h_half, w_half = getattr(teacher, _SHAPE_ATTR) expected_s = h_half * w_half - if s != expected_s: + if s < expected_s or (not target_prefix_only and s != expected_s): + expected_description = ( + f"at least {expected_s}" if target_prefix_only else str(expected_s) + ) raise RuntimeError( f"QwenImage feature-capture got hidden_states seq_len={s} but expected " - f"{expected_s} = (h_lat // 2) * (w_lat // 2). Did the input resolution " - f"drift from the attach_feature_capture-time setting?" + f"{expected_description} target tokens, where {expected_s} = " + "(h_lat // 2) * (w_lat // 2). Did the input resolution drift from " + "the attach_feature_capture-time setting?" ) + if target_prefix_only: + hidden = hidden[:, :expected_s] feat = hidden.permute(0, 2, 1).reshape(b, c, h_half, w_half) captured.append(feat) diff --git a/modelopt/torch/fastgen/plugins/qwen_image_edit.py b/modelopt/torch/fastgen/plugins/qwen_image_edit.py new file mode 100644 index 00000000000..4d56801c738 --- /dev/null +++ b/modelopt/torch/fastgen/plugins/qwen_image_edit.py @@ -0,0 +1,249 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Qwen-Image-Edit plumbing for DMD2. + +``QwenImageEditPlusPipeline`` conditions the transformer in two complementary ways: + +* the Qwen2.5-VL prompt embedding contains the edit instruction and visual context; and +* one or more clean VAE reference-image latents are packed and appended after the noisy + target-image tokens. + +Only the target image is diffused. DMD2 therefore keeps its external latent contract as +``[B, C, H, W]`` and forwards reference latents through the model kwargs under +``conditioning_latents``. This plugin packs ``[target, reference_1, ...]`` for every model +call, constructs the matching ``img_shapes``, and discards the reference-token suffix from +the model prediction before returning to the shared DMD math. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch +from torch import nn + +from .qwen_image import ( + QwenImageDMDPipeline, + pack_latents, + unpack_latents, + update_feature_capture_shape, +) +from .qwen_image import attach_feature_capture as _attach_qwen_feature_capture +from .qwen_image import remove_feature_capture as _remove_qwen_feature_capture + +if TYPE_CHECKING: + from ..config import DMDConfig + +__all__ = [ + "QwenImageEditDMDPipeline", +] + + +class QwenImageEditDMDPipeline(QwenImageDMDPipeline): + """DMD2 pipeline for Qwen-Image-Edit's target-plus-reference token layout. + + ``conditioning_latents`` must be supplied to each ``compute_*_loss`` call as a + non-empty list or tuple. Each entry is one clean reference image with shape + ``[B, C, H_ref, W_ref]``. Reference images may have different spatial shapes from + each other and from the target, but their batch/channel/device/dtype must match the + target latent. The reference order must match the order used when constructing the + multimodal prompt embedding. + """ + + def __init__( + self, + student: nn.Module, + teacher: nn.Module, + fake_score: nn.Module, + config: DMDConfig, + *, + discriminator: nn.Module | None = None, + guidance: float | None = None, + ) -> None: + """Initialize the shared Qwen pipeline and retain its timestep/guidance checks.""" + super().__init__( + student=student, + teacher=teacher, + fake_score=fake_score, + config=config, + discriminator=discriminator, + guidance=guidance, + ) + + @staticmethod + def _validate_conditioning_latents( + conditioning_latents: Any, + target: torch.Tensor, + ) -> list[torch.Tensor]: + """Validate and normalize the reference-latent sequence for one model call.""" + if not isinstance(conditioning_latents, (list, tuple)) or not conditioning_latents: + raise ValueError( + "QwenImageEditDMDPipeline requires non-empty `conditioning_latents` as a " + "list or tuple of [B, C, H, W] tensors." + ) + + b, c, _h, _w = target.shape + normalized: list[torch.Tensor] = [] + for index, latent in enumerate(conditioning_latents): + if not torch.is_tensor(latent): + raise TypeError( + f"conditioning_latents[{index}] must be a Tensor, got {type(latent).__name__}." + ) + if latent.ndim != 4: + raise ValueError( + f"conditioning_latents[{index}] must have shape [B, C, H, W], got " + f"{latent.ndim}D tensor with shape {tuple(latent.shape)}." + ) + if latent.shape[0] != b or latent.shape[1] != c: + raise ValueError( + f"conditioning_latents[{index}] batch/channels {tuple(latent.shape[:2])} " + f"must match target {(b, c)}." + ) + if latent.shape[2] % 2 or latent.shape[3] % 2: + raise ValueError( + f"conditioning_latents[{index}] requires even spatial dims for Qwen " + f"packing, got H={latent.shape[2]}, W={latent.shape[3]}." + ) + if latent.device != target.device: + raise ValueError( + f"conditioning_latents[{index}] is on {latent.device}, but target is on " + f"{target.device}." + ) + if latent.dtype != target.dtype: + raise ValueError( + f"conditioning_latents[{index}] has dtype {latent.dtype}, but target has " + f"dtype {target.dtype}." + ) + normalized.append(latent) + return normalized + + def _call_model( + self, + model: nn.Module, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + **model_kwargs: Any, + ) -> torch.Tensor: + """Pack target + references, call the transformer, and return target prediction only.""" + if hidden_states.ndim != 4: + raise ValueError( + "QwenImageEditDMDPipeline._call_model expects 4D hidden_states " + f"[B, C, H, W] (got {hidden_states.ndim}D)." + ) + b, _c, h, w = hidden_states.shape + + call_kwargs: dict[str, Any] = dict(model_kwargs) + conditioning_latents = self._validate_conditioning_latents( + call_kwargs.pop("conditioning_latents", None), hidden_states + ) + + target_packed = pack_latents(hidden_states) + update_feature_capture_shape(model, h, w) + conditioning_packed = [pack_latents(latent) for latent in conditioning_latents] + packed = torch.cat([target_packed, *conditioning_packed], dim=1) + target_num_patches = target_packed.shape[1] + + per_sample_shapes = [(1, h // 2, w // 2)] + [ + (1, latent.shape[2] // 2, latent.shape[3] // 2) for latent in conditioning_latents + ] + img_shapes = [list(per_sample_shapes) for _ in range(b)] + + # These values are owned by this wrapper. Drop caller copies so duplicate kwargs + # cannot leak through to the diffusers transformer. + call_kwargs.pop("hidden_states", None) + encoder_hidden_states_mask = call_kwargs.pop("encoder_hidden_states_mask", None) + call_kwargs.pop("img_shapes", None) + call_kwargs.pop("guidance", None) + call_kwargs.pop("return_dict", None) + # Stable Diffusers derives text lengths from encoder_hidden_states_mask. + call_kwargs.pop("txt_seq_lens", None) + + guidance = None + if self._guidance_value is not None: + guidance = torch.full( + (b,), + float(self._guidance_value), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + out = model( + hidden_states=packed, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_mask=encoder_hidden_states_mask, + img_shapes=img_shapes, + guidance=guidance, + return_dict=False, + **call_kwargs, + ) + + if isinstance(out, tuple): + raw_packed = out[0] + elif isinstance(out, torch.Tensor): + raw_packed = out + elif hasattr(out, "sample"): + raw_packed = out.sample + else: + raise TypeError( + "QwenImageEditDMDPipeline._call_model could not extract a tensor from " + f"output of type {type(out).__name__!r}." + ) + + if raw_packed.ndim != 3: + raise ValueError( + "QwenImageEditDMDPipeline expected packed model output [B, tokens, C*4], " + f"got shape {tuple(raw_packed.shape)}." + ) + if raw_packed.shape[0] != b: + raise ValueError( + f"Packed model output batch {raw_packed.shape[0]} does not match target batch {b}." + ) + if raw_packed.shape[1] < target_num_patches: + raise ValueError( + f"Packed model output has {raw_packed.shape[1]} tokens but the target prefix " + f"requires {target_num_patches}." + ) + + # QwenImageEditPlusPipeline treats only the leading target tokens as the denoising + # prediction. The appended reference-token outputs are conditioning-only. + target_prediction = raw_packed[:, :target_num_patches] + return unpack_latents(target_prediction, h, w) + + +def attach_feature_capture( + teacher: nn.Module, + feature_indices: list[int], + h_lat: int, + w_lat: int, + *, + blocks_attr: str = "transformer_blocks", +) -> None: + """Capture only the target-token prefix from Qwen-Image-Edit teacher blocks.""" + _attach_qwen_feature_capture( + teacher, + feature_indices, + h_lat, + w_lat, + blocks_attr=blocks_attr, + target_prefix_only=True, + ) + + +def remove_feature_capture(teacher: nn.Module) -> None: + """Remove feature hooks installed through :func:`attach_feature_capture`.""" + _remove_qwen_feature_capture(teacher) diff --git a/modelopt/torch/quantization/plugins/diffusion/diffusers.py b/modelopt/torch/quantization/plugins/diffusion/diffusers.py index f2f6a702479..fdb5e3443c7 100644 --- a/modelopt/torch/quantization/plugins/diffusion/diffusers.py +++ b/modelopt/torch/quantization/plugins/diffusion/diffusers.py @@ -142,9 +142,16 @@ def _quantized_sdpa(self, *args, **kwargs): k_quantized_scale = self.k_bmm_quantizer._get_amax(key) v_quantized_scale = self.v_bmm_quantizer._get_amax(value) - # We don't need to calibrate the output of softmax - return self.bmm2_output_quantizer( - fp8_sdpa( + # We don't need to calibrate the output of softmax. + # ``FP8SDPA`` is an export-only autograd Function: it exists solely to attach the ONNX + # ``symbolic`` (export_fp8_mha), and its forward is just + # ``original_scaled_dot_product_attention``. It implements no ``backward``, so routing + # through it at runtime makes quantized attention non-differentiable and breaks training + # (QAT) -- ``loss.backward()`` raises "must implement either the backward or vjp method". + # Use it only during ONNX export; at runtime call SDPA directly (identical forward math, + # with q/k/v already fake-quantized above) so autograd works. + if torch.onnx.is_in_onnx_export(): + attn_output = fp8_sdpa( query, key, value, @@ -157,7 +164,19 @@ def _quantized_sdpa(self, *args, **kwargs): else "Half", self._disable_fp8_mha if hasattr(self, "_disable_fp8_mha") else True, ) - ) + else: + # Pass attn_mask/dropout_p/is_causal/scale as keywords (``scale`` is keyword-only in + # recent torch), mirroring FP8SDPA.forward's own call to SDPA. + attn_output = original_scaled_dot_product_attention( + query, + key, + value, + attn_mask=param_dict["attn_mask"], + dropout_p=param_dict["dropout_p"], + is_causal=param_dict["is_causal"], + scale=param_dict["scale"], + ) + return self.bmm2_output_quantizer(attn_output) class _QuantAttention(_QuantFunctionalMixin): diff --git a/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py b/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py new file mode 100644 index 00000000000..2747012ec59 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_quant_state_roundtrip.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Regression test for the DMD2 QAT (restore-only) quantizer-state restore. + +The fastgen QAT path NEVER calibrates during training -- it RESTORES a ModelOpt quantizer +state (recipe + frozen amax) saved by the calibration example and re-applies it on every +(re)start. This test pins the guarantee the recipe depends on, on a tiny CPU model (no +GPU, milliseconds): ``dmd2_recipe.restore_quantizer_state`` onto a *fresh* model with +DIFFERENT weights reproduces the amax bit-identically and leaves that model's weights +untouched -- i.e. amax stays exactly as calibrated and the warm-started student weights +are preserved. + +The on-disk state is built here with ModelOpt's own idiom (the same one the calibration +example's ``--quantized-torch-ckpt-save-path`` uses), so the test also pins format +compatibility with that file. + +Dependency-guarded with ``importorskip`` so it skips where torch / modelopt are absent. +""" + +from __future__ import annotations + +import pathlib +import sys + +import pytest + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) + +torch = pytest.importorskip("torch") +mtq = pytest.importorskip("modelopt.torch.quantization") +mto = pytest.importorskip("modelopt.torch.opt") +dmd2_recipe = pytest.importorskip("dmd2_recipe") + +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.utils.core_utils import get_quantizer_state_dict + + +def _tiny_model(seed: int) -> torch.nn.Module: + torch.manual_seed(seed) + return torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.ReLU(), torch.nn.Linear(8, 4)) + + +def _amax_by_name(model: torch.nn.Module) -> dict[str, torch.Tensor]: + return { + name: module.amax.detach().clone() + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and module.amax is not None + } + + +def test_restore_quantizer_state_is_bit_identical_and_weight_free(tmp_path): + # Calibrate a tiny model (this is the ONLY place quantize/calibration happens -- the + # calibration example; the trainer never does this) and save its quantizer state in the + # weight-free format the calibration example writes (mto.modelopt_state + amax). + model = _tiny_model(seed=0) + calib = torch.randn(16, 8) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(calib)) + src_amax = _amax_by_name(model) + assert src_amax, "expected at least one calibrated TensorQuantizer amax" + + state = mto.modelopt_state(model) + state["modelopt_state_weights"] = get_quantizer_state_dict(model) + path = tmp_path / "transformer.pt" + torch.save(state, str(path)) + + # Restore onto a FRESH model with DIFFERENT weights; amax must come back + # bit-identically and the fresh model's weights must be untouched. + fresh = _tiny_model(seed=999) + before = {n: p.detach().clone() for n, p in fresh.named_parameters()} + dmd2_recipe.restore_quantizer_state(fresh, str(path)) + + restored_amax = _amax_by_name(fresh) + assert set(restored_amax) == set(src_amax) + for name, amax in src_amax.items(): + assert torch.equal(restored_amax[name], amax), f"amax mismatch at {name}" + + for n, p in fresh.named_parameters(): + if n in before: + assert torch.equal(p.detach(), before[n]), f"restore changed weight {n}" diff --git a/tests/examples/diffusers/fastgen/test_qwen_image_edit_data.py b/tests/examples/diffusers/fastgen/test_qwen_image_edit_data.py new file mode 100644 index 00000000000..e4f322cc0c7 --- /dev/null +++ b/tests/examples/diffusers/fastgen/test_qwen_image_edit_data.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Focused CPU tests for the Qwen-Image-Edit preprocessing/data contracts.""" + +from __future__ import annotations + +import io +import json +import sys +import tarfile +from pathlib import Path + +import pytest +from PIL import Image + +_REPO_ROOT = Path(__file__).resolve().parents[4] +_FASTGEN_DIR = _REPO_ROOT / "examples" / "diffusers" / "fastgen" +if str(_FASTGEN_DIR) not in sys.path: + sys.path.insert(0, str(_FASTGEN_DIR)) + +import preprocess_qwen_image_edit as edit_preprocess + + +def _jpeg_bytes(color: tuple[int, int, int]) -> bytes: + buffer = io.BytesIO() + Image.new("RGB", (32, 24), color=color).save(buffer, format="JPEG") + return buffer.getvalue() + + +def _add_tar_bytes(archive: tarfile.TarFile, name: str, payload: bytes) -> None: + member = tarfile.TarInfo(name) + member.size = len(payload) + archive.addfile(member, io.BytesIO(payload)) + + +def test_spatialedit_tar_parser_maps_source_target_and_instruction(tmp_path): + shard = tmp_path / "object_rotation" / "worker0-000000.tar" + shard.parent.mkdir() + metadata = { + "conversations": [ + {"from": "human", "value": "\nRotate the object to the right."}, + {"from": "gpt", "value": "\n"}, + ], + "meta": {"sample_id": "sample-7"}, + } + with tarfile.open(shard, "w") as archive: + _add_tar_bytes(archive, "key.0.jpg", _jpeg_bytes((255, 0, 0))) + _add_tar_bytes(archive, "key.1.jpg", _jpeg_bytes((0, 0, 255))) + _add_tar_bytes(archive, "key.json", json.dumps(metadata).encode()) + + sample = next(edit_preprocess.iter_spatialedit_samples(tmp_path)) + + # The WebDataset key is pair-unique; legacy ``meta.sample_id`` is only asset-unique. + assert sample.sample_id == "key" + assert sample.prompt == "Rotate the object to the right." + assert len(sample.conditioning_images) == 1 + assert sample.conditioning_paths[0].endswith("::key.0.jpg") + assert sample.target_path.endswith("::key.1.jpg") + # JPEG is lossy, so compare dominant channels rather than exact values. + assert sample.conditioning_images[0].getpixel((0, 0))[0] > 200 + assert sample.target_image.getpixel((0, 0))[2] > 200 + + +def test_jsonl_parser_supports_archive_descriptors_and_normalized_field_names(tmp_path): + archive_path = tmp_path / "images.tar" + with tarfile.open(archive_path, "w") as archive: + _add_tar_bytes(archive, "ref1.jpg", _jpeg_bytes((255, 0, 0))) + _add_tar_bytes(archive, "ref2.jpg", _jpeg_bytes((0, 255, 0))) + _add_tar_bytes(archive, "target.jpg", _jpeg_bytes((0, 0, 255))) + manifest = tmp_path / "data.jsonl" + manifest.write_text( + json.dumps( + { + "id": "multi-ref", + "conditioning_images": [ + {"archive": "images.tar", "member": "ref1.jpg"}, + {"archive": "images.tar", "member": "ref2.jpg"}, + ], + "reference_image": {"archive": "images.tar", "member": "target.jpg"}, + "prompt": "\nCombine both references.", + } + ) + + "\n" + ) + + sample = next(edit_preprocess.iter_jsonl_samples(manifest)) + + assert sample.sample_id == "multi-ref" + assert sample.prompt == "Combine both references." + assert sample.negative_prompt == " " + assert len(sample.conditioning_images) == 2 + assert sample.target_path.endswith("images.tar::target.jpg") + + +def test_launcher_cli_aliases_map_to_canonical_arguments(): + args = edit_preprocess.build_parser().parse_args( + [ + "--input-dir", + "raw", + "--output-dir", + "cache", + "--gpu-id", + "3", + "--shard-idx", + "1", + "--shard-count", + "4", + "--max-samples", + "25", + ] + ) + + assert args.webdataset_root == Path("raw") + assert args.gpu_id == 3 + assert (args.shard_rank, args.shard_world, args.limit) == (1, 4, 25) + + +def test_edit_collate_pads_multimodal_positive_and_negative_sequences(): + torch = pytest.importorskip("torch") + pytest.importorskip("nemo_automodel") + from fastgen_data import collate_fn_image_to_image + + def sample(pos_length: int, neg_length: int, sample_id: str): + return { + "latent": torch.randn(16, 8, 8), + "conditioning_latents": [torch.randn(16, 8, 8)], + "prompt_embeds": torch.randn(pos_length, 32), + "prompt_embeds_mask": torch.ones(pos_length, dtype=torch.long), + "negative_prompt_embeds": torch.randn(neg_length, 32), + "negative_prompt_embeds_mask": torch.ones(neg_length, dtype=torch.long), + "crop_resolution": torch.tensor([64, 64]), + "original_resolution": torch.tensor([64, 64]), + "crop_offset": torch.tensor([0, 0]), + "prompt": "edit", + "negative_prompt": " ", + "image_path": f"{sample_id}.jpg", + "conditioning_image_paths": [f"{sample_id}-ref.jpg"], + "conditioning_resolutions": [(64, 64)], + "target_latent_shape": (16, 8, 8), + "conditioning_latent_shapes": [(16, 8, 8)], + "sample_id": sample_id, + "bucket_id": 0, + "aspect_ratio": 1.0, + } + + output = collate_fn_image_to_image([sample(3, 2, "a"), sample(5, 4, "b")]) + + assert output["image_latents"].shape == (2, 16, 8, 8) + assert len(output["conditioning_latents"]) == 1 + assert output["conditioning_latents"][0].shape == (2, 16, 8, 8) + assert output["text_embeddings"].shape == (2, 5, 32) + assert output["text_embeddings_mask"].tolist() == [ + [1, 1, 1, 0, 0], + [1, 1, 1, 1, 1], + ] + assert output["negative_text_embeddings"].shape == (2, 4, 32) + assert output["negative_text_embeddings_mask"].tolist() == [ + [1, 1, 0, 0], + [1, 1, 1, 1], + ] + + +def test_qwen_image_edit_processor_is_registered(): + pytest.importorskip("torch") + from preprocess.processors import ProcessorRegistry + + assert ProcessorRegistry.is_registered("qwen_image_edit") diff --git a/tests/unit/torch/fastgen/test_qwen_image_edit_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_edit_plugin.py new file mode 100644 index 00000000000..2cb379d8f66 --- /dev/null +++ b/tests/unit/torch/fastgen/test_qwen_image_edit_plugin.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Unit tests for the Qwen-Image-Edit DMD2 target/reference-token wrapper.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from modelopt.torch.fastgen import DMDConfig +from modelopt.torch.fastgen.config import SampleTimestepConfig +from modelopt.torch.fastgen.plugins.qwen_image import ( + attach_feature_capture as attach_t2i_feature_capture, +) +from modelopt.torch.fastgen.plugins.qwen_image import ( + pack_latents, + remove_feature_capture, + update_feature_capture_shape, +) +from modelopt.torch.fastgen.plugins.qwen_image_edit import QwenImageEditDMDPipeline +from modelopt.torch.fastgen.plugins.qwen_image_edit import ( + attach_feature_capture as attach_edit_feature_capture, +) + + +class _CapturingModel(nn.Module): + """Record Qwen kwargs and return the packed input in a requested output style.""" + + def __init__(self, style: str = "tensor") -> None: + super().__init__() + self.style = style + self.last_kwargs: dict[str, object] = {} + + def forward(self, **kwargs): + self.last_kwargs = dict(kwargs) + output = kwargs["hidden_states"] + if self.style == "tensor": + return output + if self.style == "tuple": + return (output,) + if self.style == "sample": + return SimpleNamespace(sample=output) + raise ValueError(self.style) + + +class _CurrentDiffusersSignatureModel(nn.Module): + """Approximate current Diffusers Qwen forward, which removed ``txt_seq_lens``.""" + + def __init__(self) -> None: + super().__init__() + self.called = False + + def forward( + self, + hidden_states, + encoder_hidden_states, + encoder_hidden_states_mask, + timestep, + img_shapes, + guidance=None, + attention_kwargs=None, + return_dict=True, + ): + self.called = True + return hidden_states + + +class _GenericForwardWrapper(_CurrentDiffusersSignatureModel): + """Mimic a distributed wrapper that exposes ``**kwargs`` around an explicit model API.""" + + def forward(self, *args, **kwargs): + return super().forward(*args, **kwargs) + + +def _make_pipeline(student: nn.Module, *, guidance: float | None = None): + return QwenImageEditDMDPipeline( + student=student, + teacher=nn.Identity(), + fake_score=nn.Identity(), + config=DMDConfig(num_train_timesteps=None), + discriminator=None, + guidance=guidance, + ) + + +@pytest.mark.parametrize("style", ["tensor", "tuple", "sample"]) +def test_call_model_appends_multiple_references_and_crops_target_prefix(style): + """Packed references follow the target, while only target output tokens are unpacked.""" + b, c, h, w = 2, 4, 8, 10 + target = torch.arange(b * c * h * w, dtype=torch.float32).reshape(b, c, h, w) + references = [ + torch.full((b, c, 6, 8), 10_000.0), + torch.full((b, c, 4, 12), 20_000.0), + ] + model = _CapturingModel(style) + pipe = _make_pipeline(model, guidance=1.25) + timestep = torch.tensor([0.25, 0.75]) + text = torch.randn(b, 7, 16) + text_mask = torch.tensor([[1, 1, 1, 0, 0, 0, 0], [1, 1, 1, 1, 1, 0, 0]]) + + output = pipe._call_model( + model, + target, + timestep, + encoder_hidden_states=text, + encoder_hidden_states_mask=text_mask, + conditioning_latents=references, + ) + + expected_packed = torch.cat( + [pack_latents(target), *(pack_latents(x) for x in references)], dim=1 + ) + kwargs = model.last_kwargs + assert torch.equal(kwargs["hidden_states"], expected_packed) + assert kwargs["img_shapes"] == [ + [(1, 4, 5), (1, 3, 4), (1, 2, 6)], + [(1, 4, 5), (1, 3, 4), (1, 2, 6)], + ] + assert "txt_seq_lens" not in kwargs + assert torch.equal(kwargs["timestep"], timestep) + assert torch.equal(kwargs["guidance"], torch.full((b,), 1.25)) + assert kwargs["return_dict"] is False + assert "conditioning_latents" not in kwargs + + # The capturing model echoes target+reference tokens. Cropping the prefix before + # unpacking must recover the target bit-exactly rather than trying to unpack all tokens. + assert torch.equal(output, target) + + +def test_conditioning_latent_validation_errors_are_actionable(): + b, c, h, w = 1, 4, 8, 8 + target = torch.randn(b, c, h, w) + text = torch.randn(b, 2, 8) + timestep = torch.tensor([0.5]) + pipe = _make_pipeline(_CapturingModel()) + + def call(conditioning_latents): + return pipe._call_model( + pipe.student, + target, + timestep, + encoder_hidden_states=text, + conditioning_latents=conditioning_latents, + ) + + with pytest.raises(ValueError, match="non-empty.*conditioning_latents"): + call(None) + with pytest.raises(ValueError, match="list or tuple"): + call(torch.randn_like(target)) + with pytest.raises(TypeError, match=r"conditioning_latents\[0\].*Tensor"): + call(["not-a-tensor"]) + with pytest.raises(ValueError, match=r"conditioning_latents\[0\].*\[B, C, H, W\]"): + call([torch.randn(c, h, w)]) + with pytest.raises(ValueError, match="batch/channels"): + call([torch.randn(2, c, h, w)]) + with pytest.raises(ValueError, match="even spatial"): + call([torch.randn(b, c, h - 1, w)]) + with pytest.raises(ValueError, match="dtype"): + call([torch.randn(b, c, h, w, dtype=torch.bfloat16)]) + + +def test_current_diffusers_signature_does_not_receive_removed_txt_seq_lens(): + model = _GenericForwardWrapper() + pipe = _make_pipeline(model) + target = torch.randn(1, 4, 8, 8) + reference = torch.randn(1, 4, 8, 8) + text = torch.randn(1, 3, 8) + mask = torch.ones(1, 3, dtype=torch.long) + + output = pipe._call_model( + model, + target, + torch.tensor([0.5]), + encoder_hidden_states=text, + encoder_hidden_states_mask=mask, + conditioning_latents=[reference], + ) + + assert model.called + assert output.shape == target.shape + + +class _TinyEditTransformer(nn.Module): + """Grad-capable packed-token transformer used for an end-to-end DMD loss call.""" + + def __init__(self, packed_dim: int = 16) -> None: + super().__init__() + self.proj = nn.Linear(packed_dim, packed_dim) + self.seen_token_counts: list[int] = [] + + def forward(self, hidden_states, **_kwargs): + self.seen_token_counts.append(hidden_states.shape[1]) + return self.proj(hidden_states) + + +def test_shared_dmd_losses_forward_references_to_student_teacher_and_fake_score(): + """All DMD branches receive the fixed reference suffix through ``model_kwargs``.""" + torch.manual_seed(0) + student = _TinyEditTransformer() + teacher = _TinyEditTransformer() + fake_score = _TinyEditTransformer() + config = DMDConfig( + pred_type="flow", + num_train_timesteps=None, + student_sample_steps=1, + guidance_scale=None, + gan_loss_weight_gen=0.0, + sample_t_cfg=SampleTimestepConfig(time_dist_type="uniform", min_t=0.001, max_t=0.999), + ema=None, + ) + pipe = QwenImageEditDMDPipeline(student, teacher, fake_score, config) + target = torch.randn(1, 4, 8, 8) # 16 target patches + noise = torch.randn_like(target) + reference = torch.randn(1, 4, 4, 8) # 8 reference patches + text = torch.randn(1, 3, 8) + + student_losses = pipe.compute_student_loss( + target, + noise, + encoder_hidden_states=text, + conditioning_latents=[reference], + ) + assert torch.isfinite(student_losses["total"]) + student_losses["total"].backward() + assert any(p.grad is not None for p in student.parameters()) + assert student.seen_token_counts == [24] + assert teacher.seen_token_counts == [24] + assert fake_score.seen_token_counts == [24] + + fake_score.zero_grad(set_to_none=True) + fake_losses = pipe.compute_fake_score_loss( + target, + noise, + encoder_hidden_states=text, + conditioning_latents=(reference,), + ) + assert torch.isfinite(fake_losses["total"]) + fake_losses["total"].backward() + assert any(p.grad is not None for p in fake_score.parameters()) + assert student.seen_token_counts[-1] == 24 + assert fake_score.seen_token_counts[-1] == 24 + + +class _TupleBlock(nn.Module): + def forward(self, hidden_states): + return torch.empty(0), hidden_states + + +class _TeacherWithBlocks(nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer_blocks = nn.ModuleList([_TupleBlock()]) + + +def test_edit_feature_capture_keeps_target_prefix_and_t2i_remains_strict(): + b, target_h, target_w, channels = 1, 8, 6, 5 + target_tokens = (target_h // 2) * (target_w // 2) + hidden = torch.arange(b * (target_tokens + 7) * channels, dtype=torch.float32).reshape( + b, target_tokens + 7, channels + ) + + edit_teacher = _TeacherWithBlocks() + attach_edit_feature_capture(edit_teacher, [0], target_h, target_w) + edit_teacher.transformer_blocks[0](hidden) + captured = edit_teacher._fastgen_captured + assert len(captured) == 1 + expected = ( + hidden[:, :target_tokens] + .permute(0, 2, 1) + .reshape(b, channels, target_h // 2, target_w // 2) + ) + assert torch.equal(captured[0], expected) + + # The installed hooks must follow later multiresolution batches instead of retaining + # the base-resolution shape present at hook registration time. + captured.clear() + dynamic_h, dynamic_w = 4, 8 + dynamic_tokens = (dynamic_h // 2) * (dynamic_w // 2) + dynamic_hidden = hidden[:, : dynamic_tokens + 3] + update_feature_capture_shape(edit_teacher, dynamic_h, dynamic_w) + edit_teacher.transformer_blocks[0](dynamic_hidden) + dynamic_expected = ( + dynamic_hidden[:, :dynamic_tokens] + .permute(0, 2, 1) + .reshape(b, channels, dynamic_h // 2, dynamic_w // 2) + ) + assert torch.equal(captured[0], dynamic_expected) + remove_feature_capture(edit_teacher) + + # The text-to-image helper keeps exact-length validation by default, preventing + # accidental resolution drift from being silently interpreted as edit references. + t2i_teacher = _TeacherWithBlocks() + attach_t2i_feature_capture(t2i_teacher, [0], target_h, target_w) + with pytest.raises(RuntimeError, match="seq_len"): + t2i_teacher.transformer_blocks[0](hidden) + remove_feature_capture(t2i_teacher) diff --git a/tests/unit/torch/fastgen/test_qwen_image_plugin.py b/tests/unit/torch/fastgen/test_qwen_image_plugin.py index 498b6ce5f9f..07bd75fdce4 100644 --- a/tests/unit/torch/fastgen/test_qwen_image_plugin.py +++ b/tests/unit/torch/fastgen/test_qwen_image_plugin.py @@ -33,6 +33,7 @@ from torch import nn from modelopt.torch.fastgen import DMDConfig +from modelopt.torch.fastgen.plugins import qwen_image as qwen_image_plugin from modelopt.torch.fastgen.plugins.qwen_image import ( QwenImageDMDPipeline, build_img_shapes, @@ -162,12 +163,13 @@ def _make_pipeline(student: nn.Module) -> QwenImageDMDPipeline: ) -def test_call_model_forwards_qwen_kwargs(): +def test_call_model_forwards_qwen_kwargs(monkeypatch): """``_call_model`` must forward the exact Qwen signature (hidden_states packed to ``[B, num_patches, 64]``, encoder_hidden_states verbatim, - encoder_hidden_states_mask verbatim, txt_seq_lens derived from the mask, + encoder_hidden_states_mask verbatim (Diffusers derives sequence lengths from it), img_shapes as ``[[(1, h//2, w//2)]] * B``, guidance=None, return_dict=False, timestep verbatim with no /1000 rescale).""" + monkeypatch.setattr(qwen_image_plugin, "_DIFFUSERS_NEEDS_TXT_SEQ_LENS", False) b, c, h, w = 2, 16, 32, 32 student = _CapturingModel(out_shape=(b, (h // 2) * (w // 2), c * 4), style="tensor") pipe = _make_pipeline(student) @@ -191,7 +193,7 @@ def test_call_model_forwards_qwen_kwargs(): assert tuple(kw["hidden_states"].shape) == (b, (h // 2) * (w // 2), c * 4) assert tuple(kw["encoder_hidden_states"].shape) == (b, 512, 3584) assert torch.equal(kw["encoder_hidden_states_mask"], mask) - assert kw["txt_seq_lens"] == [37, 42] + assert "txt_seq_lens" not in kw assert kw["img_shapes"] == [[(1, h // 2, w // 2)]] * b assert kw["guidance"] is None assert kw["return_dict"] is False @@ -199,6 +201,27 @@ def test_call_model_forwards_qwen_kwargs(): assert tuple(out.shape) == (b, c, h, w) +def test_call_model_forwards_legacy_txt_seq_lens(monkeypatch): + """Diffusers 0.35/0.36 still needs lengths derived from the attention mask.""" + monkeypatch.setattr(qwen_image_plugin, "_DIFFUSERS_NEEDS_TXT_SEQ_LENS", True) + b, c, h, w = 2, 16, 8, 8 + student = _CapturingModel(out_shape=(b, (h // 2) * (w // 2), c * 4)) + pipe = _make_pipeline(student) + mask = torch.zeros(b, 9, dtype=torch.long) + mask[0, :4] = 1 + mask[1, :7] = 1 + + pipe._call_model( + student, + torch.randn(b, c, h, w), + torch.tensor([0.25, 0.5]), + encoder_hidden_states=torch.randn(b, 9, 32), + encoder_hidden_states_mask=mask, + ) + + assert student.last_kwargs["txt_seq_lens"] == [4, 7] + + @pytest.mark.parametrize("style", ["tensor", "tuple", "sample"]) def test_call_model_unpacks_return_styles(style): """``_call_model`` must unpack ``tensor`` / ``tuple`` / ``.sample`` return