[NPU] Fix HunyuanImage3 MoE EP+SP parallel mapping (#4344)
Signed-off-by: Celeste-jq <591998922@qq.com>
This commit is contained in:
Jiaqian Liu authored and GitHub committed
1 parent
b0f6b9de96
commit
a25636d780
5 files changed
+362
-55
No files matched your search
@@ -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
|
||||
Reference in new issue
Block a user