From 1fe8e71433875bcb62f7f8c0a5e5bb19f8d90768 Mon Sep 17 00:00:00 2001 From: XU Mingshi <91017482+mxuax@users.noreply.github.com> Date: Tue, 27 Jan 2026 11:33:38 +0800 Subject: [PATCH] [Diffusion][Feature] Non-Intrusive Sequence Parallelism (SP) Support for Wan2.2 (#966) Signed-off-by: mxuax Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com> Co-authored-by: Hongsheng Liu --- .../diffusion/parallelism_acceleration.md | 2 +- .../text_to_video/text_to_video.py | 44 +++++++++++++++++++ .../models/wan2_2/wan2_2_transformer.py | 40 +++++++++++++++++ vllm_omni/diffusion/registry.py | 15 +++++-- 4 files changed, 97 insertions(+), 4 deletions(-) diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index 007976f1a..3dc9af454 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -48,7 +48,7 @@ The following table shows which models are currently supported by parallelism me | Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel | |-------|------------------|------------|---------|--------------------------| -| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ | ❌ | +| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ❌ | ### Tensor Parallelism diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 066714c54..c655fff14 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -3,11 +3,13 @@ import argparse import os +import time from pathlib import Path import numpy as np import torch +from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.entrypoints.omni import Omni from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.platform_utils import detect_device_type, is_npu @@ -40,6 +42,23 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Enable CPU offloading for diffusion models.", ) + parser.add_argument( + "--ulysses_degree", + type=int, + default=1, + help="Number of GPUs used for ulysses sequence parallelism.", + ) + parser.add_argument( + "--ring_degree", + type=int, + default=1, + help="Number of GPUs used for ring sequence parallelism.", + ) + parser.add_argument( + "--enforce_eager", + action="store_true", + help="Disable torch.compile and force eager execution.", + ) return parser.parse_args() @@ -52,6 +71,13 @@ def main(): vae_use_slicing = is_npu() vae_use_tiling = is_npu() + # Configure parallel settings (only SP is supported for Wan) + # Note: cfg_parallel and tensor_parallel are not implemented for Wan models + parallel_config = DiffusionParallelConfig( + ulysses_degree=args.ulysses_degree, + ring_degree=args.ring_degree, + ) + # Check if profiling is requested via environment variable profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR")) @@ -62,12 +88,25 @@ def main(): boundary_ratio=args.boundary_ratio, flow_shift=args.flow_shift, enable_cpu_offload=args.enable_cpu_offload, + parallel_config=parallel_config, + enforce_eager=args.enforce_eager, ) if profiler_enabled: print("[Profiler] Starting profiling...") omni.start_profile() + # Print generation configuration + print(f"\n{'=' * 60}") + print("Generation Configuration:") + print(f" Model: {args.model}") + print(f" Inference steps: {args.num_inference_steps}") + print(f" Frames: {args.num_frames}") + print(f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}") + print(f" Video size: {args.width}x{args.height}") + print(f"{'=' * 60}\n") + + generation_start = time.perf_counter() frames = omni.generate( args.prompt, negative_prompt=args.negative_prompt, @@ -79,6 +118,11 @@ def main(): num_inference_steps=args.num_inference_steps, num_frames=args.num_frames, ) + generation_end = time.perf_counter() + generation_time = generation_end - generation_start + + # Print profiling results + print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)") # Extract video frames from OmniRequestOutput if isinstance(frames, list) and len(frames) > 0: diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index f56833e48..deb626a97 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -17,6 +17,10 @@ from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinea from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, +) logger = init_logger(__name__) @@ -497,6 +501,16 @@ class WanTransformer3DModel(nn.Module): This is an optimized version of the diffusers WanTransformer3DModel that uses vLLM's efficient QKVParallelLinear and RMSNorm implementations. + Sequence Parallelism: + This model supports non-intrusive SP via _sp_plan. The plan specifies: + - RoPE (cos/sin) splitting via rope module's split_output + - hidden_states splitting at first transformer block input + - Output gathering at proj_out layer + + The video sequence (flattened patches) is parallelized across GPUs. + + Note: Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in diffusers. + Args: patch_size: 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) num_attention_heads: Number of attention heads @@ -520,6 +534,29 @@ class WanTransformer3DModel(nn.Module): "to_qkv": ["to_q", "to_k", "to_v"], } + # Sequence Parallelism for Wan (following diffusers' _cp_plan pattern) + # + # The _sp_plan specifies sharding/gathering at module boundaries: + # - rope: Split both RoPE outputs (freqs_cos, freqs_sin) via split_output=True + # - blocks.0: Split hidden_states input at the first transformer block + # - proj_out: Gather outputs after the final projection layer + # + # Note: _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism) + _sp_plan = { + # Shard RoPE embeddings after rope module computes them + "rope": { + 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # freqs_cos [1, seq, 1, dim] + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # freqs_sin [1, seq, 1, dim] + }, + # Shard hidden_states at first transformer block input + # (after patch_embedding + flatten + transpose) + "blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), # [B, seq, dim] + }, + # Gather at proj_out (final linear projection before unpatchify) + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + def __init__( self, patch_size: tuple[int, int, int] = (1, 2, 2), @@ -613,8 +650,11 @@ class WanTransformer3DModel(nn.Module): post_patch_height = height // p_h post_patch_width = width // p_w + # Compute RoPE embeddings (sharded by _sp_plan via split_output=True) rotary_emb = self.rope(hidden_states) + # Patch embedding and flatten to sequence + # (hidden_states is sharded at blocks.0 input by _sp_plan) hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 466a58e2c..d4429aff6 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -164,7 +164,10 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - return # Find transformer model(s) in the pipeline that have _sp_plan - transformer_attrs = ["transformer", "dit", "unet"] + # Include transformer_2 for two-stage models (e.g., Wan MoE) + transformer_attrs = ["transformer", "transformer_2", "dit", "unet"] + applied_count = 0 + for attr in transformer_attrs: if not hasattr(model, attr): continue @@ -190,11 +193,17 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - else ("ulysses" if sp_config.ulysses_degree > 1 else "ring") ) logger.info( - f"Applying sequence parallelism to {transformer.__class__.__name__} " + f"Applying sequence parallelism to {transformer.__class__.__name__} ({attr}) " f"(sp_size={sp_size}, mode={mode}, ulysses={sp_config.ulysses_degree}, ring={sp_config.ring_degree})" ) apply_sequence_parallel(transformer, sp_config, plan) - return # Only apply to first transformer found + applied_count += 1 + + if applied_count == 0: + logger.warning( + f"Sequence parallelism is enabled (sp_size={sp_size}) but no transformer with _sp_plan found. " + "SP hooks not applied. Consider adding _sp_plan to your transformer model." + ) except Exception as e: logger.warning(f"Failed to apply sequence parallelism: {e}. Continuing without SP hooks.")