From 29b18ea298dff66392aa33c0f68cae6aeca6da4d Mon Sep 17 00:00:00 2001 From: SYLAR <125541396+lishunyang12@users.noreply.github.com> Date: Mon, 8 Jun 2026 02:53:06 +0800 Subject: [PATCH] [Quant] ModelOpt FP8 for Wan2.2 & HunyuanVideo-1.5 video-gen (#3305) Signed-off-by: WeiQing Chen --- .../quantization/check_modelopt_fp8_export.py | 247 +++++++++ .../quantize_hunyuanvideo_15_modelopt_fp8.py | 448 +++++++++++++++++ .../quantize_wan2_2_modelopt_fp8.py | 471 ++++++++++++++++++ .../accuracy/test_videogen_modelopt_quant.py | 273 ++++++++++ .../model_loader/diffusers_loader.py | 9 +- .../hunyuan_video_15_transformer.py | 44 +- .../pipeline_hunyuan_video_1_5.py | 4 +- .../hunyuan_video_15_dit_fp8.yaml | 38 ++ .../stage_configs/wan2_2_ti2v_dit_fp8.yaml | 41 ++ vllm_omni/patch.py | 40 ++ 10 files changed, 1607 insertions(+), 8 deletions(-) create mode 100644 examples/quantization/check_modelopt_fp8_export.py create mode 100644 examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py create mode 100644 examples/quantization/quantize_wan2_2_modelopt_fp8.py create mode 100644 tests/e2e/accuracy/test_videogen_modelopt_quant.py create mode 100644 vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml create mode 100644 vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml diff --git a/examples/quantization/check_modelopt_fp8_export.py b/examples/quantization/check_modelopt_fp8_export.py new file mode 100644 index 000000000..90d952a8f --- /dev/null +++ b/examples/quantization/check_modelopt_fp8_export.py @@ -0,0 +1,247 @@ +#!/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 new file mode 100644 index 000000000..e6e3dd5c5 --- /dev/null +++ b/examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py @@ -0,0 +1,448 @@ +#!/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 new file mode 100644 index 000000000..6b6cc42a4 --- /dev/null +++ b/examples/quantization/quantize_wan2_2_modelopt_fp8.py @@ -0,0 +1,471 @@ +#!/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() diff --git a/tests/e2e/accuracy/test_videogen_modelopt_quant.py b/tests/e2e/accuracy/test_videogen_modelopt_quant.py new file mode 100644 index 000000000..e91c15175 --- /dev/null +++ b/tests/e2e/accuracy/test_videogen_modelopt_quant.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end ModelOpt FP8 quality accuracy for video-gen DiTs. + +Modeled on ``test_hunyuan_image3.py``'s ``test_quantized_dit_matches_bf16_accuracy``: +for each calibrated FP8 checkpoint the test generates one video with BF16 and one +with the FP8 DiT under the same seed, scores prompt-faithfulness with CLIP on the +middle frame, then asserts both ``CLIP >= absolute_floor`` and +``CLIP_drop <= drop_threshold``. The test is gated by an opt-in env var (model +paths are local) so it never runs in generic CI. + +To run:: + + VIDEOGEN_RUN_QUANT_ACCURACY=1 \\ + WAN22_A14B_BF16_MODEL=Wan-AI/Wan2.2-T2V-A14B-Diffusers \\ + WAN22_A14B_FP8_MODEL=/path/to/wan22-a14b-modelopt-fp8 \\ + HV15_BF16_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\ + HV15_FP8_MODEL=/path/to/hv15-modelopt-fp8 \\ + pytest -s -v tests/e2e/accuracy/test_videogen_modelopt_quant.py +""" + +from __future__ import annotations + +import gc +import json +import os +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import torch +from PIL import Image + +from tests.e2e.accuracy.helpers import CLIPScorer, model_output_dir + +pytestmark = [pytest.mark.local_model, pytest.mark.diffusion] + +# --------------------------------------------------------------------------- +# Tunables (override via env to suit local hardware / calibration budget) +# --------------------------------------------------------------------------- +SEED = 42 +HEIGHT = int(os.environ.get("VIDEOGEN_QUANT_HEIGHT", "480")) +WIDTH = int(os.environ.get("VIDEOGEN_QUANT_WIDTH", "832")) +NUM_FRAMES = int(os.environ.get("VIDEOGEN_QUANT_NUM_FRAMES", "17")) +NUM_INFERENCE_STEPS = int(os.environ.get("VIDEOGEN_QUANT_STEPS", "20")) + +WAN_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic." +HV_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic." + +# CLIP gates: absolute floor + max drop from the BF16 reference. +CLIP_ABSOLUTE_FLOOR = float(os.environ.get("VIDEOGEN_QUANT_CLIP_FLOOR", "20.0")) +CLIP_DROP_FP8 = float(os.environ.get("VIDEOGEN_QUANT_CLIP_DROP_FP8", "7.0")) + +RUN_ENV = "VIDEOGEN_RUN_QUANT_ACCURACY" +_TRUE = {"1", "true", "yes", "on"} + + +# --------------------------------------------------------------------------- +@dataclass(frozen=True) +class _Spec: + """A model under test (BF16 + its FP8 variant).""" + + family: str + bf16_env: str + fp8_env: str + prompt: str + guidance_scale: float + guidance_scale_2: float | None = None + boundary_ratio: float | None = None + + +WAN_SPEC = _Spec( + family="wan22_a14b_t2v", + bf16_env="WAN22_A14B_BF16_MODEL", + fp8_env="WAN22_A14B_FP8_MODEL", + prompt=WAN_PROMPT, + guidance_scale=3.5, + guidance_scale_2=4.0, + boundary_ratio=0.875, +) + +HV_SPEC = _Spec( + family="hunyuanvideo15_t2v", + bf16_env="HV15_BF16_MODEL", + fp8_env="HV15_FP8_MODEL", + prompt=HV_PROMPT, + guidance_scale=6.0, +) + + +@dataclass(frozen=True) +class _QuantCase: + name: str + spec: _Spec + clip_drop: float + + +def _cases() -> list[pytest.ParameterSet]: + cases = [ + _QuantCase("wan_fp8", WAN_SPEC, CLIP_DROP_FP8), + _QuantCase("hv_fp8", HV_SPEC, CLIP_DROP_FP8), + ] + params: list[pytest.ParameterSet] = [] + opted_in = os.environ.get(RUN_ENV, "").lower() in _TRUE + for case in cases: + marks: list[Any] = [] + if not opted_in: + marks.append(pytest.mark.skip(reason=f"Set {RUN_ENV}=1 to run video-gen quant accuracy.")) + if not os.environ.get(case.spec.bf16_env): + marks.append(pytest.mark.skip(reason=f"Set {case.spec.bf16_env} to run {case.name}.")) + if not os.environ.get(case.spec.fp8_env): + marks.append(pytest.mark.skip(reason=f"Set {case.spec.fp8_env} to run {case.name}.")) + params.append(pytest.param(case, id=case.name, marks=marks)) + return params + + +# --------------------------------------------------------------------------- +def _middle_frame(frames: list[Image.Image] | list[np.ndarray] | np.ndarray) -> Image.Image: + if isinstance(frames, np.ndarray): + idx = frames.shape[0] // 2 + arr = frames[idx] + if arr.dtype != np.uint8: + arr = (arr * 255).clip(0, 255).astype(np.uint8) + return Image.fromarray(arr).convert("RGB") + if not frames: + raise AssertionError("Empty video output") + mid = frames[len(frames) // 2] + if isinstance(mid, np.ndarray): + if mid.dtype != np.uint8: + mid = (mid * 255).clip(0, 255).astype(np.uint8) + return Image.fromarray(mid).convert("RGB") + return mid.convert("RGB") + + +def _extract_frames(result) -> list[Image.Image]: + """Pull a frame list out of an Omni text-to-video result.""" + from vllm_omni.entrypoints.omni import OmniRequestOutput + + if isinstance(result, list): + result = result[0] if result else None + if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result.request_output is not None: + result = result.request_output + images = getattr(result, "images", None) + if not images: + raise AssertionError("Pipeline output had no images") + if isinstance(images, list) and images and isinstance(images[0], tuple): + # (frames, fps) + images = images[0][0] + return list(images) + + +def _generate_video( + *, + model: str, + spec: _Spec, + quantization: str | None, + output_path: Path, +) -> tuple[Image.Image, float]: + """Generate one video and return its middle frame + elapsed wall time.""" + from vllm_omni.diffusion.data import DiffusionParallelConfig + from vllm_omni.entrypoints.omni import Omni + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + from vllm_omni.platforms import current_omni_platform + + kwargs: dict[str, Any] = dict( + model=model, + parallel_config=DiffusionParallelConfig(), + enforce_eager=True, + vae_use_tiling=True, + ) + if quantization: + kwargs["quantization"] = quantization + if spec.boundary_ratio is not None: + kwargs["boundary_ratio"] = spec.boundary_ratio + omni = Omni(**kwargs) + + try: + gen = torch.Generator(device=current_omni_platform.device_type).manual_seed(SEED) + sampling = dict( + height=HEIGHT, + width=WIDTH, + generator=gen, + guidance_scale=spec.guidance_scale, + num_inference_steps=NUM_INFERENCE_STEPS, + num_frames=NUM_FRAMES, + seed=SEED, + ) + if spec.guidance_scale_2 is not None: + sampling["guidance_scale_2"] = spec.guidance_scale_2 + + import time as _time + + t0 = _time.perf_counter() + out = omni.generate({"prompt": spec.prompt}, OmniDiffusionSamplingParams(**sampling)) + elapsed = _time.perf_counter() - t0 + + frames = _extract_frames(out) + mid = _middle_frame(frames) + mid.save(output_path) + return mid, elapsed + finally: + del omni + gc.collect() + if torch.cuda.is_available(): + torch.accelerator.empty_cache() + + +@pytest.mark.parametrize("case", _cases()) +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.accelerator.device_count() < 1, + reason="Needs at least 1 GPU.", +) +def test_quantized_videogen_dit_matches_bf16_clip(case: _QuantCase, accuracy_artifact_root: Path) -> None: + """FP8 DiT should preserve prompt-aligned video quality vs BF16.""" + out_dir = model_output_dir(accuracy_artifact_root, case.spec.family + "-quant") + + with tempfile.TemporaryDirectory(): + bf16_model = os.environ[case.spec.bf16_env] + quant_model = os.environ[case.spec.fp8_env] + + bf16_frame, bf16_time = _generate_video( + model=bf16_model, + spec=case.spec, + quantization=None, + output_path=out_dir / "bf16.png", + ) + quant_frame, quant_time = _generate_video( + model=quant_model, + spec=case.spec, + quantization="fp8", + output_path=out_dir / f"{case.name}.png", + ) + + clip = CLIPScorer() + bf16_clip = clip.score(bf16_frame, case.spec.prompt) + quant_clip = clip.score(quant_frame, case.spec.prompt) + clip_drop = bf16_clip - quant_clip + speedup = (bf16_time / quant_time) if quant_time else float("nan") + + metrics = { + "case": case.name, + "family": case.spec.family, + "quantization": "fp8", + "bf16_model": bf16_model, + "quant_model": quant_model, + "prompt": case.spec.prompt, + "height": HEIGHT, + "width": WIDTH, + "num_frames": NUM_FRAMES, + "num_inference_steps": NUM_INFERENCE_STEPS, + "bf16_elapsed_s": bf16_time, + "quant_elapsed_s": quant_time, + "speedup_vs_bf16": speedup, + "bf16_clip_score": bf16_clip, + "quant_clip_score": quant_clip, + "clip_score_drop": clip_drop, + } + metrics_path = out_dir / f"{case.name}_metrics.json" + metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8") + + print(f"\n[{case.name}] BF16 {bf16_time:.2f}s | FP8 {quant_time:.2f}s (x{speedup:.2f})") + print(f" CLIP BF16={bf16_clip:.2f} FP8={quant_clip:.2f} drop={clip_drop:.2f}") + print(f" metrics={metrics_path}") + + assert quant_clip >= CLIP_ABSOLUTE_FLOOR, ( + f"{case.name} CLIP below floor: got {quant_clip:.2f}, expected >= {CLIP_ABSOLUTE_FLOOR:.2f}" + ) + assert clip_drop <= case.clip_drop, ( + f"{case.name} CLIP drop too large: got {clip_drop:.2f}, expected <= {case.clip_drop:.2f}" + ) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index c9f81b64d..e83604fc5 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -433,8 +433,13 @@ class DiffusersPipelineLoader: # that have loaded weights tracking currently. if loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights - # NOTE: if the model is quantized, ignore not_loaded check for scale weights - weights_scale_not_loaded = {name for name in weights_not_loaded if name.endswith("weight_scale")} + # NOTE: if the model is quantized, ignore not_loaded check for scale + # weights. ModelOpt FP8 carries a per-tensor `weight_scale` and a + # static activation `input_scale`, which the quant method may + # fold/track differently than plain parameters. + weights_scale_not_loaded = { + name for name in weights_not_loaded if name.endswith(("weight_scale", "input_scale")) + } weights_not_loaded = weights_not_loaded - weights_scale_not_loaded if weights_not_loaded: self._check_unloaded_weights(weights_not_loaded) diff --git a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py index 1f06799c3..9af9e8488 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any import torch import torch.nn.functional as F @@ -31,6 +31,9 @@ from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + logger = init_logger(__name__) @@ -332,6 +335,8 @@ class HunyuanVideo15Attention(nn.Module): out_bias: bool = True, eps: float = 1e-6, out_dim: int | None = None, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() @@ -350,6 +355,8 @@ class HunyuanVideo15Attention(nn.Module): head_size=self.head_dim, total_num_heads=self.heads, bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.to_qkv", ) self.to_out = nn.ModuleList( @@ -360,6 +367,8 @@ class HunyuanVideo15Attention(nn.Module): bias=out_bias, input_is_parallel=True, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_out.0", ), nn.Identity(), # placeholder for dropout (none used) ] @@ -374,6 +383,8 @@ class HunyuanVideo15Attention(nn.Module): head_size=self.head_dim, total_num_heads=self.heads, bias=added_proj_bias, + quant_config=quant_config, + prefix=f"{prefix}.add_kv_proj", ) self.to_add_out = RowParallelLinear( @@ -382,6 +393,8 @@ class HunyuanVideo15Attention(nn.Module): bias=out_bias, input_is_parallel=True, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_add_out", ) self.rope = RotaryEmbedding(is_neox_style=False) @@ -401,6 +414,8 @@ class HunyuanVideo15Attention(nn.Module): image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, hidden_states_mask: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # Ensure contiguous for FP8 quantized linear layers + hidden_states = hidden_states.contiguous() qkv, _ = self.to_qkv(hidden_states) q_size = self.to_qkv.num_heads * self.head_dim kv_size = self.to_qkv.num_kv_heads * self.head_dim @@ -421,6 +436,7 @@ class HunyuanVideo15Attention(nn.Module): key = self.rope(key, cos, sin) if encoder_hidden_states is not None: + encoder_hidden_states = encoder_hidden_states.contiguous() encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states) add_q_size = self.add_kv_proj.num_heads * self.head_dim add_kv_size = self.add_kv_proj.num_kv_heads * self.head_dim @@ -492,6 +508,8 @@ class HunyuanVideo15TransformerBlock(nn.Module): attention_head_dim: int, mlp_ratio: float = 4.0, qk_norm: str = "rms_norm", + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ) -> None: super().__init__() hidden_size = num_attention_heads * attention_head_dim @@ -507,13 +525,23 @@ class HunyuanVideo15TransformerBlock(nn.Module): out_dim=hidden_size, bias=True, eps=1e-6, + quant_config=quant_config, + prefix=f"{prefix}.attn", ) self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.ff = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio) + self.ff = FeedForward( + dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio, quant_config=quant_config, prefix=f"{prefix}.ff" + ) self.norm2_context = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.ff_context = FeedForward(dim=hidden_size, dim_out=hidden_size, mult=mlp_ratio) + self.ff_context = FeedForward( + dim=hidden_size, + dim_out=hidden_size, + mult=mlp_ratio, + quant_config=quant_config, + prefix=f"{prefix}.ff_context", + ) def forward( self, @@ -601,6 +629,7 @@ class HunyuanVideo15Transformer3DModel(nn.Module): target_size: int = 640, task_type: str = "i2v", use_meanflow: bool = False, + quant_config: "QuantizationConfig | None" = None, ): super().__init__() @@ -632,9 +661,14 @@ class HunyuanVideo15Transformer3DModel(nn.Module): self.transformer_blocks = nn.ModuleList( [ HunyuanVideo15TransformerBlock( - num_attention_heads, attention_head_dim, mlp_ratio=mlp_ratio, qk_norm=qk_norm + num_attention_heads, + attention_head_dim, + mlp_ratio=mlp_ratio, + qk_norm=qk_norm, + quant_config=quant_config, + prefix=f"transformer_blocks.{i}", ) - for _ in range(num_layers) + for i in range(num_layers) ] ) diff --git a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py index 038898a16..0684f6749 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py +++ b/vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py @@ -143,7 +143,9 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, Diff self.scheduler._shift = od_config.flow_shift transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HunyuanVideo15Transformer3DModel) - self.transformer = HunyuanVideo15Transformer3DModel(od_config=od_config, **transformer_kwargs) + self.transformer = HunyuanVideo15Transformer3DModel( + od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs + ) # Check if model uses meanflow (distilled variants) self.use_meanflow = getattr(od_config.tf_model_config, "use_meanflow", False) diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml new file mode 100644 index 000000000..097afc020 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml @@ -0,0 +1,38 @@ +# Stage config for running HunyuanVideo-1.5 DiT with ModelOpt FP8 auto-detect. +# Single GPU. Bump `tensor_parallel_size` and `devices` for multi-GPU TP. +# +# Use with a ModelOpt FP8 checkpoint (e.g. produced by +# examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py). + +stage_args: + - stage_id: 0 + stage_type: diffusion + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: dit + model_class_name: HunyuanVideo15Pipeline + max_num_seqs: 1 + enforce_eager: true + trust_remote_code: true + distributed_executor_backend: "mp" + # HunyuanVideo-1.5 peak VRAM is dominated by the single-shot VAE decode of + # the full video (the DiT contributes little), so tiling the decode is the + # main memory lever — far more than FP8 (e.g. 480x832x33: ~117 -> ~47 GiB). + # Default it on here; override with --vae-use-tiling/false if you have headroom. + vae_use_tiling: true + parallel_config: + tensor_parallel_size: 1 + + final_output: true + final_output_type: video + is_comprehension: false + default_sampling_params: + seed: 42 + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 diff --git a/vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml b/vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml new file mode 100644 index 000000000..1b1c40537 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml @@ -0,0 +1,41 @@ +# Stage config for running Wan2.2 TI2V-5B DiT with ModelOpt FP8 auto-detect. +# Single GPU (TI2V-5B fits 80GB BF16; FP8 drops by ~half). +# For the A14B MoE variants, bump `tensor_parallel_size` and `devices`. +# +# Use with a ModelOpt FP8 checkpoint (e.g. produced by +# examples/quantization/quantize_wan2_2_modelopt_fp8.py). + +stage_args: + - stage_id: 0 + stage_type: diffusion + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: dit + # Wan2.2 TI2V-5B is served by the base Wan22Pipeline, which auto-detects + # expand_timesteps (TI2V-5B mode) from model_index.json. (The standalone + # Wan22TI2VPipeline was consolidated into Wan22Pipeline upstream.) + model_class_name: Wan22Pipeline + max_num_seqs: 1 + enforce_eager: true + trust_remote_code: true + distributed_executor_backend: "mp" + # Tile the VAE decode by default. Wan2.2's 480p peak is DiT-dominated (so + # FP8 is the main lever here), but tiling adds headroom for higher + # resolutions / longer clips at negligible cost. Override with --vae-use-tiling/false. + vae_use_tiling: true + parallel_config: + tensor_parallel_size: 1 + + final_output: true + final_output_type: video + is_comprehension: false + default_sampling_params: + seed: 42 + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 29cc56e0f..c343276b1 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -217,3 +217,43 @@ def _patch_scaled_mm_fp8_contiguous_activation(): _patch_scaled_mm_fp8_contiguous_activation() + + +def _patch_flashinfer_fp8_scaled_mm_output_shape(): + """Restore the N-D output shape for the FlashInfer FP8 ScaledMM kernel. + + ``FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm`` returns the raw 2-D + GEMM result and ignores ``output_shape``, unlike the CUTLASS / PyTorch + ScaledMM kernels which reshape to it. A 3-D activation ``(B, S, D)`` thus + collapses to ``(B*S, D)``, breaking diffusion DiTs that reshape the linear + output by absolute dim (e.g. Wan2.2 ``qkv.unflatten(2, ...)``) with + ``IndexError: Dimension out of range``. It only bites >2-D inputs, so LLM + (token-flattened, 2-D) paths are unaffected. + + Carried here because the upstream fix may not have landed yet; this override + becomes a harmless no-op once vLLM honors ``output_shape`` itself. + """ + try: + from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import ( + FlashInferFP8ScaledMMLinearKernel, + ) + except ImportError: + return + + _original_apply_scaled_mm = FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm + if getattr(_original_apply_scaled_mm, "_omni_output_shape_patched", False): + return + + def _patched_apply_scaled_mm(self, *, A, B, out_dtype, As, Bs, bias, output_shape): # noqa: N803 + out = _original_apply_scaled_mm( + self, A=A, B=B, out_dtype=out_dtype, As=As, Bs=Bs, bias=bias, output_shape=output_shape + ) + if tuple(out.shape) != tuple(output_shape): + out = out.view(*output_shape) + return out + + _patched_apply_scaled_mm._omni_output_shape_patched = True + FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm = _patched_apply_scaled_mm + + +_patch_flashinfer_fp8_scaled_mm_output_shape()