@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cc2f9b32c8

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 21:00 UTC

[Distributed] Add enable_expert_parallel arg (#14305)

Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com>

Tyler Michael Smith committed Mar 6, 2025 at 13:54 UTC cc2f9b32c858fb9f5acbcbfaad518036b93456c2

5 files changed +27 -21

examples/offline_inference/data_parallel.py

+3 -3

@@ -1,7 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # usage:

3 -# VLLM_TEST_ENABLE_EP=1 VLLM_USE_V1=1 \

4 -# python examples/offline_inference/data_parallel.py

3 +# VLLM_USE_V1=1 python examples/offline_inference/data_parallel.py

4 # we need to have a launcher to create multiple data parallel

5 # ranks. And each rank will create a vLLM instance to process its own prompts.

6 import os

@@ -55,7 +54,8 @@ def main(dp_size, dp_rank, dp_master_ip, dp_master_port, GPUs_per_dp_rank):

54 # Create an LLM.

55 llm = LLM(model="ibm-research/PowerMoE-3b",

56 tensor_parallel_size=GPUs_per_dp_rank,

58 - enforce_eager=True)

57 + enforce_eager=True,

58 + enable_expert_parallel=True)

59 outputs = llm.generate(prompts, sampling_params)

60 # Print the outputs.

61 for output in outputs:

vllm/config.py

+2 -1

@@ -754,7 +754,7 @@ class ModelConfig:

754 " must be divisible by tensor parallel size "

755 f"({tensor_parallel_size}).")

756

757 - if envs.VLLM_TEST_ENABLE_EP:

757 + if parallel_config.enable_expert_parallel:

758 self._verify_with_expert_parallelism()

759

760 pipeline_parallel_size = parallel_config.pipeline_parallel_size

@@ -1334,6 +1334,7 @@ class ParallelConfig:

1334 # IP of the data parallel master.

1335 data_parallel_master_ip: str = "127.0.0.1"

1336 data_parallel_master_port: int = 29500 # Port of the data parallel master.

1337 + enable_expert_parallel: bool = False # Use EP instead of TP for MoE layers.

1338

1339 # Maximum number of multiple batches

1340 # when load model sequentially. To avoid RAM OOM when using tensor

vllm/engine/arg_utils.py

@@ -114,6 +114,7 @@ class EngineArgs:

114 # number of P/D disaggregation (or other disaggregation) workers

115 pipeline_parallel_size: int = 1

116 tensor_parallel_size: int = 1

117 + enable_expert_parallel: bool = False

118 max_parallel_loading_workers: Optional[int] = None

119 block_size: Optional[int] = None

120 enable_prefix_caching: Optional[bool] = None

@@ -440,6 +441,11 @@ class EngineArgs:

441 type=int,

442 default=EngineArgs.tensor_parallel_size,

443 help='Number of tensor parallel replicas.')

444 + parser.add_argument(

445 + '--enable-expert-parallel',

446 + action='store_true',

447 + help='Use expert parallelism instead of tensor parallelism '

448 + 'for MoE layers.')

449 parser.add_argument(

450 '--max-parallel-loading-workers',

451 type=int,

@@ -1207,6 +1213,7 @@ class EngineArgs:

1213 parallel_config = ParallelConfig(

1214 pipeline_parallel_size=self.pipeline_parallel_size,

1215 tensor_parallel_size=self.tensor_parallel_size,

1216 + enable_expert_parallel=self.enable_expert_parallel,

1217 max_parallel_loading_workers=self.max_parallel_loading_workers,

1218 disable_custom_all_reduce=self.disable_custom_all_reduce,

1219 tokenizer_pool_config=TokenizerPoolConfig.create_config(

vllm/envs.py

-7

@@ -86,7 +86,6 @@ if TYPE_CHECKING:

86 VLLM_MLA_PERFORM_MATRIX_ABSORPTION: bool = True

87 VLLM_MLA_DISABLE_REQUANTIZATION: bool = False

88 VLLM_MLA_CUDA_MEM_ALIGN_KV_CACHE: bool = True

89 - VLLM_TEST_ENABLE_EP: bool = False

89 VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON: bool = False

90 VLLM_RAY_PER_WORKER_GPUS: float = 1.0

91 VLLM_RAY_BUNDLE_INDICES: str = ""

@@ -579,12 +578,6 @@ environment_variables: dict[str, Callable[[], Any]] = {

578 lambda: bool(int(os.getenv("VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON", "0"))

579 ),

580

582 - # If set, vLLM will use the experimental expert parallel implementation on

583 - # the FusedMoE layer, using tensor parallelism size as expert parallelism

584 - # size.

585 - "VLLM_TEST_ENABLE_EP":

586 - lambda: bool(int(os.getenv("VLLM_TEST_ENABLE_EP", "0"))),

587 -

581 # Number of GPUs per worker in Ray, if it is set to be a fraction,

582 # it allows ray to schedule multiple actors on a single GPU,

583 # so that users can colocate other actors on the same GPUs as vLLM.

vllm/model_executor/layers/fused_moe/layer.py

+15 -10

@@ -7,7 +7,6 @@ from typing import Callable, List, Optional, Tuple

7 import torch

8 from torch.nn.parameter import UninitializedParameter

10 -import vllm.envs as envs

10 from vllm.config import get_current_vllm_config

11 from vllm.distributed import (get_dp_group, get_tensor_model_parallel_rank,

12 get_tensor_model_parallel_world_size,

@@ -342,14 +341,6 @@ class FusedMoE(torch.nn.Module):

341 if params_dtype is None:

342 params_dtype = torch.get_default_dtype()

343

345 - # For smuggling this layer into the fused moe custom op

346 - compilation_config = get_current_vllm_config().compilation_config

347 - if prefix in compilation_config.static_forward_context:

348 - raise ValueError("Duplicate layer name: {}".format(prefix))

349 - compilation_config.static_forward_context[prefix] = self

350 - self.layer_name = prefix

351 - self.use_direct_call = not envs.VLLM_TEST_ENABLE_EP

352 -

344 # Note: here we guard against accessing the TP and DP groups when

345 # uninitialized (this happens when testing)

346 self.tp_size = (tp_size if tp_size is not None else

@@ -361,7 +352,21 @@ class FusedMoE(torch.nn.Module):

352 if self.dp_size == 1 else get_dp_group().rank_in_group)

353 self.global_num_experts = num_experts

354

364 - if envs.VLLM_TEST_ENABLE_EP:

355 + # Use expert parallelism instead of tensor parallelism?

356 + vllm_config = get_current_vllm_config()

357 + use_ep = (vllm_config.parallel_config.enable_expert_parallel

358 + and self.tp_size > 1)

359 +

360 + # For smuggling this layer into the fused moe custom op

361 + self.use_direct_call = self.dp_size == 1

362 + if self.use_direct_call:

363 + compilation_config = vllm_config.compilation_config

364 + if prefix in compilation_config.static_forward_context:

365 + raise ValueError("Duplicate layer name: {}".format(prefix))

366 + compilation_config.static_forward_context[prefix] = self

367 + self.layer_name = prefix

368 +

369 + if use_ep:

370 # Set TP size to 1 to adjust for EP and adjust EP size and rank

371 # for DP attention.

372 self.ep_rank = tp_rank + self.tp_size * self.dp_rank