[Core] Validate HSDP shard size before initialization

Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>

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

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)

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

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