add fp8 quantization support of Dit for z-image

Signed-off-by: lishunyang <[email protected]>

Author: lishunyang12

examples/offline_inference/text_to_image/text_to_image.py

@@ -113,6 +113,15 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8"],
+        help="Quantization method for the transformer. "
+        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs). "
+        "Default: None (no quantization, uses BF16).",
+    )
     parser.add_argument(
         "--vae_use_slicing",
         action="store_true",
@@ -180,6 +189,7 @@ def main():
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
+        quantization=args.quantization,
     )
 
     if profiler_enabled:
@@ -192,6 +202,7 @@ def main():
     print(f"  Model: {args.model}")
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
+    print(f"  Quantization: {args.quantization if args.quantization else 'None (BF16)'}")
     print(
         f"  Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
         f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"

tests/diffusion/quantization/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/diffusion/quantization/test_fp8_config.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for FP8 quantization config."""
+
+import pytest
+
+
+def test_fp8_config_creation():
+    """Test that FP8 config can be created."""
+    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+
+    config = get_diffusion_quant_config("fp8")
+    assert config is not None
+    assert config.get_name() == "fp8"
+
+
+def test_vllm_config_extraction():
+    """Test that vLLM config can be extracted from diffusion config."""
+    from vllm_omni.diffusion.quantization import (
+        get_diffusion_quant_config,
+        get_vllm_quant_config_for_layers,
+    )
+
+    diff_config = get_diffusion_quant_config("fp8")
+    vllm_config = get_vllm_quant_config_for_layers(diff_config)
+    assert vllm_config is not None
+    assert vllm_config.activation_scheme == "dynamic"
+
+
+def test_none_quantization():
+    """Test that None quantization returns None config."""
+    from vllm_omni.diffusion.quantization import (
+        get_diffusion_quant_config,
+        get_vllm_quant_config_for_layers,
+    )
+
+    config = get_diffusion_quant_config(None)
+    assert config is None
+    vllm_config = get_vllm_quant_config_for_layers(config)
+    assert vllm_config is None
+
+
+def test_invalid_quantization():
+    """Test that invalid quantization method raises error."""
+    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+
+    with pytest.raises(ValueError, match="Unknown quantization method"):
+        get_diffusion_quant_config("invalid_method")
+
+
+def test_fp8_config_with_custom_params():
+    """Test FP8 config with custom parameters."""
+    from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+
+    config = get_diffusion_quant_config(
+        "fp8",
+        activation_scheme="static",
+        ignored_layers=["proj_out"],
+    )
+    assert config is not None
+    assert config.activation_scheme == "static"
+    assert "proj_out" in config.ignored_layers
+
+
+def test_supported_methods():
+    """Test that supported methods list is correct."""
+    from vllm_omni.diffusion.quantization import SUPPORTED_QUANTIZATION_METHODS
+
+    assert "fp8" in SUPPORTED_QUANTIZATION_METHODS

vllm_omni/diffusion/data.py

@@ -16,6 +16,9 @@
 
 from vllm_omni.diffusion.utils.network_utils import is_port_available
 
+# Import after TYPE_CHECKING to avoid circular imports at runtime
+# The actual import is deferred to __post_init__ to avoid import order issues
+
 logger = init_logger(__name__)
 
 
@@ -358,6 +361,11 @@ class OmniDiffusionConfig:
     # Omni configuration (injected from stage config)
     omni_kv_config: dict[str, Any] = field(default_factory=dict)
 
+    # Quantization settings
+    # Supported methods: "fp8" (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs)
+    quantization: str | None = None
+    quantization_config: Any | None = None  # DiffusionQuantizationConfig or dict
+
     def settle_port(self, port: int, port_inc: int = 42, max_attempts: int = 100) -> int:
         """
         Find an available port with retry logic.
@@ -444,6 +452,16 @@ def __post_init__(self):
             # If it's neither dict nor DiffusionCacheConfig, convert to empty config
             self.cache_config = DiffusionCacheConfig()
 
+        # Convert quantization config (deferred import to avoid circular imports)
+        if self.quantization is not None or self.quantization_config is not None:
+            from vllm_omni.diffusion.quantization import get_diffusion_quant_config
+
+            if isinstance(self.quantization_config, dict):
+                quant_method = self.quantization_config.pop("method", self.quantization)
+                self.quantization_config = get_diffusion_quant_config(quant_method, **self.quantization_config)
+            elif self.quantization_config is None and self.quantization is not None:
+                self.quantization_config = get_diffusion_quant_config(self.quantization)
+
         if self.max_cpu_loras is None:
             self.max_cpu_loras = 1
         elif self.max_cpu_loras < 1:

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -37,6 +37,7 @@
 from vllm_omni.diffusion.models.z_image.z_image_transformer import (
     ZImageTransformer2DModel,
 )
+from vllm_omni.diffusion.quantization import get_vllm_quant_config_for_layers
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
@@ -173,7 +174,9 @@ def __init__(
         self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
             self._execution_device
         )
-        self.transformer = ZImageTransformer2DModel()
+        # Get vLLM quantization config for linear layers
+        quant_config = get_vllm_quant_config_for_layers(od_config.quantization_config)
+        self.transformer = ZImageTransformer2DModel(quant_config=quant_config)
         self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 
         # Note: Context parallelism is applied centrally in registry.initialize_model()

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -18,6 +18,7 @@
 
 import math
 from collections.abc import Iterable
+from typing import TYPE_CHECKING
 
 import torch
 import torch.nn as nn
@@ -32,6 +33,11 @@
 )
 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,
