From 99234881f5da5c5442fdeae85eaf013be77ac3e2 Mon Sep 17 00:00:00 2001 From: Yi Liu Date: Tue, 30 Jun 2026 16:27:28 +0800 Subject: [PATCH] [Quant][AutoRound] Add W4A16 support for Qwen-Image (#3588) Signed-off-by: yiliu30 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...st_qwen_image_autoround_w4a16_expansion.py | 101 +++++++++++++++++ .../diffusion/models/flux/flux_transformer.py | 24 +--- .../models/qwen_image/pipeline_qwen_image.py | 4 +- .../qwen_image/pipeline_qwen_image_edit.py | 14 ++- .../pipeline_qwen_image_edit_plus.py | 6 +- .../qwen_image/pipeline_qwen_image_layered.py | 6 +- .../qwen_image/qwen_image_transformer.py | 106 +++++++++++------- vllm_omni/quantization/component_config.py | 24 ++++ 8 files changed, 215 insertions(+), 70 deletions(-) create mode 100644 tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py diff --git a/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py new file mode 100644 index 000000000..8d7d31882 --- /dev/null +++ b/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E test for Qwen-Image AutoRound W4A16 quantized inference. + +Verifies that the W4A16 quantized Qwen-Image checkpoint loads end-to-end +(quant_config propagated correctly) and produces a valid image. +""" + +import os +from pathlib import Path + +import numpy as np +import pytest +import torch + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + +QUANTIZED_MODEL = os.environ.get( + "QWEN_IMAGE_AUTOROUND_MODEL", + "INC4AI/Qwen-Image-AutoRound-W4A16", +) + +HEIGHT = 512 +WIDTH = 512 +NUM_STEPS = 2 + + +def _sampling_params() -> OmniDiffusionSamplingParams: + return OmniDiffusionSamplingParams( + height=HEIGHT, + width=WIDTH, + num_inference_steps=NUM_STEPS, + true_cfg_scale=1.0, + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), + ) + + +def _first_request_images(outputs) -> list: + first_output = outputs[0] + assert first_output.final_output_type == "image" + req_out = first_output.request_output + assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images") + return req_out.images + + +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}) +def test_qwen_image_autoround_w4a16_load(): + """Load the W4A16 quantized Qwen-Image model and run a minimal generation. + + Verifies that quant_config is propagated to all transformer blocks + and produces a valid, non-blank image. + """ + with OmniRunner(QUANTIZED_MODEL, enforce_eager=True) as runner: + outputs = runner.omni.generate( + "a cup of coffee on a table", + _sampling_params(), + ) + images = _first_request_images(outputs) + assert len(images) >= 1, "Expected at least one generated image" + img = images[0] + assert img.width == WIDTH and img.height == HEIGHT + arr = np.array(img) + assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)" + + +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}) +def test_qwen_image_autoround_w4a16_generate(): + """Full generation: 512×512, 20 steps, CFG=5.0. + + Validates that realistic quantized inference (multi-step with CFG + negative prompts and longer denoising) produces a valid non-blank image. + """ + params = OmniDiffusionSamplingParams( + height=HEIGHT, + width=WIDTH, + num_inference_steps=20, + true_cfg_scale=5.0, + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), + ) + with OmniRunner(QUANTIZED_MODEL, enforce_eager=True) as runner: + outputs = runner.omni.generate( + "a cup of coffee on a wooden table, morning sunlight, photorealistic", + params, + ) + images = _first_request_images(outputs) + assert len(images) >= 1, "Expected at least one generated image" + img = images[0] + + # Save for manual inspection + output_path = Path(__file__).resolve().parents[4] / "ar_coffee_full.png" + img.save(str(output_path)) + + assert img.width == WIDTH and img.height == HEIGHT + arr = np.array(img) + assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)" diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index 9d89368be..952ee9eb4 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/vllm_omni/diffusion/models/flux/flux_transformer.py @@ -27,6 +27,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig +from vllm_omni.quantization.component_config import safe_quant_config if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -44,23 +45,6 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding, apply_rope_to_qk logger = init_logger(__name__) -def _safe_quant_config(quant_config: "QuantizationConfig | None") -> "QuantizationConfig | None": - """Return quant_config only if it is safe to propagate here, else None. - - Dual-stream transformer_blocks, norm modulation layers, and norm_out are - kept at full precision for FP8 (see #2728). Offline quantization (e.g. - INC/AutoRound W4A16) needs the config propagated so packed weights load - correctly. - """ - if quant_config is None: - return None - from vllm.model_executor.layers.quantization.inc import INCConfig - - if isinstance(quant_config, INCConfig): - return quant_config - return None - - class ColumnParallelApproxGELU(nn.Module): def __init__( self, @@ -401,7 +385,7 @@ class FluxSingleTransformerBlock(nn.Module): super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) - self.norm = AdaLayerNormZeroSingle(dim, quant_config=_safe_quant_config(quant_config), prefix=f"{prefix}.norm") + self.norm = AdaLayerNormZeroSingle(dim, quant_config=safe_quant_config(quant_config), prefix=f"{prefix}.norm") self.proj_mlp = ReplicatedLinear( dim, self.mlp_hidden_dim, @@ -596,7 +580,7 @@ class FluxTransformer2DModel(nn.Module): dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, - quant_config=_safe_quant_config(quant_config), + quant_config=safe_quant_config(quant_config), prefix=f"transformer_blocks.{i}", ) for i in range(num_layers) @@ -621,7 +605,7 @@ class FluxTransformer2DModel(nn.Module): self.inner_dim, elementwise_affine=False, eps=1e-6, - quant_config=_safe_quant_config(quant_config), + quant_config=safe_quant_config(quant_config), prefix="norm_out", ) self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 974436b21..43cd4bdc4 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -332,7 +332,9 @@ class QwenImagePipeline( ).to(self.device) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) self.transformer = QwenImageTransformer2DModel( - od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs + od_config=od_config, + quant_config=od_config.quantization_config, + **transformer_kwargs, ) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index 41475c65d..e594d5ad8 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py @@ -36,7 +36,6 @@ from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import calculate_ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( QwenImageTransformer2DModel, ) -from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.prompt_utils import ( @@ -257,6 +256,7 @@ class QwenImageEditPipeline( prefetch_subfolders( model, qwen_subfolders, + local_files_only=local_files_only, ) self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( @@ -282,7 +282,11 @@ class QwenImageEditPipeline( local_files_only=local_files_only, ).to(self.device) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) - self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs) + self.transformer = QwenImageTransformer2DModel( + od_config=od_config, + quant_config=od_config.quantization_config, + **transformer_kwargs, + ) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.processor = from_pretrained_with_prefetch( Qwen2VLProcessor.from_pretrained, @@ -846,8 +850,10 @@ class QwenImageEditPipeline( if self.attention_kwargs is None: self._attention_kwargs = {} - txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds) - negative_txt_seq_lens = txt_seq_lens_from_embeds(negative_prompt_embeds) + txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None + negative_txt_seq_lens = ( + negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None + ) latents = self.diffuse( prompt_embeds, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index 17f4a8323..bf63805b3 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -250,7 +250,11 @@ class QwenImageEditPlusPipeline( ).to(self.device) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) - self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs) + self.transformer = QwenImageTransformer2DModel( + od_config=od_config, + quant_config=od_config.quantization_config, + **transformer_kwargs, + ) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.processor = from_pretrained_with_prefetch( Qwen2VLProcessor.from_pretrained, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index 10363f42b..f7a9fc4b1 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py @@ -275,7 +275,11 @@ class QwenImageLayeredPipeline( ] transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) - self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs) + self.transformer = QwenImageTransformer2DModel( + od_config=od_config, + quant_config=od_config.quantization_config, + **transformer_kwargs, + ) # Pipeline configuration & processing parameters self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index 1070f1369..07015a4f4 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -17,6 +17,7 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, QKVParallelLinear, @@ -25,6 +26,8 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.quantization.component_config import safe_quant_config + if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, @@ -48,8 +51,23 @@ from vllm_omni.diffusion.layers.rope import RotaryEmbedding logger = init_logger(__name__) -def _join_prefix(prefix: str, suffix: str) -> str: - return f"{prefix}.{suffix}" if prefix else suffix +def _normalize_qwen_image_weight_name(name: str) -> str: + name = name.removeprefix("transformer.") + if ".to_out.0." in name: + name = name.replace(".to_out.0.", ".to_out.") + return name + + +def _resolve_qwen_image_lookup_name( + name: str, + stacked_params_mapping: list[tuple[str, str, str]], +) -> tuple[str, str | None]: + lookup_name = _normalize_qwen_image_weight_name(name) + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in lookup_name or param_name in lookup_name: + continue + return lookup_name.replace(weight_name, param_name), shard_id + return lookup_name, None class ImageRopePrepare(nn.Module): @@ -444,7 +462,7 @@ class ColumnParallelApproxGELU(nn.Module): gather_output=False, return_bias=False, quant_config=quant_config, - prefix=prefix, + prefix=f"{prefix}.proj" if prefix else "proj", ) self.approximate = approximate @@ -479,7 +497,7 @@ class FeedForward(nn.Module): approximate="tanh", bias=bias, quant_config=quant_config, - prefix=_join_prefix(prefix, "net.0.proj"), + prefix=f"{prefix}.net.0", ), nn.Identity(), # placeholder for weight loading RowParallelLinear( @@ -488,7 +506,7 @@ class FeedForward(nn.Module): input_is_parallel=True, return_bias=False, quant_config=quant_config, - prefix=_join_prefix(prefix, "net.2"), + prefix=f"{prefix}.net.2", ), ] @@ -532,13 +550,13 @@ class QwenImageCrossAttention(nn.Module): head_size=self.head_dim, total_num_heads=num_heads, quant_config=quant_config, - prefix=_join_prefix(prefix, "to_qkv"), + prefix=f"{prefix}.to_qkv", ) self.query_num_heads = self.to_qkv.num_heads self.kv_num_heads = self.to_qkv.num_kv_heads - self.norm_q = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() - self.norm_k = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() + self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity() self.inner_dim = out_dim if out_dim is not None else head_dim * self.total_num_heads @@ -548,7 +566,7 @@ class QwenImageCrossAttention(nn.Module): head_size=head_dim, total_num_heads=num_heads, quant_config=quant_config, - prefix=_join_prefix(prefix, "add_kv_proj"), + prefix=f"{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 @@ -561,7 +579,7 @@ class QwenImageCrossAttention(nn.Module): input_is_parallel=True, return_bias=False, quant_config=quant_config, - prefix=_join_prefix(prefix, "to_add_out"), + prefix=f"{prefix}.to_add_out", ) assert not pre_only @@ -572,11 +590,11 @@ class QwenImageCrossAttention(nn.Module): input_is_parallel=True, return_bias=False, quant_config=quant_config, - prefix=_join_prefix(prefix, "to_out"), + prefix=f"{prefix}.to_out.0", ) - self.norm_added_q = nn.RMSNorm(head_dim, eps=eps) - self.norm_added_k = nn.RMSNorm(head_dim, eps=eps) + self.norm_added_q = RMSNorm(head_dim, eps=eps) + self.norm_added_k = RMSNorm(head_dim, eps=eps) self.attn = Attention( num_heads=self.query_num_heads, @@ -713,22 +731,19 @@ class QwenImageTransformerBlock(nn.Module): self.dim = dim self.num_attention_heads = num_attention_heads self.attention_head_dim = attention_head_dim + txt_mod_quant_config = safe_quant_config(quant_config) # Image processing modules. - # Modulation linear is kept unquantized (quant_config=None) — it - # produces shift/scale/gate values that are precision-sensitive - # (see #2728). Use column TP with gather_output=True so weights are - # sharded while downstream modulation still receives full [B, 6 * dim]. + # The re-quantized W4A16 checkpoint keeps img_mod.1 in full precision. self.img_mod = nn.Sequential( nn.SiLU(), - ColumnParallelLinear( + ReplicatedLinear( dim, 6 * dim, bias=True, - gather_output=True, return_bias=False, quant_config=None, - prefix=_join_prefix(prefix, "img_mod.1"), + prefix=f"{prefix}.img_mod.1", ), ) self.img_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) @@ -739,27 +754,27 @@ class QwenImageTransformerBlock(nn.Module): context_pre_only=False, head_dim=attention_head_dim, quant_config=quant_config, - prefix=_join_prefix(prefix, "attn"), + prefix=f"{prefix}.attn", ) self.img_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) self.img_mlp = FeedForward( dim=dim, dim_out=dim, quant_config=quant_config, - prefix=_join_prefix(prefix, "img_mlp"), + prefix=f"{prefix}.img_mlp", ) # Text processing modules. + # AutoRound keeps txt_mod.1 quantized inside transformer_blocks. self.txt_mod = nn.Sequential( nn.SiLU(), - ColumnParallelLinear( + ReplicatedLinear( dim, 6 * dim, bias=True, - gather_output=True, return_bias=False, - quant_config=None, - prefix=_join_prefix(prefix, "txt_mod.1"), + quant_config=txt_mod_quant_config, + prefix=f"{prefix}.txt_mod.1", ), ) self.txt_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) @@ -769,7 +784,7 @@ class QwenImageTransformerBlock(nn.Module): dim=dim, dim_out=dim, quant_config=quant_config, - prefix=_join_prefix(prefix, "txt_mlp"), + prefix=f"{prefix}.txt_mlp", ) self.zero_cond_t = zero_cond_t @@ -979,6 +994,7 @@ class QwenImageTransformer2DModel(CachedTransformer): self.out_channels = out_channels or in_channels self.inner_dim = num_attention_heads * attention_head_dim self.guidance_embeds = guidance_embeds + self.quant_config = quant_config if not use_layer3d_rope: self.pos_embed = QwenEmbedRope(theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True) @@ -991,7 +1007,7 @@ class QwenImageTransformer2DModel(CachedTransformer): quant_config=quant_config, ) - self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6) + self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6) # Entry projections (image/text) are kept full precision — # small sensitive layers at the network boundary (see #2728). @@ -1225,22 +1241,26 @@ class QwenImageTransformer2DModel(CachedTransformer): loaded_params: set[str] = set() for name, loaded_weight in weights: - original_name = name - lookup_name = name - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in original_name or param_name in original_name: - continue - lookup_name = original_name.replace(weight_name, param_name) - param = params_dict[lookup_name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - if lookup_name not in params_dict and ".to_out.0." in lookup_name: - lookup_name = lookup_name.replace(".to_out.0.", ".to_out.") - param = params_dict[lookup_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) + original_name = name.removeprefix("transformer.") + lookup_name, shard_id = _resolve_qwen_image_lookup_name( + original_name, + stacked_params_mapping, + ) + + if lookup_name.endswith(".bias") and lookup_name not in params_dict: + continue + + param = params_dict.get(lookup_name) + if param is None: + logger.debug("Skipping unexpected Qwen-Image transformer weight %s", original_name) + continue + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if shard_id is None: weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(original_name) loaded_params.add(lookup_name) return loaded_params diff --git a/vllm_omni/quantization/component_config.py b/vllm_omni/quantization/component_config.py index 25d1abf1e..87564ef46 100644 --- a/vllm_omni/quantization/component_config.py +++ b/vllm_omni/quantization/component_config.py @@ -48,6 +48,30 @@ def resolve_encoder_quant_config( return quant_config +def safe_quant_config( + quant_config: QuantizationConfig | None, +) -> QuantizationConfig | None: + """Return *quant_config* only if it is safe for norm/modulation layers. + + Norm and modulation layers (LayerNorm, RMSNorm, AdaLayerNorm, img_mod, + txt_mod, etc.) produce precision-sensitive shift/scale/gate values and + should not receive FP8 quant configs (see #2728). Pre-quantized methods + like INC/AutoRound W4A16 need the config propagated so packed weights + load correctly. + + This is the inverse of :func:`resolve_encoder_quant_config`: that function + strips pre-quantized configs from encoders, while this one strips + *every config except* pre-quantized configs from norm/mod layers. + """ + if quant_config is None: + return None + from vllm.model_executor.layers.quantization.inc import INCConfig + + if isinstance(quant_config, INCConfig): + return quant_config + return None + + class ComponentQuantizationConfig(QuantizationConfig): """Routes quantization to different configs by layer prefix."""