[Bugfix] Fix Parallel State Initialization For MoE + Diffusion (#7676)

Signed-off-by: Alex Brooks <albrooks@redhat.com>

tests/diffusion/distributed/test_parallel_state_sp_groups.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for SP subgroup construction in parallel_state.py."""
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for SP subgroup construction for diffusion parallel state."""
 
 from __future__ import annotations
 
@@ -8,12 +8,55 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+import vllm.distributed.parallel_state as vllm_parallel_state
 
+from tests.helpers.runtime import get_distributed_init_method
+from vllm_omni.diffusion.distributed import parallel_state as omni_parallel_state
 from vllm_omni.diffusion.distributed.parallel_state import RankGenerator, set_seq_parallel_pg
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.core_model, pytest.mark.cpu]
 
 
+def ensure_parallel_state_initialized():
+    """Ensure torch distributed is initialized and that all world / count vars are set."""
+    assert torch.distributed.is_initialized()
+    assert omni_parallel_state._WORLD is not None
+    assert vllm_parallel_state._WORLD is not None
+    assert vllm_parallel_state._NODE_COUNT is not None
+
+
+def ensure_parallel_state_not_initialized():
+    """Ensure torch distributed is not initialized and that all world / count vars are unset."""
+    assert not torch.distributed.is_initialized()
+    assert omni_parallel_state._WORLD is None
+    assert vllm_parallel_state._WORLD is None
+    assert vllm_parallel_state._NODE_COUNT is None
+
+
+def test_omni_manages_vllm_distributed_state(monkeypatch):
+    """Verify Omni initializes and tears down vLLM's distributed state properly.
+
+    This is a regression test for ensuring vLLM Omni also sets up vLLM's vars correctly
+    to avoid potential misalignment in cases where Omni uses vLLM's native coordinator,
+    e.g., MoE + diffusion.
+    """
+    ensure_parallel_state_not_initialized()
+    monkeypatch.setattr(omni_parallel_state.current_omni_platform, "get_device_count", lambda: 1)
+    monkeypatch.setattr(omni_parallel_state.current_omni_platform, "set_device", lambda _device: None)
+    omni_parallel_state.init_distributed_environment(
+        world_size=1,
+        rank=0,
+        local_rank=0,
+        distributed_init_method=get_distributed_init_method(),
+        backend="gloo",
+    )
+    try:
+        ensure_parallel_state_initialized()
+    finally:
+        omni_parallel_state.destroy_distributed_env()
+    ensure_parallel_state_not_initialized()
+
+
 def _fake_new_group_factory(created_groups: list[SimpleNamespace]):
     def _fake_new_group(ranks, *args, **kwargs):
         group = SimpleNamespace(ranks=list(ranks))

vllm_omni/diffusion/distributed/parallel_state.py

@@ -438,6 +438,17 @@ def init_distributed_environment(
             "world group already initialized with a different world size"
         )
 
+    # Now, initialize vLLM's distributed environment; we need to do this before initializing
+    # model parallel in Omni. Otherwise, if we use vLLM's native coordinator, we may run into
+    # issues with _NODE_COUNT / _WORLD (in vLLM, not the Omni _WORLD var above) being unset.
+    vllm_parallel_state.init_distributed_environment(
+        world_size=world_size,
+        rank=rank,
+        distributed_init_method=distributed_init_method,
+        local_rank=local_rank,
+        backend=backend,
+    )
+
 
 def model_parallel_is_initialized():
     """Check if tensor and pipeline parallel groups are initialized."""
@@ -1061,8 +1072,8 @@ def destroy_distributed_environment():
     if _WORLD:
         _WORLD.destroy()
     _WORLD = None
-    if torch.distributed.is_initialized():
-        torch.distributed.destroy_process_group()
+    # We only need to destroy / clear the Omni world; this takes care of the torch process groups
+    vllm_parallel_state.destroy_distributed_environment()
 
 
 def destroy_distributed_env():