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

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

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"]