diff --git a/docs/.nav.yml b/docs/.nav.yml index 94de4dc94..0228c7139 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -78,6 +78,7 @@ nav: - Online Quantization: user_guide/quantization/online.md - FP8 W8A8: user_guide/quantization/fp8.md - Int8 W8A8: user_guide/quantization/int8.md + - ModelOpt: user_guide/quantization/modelopt.md - GGUF: user_guide/quantization/gguf.md - AutoRound: user_guide/quantization/autoround.md - msModelSlim: user_guide/quantization/msmodelslim.md diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index 7373a39ff..c8f14548c 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -2,11 +2,11 @@ ## Overview -FP8 quantization converts BF16/FP16 weights to FP8 at model load time, or loads -a checkpoint whose target stage already declares an FP8 quantization config. -Online activation scaling is the default and does not require calibration. -Static activation scaling is supported when calibrated scale information is -available. +FP8 quantization converts BF16/FP16 weights to FP8 at model load time. Online +activation scaling is the default and does not require calibration. Static +activation scaling is supported when calibrated scale information is available. +For ModelOpt-produced pre-quantized checkpoints, see +[ModelOpt Quantization](modelopt.md). Some architectures can quantize all linear layers. Others have quality-sensitive layers that should stay in BF16 through `ignored_layers`. @@ -46,7 +46,7 @@ guide. FP8 on Ampere may use a weight-only path where available. | Model | Scope | Format | Status | |-------|-------|--------|--------| -| Qwen3-Omni | Thinker language-model stage | ModelOpt `quant_algo=FP8` | Tested for thinker memory reduction | +| Qwen3-Omni | Thinker language-model stage | [ModelOpt](modelopt.md) `quant_algo=FP8` | Tested for thinker memory reduction | | Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated | Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a diff --git a/docs/user_guide/quantization/modelopt.md b/docs/user_guide/quantization/modelopt.md new file mode 100644 index 000000000..8a4389015 --- /dev/null +++ b/docs/user_guide/quantization/modelopt.md @@ -0,0 +1,138 @@ +# ModelOpt Quantization + +## Overview + +ModelOpt quantization loads checkpoints produced by NVIDIA ModelOpt. The +quantized weights and scale tensors are generated before serving, so inference +does not run online calibration or convert a BF16 checkpoint at startup. + +vLLM-Omni currently validates the ModelOpt FP8 checkpoint path for diffusion +transformers. The loader auto-detects supported ModelOpt FP8 checkpoint configs +and keeps non-transformer components, such as the tokenizer, scheduler, text +encoder, and VAE, on the base checkpoint unless a model-specific guide says +otherwise. + +!!! note + `--force-cutlass-fp8` is an explicit runtime override for diffusion + checkpoints that already carry a supported ModelOpt FP8 config. It does not + quantize BF16 checkpoints and it does not apply to online `--quantization + fp8`. The flag only takes effect for ModelOpt FP8 diffusion stages on CUDA + SM89+ devices; other platforms and non-ModelOpt FP8 paths fall back to the + normal vLLM kernel selection. + +## Supported ModelOpt Checkpoint Formats + +vLLM-Omni treats ModelOpt checkpoints as pre-quantized checkpoints. The +checkpoint config must identify ModelOpt as the quantization method or producer, +and the quantization algorithm must be one of the validated FP8 algorithms. + +| Checkpoint field | Supported value | +|------------------|-----------------| +| `method` / `quant_method` | `modelopt` | +| `producer.name` | `modelopt` | +| `quant_algo` | `FP8`, `FP8_PER_CHANNEL_PER_TOKEN` | + +Other ModelOpt algorithms, such as NVFP4, are not enabled by this diffusion +FP8 path until they have separate model and quality validation. + +## Hardware Support + +| Device | Support | +|--------|---------| +| NVIDIA Blackwell GPU (SM 100+) | ✅ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | +| NVIDIA Ampere GPU (SM 80+) | ⭕ | +| AMD ROCm | ⭕ | +| Intel XPU | ⭕ | +| Ascend NPU | ❌ | + +Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. +The optional CUTLASS FP8 runtime override requires CUDA SM89+. + +## Model Type Support + +### Diffusion Model + +| Model | HF checkpoint | Scope | Status | +|-------|---------------|-------|--------| +| Qwen-Image 2512 | `feizhai123/qwen-image-2512-modelopt-fp8-dynamic-all` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints | +| Z-Image | `feizhai123/z-image-modelopt-fp8-conservative` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints | +| FLUX.2-dev | `feizhai123/flux2-dev-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints | +| FLUX.2-klein 4B | `feizhai123/flux2-klein-4b-modelopt-fp8` | Diffusion transformer | Validated for ModelOpt FP8 checkpoints | +| HunyuanImage-3.0 | `feizhai123/hunyuan-image3-modelopt-fp8` | MoE diffusion transformer | Validated for ModelOpt FP8 checkpoints | +| Wan2.2 | Not available | Diffusion transformer | Not validated | + +### Multi-Stage Omni/TTS Model + +| Model | Scope | Status | +|-------|-------|--------| +| Qwen3-Omni | Thinker language-model stage | ModelOpt FP8 checkpoint path | +| Qwen3-TTS | TTS language-model stage | Not validated | + +Audio encoder, vision encoder, talker, and code2wav stages stay in BF16 unless +a model-specific guide documents otherwise. + +### Multi-Stage Diffusion Model + +ModelOpt checkpoints must be routed to the stage whose checkpoint contains the +ModelOpt `quantization_config`. BAGEL and GLM-Image are not listed as validated +ModelOpt targets yet. + +## Configuration + +For pre-quantized ModelOpt FP8 checkpoints, no `--quantization fp8` flag is +needed. The checkpoint config selects the ModelOpt path. + +Online serving: + +```bash +vllm serve \ + --omni \ + --tensor-parallel-size \ + --force-cutlass-fp8 +``` + +Offline inference: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model \ + --tensor-parallel-size \ + --prompt "a red ceramic teapot on a wooden table" \ + --height 1024 \ + --width 1024 \ + --num-inference-steps 20 \ + --seed 42 \ + --output outputs/modelopt_fp8.png +``` + +Python API: + +```python +from vllm_omni import Omni + +omni = Omni( + model="", + tensor_parallel_size=2, + force_cutlass_fp8=True, +) +``` + +## Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `force_cutlass_fp8` / `--force-cutlass-fp8` | bool | `False` | Force CUTLASS FP8 linear kernels for supported ModelOpt FP8 diffusion stages on CUDA SM89+ | + +## Validation and Notes + +1. Compare the ModelOpt FP8 checkpoint against the BF16 baseline with the same + prompt, resolution, seed, and inference steps. +2. Use `tests/diffusion/quantization/test_quantization_quality.py` with + `VLLM_OMNI_QUALITY_CONFIGS` to validate local baseline and quantized model + paths. +3. Report LPIPS, PSNR, MAE, throughput, latency, and peak memory when adding a + new validated ModelOpt diffusion checkpoint. +4. Keep `--quantization fp8` for online FP8 from BF16 checkpoints; use this + ModelOpt path only when the checkpoint already contains ModelOpt FP8 weights + and scales. diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index 762fb2fc6..ed7a1b8cf 100644 --- a/docs/user_guide/quantization/overview.md +++ b/docs/user_guide/quantization/overview.md @@ -10,18 +10,18 @@ type has a different quantization scope. | Mode | Guide | Description | Methods | |------|-------|-------------|---------| | Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8 | -| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | GGUF, AutoRound, msModelSlim, serialized Int8 | +| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8 | ## Hardware Support -| Device | FP8 W8A8 | Int8 W8A8 | GGUF | AutoRound | msModelSlim | -|--------|----------|-----------|------|-----------|-------------| -| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ❌ | -| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ❌ | -| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ✅ | ❌ | -| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ❌ | -| Intel XPU | ⭕ | ⭕ | ⭕ | ✅ | ❌ | -| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ | +| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | GGUF | AutoRound | msModelSlim | +|--------|----------|-----------|----------|------|-----------|-------------| +| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | +| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ✅ | ✅ | ❌ | +| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ | +| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | +| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide. FP8 on Ampere may use a weight-only path where available. @@ -39,6 +39,7 @@ otherwise. |--------|-------|------|----------------|--------| | FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models | | Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image | +| ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints | | GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists | | AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven | | msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path | @@ -52,7 +53,7 @@ in BF16 unless the model guide explicitly adds support. | Method | Guide | Scope | Example models | Status | |--------|-------|-------|----------------|--------| -| FP8 | [FP8](fp8.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path | +| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path | | Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated | | AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints | @@ -67,6 +68,7 @@ attached to the intended stage rather than applied globally. |--------|-------|-------|----------------|--------| | FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation | | Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation | +| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation | | GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed | | AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed | | msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation | @@ -94,7 +96,7 @@ config = build_quant_config({ | Component | Default quantized? | Notes | |-----------|--------------------|-------| -| Diffusion transformer | Yes | Primary target for FP8, Int8, GGUF, AutoRound, and msModelSlim | +| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, GGUF, AutoRound, and msModelSlim | | Text encoder | No | Keep BF16 unless a method-specific guide documents support | | VAE | No | Keep BF16; storage-only paths are method-specific | | Scheduler/tokenizer | No | Loaded from the base model repository | diff --git a/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py b/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py new file mode 100644 index 000000000..513e80d28 --- /dev/null +++ b/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( + ModelOptFp8CheckpointAdapter, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _PackedModelOptModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer = nn.Module() + self.transformer.block = nn.Module() + self.transformer.block.to_qkv = nn.Linear(2, 2, bias=False) + + +class _QuantizedPackedModelOptModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer = nn.Module() + self.transformer.block = nn.Module() + self.transformer.block.to_qkv = nn.Module() + self.transformer.block.to_qkv.register_parameter( + "weight", + nn.Parameter(torch.empty(2, 2, dtype=torch.float8_e4m3fn), requires_grad=False), + ) + self.transformer.block.to_qkv.register_parameter( + "weight_scale", + nn.Parameter(torch.empty(1), requires_grad=False), + ) + self.transformer.block.to_qkv.register_parameter( + "input_scale", + nn.Parameter(torch.empty(1), requires_grad=False), + ) + + +def _make_source() -> SimpleNamespace: + return SimpleNamespace( + subfolder="transformer", + prefix="transformer.", + ) + + +def test_modelopt_adapter_dequantizes_fp8_weight_for_full_precision_target(): + model = _PackedModelOptModel() + adapter = ModelOptFp8CheckpointAdapter(model, _make_source()) + fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn) + scale = torch.tensor([0.5], dtype=torch.float32) + + adapted = list( + adapter.adapt( + iter( + [ + ("transformer.block.to_q.weight_scale", scale), + ("transformer.block.to_q.input_scale", torch.tensor([1.0])), + ("transformer.block.to_q.weight", fp8_weight), + ] + ) + ) + ) + + assert [name for name, _ in adapted] == ["transformer.block.to_q.weight"] + assert adapted[0][1].dtype == model.transformer.block.to_qkv.weight.dtype + assert torch.allclose(adapted[0][1], fp8_weight.to(torch.float32) * scale) + + +def test_modelopt_adapter_keeps_scale_tensors_for_quantized_target(): + model = _QuantizedPackedModelOptModel() + adapter = ModelOptFp8CheckpointAdapter(model, _make_source()) + scale = torch.tensor([0.5], dtype=torch.float32) + + adapted = list( + adapter.adapt( + iter( + [ + ("transformer.block.to_q.weight_scale", scale), + ("transformer.block.to_q.input_scale", torch.tensor([1.0])), + ] + ) + ) + ) + + assert [name for name, _ in adapted] == [ + "transformer.block.to_q.weight_scale", + "transformer.block.to_q.input_scale", + ] diff --git a/tests/diffusion/quantization/test_fp8_config.py b/tests/diffusion/quantization/test_fp8_config.py index 574af7a66..a1f79030e 100644 --- a/tests/diffusion/quantization/test_fp8_config.py +++ b/tests/diffusion/quantization/test_fp8_config.py @@ -54,6 +54,25 @@ def test_build_quant_config_dict_not_mutated(): assert original == copy +def test_build_quant_config_modelopt_fp8_config_json(): + from vllm.model_executor.layers.quantization.modelopt import ModelOptFp8Config + + from vllm_omni.quantization import build_quant_config + + config = build_quant_config( + { + "quant_method": "modelopt", + "quant_algo": "FP8", + "ignore": ["proj_out"], + "producer": {"name": "modelopt"}, + } + ) + + assert isinstance(config, ModelOptFp8Config) + assert config.get_name() == "modelopt" + assert config.is_checkpoint_fp8_serialized + + def test_build_quant_config_per_component(): from vllm_omni.quantization import ComponentQuantizationConfig, build_quant_config @@ -91,7 +110,7 @@ def test_flat_dict_not_misdetected_as_per_component(): as a per-component dict — it should raise ValueError for missing 'method'.""" from vllm_omni.quantization import build_quant_config - with pytest.raises(ValueError, match="must have a 'method' key"): + with pytest.raises(ValueError, match="must have a 'method' or 'quant_method' key"): build_quant_config({"activation_scheme": "static"}) @@ -194,10 +213,30 @@ def test_integration_per_component(): assert config.quantization_config.component_configs["vae"] is None +def test_transformer_config_auto_detects_modelopt_fp8(): + from vllm.model_executor.layers.quantization.modelopt import ModelOptFp8Config + + from vllm_omni.diffusion.data import TransformerConfig + + config = TransformerConfig.from_dict( + { + "_class_name": "FluxTransformer2DModel", + "quantization_config": { + "quant_method": "modelopt", + "quant_algo": "FP8", + "ignore": ["proj_out"], + }, + } + ) + + assert isinstance(config.quant_config, ModelOptFp8Config) + assert config.quant_method == "modelopt" + + def test_supported_methods_includes_vllm(): from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS - for method in ["fp8", "gguf", "awq", "gptq", "bitsandbytes", "modelopt", "modelopt_fp4"]: + for method in ["fp8", "gguf", "awq", "gptq", "bitsandbytes", "modelopt"]: assert method in SUPPORTED_QUANTIZATION_METHODS, f"{method} missing" @@ -243,14 +282,10 @@ def test_per_component_routing_with_default(): assert resolved.get_name() == "fp8" -@pytest.mark.parametrize( - "quant_algo", - ["FP8", "NVFP4"], - ids=["modelopt_fp8", "modelopt_nvfp4"], -) +@pytest.mark.parametrize("quant_algo", ["FP8"], ids=["modelopt_fp8"]) def test_omni_convertor_thinker_finds_text_config_quant(quant_algo): """Thinker stage should discover quantization_config from - thinker_config.text_config for both FP8 and NVFP4 modelopt checkpoints.""" + thinker_config.text_config for verified modelopt FP8 checkpoints.""" from types import SimpleNamespace from vllm_omni.config.model import OmniModelArchConfigConvertor diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index 05dc868d0..6288bde5e 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Quantization quality gate — validates that quantized models stay within -acceptable perceptual distance (LPIPS) of the BF16 baseline. +Quantization quality gate for diffusion models. Developers adding a new quantization method should: 1. Add their method + model to QUALITY_CONFIGS below @@ -11,7 +10,7 @@ Developers adding a new quantization method should: 4. Paste the output table into their PR description The test generates outputs with both BF16 and the quantized method using the -same seed, computes LPIPS, and fails if it exceeds the threshold. +same seed, computes similarity metrics, and fails if LPIPS exceeds the threshold. Requirements: pip install lpips @@ -21,12 +20,22 @@ Example — run only FP8 tests: Example — run a specific model: pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image" + +Example — validate a local BF16 baseline against a local pre-quantized checkpoint: + export VLLM_OMNI_QUALITY_CONFIGS=/tmp/modelopt_quality_cases.json + pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "qwen_image_2512" + +Optional artifact dump: + export VLLM_OMNI_QUALITY_OUTPUT_DIR=/tmp/modelopt_quality_outputs """ from __future__ import annotations import gc +import json +import os from dataclasses import dataclass +from pathlib import Path import numpy as np import pytest @@ -44,17 +53,49 @@ class QualityTestConfig: """Defines a single quantization quality test case.""" id: str # pytest ID, e.g. "fp8_z_image" - model: str # HF model name - quantization: str # quantization method, e.g. "fp8" task: str # "t2i" or "t2v" prompt: str # generation prompt max_lpips: float # fail threshold — higher = more lenient + model: str | None = None # HF model name + quantization: str | None = None # quantization method, e.g. "fp8" + baseline_model: str | None = None # explicit BF16/local baseline path + quantized_model: str | None = None # explicit quantized/local model path height: int = 1024 width: int = 1024 num_inference_steps: int = 20 # keep low for CI speed num_frames: int = 5 # only for t2v seed: int = 42 gpu: str = "H100" # minimum GPU requirement + negative_prompt: str = "" + guidance_scale: float | None = None + + def baseline_ref(self) -> str: + return self.baseline_model or self.model or "" + + def quantized_ref(self) -> str: + if self.quantized_model is not None: + return self.quantized_model + return self.model or "" + + def quantization_ref(self) -> str | None: + if self.quantized_model is not None: + return None + return self.quantization + + def validate(self) -> None: + uses_explicit_models = self.baseline_model is not None or self.quantized_model is not None + uses_model_plus_method = self.model is not None or self.quantization is not None + + if uses_explicit_models and uses_model_plus_method: + raise ValueError(f"{self.id}: explicit baseline/quantized paths cannot be mixed with model/quantization") + + if uses_explicit_models: + if self.baseline_model is None or self.quantized_model is None: + raise ValueError(f"{self.id}: baseline_model and quantized_model must be provided together") + return + + if self.model is None or self.quantization is None: + raise ValueError(f"{self.id}: expected either model+quantization or baseline_model+quantized_model") # Add new quantization methods / models here. @@ -91,6 +132,53 @@ QUALITY_CONFIGS = [ ] +def _load_extra_quality_configs() -> list[QualityTestConfig]: + config_path = os.environ.get("VLLM_OMNI_QUALITY_CONFIGS") + if not config_path: + return [] + + payload = json.loads(Path(config_path).read_text(encoding="utf-8")) + if not isinstance(payload, list): + raise ValueError("VLLM_OMNI_QUALITY_CONFIGS must point to a JSON list") + + configs: list[QualityTestConfig] = [] + for item in payload: + if not isinstance(item, dict): + raise ValueError("Each quality config entry must be a JSON object") + config = QualityTestConfig(**item) + config.validate() + configs.append(config) + return configs + + +def _all_quality_configs() -> list[QualityTestConfig]: + configs = [*QUALITY_CONFIGS, *_load_extra_quality_configs()] + for config in configs: + config.validate() + return configs + + +def _output_path(output_dir: Path, config: QualityTestConfig, label: str, suffix: str) -> Path: + safe_id = config.id.replace("/", "_") + return output_dir / f"{safe_id}_{label}{suffix}" + + +def _maybe_save_output(output_dir: Path | None, config: QualityTestConfig, label: str, output) -> None: + if output_dir is None: + return + + output_dir.mkdir(parents=True, exist_ok=True) + if config.task == "t2i": + output.save(_output_path(output_dir, config, label, ".png")) + return + + if isinstance(output, np.ndarray): + np.save(_output_path(output_dir, config, label, ".npy"), output) + return + + raise TypeError(f"Unsupported output type for saving: {type(output)!r}") + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -107,12 +195,13 @@ def _generate_image(omni, config: QualityTestConfig): torch.accelerator.reset_peak_memory_stats() outputs = omni.generate( - {"prompt": config.prompt}, + {"prompt": config.prompt, "negative_prompt": config.negative_prompt}, OmniDiffusionSamplingParams( height=config.height, width=config.width, generator=generator, num_inference_steps=config.num_inference_steps, + guidance_scale=config.guidance_scale, ), ) @@ -138,13 +227,14 @@ def _generate_video(omni, config: QualityTestConfig): torch.accelerator.reset_peak_memory_stats() outputs = omni.generate( - {"prompt": config.prompt, "negative_prompt": ""}, + {"prompt": config.prompt, "negative_prompt": config.negative_prompt}, OmniDiffusionSamplingParams( height=config.height, width=config.width, generator=generator, num_inference_steps=config.num_inference_steps, num_frames=config.num_frames, + guidance_scale=config.guidance_scale, ), ) @@ -186,6 +276,34 @@ def _compute_lpips(baseline, quantized, task: str) -> float: return compute_lpips_video(baseline, quantized) +def _to_float_array(output, task: str) -> np.ndarray: + if task == "t2i": + array = np.asarray(output.convert("RGB"), dtype=np.float32) / 255.0 + else: + array = np.asarray(output, dtype=np.float32) + if array.max() > 1.0 or array.min() < 0.0: + array = np.clip(array, 0.0, 255.0) / 255.0 + else: + array = np.clip(array, 0.0, 1.0) + return array + + +def _compute_psnr_and_mae(baseline, quantized, task: str) -> tuple[float, float]: + baseline_array = _to_float_array(baseline, task) + quantized_array = _to_float_array(quantized, task) + if baseline_array.shape != quantized_array.shape: + raise ValueError( + "Output shapes do not match for metric computation: " + f"baseline={baseline_array.shape}, quantized={quantized_array.shape}" + ) + + diff = baseline_array - quantized_array + mae = float(np.mean(np.abs(diff))) + mse = float(np.mean(np.square(diff))) + psnr = float("inf") if mse == 0.0 else float(20.0 * np.log10(1.0 / np.sqrt(mse))) + return psnr, mae + + def _unload(omni): del omni gc.collect() @@ -199,6 +317,7 @@ def _unload(omni): # --------------------------------------------------------------------------- _marks = hardware_marks(res={"cuda": "H100"}) +_OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUALITY_OUTPUT_DIR" in os.environ else None def _quality_param(c: QualityTestConfig): @@ -214,7 +333,7 @@ def _quality_param(c: QualityTestConfig): @pytest.mark.diffusion @pytest.mark.parametrize( "config", - [_quality_param(c) for c in QUALITY_CONFIGS], + [_quality_param(c) for c in _all_quality_configs()], ) def test_quantization_quality(config: QualityTestConfig): """Validate that quantized output stays within LPIPS threshold of BF16.""" @@ -223,31 +342,44 @@ def test_quantization_quality(config: QualityTestConfig): generate_fn = _generate_video if config.task == "t2v" else _generate_image # --- BF16 baseline --- - omni_bl = Omni(model=config.model) + omni_bl = Omni(model=config.baseline_ref()) baseline_out, bl_mem = generate_fn(omni_bl, config) _unload(omni_bl) + _maybe_save_output(_OUTPUT_DIR, config, "baseline", baseline_out) # --- Quantized --- - omni_qt = Omni(model=config.model, quantization_config=config.quantization) + quantization = config.quantization_ref() + if quantization is None: + omni_qt = Omni(model=config.quantized_ref()) + else: + omni_qt = Omni(model=config.quantized_ref(), quantization_config=quantization) quant_out, qt_mem = generate_fn(omni_qt, config) _unload(omni_qt) + _maybe_save_output(_OUTPUT_DIR, config, "quantized", quant_out) - # --- LPIPS --- + # --- Similarity metrics --- lpips_score = _compute_lpips(baseline_out, quant_out, config.task) + psnr_score, mae_score = _compute_psnr_and_mae(baseline_out, quant_out, config.task) + assert lpips_score <= config.max_lpips, ( + f"LPIPS {lpips_score:.4f} exceeds threshold {config.max_lpips} " + f"for {config.quantization_ref() or 'pre-quantized checkpoint'} on {config.quantized_ref()}" + ) # --- Report --- mem_reduction = (bl_mem - qt_mem) / bl_mem * 100 if bl_mem > 0 else 0 print(f"\n{'=' * 60}") print(f"Quantization Quality: {config.id}") print(f"{'=' * 60}") - print(f" Model: {config.model}") - print(f" Method: {config.quantization}") + print(f" Baseline: {config.baseline_ref()}") + print(f" Quantized: {config.quantized_ref()}") + print(f" Method: {config.quantization_ref() or 'pre-quantized checkpoint'}") print(f" LPIPS: {lpips_score:.4f} (threshold: {config.max_lpips})") + print(f" PSNR: {psnr_score:.4f} dB (higher is better)") + print(f" MAE: {mae_score:.6f} (lower is better)") print(f" BF16 memory: {bl_mem:.2f} GiB") print(f" Quant memory: {qt_mem:.2f} GiB ({mem_reduction:.0f}% reduction)") print(f" Result: {'PASS' if lpips_score <= config.max_lpips else 'FAIL'}") print(f"{'=' * 60}\n") - assert lpips_score <= config.max_lpips, ( - f"LPIPS {lpips_score:.4f} exceeds threshold {config.max_lpips} for {config.quantization} on {config.model}" - ) + assert np.isfinite(psnr_score) or np.isinf(psnr_score), f"PSNR is invalid for {config.id}: {psnr_score}" + assert np.isfinite(mae_score), f"MAE is not finite for {config.id}: {mae_score}" diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index fc08c5f7f..88b4d8c06 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -14,7 +14,7 @@ import pytest import torch from pytest_mock import MockerFixture -from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker +from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker, _make_diffusion_vllm_model_config pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] @@ -78,6 +78,34 @@ class TestDiffusionWorkerLoadWeights: assert result == set() +def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods(): + from types import SimpleNamespace + + from vllm_omni.quantization import build_quant_config + + od_config = SimpleNamespace( + model="dummy", + dtype=torch.bfloat16, + quantization_config=build_quant_config( + { + "quant_method": "modelopt", + "quant_algo": "FP8", + "ignore": [], + } + ), + tf_model_config=SimpleNamespace(), + enforce_eager=True, + is_moe=False, + ) + + model_config = _make_diffusion_vllm_model_config(od_config) + + assert model_config.dtype is torch.bfloat16 + assert model_config.quantization == "modelopt" + assert model_config.quantization_config is od_config.quantization_config + assert model_config.is_quantized() + + class TestDiffusionWorkerSleep: """Test DiffusionWorker.sleep method.""" diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 9491b759a..80171e1d6 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -208,10 +208,11 @@ class TransformerConfig: quant_method: str | None = None quant_config: QuantizationConfig | None = None disk_qc = params.get("quantization_config") - if isinstance(disk_qc, dict) and "quant_method" in disk_qc: - quant_method = disk_qc["quant_method"] - kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"} - quant_config = build_quant_config(quant_method, **kwargs) + if isinstance(disk_qc, dict): + raw_quant_method = disk_qc.get("quant_method", disk_qc.get("method")) + quant_config = build_quant_config(disk_qc) + if quant_config is not None: + quant_method = raw_quant_method if raw_quant_method is not None else quant_config.get_name() return cls(params=params, quant_method=quant_method, quant_config=quant_config) @@ -513,6 +514,10 @@ class OmniDiffusionConfig: # str is resolved to {"method": } internally. # Per-component: {"transformer": {"method": "fp8"}, "vae": None} quantization_config: str | QuantizationConfig | dict[str, Any] | None = None + # Explicit runtime override for ModelOpt FP8 diffusion checkpoints. This + # does not enable FP8 by itself; it only selects CUTLASS once the checkpoint + # has already resolved to vLLM's ModelOpt FP8 linear method. + force_cutlass_fp8: bool = False # Diffusion pipeline Profiling config enable_diffusion_pipeline_profiler: bool = False @@ -630,14 +635,9 @@ class OmniDiffusionConfig: # Auto-detect quantization from TransformerConfig if not explicitly set. # This covers the case where tf_model_config is passed at construction - # time. For late (post-construction) assignment, callers should use + # time. For late (post-construction) assignment, callers should use # set_tf_model_config() which propagates quant_config automatically. - if self.quantization_config is None and self.tf_model_config.quant_config is not None: - self.quantization_config = self.tf_model_config.quant_config - logger.info( - "Auto-detected quantization '%s' from model config", - self.tf_model_config.quant_method, - ) + self._propagate_quantization_from_tf_config(self.tf_model_config) # Resolve quantization_config: str/dict -> QuantizationConfig via build_quant_config. if self.quantization_config is not None: @@ -664,6 +664,32 @@ class OmniDiffusionConfig: "valid together with diffusion_load_format=diffusers" ) + def _propagate_quantization_from_tf_config(self, tf_config: "TransformerConfig") -> None: + if tf_config.quant_config is None: + return + + is_checkpoint_fp8 = bool(getattr(tf_config.quant_config, "is_checkpoint_fp8_serialized", False)) + should_use_checkpoint_config = self.quantization_config is None or ( + is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config) + ) + if should_use_checkpoint_config: + self.quantization_config = tf_config.quant_config + logger.info( + "Auto-detected quantization '%s' from model config", + tf_config.quant_method, + ) + + @staticmethod + def _is_generic_fp8_quant_config(quant_config: object) -> bool: + if isinstance(quant_config, str): + return quant_config.lower() == "fp8" + if isinstance(quant_config, Mapping): + method = quant_config.get("method", quant_config.get("quant_method")) + return isinstance(method, str) and method.lower() == "fp8" + if hasattr(quant_config, "get_name"): + return quant_config.get_name() == "fp8" + return False + def set_tf_model_config(self, tf_config: "TransformerConfig") -> None: """Assign `tf_model_config` and propagate quantization if detected. @@ -679,12 +705,7 @@ class OmniDiffusionConfig: `TransformerConfig.from_dict`. """ self.tf_model_config = tf_config - if self.quantization_config is None and tf_config.quant_config is not None: - self.quantization_config = tf_config.quant_config - logger.info( - "Auto-detected quantization '%s' from model config", - tf_config.quant_method, - ) + self._propagate_quantization_from_tf_config(tf_config) def update_multimodal_support(self) -> None: # Resolve serving-visible multimodal behavior from shared metadata @@ -716,50 +737,35 @@ class OmniDiffusionConfig: # Skip transformer config loading for diffusers adapter # (non-DiT models don't have a separate transformer folder/config) if self.diffusion_load_format == "diffusers": - self.tf_model_config = TransformerConfig() - try: - diffusers_pipeline_cls_name = config_dict["_class_name"] - self.diffusers_pipeline_cls = getattr(diffusers, diffusers_pipeline_cls_name) - except (KeyError, AttributeError) as exc: - logger.warning( - f"Could not find valid _class_name for diffusers pipeline in model_index.json: {exc}. " - "Without knowing the underlying diffusers pipeline class, the dummy run will input only " - "text prompt, which may cause errors for pipelines that require additional inputs." - ) + self.set_tf_model_config(TransformerConfig()) else: tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model) - self.tf_model_config = TransformerConfig.from_dict(tf_config_dict) + self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict)) else: raise FileNotFoundError("model_index.json not found") except (AttributeError, OSError, ValueError, FileNotFoundError): # Skip transformer config loading for diffusers adapter # (non-DiT models don't have a separate transformer folder/config) if self.diffusion_load_format == "diffusers": - self.tf_model_config = TransformerConfig() - logger.warning( - "Could not find valid model_index.json per diffusers format. " - "This model is suspectedly unsupported by the diffusers backend. " - "Also, without knowing the underlying diffusers pipeline class from model_index.json, " - "the dummy run will input only text prompt, which may cause errors for pipelines " - "that require additional inputs." - ) + self.set_tf_model_config(TransformerConfig()) + self.update_multimodal_support() else: cfg = get_hf_file_to_dict("config.json", self.model) if cfg is None: raise ValueError(f"Could not find config.json or model_index.json for model {self.model}") - self.tf_model_config = TransformerConfig.from_dict(cfg) + self.set_tf_model_config(TransformerConfig.from_dict(cfg)) model_type = cfg.get("model_type") architectures = cfg.get("architectures") or [] if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: self.model_class_name = "BagelPipeline" - self.tf_model_config = TransformerConfig() + self.set_tf_model_config(TransformerConfig()) self.update_multimodal_support() elif model_type == "nextstep": if self.model_class_name is None: self.model_class_name = "NextStep11Pipeline" - self.tf_model_config = TransformerConfig() + self.set_tf_model_config(TransformerConfig()) self.update_multimodal_support() elif model_type == "s2v": if self.model_class_name is None: @@ -794,23 +800,10 @@ class OmniDiffusionConfig: cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER") kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none" - # Falsy-value check for not-None fields (convert potential None values in YAML config to empty containers) - if "diffusers_load_kwargs" in kwargs and kwargs["diffusers_load_kwargs"] is None: - kwargs["diffusers_load_kwargs"] = {} - if "diffusers_call_kwargs" in kwargs and kwargs["diffusers_call_kwargs"] is None: - kwargs["diffusers_call_kwargs"] = {} - - # Forward top-level parallel knobs (e.g. --tensor-parallel-size from CLI) - # into parallel_config so the diffusion engine sees them. - par = kwargs.get("parallel_config", {}) - if isinstance(par, Mapping): - par = dict(par) - if par.get("tensor_parallel_size") is None: - par.pop("tensor_parallel_size", None) - tensor_parallel_size = kwargs.get("tensor_parallel_size") - if tensor_parallel_size is not None and "tensor_parallel_size" not in par: - par["tensor_parallel_size"] = tensor_parallel_size - kwargs["parallel_config"] = par + # Convert optional YAML null values to empty containers. + for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"): + if key in kwargs and kwargs[key] is None: + kwargs[key] = {} # Filter kwargs to only include valid fields valid_fields = {f.name for f in fields(cls)} diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 127e1b2cb..61b62b2f4 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -65,6 +65,22 @@ def apply_rotary_emb_mindiesd( return rotary_position_embedding(x, cos, sin, rotated_mode="rotated_half", head_first=False, fused=True) +def _ensure_batch_dim(x: torch.Tensor) -> tuple[torch.Tensor, bool]: + # Upstream fused rotary kernels expect ``x`` shaped as + # ``[batch_size, seq_len, nheads, headdim]``. Some omni diffusion call + # sites pass ``[seq_len, nheads, headdim]`` instead, so normalize to 4D + # here before entering the fused path. + if x.dim() == 3: + return x.unsqueeze(0), True + return x, False + + +def _restore_batch_dim(x: torch.Tensor, squeezed: bool) -> torch.Tensor: + if squeezed: + return x.squeeze(0) + return x + + class RotaryEmbedding(CustomOp): """ rotary positional embedding. @@ -98,12 +114,14 @@ class RotaryEmbedding(CustomOp): cos = cos[0] sin = sin[0] - return apply_rotary_emb( + x, squeezed = _ensure_batch_dim(x) + output = apply_rotary_emb( x, cos, sin, interleaved=self.interleaved, ) + return _restore_batch_dim(output, squeezed) def forward_hip( self, @@ -119,12 +137,14 @@ class RotaryEmbedding(CustomOp): cos = cos[0] sin = sin[0] - return self.apply_rotary_emb_flash_attn( + x, squeezed = _ensure_batch_dim(x) + output = self.apply_rotary_emb_flash_attn( x, cos, sin, interleaved=self.interleaved, ) + return _restore_batch_dim(output, squeezed) def forward_npu( self, diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py new file mode 100644 index 000000000..cc20b2da4 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/__init__.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from torch import nn + +from .modelopt import ModelOptFp8CheckpointAdapter + + +def get_checkpoint_adapter( + model: nn.Module, + source: object, + quant_config: object | None, + use_safetensors: bool, +) -> ModelOptFp8CheckpointAdapter | None: + if ModelOptFp8CheckpointAdapter.is_compatible(source, quant_config, use_safetensors): + return ModelOptFp8CheckpointAdapter(model, source) + return None + + +__all__ = [ + "ModelOptFp8CheckpointAdapter", + "get_checkpoint_adapter", +] diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py new file mode 100644 index 000000000..6102a55e6 --- /dev/null +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Generator, Iterable +from dataclasses import dataclass, field + +import torch +from torch import nn +from vllm.logger import init_logger +from vllm.model_executor.models.utils import WeightsMapper +from vllm.model_executor.utils import get_packed_modules_mapping + +logger = init_logger(__name__) + +MODEL_OPT_SCALE_SUFFIXES = (".input_scale", ".weight_scale", ".weight_scale_inv") +DEFAULT_PACKED_MODULES_MAPPING = { + "to_qkv": ("to_q", "to_k", "to_v"), + "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"), + "w13": ("w1", "w3"), +} +FP8_DTYPES = tuple( + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e5m2", None), + getattr(torch, "float8_e4m3fnuz", None), + getattr(torch, "float8_e5m2fnuz", None), + ) + if dtype is not None +) + + +@dataclass +class _AdaptState: + scale_tensors: dict[str, torch.Tensor] = field(default_factory=dict) + pending_weights: dict[str, list[tuple[str, torch.Tensor, torch.dtype]]] = field(default_factory=dict) + skipped_scales: int = 0 + dequantized_weights: int = 0 + + +class ModelOptFp8CheckpointAdapter: + def __init__(self, model: nn.Module, source: object): + self._loadable_tensors = self._get_model_loadable_tensors(model) + self._weights_mapper = self._get_weights_mapper(model) + self._source_label = getattr(source, "prefix", "") or getattr(source, "subfolder", None) or "model" + + @classmethod + def is_compatible( + cls, + source: object, + quant_config: object | None, + use_safetensors: bool, + ) -> bool: + return use_safetensors and cls._is_transformer_source(source) and cls._is_checkpoint_quant_config(quant_config) + + @staticmethod + def _is_transformer_source(source: object) -> bool: + if getattr(source, "subfolder", None) == "transformer": + return True + return str(getattr(source, "prefix", "")).startswith("transformer.") + + @staticmethod + def _is_checkpoint_quant_config(quant_config: object | None) -> bool: + return ( + quant_config is not None + and hasattr(quant_config, "get_name") + and quant_config.get_name() == "modelopt" + and bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False)) + ) + + @staticmethod + def _get_model_loadable_tensors(model: nn.Module) -> dict[str, torch.Tensor]: + loadable_tensors: dict[str, torch.Tensor] = {name: param for name, param in model.named_parameters()} + loadable_tensors.update({name: buffer for name, buffer in model.named_buffers()}) + return loadable_tensors + + @staticmethod + def _is_scale(name: str) -> bool: + return name.endswith(MODEL_OPT_SCALE_SUFFIXES) + + @staticmethod + def _is_fp8_tensor(tensor: torch.Tensor) -> bool: + return tensor.dtype in FP8_DTYPES + + @staticmethod + def _get_weight_scale_name(weight_name: str) -> str | None: + if weight_name.endswith(".weight"): + return weight_name[: -len(".weight")] + ".weight_scale" + return None + + @classmethod + def _get_weights_mapper(cls, model: nn.Module) -> WeightsMapper: + mapping = { + packed_name: tuple(shard_names) for packed_name, shard_names in DEFAULT_PACKED_MODULES_MAPPING.items() + } + mapping.update( + { + str(packed_name): tuple(str(shard_name) for shard_name in shard_names) + for packed_name, shard_names in get_packed_modules_mapping(model).items() + } + ) + + orig_to_new_substr = {".to_out.0.": ".to_out."} + orig_to_new_prefix: dict[str, str] = {} + for packed_name, shard_names in mapping.items(): + for shard_name in shard_names: + orig_to_new_substr[f".{shard_name}."] = f".{packed_name}." + orig_to_new_prefix[f"{shard_name}."] = f"{packed_name}." + return WeightsMapper( + orig_to_new_substr=orig_to_new_substr, + orig_to_new_prefix=orig_to_new_prefix, + ) + + def _resolve_target_name(self, name: str) -> str | None: + if name in self._loadable_tensors: + return name + + for candidate in self._weights_mapper.apply_list([name]): + if candidate != name and candidate in self._loadable_tensors: + return candidate + return None + + @staticmethod + def _reshape_weight_scale(scale: torch.Tensor, weight_shape: torch.Size) -> torch.Tensor: + if scale.numel() == 1: + return scale.reshape(()) + if len(weight_shape) == 2 and scale.ndim == 1 and scale.shape[0] == weight_shape[0]: + return scale.reshape(-1, 1) + if tuple(scale.shape) == tuple(weight_shape): + return scale + if ( + len(weight_shape) == 2 + and scale.ndim == 4 + and scale.shape[1] == 1 + and scale.shape[3] == 1 + and weight_shape[0] % scale.shape[0] == 0 + and weight_shape[1] % scale.shape[2] == 0 + ): + block_n = weight_shape[0] // scale.shape[0] + block_k = weight_shape[1] // scale.shape[2] + return scale.expand(scale.shape[0], block_n, scale.shape[2], block_k).reshape(weight_shape) + raise ValueError(f"Unsupported ModelOpt FP8 weight_scale shape {tuple(scale.shape)} for weight {weight_shape}") + + def _dequantize_weight( + self, + name: str, + loaded_weight: torch.Tensor, + state: _AdaptState, + target_dtype: torch.dtype, + ) -> torch.Tensor: + scale_name = self._get_weight_scale_name(name) + if scale_name is None or scale_name not in state.scale_tensors: + raise ValueError(f"Missing ModelOpt FP8 weight_scale for full-precision target weight {name!r}") + + scale = state.scale_tensors[scale_name].to(dtype=target_dtype, device=loaded_weight.device) + scale = self._reshape_weight_scale(scale, loaded_weight.shape) + return loaded_weight.to(dtype=target_dtype) * scale + + def _flush_pending_weights( + self, + scale_name: str, + state: _AdaptState, + ) -> Generator[tuple[str, torch.Tensor], None, None]: + for weight_name, weight_tensor, target_dtype in state.pending_weights.pop(scale_name, []): + yield weight_name, self._dequantize_weight(weight_name, weight_tensor, state, target_dtype) + state.dequantized_weights += 1 + + def _handle_scale_tensor( + self, + name: str, + tensor: torch.Tensor, + target_name: str | None, + state: _AdaptState, + ) -> Generator[tuple[str, torch.Tensor], None, None]: + state.scale_tensors[name] = tensor + if target_name is None: + state.skipped_scales += 1 + else: + yield name, tensor + yield from self._flush_pending_weights(name, state) + + def _target_dtype_for_dequantization( + self, + tensor: torch.Tensor, + target_name: str | None, + ) -> torch.dtype | None: + if target_name is None or not self._is_fp8_tensor(tensor): + return None + + target_dtype = self._loadable_tensors[target_name].dtype + if target_dtype in FP8_DTYPES: + return None + return target_dtype + + def _maybe_dequantize_or_defer_weight( + self, + name: str, + tensor: torch.Tensor, + target_dtype: torch.dtype, + state: _AdaptState, + ) -> torch.Tensor | None: + scale_name = self._get_weight_scale_name(name) + if scale_name is None: + raise ValueError(f"Missing ModelOpt FP8 weight_scale name for weight {name!r}") + + if scale_name not in state.scale_tensors: + state.pending_weights.setdefault(scale_name, []).append((name, tensor, target_dtype)) + return None + + state.dequantized_weights += 1 + return self._dequantize_weight(name, tensor, state, target_dtype) + + @staticmethod + def _check_pending_weights(state: _AdaptState) -> None: + if not state.pending_weights: + return + + missing_scale_names = ", ".join(repr(name) for name in sorted(state.pending_weights)) + raise ValueError(f"Missing ModelOpt FP8 weight_scale for full-precision target weights: {missing_scale_names}") + + def _log_adaptation_summary(self, state: _AdaptState) -> None: + if not state.skipped_scales and not state.dequantized_weights: + return + + logger.info_once( + "Adapted ModelOpt FP8 %s weights: dequantized %d full-precision weights, skipped %d scale tensors", + self._source_label, + state.dequantized_weights, + state.skipped_scales, + ) + + def adapt( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Generator[tuple[str, torch.Tensor], None, None]: + state = _AdaptState() + + for name, tensor in weights: + target_name = self._resolve_target_name(name) + if self._is_scale(name): + yield from self._handle_scale_tensor(name, tensor, target_name, state) + continue + + target_dtype = self._target_dtype_for_dequantization(tensor, target_name) + if target_dtype is not None: + tensor = self._maybe_dequantize_or_defer_weight(name, tensor, target_dtype, state) + if tensor is None: + continue + yield name, tensor + + self._check_pending_weights(state) + self._log_adaptation_summary(state) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index eb89e7e81..1d2c9e357 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -7,7 +7,7 @@ import re import time from collections.abc import Generator, Iterable from pathlib import Path -from typing import TYPE_CHECKING, cast +from typing import cast import torch from huggingface_hub import hf_hub_download @@ -32,14 +32,14 @@ from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig +from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( + get_checkpoint_adapter, +) from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.registry import initialize_model -if TYPE_CHECKING: - from vllm_omni.diffusion.data import OmniDiffusionConfig - logger = init_logger(__name__) @@ -177,7 +177,11 @@ class DiffusersPipelineLoader: return hf_folder, hf_weights_files, use_safetensors - def _get_weights_iterator(self, source: "ComponentSource") -> Generator[tuple[str, torch.Tensor], None, None]: + def _get_weights_iterator( + self, + source: "ComponentSource", + model: nn.Module | None = None, + ) -> Generator[tuple[str, torch.Tensor], None, None]: """Get an iterator for the model weights based on the load format.""" _, hf_weights_files, use_safetensors = self._prepare_weights( source.model_or_path, @@ -213,7 +217,34 @@ class DiffusersPipelineLoader: if self.counter_before_loading_weights == 0.0: self.counter_before_loading_weights = time.perf_counter() # Apply the prefix. - return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) + prefixed_weights_iterator = ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) + if model is not None: + checkpoint_adapter = self._get_checkpoint_adapter(model, source, use_safetensors) + if checkpoint_adapter is not None: + return checkpoint_adapter.adapt(prefixed_weights_iterator) + return prefixed_weights_iterator + + def _get_source_quant_config(self, source: "ComponentSource") -> object | None: + if self.od_config is None: + return None + + quant_config = self.od_config.quantization_config + if hasattr(quant_config, "resolve"): + return quant_config.resolve(source.prefix.rstrip(".")) + return quant_config + + def _get_checkpoint_adapter( + self, + model: nn.Module, + source: "ComponentSource", + use_safetensors: bool, + ): + return get_checkpoint_adapter( + model=model, + source=source, + quant_config=self._get_source_quant_config(source), + use_safetensors=use_safetensors, + ) def get_all_weights( self, @@ -221,7 +252,7 @@ class DiffusersPipelineLoader: ) -> Generator[tuple[str, torch.Tensor], None, None]: sources = self._get_weight_sources(model) for source in sources: - yield from self._get_weights_iterator(source) + yield from self._get_weights_iterator(source, model=model) def _get_weight_sources(self, model: nn.Module) -> tuple["ComponentSource", ...]: return tuple( @@ -269,7 +300,7 @@ class DiffusersPipelineLoader: """Load a model with the given configurations.""" if load_format is None: load_format = "default" - + self.od_config = od_config # CPU offload + FP8: load weights on device for FP8 quantization if load_device == "cpu" and od_config.quantization_config is not None: load_device = device.type @@ -291,13 +322,12 @@ class DiffusersPipelineLoader: model_cls = resolve_obj_by_qualname(custom_pipeline_name) model = model_cls(od_config=od_config) else: - # 'dummy' format should not call this function at all raise ValueError(f"Unknown load_format: {load_format}") logger.debug("Loading weights on %s ...", load_device) if load_format == "diffusers": # DiffusersAdapterPipeline.load_weights() calls - # DiffusionPipeline.from_pretrained() internally — it does - # NOT use our native (customized) pipeline classes. + # DiffusionPipeline.from_pretrained() internally; it does + # not use our native customized pipeline classes. cast(DiffusersAdapterPipeline, model).load_weights() elif self._is_gguf_quantization(od_config): self._load_weights_with_gguf(model, od_config) @@ -449,9 +479,7 @@ class DiffusersPipelineLoader: def _get_model_loadable_names(self, model: nn.Module) -> set[str]: # Avoid model.state_dict() here because GGUF uses UninitializedParameter # which raises during detach(). Collect names directly. - names = {name for name, _ in model.named_parameters()} - names.update(name for name, _ in model.named_buffers()) - return names + return {name for name, _ in model.named_parameters()} | {name for name, _ in model.named_buffers()} def _resolve_gguf_model_path(self, gguf_model: str, revision: str | None) -> str: if os.path.isfile(gguf_model): diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py index 0a4452197..ef2cdd025 100644 --- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py +++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterable from types import SimpleNamespace -from typing import Any +from typing import TYPE_CHECKING, Any import torch from diffusers.models.embeddings import ( @@ -27,6 +27,13 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.layers.rope import RotaryEmbedding +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + + +def _join_prefix(prefix: str, name: str) -> str: + return f"{prefix}.{name}" if prefix else name + class Flux2SwiGLU(nn.Module): """SwiGLU activation used by Flux2.""" @@ -48,6 +55,8 @@ class Flux2FeedForward(nn.Module): mult: float = 3.0, inner_dim: int | None = None, bias: bool = False, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() if inner_dim is None: @@ -59,6 +68,8 @@ class Flux2FeedForward(nn.Module): [inner_dim, inner_dim], bias=bias, return_bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "linear_in"), ) self.act_fn = Flux2SwiGLU() self.linear_out = RowParallelLinear( @@ -67,6 +78,8 @@ class Flux2FeedForward(nn.Module): bias=bias, input_is_parallel=True, return_bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "linear_out"), ) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -89,6 +102,8 @@ class Flux2Attention(nn.Module): eps: float = 1e-5, out_dim: int = None, elementwise_affine: bool = True, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() self.head_dim = dim_head @@ -104,6 +119,8 @@ class Flux2Attention(nn.Module): head_size=self.head_dim, total_num_heads=self.heads, bias=bias, + quant_config=quant_config, + prefix=_join_prefix(prefix, "to_qkv"), ) self.query_num_heads = self.to_qkv.num_heads self.kv_num_heads = self.to_qkv.num_kv_heads @@ -119,6 +136,8 @@ class Flux2Attention(nn.Module): bias=out_bias, input_is_parallel=True, return_bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "to_out.0"), ), nn.Dropout(dropout), ] @@ -132,6 +151,8 @@ class Flux2Attention(nn.Module): head_size=self.head_dim, total_num_heads=self.heads, bias=added_proj_bias, + quant_config=quant_config, + prefix=_join_prefix(prefix, "add_kv_proj"), ) self.add_query_num_heads = self.add_kv_proj.num_heads self.add_kv_num_heads = self.add_kv_proj.num_kv_heads @@ -141,6 +162,8 @@ class Flux2Attention(nn.Module): bias=out_bias, input_is_parallel=True, return_bias=False, + quant_config=quant_config, + prefix=_join_prefix(prefix, "to_add_out"), ) self.rope = RotaryEmbedding(is_neox_style=False) @@ -238,6 +261,8 @@ class Flux2ParallelSelfAttention(nn.Module): elementwise_affine: bool = True, mlp_ratio: float = 4.0, mlp_mult_factor: int = 2, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() self.head_dim = dim_head @@ -256,6 +281,8 @@ class Flux2ParallelSelfAttention(nn.Module): self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor, bias=bias, gather_output=True, + quant_config=quant_config, + prefix=_join_prefix(prefix, "to_qkv_mlp_proj"), ) self.mlp_act_fn = Flux2SwiGLU() @@ -267,6 +294,8 @@ class Flux2ParallelSelfAttention(nn.Module): self.out_dim, bias=out_bias, gather_output=True, + quant_config=quant_config, + prefix=_join_prefix(prefix, "to_out"), ) self.rope = RotaryEmbedding(is_neox_style=False) self.attn = Attention( @@ -329,6 +358,8 @@ class Flux2SingleTransformerBlock(nn.Module): mlp_ratio: float = 3.0, eps: float = 1e-6, bias: bool = False, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) @@ -342,6 +373,8 @@ class Flux2SingleTransformerBlock(nn.Module): eps=eps, mlp_ratio=mlp_ratio, mlp_mult_factor=2, + quant_config=quant_config, + prefix=_join_prefix(prefix, "attn"), ) def forward( @@ -389,6 +422,8 @@ class Flux2TransformerBlock(nn.Module): mlp_ratio: float = 3.0, eps: float = 1e-6, bias: bool = False, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) @@ -406,13 +441,29 @@ class Flux2TransformerBlock(nn.Module): added_proj_bias=bias, out_bias=bias, eps=eps, + quant_config=quant_config, + prefix=_join_prefix(prefix, "attn"), ) self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) - self.ff = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias) + self.ff = Flux2FeedForward( + dim=dim, + dim_out=dim, + mult=mlp_ratio, + bias=bias, + quant_config=quant_config, + prefix=_join_prefix(prefix, "ff"), + ) self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) - self.ff_context = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias) + self.ff_context = Flux2FeedForward( + dim=dim, + dim_out=dim, + mult=mlp_ratio, + bias=bias, + quant_config=quant_config, + prefix=_join_prefix(prefix, "ff_context"), + ) def forward( self, @@ -576,6 +627,7 @@ class Flux2Transformer2DModel(nn.Module): rope_theta: int = 2000, eps: float = 1e-6, guidance_embeds: bool = True, + quant_config: "QuantizationConfig | None" = None, ): super().__init__() self.guidance_embeds = guidance_embeds @@ -624,8 +676,10 @@ class Flux2Transformer2DModel(nn.Module): mlp_ratio=mlp_ratio, eps=eps, bias=False, + quant_config=quant_config, + prefix=f"transformer_blocks.{i}", ) - for _ in range(num_layers) + for i in range(num_layers) ] ) @@ -638,8 +692,10 @@ class Flux2Transformer2DModel(nn.Module): mlp_ratio=mlp_ratio, eps=eps, bias=False, + quant_config=quant_config, + prefix=f"single_transformer_blocks.{i}", ) - for _ in range(num_single_layers) + for i in range(num_single_layers) ] ) diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py index 404f05b60..44d991ad3 100644 --- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py +++ b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py @@ -378,7 +378,7 @@ class Flux2Pipeline(nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarM self._execution_device ) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel) - self.transformer = Flux2Transformer2DModel(**transformer_kwargs) + self.transformer = Flux2Transformer2DModel(quant_config=od_config.quantization_config, **transformer_kwargs) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py index 9cf2fb756..a708a8ece 100644 --- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py +++ b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py @@ -53,6 +53,10 @@ if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig +def _join_prefix(prefix: str, name: str) -> str: + return f"{prefix}.{name}" if prefix else name + + class Flux2SwiGLU(nn.Module): """SwiGLU activation used by Flux2.""" @@ -74,6 +78,7 @@ class Flux2FeedForward(nn.Module): inner_dim: int | None = None, bias: bool = False, quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() if inner_dim is None: @@ -86,6 +91,7 @@ class Flux2FeedForward(nn.Module): bias=bias, return_bias=False, quant_config=quant_config, + prefix=_join_prefix(prefix, "linear_in"), ) self.act_fn = Flux2SwiGLU() self.linear_out = RowParallelLinear( @@ -95,6 +101,7 @@ class Flux2FeedForward(nn.Module): input_is_parallel=True, return_bias=False, quant_config=quant_config, + prefix=_join_prefix(prefix, "linear_out"), ) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -119,6 +126,7 @@ class Flux2Attention(nn.Module): out_dim: int = None, elementwise_affine: bool = True, quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() self.parallel_config = parallel_config @@ -136,6 +144,7 @@ class Flux2Attention(nn.Module): total_num_heads=self.heads, bias=bias, quant_config=quant_config, + prefix=_join_prefix(prefix, "to_qkv"), ) self.query_num_heads = self.to_qkv.num_heads self.kv_num_heads = self.to_qkv.num_kv_heads @@ -152,6 +161,7 @@ class Flux2Attention(nn.Module): input_is_parallel=True, return_bias=False, quant_config=quant_config, + prefix=_join_prefix(prefix, "to_out.0"), ), nn.Dropout(dropout), ] @@ -166,6 +176,7 @@ class Flux2Attention(nn.Module): total_num_heads=self.heads, bias=added_proj_bias, quant_config=quant_config, + prefix=_join_prefix(prefix, "add_kv_proj"), ) self.add_query_num_heads = self.add_kv_proj.num_heads self.add_kv_num_heads = self.add_kv_proj.num_kv_heads @@ -176,6 +187,7 @@ class Flux2Attention(nn.Module): input_is_parallel=True, return_bias=False, quant_config=quant_config, + prefix=_join_prefix(prefix, "to_add_out"), ) self.rope = RotaryEmbedding(is_neox_style=False) @@ -330,6 +342,7 @@ class Flux2ParallelSelfAttention(nn.Module): mlp_ratio: float = 4.0, mlp_mult_factor: int = 2, quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() self.parallel_config = parallel_config @@ -350,6 +363,7 @@ class Flux2ParallelSelfAttention(nn.Module): bias=bias, gather_output=True, quant_config=quant_config, + prefix=_join_prefix(prefix, "to_qkv_mlp_proj"), ) self.mlp_act_fn = Flux2SwiGLU() @@ -362,6 +376,7 @@ class Flux2ParallelSelfAttention(nn.Module): bias=out_bias, gather_output=True, quant_config=quant_config, + prefix=_join_prefix(prefix, "to_out"), ) self.rope = RotaryEmbedding(is_neox_style=False) self.attn = Attention( @@ -463,6 +478,7 @@ class Flux2SingleTransformerBlock(nn.Module): eps: float = 1e-6, bias: bool = False, quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) @@ -478,6 +494,7 @@ class Flux2SingleTransformerBlock(nn.Module): mlp_ratio=mlp_ratio, mlp_mult_factor=2, quant_config=quant_config, + prefix=_join_prefix(prefix, "attn"), ) def forward( @@ -534,6 +551,7 @@ class Flux2TransformerBlock(nn.Module): eps: float = 1e-6, bias: bool = False, quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) @@ -551,10 +569,18 @@ class Flux2TransformerBlock(nn.Module): out_bias=bias, eps=eps, quant_config=quant_config, + prefix=_join_prefix(prefix, "attn"), ) self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) - self.ff = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias, quant_config=quant_config) + self.ff = Flux2FeedForward( + dim=dim, + dim_out=dim, + mult=mlp_ratio, + bias=bias, + quant_config=quant_config, + prefix=_join_prefix(prefix, "ff"), + ) self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) self.ff_context = Flux2FeedForward( @@ -563,6 +589,7 @@ class Flux2TransformerBlock(nn.Module): mult=mlp_ratio, bias=bias, quant_config=quant_config, + prefix=_join_prefix(prefix, "ff_context"), ) def forward( @@ -835,8 +862,9 @@ class Flux2Transformer2DModel(nn.Module): eps=eps, bias=False, quant_config=quant_config, + prefix=f"transformer_blocks.{i}", ) - for _ in range(num_layers) + for i in range(num_layers) ] ) @@ -851,8 +879,9 @@ class Flux2Transformer2DModel(nn.Module): eps=eps, bias=False, quant_config=quant_config, + prefix=f"single_transformer_blocks.{i}", ) - for _ in range(num_single_layers) + for i in range(num_single_layers) ] ) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py index 63ba0d724..fe204502b 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py @@ -27,8 +27,16 @@ def _set_forward_context_num_tokens(num_tokens: int) -> None: class HunyuanFusedMoEDefault(FusedMoE): def __init__(self, *, prefix: str = "", **kwargs: Any) -> None: + # Current vLLM FusedMoE handles output reduction internally. + kwargs.pop("reduce_results", None) super().__init__(prefix=prefix, **kwargs) self._prefix = prefix + self._init_hook_handle = self.register_forward_pre_hook(self._initialize_kernel_hook, with_kwargs=True) + + def _initialize_kernel_hook(self, module: Any, args: Any, kwargs: Any) -> None: + if self.quant_method and getattr(self.quant_method, "moe_kernel", None) is None: + self.quant_method.process_weights_after_loading(self) + self._init_hook_handle.remove() def forward(self, hidden_states: Any, router_logits: Any) -> Any: _set_forward_context_num_tokens(hidden_states.shape[0]) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index a4ba514d9..06cc6cfde 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -2035,6 +2035,44 @@ class HunyuanImage3Model(nn.Module): return True return False + def is_scalar_quant_scale(name: str, tensor: torch.Tensor) -> bool: + return tensor.numel() == 1 and name.endswith((".input_scale", ".weight_scale")) + + def load_split_param( + name: str, + tensor: torch.Tensor, + den: int, + split_param: list[tuple[str | int, int]], + func: Callable[[torch.Tensor], torch.Tensor] | None, + ) -> None: + param = params_dict[name] + weight_loader = param.weight_loader + if is_scalar_quant_scale(name, tensor): + for shard_id, _ in split_param: + weight_loader(param, tensor, shard_id) + return + + assert tensor.shape[0] % den == 0 + units = tensor.shape[0] // den + offset = 0 + tensor = func(tensor) if func else tensor + for shard_id, num in split_param: + new_offset = offset + num * units + weight_loader(param, tensor[offset:new_offset], shard_id) + offset = new_offset + + def get_loaded_weight_shard( + name: str, + tensor: torch.Tensor, + offset: int, + den: int, + ) -> torch.Tensor: + if is_scalar_quant_scale(name, tensor): + return tensor + assert tensor.shape[0] % den == 0 + units = tensor.shape[0] // den + return tensor[offset * units : offset * units + units] + for name, loaded_weight in weights: # print(f"Loading weight name: {name}, tp_rank: {tp_rank}", flush=True) if contains_unexpected_keyword(name, unexpected_keywords): @@ -2112,19 +2150,7 @@ class HunyuanImage3Model(nn.Module): if is_pp_missing_parameter(name, self): continue - assert loaded_weight.shape[0] % den == 0 - units = loaded_weight.shape[0] // den - param = params_dict[name] - weight_loader = param.weight_loader - offset = 0 - for shard_id, num in split_param: - new_offset = offset + num * units - if func: - weight_loader(param, func(loaded_weight)[offset:new_offset], shard_id) - else: - weight_loader(param, loaded_weight[offset:new_offset], shard_id) - offset = new_offset - + load_split_param(name, loaded_weight, den, split_param, func) break else: # Skip loading extra bias for GPTQ models. @@ -2154,12 +2180,11 @@ class HunyuanImage3Model(nn.Module): continue param = params_dict[name_mapped] weight_loader = cast(Callable[..., bool], param.weight_loader) - assert loaded_weight.shape[0] % den == 0 - units = loaded_weight.shape[0] // den + loaded_weight_shard = get_loaded_weight_shard(name, loaded_weight, offset, den) success = weight_loader( param, - loaded_weight[offset * units : offset * units + units], + loaded_weight_shard, name_mapped, shard_id=shard_id, expert_id=expert_id, diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index 527c472eb..1399eb0f0 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -209,6 +209,18 @@ def validate_zimage_tp_constraints( return ffn_hidden_dim, final_out_dims, supported_tp_candidates +def _join_prefix(prefix: str, name: str) -> str: + return f"{prefix}.{name}" if prefix else name + + +def _restore_linear_output_shape(output: torch.Tensor, input_: torch.Tensor) -> torch.Tensor: + if input_.dim() > 2 and output.dim() == 2: + token_count = math.prod(input_.shape[:-1]) + if output.shape[0] == token_count: + return output.reshape(*input_.shape[:-1], output.shape[-1]) + return output + + class TimestepEmbedder(nn.Module): def __init__( self, out_size, mid_size=None, frequency_embedding_size=256, quant_config: "QuantizationConfig | None" = None @@ -256,6 +268,7 @@ class ZImageAttention(nn.Module): qk_norm: bool = True, eps: float = 1e-6, quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ) -> None: super().__init__() self.dim = dim @@ -271,7 +284,7 @@ class ZImageAttention(nn.Module): total_num_kv_heads=num_kv_heads, bias=False, quant_config=quant_config, - prefix="to_qkv", + prefix=_join_prefix(prefix, "to_qkv"), ) assert qk_norm is True @@ -290,7 +303,7 @@ class ZImageAttention(nn.Module): input_is_parallel=True, return_bias=False, quant_config=quant_config, - prefix="to_out", + prefix=_join_prefix(prefix, "to_out.0"), ) ] ) @@ -312,6 +325,7 @@ class ZImageAttention(nn.Module): sin: torch.Tensor, ): qkv, _ = self.to_qkv(hidden_states) + qkv = _restore_linear_output_shape(qkv, hidden_states) q_size = self.to_qkv.num_heads * self.head_dim kv_size = self.to_qkv.num_kv_heads * self.head_dim query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) @@ -344,7 +358,9 @@ class ZImageAttention(nn.Module): hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(dtype) + to_out_input = hidden_states hidden_states = self.to_out[0](hidden_states) + hidden_states = _restore_linear_output_shape(hidden_states, to_out_input) return hidden_states @@ -364,7 +380,7 @@ class FeedForward(nn.Module): bias=False, return_bias=False, quant_config=quant_config, - prefix=prefix, + prefix=_join_prefix(prefix, "w13"), ) self.act = SiluAndMul() self.w2 = RowParallelLinear( @@ -374,11 +390,15 @@ class FeedForward(nn.Module): input_is_parallel=True, return_bias=False, quant_config=quant_config, - prefix=prefix, + prefix=_join_prefix(prefix, "w2"), ) def forward(self, x): - return self.w2(self.act(self.w13(x))) + hidden_states = self.w13(x) + hidden_states = _restore_linear_output_shape(hidden_states, x) + hidden_states = self.act(hidden_states) + hidden_states = self.w2(hidden_states) + return _restore_linear_output_shape(hidden_states, x) class ZImageTransformerBlock(nn.Module): @@ -392,6 +412,7 @@ class ZImageTransformerBlock(nn.Module): qk_norm: bool, modulation=True, quant_config: "QuantizationConfig | None" = None, + prefix: str = "", ): super().__init__() self.dim = dim @@ -403,10 +424,14 @@ class ZImageTransformerBlock(nn.Module): qk_norm=qk_norm, eps=1e-5, quant_config=quant_config, + prefix=_join_prefix(prefix, "attention"), ) self.feed_forward = FeedForward( - dim=dim, hidden_dim=int(dim / 3 * 8), quant_config=quant_config, prefix="feed_forward" + dim=dim, + hidden_dim=int(dim / 3 * 8), + quant_config=quant_config, + prefix=_join_prefix(prefix, "feed_forward"), ) self.layer_id = layer_id @@ -711,6 +736,7 @@ class ZImageTransformer2DModel(CachedTransformer): qk_norm, modulation=True, quant_config=quant_config, + prefix=f"noise_refiner.{layer_id}", ) for layer_id in range(n_refiner_layers) ] @@ -726,6 +752,7 @@ class ZImageTransformer2DModel(CachedTransformer): qk_norm, modulation=False, quant_config=quant_config, + prefix=f"context_refiner.{layer_id}", ) for layer_id in range(n_refiner_layers) ] @@ -757,6 +784,7 @@ class ZImageTransformer2DModel(CachedTransformer): norm_eps, qk_norm, quant_config=quant_config, + prefix=f"layers.{layer_id}", ) for layer_id in range(n_layers) ] diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 5140e89b7..72ea62045 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -11,9 +11,9 @@ to DiffusionModelRunner. import gc import multiprocessing as mp import os -from collections.abc import Iterable -from contextlib import AbstractContextManager, nullcontext -from types import SimpleNamespace +from collections.abc import Iterable, Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from dataclasses import dataclass from typing import Any import torch @@ -22,7 +22,7 @@ from vllm.config import CompilationConfig, DeviceConfig, VllmConfig, set_current from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.logger import init_logger from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler -from vllm.transformers_utils.config import get_config, get_hf_text_config +from vllm.transformers_utils.config import get_hf_text_config from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.mem_utils import GiB_bytes from vllm.v1.worker.workspace import init_workspace_manager @@ -54,6 +54,78 @@ from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory logger = init_logger(__name__) +@dataclass +class _DiffusionVllmModelConfig: + model: str + dtype: torch.dtype + quantization: str | None = None + quantization_config: Any | None = None + hf_config: Any | None = None + hf_text_config: Any | None = None + multimodal_config: Any | None = None + enforce_eager: bool = False + disable_cascade_attn: bool = False + enable_return_routed_experts: bool = False + is_moe: bool = False + + def is_quantized(self) -> bool: + return self.quantization is not None + + def is_model_moe(self) -> bool: + return self.is_moe + + def is_nvfp4_quantized(self) -> bool: + return self.quantization == "modelopt_fp4" + + +def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig: + quant_config = getattr(od_config, "quantization_config", None) + quantization = quant_config.get_name() if quant_config is not None and hasattr(quant_config, "get_name") else None + hf_config = getattr(od_config, "tf_model_config", None) + hf_text_config = get_hf_text_config(hf_config) if hasattr(hf_config, "get_text_config") else hf_config + return _DiffusionVllmModelConfig( + model=od_config.model, + dtype=od_config.dtype, + quantization=quantization, + quantization_config=quant_config, + hf_config=hf_config, + hf_text_config=hf_text_config, + enforce_eager=getattr(od_config, "enforce_eager", False), + is_moe=bool(getattr(od_config, "is_moe", False)), + ) + + +@contextmanager +def _force_cutlass_fp8_linear_kernel(quant_config: object | None) -> Iterator[None]: + from vllm.model_executor.layers.quantization import modelopt as vllm_modelopt + + linear_method_cls = getattr(quant_config, "LinearMethodCls", None) + if linear_method_cls in { + vllm_modelopt.ModelOptFp8LinearMethod, + vllm_modelopt.ModelOptFp8PcPtLinearMethod, + }: + from vllm.platforms import current_platform + + if current_platform.is_cuda() and current_platform.has_device_capability(89): + from vllm.model_executor.kernels.linear import CutlassFP8ScaledMMLinearKernel + + original_init_fp8_linear_kernel = vllm_modelopt.init_fp8_linear_kernel + + def init_fp8_linear_kernel_with_cutlass(*args: Any, **kwargs: Any) -> Any: + kwargs.setdefault("force_kernel", CutlassFP8ScaledMMLinearKernel) + return original_init_fp8_linear_kernel(*args, **kwargs) + + vllm_modelopt.init_fp8_linear_kernel = init_fp8_linear_kernel_with_cutlass + logger.info("Using CUTLASS FP8 linear kernels for this ModelOpt FP8 diffusion stage.") + try: + yield + finally: + vllm_modelopt.init_fp8_linear_kernel = original_init_fp8_linear_kernel + return + + yield + + class DiffusionWorker: """ A worker that manages GPU infrastructure and delegates to the model runner. @@ -124,19 +196,7 @@ class DiffusionWorker: vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel vllm_config.profiler_config = self.od_config.profiler_config - try: - hf_config = get_config(self.od_config.model, trust_remote_code=self.od_config.trust_remote_code) - except ValueError: - hf_config = None - logger.info("Skipping hf_config loading for diffusion model %r", self.od_config.model_class_name) - hf_text_config = get_hf_text_config(hf_config) if hf_config is not None else None - vllm_config.model_config = SimpleNamespace( - hf_config=hf_config, - hf_text_config=hf_text_config, - enforce_eager=self.od_config.enforce_eager, - dtype=self.od_config.dtype, - enable_return_routed_experts=False, - ) + vllm_config.model_config = _make_diffusion_vllm_model_config(self.od_config) # type: ignore[assignment] vllm_config.quant_config = self.od_config.quantization_config # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping. # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion. @@ -191,9 +251,15 @@ class DiffusionWorker: """Load the diffusion model using DiffusionModelRunner.""" load_format = kwargs.get("load_format", load_format) custom_pipeline_name = kwargs.get("custom_pipeline_name", custom_pipeline_name) + cutlass_fp8_context = ( + _force_cutlass_fp8_linear_kernel(self.od_config.quantization_config) + if getattr(self.od_config, "force_cutlass_fp8", False) + else nullcontext() + ) with ( set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config), set_current_vllm_config(self.vllm_config), + cutlass_fp8_context, ): self.model_runner.load_model( memory_pool_context_fn=self._maybe_get_memory_pool_context, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 1ce4f8d82..f01094bef 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -140,6 +140,7 @@ class OmniEngineArgs(EngineArgs): async_chunk: bool = False omni_kv_config: dict | None = None quantization_config: Any | None = None + force_cutlass_fp8: bool | None = None worker_type: str | None = None task_type: str | None = None worker_cls: str = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index e670a52d8..611f48bdc 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1467,6 +1467,7 @@ class AsyncOmniEngine: "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), "quantization": kwargs.get("quantization", None), + "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)), "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), "enable_ar_profiler": kwargs.get("enable_ar_profiler", False), **( diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 9f9eef2ba..6880cf7c5 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -338,6 +338,16 @@ class OmniServeCommand(CLISubcommand): 'Example: \'{"method":"gguf","gguf_model":"/path/to/model.gguf"}\'.' ), ) + omni_config_group.add_argument( + "--force-cutlass-fp8", + action="store_true", + default=None, + help=( + "Diffusion-only runtime override for ModelOpt FP8 checkpoints: " + "force CUTLASS FP8 linear kernels on CUDA SM89+ devices. " + "Ignored for BF16, non-ModelOpt FP8, ROCm, and older CUDA GPUs." + ), + ) # HSDP (Hybrid Sharded Data Parallel) parameters omni_config_group.add_argument( diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index 75e4adb7e..88494eda4 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -132,7 +132,7 @@ class HunyuanModel(HunYuanModel): # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) fused_moe_expert_mapping = fused_moe_make_expert_params_mapping( - self, + model=self, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", diff --git a/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py b/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py index 730472e99..ca7d00f50 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py @@ -416,10 +416,8 @@ class BailingMoeV2SparseMoeBlock(nn.Module): dim=-1, ) - # FusedMoE expects 2D hidden_states. hidden_states_2d = hidden_states.view(-1, h) final_hidden_states = self.experts(hidden_states_2d, packed_routing) - final_hidden_states = final_hidden_states.view(bsz, seq_len, h) return final_hidden_states.squeeze(0) if input_is_2d else final_hidden_states diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index 5af47a1f7..de9e6033b 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -60,17 +60,83 @@ _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { "int8": _build_int8, "inc": _build_inc, "auto-round": _build_inc, + "auto_round": _build_inc, } SUPPORTED_QUANTIZATION_METHODS: list[str] = list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys()))) +_MODEL_OPT_METHODS = { + "modelopt", +} +_MODEL_OPT_FP8_ALGOS = { + "FP8", + "FP8_PER_CHANNEL_PER_TOKEN", +} + + +def _normalize_method_name(method: Any) -> str: + return str(method).lower().replace("-", "_") + + +def _detect_modelopt_method(config: Mapping[str, Any]) -> str | None: + quantization = config.get("quantization") + if isinstance(quantization, Mapping): + quant_algo = str(quantization.get("quant_algo", "")).upper() + else: + quant_algo = str(config.get("quant_algo", "")).upper() + + method = config.get("method", config.get("quant_method")) + normalized_method = _normalize_method_name(method) if method is not None else None + + producer = config.get("producer") + is_modelopt_config = normalized_method in _MODEL_OPT_METHODS or ( + isinstance(producer, Mapping) and str(producer.get("name", "")).lower() == "modelopt" + ) + + if not is_modelopt_config: + return None + + if quant_algo: + if quant_algo in _MODEL_OPT_FP8_ALGOS: + return "modelopt" + return None + + if method is not None: + if normalized_method in _MODEL_OPT_METHODS: + return normalized_method + + return None + + +def _build_modelopt_from_config(method: str, config: Mapping[str, Any]) -> QuantizationConfig: + config_cls = get_quantization_config(method) + normalized_config = dict(config) + normalized_config.setdefault("quant_method", method) + return config_cls.from_config(normalized_config) + + +def _pop_method_name(spec: dict[str, Any]) -> str | None: + method = spec.pop("method", None) + if method is None: + method = spec.pop("quant_method", None) + return method + + +def _build_from_method_and_config(method: str, config: Mapping[str, Any]) -> QuantizationConfig: + normalized_config = {"quant_method": method, **config} + modelopt_method = _detect_modelopt_method(normalized_config) + if modelopt_method is not None: + return _build_modelopt_from_config(modelopt_method, normalized_config) + return _build_single(method, **config) + + def _build_single(method: str, **kwargs: Any) -> QuantizationConfig: """Build a single QuantizationConfig by method name. Resolution: _OVERRIDES first, then vLLM registry via from_config(). """ - method = method.lower() + method = _normalize_method_name(method) if method in _OVERRIDES: return _OVERRIDES[method](**kwargs) @@ -92,16 +158,17 @@ def _build_single(method: str, **kwargs: Any) -> QuantizationConfig: def _is_per_component_dict(spec: dict[str, Any]) -> bool: """Check if a dict describes per-component quantization. - A per-component dict has no "method" key and all values are + A per-component dict has no "method" / "quant_method" key and all values are str, dict, or None. To avoid misdetecting a flat config with all-string values (e.g. {"activation_scheme": "static"}), we - require at least one value to be None or a dict with "method". + require at least one value to be None or a dict with "method" / + "quant_method". """ - if "method" in spec: + if "method" in spec or "quant_method" in spec: return False if not all(isinstance(v, (dict, str, type(None))) for v in spec.values()): return False - return any(v is None or (isinstance(v, dict) and "method" in v) for v in spec.values()) + return any(v is None or (isinstance(v, dict) and ("method" in v or "quant_method" in v)) for v in spec.values()) def _build_component_config(spec: dict[str, Any]) -> ComponentQuantizationConfig: @@ -116,10 +183,10 @@ def _build_component_config(spec: dict[str, Any]) -> ComponentQuantizationConfig config = _build_single(value) elif isinstance(value, dict): value = dict(value) # avoid mutating caller's dict - method = value.pop("method", None) + method = _pop_method_name(value) if method is None: - raise ValueError(f"Component '{prefix}' config dict must have a 'method' key") - config = _build_single(method, **value) + raise ValueError(f"Component '{prefix}' config dict must have a 'method' or 'quant_method' key") + config = _build_from_method_and_config(method, value) else: raise TypeError(f"Component '{prefix}' config must be str, dict, or None, got {type(value).__name__}") @@ -164,14 +231,19 @@ def build_quant_config( if _is_per_component_dict(spec): return _build_component_config(spec) - method = spec.pop("method", None) + modelopt_method = _detect_modelopt_method(spec) + if modelopt_method is not None: + logger.info("Building quantization config: %s", modelopt_method) + return _build_modelopt_from_config(modelopt_method, spec) + + method = _pop_method_name(spec) if method is None: raise ValueError( - "Dict quantization config must have a 'method' key or " + "Dict quantization config must have a 'method' or 'quant_method' key or " "be a per-component config with component prefixes as keys." ) merged = {**spec, **kwargs} logger.info("Building quantization config: %s", method) - return _build_single(method, **merged) + return _build_from_method_and_config(method, merged) raise TypeError(f"quantization config must be str, dict, QuantizationConfig, or None, got {type(spec).__name__}")