[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

@@ -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
 
 import importlib.util
 import shutil
@@ -31,6 +31,7 @@ _GROUP_GETTERS = (
     "get_cfg_group",
     "get_tp_group",
     "get_fs_group",
+    "get_hsdp_replicate_group",
     "get_ep_group",
 )
 
@@ -61,7 +62,7 @@ def test_setup_uses_default_name_before_model_parallel_init(
     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)
+    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",
@@ -73,18 +74,20 @@ def test_setup_uses_default_name_before_model_parallel_init(
 
 
 @pytest.mark.parametrize(
-    ("groups", "enable_ep", "expected_name"),
+    ("groups", "enable_ep", "use_hsdp", "hsdp_replicate_size", "expected_name"),
     [
-        ({}, False, "DiffusionWorker"),
-        ({"get_tp_group": _FakeGroup(2, 1)}, False, "DiffusionWorker_TP1"),
-        ({"get_dp_group": _FakeGroup(4, 2)}, False, "DiffusionWorker_DP2"),
-        ({"get_pp_group": _FakeGroup(2, 1)}, False, "DiffusionWorker_PP1"),
+        ({}, 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",
         ),
         (
@@ -94,22 +97,39 @@ def test_setup_uses_default_name_before_model_parallel_init(
                 "get_tp_group": _FakeGroup(2, 1),
             },
             False,
+            False,
+            1,
             "DiffusionWorker_DP1_SP2_TP1",
         ),
         (
-            {
-                "get_dp_group": _FakeGroup(4, 3),
-                "get_fs_group": _FakeGroup(2, 1),
-            },
+            {"get_fs_group": _FakeGroup(4, 3)},
             False,
-            "DiffusionWorker_DP3_FS1",
+            True,
+            1,
+            "DiffusionWorker_FS3",
         ),
         (
-            {"get_ep_group": _FakeGroup(4, 2)},
+            {
+                "get_sp_group": _FakeGroup(2, 1),
+                "get_fs_group": _FakeGroup(4, 3),
+            },
+            False,
             True,
-            "DiffusionWorker_EP2",
+            1,
+            "DiffusionWorker_SP1_FS3",
         ),
-        ({}, True, "DiffusionWorker"),
+        (
+            {
+                "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",
@@ -118,7 +138,9 @@ def test_setup_uses_default_name_before_model_parallel_init(
         "pp-only",
         "tp-cfg",
         "dp-sp-tp",
-        "standalone-hsdp",
+        "hsdp-only",
+        "hsdp-sp",
+        "hsdp-replicated",
         "expert-parallel",
         "singleton-ep",
     ],
@@ -127,6 +149,8 @@ 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(
@@ -138,13 +162,26 @@ def test_setup_uses_initialized_parallel_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)
+    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:
@@ -160,7 +197,7 @@ def test_missing_setproctitle_is_non_fatal(mocker: MockerFixture) -> None:
     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)
+    worker_module._setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=False, use_hsdp=False)
 
     decorate_logs.assert_called_once_with("DiffusionWorker")
 

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

@@ -51,6 +51,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
     get_dp_group,
     get_fs_group,
+    get_hsdp_replicate_group,
     get_pp_group,
     get_sp_group,
     init_distributed_environment,
@@ -110,38 +111,6 @@ def _all_gather_rank_values(value: Any) -> list[Any]:
     return values
 
 
-def _setup_diffusion_worker_proc_title_and_log_prefix(enable_ep: bool) -> 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()
-        fs_group = get_fs_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 fs_group.world_size > 1:
-            process_name += f"_FS{fs_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)
-
-
 def _run_and_gather_rank_values(operation: str, func: Callable[[], Any]) -> list[Any]:
     """Run one rank-local probe without stranding peers on local failure."""
 
@@ -158,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
@@ -334,10 +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)
+            _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
@@ -1440,7 +1455,11 @@ class WorkerProc:
 
         set_death_signal(signal.SIGTERM)
 
-        _setup_diffusion_worker_proc_title_and_log_prefix(enable_ep=od_config.parallel_config.enable_expert_parallel)
+        _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