[WAN_S2V] Enable HSDP for wan s2v (#4458)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

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)      |       ❌        |        ❌         |

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

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

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()
 

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(