From a25636d7803a2a5da22cb5beae2363eef1f09fad Mon Sep 17 00:00:00 2001 From: Jiaqian Liu <61532106+Celeste-jq@users.noreply.github.com> Date: Thu, 2 Jul 2026 12:14:12 +0800 Subject: [PATCH] [NPU] Fix HunyuanImage3 MoE EP+SP parallel mapping (#4344) Signed-off-by: Celeste-jq <591998922@qq.com> --- .../test_expert_parallel_layout.py | 282 ++++++++++++++++++ .../diffusion/distributed/parallel_state.py | 81 +++-- .../hunyuan_image3_transformer.py | 5 +- .../diffusion/worker/diffusion_worker.py | 15 +- .../platforms/npu/models/hunyuan_fused_moe.py | 34 +-- 5 files changed, 362 insertions(+), 55 deletions(-) create mode 100644 tests/diffusion/distributed/test_expert_parallel_layout.py diff --git a/tests/diffusion/distributed/test_expert_parallel_layout.py b/tests/diffusion/distributed/test_expert_parallel_layout.py new file mode 100644 index 000000000..8fc1db605 --- /dev/null +++ b/tests/diffusion/distributed/test_expert_parallel_layout.py @@ -0,0 +1,282 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest + +import vllm_omni.diffusion.distributed.parallel_state as parallel_state + + +class _FakeGroup: + def __init__( + self, + group_ranks: list[list[int]], + local_rank: int, + parallel_mode: str, + **kwargs, + ) -> None: + self.group_ranks = group_ranks + self.parallel_mode = parallel_mode + self.device_group = object() + self.ulysses_group = kwargs.get("ulysses_group") + self.ring_group = kwargs.get("ring_group") + self.local_group = next(group for group in group_ranks if local_rank in group) + self.world_size = len(self.local_group) + self.rank_in_group = self.local_group.index(local_rank) + + def destroy(self) -> None: + pass + + +@pytest.mark.cpu +@pytest.mark.core_model +def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch): + """MoE+EP rank layout should map SP->PCP, CFG*DP->DP, and TP*SP*CFG*DP->EP.""" + local_rank = 0 + world_size = 32 + created_groups: list[_FakeGroup] = [] + + def fake_init_model_parallel_group( + group_ranks, + local_rank, + backend, + parallel_mode=None, + group_name=None, + **kwargs, + ): + del backend, group_name + group = _FakeGroup( + [list(ranks) for ranks in group_ranks], + local_rank, + parallel_mode or "", + **kwargs, + ) + created_groups.append(group) + return group + + fake_world_group = SimpleNamespace( + rank_in_group=local_rank, + local_rank=local_rank, + device_group=object(), + ) + fake_forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(is_moe=True)) + 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.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo") + monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks)) + monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) + monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context) + monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) + monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) + + for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_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) + + parallel_state.initialize_model_parallel( + tensor_parallel_size=2, + sequence_parallel_size=2, + ulysses_degree=2, + ring_degree=1, + pipeline_parallel_size=2, + cfg_parallel_size=2, + data_parallel_size=2, + enable_expert_parallel=True, + backend="gloo", + ) + + assert parallel_state.vllm_parallel_state._PCP is parallel_state._SP + assert parallel_state.vllm_parallel_state._PCP.world_size == 2 + assert parallel_state._DP.world_size == 2 + assert parallel_state.vllm_parallel_state._DP is not parallel_state._DP + assert parallel_state.vllm_parallel_state._DP.world_size == 4 + assert parallel_state.vllm_parallel_state._EP.world_size == 16 + assert parallel_state.vllm_parallel_state._TP.world_size == 2 + assert parallel_state._PP.world_size == 2 + + assert parallel_state.vllm_parallel_state._PCP.local_group == [0, 2] + assert parallel_state.vllm_parallel_state._DP.local_group == [0, 8, 16, 24] + assert parallel_state.vllm_parallel_state._EP.local_group == [ + 0, + 1, + 2, + 3, + 8, + 9, + 10, + 11, + 16, + 17, + 18, + 19, + 24, + 25, + 26, + 27, + ] + assert parallel_state.get_expert_parallel_group_ranks() == [ + [ + 0, + 1, + 2, + 3, + 8, + 9, + 10, + 11, + 16, + 17, + 18, + 19, + 24, + 25, + 26, + 27, + ], + [ + 4, + 5, + 6, + 7, + 12, + 13, + 14, + 15, + 20, + 21, + 22, + 23, + 28, + 29, + 30, + 31, + ], + ] + + ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "expert"] + assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group] + + +@pytest.mark.cpu +@pytest.mark.core_model +def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch): + """vLLM DP should keep diffusion DP when expert parallelism is not enabled.""" + local_rank = 0 + world_size = 8 + + def fake_init_model_parallel_group( + group_ranks, + local_rank, + backend, + parallel_mode=None, + group_name=None, + **kwargs, + ): + del backend, group_name + return _FakeGroup( + [list(ranks) for ranks in group_ranks], + local_rank, + parallel_mode or "", + **kwargs, + ) + + fake_world_group = SimpleNamespace( + rank_in_group=local_rank, + local_rank=local_rank, + 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.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo") + monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks)) + monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) + monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) + monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) + + for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_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) + + parallel_state.initialize_model_parallel( + tensor_parallel_size=2, + sequence_parallel_size=1, + ulysses_degree=1, + ring_degree=1, + pipeline_parallel_size=1, + cfg_parallel_size=2, + data_parallel_size=2, + enable_expert_parallel=False, + backend="gloo", + ) + + assert parallel_state._DP.world_size == 2 + assert parallel_state.vllm_parallel_state._DP is parallel_state._DP + assert parallel_state.vllm_parallel_state._DP.world_size == 2 + assert parallel_state.vllm_parallel_state._DP.local_group == [0, 4] + assert parallel_state.vllm_parallel_state._PCP is None + assert parallel_state.vllm_parallel_state._EP is None + assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None + + +@pytest.mark.cpu +@pytest.mark.core_model +def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch): + """Non-MoE diffusion configs should not create vLLM PCP/DP/EP remap state.""" + local_rank = 0 + world_size = 8 + + def fake_init_model_parallel_group( + group_ranks, + local_rank, + backend, + parallel_mode=None, + group_name=None, + **kwargs, + ): + del backend, group_name + return _FakeGroup( + [list(ranks) for ranks in group_ranks], + local_rank, + parallel_mode or "", + **kwargs, + ) + + fake_world_group = SimpleNamespace( + rank_in_group=local_rank, + local_rank=local_rank, + device_group=object(), + ) + fake_forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(is_moe=False)) + 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.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo") + monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks)) + monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) + monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context) + monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) + monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) + + for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_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(RuntimeError, match="Expert parallelism enabled for a non-MoE model"): + parallel_state.initialize_model_parallel( + tensor_parallel_size=2, + sequence_parallel_size=2, + ulysses_degree=2, + ring_degree=1, + pipeline_parallel_size=1, + cfg_parallel_size=2, + data_parallel_size=1, + enable_expert_parallel=True, + backend="gloo", + ) + + assert parallel_state.vllm_parallel_state._PCP is None + assert parallel_state.vllm_parallel_state._EP is None + assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 163f9dd45..5904ae94f 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -35,7 +35,6 @@ from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm_omni.diffusion import envs -from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.platforms import current_omni_platform @@ -61,6 +60,11 @@ _DP: GroupCoordinator | None = None _FS: GroupCoordinator | None = None # Fully Sharded (HSDP shard dimension) _DIT: GroupCoordinator | None = None +# Rank-layout metadata for expert parallelism. This is not a process group; +# it is reused by platform-specific runtimes that must build companion groups +# with the same rank layout as vLLM EP. +_EXPERT_PARALLEL_GROUP_RANKS: list[list[int]] | None = None + def generate_masked_orthogonal_rank_groups( world_size: int, parallel_size: list[int], mask: list[bool] @@ -190,7 +194,6 @@ class RankGenerator: self.fs = fs self.rank_offset = rank_offset self.world_size = tp * sp * pp * cfg * dp - self.ep = tp * sp * cfg * dp # EP level exclude PP self.name_to_size = { "tp": self.tp, @@ -253,15 +256,6 @@ class RankGenerator: ranks.append(group) return ranks - if token == "ep": - ranks = [] - num_pp_stages = self.pp - for i in range(num_pp_stages): - start = i * self.ep + self.rank_offset - end = start + self.ep - ranks.append(list(range(start, end))) - return ranks - mask = self.get_mask(self.order, token) ranks = generate_masked_orthogonal_rank_groups(self.world_size, self.ordered_size, mask) if self.rank_offset > 0: @@ -309,6 +303,12 @@ def get_ring_parallel_rank(): return get_sp_group().ring_rank +def get_expert_parallel_group_ranks() -> list[list[int]]: + assert vllm_parallel_state._EP is not None, "expert parallel group is not initialized" + assert _EXPERT_PARALLEL_GROUP_RANKS is not None, "expert parallel group ranks are not initialized" + return _EXPERT_PARALLEL_GROUP_RANKS + + # PP def get_pp_group() -> PipelineGroupCoordinator: assert _PP is not None, "pipeline model parallel group is not initialized" @@ -782,6 +782,13 @@ def initialize_model_parallel( fs=fully_shard_degree, order="tp-sp-pp-cfg-dp", ) + use_moe_parallel_mapping = False + if enable_expert_parallel: + od_config = get_forward_context().omni_diffusion_config + use_moe_parallel_mapping = bool(od_config and od_config.is_moe) + if not use_moe_parallel_mapping: + raise RuntimeError("Expert parallelism enabled for a non-MoE model") + sp_group_ranks = rank_generator.get_ranks("sp") global _DP assert _DP is None, "data parallel group is already initialized" @@ -828,6 +835,10 @@ def initialize_model_parallel( ulysses_group=ulysses_pg, ring_group=ring_pg, ) + if use_moe_parallel_mapping: + # Diffusion normally uses its own SP group. Map it to vLLM PCP only for + # expert-parallel runtimes that rely on vLLM FusedMoE group semantics. + vllm_parallel_state._PCP = _SP assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized" vllm_parallel_state._TP = init_model_parallel_group( @@ -836,8 +847,17 @@ def initialize_model_parallel( backend=backend, parallel_mode="tensor", ) + if use_moe_parallel_mapping and cfg_parallel_size > 1: + # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP + # only when constructing the vLLM EP layout for expert-parallel paths. + vllm_parallel_state._DP = init_model_parallel_group( + group_ranks=rank_generator.get_ranks("cfg-dp"), + local_rank=get_world_group().local_rank, + backend=backend, + parallel_mode="data", + ) - global _FS + global _FS, _EXPERT_PARALLEL_GROUP_RANKS assert _FS is None, "fully shard group is already initialized" _FS = init_model_parallel_group( group_ranks=rank_generator.get_ranks("fs", independent_ranks=True), @@ -846,34 +866,40 @@ def initialize_model_parallel( parallel_mode="fully_shard", ) - if enable_expert_parallel: - od_config: OmniDiffusionConfig | None = get_forward_context().omni_diffusion_config - if od_config and od_config.is_moe: - vllm_parallel_state._EP = init_model_parallel_group( - group_ranks=rank_generator.get_ranks("ep"), - local_rank=get_world_group().local_rank, - backend=backend, - parallel_mode="expert", - ) - else: - raise RuntimeError("Expert parallelism enabled for a non-MoE model ") + _EXPERT_PARALLEL_GROUP_RANKS = None + if use_moe_parallel_mapping: + ep_group_ranks = rank_generator.get_ranks("tp-sp-cfg-dp") + vllm_parallel_state._EP = init_model_parallel_group( + group_ranks=ep_group_ranks, + local_rank=get_world_group().local_rank, + backend=backend, + parallel_mode="expert", + ) + _EXPERT_PARALLEL_GROUP_RANKS = ep_group_ranks init_dit_group(dit_parallel_size, backend) def destroy_model_parallel(): """Set the groups to none and destroy them.""" - global _DP + global _DP, _CFG, _SP, _PP, _FS, _EXPERT_PARALLEL_GROUP_RANKS + + if vllm_parallel_state._DP and vllm_parallel_state._DP is not _DP: + vllm_parallel_state._DP.destroy() + vllm_parallel_state._DP = None + if _DP: _DP.destroy() _DP = None - global _CFG if _CFG: _CFG.destroy() _CFG = None - global _SP + if vllm_parallel_state._PCP and vllm_parallel_state._PCP is not _SP: + vllm_parallel_state._PCP.destroy() + vllm_parallel_state._PCP = None + if _SP: _SP.destroy() _SP = None @@ -885,13 +911,12 @@ def destroy_model_parallel(): if vllm_parallel_state._EP: vllm_parallel_state._EP.destroy() vllm_parallel_state._EP = None + _EXPERT_PARALLEL_GROUP_RANKS = None - global _PP if _PP: _PP.destroy() _PP = None - global _FS if _FS: _FS.destroy() _FS = None diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index e4c9e0b2a..1212497e6 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -29,7 +29,7 @@ from transformers.modeling_outputs import ( CausalLMOutputWithPast, ) from transformers.modeling_utils import PreTrainedModel -from vllm.config import CacheConfig +from vllm.config import CacheConfig, get_current_vllm_config from vllm.distributed import ( get_tensor_model_parallel_world_size, ) @@ -1624,6 +1624,7 @@ class HunYuanSparseMoeBlock(nn.Module): else: self.shared_mlp = None + enable_expert_parallel = get_current_vllm_config().parallel_config.enable_expert_parallel self.experts = HunyuanFusedMoE( shared_experts=self.shared_mlp, num_experts=self.n_routed_experts, @@ -1635,7 +1636,7 @@ class HunYuanSparseMoeBlock(nn.Module): prefix=f"{prefix}.experts", enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, - pcp_size=1, + pcp_size=None if enable_expert_parallel else 1, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 0e1542bef..172f48cf1 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -246,9 +246,18 @@ class DiffusionWorker: # Create vllm_config for parallel configuration. Pass explicit device_config # so DeviceConfig does not rely on current_platform in worker subprocesses. vllm_config = _create_diffusion_worker_vllm_config(self.device, self.od_config) - vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size - vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size - vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel + parallel_config = self.od_config.parallel_config + vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size + vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size + if parallel_config.enable_expert_parallel and self.od_config.is_moe: + # Diffusion uses its own DP/CFG/SP groups normally. vLLM groups are + # only remapped for expert-parallel runtimes that consume vLLM's + # FusedMoE/EP semantics. + vllm_config.parallel_config.data_parallel_size = ( + parallel_config.data_parallel_size * parallel_config.cfg_parallel_size + ) + vllm_config.parallel_config.prefill_context_parallel_size = parallel_config.sequence_parallel_size + vllm_config.parallel_config.enable_expert_parallel = parallel_config.enable_expert_parallel vllm_config.profiler_config = self.od_config.profiler_config vllm_config.model_config = _make_diffusion_vllm_model_config(self.od_config) # type: ignore[assignment] vllm_config.quant_config = self.od_config.quantization_config diff --git a/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py b/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py index c53fad2e2..251eac2cf 100644 --- a/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py +++ b/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py @@ -7,7 +7,6 @@ import torch import vllm.forward_context as _vllm_fc from vllm.config import VllmConfig from vllm.distributed import get_ep_group -from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size from vllm.distributed.parallel_state import ( init_model_parallel_group as vllm_init_model_parallel_group, ) @@ -17,7 +16,7 @@ from vllm_ascend.ops.fused_moe.moe_comm_method import _MoECommMethods from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type from vllm_omni.diffusion.distributed.parallel_state import ( - get_data_parallel_world_size, + get_expert_parallel_group_ranks, get_world_group, ) from vllm_omni.diffusion.forward_context import get_forward_context as omni_get_ctx @@ -42,20 +41,14 @@ def _set_hunyuan_fused_moe_forward_context(num_tokens: int) -> None: def _init_mc2_group_for_diffusion( - world_size: int, - data_parallel_size: int, - tensor_parallel_size: int, backend: str, local_rank: int, + group_ranks: list[list[int]], ) -> None: import vllm_ascend.distributed.parallel_state as vllm_ascend_parallel_state if getattr(vllm_ascend_parallel_state, "_MC2", None) is not None: return - all_ranks = torch.arange(world_size).reshape(-1, data_parallel_size * tensor_parallel_size) - group_ranks = all_ranks.unbind(0) - group_ranks = [x.tolist() for x in group_ranks] - vllm_ascend_parallel_state._MC2 = vllm_init_model_parallel_group( group_ranks, local_rank, @@ -82,20 +75,17 @@ def _select_moe_comm_method(vllm_config: VllmConfig) -> MoECommType | None: def prepare_hunyuan_fused_moe_runtime() -> None: - world_size = torch.distributed.get_world_size() - data_parallel_size = get_data_parallel_world_size() - tensor_parallel_size = get_tensor_model_parallel_world_size() - backend = torch.distributed.get_backend(get_world_group().device_group) - local_rank = get_world_group().local_rank - _init_mc2_group_for_diffusion( - world_size=world_size, - data_parallel_size=data_parallel_size, - tensor_parallel_size=tensor_parallel_size, - backend=backend, - local_rank=local_rank, - ) + vllm_config = omni_get_ctx().vllm_config + if vllm_config.parallel_config.enable_expert_parallel: + backend = torch.distributed.get_backend(get_world_group().device_group) + local_rank = get_world_group().local_rank + _init_mc2_group_for_diffusion( + backend=backend, + local_rank=local_rank, + group_ranks=get_expert_parallel_group_ranks(), + ) - moe_comm_type = _select_moe_comm_method(vllm_config=omni_get_ctx().vllm_config) + moe_comm_type = _select_moe_comm_method(vllm_config=vllm_config) _ensure_forward_context_attr("num_tokens", int | None, None) _ensure_forward_context_attr("in_profile_run", bool, False) _ensure_forward_context_attr("moe_comm_type", MoECommType | None, moe_comm_type)