From eb81914c644cbe785b3dbdd5e6efa390af7c5808 Mon Sep 17 00:00:00 2001 From: "Chendi.Xue" Date: Wed, 24 Jun 2026 01:04:04 -0500 Subject: [PATCH] [S2V] Add sequence parallelism support to WanS2VTransformer3DModel (#4276) Signed-off-by: Chendi Xue --- docker/Dockerfile.xpu | 2 +- docs/models/supported_models.md | 2 +- docs/user_guide/diffusion_features.md | 2 +- .../models/wan2_2/pipeline_wan2_2_s2v.py | 8 +- .../models/wan2_2/wan2_2_s2v_transformer.py | 76 +++++++++++++++++++ 5 files changed, 85 insertions(+), 5 deletions(-) diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 0dc703a98..f372a1239 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -101,7 +101,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \ export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" -ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/" +ENV LD_LIBRARY_PATH="/opt/intel/oneapi/ccl/2021.15/lib:$LD_LIBRARY_PATH:/usr/local/lib/" ENV VLLM_TARGET_DEVICE=xpu ENV VLLM_WORKER_MULTIPROC_METHOD=spawn diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 75b3ea1ad..86e526a8d 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -37,7 +37,7 @@ th { | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | -| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | | | `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 6c96015d4..0f96aa2c6 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -146,7 +146,7 @@ The following tables show which models support each feature: | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | |------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:| | **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ | -| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ | +| **Wan2.2-S2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ | | **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py index 1bd32bb42..fd9df2676 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py @@ -606,13 +606,17 @@ class Wan22S2VPipeline( t5_checkpoint = os.path.join(model_path, "models_t5_umt5-xxl-enc-bf16.pth") self.text_encoder = UMT5EncoderModel(_WAN_UMT5_CONFIG) - self.text_encoder = _load_wan_t5_as_umt5(self.text_encoder, t5_checkpoint, dtype=dtype).to(self.device) + _cpu_offload = self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload + self.text_encoder = _load_wan_t5_as_umt5(self.text_encoder, t5_checkpoint, dtype=dtype) + if not _cpu_offload: + self.text_encoder = self.text_encoder.to(self.device) # -- VAE (original Wan2.1 VAE, loaded via diffusers from_single_file) -- vae_pth = os.path.join(model_path, "Wan2.1_VAE.pth") self.vae = DistributedAutoencoderKLWan.from_single_file(vae_pth, torch_dtype=dtype) self.vae.init_distributed() - self.vae = self.vae.to(self.device) + if not _cpu_offload: + self.vae = self.vae.to(self.device) # -- Audio encoder (wav2vec2) -- wav2vec_path = os.path.join(model_path, "wav2vec2-large-xlsr-53-english") diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py index 8fef41f74..7d18eb214 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py @@ -31,6 +31,7 @@ from vllm.model_executor.layers.linear import ( 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 from vllm_omni.diffusion.layers.rope import RotaryEmbeddingS2VGrid, RotaryEmbeddingWanS2V, WanS2VRotaryPosEmbed from vllm_omni.platforms import current_omni_platform @@ -1058,6 +1059,32 @@ def _sample_indices(num_frames, stride, expand_ratio, c): return indices +# --------------------------------------------------------------------------- +# SP helper modules +# --------------------------------------------------------------------------- + + +class WanS2VBlocksPrepare(nn.Module): + """Pass-through module providing SP split hook point before transformer blocks. + + Returns (hidden_states, freqs) unchanged. The SP hook shards both outputs + via split_output=True along the sequence dimension. + """ + + def forward(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return hidden_states, freqs + + +class WanS2VBlocksOutputGather(nn.Module): + """Pass-through module providing SP gather hook point after transformer blocks. + + Returns hidden_states unchanged. The SP gather hook collects from all ranks. + """ + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + # --------------------------------------------------------------------------- # Main model # --------------------------------------------------------------------------- @@ -1087,6 +1114,17 @@ class WanS2VTransformer3DModel(nn.Module): _hsdp_shard_conditions = [_is_transformer_block] + _sp_plan = { + # Shard both hidden_states and freqs via blocks_prepare split_output + # hidden_states: [B, seq, dim], freqs: [B, seq, 1, head_dim//2] (complex) + "blocks_prepare": { + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), + }, + # Gather hidden_states after transformer blocks + "blocks_output_gather": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + def __init__( self, cond_dim=0, @@ -1157,6 +1195,10 @@ class WanS2VTransformer3DModel(nn.Module): ] ) + # SP hook modules (pass-through; hooks shard/gather at these boundaries) + self.blocks_prepare = WanS2VBlocksPrepare() + self.blocks_output_gather = WanS2VBlocksOutputGather() + # Output head self.norm_out = FP32LayerNorm(dim, eps, elementwise_affine=False) self.proj_out = nn.Linear(dim, math.prod(patch_size) * out_dim) @@ -1396,8 +1438,21 @@ class WanS2VTransformer3DModel(nn.Module): ] return x, seq_lens, rope_embs, mask_input + def _is_sp_active(self) -> bool: + from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available + + if not is_forward_context_available(): + return False + return get_forward_context().sp_active + def after_transformer_block(self, block_idx, hidden_states): if block_idx in self.audio_injector.injected_block_id.keys(): + sp_active = self._is_sp_active() + if sp_active: + from vllm_omni.diffusion.distributed.sp_sharding import sp_gather, sp_shard + + hidden_states = sp_gather(hidden_states, dim=1, validate=False) + audio_attn_id = self.audio_injector.injected_block_id[block_idx] audio_emb = self.merged_audio_emb num_frames = audio_emb.shape[1] @@ -1424,6 +1479,9 @@ class WanS2VTransformer3DModel(nn.Module): residual_out = rearrange(residual_out, "(b t) n c -> b (t n) c", t=num_frames) hidden_states[:, : self.original_seq_len].add_(residual_out) + if sp_active: + hidden_states = sp_shard(hidden_states, dim=1, validate=False) + return hidden_states # ------------------------------------------------------------------ @@ -1574,6 +1632,21 @@ class WanS2VTransformer3DModel(nn.Module): context_lens = None + # SP: shard hidden_states and freqs before entering blocks + x, self.pre_compute_freqs = self.blocks_prepare(x, self.pre_compute_freqs) + + # SP: adjust segment boundary for the local shard + if self._is_sp_active(): + from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_rank, + ) + + sp_rank = get_sequence_parallel_rank() + chunk_size = x.shape[1] + global_start = sp_rank * chunk_size + local_seg_boundary = max(0, min(e0[1] - global_start, chunk_size)) + e0 = [e0[0], local_seg_boundary] + # Transformer blocks kwargs = dict( e=e0, @@ -1588,6 +1661,9 @@ class WanS2VTransformer3DModel(nn.Module): x = block(x, **kwargs) x = self.after_transformer_block(idx, x) + # SP: gather hidden_states from all ranks + x = self.blocks_output_gather(x) + # Output norm, projection & unpatchify x = x[:, : self.original_seq_len] shift, scale = self.output_scale_shift_prepare(e)