[Feat] support HSDP for DreamID-Omni (#3138)

Signed-off-by: hanzheli <hanzheli@kuaishou.com>
Signed-off-by: fywc <hanzheli@kuaishou.com>

docs/user_guide/diffusion_features.md

@@ -145,7 +145,7 @@ The following tables show which models support each feature:
 | **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
-| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 
 **Frame Interpolation Support**
 

docs/user_guide/examples/offline_inference/x_to_video_audio.md

@@ -47,6 +47,24 @@ python x_to_video_audio.py \
 ```
 In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
 The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
+For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
+
+```python
+python x_to_video_audio.py \
+  --model /path/to/dreamid_omni \
+  --prompt "Two people walking together and singing happily" \
+  --image-path ./example0.png ./example1.png \
+  --audio-path ./example0.wav ./example1.wav \
+  --video-negative-prompt "jitter, bad hands, blur, distortion" \
+  --audio-negative-prompt "robotic, muffled, echo, distorted" \
+  --cfg-parallel-size 2 \
+  --num-inference-steps 45 \
+  --height 704 \
+  --width 1280 \
+  --use-hsdp \
+  --hsdp-shard-size 2 \
+  --output out_dreamid_omni_twoip.mp4
+```
 
 
 You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
@@ -77,6 +95,9 @@ Key arguments:
 - `--image-path`: path to the input image list.
 - `--audio-path`: path to the input audio list, indicate the timbre of the output video.
 - `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
+- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
+- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
+- `--hsdp-replicate-size`: number of HSDP replica groups.
 - `--num-inference-steps`: number of denoising steps (defaults 45).
 - `--video-negative-prompt`: negative prompt for video generation.
 - `--audio-negative-prompt`: negative prompt for audio generation.

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -46,6 +46,24 @@ python x_to_video_audio.py \
 ```
 In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
 The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
+For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
+
+```python
+python x_to_video_audio.py \
+  --model /path/to/dreamid_omni \
+  --prompt "Two people walking together and singing happily" \
+  --image-path ./example0.png ./example1.png \
+  --audio-path ./example0.wav ./example1.wav \
+  --video-negative-prompt "jitter, bad hands, blur, distortion" \
+  --audio-negative-prompt "robotic, muffled, echo, distorted" \
+  --cfg-parallel-size 2 \
+  --num-inference-steps 45 \
+  --height 704 \
+  --width 1280 \
+  --use-hsdp \
+  --hsdp-shard-size 2 \
+  --output out_dreamid_omni_twoip.mp4
+```
 
 
 You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
@@ -76,6 +94,9 @@ Key arguments:
 - `--image-path`: path to the input image list.
 - `--audio-path`: path to the input audio list, indicate the timbre of the output video.
 - `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
+- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
+- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
+- `--hsdp-replicate-size`: number of HSDP replica groups.
 - `--num-inference-steps`: number of denoising steps (defaults 45).
 - `--video-negative-prompt`: negative prompt for video generation.
 - `--audio-negative-prompt`: negative prompt for audio generation.

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -63,6 +63,23 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--use-hsdp",
+        action="store_true",
+        help="Enable HSDP for DreamID-Omni fused transformer blocks.",
+    )
+    parser.add_argument(
+        "--hsdp-shard-size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for HSDP sharding when HSDP is enabled.",
+    )
+    parser.add_argument(
+        "--hsdp-replicate-size",
+        type=int,
+        default=1,
+        help="Number of HSDP replica groups. Default 1 means pure sharding.",
+    )
     return parser.parse_args()
 
 
@@ -124,6 +141,9 @@ def main() -> None:
 
     parallel_config = DiffusionParallelConfig(
         cfg_parallel_size=args.cfg_parallel_size,
+        use_hsdp=args.use_hsdp,
+        hsdp_shard_size=args.hsdp_shard_size,
+        hsdp_replicate_size=args.hsdp_replicate_size,
     )
 
     omni = Omni(

tests/diffusion/models/dreamid_omni/test_dreamid_omni_hsdp.py

@@ -0,0 +1,27 @@
+import pytest
+import torch.nn as nn
+
+pytest.importorskip("dreamid_omni.modules.model")
+
+from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_dreamid_omni_exposes_hsdp_shard_conditions_for_fused_blocks():
+    model = object.__new__(FusionModel)
+    nn.Module.__init__(model)
+    model.fused_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
+    model.attn = nn.Identity()
+
+    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 == ["fused_blocks.0", "fused_blocks.1"]

vllm_omni/diffusion/models/dreamid_omni/__init__.py

@@ -1,8 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from .pipeline_dreamid_omni import DreamIDOmniPipeline
-
 __all__ = [
     "DreamIDOmniPipeline",
 ]
+
+
+def __getattr__(name: str):
+    if name == "DreamIDOmniPipeline":
+        from .pipeline_dreamid_omni import DreamIDOmniPipeline
+
+        return DreamIDOmniPipeline
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

vllm_omni/diffusion/models/dreamid_omni/fusion.py

@@ -241,6 +241,12 @@ class FusedBlock(nn.Module):
 class FusionModel(nn.Module):
     _layerwise_offload_blocks_attrs = ["fused_blocks"]
 
+    @staticmethod
+    def _is_fused_block(name: str, module) -> bool:
+        return "fused_blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_fused_block]
+
     def __init__(self, video_config=None, audio_config=None):
         super().__init__()
         has_video = True