[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>

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

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)
             ]