diff --git a/tests/diffusion/test_diffusion_worker_process_title.py b/tests/diffusion/test_diffusion_worker_process_title.py index da090447c..83863bb86 100644 --- a/tests/diffusion/test_diffusion_worker_process_title.py +++ b/tests/diffusion/test_diffusion_worker_process_title.py @@ -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: diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 7e1b57786..7d219ebbe 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -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 diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index a36ba8e6f..de8111c40 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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, @@ -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()