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

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

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