diff --git a/tests/diffusion/distributed/test_expert_parallel_layout.py b/tests/diffusion/distributed/test_expert_parallel_layout.py index d7b66be5d..bfbbca588 100644 --- a/tests/diffusion/distributed/test_expert_parallel_layout.py +++ b/tests/diffusion/distributed/test_expert_parallel_layout.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace @@ -194,6 +194,47 @@ def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch): assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group] +@pytest.mark.cpu +@pytest.mark.core_model +@pytest.mark.parametrize( + ("fully_shard_degree", "error_message"), + [(0, "fully_shard_degree must be positive"), (3, "must be divisible by fully_shard_degree")], +) +def test_invalid_hsdp_shard_size_fails_before_group_creation(monkeypatch, fully_shard_degree: int, error_message: str): + """Reject invalid HSDP shard sizes before allocating any process groups.""" + world_size = 4 + created_groups: list[object] = [] + + def fail_if_group_created(*args, **kwargs): + del args, kwargs + created_groups.append(object()) + raise AssertionError("a process group was created before HSDP validation") + + fake_world_group = SimpleNamespace( + rank_in_group=0, + local_rank=0, + device_group=object(), + ) + monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size) + monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) + monkeypatch.setattr(parallel_state, "init_model_parallel_group", fail_if_group_created) + + for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_HSDP_REPLICATE", "_EXPERT_PARALLEL_GROUP_RANKS"): + monkeypatch.setattr(parallel_state, name, None) + for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"): + monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) + + with pytest.raises(ValueError, match=error_message): + parallel_state.initialize_model_parallel( + fully_shard_degree=fully_shard_degree, + use_hsdp=True, + backend="gloo", + ) + + assert created_groups == [] + + @pytest.mark.cpu @pytest.mark.core_model def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch): diff --git a/tests/diffusion/test_diffusion_worker_process_title.py b/tests/diffusion/test_diffusion_worker_process_title.py new file mode 100644 index 000000000..83863bb86 --- /dev/null +++ b/tests/diffusion/test_diffusion_worker_process_title.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import importlib.util +import shutil +import subprocess +import sys +import textwrap +import time +from dataclasses import dataclass +from unittest.mock import Mock + +import pytest +from pytest_mock import MockerFixture + +import vllm_omni.diffusion.worker.diffusion_worker as worker_module + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@dataclass(frozen=True) +class _FakeGroup: + world_size: int = 1 + rank_in_group: int = 0 + + +_GROUP_GETTERS = ( + "get_dp_group", + "get_pp_group", + "get_sp_group", + "get_cfg_group", + "get_tp_group", + "get_fs_group", + "get_hsdp_replicate_group", + "get_ep_group", +) + + +def _patch_groups( + mocker: MockerFixture, + groups: dict[str, _FakeGroup], +) -> dict[str, Mock]: + patched_getters = {} + for getter_name in _GROUP_GETTERS: + patched_getters[getter_name] = mocker.patch.object( + worker_module, + getter_name, + return_value=groups.get(getter_name, _FakeGroup()), + ) + return patched_getters + + +def test_setup_uses_default_name_before_model_parallel_init( + mocker: MockerFixture, +) -> None: + mocker.patch.object( + worker_module, + "model_parallel_is_initialized", + return_value=False, + ) + group_getters = _patch_groups(mocker, {}) + set_process_title = mocker.patch.object(worker_module, "set_process_title") + decorate_logs = mocker.patch.object(worker_module, "decorate_logs") + + worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=True, use_hsdp=True) + + set_process_title.assert_called_once_with( + name="DiffusionWorker", + prefix="vLLM-Omni", + ) + decorate_logs.assert_called_once_with("DiffusionWorker") + for group_getter in group_getters.values(): + group_getter.assert_not_called() + + +@pytest.mark.parametrize( + ("groups", "enable_ep", "use_hsdp", "hsdp_replicate_size", "expected_name"), + [ + ({}, False, False, 1, "DiffusionWorker"), + ({"get_tp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_TP1"), + ({"get_dp_group": _FakeGroup(4, 2)}, False, False, 1, "DiffusionWorker_DP2"), + ({"get_pp_group": _FakeGroup(2, 1)}, False, False, 1, "DiffusionWorker_PP1"), + ( + { + "get_cfg_group": _FakeGroup(2, 1), + "get_tp_group": _FakeGroup(2, 1), + }, + False, + False, + 1, + "DiffusionWorker_CFG1_TP1", + ), + ( + { + "get_dp_group": _FakeGroup(2, 1), + "get_sp_group": _FakeGroup(4, 2), + "get_tp_group": _FakeGroup(2, 1), + }, + False, + False, + 1, + "DiffusionWorker_DP1_SP2_TP1", + ), + ( + {"get_fs_group": _FakeGroup(4, 3)}, + False, + True, + 1, + "DiffusionWorker_FS3", + ), + ( + { + "get_sp_group": _FakeGroup(2, 1), + "get_fs_group": _FakeGroup(4, 3), + }, + False, + True, + 1, + "DiffusionWorker_SP1_FS3", + ), + ( + { + "get_fs_group": _FakeGroup(1, 0), + "get_hsdp_replicate_group": _FakeGroup(2, 1), + }, + False, + True, + 2, + "DiffusionWorker_RP1", + ), + ({"get_ep_group": _FakeGroup(4, 2)}, True, False, 1, "DiffusionWorker_EP2"), + ({}, True, False, 1, "DiffusionWorker"), + ], + ids=[ + "all-singleton", + "tp-only", + "dp-only", + "pp-only", + "tp-cfg", + "dp-sp-tp", + "hsdp-only", + "hsdp-sp", + "hsdp-replicated", + "expert-parallel", + "singleton-ep", + ], +) +def test_setup_uses_initialized_parallel_groups( + mocker: MockerFixture, + groups: dict[str, _FakeGroup], + enable_ep: bool, + use_hsdp: bool, + hsdp_replicate_size: int, + expected_name: str, +) -> None: + mocker.patch.object( + worker_module, + "model_parallel_is_initialized", + return_value=True, + ) + group_getters = _patch_groups(mocker, groups) + set_process_title = mocker.patch.object(worker_module, "set_process_title") + decorate_logs = mocker.patch.object(worker_module, "decorate_logs") + + worker_module._setup_diffusion_worker_proc_title_and_log_prefix( + enable_ep=enable_ep, + use_hsdp=use_hsdp, + hsdp_replicate_size=hsdp_replicate_size, + ) + + set_process_title.assert_called_once_with( + name=expected_name, + prefix="vLLM-Omni", + ) + decorate_logs.assert_called_once_with(expected_name) + if use_hsdp: + group_getters["get_fs_group"].assert_called_once_with() + if hsdp_replicate_size > 1: + group_getters["get_hsdp_replicate_group"].assert_called_once_with() + else: + group_getters["get_hsdp_replicate_group"].assert_not_called() + else: + group_getters["get_fs_group"].assert_not_called() + group_getters["get_hsdp_replicate_group"].assert_not_called() + if enable_ep: + group_getters["get_ep_group"].assert_called_once_with() + else: + group_getters["get_ep_group"].assert_not_called() + + +def test_missing_setproctitle_is_non_fatal(mocker: MockerFixture) -> None: + mocker.patch.object( + worker_module, + "model_parallel_is_initialized", + return_value=False, + ) + mocker.patch.dict(sys.modules, {"setproctitle": None}) + decorate_logs = mocker.patch.object(worker_module, "decorate_logs") + + worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=False, use_hsdp=False) + + decorate_logs.assert_called_once_with("DiffusionWorker") + + +@pytest.mark.skipif( + sys.platform != "linux" or shutil.which("ps") is None or importlib.util.find_spec("setproctitle") is None, + reason="requires Linux ps and setproctitle", +) +def test_process_title_is_visible_through_ps() -> None: + expected_name = "vLLM-Omni::DiffusionWorker_TP1" + child_script = textwrap.dedent( + """ + import time + + from vllm.utils.system_utils import set_process_title + + + set_process_title( + name="DiffusionWorker_TP1", + prefix="vLLM-Omni", + ) + time.sleep(30) + """ + ) + process = subprocess.Popen( + [sys.executable, "-c", child_script], + stderr=subprocess.PIPE, + text=True, + ) + + observed_title = "" + try: + for _ in range(100): + if process.poll() is not None: + stderr = process.stderr.read() if process.stderr is not None else "" + pytest.fail(f"child process exited before title check: {stderr}") + result = subprocess.run( + ["ps", "-o", "args=", "-p", str(process.pid)], + check=False, + capture_output=True, + text=True, + ) + observed_title = result.stdout.strip() + if expected_name in observed_title: + break + time.sleep(0.1) + else: + pytest.fail(f"expected {expected_name!r} in process title, got {observed_title!r}") + finally: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5) diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 76d458848..85ba9cd6c 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project # Adapted from: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/distributed/parallel_state.py # Copyright 2023 The vLLM team. # Adapted from @@ -59,6 +60,8 @@ _SP: SequenceParallelGroupCoordinator | None = None _PP: PipelineGroupCoordinator | None = None _CFG: GroupCoordinator | None = None _DP: GroupCoordinator | None = None +_FS: GroupCoordinator | None = None # Fully Sharded (HSDP shard dimension) +_HSDP_REPLICATE: GroupCoordinator | None = None # HSDP replica dimension # Rank-layout metadata for expert parallelism. This is not a process group; # it is reused by platform-specific runtimes that must build companion groups @@ -350,6 +353,17 @@ def get_data_parallel_rank(): return get_dp_group().rank_in_group +# FS (Fully Shard / HSDP shard dimension) +def get_fs_group() -> GroupCoordinator: + assert _FS is not None, "fully shard group is not initialized" + return _FS + + +def get_hsdp_replicate_group() -> GroupCoordinator: + assert _HSDP_REPLICATE is not None, "HSDP replicate group is not initialized" + return _HSDP_REPLICATE + + def is_dp_last_group(): """Return True if in the last data parallel group, False otherwise.""" return ( @@ -445,6 +459,7 @@ def init_model_parallel_group( "expert", "sequence", "classifier_free_guidance", + "fully_shard", ], f"parallel_mode {parallel_mode} is not supported" if parallel_mode == "pipeline": return PipelineGroupCoordinator( @@ -679,10 +694,13 @@ def _initialize_model_parallel( allgather_degree: int = 1, tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, + fully_shard_degree: int = 1, enable_expert_parallel: bool = False, use_hsdp: bool = False, backend: str | None = None, ) -> None: + global _FS, _HSDP_REPLICATE + if backend is None: backend = current_omni_platform.dist_backend """ @@ -700,6 +718,7 @@ def _initialize_model_parallel( (causal=False only). Mutually exclusive with ulysses/ring in v1. tensor_parallel_size: number of GPUs used for tensor parallelism. pipeline_parallel_size: number of GPUs used for pipeline parallelism. + fully_shard_degree: number of GPUs used for the HSDP shard dimension. backend: distributed backend of pytorch collective comm. Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we @@ -763,6 +782,12 @@ def _initialize_model_parallel( raise ValueError("HSDP (FSDP2) requires data_parallel_size to be 1") if non_dp_size not in (1, world_size): raise ValueError(f"HSDP non-DP parallel size must be 1 or WORLD size ({world_size}), but got {non_dp_size}") + if fully_shard_degree <= 0: + raise ValueError(f"fully_shard_degree must be positive, got {fully_shard_degree}") + if world_size % fully_shard_degree != 0: + raise ValueError( + f"WORLD size ({world_size}) must be divisible by fully_shard_degree ({fully_shard_degree})" + ) data_parallel_size = 1 else: inferred_data_parallel_size = world_size // non_dp_size @@ -881,6 +906,32 @@ def _initialize_model_parallel( group_name="dp", ) + if use_hsdp: + assert _FS is None, "fully shard group is already initialized" + # HSDP builds its mesh from arange(world_size).reshape(replicate, shard), + # so each consecutive rank run is one fully-sharded group. + fs_group_ranks = [ + list(range(start, start + fully_shard_degree)) for start in range(0, world_size, fully_shard_degree) + ] + _FS = init_model_parallel_group( + group_ranks=fs_group_ranks, + local_rank=get_world_group().local_rank, + backend=backend, + parallel_mode="fully_shard", + ) + if world_size > fully_shard_degree: + # The HSDP mesh's columns are replica groups: each column contains + # the same shard position across all replica rows. + hsdp_replicate_group_ranks = [ + list(range(offset, world_size, fully_shard_degree)) for offset in range(fully_shard_degree) + ] + _HSDP_REPLICATE = init_model_parallel_group( + group_ranks=hsdp_replicate_group_ranks, + local_rank=get_world_group().local_rank, + backend=backend, + parallel_mode="fully_shard", + ) + global _EXPERT_PARALLEL_GROUP_RANKS _EXPERT_PARALLEL_GROUP_RANKS = get_rank_groups("tp-sp-cfg-dp") if use_moe_parallel_mapping: @@ -902,6 +953,7 @@ def initialize_model_parallel( allgather_degree: int = 1, tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, + fully_shard_degree: int = 1, enable_expert_parallel: bool = False, use_hsdp: bool = False, backend: str | None = None, @@ -916,6 +968,8 @@ def initialize_model_parallel( "cfg": _CFG, "sp": _SP, "pp": _PP, + "fs": _FS, + "hsdp_replicate": _HSDP_REPLICATE, "vllm_tp": vllm_parallel_state._TP, "vllm_dp": vllm_parallel_state._DP, "vllm_pp": vllm_parallel_state._PP, @@ -937,6 +991,7 @@ def initialize_model_parallel( allgather_degree=allgather_degree, tensor_parallel_size=tensor_parallel_size, pipeline_parallel_size=pipeline_parallel_size, + fully_shard_degree=fully_shard_degree, enable_expert_parallel=enable_expert_parallel, use_hsdp=use_hsdp, backend=backend, @@ -948,7 +1003,7 @@ def initialize_model_parallel( def destroy_model_parallel(): """Set the groups to none and destroy them.""" - global _DP, _CFG, _SP, _PP, _EXPERT_PARALLEL_GROUP_RANKS + global _DP, _CFG, _SP, _PP, _FS, _HSDP_REPLICATE, _EXPERT_PARALLEL_GROUP_RANKS if vllm_parallel_state._DP and vllm_parallel_state._DP is not _DP: vllm_parallel_state._DP.destroy() @@ -958,6 +1013,14 @@ def destroy_model_parallel(): _DP.destroy() _DP = None + if _FS: + _FS.destroy() + _FS = None + + if _HSDP_REPLICATE: + _HSDP_REPLICATE.destroy() + _HSDP_REPLICATE = None + if _CFG: _CFG.destroy() _CFG = None diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index f71c2bf8a..de8111c40 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -25,10 +25,12 @@ import torch.distributed as dist import zmq from vllm.config import VllmConfig, set_current_vllm_config from vllm.distributed.device_communicators.shm_broadcast import MessageQueue +from vllm.distributed.parallel_state import get_ep_group, get_tp_group from vllm.logger import init_logger from vllm.profiler.wrapper import CudaProfilerWrapper, WorkerProfiler from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.utils.mem_utils import GiB_bytes, MemorySnapshot, format_gib, memory_profiling +from vllm.utils.system_utils import decorate_logs, set_process_title from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager @@ -46,8 +48,15 @@ from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, + get_cfg_group, + get_dp_group, + get_fs_group, + get_hsdp_replicate_group, + get_pp_group, + get_sp_group, init_distributed_environment, initialize_model_parallel, + model_parallel_is_initialized, ) from vllm_omni.diffusion.forward_context import set_forward_context from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm @@ -118,6 +127,47 @@ def _run_and_gather_rank_values(operation: str, func: Callable[[], Any]) -> list return [result for _, result in rank_results] +def _setup_diffusion_worker_proc_title_and_log_prefix( + enable_ep: bool, + use_hsdp: bool, + hsdp_replicate_size: int = 1, +) -> None: + """Set the worker process title and log prefix from initialized groups.""" + process_name = "DiffusionWorker" + if model_parallel_is_initialized(): + dp_group = get_dp_group() + pp_group = get_pp_group() + sp_group = get_sp_group() + cfg_group = get_cfg_group() + tp_group = get_tp_group() + + if dp_group.world_size > 1: + process_name += f"_DP{dp_group.rank_in_group}" + if pp_group.world_size > 1: + process_name += f"_PP{pp_group.rank_in_group}" + if sp_group.world_size > 1: + process_name += f"_SP{sp_group.rank_in_group}" + if cfg_group.world_size > 1: + process_name += f"_CFG{cfg_group.rank_in_group}" + if tp_group.world_size > 1: + process_name += f"_TP{tp_group.rank_in_group}" + if use_hsdp: + fs_group = get_fs_group() + if fs_group.world_size > 1: + process_name += f"_FS{fs_group.rank_in_group}" + if hsdp_replicate_size > 1: + replicate_group = get_hsdp_replicate_group() + if replicate_group.world_size > 1: + process_name += f"_RP{replicate_group.rank_in_group}" + if enable_ep: + ep_group = get_ep_group() + if ep_group.world_size > 1: + process_name += f"_EP{ep_group.rank_in_group}" + + set_process_title(name=process_name, prefix="vLLM-Omni") + decorate_logs(process_name) + + @contextmanager def _force_cutlass_fp8_linear_kernel(quant_config: object | None) -> Iterator[None]: import vllm.model_executor.layers.quantization.modelopt as vllm_modelopt @@ -294,9 +344,15 @@ class DiffusionWorker: allgather_degree=parallel_config.allgather_degree, tensor_parallel_size=parallel_config.tensor_parallel_size, pipeline_parallel_size=parallel_config.pipeline_parallel_size, + fully_shard_degree=parallel_config.hsdp_shard_size if parallel_config.use_hsdp else 1, enable_expert_parallel=parallel_config.enable_expert_parallel, use_hsdp=parallel_config.use_hsdp, ) + _setup_diffusion_worker_proc_title_and_log_prefix( + enable_ep=parallel_config.enable_expert_parallel, + use_hsdp=parallel_config.use_hsdp, + hsdp_replicate_size=parallel_config.hsdp_replicate_size, + ) if ( getattr(self.od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) is DiffusionKVCacheMode.PAGED_SCHEDULER @@ -1399,13 +1455,11 @@ class WorkerProc: set_death_signal(signal.SIGTERM) - # Set process title for visibility in nvidia-smi / htop (optional, non-fatal) - try: - import setproctitle - - setproctitle.setproctitle(f"vLLM-Omni::DiffusionWorker-{rank}") - except ImportError: - pass # setproctitle not installed, skip process title setting + _setup_diffusion_worker_proc_title_and_log_prefix( + enable_ep=od_config.parallel_config.enable_expert_parallel, + use_hsdp=od_config.parallel_config.use_hsdp, + hsdp_replicate_size=od_config.parallel_config.hsdp_replicate_size, + ) load_omni_general_plugins() worker_proc = None