From ef71d74c81b1bfd20e16a9ef1492618bed6d9f85 Mon Sep 17 00:00:00 2001 From: Yi Liu Date: Mon, 1 Jun 2026 12:06:02 +0800 Subject: [PATCH] [Bugfix] Fix FLUX W4A16/AutoRound quant_config propagation (#3587) Signed-off-by: yiliu30 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...est_flux1_dev_autoround_w4a16_expansion.py | 27 ++++++++++++++--- .../diffusion/models/flux/flux_transformer.py | 30 +++++++++++++------ 2 files changed, 44 insertions(+), 13 deletions(-) diff --git a/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py index 1c81d5a4c..a9b24ef9d 100644 --- a/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py @@ -26,10 +26,6 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform -pytestmark = pytest.mark.skip( - reason="Current Hugging Face account is not yet authorized to access this model.", -) - QUANTIZED_MODEL = "vllm-project-org/FLUX.1-dev-AutoRound-w4a16" BASELINE_MODEL = "black-forest-labs/FLUX.1-dev" @@ -63,6 +59,27 @@ def _first_request_images(outputs) -> list: return req_out.images +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}) +def test_flux_autoround_w4a16_load(): + """Load the W4A16 quantized FLUX model and run a minimal generation. + + Verifies that the quantized checkpoint loads end-to-end (quant_config + propagated to all transformer blocks) and produces a valid 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)" + + def _generate_image(model_name: str, **extra_kwargs) -> tuple[list, float]: """Load a FLUX model, generate one image, return (images, peak_memory_mb).""" gc.collect() @@ -87,6 +104,7 @@ def _generate_image(model_name: str, **extra_kwargs) -> tuple[list, float]: return images, peak +@pytest.mark.skip(reason="BF16 baseline model cannot be loaded correctly yet.") @pytest.mark.full_model @pytest.mark.diffusion @hardware_test(res={"cuda": "L4"}) @@ -105,6 +123,7 @@ def test_flux_autoround_w4a16_generates_image(omni_runner: OmniRunner): assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)" +@pytest.mark.skip(reason="BF16 baseline model cannot be loaded correctly yet.") @pytest.mark.full_model @pytest.mark.diffusion @hardware_test(res={"cuda": "L4"}) diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index dd88ee76c..88d8126c8 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/vllm_omni/diffusion/models/flux/flux_transformer.py @@ -29,6 +29,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.layers.adalayernorm import ( @@ -41,6 +42,23 @@ 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, @@ -381,9 +399,7 @@ class FluxSingleTransformerBlock(nn.Module): super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) - # Modulation linear kept full precision; shift/scale/gate outputs - # are multiplied into the residual stream every block (see #2728). - self.norm = AdaLayerNormZeroSingle(dim, quant_config=None, 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, @@ -565,16 +581,13 @@ class FluxTransformer2DModel(nn.Module): self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim) self.x_embedder = nn.Linear(in_channels, self.inner_dim) - # Dual-stream blocks kept full precision — FP8 on their joint - # attention path causes noise on FLUX (#2728). Single-stream - # blocks (38 vs 19) still get FP8 for memory savings. self.transformer_blocks = nn.ModuleList( [ FluxTransformerBlock( dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, - quant_config=None, + quant_config=_safe_quant_config(quant_config), prefix=f"transformer_blocks.{i}", ) for i in range(num_layers) @@ -594,13 +607,12 @@ class FluxTransformer2DModel(nn.Module): ] ) - # Final modulation feeds proj_out; keep full precision (see #2728). self.norm_out = AdaLayerNormContinuous( self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6, - quant_config=None, + 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)