+    )
+
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.cache.base import CachedTransformer
 from vllm_omni.diffusion.distributed.sp_plan import (
@@ -250,6 +256,7 @@ def __init__(
         num_kv_heads: int,
         qk_norm: bool = True,
         eps: float = 1e-6,
+        quant_config: "QuantizationConfig | None" = None,
     ) -> None:
         super().__init__()
         self.dim = dim
@@ -264,6 +271,7 @@ def __init__(
             total_num_heads=num_heads,
             total_num_kv_heads=num_kv_heads,
             bias=False,
+            quant_config=quant_config,
         )
 
         assert qk_norm is True
@@ -281,6 +289,7 @@ def __init__(
                     bias=False,
                     input_is_parallel=True,
                     return_bias=False,
+                    quant_config=quant_config,
                 )
             ]
         )
@@ -343,13 +352,19 @@ def forward(
 
 
 class FeedForward(nn.Module):
-    def __init__(self, dim: int, hidden_dim: int):
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        quant_config: "QuantizationConfig | None" = None,
+    ):
         super().__init__()
         self.w13 = MergedColumnParallelLinear(
             dim,
             [hidden_dim] * 2,
             bias=False,
             return_bias=False,
+            quant_config=quant_config,
         )
         self.act = SiluAndMul()
         self.w2 = RowParallelLinear(
@@ -358,6 +373,7 @@ def __init__(self, dim: int, hidden_dim: int):
             bias=False,
             input_is_parallel=True,
             return_bias=False,
+            quant_config=quant_config,
         )
 
     def forward(self, x):
@@ -374,6 +390,7 @@ def __init__(
         norm_eps: float,
         qk_norm: bool,
         modulation=True,
+        quant_config: "QuantizationConfig | None" = None,
     ):
         super().__init__()
         self.dim = dim
@@ -384,9 +401,14 @@ def __init__(
             num_kv_heads=n_kv_heads,
             qk_norm=qk_norm,
             eps=1e-5,
+            quant_config=quant_config,
         )
 
-        self.feed_forward = FeedForward(dim=dim, hidden_dim=int(dim / 3 * 8))
+        self.feed_forward = FeedForward(
+            dim=dim,
+            hidden_dim=int(dim / 3 * 8),
+            quant_config=quant_config,
+        )
         self.layer_id = layer_id
 
         self.attention_norm1 = RMSNorm(dim, eps=norm_eps)
@@ -589,6 +611,7 @@ def __init__(
         t_scale=1000.0,
         axes_dims=[32, 48, 48],
         axes_lens=[1024, 512, 512],
+        quant_config: "QuantizationConfig | None" = None,
     ) -> None:
         super().__init__()
         self.dtype = torch.bfloat16
@@ -648,6 +671,7 @@ def __init__(
                     norm_eps,
                     qk_norm,
                     modulation=True,
+                    quant_config=quant_config,
                 )
                 for layer_id in range(n_refiner_layers)
             ]
@@ -662,6 +686,7 @@ def __init__(
                     norm_eps,
                     qk_norm,
                     modulation=False,
+                    quant_config=quant_config,
                 )
                 for layer_id in range(n_refiner_layers)
             ]
@@ -677,7 +702,15 @@ def __init__(
 
         self.layers = nn.ModuleList(
             [
-                ZImageTransformerBlock(layer_id, dim, n_heads, n_kv_heads, norm_eps, qk_norm)
+                ZImageTransformerBlock(
+                    layer_id,
+                    dim,
+                    n_heads,
+                    n_kv_heads,
+                    norm_eps,
+                    qk_norm,
+                    quant_config=quant_config,
+                )
                 for layer_id in range(n_layers)
             ]
         )