diff --git a/tests/config/test_explicit_expert_parallel.py b/tests/config/test_explicit_expert_parallel.py new file mode 100644 index 000000000000..323e4cedbd5f --- /dev/null +++ b/tests/config/test_explicit_expert_parallel.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm.config import ParallelConfig +from vllm.distributed.parallel_state import _get_ep_group_ranks + + +def test_explicit_ep_groups_preserve_tp_lanes() -> None: + all_ranks = torch.arange(8).reshape(1, 2, 1, 1, 4) + + groups = _get_ep_group_ranks( + all_ranks, + data_parallel_size=2, + prefill_context_model_parallel_size=1, + tensor_model_parallel_size=4, + expert_parallel_size=2, + ) + + assert groups == [[0, 4], [1, 5], [2, 6], [3, 7]] + + +def test_explicit_ep_disables_flattened_sequence_parallel() -> None: + config = ParallelConfig( + tensor_parallel_size=4, + data_parallel_size=2, + enable_expert_parallel=True, + expert_parallel_size=2, + ) + + assert not config.use_sequence_parallel_moe + + +@pytest.mark.parametrize( + "overrides", + [ + {"enable_expert_parallel": False}, + {"pipeline_parallel_size": 2}, + {"prefill_context_parallel_size": 2}, + {"expert_parallel_size": 4}, + {"all2all_backend": "deepep_low_latency"}, + {"enable_eplb": True}, + {"enable_elastic_ep": True}, + ], +) +def test_explicit_ep_rejects_unsupported_topologies(overrides) -> None: + kwargs = { + "tensor_parallel_size": 4, + "data_parallel_size": 2, + "enable_expert_parallel": True, + "expert_parallel_size": 2, + } + kwargs.update(overrides) + + with pytest.raises(ValueError): + ParallelConfig(**kwargs) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 83c4e6fe53d7..658adef13973 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -164,6 +164,13 @@ class ParallelConfig: """Whether the deployed model is MoE (if known).""" enable_expert_parallel: bool = False """Use expert parallelism instead of tensor parallelism for MoE layers.""" + expert_parallel_size: int | None = None + """Experimental explicit expert-parallel size. + + When unset, EP retains the standard flattened ``TP * PCP * DP`` size. + The initial explicit topology supports ``EP == DP`` with ``PCP == PP == 1`` + and preserves tensor parallelism within each expert. + """ enable_ep_weight_filter: bool = False """Skip non-local expert weights during model loading when expert parallelism is active. Each rank only reads its own expert shard from @@ -682,6 +689,7 @@ def use_sequence_parallel_moe(self) -> bool: "nixl_ep", ) and self.enable_expert_parallel + and self.expert_parallel_size is None and self.tensor_parallel_size > 1 and self.data_parallel_size > 1 ) @@ -904,6 +912,38 @@ def __post_init__(self) -> None: self.data_parallel_index = self.data_parallel_rank + if self.expert_parallel_size is not None: + if not self.enable_expert_parallel: + raise ValueError( + "expert_parallel_size requires enable_expert_parallel=True." + ) + if self.pipeline_parallel_size != 1: + raise ValueError( + "Explicit expert_parallel_size does not yet support " + "pipeline parallelism." + ) + if self.prefill_context_parallel_size != 1: + raise ValueError( + "Explicit expert_parallel_size does not yet support " + "prefill context parallelism." + ) + if self.expert_parallel_size != self.data_parallel_size: + raise ValueError( + "The initial hybrid TP x EP implementation requires " + "expert_parallel_size == data_parallel_size, but got " + f"{self.expert_parallel_size} and {self.data_parallel_size}." + ) + if self.all2all_backend != "allgather_reducescatter": + raise ValueError( + "Explicit expert_parallel_size initially supports only " + "allgather_reducescatter." + ) + if self.enable_eplb or self.enable_elastic_ep: + raise ValueError( + "Explicit expert_parallel_size does not yet support EPLB " + "or elastic EP." + ) + if self.distributed_executor_backend == "external_launcher": os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" logger.info("Disabling V1 multiprocessing for external launcher.") diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 128db903d4d6..7e7879d2153b 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1748,6 +1748,28 @@ def init_distributed_environment( _INNER_DP_WORLD = _WORLD +def _get_ep_group_ranks( + all_ranks: torch.Tensor, + *, + data_parallel_size: int, + prefill_context_model_parallel_size: int, + tensor_model_parallel_size: int, + expert_parallel_size: int | None, +) -> list[list[int]]: + if expert_parallel_size is not None: + # Preserve each TP group and form EP groups across DP ranks at the + # same TP lane. TP4/DP2 yields [0,4], [1,5], [2,6], [3,7]. + group_ranks = all_ranks.transpose(1, 4).reshape(-1, expert_parallel_size) + else: + group_ranks = all_ranks.transpose(1, 2).reshape( + -1, + data_parallel_size + * prefill_context_model_parallel_size + * tensor_model_parallel_size, + ) + return [ranks.tolist() for ranks in group_ranks.unbind(0)] + + def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, @@ -1924,17 +1946,13 @@ def initialize_model_parallel( assert _EP is None, "expert parallel group is already initialized" # Don't create EP group for dense models. if config.model_config is None or config.model_config.is_moe: - group_ranks = ( - all_ranks.transpose(1, 2) - .reshape( - -1, - data_parallel_size - * prefill_context_model_parallel_size - * tensor_model_parallel_size, - ) - .unbind(0) + group_ranks = _get_ep_group_ranks( + all_ranks, + data_parallel_size=data_parallel_size, + prefill_context_model_parallel_size=(prefill_context_model_parallel_size), + tensor_model_parallel_size=tensor_model_parallel_size, + expert_parallel_size=parallel_config.expert_parallel_size, ) - group_ranks = [x.tolist() for x in group_ranks] use_all2all = parallel_config.use_all2all if enable_elastic_ep: _EP = _init_stateless_group( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 3b9ead7f7e3b..871e92d66df7 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -493,6 +493,7 @@ class EngineArgs: data_parallel_multi_port_external_lb: bool = False data_parallel_backend: DataParallelBackend = ParallelConfig.data_parallel_backend enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel + expert_parallel_size: int | None = ParallelConfig.expert_parallel_size enable_ep_weight_filter: bool = ParallelConfig.enable_ep_weight_filter moe_backend: MoEBackend = KernelConfig.moe_backend linear_backend: LinearBackend = KernelConfig.linear_backend @@ -1134,6 +1135,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "-ep", **parallel_kwargs["enable_expert_parallel"], ) + parallel_group.add_argument( + "--expert-parallel-size", + **parallel_kwargs["expert_parallel_size"], + ) parallel_group.add_argument( "--enable-ep-weight-filter", **parallel_kwargs["enable_ep_weight_filter"], @@ -2241,6 +2246,7 @@ def create_engine_config( data_parallel_hybrid_lb=self.data_parallel_hybrid_lb, is_moe_model=model_config.is_moe, enable_expert_parallel=self.enable_expert_parallel, + expert_parallel_size=self.expert_parallel_size, enable_ep_weight_filter=self.enable_ep_weight_filter, all2all_backend=self.all2all_backend, enable_elastic_ep=self.enable_elastic_ep,