diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index f0969b677..d633e7de8 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -105,7 +105,7 @@ The following tables show which models support each feature: | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | | **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | -| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | | **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | | **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | | **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | @@ -113,13 +113,13 @@ The following tables show which models support each feature: | **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | | **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | -| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | -| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | -| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | -| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | +| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | +| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | +| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | -| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ❌ | ✅ | ✅ | ❌ | +| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ | ✅ | ❌ | > Notes: > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT. diff --git a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py index 4501569aa..14e4c915b 100644 --- a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py @@ -98,6 +98,18 @@ def _get_diffusion_feature_cases(model: str): id="parallel_004", marks=PARALLEL_FEATURE_MARKS, ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--use-hsdp", + "--hsdp-shard-size", + "2", + ], + ), + id="parallel_005", + marks=PARALLEL_FEATURE_MARKS, + ), ] diff --git a/tests/e2e/online_serving/test_qwen_image_expansion.py b/tests/e2e/online_serving/test_qwen_image_expansion.py index 6d6d23601..88e56cc3e 100644 --- a/tests/e2e/online_serving/test_qwen_image_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_expansion.py @@ -107,6 +107,18 @@ def _get_diffusion_feature_cases(model: str): id="vae_patch_parallel_2", marks=PARALLEL_FEATURE_MARKS, ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--use-hsdp", + "--hsdp-shard-size", + "2", + ], + ), + id="parallel_hsdp", + marks=PARALLEL_FEATURE_MARKS, + ), ] diff --git a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py index 39b8f36b3..fc73801c0 100644 --- a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py @@ -62,6 +62,18 @@ FEATURE_CASES = [ id="cfg_parallel_001", marks=PARALLEL_FEATURE_MARKS, ), + pytest.param( + OmniServerParams( + model=MODEL, + server_args=[ + "--use-hsdp", + "--hsdp-shard-size", + "2", + ], + ), + id="parallel_hsdp", + marks=PARALLEL_FEATURE_MARKS, + ), ] diff --git a/tests/e2e/online_serving/test_zimage_expansion.py b/tests/e2e/online_serving/test_zimage_expansion.py index dfca76ca2..bed95545a 100644 --- a/tests/e2e/online_serving/test_zimage_expansion.py +++ b/tests/e2e/online_serving/test_zimage_expansion.py @@ -60,6 +60,18 @@ def _get_diffusion_feature_cases(): id="parallel_teacache_fp8_ulysses2_ring2", marks=FOUR_CARD_MARKS, ), + pytest.param( + OmniServerParams( + model=MODEL, + server_args=[ + "--use-hsdp", + "--hsdp-shard-size", + "2", + ], + ), + id="parallel_hsdp", + marks=FOUR_CARD_MARKS, + ), ] diff --git a/vllm_omni/diffusion/distributed/hsdp_utils.py b/vllm_omni/diffusion/distributed/hsdp_utils.py new file mode 100644 index 000000000..3e538d6fd --- /dev/null +++ b/vllm_omni/diffusion/distributed/hsdp_utils.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + + +def is_transformer_block_module(name: str, module: Any) -> bool: + """Return True for numbered modules under `transformer_blocks`.""" + return "transformer_blocks" in name and name.split(".")[-1].isdigit() diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py index 8b129ce2a..490e0198b 100644 --- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py +++ b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py @@ -22,6 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.base import CachedTransformer from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module logger = init_logger(__name__) @@ -724,6 +725,8 @@ class GlmImageTransformer2DModel(CachedTransformer): _repeated_blocks = ["GlmImageTransformerBlock"] + _hsdp_shard_conditions = [is_transformer_block_module] + def __init__( self, od_config: OmniDiffusionConfig, diff --git a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py index 2f7318cef..263e39e01 100644 --- a/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_video/hunyuan_video_15_transformer.py @@ -23,6 +23,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward @@ -544,11 +545,7 @@ class HunyuanVideo15Transformer3DModel(nn.Module): "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], } - @staticmethod - def _is_transformer_block(name: str, module) -> bool: - return "transformer_blocks" in name and name.split(".")[-1].isdigit() - - _hsdp_shard_conditions = [_is_transformer_block] + _hsdp_shard_conditions = [is_transformer_block_module] def __init__( self, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 505bad3d5..5056b5342 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -273,7 +273,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP ) self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only - ) + ).to(self.device) self.vae = DistributedAutoencoderKLQwenImage.from_pretrained( model, subfolder="vae", local_files_only=local_files_only ).to(self.device) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index f805a7e7c..3d0cd2a6d 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py @@ -245,7 +245,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi ) self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only - ) + ).to(self.device) self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index 8e2ba90a4..cb5a36579 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -198,7 +198,7 @@ class QwenImageEditPlusPipeline( ) self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only - ) + ).to(self.device) self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index ee2f471f5..f1d28f068 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py @@ -219,7 +219,7 @@ class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParalle ) self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only - ) + ).to(self.device) self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device ) diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index 3e9a0f0f3..c21156706 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -37,6 +37,7 @@ from vllm_omni.diffusion.attention.backends.abstract import ( from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.base import CachedTransformer from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, @@ -887,6 +888,8 @@ class QwenImageTransformer2DModel(CachedTransformer): "add_kv_proj": ["add_q_proj", "add_k_proj", "add_v_proj"], } + _hsdp_shard_conditions = [is_transformer_block_module] + # Sequence Parallelism plan (following diffusers' _cp_plan pattern) # Similar to Z-Image's UnifiedPrepare, we use ImageRopePrepare to create # a module boundary where _sp_plan can shard hidden_states and vid_freqs together. diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index ac18d5773..b9aceed2e 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -170,7 +170,7 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): self.text_encoder = AutoModel.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only - ) + ).to(self._execution_device) self.vae = DistributedAutoencoderKL.from_pretrained( model, subfolder="vae", local_files_only=local_files_only ).to(self._execution_device) diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index faeff3dce..fd8b0e490 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -580,6 +580,12 @@ class ZImageTransformer2DModel(CachedTransformer): _repeated_blocks = ["ZImageTransformerBlock"] + @staticmethod + def _is_transformer_block(name: str, module) -> bool: + return "layers" in name and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] + # Sequence Parallelism for Z-Image (following diffusers' _cp_plan pattern) # Similar to how Wan uses `rope` module's split_output to shard rotary embeddings, # Z-Image uses `unified_prepare` module's split_output to shard unified tensors.