diff --git a/tests/diffusion/distributed/test_comm.py b/tests/diffusion/distributed/test_comm.py index 21cb18821..a8d8a032a 100644 --- a/tests/diffusion/distributed/test_comm.py +++ b/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) @@ -243,7 +246,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 +260,7 @@ def _spawn_4d_identity( master_port, ), nprocs=world_size, + device_kind=device_kind, ) @@ -272,7 +276,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 +290,7 @@ def _spawn_5d_identity( master_port, ), nprocs=world_size, + device_kind=device_kind, ) @@ -299,10 +304,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 +505,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", ) diff --git a/tests/diffusion/distributed/test_pipeline_parallel.py b/tests/diffusion/distributed/test_pipeline_parallel.py index 7cebbcc53..c64e880cb 100644 --- a/tests/diffusion/distributed/test_pipeline_parallel.py +++ b/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() diff --git a/tests/helpers/cpu_collective_server.py b/tests/helpers/cpu_collective_server.py new file mode 100644 index 000000000..30cc80431 --- /dev/null +++ b/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") diff --git a/tests/helpers/distributed.py b/tests/helpers/distributed.py new file mode 100644 index 000000000..1cc79753b --- /dev/null +++ b/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