From 60ccadfbb63f2b45d28d32c9e12f5e5e53956480 Mon Sep 17 00:00:00 2001 From: "Chendi.Xue" Date: Sat, 27 Jun 2026 23:27:58 -0500 Subject: [PATCH] [WAN_S2V] Enable HSDP for wan s2v (#4458) Signed-off-by: Chendi Xue --- docs/user_guide/diffusion_features.md | 2 +- .../speech_to_video/speech_to_video.py | 26 ++ .../models/wan2_2/test_wan22_s2v_pipeline.py | 335 ++++++++++++++++++ .../models/wan2_2/pipeline_wan2_2_s2v.py | 2 +- .../models/wan2_2/wan2_2_s2v_transformer.py | 33 +- 5 files changed, 384 insertions(+), 14 deletions(-) create mode 100644 tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index fef039784..3e1bff799 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/examples/offline_inference/speech_to_video/speech_to_video.py b/examples/offline_inference/speech_to_video/speech_to_video.py index d783ee2f7..bfa6d961a 100644 --- a/examples/offline_inference/speech_to_video/speech_to_video.py +++ b/examples/offline_inference/speech_to_video/speech_to_video.py @@ -172,6 +172,29 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for ulysses sequence parallelism.", ) + parser.add_argument( + "--use-hsdp", + action="store_true", + help="Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.", + ) + parser.add_argument( + "--hsdp-shard-size", + type=int, + default=-1, + help=( + "Number of GPUs to shard model weights across within each replica group. " + "-1 (default) auto-calculates as world_size / replicate_size." + ), + ) + parser.add_argument( + "--hsdp-replicate-size", + type=int, + default=1, + help=( + "Number of replica groups for HSDP. Each replica holds a full sharded copy. " + "Default 1 means pure sharding (no replication)." + ), + ) parser.add_argument( "--enforce-eager", action="store_true", @@ -234,6 +257,9 @@ def main(): vae_patch_parallel_size=args.vae_patch_parallel_size, ring_degree=args.ring_degree, ulysses_degree=args.ulysses_degree, + use_hsdp=args.use_hsdp, + hsdp_shard_size=args.hsdp_shard_size, + hsdp_replicate_size=args.hsdp_replicate_size, ) profiler_enabled = args.profiler_config is not None diff --git a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py new file mode 100644 index 000000000..4440a8f35 --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py @@ -0,0 +1,335 @@ +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_s2v_exposes_hsdp_shard_conditions_for_transformer_blocks(): + model = object.__new__(WanS2VTransformer3DModel) + nn.Module.__init__(model) + model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(3)]) + + conditions = getattr(model, "_hsdp_shard_conditions", None) + + assert conditions is not None + assert len(conditions) == 1 + + matched = [] + for name, module in model.named_modules(): + if any(cond(name, module) for cond in conditions): + matched.append(name) + + assert matched == ["blocks.0", "blocks.1", "blocks.2"] + + +def test_s2v_hsdp_shard_condition_does_not_match_non_block_modules(): + model = object.__new__(WanS2VTransformer3DModel) + nn.Module.__init__(model) + model.blocks = nn.ModuleList([nn.Linear(4, 4)]) + model.head_indicator = nn.Linear(4, 4) + model.casual_audio_encoder = nn.Linear(4, 4) + + conditions = model._hsdp_shard_conditions + non_block_matched = [] + for name, module in model.named_modules(): + if name and "blocks" not in name: + if any(cond(name, module) for cond in conditions): + non_block_matched.append(name) + + assert non_block_matched == [] + + +def test_encode_audio_calls_unshard_reshard_when_fsdp_managed(): + model = object.__new__(WanS2VTransformer3DModel) + nn.Module.__init__(model) + model.enable_adain = False + model.casual_audio_encoder = MagicMock(return_value=torch.zeros(1, 10, 64)) + + model.unshard = MagicMock() + model.reshard = MagicMock() + + audio_input = torch.randn(1, 1, 64, 5) + motion_frames = [2, 2] + + result = model.encode_audio(audio_input, motion_frames) + + model.unshard.assert_called_once() + model.reshard.assert_called_once() + assert "audio_emb" in result + + +def test_encode_audio_reshard_called_on_exception(): + """Test that reshard() is always called even when encode_audio logic raises.""" + model = object.__new__(WanS2VTransformer3DModel) + nn.Module.__init__(model) + model.enable_adain = False + model.casual_audio_encoder = MagicMock(side_effect=RuntimeError("encoder failed")) + + model.unshard = MagicMock() + model.reshard = MagicMock() + + audio_input = torch.randn(1, 1, 64, 5) + motion_frames = [2, 2] + + with pytest.raises(RuntimeError, match="encoder failed"): + model.encode_audio(audio_input, motion_frames) + + model.unshard.assert_called_once() + model.reshard.assert_called_once() + + +def test_encode_audio_skips_unshard_reshard_when_not_fsdp(): + model = object.__new__(WanS2VTransformer3DModel) + nn.Module.__init__(model) + model.enable_adain = False + model.casual_audio_encoder = MagicMock(return_value=torch.zeros(1, 10, 64)) + + audio_input = torch.randn(1, 1, 64, 5) + motion_frames = [2, 2] + + result = model.encode_audio(audio_input, motion_frames) + + assert not hasattr(model, "unshard") + assert not hasattr(model, "reshard") + assert "audio_emb" in result + + +def test_s2v_pipeline_skips_cpu_offload_when_hsdp_enabled(): + """Test that transformer.to('cpu') is NOT called when HSDP is active.""" + from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline + + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + + od_config = MagicMock() + od_config.enable_cpu_offload = True + parallel_config = MagicMock() + parallel_config.use_hsdp = True + od_config.parallel_config = parallel_config + pipeline.od_config = od_config + + mock_transformer = MagicMock() + pipeline.transformer = mock_transformer + + # Simulate the offload decision from the forward loop + if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False): + pipeline.transformer.to("cpu") + + mock_transformer.to.assert_not_called() + + +def test_s2v_pipeline_allows_cpu_offload_when_hsdp_disabled(): + """Test that transformer.to('cpu') IS called when HSDP is not active.""" + from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline + + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + + od_config = MagicMock() + od_config.enable_cpu_offload = True + parallel_config = MagicMock() + parallel_config.use_hsdp = False + od_config.parallel_config = parallel_config + pipeline.od_config = od_config + + mock_transformer = MagicMock() + pipeline.transformer = mock_transformer + + # Simulate the offload decision from the forward loop + if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False): + pipeline.transformer.to("cpu") + + mock_transformer.to.assert_called_once_with("cpu") + + +def test_s2v_pipeline_hsdp_forward_complete_process(): + """Integration-level mock test verifying the complete forward path of the + S2V pipeline under HSDP mode. + + Verifies: + - Text encoding is invoked + - Audio encoding invokes unshard/reshard on the transformer + - Reference image encoding via VAE + - Denoising loop calls the transformer + - VAE decode is invoked + - transformer.to('cpu') is NOT called (HSDP mode) + - Final output is a DiffusionOutput with video + audio + """ + import numpy as np + import PIL.Image + + from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline + + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + + # -- Config -- + od_config = MagicMock() + od_config.enable_cpu_offload = True + od_config.enable_diffusion_pipeline_profiler = False + parallel_config = MagicMock() + parallel_config.use_hsdp = True + od_config.parallel_config = parallel_config + pipeline.od_config = od_config + + # -- Pipeline attributes -- + pipeline.device = torch.device("cpu") + pipeline._guidance_scale = 4.5 + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.vae_scale_factor_spatial = 8 + pipeline.vae_scale_factor_temporal = 4 + pipeline.resolution_divisor = 16 + pipeline.motion_frames = 73 + pipeline.drop_first_motion = True + pipeline.fps = 16 + pipeline.audio_sample_m = 0 + pipeline._DEFAULT_INFER_FRAMES = 80 + + # -- Mock text encoder -- + pipeline.tokenizer = MagicMock() + pipeline.tokenizer.return_value = MagicMock( + input_ids=torch.zeros(1, 512, dtype=torch.long), + attention_mask=torch.ones(1, 512, dtype=torch.long), + ) + mock_text_encoder = MagicMock() + mock_text_encoder.dtype = torch.bfloat16 + mock_text_encoder.return_value = MagicMock(last_hidden_state=torch.zeros(1, 512, 4096)) + pipeline.text_encoder = mock_text_encoder + + # -- Mock transformer -- + mock_transformer = MagicMock() + mock_transformer.dtype = torch.bfloat16 + mock_transformer.parameters = MagicMock(return_value=iter([torch.zeros(1)])) + + # encode_audio returns dict with audio_emb + mock_transformer.encode_audio = MagicMock(return_value={"audio_emb": torch.zeros(1, 10, 64)}) + # Forward returns noise prediction + mock_transformer.return_value = (torch.zeros(1, 16, 20, 88, 128),) + pipeline.transformer = mock_transformer + + # -- Mock VAE -- + mock_vae = MagicMock() + mock_vae.dtype = torch.bfloat16 + mock_vae.config = MagicMock() + mock_vae.config.scale_factor_temporal = 4 + mock_vae.config.scale_factor_spatial = 8 + mock_vae.config.latents_mean = [0.0] * 16 + mock_vae.config.latents_std = [1.0] * 16 + mock_vae.config.z_dim = 16 + # encode returns mock with latent_dist + mock_encode_result = MagicMock() + mock_encode_result.latent_dist = MagicMock() + mock_encode_result.latent_dist.mode = MagicMock(return_value=torch.zeros(1, 16, 1, 88, 128)) + mock_vae.encode = MagicMock(return_value=mock_encode_result) + # decode returns video + mock_vae.decode = MagicMock(return_value=(torch.zeros(1, 3, 80, 704, 1024),)) + pipeline.vae = mock_vae + + # -- Mock audio model -- + mock_audio_model = MagicMock() + mock_audio_model.device = torch.device("cpu") + mock_audio_param = torch.zeros(1) + mock_audio_model.parameters = MagicMock(return_value=iter([mock_audio_param])) + mock_audio_model.return_value = MagicMock(hidden_states=[torch.zeros(1, 100, 1024)] * 25) + pipeline.audio_model = mock_audio_model + + pipeline.audio_processor = MagicMock() + pipeline.audio_processor.return_value = MagicMock(input_values=torch.zeros(1, 16000)) + + # -- Mock scheduler -- + mock_scheduler = MagicMock() + mock_scheduler.timesteps = torch.linspace(999, 0, 5) + mock_scheduler.step = MagicMock(return_value=(torch.zeros(1, 16, 20, 88, 128),)) + pipeline.scheduler = mock_scheduler + + # -- Bind methods from the real class -- + pipeline.encode_prompt = Wan22S2VPipeline.encode_prompt.__get__(pipeline) + pipeline.encode_ref_image = Wan22S2VPipeline.encode_ref_image.__get__(pipeline) + pipeline.prepare_motion_latents = Wan22S2VPipeline.prepare_motion_latents.__get__(pipeline) + pipeline.prepare_latents = Wan22S2VPipeline.prepare_latents.__get__(pipeline) + pipeline.check_inputs = Wan22S2VPipeline.check_inputs.__get__(pipeline) + pipeline.diffuse = Wan22S2VPipeline.diffuse.__get__(pipeline) + pipeline._normalize_latents = Wan22S2VPipeline._normalize_latents.__get__(pipeline) + pipeline._denormalize_latents = Wan22S2VPipeline._denormalize_latents.__get__(pipeline) + pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean + + # -- Build request -- + from vllm_omni.diffusion.request import OmniDiffusionRequest + + sampling_params = MagicMock() + sampling_params.height = 704 + sampling_params.width = 1024 + sampling_params.num_frames = 80 + sampling_params.num_inference_steps = 5 + sampling_params.guidance_scale = 4.5 + sampling_params.guidance_scale_provided = True + sampling_params.generator = None + sampling_params.seed = 42 + + ref_image = PIL.Image.new("RGB", (1024, 704)) + audio_data = np.zeros(16000, dtype=np.float32) + + req = MagicMock(spec=OmniDiffusionRequest) + req.prompts = [ + { + "prompt": "test prompt", + "negative_prompt": "bad quality", + "multi_modal_data": {"image": ref_image, "audio": audio_data}, + "additional_information": {"audio_path": audio_data, "pose_video": None, "init_first_frame": False}, + } + ] + req.sampling_params = sampling_params + + # -- Mock methods that use complex internal state -- + pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80)) + + # Mock progress bar context + pipeline.progress_bar = MagicMock() + pipeline.progress_bar.return_value.__enter__ = MagicMock(return_value=MagicMock()) + pipeline.progress_bar.return_value.__exit__ = MagicMock(return_value=False) + + # Mock predict_noise_maybe_with_cfg from CFGParallelMixin + pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(16, 20, 88, 128)) + pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline) + + # Mock platform methods + with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as mock_platform: + mock_platform.empty_cache = MagicMock() + mock_platform.is_available = MagicMock(return_value=False) + + with patch( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.load_audio", return_value=(audio_data, 16000) + ): + result = Wan22S2VPipeline.forward(pipeline, req=req) + + # -- Assertions -- + # Text encoder was called + mock_text_encoder.assert_called() + + # Audio encoding was invoked + pipeline.encode_audio.assert_called_once() + + # VAE encode was called (for ref image and motion latents) + mock_vae.encode.assert_called() + + # VAE decode was called + mock_vae.decode.assert_called() + + # transformer.to("cpu") must NOT be called under HSDP + mock_transformer.to.assert_not_called() + + # Output is a DiffusionOutput tuple with (video, audio_waveform, sample_rate) + from vllm_omni.diffusion.data import DiffusionOutput + + assert isinstance(result, DiffusionOutput) + video, audio_waveform, audio_sr = result.output + assert video.shape[0] == 1 # batch + assert video.shape[1] == 3 # channels + assert audio_waveform is not None + assert audio_sr == 16000 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 fd9df2676..eb4ddc40c 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 @@ -1288,7 +1288,7 @@ class Wan22S2VPipeline( ) # ---- Decode this clip ---- - if self.od_config.enable_cpu_offload: + if self.od_config.enable_cpu_offload and not getattr(self.od_config.parallel_config, "use_hsdp", False): self.transformer.to("cpu") current_omni_platform.empty_cache() 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 7d18eb214..d1b74d020 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 @@ -1499,18 +1499,27 @@ class WanS2VTransformer3DModel(nn.Module): dict with 'audio_emb' (and optionally 'audio_emb_global' when enable_adain is True). """ - audio_input = torch.cat( - [audio_input[..., 0:1].repeat(1, 1, 1, motion_frames[0]), audio_input], - dim=-1, - ) - audio_emb_res = self.casual_audio_encoder(audio_input) - result = {} - if self.enable_adain: - audio_emb_global, audio_emb = audio_emb_res - result["audio_emb_global"] = audio_emb_global[:, motion_frames[1] :] - else: - audio_emb = audio_emb_res - result["audio_emb"] = audio_emb[:, motion_frames[1] :, :] + # Under HSDP, this method is called outside the FSDP forward hook. + # Unshard root-managed params so casual_audio_encoder can run. + is_fsdp = hasattr(self, "unshard") and hasattr(self, "reshard") + if is_fsdp: + self.unshard() + try: + audio_input = torch.cat( + [audio_input[..., 0:1].repeat(1, 1, 1, motion_frames[0]), audio_input], + dim=-1, + ) + audio_emb_res = self.casual_audio_encoder(audio_input) + result = {} + if self.enable_adain: + audio_emb_global, audio_emb = audio_emb_res + result["audio_emb_global"] = audio_emb_global[:, motion_frames[1] :] + else: + audio_emb = audio_emb_res + result["audio_emb"] = audio_emb[:, motion_frames[1] :, :] + finally: + if is_fsdp: + self.reshard() return result def forward(