anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feat] support HSDP for Flux family (#1900)

Browse Source 
Signed-off-by: Lancer <maruixiang6688@gmail.com>
This commit is contained in:
Lancer authored and GitHub committed 2026-03-15 22:48:12 +08:00
1 parent 284575ab82
commit 9aae75c3ee
7 files changed
+37 -20
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/diffusion/parallelism_acceleration.md
+16 -16
View File
@@ -24,22 +24,22 @@ The following table shows which models are currently supported by parallelism me
### ImageGen
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel |
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A |
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A |
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A |
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A |
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A |
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A |
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A |
| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ |
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A |
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | HSDP |
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|:----:|
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ❌ |
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ❌ |
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A | ❌ |
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A | ❌ |
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A | ❌ |
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A | ❌ |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A | ✅ |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A | ✅ |
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ✅ |
| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ |
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A | ❌ |
!!! note "TP Limitations for Diffusion Models"
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
vllm_omni/diffusion/models/flux/flux_transformer.py
+6
View File
@@ -470,6 +470,12 @@ class FluxTransformer2DModel(nn.Module):
# -- typically a transformer layer
# used for torch compile optimizations
_repeated_blocks = ["FluxTransformerBlock"]
@staticmethod
def _is_transformer_block(name: str, module) -> bool:
return ("transformer_blocks" in name or "single_transformer_blocks" in name) and name.split(".")[-1].isdigit()
_hsdp_shard_conditions = [_is_transformer_block]
packed_modules_mapping = {
"to_qkv": ["to_q", "to_k", "to_v"],
"add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"],
vllm_omni/diffusion/models/flux/pipeline_flux.py
+2 -2
View File
@@ -160,10 +160,10 @@ class FluxPipeline(nn.Module, CFGParallelMixin):
)
self.text_encoder = CLIPTextModel.from_pretrained(
model, subfolder="text_encoder", local_files_only=local_files_only
)
).to(self.device)
self.text_encoder_2 = T5EncoderModel.from_pretrained(
model, subfolder="text_encoder_2", local_files_only=local_files_only
)
).to(self.device)
self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
self.device
)
vllm_omni/diffusion/models/flux2/flux2_transformer.py
+6
View File
@@ -553,6 +553,12 @@ class Flux2Transformer2DModel(nn.Module):
_repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
@staticmethod
def _is_transformer_block(name: str, module) -> bool:
return ("transformer_blocks" in name or "single_transformer_blocks" in name) and name.split(".")[-1].isdigit()
_hsdp_shard_conditions = [_is_transformer_block]
def __init__(
self,
patch_size: int = 1,
vllm_omni/diffusion/models/flux2/pipeline_flux2.py
+1 -1
View File
@@ -366,7 +366,7 @@ class Flux2Pipeline(nn.Module, SupportImageInput):
)
self.text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
model, subfolder="text_encoder", local_files_only=local_files_only
)
).to(self._execution_device)
self.tokenizer = PixtralProcessor.from_pretrained(
model, subfolder="tokenizer", local_files_only=local_files_only
)
vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py
+5
View File
@@ -741,6 +741,11 @@ class Flux2Transformer2DModel(nn.Module):
_repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"]
@staticmethod
def _is_transformer_block(name: str, module) -> bool:
return ("transformer_blocks" in name or "single_transformer_blocks" in name) and name.split(".")[-1].isdigit()
_hsdp_shard_conditions = [_is_transformer_block]
_sp_plan = {
"": {
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True),
vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
+1 -1
View File
@@ -218,7 +218,7 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput):
model,
subfolder="text_encoder",
local_files_only=local_files_only,
)
).to(self._execution_device)
self.tokenizer = Qwen2TokenizerFast.from_pretrained(
model,
subfolder="tokenizer",