diff --git a/tests/diffusion/distributed/test_parallel_state_sp_groups.py b/tests/diffusion/distributed/test_parallel_state_sp_groups.py index c12882a22..0febb1925 100644 --- a/tests/diffusion/distributed/test_parallel_state_sp_groups.py +++ b/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)) diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 335418254..858a2e398 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/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():