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

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

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():