Synchronize CPU ring workers before group teardown

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
(cherry picked from commit 4ad71d8270)
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/diffusion/distributed/test_comm.py

@@ -14,6 +14,7 @@ from typing import Literal
 import pytest
 import torch
 
+from tests.helpers.distributed import configure_cpu_collective_worker, start_collective_workers
 from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D, SeqAllToAll5D
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -56,6 +57,8 @@ def _init_worker(
             "MASTER_PORT": str(master_port),
         }
     )
+    if device_kind == "cpu":
+        configure_cpu_collective_worker()
     backend = "gloo" if device_kind == "cpu" else None
     init_distributed_environment(backend=backend)
 
@@ -227,6 +230,10 @@ def _run_ring_p2p(
             atol=atol,
             msg=f"[Rank {local_rank}] Ring P2P data mismatch",
         )
+        if device_kind == "cpu":
+            # P2P completion only synchronizes neighbors. Keep fast CPU ranks
+            # alive until every rank has finished constructing and using groups.
+            torch.distributed.barrier()
     finally:
         destroy_distributed_env()
 
@@ -243,7 +250,7 @@ def _spawn_4d_identity(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_4d_identity,
         args=(
             world_size,
@@ -257,6 +264,7 @@ def _spawn_4d_identity(
             master_port,
         ),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
 
@@ -272,7 +280,7 @@ def _spawn_5d_identity(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_5d_identity,
         args=(
             world_size,
@@ -286,6 +294,7 @@ def _spawn_5d_identity(
             master_port,
         ),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
 
@@ -299,10 +308,11 @@ def _spawn_ring_p2p(
     device_kind: DeviceKind,
     master_port: int,
 ) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_ring_p2p,
         args=(world_size, dtype, batch_size, num_heads, head_size, device_kind, master_port),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
 
@@ -499,8 +509,9 @@ def _run_fused_qkv_matches_three_4d(
 @pytest.mark.cpu
 @pytest.mark.parametrize("world_size", [2, 4])
 def test_fused_qkv_all_to_all_matches_three_4d_exchanges_cpu(world_size: int) -> None:
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         _run_fused_qkv_matches_three_4d,
         args=(world_size, "cpu", 29613),
         nprocs=world_size,
+        device_kind="cpu",
     )

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -19,6 +19,7 @@ from vllm.utils.network_utils import get_file_store_init_method
 from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
 
 import vllm_omni.diffusion.distributed.pipeline_parallel as pp_module
+from tests.helpers.distributed import collective_context, start_collective_workers
 from tests.helpers.mark import hardware_marks
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -39,6 +40,25 @@ DeviceKind = Literal["cpu", "cuda"]
 _UNIT_MARKS = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
+@pytest.fixture(autouse=True)
+def _restore_collective_platforms():
+    """CPU baselines run in the parent process; retain its original platforms."""
+    import vllm.platforms
+
+    from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
+
+    original_platform = vllm.platforms.current_platform
+    original_omni_platforms = {
+        module: module.current_omni_platform for module in (comm, group_coordinator, parallel_state)
+    }
+    try:
+        yield
+    finally:
+        vllm.platforms.current_platform = original_platform
+        for module, platform in original_omni_platforms.items():
+            module.current_omni_platform = platform
+
+
 def _cleanup_distributed() -> None:
     """Destroy omni distributed state left dangling after a test.
 
@@ -409,6 +429,10 @@ def init_dist(
     device = _worker_device(local_rank, device_kind)
     if device_kind == "cuda":
         current_omni_platform.set_device(device)
+    else:
+        from tests.helpers.distributed import configure_cpu_collective_worker
+
+        configure_cpu_collective_worker()
     backend = "gloo" if device_kind == "cpu" else None
     init_distributed_environment(
         world_size=world_size,
@@ -500,13 +524,14 @@ def isend_irecv_worker(
 
 
 def _run_isend_irecv(pp_size: int, device_kind: DeviceKind, init_method: str) -> None:
-    mp_context = torch.multiprocessing.get_context("spawn")
+    mp_context = collective_context(device_kind)
     manager = mp_context.Manager()
     q = manager.Queue()
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         isend_irecv_worker,
         args=(pp_size, init_method, device_kind, q),
         nprocs=pp_size,
+        device_kind=device_kind,
     )
     results = {label: tensor for label, tensor in [q.get(), q.get()]}
     torch.testing.assert_close(
@@ -656,15 +681,16 @@ def _run_predict_noise(
 
     baseline_out = compute_single_gpu_baseline(test_config, dtype, do_true_cfg, device_kind)
 
-    mp_context = torch.multiprocessing.get_context("spawn")
+    mp_context = collective_context(device_kind)
     manager = mp_context.Manager()
     pp_q = manager.Queue()
 
     world_size = pp_size * cfg_size
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         predict_noise_worker,
         args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, dtype, test_config, device_kind, pp_q),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
     pp_out = pp_q.get()
@@ -900,15 +926,16 @@ def _run_scheduler_step(
 
     baseline = compute_scheduler_step_baseline(test_config, do_true_cfg, device_kind)
 
-    mp_context = torch.multiprocessing.get_context("spawn")
+    mp_context = collective_context(device_kind)
     manager = mp_context.Manager()
     q = manager.Queue()
 
     world_size = pp_size * cfg_size
-    torch.multiprocessing.spawn(
+    start_collective_workers(
         scheduler_step_worker,
         args=(world_size, init_method, pp_size, cfg_size, do_true_cfg, test_config, device_kind, q),
         nprocs=world_size,
+        device_kind=device_kind,
     )
 
     result = q.get()

tests/helpers/cpu_collective_server.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Import preload for the CPU-only collective forkserver.
+
+This module is imported in the forkserver process, before it starts workers.
+Its CUDA visibility does not alter the pytest process or CUDA parity workers.
+"""
+
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+import torch  # noqa: E402
+
+# Retain the expensive worker-module imports across CPU process groups.
+from tests.diffusion.distributed import test_comm, test_pipeline_parallel  # noqa: E402, F401
+
+if torch.cuda.is_initialized():
+    raise RuntimeError("CPU collective forkserver must not initialize CUDA")

tests/helpers/distributed.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Process and platform setup for isolated collective test workers."""
+
+from collections.abc import Callable
+from multiprocessing.context import BaseContext
+
+import torch
+
+
+def collective_context(device_kind: str) -> BaseContext:
+    """Reuse imports in a clean CPU forkserver; CUDA workers use spawn."""
+    if device_kind == "cpu" and "forkserver" in torch.multiprocessing.get_all_start_methods():
+        torch.multiprocessing.set_forkserver_preload(["tests.helpers.cpu_collective_server"])
+        return torch.multiprocessing.get_context("forkserver")
+    return torch.multiprocessing.get_context("spawn")
+
+
+def start_collective_workers(
+    worker: Callable[..., None], *, args: tuple[object, ...], nprocs: int, device_kind: str
+) -> None:
+    context = collective_context(device_kind)
+    if context.get_start_method() == "spawn":
+        torch.multiprocessing.spawn(worker, args=args, nprocs=nprocs)
+    else:
+        torch.multiprocessing.start_processes(worker, args=args, nprocs=nprocs, start_method=context.get_start_method())
+
+
+def configure_cpu_collective_worker() -> None:
+    """Select CPU platforms inside a spawned worker that only uses Gloo.
+
+    Upstream coordinators resolve local ranks through the selected
+    platform's visible devices. CPU ranks can outnumber visible GPUs, so they
+    must use the CPU platform even when the test host has CUDA installed.
+    The worker still creates real process groups and runs real collectives.
+    """
+    import vllm.platforms
+    from vllm.platforms.cpu import CpuPlatform
+
+    from vllm_omni.diffusion.distributed import comm, group_coordinator, parallel_state
+    from vllm_omni.platforms.interface import UnspecifiedOmniPlatform
+
+    class CpuCollectivePlatform(UnspecifiedOmniPlatform):
+        dist_backend = "gloo"
+
+        @classmethod
+        def get_device_count(cls) -> int:
+            # CPU tensors share one logical device regardless of process rank.
+            return 1
+
+        @classmethod
+        def set_device(cls, device: torch.device) -> None:
+            assert device.type == "cpu"
+
+        @classmethod
+        def synchronize(cls) -> None:
+            pass
+
+    vllm.platforms.current_platform = CpuPlatform()
+    platform = CpuCollectivePlatform()
+    for module in (comm, group_coordinator, parallel_state):
+        module.current_omni_platform = platform

tests/worker_v2/test_omni_gpu_model_runner.py

@@ -13,7 +13,6 @@ from vllm import SamplingParams
 from vllm.config.compilation import CUDAGraphMode
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
-from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 
 from vllm_omni.config.model import OmniModelConfig
@@ -164,28 +163,6 @@ def test_init_model_state_factory_dispatches_omni_only(monkeypatch, flag, expect
         assert state is upstream.return_value
 
 
-@pytest.mark.parametrize("omni_stage", [False, True])
-@pytest.mark.parametrize("specialized", [False, True])
-def test_auxiliary_default_state_has_omni_lifecycle_but_preserves_upstream_states(monkeypatch, omni_stage, specialized):
-    class SpecializedState(DefaultModelState):
-        pass
-
-    state_cls = SpecializedState if specialized else DefaultModelState
-    upstream_state = object.__new__(state_cls)
-    monkeypatch.setattr("vllm_omni.worker_v2.model_states._upstream_init_model_state", lambda *_args: upstream_state)
-    monkeypatch.setattr(OmniModelState, "__init__", lambda *_args: None)
-    model_config = object.__new__(OmniModelConfig) if omni_stage else SimpleNamespace()
-    state = init_omni_model_state(
-        SimpleNamespace(model_config=model_config), SimpleNamespace(), None, torch.device("cpu")
-    )
-    if omni_stage and not specialized:
-        assert isinstance(state, OmniModelState)
-        assert callable(state.run_preprocess)
-        assert callable(state.postprocess_model_output)
-    else:
-        assert state is upstream_state
-
-
 def test_finish_requests_notifies_model_and_cleans_only_known_slots(monkeypatch):
     runner = _make_runner()
     calls = []

vllm_omni/worker_v2/model_states/__init__.py

@@ -4,8 +4,7 @@
 """Omni-aware ``init_model_state`` factory.
 
 Extends the upstream v2 factory with Omni architecture dispatch.
-Non-Omni architectures use upstream dispatch; its default state acquires
-Omni lifecycle support when owned by an Omni stage.
+Non-Omni architectures fall through to the upstream ``init_model_state``.
 """
 
 from __future__ import annotations
@@ -17,11 +16,8 @@ from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.model_states import (
     init_model_state as _upstream_init_model_state,
 )
-from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 
-from vllm_omni.config.model import OmniModelConfig
-
 
 def init_omni_model_state(
     vllm_config: VllmConfig,
@@ -32,8 +28,8 @@ def init_omni_model_state(
     """Create the appropriate ``ModelState`` for *model*.
 
     Returns an ``OmniModelState`` when the model declares Omni lifecycle
-    capabilities. Otherwise delegates to the upstream v2 factory, wrapping
-    its default state for auxiliary Omni stages.
+    capabilities; otherwise
+    delegates to the upstream v2 factory.
     """
     factory = getattr(model, "create_omni_model_state", None)
     if factory is not None:
@@ -53,12 +49,4 @@ def init_omni_model_state(
         )
 
         return OmniModelState(vllm_config, model, encoder_cache, device)
-    state = _upstream_init_model_state(vllm_config, model, encoder_cache, device)
-    if type(state) is DefaultModelState and isinstance(getattr(vllm_config, "model_config", None), OmniModelConfig):
-        # Auxiliary stages such as the forced aligner have no model-owned
-        # hooks, but Omni's runner still needs its request/output lifecycle.
-        # Keep specialized upstream states and plain vLLM configs intact.
-        from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
-
-        return OmniModelState(vllm_config, model, encoder_cache, device)
-    return state
+    return _upstream_init_model_state(vllm_config, model, encoder_cache, device)