From 6cd1ff7e645030e6e76adbe2ee2b15cf8bda968a Mon Sep 17 00:00:00 2001 From: Diya Peng <50388704+dpeng123@users.noreply.github.com> Date: Mon, 20 Jul 2026 01:16:31 -0700 Subject: [PATCH] [Quantization] Add FLUX.2-dev Mistral TE online FP8 support (#5136) Signed-off-by: diyapeng <50388704+diyapeng@users.noreply.github.com> Co-authored-by: diyapeng <50388704+diyapeng@users.noreply.github.com> --- docs/user_guide/quantization/fp8.md | 1 + .../flux2/test_pipeline_flux2_quantization.py | 50 +++++++++++++++++ .../test_mistral_encoder_tp.py | 55 +++++++++++++++++++ .../quantization/test_quantization_quality.py | 25 ++++++++- .../diffusion/models/flux2/pipeline_flux2.py | 37 ++++++++++++- .../models/mistral_encoder/mistral_encoder.py | 25 ++++++++- 6 files changed, 185 insertions(+), 8 deletions(-) create mode 100644 tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index 6114812db..c76cfd32e 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -88,6 +88,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)]) | Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | | | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ | | FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | | +| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ | | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None | | | HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | | | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | | diff --git a/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py b/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py new file mode 100644 index 000000000..3900c0c94 --- /dev/null +++ b/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py @@ -0,0 +1,50 @@ +from unittest.mock import Mock, patch + +import pytest +import torch + +from vllm_omni.diffusion.models.flux2.pipeline_flux2 import ( + Flux2Pipeline, + _resolve_component_quant_config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_resolve_component_quant_config_routes_components(): + text_encoder_config = object() + component_config = Mock() + component_config.resolve.side_effect = { + "text_encoder": text_encoder_config, + "transformer": None, + }.get + + assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config + assert _resolve_component_quant_config(component_config, "transformer") is None + + +def test_resolve_component_quant_config_preserves_global_config(): + global_config = object() + assert _resolve_component_quant_config(global_config, "transformer") is global_config + assert _resolve_component_quant_config(None, "text_encoder") is None + + +def test_offload_transformer_before_text_encoder_weights(): + pipeline = object.__new__(Flux2Pipeline) + pipeline.transformer = Mock() + weights = [ + ("transformer.layer.weight", torch.empty(1)), + ("text_encoder.layer.weight", torch.empty(1)), + ("text_encoder.layer.bias", torch.empty(1)), + ] + + iterator = iter(pipeline._offload_transformer_before_text_encoder(weights)) + assert next(iterator)[0] == "transformer.layer.weight" + pipeline.transformer.to.assert_not_called() + + with patch("vllm_omni.diffusion.models.flux2.pipeline_flux2.current_omni_platform.empty_cache") as empty_cache: + assert next(iterator)[0] == "text_encoder.layer.weight" + assert next(iterator)[0] == "text_encoder.layer.bias" + + pipeline.transformer.to.assert_called_once_with("cpu") + empty_cache.assert_called_once_with() diff --git a/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py b/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py index 27d1e757e..20024abe7 100644 --- a/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py +++ b/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py @@ -5,6 +5,7 @@ from types import SimpleNamespace import pytest import torch from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm_omni.diffusion.models.mistral_encoder.mistral_encoder import ( MistralEncoderModel, @@ -491,3 +492,57 @@ class TestComputeLogits: assert logits.dim() == 3 assert logits.shape[0] == 1 assert logits.shape[1] == 4 + + +class TestQuantConfig: + """Verify quant_config is threaded into weight linears only, with correct prefixes.""" + + def test_quant_config_on_weight_linears(self, mocker): + quant_config = mocker.MagicMock(name="QuantizationConfig") + # Keep construction on the CPU-safe unquantized path while still + # verifying LinearBase consults the supplied quant_config. + quant_config.get_quant_method.return_value = UnquantizedLinearMethod() + + config = _make_config(num_hidden_layers=1) + model = MistralEncoderModel(config, prefix="text_encoder", quant_config=quant_config) + + layer = model.language_model.model.layers[0] + for linear in ( + layer.self_attn.qkv_proj, + layer.self_attn.o_proj, + layer.mlp.gate_up_proj, + layer.mlp.down_proj, + ): + assert linear.quant_config is quant_config + + # get_quant_method should be consulted for the four weight linears. + assert quant_config.get_quant_method.call_count == 4 + seen_prefixes = set() + for call in quant_config.get_quant_method.call_args_list: + if "prefix" in call.kwargs: + seen_prefixes.add(call.kwargs["prefix"]) + elif len(call.args) > 1: + seen_prefixes.add(call.args[1]) + assert any(p.endswith("self_attn.qkv_proj") for p in seen_prefixes) + assert any(p.endswith("self_attn.o_proj") for p in seen_prefixes) + assert any(p.endswith("mlp.gate_up_proj") for p in seen_prefixes) + assert any(p.endswith("mlp.down_proj") for p in seen_prefixes) + + # Embeddings / lm_head stay unquantized (quant_config never passed). + embed = model.language_model.model.embed_tokens + lm_head = model.language_model.lm_head + assert getattr(embed, "quant_config", None) is None + assert getattr(lm_head, "quant_config", None) is None + + def test_layer_prefixes_include_text_encoder(self): + config = _make_config(num_hidden_layers=2) + model = MistralEncoderModel(config, prefix="text_encoder") + + layer0 = model.language_model.model.layers[0] + assert layer0.self_attn.qkv_proj.prefix.startswith("text_encoder.language_model.model.layers.0") + assert layer0.self_attn.o_proj.prefix.startswith("text_encoder.language_model.model.layers.0") + assert layer0.mlp.gate_up_proj.prefix.startswith("text_encoder.language_model.model.layers.0") + assert layer0.mlp.down_proj.prefix.startswith("text_encoder.language_model.model.layers.0") + + layer1 = model.language_model.model.layers[1] + assert "layers.1" in layer1.self_attn.qkv_proj.prefix diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index 5f0dc4d62..9b53d8941 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -82,6 +82,7 @@ class QualityTestConfig: gpu: str = "H100" # minimum GPU requirement negative_prompt: str = "" guidance_scale: float | None = None + enable_cpu_offload: bool = False def baseline_ref(self) -> str: return self.baseline_model or self.model or "" @@ -181,6 +182,18 @@ QUALITY_CONFIGS = [ max_lpips=0.20, num_inference_steps=10, ), + QualityTestConfig( + id="fp8_flux2_dev_text_encoder", + model="black-forest-labs/FLUX.2-dev", + quantization={"text_encoder": "fp8", "transformer": None, "vae": None}, + task="t2i", + prompt="a cup of coffee on a wooden table, morning light", + max_lpips=0.15, + num_inference_steps=10, + enable_cpu_offload=True, + height=1024, + width=1024, + ), QualityTestConfig( id="fp8_qwen_image", model="Qwen/Qwen-Image", @@ -459,7 +472,10 @@ def test_quantization_quality(config: QualityTestConfig): generate_fn = _generate_video if config.task == "t2v" else _generate_image # --- BF16 baseline --- - omni_bl = Omni(model=config.baseline_ref()) + bl_kwargs: dict = {"model": config.baseline_ref()} + if config.enable_cpu_offload: + bl_kwargs["enable_cpu_offload"] = True + omni_bl = Omni(**bl_kwargs) baseline_out, bl_mem = generate_fn(omni_bl, config) omni_bl.shutdown() del omni_bl @@ -468,10 +484,13 @@ def test_quantization_quality(config: QualityTestConfig): # --- Quantized --- quantization = config.quantization_ref() + qt_kwargs: dict = {"model": config.quantized_ref()} + if config.enable_cpu_offload: + qt_kwargs["enable_cpu_offload"] = True if quantization is None: - omni_qt = Omni(model=config.quantized_ref()) + omni_qt = Omni(**qt_kwargs) else: - omni_qt = Omni(model=config.quantized_ref(), quantization_config=quantization) + omni_qt = Omni(**qt_kwargs, quantization_config=quantization) quant_out, qt_mem = generate_fn(omni_qt, config) omni_qt.shutdown() del omni_qt diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py index f88e582e3..351d94f7e 100644 --- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py +++ b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py @@ -37,10 +37,17 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) +def _resolve_component_quant_config(quant_config, component: str): + if hasattr(quant_config, "resolve"): + return quant_config.resolve(component) + return quant_config + + class Flux2ImageProcessor(VaeImageProcessor): """Image processor to preprocess the reference image for Flux2.""" @@ -393,10 +400,14 @@ class Flux2Pipeline( text_encoder_config = AutoConfig.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only ) + te_quant_config = _resolve_component_quant_config(od_config.quantization_config, "text_encoder") self.text_encoder = MistralEncoderModel( text_encoder_config, prefix="text_encoder", - ).to(self._execution_device) + quant_config=te_quant_config, + ) + if not any(param.is_meta for param in self.text_encoder.parameters()): + self.text_encoder.to(self._execution_device) self.tokenizer = PixtralProcessor.from_pretrained( model, subfolder="tokenizer", local_files_only=local_files_only ) @@ -409,8 +420,9 @@ class Flux2Pipeline( self._execution_device ) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel) + transformer_quant_config = _resolve_component_quant_config(od_config.quantization_config, "transformer") self.transformer = Flux2Transformer2DModel( - quant_config=od_config.quantization_config, od_config=od_config, **transformer_kwargs + quant_config=transformer_quant_config, od_config=od_config, **transformer_kwargs ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 @@ -1148,5 +1160,26 @@ class Flux2Pipeline( return DiffusionOutput(output=image) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + text_encoder_quant_config = _resolve_component_quant_config(self.od_config.quantization_config, "text_encoder") + transformer_quant_config = _resolve_component_quant_config(self.od_config.quantization_config, "transformer") + if ( + self.od_config.enable_cpu_offload + and text_encoder_quant_config is not None + and transformer_quant_config is None + ): + weights = self._offload_transformer_before_text_encoder(weights) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) + + def _offload_transformer_before_text_encoder( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[tuple[str, torch.Tensor]]: + transformer_offloaded = False + for name, weight in weights: + if not transformer_offloaded and name.startswith("text_encoder."): + self.transformer.to("cpu") + current_omni_platform.empty_cache() + transformer_offloaded = True + logger.info("Offloaded the unquantized transformer before text encoder quantization") + yield name, weight diff --git a/vllm_omni/diffusion/models/mistral_encoder/mistral_encoder.py b/vllm_omni/diffusion/models/mistral_encoder/mistral_encoder.py index abeacff2a..518453a17 100644 --- a/vllm_omni/diffusion/models/mistral_encoder/mistral_encoder.py +++ b/vllm_omni/diffusion/models/mistral_encoder/mistral_encoder.py @@ -17,7 +17,7 @@ from __future__ import annotations import logging from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any import torch import torch.nn as nn @@ -37,6 +37,9 @@ from vllm.model_executor.layers.linear import ( from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead, VocabParallelEmbedding from vllm.model_executor.model_loader.weight_utils import default_weight_loader +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + logger = logging.getLogger(__name__) @@ -118,6 +121,7 @@ class MistralEncoderAttention(nn.Module): num_kv_heads: int, head_dim: int, prefix: str = "", + quant_config: QuantizationConfig | None = None, ): super().__init__() self.hidden_size = hidden_size @@ -140,6 +144,7 @@ class MistralEncoderAttention(nn.Module): total_num_heads=num_heads, total_num_kv_heads=num_kv_heads, bias=False, + quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) @@ -147,6 +152,7 @@ class MistralEncoderAttention(nn.Module): input_size=num_heads * head_dim, output_size=hidden_size, bias=False, + quant_config=quant_config, prefix=f"{prefix}.o_proj", ) @@ -194,18 +200,21 @@ class MistralEncoderMLP(nn.Module): hidden_size: int, intermediate_size: int, prefix: str = "", + quant_config: QuantizationConfig | None = None, ): super().__init__() self.gate_up_proj = MergedColumnParallelLinear( input_size=hidden_size, output_sizes=[intermediate_size, intermediate_size], bias=False, + quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( input_size=intermediate_size, output_size=hidden_size, bias=False, + quant_config=quant_config, prefix=f"{prefix}.down_proj", ) self.act_fn = SiluAndMul() @@ -227,6 +236,7 @@ class MistralEncoderLayer(nn.Module): intermediate_size: int, rms_norm_eps: float, prefix: str = "", + quant_config: QuantizationConfig | None = None, ): super().__init__() self.self_attn = MistralEncoderAttention( @@ -235,11 +245,13 @@ class MistralEncoderLayer(nn.Module): num_kv_heads=num_kv_heads, head_dim=head_dim, prefix=f"{prefix}.self_attn", + quant_config=quant_config, ) self.mlp = MistralEncoderMLP( hidden_size=hidden_size, intermediate_size=intermediate_size, prefix=f"{prefix}.mlp", + quant_config=quant_config, ) self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) @@ -295,7 +307,12 @@ class MistralEncoderModel(nn.Module): parallel layers for TP but simple SDPA for attention (no PagedAttention). """ - def __init__(self, config: PretrainedConfig, prefix: str = ""): + def __init__( + self, + config: PretrainedConfig, + prefix: str = "", + quant_config: QuantizationConfig | None = None, + ): super().__init__() self._processor = None self._system_message_t2i: str | None = None @@ -339,6 +356,7 @@ class MistralEncoderModel(nn.Module): m = self.language_model.model m.embed_tokens = VocabParallelEmbedding(self.vocab_size, self.hidden_size) + layer_prefix_root = f"{prefix}.language_model.model.layers" if prefix else "language_model.model.layers" m.layers = nn.ModuleList( [ @@ -349,7 +367,8 @@ class MistralEncoderModel(nn.Module): head_dim=self.head_dim, intermediate_size=self.intermediate_size, rms_norm_eps=self.rms_norm_eps, - prefix=f"language_model.model.layers.{i}", + prefix=f"{layer_prefix_root}.{i}", + quant_config=quant_config, ) for i in range(self.num_layers) ]