[Bugfix] Use dedicated WORLD group for distributed VAE communication (#6401)

Signed-off-by: david6666666 <530634352@qq.com>

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from types import SimpleNamespace
 
 import pytest
@@ -58,6 +59,11 @@ class E2EOperator:
         return torch.cat(tiles, dim=0)
 
 
+@dataclass
+class FakeWorldGroup:
+    device_group: object
+
+
 class DummyMixin(DistributedVaeMixin):
     def __init__(self):
         self.use_tiling = True
@@ -74,6 +80,16 @@ def mock_dist(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(dist, "broadcast", lambda *args, **kwargs: None)
 
 
+@pytest.fixture(autouse=True)
+def mock_world_group(monkeypatch: pytest.MonkeyPatch):
+    group = object()
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_world_group",
+        lambda: FakeWorldGroup(device_group=group),
+    )
+    return group
+
+
 @pytest.fixture(autouse=True)
 def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(DistributedVaeExecutor, "gather_tensors", lambda self, x: [x])
@@ -85,6 +101,11 @@ def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch):
 # ============================
 
 
+def test_uses_dedicated_world_device_group(mock_world_group):
+    executor = DistributedVaeExecutor()
+    assert executor.group is mock_world_group
+
+
 def test_balance_tasks():
     executor = DistributedVaeExecutor()
     workloads = [2, 5, 13, 8, 2, 4]  # each is 17