From 50f4e675c195ad8bcfdeba2fe34d3217f97ab422 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Mon, 10 Aug 2026 14:13:30 +0800 Subject: [PATCH] docs: remove quantization example sources Signed-off-by: Hongsheng Liu --- .../quantization/check_modelopt_fp8_export.py | 247 --------- .../quantize_hunyuanvideo_15_modelopt_fp8.py | 448 ----------------- .../quantize_wan2_2_modelopt_fp8.py | 471 ------------------ 3 files changed, 1166 deletions(-) delete mode 100644 examples/quantization/check_modelopt_fp8_export.py delete mode 100644 examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py delete mode 100644 examples/quantization/quantize_wan2_2_modelopt_fp8.py diff --git a/examples/quantization/check_modelopt_fp8_export.py b/examples/quantization/check_modelopt_fp8_export.py deleted file mode 100644 index 90d952a8f..000000000 --- a/examples/quantization/check_modelopt_fp8_export.py +++ /dev/null @@ -1,247 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Verify a ModelOpt FP8 diffusers checkpoint exported by -quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py. - -Three checks: - A. transformer/config.json has a sane quantization_config block. - B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors. - C. transformer disk size is materially smaller than a BF16 baseline. - -Example: - python examples/quantization/check_modelopt_fp8_export.py \\ - --output ./hv15-480p-modelopt-fp8 - - # Optional: compare disk size against a local or HF BF16 baseline. - python examples/quantization/check_modelopt_fp8_export.py \\ - --output ./hv15-480p-modelopt-fp8 \\ - --baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v -""" - -from __future__ import annotations - -import argparse -import json -import sys -from collections import Counter -from pathlib import Path - -SUPPORTED_ALGOS = ("FP8",) - - -def _check_config(transformer_dir: Path) -> tuple[int, str | None]: - """Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn.""" - cfg_path = transformer_dir / "config.json" - if not cfg_path.exists(): - print(f"[FAIL] {cfg_path} missing.") - return 1, None - - with cfg_path.open(encoding="utf-8") as f: - cfg = json.load(f) - - qc = cfg.get("quantization_config") - if not isinstance(qc, dict): - print(f"[FAIL] No `quantization_config` block in {cfg_path}.") - return 1, None - - print(f"[A] quantization_config from {cfg_path}:") - print(json.dumps(qc, indent=2)) - - quant_algo = qc.get("quant_algo") - issues = [] - if qc.get("quant_method") != "modelopt": - issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')") - if quant_algo not in SUPPORTED_ALGOS: - issues.append( - f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)" - ) - - if issues: - print("[A] WARN — config looks incomplete:") - for issue in issues: - print(f" - {issue}") - return 2, quant_algo - print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).") - return 0, quant_algo - - -def _read_safetensors_header(path: Path) -> dict: - """Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors. - - Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}. - Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ... - """ - import struct - - with open(path, "rb") as f: - header_len = struct.unpack(" str: - """Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes. - - ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1 - interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 — - and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block. - """ - if not weight_scale_shapes: - return "no weight_scale tensors found" - - def meaningful_dims(shape: list[int]) -> int: - return sum(1 for d in shape if d > 1) - - per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0) - per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1) - per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2) - total = len(weight_scale_shapes) - if per_tensor == total: - return "per-tensor (all scalar scales)" - if per_channel == total: - return "per-channel (1 meaningful dim)" - if per_block == total: - return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)" - return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}" - - -def _check_safetensors(transformer_dir: Path) -> int: - """Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header.""" - files = sorted(transformer_dir.glob("*.safetensors")) - if not files: - print(f"[FAIL] No *.safetensors in {transformer_dir}.") - return 1 - - header_dtype_counts: Counter[str] = Counter() - sample_quant_weight_keys: list[str] = [] - sample_scale_keys: list[str] = [] - weight_scale_shapes: list[list[int]] = [] - sample_weight_scale_entries: list[tuple[str, list[int]]] = [] - - for f in files: - try: - header = _read_safetensors_header(f) - except Exception as exc: - print(f"[B] WARN — could not parse header of {f}: {exc}") - continue - for k, info in header.items(): - dtype = info.get("dtype", "?") - header_dtype_counts[dtype] += 1 - # FP8 stores weights as F8_E4M3 directly. - if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5: - sample_quant_weight_keys.append(k) - if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5: - sample_scale_keys.append(k) - if k.endswith(".weight_scale"): - weight_scale_shapes.append(info.get("shape", [])) - if len(sample_weight_scale_entries) < 5: - sample_weight_scale_entries.append((k, info.get("shape", []))) - - print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):") - for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]): - marker = " <-- FP8" if dtype.startswith("F8") else "" - print(f" {dtype:10s} {count:>6d}{marker}") - - quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8")) - if quant_count == 0: - print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.") - return 1 - - print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.") - if sample_quant_weight_keys: - print(f" sample quantized weight tensors: {sample_quant_weight_keys[:3]}") - if sample_scale_keys: - print(f" sample scale tensors: {sample_scale_keys[:3]}") - print(" (Note: torch's get_tensor() may return these as bf16 views on some versions —") - print(" irrelevant; vLLM's loader uses native FP8 ops.)") - - # Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D). - print(f"\n weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}") - for key, shape in sample_weight_scale_entries[:3]: - print(f" {key}: shape {shape}") - return 0 - - -def _disk_size_gib(p: Path) -> float: - return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3) - - -def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int: - """Returns 0 always (informational only).""" - quant_size = _disk_size_gib(transformer_dir) - print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB") - - if baseline is None: - print("[C] SKIP — pass --baseline to compare against BF16.") - return 0 - - baseline_path = Path(baseline) - if not baseline_path.exists(): - # Try HF download. - try: - from huggingface_hub import snapshot_download - except ImportError: - print("[C] SKIP — huggingface_hub not installed and baseline not a local path.") - return 0 - print(f" Downloading baseline transformer from HF: {baseline}") - baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"])) - - bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path - bf16_size = _disk_size_gib(bf16_dir) - if bf16_size == 0: - print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}") - return 0 - - # Expected reduction for FP8: ~50%. - min_reduction = 30 - reduction = (1 - quant_size / bf16_size) * 100 - print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})") - print(f"[C] Disk reduction: {reduction:.1f}% (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)") - if reduction < min_reduction: - print( - f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization." - ) - return 0 - - -def main() -> None: - p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.") - p.add_argument( - "--baseline", - default=None, - help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.", - ) - args = p.parse_args() - - out_root = Path(args.output).expanduser().resolve() - transformer_dir = out_root / "transformer" - if not transformer_dir.exists(): - print(f"[FAIL] {transformer_dir} does not exist.") - sys.exit(1) - - print(f"Checking: {out_root}\n") - - fail = 0 - config_status, _quant_algo = _check_config(transformer_dir) - fail |= config_status - fail |= _check_safetensors(transformer_dir) - _check_size_vs_baseline(transformer_dir, args.baseline) - - print() - if fail == 0: - print("=" * 60) - print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.") - elif fail == 1: - print("=" * 60) - print("FAILURES detected — calibration may need to be re-run.") - sys.exit(1) - else: - print("=" * 60) - print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.") - - -if __name__ == "__main__": - main() diff --git a/examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py b/examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py deleted file mode 100644 index e6e3dd5c5..000000000 --- a/examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py +++ /dev/null @@ -1,448 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint. - -Calibrates the DiT transformer using a small video prompt set and exports a -diffusers-style directory whose transformer carries ModelOpt FP8 metadata. -The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter -(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py). - -Layers kept full precision match the #2728 / #2795 pattern: modulation, -AdaLayerNorm, entry/exit projections, embeddings, the token refiner path, -and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention -empirically degrades under FP8 (see #2920 ablation). - -Example: - python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\ - --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\ - --output ./hv15-480p-modelopt-fp8 \\ - --overwrite -""" - -from __future__ import annotations - -import argparse -import copy -import json -import re -import shutil -import sys -from pathlib import Path -from typing import Any - -import torch -from diffusers import DiffusionPipeline - -DEFAULT_PROMPTS = [ - "A dog running across a field of golden wheat.", - "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.", - "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.", - "A crackling campfire at night under a starry sky, sparks rising into the dark.", - "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.", - "A close-up of a blooming rose covered in morning dew, soft natural light.", - "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.", - "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.", -] - - -def _build_parser() -> argparse.ArgumentParser: - p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.") - p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.") - p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16") - p.add_argument("--height", type=int, default=480) - p.add_argument("--width", type=int, default=832) - p.add_argument( - "--num-frames", - type=int, - default=33, - help="Frames per calibration sample. 33 matches the typical short benchmark.", - ) - p.add_argument("--guidance-scale", type=float, default=6.0) - p.add_argument( - "--calib-steps", - type=int, - default=10, - help="Denoising steps per calibration prompt (10 is enough for amax statistics).", - ) - p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.") - p.add_argument("--seed", type=int, default=42) - p.add_argument( - "--prompt", - action="append", - default=[], - help="Custom calibration prompt. Repeat to provide multiple.", - ) - p.add_argument( - "--quantize-mha", - action="store_true", - help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.", - ) - p.add_argument( - "--weight-block-size", - type=str, - default=None, - help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). " - "Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at " - "negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.", - ) - p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.") - return p - - -def _parse_block_size(spec: str | None) -> list[int] | None: - if spec is None: - return None - parts = [int(x) for x in spec.split(",") if x.strip()] - if len(parts) != 2: - raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}") - return parts - - -def _require_modelopt() -> Any: - try: - import modelopt.torch.quantization as mtq - except ModuleNotFoundError as exc: - raise SystemExit( - "NVIDIA ModelOpt is not installed. Install with:\n" - " pip install 'nvidia-modelopt[all]'\n" - f"Original error: {exc}" - ) from exc - return mtq - - -def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]: - model_path = args.model - output_dir = Path(args.output).expanduser().resolve() - if output_dir.exists(): - if not args.overwrite: - raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.") - shutil.rmtree(output_dir) - return model_path, output_dir - - -def _select_dtype(name: str) -> torch.dtype: - return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name] - - -def _build_prompts(args: argparse.Namespace) -> list[str]: - prompts = args.prompt or DEFAULT_PROMPTS - if args.calib_size <= 0: - raise SystemExit("--calib-size must be positive.") - if len(prompts) < args.calib_size: - repeats = (args.calib_size + len(prompts) - 1) // len(prompts) - prompts = (prompts * repeats)[: args.calib_size] - return prompts[: args.calib_size] - - -# Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern). -# - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding -# - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation -# - proj_out: final output projection -# - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear -def _filter_func_hv15(name: str) -> bool: - pattern = re.compile( - r"(proj_out.*|" - r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|" - r"time_embed|cond_type_embed|" - r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|" - r"token_refiner).*)" - ) - return pattern.match(name) is not None - - -def _mha_filter_func(name: str) -> bool: - pattern = re.compile( - r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*" - ) - return pattern.match(name) is not None - - -def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None: - if not hasattr(mtq, "disable_quantizer"): - return - mtq.disable_quantizer(backbone, _filter_func_hv15) - if not quantize_mha: - mtq.disable_quantizer(backbone, _mha_filter_func) - - -def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline: - pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype) - if hasattr(pipe, "set_progress_bar_config"): - pipe.set_progress_bar_config(disable=True) - pipe.to("cuda") - return pipe - - -def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]): - generator = torch.Generator(device="cuda") - - # Try to set guidance on the pipeline's guider object up front (modern - # diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls - # back silently — calibration uses whatever default the pipeline ships with. - guider = getattr(pipe, "guider", None) - if guider is not None and hasattr(guider, "guidance_scale"): - try: - guider.guidance_scale = args.guidance_scale - except Exception: - pass - - base_kwargs = dict( - height=args.height, - width=args.width, - num_frames=args.num_frames, - num_inference_steps=args.calib_steps, - output_type="latent", - ) - - def forward_loop(*_unused_args, **_unused_kwargs) -> None: - with torch.inference_mode(): - for idx, prompt in enumerate(prompts): - generator.manual_seed(args.seed + idx) - # Try with guidance_scale first; fall back without on TypeError - # for pipelines (like HV-1.5) that take CFG via guider config. - try: - pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs) - except TypeError as exc: - if "guidance_scale" not in str(exc): - raise - pipe(prompt=prompt, generator=generator, **base_kwargs) - - return forward_loop - - -def _summarize_export(output_dir: Path) -> None: - cfg_path = output_dir / "transformer" / "config.json" - if not cfg_path.exists(): - print(f"[warn] {cfg_path} missing.", file=sys.stderr) - return - with cfg_path.open(encoding="utf-8") as f: - cfg = json.load(f) - qc = cfg.get("quantization_config") - if not isinstance(qc, dict): - print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr) - return - print("Export summary:") - print(f" quant_method: {qc.get('quant_method')}") - print(f" quant_algo: {qc.get('quant_algo')}") - producer = qc.get("producer") - if isinstance(producer, dict): - print(f" producer: {producer.get('name')} {producer.get('version')}") - print(f" config path: {cfg_path}") - - -def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int: - """Convert in-memory weights of quantized modules to actual FP8 storage. - - `export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't - in ModelOpt's recognized-model registry), so we must call the per-weight - export helper ourselves. Same workaround as the HunyuanImage-3 calibration - helper. - """ - from modelopt.torch.export.quant_utils import ( - QUANTIZATION_NONE, - get_quantization_format, - quantizer_attr_names, - weight_attr_names, - ) - from modelopt.torch.export.unified_export_hf import _export_quantized_weight - - exported = 0 - for name, module in backbone.named_modules(): - try: - quantization_format = get_quantization_format(module) - except Exception as exc: - print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr) - continue - if quantization_format == QUANTIZATION_NONE: - continue - for weight_name in weight_attr_names(module): - quantizer_attrs = quantizer_attr_names(weight_name) - weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None) - if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False): - continue - _export_quantized_weight(module, dtype, weight_name) - exported += 1 - return exported - - -def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict: - """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913). - - Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When - `weight_block_size` is given, advertise block-wise weight quantization in - the saved metadata (so consumers know to expect multi-element scale tensors). - """ - weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"} - if weight_block_size is not None: - weights_cfg["strategy"] = "block" - weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}" - return { - "config_groups": { - "group_0": { - "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"}, - "weights": weights_cfg, - "targets": ["Linear"], - } - }, - "ignore": [ - "context_embedder*", - "context_embedder_2*", - "cond_type_embed*", - "image_embedder*", - "norm1.linear*", - "norm1_context.linear*", - "norm2*", - "norm2_context*", - "norm_out*", - "proj_out*", - "time_embed*", - "token_refiner*", - "x_embedder*", - ], - "producer": {"name": "modelopt"}, - "quant_algo": "FP8", - "quant_method": "modelopt", - } - - -def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None: - """Inject quant_algo: FP8 + config_groups into transformer/config.json so - vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.""" - cfg_path = output_dir / "transformer" / "config.json" - with cfg_path.open(encoding="utf-8") as f: - cfg = json.load(f) - - new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size) - existing = cfg.get("quantization_config") - if isinstance(existing, dict): - producer = existing.get("producer") - if isinstance(producer, dict): - new_qc["producer"] = producer - - cfg["quantization_config"] = new_qc - with cfg_path.open("w", encoding="utf-8") as f: - json.dump(cfg, f, indent=2) - - -def _save_pipeline_with_fp8_transformer( - pipe: DiffusionPipeline, - model_path: str, - output_dir: Path, - max_shard_size: str = "5GB", -) -> None: - """Save the pipeline with the (now FP8) transformer. - - Copies the source directory verbatim except for `transformer/`, then - saves the transformer with quantizers hidden so the state dict contains - only the FP8 weights + scale tensors. - """ - from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict - - src = Path(model_path) - if not src.exists(): - from huggingface_hub import snapshot_download - - src = Path(snapshot_download(model_path)) - - if output_dir.exists(): - shutil.rmtree(output_dir) - shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer")) - - transformer_out = output_dir / "transformer" - # `hide_quantizers_from_state_dict` walks named_modules(); pass the actual - # nn.Module (transformer), not the diffusers Pipeline wrapper. - with hide_quantizers_from_state_dict(pipe.transformer): - pipe.transformer.save_pretrained( - str(transformer_out), - safe_serialization=True, - max_shard_size=max_shard_size, - ) - - -def main() -> None: - args = _build_parser().parse_args() - if not torch.cuda.is_available(): - raise SystemExit("CUDA is required for ModelOpt FP8 quantization.") - - mtq = _require_modelopt() - model_path, output_dir = _ensure_paths(args) - dtype = _select_dtype(args.dtype) - prompts = _build_prompts(args) - - print("Quantization plan:") - weight_block_size = _parse_block_size(args.weight_block_size) - - print(f" input: {args.model}") - print(f" output: {output_dir}") - print(f" dtype: {dtype}") - print(f" height/width: {args.height}x{args.width}") - print(f" num_frames: {args.num_frames}") - print(f" calib_size: {len(prompts)}") - print(f" calib_steps: {args.calib_steps}") - print(f" quantize_mha: {args.quantize_mha}") - print( - f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}" - ) - - pipe = _load_pipeline(model_path, dtype) - backbone = pipe.transformer - - quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) - if weight_block_size is not None: - # Switch from per-tensor (default) to block-wise weight quantization. - # ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer. - quant_config["quant_cfg"]["*weight_quantizer"] = { - "num_bits": (4, 3), # E4M3 (FP8 weights, same as default) - "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]}, - } - print( - f" -> overriding weight quantizer with block_sizes={weight_block_size} " - f"({weight_block_size[0]}x{weight_block_size[1]} tiles)" - ) - - forward_loop = _build_forward_loop(pipe, args, prompts) - quantized = mtq.quantize(backbone, quant_config, forward_loop) - if quantized is not None: - pipe.transformer = quantized - backbone = quantized - - _disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha) - - print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,") - print("so we have to call the per-weight export helper ourselves)...") - exported = _force_export_quantized_weights(backbone, dtype) - print(f" -> {exported} weights converted to FP8 in memory") - if exported == 0: - raise SystemExit( - "No quantized weights were exported. Calibration may have skipped every layer " - "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any " - "weight quantizers." - ) - - print("\nSaving pipeline with FP8 transformer...") - _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir) - _patch_quant_config(output_dir, weight_block_size=weight_block_size) - print(f"Saved to: {output_dir}") - _summarize_export(output_dir) - - print("\nNext: validate the checkpoint with vllm-omni:") - print( - " python examples/offline_inference/text_to_video/text_to_video.py \\\n" - f" --model {output_dir} \\\n" - " --quantization fp8 \\\n" - " --prompt 'A dog running across a field of golden wheat.' \\\n" - f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n" - " --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n" - " --output outputs/hv15_modelopt_fp8.mp4 \\\n" - " --enforce-eager" - ) - print( - "\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the " - "checkpoint's config.json has modelopt metadata.)" - ) - - -if __name__ == "__main__": - main() diff --git a/examples/quantization/quantize_wan2_2_modelopt_fp8.py b/examples/quantization/quantize_wan2_2_modelopt_fp8.py deleted file mode 100644 index 6b6cc42a4..000000000 --- a/examples/quantization/quantize_wan2_2_modelopt_fp8.py +++ /dev/null @@ -1,471 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint. - -Calibrates the DiT transformer using a small video prompt set and exports a -diffusers-style directory whose transformer carries ModelOpt FP8 metadata. -The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter -(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py). - -Layers kept full precision match the #2728 / #2795 pattern: condition embedder -(time/text/image), patch embedding, modulation (scale_shift_table), final -norm + proj_out, and sequence-parallel helpers. All attention + FFN linears -are quantized — static calibration handles the numerics that online FP8 -couldn't (see #2920 ablation). - -Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that -fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here. - -Example: - python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\ - --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\ - --output ./wan22-ti2v-modelopt-fp8 \\ - --overwrite -""" - -from __future__ import annotations - -import argparse -import copy -import json -import re -import shutil -import sys -from pathlib import Path -from typing import Any - -import torch -from diffusers import DiffusionPipeline - -DEFAULT_PROMPTS = [ - "A dog running across a field of golden wheat.", - "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.", - "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.", - "A crackling campfire at night under a starry sky, sparks rising into the dark.", - "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.", - "A close-up of a blooming rose covered in morning dew, soft natural light.", - "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.", - "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.", -] - - -def _build_parser() -> argparse.ArgumentParser: - p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.") - p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.") - p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16") - p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).") - p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).") - p.add_argument( - "--num-frames", - type=int, - default=49, - help="Frames per calibration sample. 49 matches the typical short benchmark; " - "use 17 to reduce memory pressure during calibration.", - ) - p.add_argument("--guidance-scale", type=float, default=5.0) - p.add_argument( - "--calib-steps", - type=int, - default=10, - help="Denoising steps per calibration prompt (10 is enough for amax statistics).", - ) - p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.") - p.add_argument("--seed", type=int, default=42) - p.add_argument( - "--prompt", - action="append", - default=[], - help="Custom calibration prompt. Repeat to provide multiple.", - ) - p.add_argument( - "--quantize-mha", - action="store_true", - help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention " - "sequences amplified FP8 drift in the online ablation (see #2920).", - ) - p.add_argument( - "--weight-block-size", - type=str, - default=None, - help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. " - "Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 " - "for the HV-1.5 investigation status before relying on this.", - ) - p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.") - return p - - -def _parse_block_size(spec: str | None) -> list[int] | None: - if spec is None: - return None - parts = [int(x) for x in spec.split(",") if x.strip()] - if len(parts) != 2: - raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}") - return parts - - -def _require_modelopt() -> Any: - try: - import modelopt.torch.quantization as mtq - except ModuleNotFoundError as exc: - raise SystemExit( - "NVIDIA ModelOpt is not installed. Install with:\n" - " pip install 'nvidia-modelopt[all]'\n" - f"Original error: {exc}" - ) from exc - return mtq - - -def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]: - model_path = args.model - output_dir = Path(args.output).expanduser().resolve() - if output_dir.exists(): - if not args.overwrite: - raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.") - shutil.rmtree(output_dir) - return model_path, output_dir - - -def _select_dtype(name: str) -> torch.dtype: - return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name] - - -def _build_prompts(args: argparse.Namespace) -> list[str]: - prompts = args.prompt or DEFAULT_PROMPTS - if args.calib_size <= 0: - raise SystemExit("--calib-size must be positive.") - if len(prompts) < args.calib_size: - repeats = (args.calib_size + len(prompts) - 1) // len(prompts) - prompts = (prompts * repeats)[: args.calib_size] - return prompts[: args.calib_size] - - -# Layers to KEEP at full precision. Wan2.2's module naming: -# - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V) -# - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip) -# - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard) -# - norm_out: AdaLayerNorm final -# - proj_out: final nn.Linear -# - timestep_proj_prepare / output_scale_shift_prepare: SP helpers -def _filter_func_wan22(name: str) -> bool: - pattern = re.compile( - r"(proj_out.*|" - r".*(condition_embedder|patch_embedding|" - r"norm_out|scale_shift_table|" - r"timestep_proj_prepare|output_scale_shift_prepare).*)" - ) - return pattern.match(name) is not None - - -def _mha_filter_func(name: str) -> bool: - pattern = re.compile( - r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*" - ) - return pattern.match(name) is not None - - -def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None: - if not hasattr(mtq, "disable_quantizer"): - return - mtq.disable_quantizer(backbone, _filter_func_wan22) - if not quantize_mha: - mtq.disable_quantizer(backbone, _mha_filter_func) - - -def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline: - pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype) - if hasattr(pipe, "set_progress_bar_config"): - pipe.set_progress_bar_config(disable=True) - pipe.to("cuda") - return pipe - - -def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]): - generator = torch.Generator(device="cuda") - - # Try setting guidance on the pipeline's guider if present (newer diffusers APIs). - guider = getattr(pipe, "guider", None) - if guider is not None and hasattr(guider, "guidance_scale"): - try: - guider.guidance_scale = args.guidance_scale - except Exception: - pass - - base_kwargs = dict( - height=args.height, - width=args.width, - num_frames=args.num_frames, - num_inference_steps=args.calib_steps, - output_type="latent", - ) - - def forward_loop(*_unused_args, **_unused_kwargs) -> None: - with torch.inference_mode(): - for idx, prompt in enumerate(prompts): - generator.manual_seed(args.seed + idx) - # Try with guidance_scale first; fall back without on TypeError - # for pipelines that take CFG via guider config only. - try: - pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs) - except TypeError as exc: - if "guidance_scale" not in str(exc): - raise - pipe(prompt=prompt, generator=generator, **base_kwargs) - - return forward_loop - - -def _summarize_export(output_dir: Path) -> None: - cfg_path = output_dir / "transformer" / "config.json" - if not cfg_path.exists(): - print(f"[warn] {cfg_path} missing.", file=sys.stderr) - return - with cfg_path.open(encoding="utf-8") as f: - cfg = json.load(f) - qc = cfg.get("quantization_config") - if not isinstance(qc, dict): - print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr) - return - print("Export summary:") - print(f" quant_method: {qc.get('quant_method')}") - print(f" quant_algo: {qc.get('quant_algo')}") - producer = qc.get("producer") - if isinstance(producer, dict): - print(f" producer: {producer.get('name')} {producer.get('version')}") - print(f" config path: {cfg_path}") - - -def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int: - """Convert in-memory weights of quantized modules to actual FP8 storage. - - `export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't - in ModelOpt's recognized-model registry), so we must call the per-weight - export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3 - calibration helpers. - """ - from modelopt.torch.export.quant_utils import ( - QUANTIZATION_NONE, - get_quantization_format, - quantizer_attr_names, - weight_attr_names, - ) - from modelopt.torch.export.unified_export_hf import _export_quantized_weight - - exported = 0 - for name, module in backbone.named_modules(): - try: - quantization_format = get_quantization_format(module) - except Exception as exc: - print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr) - continue - if quantization_format == QUANTIZATION_NONE: - continue - for weight_name in weight_attr_names(module): - quantizer_attrs = quantizer_attr_names(weight_name) - weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None) - if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False): - continue - _export_quantized_weight(module, dtype, weight_name) - exported += 1 - return exported - - -def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict: - """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).""" - weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"} - if weight_block_size is not None: - weights_cfg["strategy"] = "block" - weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}" - return { - "config_groups": { - "group_0": { - "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"}, - "weights": weights_cfg, - "targets": ["Linear"], - } - }, - "ignore": [ - "condition_embedder*", - "norm_out*", - "output_scale_shift_prepare*", - "patch_embedding*", - "proj_out*", - "scale_shift_table*", - "timestep_proj_prepare*", - ], - "producer": {"name": "modelopt"}, - "quant_algo": "FP8", - "quant_method": "modelopt", - } - - -def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None: - """Inject quant_algo: FP8 + config_groups into each transformer's config.json - so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8. - - Patches both ``transformer`` and ``transformer_2`` (A14B) when present.""" - for sub in ("transformer", "transformer_2"): - cfg_path = output_dir / sub / "config.json" - if not cfg_path.exists(): - continue - with cfg_path.open(encoding="utf-8") as f: - cfg = json.load(f) - - new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size) - existing = cfg.get("quantization_config") - if isinstance(existing, dict): - producer = existing.get("producer") - if isinstance(producer, dict): - new_qc["producer"] = producer - - cfg["quantization_config"] = new_qc - with cfg_path.open("w", encoding="utf-8") as f: - json.dump(cfg, f, indent=2) - - -def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]: - """Return [(attr_name, module)] for every DiT on the pipeline. - - Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2`` - (high-noise). TI2V-5B has only ``transformer``. - """ - out = [] - for attr in ("transformer", "transformer_2"): - mod = getattr(pipe, attr, None) - if mod is not None: - out.append((attr, mod)) - return out - - -def _save_pipeline_with_fp8_transformer( - pipe: DiffusionPipeline, - model_path: str, - output_dir: Path, - max_shard_size: str = "5GB", -) -> None: - """Copy source dir verbatim minus transformer dirs, then save each quantized DiT. - - Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well - as the single-DiT TI2V-5B. - """ - from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict - - src = Path(model_path) - if not src.exists(): - from huggingface_hub import snapshot_download - - src = Path(snapshot_download(model_path)) - - if output_dir.exists(): - shutil.rmtree(output_dir) - shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2")) - - for attr, backbone in _list_transformers(pipe): - transformer_out = output_dir / attr - # Pass the nn.Module (transformer), not the Pipeline wrapper. - with hide_quantizers_from_state_dict(backbone): - backbone.save_pretrained( - str(transformer_out), - safe_serialization=True, - max_shard_size=max_shard_size, - ) - - -def main() -> None: - args = _build_parser().parse_args() - if not torch.cuda.is_available(): - raise SystemExit("CUDA is required for ModelOpt FP8 quantization.") - - mtq = _require_modelopt() - model_path, output_dir = _ensure_paths(args) - dtype = _select_dtype(args.dtype) - prompts = _build_prompts(args) - weight_block_size = _parse_block_size(args.weight_block_size) - - print("Quantization plan:") - print(f" input: {args.model}") - print(f" output: {output_dir}") - print(f" dtype: {dtype}") - print(f" height/width: {args.height}x{args.width}") - print(f" num_frames: {args.num_frames}") - print(f" calib_size: {len(prompts)}") - print(f" calib_steps: {args.calib_steps}") - print(f" quantize_mha: {args.quantize_mha}") - print( - f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}" - ) - - pipe = _load_pipeline(model_path, dtype) - transformers = _list_transformers(pipe) - if not transformers: - raise SystemExit("Pipeline has no transformer or transformer_2 attribute.") - print(f" found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}") - - quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) - if weight_block_size is not None: - quant_config["quant_cfg"]["*weight_quantizer"] = { - "num_bits": (4, 3), - "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]}, - } - print( - f" -> overriding weight quantizer with block_sizes={weight_block_size} " - f"({weight_block_size[0]}x{weight_block_size[1]} tiles)" - ) - - forward_loop = _build_forward_loop(pipe, args, prompts) - - # Quantize each DiT. The forward_loop runs the full pipeline, so quantizers - # on every attached transformer (A14B: transformer + transformer_2) get - # calibrated during the same passes. - for attr, backbone in transformers: - print(f"\nQuantizing {attr}...") - quantized = mtq.quantize(backbone, quant_config, forward_loop) - if quantized is not None: - setattr(pipe, attr, quantized) - _disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha) - - # Export weights to FP8 only AFTER all transformers are calibrated. Exporting - # inside the loop would convert transformer's weights to real FP8, and the - # next transformer's forward_loop (full pipeline) would then fake-quantize an - # already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn". - print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,") - print("so we have to call the per-weight export helper ourselves)...") - total_exported = 0 - for attr, backbone in _list_transformers(pipe): - exported = _force_export_quantized_weights(backbone, dtype) - print(f" -> {exported} weights converted to FP8 in {attr}") - total_exported += exported - - if total_exported == 0: - raise SystemExit( - "No quantized weights were exported. Calibration may have skipped every layer " - "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any " - "weight quantizers." - ) - - print("\nSaving pipeline with FP8 transformer(s)...") - _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir) - _patch_quant_config(output_dir, weight_block_size=weight_block_size) - print(f"Saved to: {output_dir}") - _summarize_export(output_dir) - - print("\nNext: validate the checkpoint with vllm-omni:") - print( - " python examples/offline_inference/text_to_video/text_to_video.py \\\n" - f" --model {output_dir} \\\n" - " --quantization fp8 \\\n" - " --prompt 'A dog running across a field of golden wheat.' \\\n" - f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n" - " --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n" - " --output outputs/wan22_modelopt_fp8.mp4" - ) - print( - "\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the " - "checkpoint's config.json has modelopt metadata.)" - ) - - -if __name__ == "__main__": - main()