From 9d9e6abae666d9e4651273c614e2b6bd72ca013d Mon Sep 17 00:00:00 2001 From: "rongfu.leng" Date: Sun, 29 Mar 2026 14:35:56 +0800 Subject: [PATCH] [Feat] add helios fp8 quantization (#1916) Signed-off-by: rongfu.leng Co-authored-by: Hongsheng Liu --- docs/user_guide/diffusion/quantization/fp8.md | 1 + examples/offline_inference/helios/end2end.py | 7 +++ .../models/helios/helios_transformer.py | 46 +++++++++++++------ .../models/helios/pipeline_helios.py | 17 +++++-- 4 files changed, 54 insertions(+), 17 deletions(-) diff --git a/docs/user_guide/diffusion/quantization/fp8.md b/docs/user_guide/diffusion/quantization/fp8.md index 338781fa2..9906631b6 100644 --- a/docs/user_guide/diffusion/quantization/fp8.md +++ b/docs/user_guide/diffusion/quantization/fp8.md @@ -65,6 +65,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to | Flux | `black-forest-labs/FLUX.1-dev` | All layers | None | | HunyuanImage-3 | `tencent/HunyuanImage3` | All layers | None | | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None | +| Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | All layers | None | ## Combining with Other Features diff --git a/examples/offline_inference/helios/end2end.py b/examples/offline_inference/helios/end2end.py index dae707c2c..88c3b865d 100644 --- a/examples/offline_inference/helios/end2end.py +++ b/examples/offline_inference/helios/end2end.py @@ -178,6 +178,12 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable diffusion pipeline profiler to display stage durations.", ) + parser.add_argument( + "--quantization", + type=str, + default=None, + help="Quantization method (e.g. 'fp8').", + ) # Memory & parallelism parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.") @@ -213,6 +219,7 @@ def main(): parallel_config=parallel_config, enforce_eager=args.enforce_eager, enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, + quantization=args.quantization, ) # Validate I2V / V2V arguments diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py index 17b9ec37e..812da7db1 100644 --- a/vllm_omni/diffusion/models/helios/helios_transformer.py +++ b/vllm_omni/diffusion/models/helios/helios_transformer.py @@ -5,7 +5,7 @@ import math from collections.abc import Iterable from functools import lru_cache -from typing import Any +from typing import TYPE_CHECKING, Any import torch import torch.nn as nn @@ -29,6 +29,11 @@ from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelOutput, ) +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + ) + logger = init_logger(__name__) @@ -96,7 +101,15 @@ class DistributedRMSNorm(nn.Module): class ColumnParallelGELU(nn.Module): """Column parallel linear with GELU activation.""" - def __init__(self, dim_in: int, dim_out: int, *, approximate: str = "tanh", bias: bool = True): + def __init__( + self, + dim_in: int, + dim_out: int, + *, + approximate: str = "tanh", + bias: bool = True, + quant_config: "QuantizationConfig | None" = None, + ): super().__init__() self.proj = ColumnParallelLinear( dim_in, @@ -104,6 +117,7 @@ class ColumnParallelGELU(nn.Module): bias=bias, gather_output=False, return_bias=False, + quant_config=quant_config, ) self.approximate = approximate @@ -121,18 +135,15 @@ class HeliosFeedForward(nn.Module): inner_dim: int, dim_out: int | None = None, bias: bool = True, + quant_config: "QuantizationConfig | None" = None, ) -> None: super().__init__() dim_out = dim_out or dim - self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias) + self.net_0 = ColumnParallelGELU(dim, inner_dim, approximate="tanh", bias=bias, quant_config=quant_config) self.net_1 = nn.Identity() self.net_2 = RowParallelLinear( - inner_dim, - dim_out, - bias=bias, - input_is_parallel=True, - return_bias=False, + inner_dim, dim_out, bias=bias, input_is_parallel=True, return_bias=False, quant_config=quant_config ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -262,6 +273,7 @@ class HeliosSelfAttention(nn.Module): dropout: float = 0.0, is_amplify_history: bool = False, history_scale_mode: str = "per_head", + quant_config: "QuantizationConfig | None" = None, ): super().__init__() @@ -275,6 +287,7 @@ class HeliosSelfAttention(nn.Module): head_size=head_dim, total_num_heads=num_heads, bias=True, + quant_config=quant_config, ) self.num_heads = self.to_qkv.num_heads @@ -290,6 +303,7 @@ class HeliosSelfAttention(nn.Module): bias=True, input_is_parallel=True, return_bias=False, + quant_config=quant_config, ) self.dropout = nn.Dropout(dropout) @@ -368,6 +382,7 @@ class HeliosCrossAttention(nn.Module): head_dim: int, eps: float = 1e-5, dropout: float = 0.0, + quant_config: "QuantizationConfig | None" = None, ): super().__init__() @@ -382,6 +397,7 @@ class HeliosCrossAttention(nn.Module): bias=True, gather_output=False, return_bias=False, + quant_config=quant_config, ) self.to_k = ColumnParallelLinear( dim, @@ -389,6 +405,7 @@ class HeliosCrossAttention(nn.Module): bias=True, gather_output=False, return_bias=False, + quant_config=quant_config, ) self.to_v = ColumnParallelLinear( dim, @@ -396,6 +413,7 @@ class HeliosCrossAttention(nn.Module): bias=True, gather_output=False, return_bias=False, + quant_config=quant_config, ) tp_size = get_tensor_model_parallel_world_size() @@ -411,6 +429,7 @@ class HeliosCrossAttention(nn.Module): bias=True, input_is_parallel=True, return_bias=False, + quant_config=quant_config, ) self.dropout = nn.Dropout(dropout) @@ -461,6 +480,7 @@ class HeliosTransformerBlock(nn.Module): guidance_cross_attn: bool = False, is_amplify_history: bool = False, history_scale_mode: str = "per_head", + quant_config: "QuantizationConfig | None" = None, ): super().__init__() @@ -475,19 +495,17 @@ class HeliosTransformerBlock(nn.Module): eps=eps, is_amplify_history=is_amplify_history, history_scale_mode=history_scale_mode, + quant_config=quant_config, ) # 2. Cross-attention self.attn2 = HeliosCrossAttention( - dim=dim, - num_heads=num_heads, - head_dim=head_dim, - eps=eps, + dim=dim, num_heads=num_heads, head_dim=head_dim, eps=eps, quant_config=quant_config ) self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() # 3. Feed-forward - self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim) + self.ffn = HeliosFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim, quant_config=quant_config) self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) @@ -602,6 +620,7 @@ class HeliosTransformer3DModel(nn.Module): has_multi_term_memory_patch: bool = True, is_amplify_history: bool = False, history_scale_mode: str = "per_head", + quant_config: "QuantizationConfig | None" = None, ) -> None: super().__init__() @@ -689,6 +708,7 @@ class HeliosTransformer3DModel(nn.Module): guidance_cross_attn=guidance_cross_attn, is_amplify_history=is_amplify_history, history_scale_mode=history_scale_mode, + quant_config=quant_config, ) for _ in range(num_layers) ] diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index abccbec1a..df709a515 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -9,7 +9,7 @@ import logging import math import os from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -31,6 +31,11 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.platforms import current_omni_platform +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + ) + logger = logging.getLogger(__name__) @@ -80,7 +85,9 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc return load_json_config(model_path, subfolder, "config.json", local_files_only) -def create_transformer_from_config(config: dict) -> HeliosTransformer3DModel: +def create_transformer_from_config( + config: dict, quant_config: QuantizationConfig | None = None +) -> HeliosTransformer3DModel: kwargs = {} key_map = [ @@ -112,7 +119,7 @@ def create_transformer_from_config(config: dict) -> HeliosTransformer3DModel: val = tuple(val) kwargs[key] = val - return HeliosTransformer3DModel(**kwargs) + return HeliosTransformer3DModel(quant_config=quant_config, **kwargs) def get_helios_post_process_func( @@ -192,7 +199,9 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPip ).to(self.device) transformer_config = load_transformer_config(model, "transformer", local_files_only) - self.transformer = create_transformer_from_config(transformer_config) + self.transformer = create_transformer_from_config( + transformer_config, quant_config=od_config.quantization_config + ) # Read scheduler config to determine scheduler type sched_cfg = load_json_config(model, "scheduler", "scheduler_config.json", local_files_only)