anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Core] Identify HSDP replica ranks in worker titles

Browse Source 
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
This commit is contained in:
hsliuustc0106 committed 2026-08-28 18:27:56 +08:00
1 parent 65ac6794f9
commit 5c8e0e85d8
3 files changed
+71 -11
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/test_diffusion_worker_process_title.py
+34 -8
View File
@@ -31,6 +31,7 @@ _GROUP_GETTERS = (
"get_cfg_group",
"get_tp_group",
"get_fs_group",
"get_hsdp_replicate_group",
"get_ep_group",
)
@@ -73,12 +74,12 @@ def test_setup_uses_default_name_before_model_parallel_init(
@pytest.mark.parametrize(
("groups", "enable_ep", "use_hsdp", "expected_name"),
("groups", "enable_ep", "use_hsdp", "hsdp_replicate_size", "expected_name"),
[
({}, False, False, "DiffusionWorker"),
({"get_tp_group": _FakeGroup(2, 1)}, False, False, "DiffusionWorker_TP1"),
({"get_dp_group": _FakeGroup(4, 2)}, False, False, "DiffusionWorker_DP2"),
({"get_pp_group": _FakeGroup(2, 1)}, False, 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),
@@ -86,6 +87,7 @@ def test_setup_uses_default_name_before_model_parallel_init(
},
False,
False,
1,
"DiffusionWorker_CFG1_TP1",
),
(
@@ -96,12 +98,14 @@ def test_setup_uses_default_name_before_model_parallel_init(
},
False,
False,
1,
"DiffusionWorker_DP1_SP2_TP1",
),
(
{"get_fs_group": _FakeGroup(4, 3)},
False,
True,
1,
"DiffusionWorker_FS3",
),
(
@@ -111,10 +115,21 @@ def test_setup_uses_default_name_before_model_parallel_init(
},
False,
True,
1,
"DiffusionWorker_SP1_FS3",
),
({"get_ep_group": _FakeGroup(4, 2)}, True, False, "DiffusionWorker_EP2"),
({}, True, False, "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",
@@ -125,6 +140,7 @@ def test_setup_uses_default_name_before_model_parallel_init(
"dp-sp-tp",
"hsdp-only",
"hsdp-sp",
"hsdp-replicated",
"expert-parallel",
"singleton-ep",
],
@@ -134,6 +150,7 @@ def test_setup_uses_initialized_parallel_groups(
groups: dict[str, _FakeGroup],
enable_ep: bool,
use_hsdp: bool,
hsdp_replicate_size: int,
expected_name: str,
) -> None:
mocker.patch.object(
@@ -145,7 +162,11 @@ 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, use_hsdp=use_hsdp)
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,
@@ -154,8 +175,13 @@ def test_setup_uses_initialized_parallel_groups(
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:
vllm_omni/diffusion/distributed/parallel_state.py
+25 -2
View File
@@ -61,6 +61,7 @@ _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
@@ -358,6 +359,11 @@ def get_fs_group() -> GroupCoordinator:
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 (
@@ -693,7 +699,7 @@ def _initialize_model_parallel(
use_hsdp: bool = False,
backend: str | None = None,
) -> None:
global _FS
global _FS, _HSDP_REPLICATE
if backend is None:
backend = current_omni_platform.dist_backend
@@ -914,6 +920,18 @@ def _initialize_model_parallel(
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")
@@ -952,6 +970,7 @@ def initialize_model_parallel(
"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,
@@ -985,7 +1004,7 @@ def initialize_model_parallel(
def destroy_model_parallel():
"""Set the groups to none and destroy them."""
global _DP, _CFG, _SP, _PP, _FS, _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()
@@ -999,6 +1018,10 @@ def destroy_model_parallel():
_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
+12 -1
View File
@@ -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,
@@ -126,7 +127,11 @@ 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) -> None:
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():
@@ -150,6 +155,10 @@ def _setup_diffusion_worker_proc_title_and_log_prefix(enable_ep: bool, 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:
@@ -342,6 +351,7 @@ class DiffusionWorker:
_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)
@@ -1448,6 +1458,7 @@ class WorkerProc:
_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()