@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e50c454672
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 13, 18:15 UTC

[BugFix] Support EP/DP + EPLB with MTP (#25311)

Signed-off-by: ilmarkov <markovilya197@gmail.com> Signed-off-by: Sage Moore <sage@neuralmagic.com> Co-authored-by: Sage Moore <sage@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

Ilya Markov committed Nov 5, 2025 at 16:22 UTC e50c45467215f96068d95736b08d8a25f624e67d
27 files changed +956 -528
.buildkite/test-pipeline.yaml
+3 -2
@@ -232,8 +232,8 @@ steps:
232 commands:
233 - pytest -v -s distributed/test_eplb_algo.py
234
235 -- label: EPLB Execution Test # 5min
236 - timeout_in_minutes: 15
235 +- label: EPLB Execution Test # 10min
236 + timeout_in_minutes: 20
237 working_dir: "/vllm-workspace/tests"
238 num_gpus: 4
239 source_file_dependencies:
@@ -241,6 +241,7 @@ steps:
241 - tests/distributed/test_eplb_execute.py
242 commands:
243 - pytest -v -s distributed/test_eplb_execute.py
244 + - pytest -v -s distributed/test_eplb_spec_decode.py
245
246 - label: Metrics, Tracing Test # 12min
247 timeout_in_minutes: 20
tests/distributed/test_eplb_spec_decode.py new
+96
@@ -0,0 +1,96 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from __future__ import annotations
4 +
5 +import lm_eval
6 +import pytest
7 +
8 +from tests.utils import large_gpu_mark
9 +
10 +
11 +def get_model_args(
12 + model_name: str,
13 + spec_model_name: str,
14 + spec_method: str,
15 + tp_size: int,
16 + model_max_len: int,
17 +) -> dict:
18 + speculative_config = {
19 + "method": spec_method,
20 + "model": spec_model_name,
21 + "num_speculative_tokens": 1,
22 + "max_model_len": model_max_len,
23 + }
24 +
25 + model_args = {
26 + "pretrained": model_name,
27 + "dtype": "auto",
28 + "add_bos_token": True,
29 + "tensor_parallel_size": tp_size,
30 + "gpu_memory_utilization": 0.7,
31 + "speculative_config": speculative_config,
32 + "enable_expert_parallel": True,
33 + "num_redundant_experts": tp_size,
34 + "eplb_window_size": 128,
35 + "eplb_step_interval": 1024,
36 + "eplb_log_balancedness": False,
37 + "enable_eplb": True,
38 + "max_model_len": model_max_len,
39 + }
40 + return model_args
41 +
42 +
43 +@pytest.mark.parametrize(
44 + "model_setup",
45 + [
46 + pytest.param(
47 + ("mtp", "Qwen/Qwen3-Next-80B-A3B-Instruct", None, 4, 0.86),
48 + marks=large_gpu_mark(min_gb=80),
49 + ),
50 + pytest.param(
51 + (
52 + "eagle",
53 + "meta-llama/Llama-4-Scout-17B-16E-Instruct",
54 + "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
55 + 4,
56 + 0.92,
57 + ),
58 + marks=pytest.mark.skip(reason="Skipping due to CI OOM issues"),
59 + ),
60 + ],
61 + ids=["qwen3_next_mtp", "llama4_eagle"],
62 +)
63 +def test_eplb_spec_decode(
64 + monkeypatch: pytest.MonkeyPatch,
65 + model_setup: tuple[str, str, str, int, float],
66 +):
67 + """
68 + Test the correctness of EPLB speculative decoding with GSM8K dataset.
69 + Applicable to MoE models with mtp or eagle spec decode.
70 + """
71 + method, model_name, spec_model_name, tp_size, expected_gsm8k_value = model_setup
72 +
73 + TASK = "gsm8k"
74 + FILTER = "exact_match,strict-match"
75 + RTOL = 0.03
76 +
77 + model_args = get_model_args(
78 + model_name=model_name,
79 + spec_model_name=spec_model_name,
80 + spec_method=method,
81 + tp_size=tp_size,
82 + model_max_len=4096,
83 + )
84 +
85 + results = lm_eval.simple_evaluate(
86 + model="vllm",
87 + model_args=model_args,
88 + tasks=TASK,
89 + batch_size=64,
90 + num_fewshot=8,
91 + )
92 + measured_value = results["results"][TASK][FILTER]
93 + assert (
94 + measured_value - RTOL < expected_gsm8k_value
95 + and measured_value + RTOL > expected_gsm8k_value
96 + ), f"Expected: {expected_gsm8k_value} | Measured: {measured_value}"
vllm/distributed/eplb/eplb_state.py
+375 -184
@@ -33,7 +33,7 @@ from dataclasses import dataclass
33 import torch
34 from torch.distributed import ProcessGroup, all_reduce
35
36 -from vllm.config import ParallelConfig
36 +from vllm.config import ModelConfig, ParallelConfig
37 from vllm.distributed.parallel_state import (
38 get_ep_group,
39 get_node_count,
@@ -50,7 +50,7 @@ logger = init_logger(__name__)
50
51
52 @dataclass
53 -class EplbState:
53 +class EplbModelState:
54 """EPLB metrics."""
55
56 physical_to_logical_map: torch.Tensor
@@ -130,34 +130,46 @@ class EplbState:
130 See:
131 https://github.com/vllm-project/vllm/pull/22167#pullrequestreview-3086143856
132 """
133 - expert_load_window_step: int = 0
134 - """
135 - Current step in the sliding window.
133 + model_name: str
134 + model: MixtureOfExperts
135
137 - Different from `expert_rearrangement_step`, each EP rank may have its own
138 - `expert_load_window_step`.
139 - """
140 - expert_load_window_size: int = 0
136 +
137 +class EplbState:
138 """
142 - Size of the expert load sliding window.
143 - This is a constant and is taken from the config.
139 + EplbState of each expert parallel model. Key is the model config hash.
140 """
141
146 - expert_rearrangement_step: int = 0
147 - """
148 - Steps after last rearrangement.
149 - Will trigger a rearrangement if it exceeds the threshold.
142 + def __init__(self, parallel_config: ParallelConfig, device: torch.device):
143 + self.parallel_config = parallel_config
144 + self.device = device
145 + self.model_states: dict[str, EplbModelState] = {}
146 + """
147 + Current step in the sliding window.
148
151 - NOTE: Keep in mind that all EP ranks need to have the same
152 - `expert_rearrangement_step` value to ensure synchronization.
153 - Otherwise, the rearrangement will hang at collective
154 - communication calls.
155 - """
156 - expert_rearrangement_step_interval: int = 0
157 - """
158 - Interval for expert rearrangement steps.
159 - This is a constant and is taken from the config.
160 - """
149 + Different from `expert_rearrangement_step`,
150 + each EP rank may have its own `expert_load_window_step`.
151 + """
152 + self.expert_load_window_step: int = 0
153 + """
154 + Size of the expert load sliding window.
155 + This is a constant and is taken from the config.
156 + """
157 + self.expert_load_window_size: int = 0
158 + """
159 + Steps after last rearrangement.
160 + Will trigger a rearrangement if it exceeds the threshold.
161 +
162 + NOTE: Keep in mind that all EP ranks need to have the same
163 + `expert_rearrangement_step` value to ensure synchronization.
164 + Otherwise, the rearrangement will hang at collective
165 + communication calls.
166 + """
167 + self.expert_rearrangement_step: int = 0
168 + """
169 + Interval for expert rearrangement steps.
170 + This is a constant and is taken from the config.
171 + """
172 + self.expert_rearrangement_step_interval: int = 0
173
174 @staticmethod
175 def build_initial_global_physical_to_logical_map(
@@ -179,26 +191,63 @@ class EplbState:
191 ]
192 return global_physical_to_logical_map
193
182 - @classmethod
183 - def build(
184 - cls,
194 + def validate_ep_configuration(self, new_model: MixtureOfExperts):
195 + """
196 + Validate that the expert parallel configuration of
197 + the new model is the same as the existing models.
198 + """
199 + if len(self.model_states) > 0:
200 + model = next(iter(self.model_states.values())).model
201 + if (
202 + model.num_routed_experts != new_model.num_routed_experts
203 + or model.num_redundant_experts != new_model.num_redundant_experts
204 + or model.num_physical_experts != new_model.num_physical_experts
205 + or model.num_logical_experts != new_model.num_logical_experts
206 + or model.num_expert_groups != new_model.num_expert_groups
207 + ):
208 + raise RuntimeError(
209 + "Model: {} "
210 + "with config {} "
211 + "{} {} {} {} "
212 + "mismatch with new model {} "
213 + "with config {} "
214 + "{} {} {} {}".format(
215 + type(model),
216 + model.num_routed_experts,
217 + model.num_redundant_experts,
218 + model.num_physical_experts,
219 + model.num_logical_experts,
220 + model.num_expert_groups,
221 + type(new_model),
222 + new_model.num_routed_experts,
223 + new_model.num_redundant_experts,
224 + new_model.num_physical_experts,
225 + new_model.num_logical_experts,
226 + new_model.num_expert_groups,
227 + )
228 + )
229 +
230 + def add_model(
231 + self,
232 model: MixtureOfExperts,
186 - device: torch.device,
187 - parallel_config: ParallelConfig,
233 + model_config: ModelConfig,
234 global_expert_load: torch.Tensor | None = None,
235 old_global_expert_indices: torch.Tensor | None = None,
236 rank_mapping: dict[int, int] | None = None,
191 - ) -> "EplbState":
237 + ):
238 """
239 Build the initial EPLB state.
240 """
195 - physical_to_logical_map_list = cls.build_initial_global_physical_to_logical_map(
196 - model.num_routed_experts,
197 - model.num_redundant_experts,
241 + self.validate_ep_configuration(model)
242 + physical_to_logical_map_list = (
243 + EplbState.build_initial_global_physical_to_logical_map(
244 + model.num_routed_experts,
245 + model.num_redundant_experts,
246 + )
247 )
248 physical_to_logical_map = torch.tensor(
249 physical_to_logical_map_list,
201 - device=device,
250 + device=self.device,
251 )
252 # Assuming 8 GPUs per node, this supports up to
253 # (1023 + 1) / 8 = 128 nodes for now.
@@ -212,11 +261,11 @@ class EplbState:
261 logical_to_physical_map = torch.full(
262 (model.num_logical_experts, max_slots_per_logical_expert),
263 -1,
215 - device=device,
264 + device=self.device,
265 )
266 logical_replica_count = torch.zeros(
267 (model.num_logical_experts,),
219 - device=device,
268 + device=self.device,
269 dtype=torch.long,
270 )
271
@@ -255,18 +304,25 @@ class EplbState:
304 expert_load_pass = torch.zeros(
305 (model.num_moe_layers, model.num_physical_experts),
306 dtype=torch.int32,
258 - device=device,
307 + device=self.device,
308 )
260 - expert_load_window_size = parallel_config.eplb_config.window_size
309 + self.expert_load_window_size = self.parallel_config.eplb_config.window_size
310 expert_load_window = torch.zeros(
262 - (expert_load_window_size, model.num_moe_layers, model.num_physical_experts),
311 + (
312 + self.expert_load_window_size,
313 + model.num_moe_layers,
314 + model.num_physical_experts,
315 + ),
316 dtype=torch.int32,
264 - device=device,
317 + device=self.device,
318 )
319
320 # Set the initial progress of rearrangement to 3/4
268 - eplb_step_interval = parallel_config.eplb_config.step_interval
269 - expert_rearrangement_step = max(0, eplb_step_interval - eplb_step_interval // 4)
321 + eplb_step_interval = self.parallel_config.eplb_config.step_interval
322 + self.expert_rearrangement_step = max(
323 + 0, eplb_step_interval - eplb_step_interval // 4
324 + )
325 + self.expert_rearrangement_step_interval = eplb_step_interval
326
327 if global_expert_load is not None:
328 ep_group = get_ep_group().device_group
@@ -309,7 +365,7 @@ class EplbState:
365 (0, logical_to_physical_map.shape[-1] - max_physical_slots),
366 value=-1,
367 )
312 - physical_to_logical_map = new_physical_to_logical_map.to(device)
368 + physical_to_logical_map = new_physical_to_logical_map.to(self.device)
369 logical_to_physical_map.copy_(new_logical_to_physical_map)
370 logical_replica_count.copy_(new_logical_replica_count)
371
@@ -327,22 +383,20 @@ class EplbState:
383 False,
384 rank_mapping,
385 )
330 - expert_rearrangement_step = 0
386 + self.expert_rearrangement_step = 0
387
332 - return cls(
388 + self.model_states[model_config.compute_hash()] = EplbModelState(
389 physical_to_logical_map,
390 logical_to_physical_map,
391 logical_replica_count,
392 expert_load_pass,
393 expert_load_window,
338 - expert_load_window_size=expert_load_window_size,
339 - expert_rearrangement_step=expert_rearrangement_step,
340 - expert_rearrangement_step_interval=eplb_step_interval,
394 + model_config.model,
395 + model,
396 )
397
398 def step(
399 self,
345 - model: MixtureOfExperts,
400 is_dummy: bool = False,
401 is_profile: bool = False,
402 log_stats: bool = False,
@@ -351,7 +405,6 @@ class EplbState:
405 Step the EPLB state.
406
407 Args:
354 - model (MixtureOfExperts): The MoE model.
408 is_dummy (bool): If `True`, this is a dummy step and the load
409 metrics recorded in this forward pass will not count.
410 Defaults to `False`.
@@ -369,60 +422,66 @@ class EplbState:
422 """
423
424 if is_profile:
372 - self.rearrange(model, is_profile=True)
425 + self.rearrange(is_profile=True)
426 return
427
428 if is_dummy:
429 # Do not record load metrics for dummy steps
377 - self.expert_load_pass.zero_()
430 + for eplb_model_state in self.model_states.values():
431 + eplb_model_state.expert_load_pass.zero_()
432
433 if log_stats:
380 - # total_expert_load_pass: (num_moe_layers, num_physical_experts)
381 - total_expert_load_pass = self.expert_load_pass.clone()
382 -
383 - # Collect load metrics from all ranks
434 + # Sync the expert load pass for each model (main and drafter).
435 + # expert_load_pass: (num_moe_layers, num_physical_experts)
436 + expert_load_pass_list = self._sync_load_pass()
437 ep_group = get_ep_group().device_group
385 - all_reduce(total_expert_load_pass, group=ep_group)
386 -
387 - # num_tokens_per_rank: (num_moe_layers, num_ranks)
388 - num_tokens_per_rank = (
389 - total_expert_load_pass.reshape(
390 - total_expert_load_pass.shape[0], ep_group.size(), -1
438 + for expert_load_pass, eplb_model_state in zip(
439 + expert_load_pass_list, self.model_states.values()
440 + ):
441 + # num_tokens_per_rank: (num_moe_layers, num_ranks)
442 + num_tokens_per_rank = (
443 + expert_load_pass.reshape(
444 + expert_load_pass.shape[0], ep_group.size(), -1
445 + )
446 + .sum(dim=-1)
447 + .float()
448 )
392 - .sum(dim=-1)
393 - .float()
394 - )
449
396 - # Compute balancedness ratio:
397 - # for each layer:
398 - # (mean load across ranks) / (max load across ranks)
399 - avg_tokens_tensor = num_tokens_per_rank.mean(dim=0).sum(dim=0)
400 - max_tokens_tensor = num_tokens_per_rank.max(dim=0).values.sum(dim=0)
401 -
402 - # Just to make type checker happy
403 - tokens_tensors: list[float] = torch.stack(
404 - [avg_tokens_tensor, max_tokens_tensor]
405 - ).tolist()
406 - avg_tokens, max_tokens = tokens_tensors
407 - balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0
408 -
409 - if ep_group.rank() == 0:
410 - logger.info(
411 - "EPLB step: avg_tokens=%.2f, max_tokens=%d, balancedness=%.4f",
412 - avg_tokens,
413 - max_tokens,
414 - balancedness,
415 - )
450 + # Compute balancedness ratio:
451 + # for each layer:
452 + # (mean load across ranks) / (max load across ranks)
453 + avg_tokens_tensor = num_tokens_per_rank.mean(dim=0).sum(dim=0)
454 + max_tokens_tensor = num_tokens_per_rank.max(dim=0).values.sum(dim=0)
455 +
456 + # Just to make type checker happy
457 + tokens_tensors: list[float] = torch.stack(
458 + [avg_tokens_tensor, max_tokens_tensor]
459 + ).tolist()
460 + avg_tokens, max_tokens = tokens_tensors
461 + balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0
462 +
463 + if ep_group.rank() == 0:
464 + logger.info(
465 + "EPLB step: %d for model %s: avg_tokens=%.2f, "
466 + "max_tokens=%d, balancedness=%.4f",
467 + self.expert_rearrangement_step,
468 + eplb_model_state.model_name,
469 + avg_tokens,
470 + max_tokens,
471 + balancedness,
472 + )
473
474 # Update the expert load sliding window
475 if not is_dummy:
419 - self.expert_load_window[self.expert_load_window_step] = (
420 - self.expert_load_pass.clone()
421 - )
476 + for eplb_model_state in self.model_states.values():
477 + eplb_model_state.expert_load_window[self.expert_load_window_step] = (
478 + eplb_model_state.expert_load_pass.clone()
479 + )
480 + eplb_model_state.expert_load_pass.zero_()
481 +
482 self.expert_load_window_step += 1
483 if self.expert_load_window_step >= self.expert_load_window_size:
484 self.expert_load_window_step = 0
425 - self.expert_load_pass.zero_()
485
486 # Step the expert rearrangement step
487 # Note that even if this is a dummy step, we still increment the
@@ -431,18 +490,30 @@ class EplbState:
490 self.expert_rearrangement_step += 1
491 if self.expert_rearrangement_step >= self.expert_rearrangement_step_interval:
492 self.expert_rearrangement_step = 0
434 - self.rearrange(model)
493 + self.rearrange()
494
495 def rearrange(
496 self,
438 - model: MixtureOfExperts,
497 is_profile: bool = False,
498 execute_shuffle: bool = True,
441 - global_expert_load: torch.Tensor | None = None,
499 + global_expert_loads: list[torch.Tensor] | None = None,
500 rank_mapping: dict[int, int] | None = None,
501 ) -> torch.Tensor | None:
502 """
503 Rearrange the experts according to the current load.
504 +
505 + Args:
506 + is_profile (bool): If `True`, perform a dummy rearrangement.
507 + This is used in `profile_run` to reserve enough memory,
508 + no memory movement will be performed. Default is False.
509 + execute_shuffle (bool): If `True`, execute the shuffle
510 + in elastic expert parallel (EEP). Default is True.
511 + global_expert_loads (list[torch.Tensor] | None): The global expert
512 + loads when scaling is done in EEP.
513 + List of expert loads for the main and drafter
514 + (when spec decode is used) models.
515 + rank_mapping (dict[int, int] | None): The rank mapping
516 + when scaling is done in EEP.
517 """
518
519 ep_group = get_ep_group().device_group
@@ -455,53 +526,71 @@ class EplbState:
526 time_start = time.perf_counter()
527 logger.info("Rearranging experts %s...", "(profile)" if is_profile else "")
528
458 - if global_expert_load is None:
529 + if global_expert_loads is None:
530 # Map the physical expert load to global logical experts
460 - logical_expert_load_window = torch.zeros(
461 - self.expert_load_window_size,
462 - model.num_moe_layers,
463 - model.num_logical_experts,
464 - dtype=self.expert_load_window.dtype,
465 - device=self.expert_load_window.device,
466 - )
467 - logical_expert_load_window.scatter_add_(
468 - dim=-1,
469 - index=self.physical_to_logical_map.unsqueeze(0)
470 - .expand_as(self.expert_load_window)
471 - .long(),
472 - src=self.expert_load_window,
473 - )
474 -
531 + global_expert_load_windows = []
532 if not execute_shuffle:
476 - metadata = torch.tensor(
477 - [
478 - model.num_moe_layers,
479 - model.num_logical_experts,
480 - self.physical_to_logical_map.shape[1],
481 - ],
482 - dtype=torch.int32,
483 - device="cpu",
533 + num_models = torch.tensor(
534 + [len(self.model_states)], dtype=torch.int32, device="cpu"
535 )
536 torch.distributed.broadcast(
486 - metadata, group=get_ep_group().cpu_group, group_src=0
537 + num_models, group=get_ep_group().cpu_group, group_src=0
538 )
539
489 - # Perform all-reduce to get the expert load across all ranks
490 - global_expert_load_window = logical_expert_load_window.sum(dim=0)
491 - all_reduce(global_expert_load_window, group=ep_group)
540 + for eplb_model_state in self.model_states.values():
541 + logical_expert_load_window = torch.zeros(
542 + self.expert_load_window_size,
543 + eplb_model_state.model.num_moe_layers,
544 + eplb_model_state.model.num_logical_experts,
545 + dtype=eplb_model_state.expert_load_window.dtype,
546 + device=eplb_model_state.expert_load_window.device,
547 + )
548 + logical_expert_load_window.scatter_add_(
549 + dim=-1,
550 + index=eplb_model_state.physical_to_logical_map.unsqueeze(0)
551 + .expand_as(eplb_model_state.expert_load_window)
552 + .long(),
553 + src=eplb_model_state.expert_load_window,
554 + )
555
556 + if not execute_shuffle:
557 + metadata = torch.tensor(
558 + [
559 + eplb_model_state.model.num_moe_layers,
560 + eplb_model_state.model.num_logical_experts,
561 + eplb_model_state.physical_to_logical_map.shape[1],
562 + ],
563 + dtype=torch.int32,
564 + device="cpu",
565 + )
566 + torch.distributed.broadcast(
567 + metadata, group=get_ep_group().cpu_group, group_src=0
568 + )
569 +
570 + global_expert_load_window = logical_expert_load_window.sum(dim=0)
571 + global_expert_load_windows.append(global_expert_load_window)
572 + # Perform all-reduce to get the expert load across all ranks for each model
573 + global_expert_load_windows = self._allreduce_list(
574 + global_expert_load_windows
575 + )
576 if not execute_shuffle:
494 - # (num_moe_layers, old_num_physical_experts)
495 - old_global_expert_indices = self.physical_to_logical_map
496 - torch.distributed.broadcast(
497 - old_global_expert_indices, group=ep_group, group_src=0
498 - )
499 - return global_expert_load_window
577 + for eplb_model_state, global_expert_load_window in zip(
578 + self.model_states.values(), global_expert_load_windows
579 + ):
580 + # (num_moe_layers, old_num_physical_experts)
581 + old_global_expert_indices = eplb_model_state.physical_to_logical_map
582 + torch.distributed.broadcast(
583 + old_global_expert_indices, group=ep_group, group_src=0
584 + )
585 + if not execute_shuffle:
586 + return global_expert_load_windows
587 else:
588 assert execute_shuffle
502 - global_expert_load_window = global_expert_load
589 + global_expert_load_windows = global_expert_loads
590
591 # TODO(bowen): Treat differently for prefill and decode nodes
592 + eplb_model_state = next(iter(self.model_states.values()))
593 + model = eplb_model_state.model
594 num_replicas = model.num_physical_experts
595 num_groups = model.num_expert_groups
596 if rank_mapping is not None and len(rank_mapping) == ep_group.size():
@@ -526,48 +615,64 @@ class EplbState:
615 f"{num_gpus=}, {num_nodes=}"
616 )
617
529 - # Get new expert mappings
530 - (
531 - new_physical_to_logical_map,
532 - new_logical_to_physical_map,
533 - new_logical_replica_count,
534 - ) = rebalance_experts(
535 - global_expert_load_window,
536 - num_replicas,
537 - num_groups,
538 - num_nodes,
539 - num_gpus,
540 - )
618 + for eplb_model_state, global_expert_load_window in zip(
619 + self.model_states.values(), global_expert_load_windows
620 + ):
621 + # Get new expert mappings for the model
622 + (
623 + new_physical_to_logical_map,
624 + new_logical_to_physical_map,
625 + new_logical_replica_count,
626 + ) = rebalance_experts(
627 + global_expert_load_window,
628 + num_replicas,
629 + num_groups,
630 + num_nodes,
631 + num_gpus,
632 + )
633
542 - # Update expert weights
543 - rearrange_expert_weights_inplace(
544 - self.physical_to_logical_map,
545 - new_physical_to_logical_map,
546 - model.expert_weights,
547 - ep_group,
548 - is_profile,
549 - rank_mapping,
550 - )
634 + # Update expert weights
635 + rearrange_expert_weights_inplace(
636 + eplb_model_state.physical_to_logical_map,
637 + new_physical_to_logical_map,
638 + eplb_model_state.model.expert_weights,
639 + ep_group,
640 + is_profile,
641 + rank_mapping,
642 + )
643
552 - if not is_profile:
553 - if (
554 - self.physical_to_logical_map.shape[1]
555 - != new_physical_to_logical_map.shape[1]
556 - ):
557 - self.physical_to_logical_map = new_physical_to_logical_map.to(
558 - self.physical_to_logical_map.device
644 + if not is_profile:
645 + if (
646 + eplb_model_state.physical_to_logical_map.shape[1]
647 + != new_physical_to_logical_map.shape[1]
648 + ):
649 + eplb_model_state.physical_to_logical_map = (
650 + new_physical_to_logical_map.to(
651 + eplb_model_state.physical_to_logical_map.device
652 + )
653 + )
654 + else:
655 + eplb_model_state.physical_to_logical_map.copy_(
656 + new_physical_to_logical_map
657 + )
658 + max_physical_slots = new_logical_to_physical_map.shape[-1]
659 + assert (
660 + max_physical_slots
661 + <= eplb_model_state.logical_to_physical_map.shape[-1]
662 )
560 - else:
561 - self.physical_to_logical_map.copy_(new_physical_to_logical_map)
562 - max_physical_slots = new_logical_to_physical_map.shape[-1]
563 - assert max_physical_slots <= self.logical_to_physical_map.shape[-1]
564 - new_logical_to_physical_map = torch.nn.functional.pad(
565 - new_logical_to_physical_map,
566 - (0, self.logical_to_physical_map.shape[-1] - max_physical_slots),
567 - value=-1,
568 - )
569 - self.logical_to_physical_map.copy_(new_logical_to_physical_map)
570 - self.logical_replica_count.copy_(new_logical_replica_count)
663 + new_logical_to_physical_map = torch.nn.functional.pad(
664 + new_logical_to_physical_map,
665 + (
666 + 0,
667 + eplb_model_state.logical_to_physical_map.shape[-1]
668 + - max_physical_slots,
669 + ),
670 + value=-1,
671 + )
672 + eplb_model_state.logical_to_physical_map.copy_(
673 + new_logical_to_physical_map
674 + )
675 + eplb_model_state.logical_replica_count.copy_(new_logical_replica_count)
676
677 if is_main_rank:
678 assert time_start is not None
@@ -581,32 +686,118 @@ class EplbState:
686 return None
687
688 @staticmethod
584 - def recv_state() -> tuple[torch.Tensor, torch.Tensor]:
689 + def recv_state() -> tuple[list[torch.Tensor], list[torch.Tensor]]:
690 """
691 Receive the expert load and old placement from the master rank.
692 """
693 ep_group = get_ep_group()
589 - metadata = torch.empty(3, dtype=torch.int32, device="cpu")
590 - torch.distributed.broadcast(metadata, group=ep_group.cpu_group, group_src=0)
591 - num_moe_layers, num_logical_experts, num_old_physical_experts = (
592 - metadata.tolist()
694 + num_models = torch.empty(1, dtype=torch.int32, device="cpu")
695 + torch.distributed.broadcast(num_models, group=ep_group.cpu_group, group_src=0)
696 + num_models = num_models.item()
697 + global_expert_loads = []
698 + old_global_expert_indices_per_model = []
699 + for _ in range(num_models):
700 + metadata = torch.empty(3, dtype=torch.int32, device="cpu")
701 + torch.distributed.broadcast(metadata, group=ep_group.cpu_group, group_src=0)
702 + num_moe_layers, num_logical_experts, num_old_physical_experts = (
703 + metadata.tolist()
704 + )
705 + global_expert_load = torch.zeros(
706 + (num_moe_layers, num_logical_experts),
707 + dtype=torch.int64,
708 + device=ep_group.device,
709 + )
710 + all_reduce(global_expert_load, group=ep_group.device_group)
711 + old_global_expert_indices = torch.empty(
712 + (num_moe_layers, num_old_physical_experts),
713 + dtype=torch.int64,
714 + device=ep_group.device,
715 + )
716 + torch.distributed.broadcast(
717 + old_global_expert_indices,
718 + group=ep_group.device_group,
719 + group_src=0,
720 + )
721 + global_expert_loads.append(global_expert_load)
722 + old_global_expert_indices_per_model.append(old_global_expert_indices)
723 + return global_expert_loads, old_global_expert_indices_per_model
724 +
725 + @classmethod
726 + def get_eep_state(
727 + cls, parallel_config: ParallelConfig
728 + ) -> tuple[
729 + list[torch.Tensor] | None,
730 + list[torch.Tensor] | None,
731 + dict[int, int] | None,
732 + ]:
733 + num_local_physical_experts = torch.empty(1, dtype=torch.int32, device="cpu")
734 + torch.distributed.broadcast(
735 + num_local_physical_experts,
736 + group=get_ep_group().cpu_group,
737 + group_src=0,
738 )
594 - global_expert_load = torch.zeros(
595 - (num_moe_layers, num_logical_experts),
596 - dtype=torch.int64,
597 - device=ep_group.device,
739 + num_local_physical_experts = int(num_local_physical_experts.item())
740 + new_ep_size = get_ep_group().world_size
741 + global_expert_loads, old_global_expert_indices_per_model = (
742 + EplbState.recv_state()
743 )
599 - all_reduce(global_expert_load, group=ep_group.device_group)
600 - old_global_expert_indices = torch.empty(
601 - (num_moe_layers, num_old_physical_experts),
602 - dtype=torch.int64,
603 - device=ep_group.device,
744 +
745 + # EP configuration for all models has to be the same so as eplb config
746 + num_logical_experts = global_expert_loads[0].shape[1]
747 + parallel_config.eplb_config.num_redundant_experts = (
748 + num_local_physical_experts * new_ep_size - num_logical_experts
749 )
605 - torch.distributed.broadcast(
606 - old_global_expert_indices, group=ep_group.device_group, group_src=0
750 + assert (
751 + old_global_expert_indices_per_model[0].shape[1] % num_local_physical_experts
752 + == 0
753 + )
754 + old_ep_size = (
755 + old_global_expert_indices_per_model[0].shape[1]
756 + // num_local_physical_experts
757 + )
758 + rank_mapping = {old_ep_rank: old_ep_rank for old_ep_rank in range(old_ep_size)}
759 + return (
760 + global_expert_loads,
761 + old_global_expert_indices_per_model,
762 + rank_mapping,
763 )
764
609 - return global_expert_load, old_global_expert_indices
765 + def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]:
766 + """
767 + All-reduce a list of tensors.
768 + """
769 + if len(tensor_list) == 1:
770 + all_reduce(tensor_list[0], group=get_ep_group().device_group)
771 + return tensor_list
772 + assert all(t.dim() == 2 for t in tensor_list), "All tensors must be 2D."
773 + assert all(t.shape[1] == tensor_list[0].shape[1] for t in tensor_list), (
774 + "All tensors must have the same shape[1]."
775 + )
776 + # Concatenate, all_reduce, then unpack to original shapes.
777 + # We assume all tensors are 2D and shape[1] (num_physical_experts)
778 + # is the same across all models.
779 + shapes = [t.shape for t in tensor_list]
780 + concat_tensor = torch.cat(tensor_list, dim=0)
781 +
782 + ep_group = get_ep_group().device_group
783 + all_reduce(concat_tensor, group=ep_group)
784 +
785 + all_reduce_list = []
786 + offset = 0
787 + for shape in shapes:
788 + all_reduce_list.append(concat_tensor[offset : offset + shape[0], :])
789 + offset += shape[0]
790 + return all_reduce_list
791 +
792 + def _sync_load_pass(self) -> list[torch.Tensor]:
793 + """
794 + Sync the expert load pass across all ranks for log stats.
795 + Doesn't update the expert load pass in eplb_model_state.
796 + """
797 + load_pass_list = []
798 + for eplb_model_state in self.model_states.values():
799 + load_pass_list.append(eplb_model_state.expert_load_pass.clone())
800 + return self._allreduce_list(load_pass_list)
801
802
803 def _node_count_with_rank_mapping(
vllm/entrypoints/openai/tool_parsers/abstract_tool_parser.py
+1 -1
@@ -226,7 +226,7 @@ class ToolParserManager:
226
227 if isinstance(name, str):
228 names = [name]
229 - elif is_list_of(name, str):
229 + elif name is not None and is_list_of(name, str):
230 names = name
231 else:
232 names = [class_name]
vllm/model_executor/models/deepseek_eagle.py
+7
@@ -24,9 +24,12 @@ from vllm.model_executor.models.deepseek_v2 import (
24 DeepseekV2DecoderLayer,
25 DeepseekV3ForCausalLM,
26 )
27 +from vllm.utils import init_logger
28
29 from .utils import AutoWeightsLoader, maybe_prefix
30
31 +logger = init_logger(__name__)
32 +
33
34 @support_torch_compile
35 class DeepseekV2Model(nn.Module):
@@ -215,6 +218,10 @@ class EagleDeepseekV3ForCausalLM(DeepseekV3ForCausalLM):
218 self.config.vocab_size, scale=logit_scale
219 )
220
221 + # Set MoE hyperparameters
222 + self.num_moe_layers = self.config.num_hidden_layers
223 + self.set_moe_parameters()
224 +
225 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
226 return self.model.get_input_embeddings(input_ids)
227
vllm/model_executor/models/deepseek_mtp.py
+27 -1
@@ -8,6 +8,7 @@ from transformers import PretrainedConfig
8
9 from vllm.compilation.decorators import support_torch_compile
10 from vllm.config import VllmConfig
11 +from vllm.logger import init_logger
12 from vllm.model_executor.layers.fused_moe import FusedMoE
13 from vllm.model_executor.layers.layernorm import RMSNorm
14 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -25,11 +26,15 @@ from vllm.sequence import IntermediateTensors
26
27 from .deepseek_v2 import (
28 DeepseekV2DecoderLayer,
29 + DeepseekV2MixtureOfExperts,
30 + DeepseekV2MoE,
31 get_spec_layer_idx_from_weight_name,
32 )
33 from .interfaces import SupportsPP
34 from .utils import maybe_prefix
35
36 +logger = init_logger(__name__)
37 +
38
39 class SharedHead(nn.Module):
40 def __init__(
@@ -119,6 +124,7 @@ class DeepSeekMultiTokenPredictor(nn.Module):
124 self.mtp_start_layer_idx = config.num_hidden_layers
125 self.num_mtp_layers = config.num_nextn_predict_layers
126 # to map the exact layer index from weights
127 +
128 self.layers = torch.nn.ModuleDict(
129 {
130 str(idx): DeepSeekMultiTokenPredictorLayer(
@@ -172,13 +178,33 @@ class DeepSeekMultiTokenPredictor(nn.Module):
178
179
180 @support_torch_compile
175 -class DeepSeekMTP(nn.Module, SupportsPP):
181 +class DeepSeekMTP(nn.Module, SupportsPP, DeepseekV2MixtureOfExperts):
182 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
183 super().__init__()
184 self.config = vllm_config.model_config.hf_config
185 self.model = DeepSeekMultiTokenPredictor(
186 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
187 )
188 + # Set MoE hyperparameters
189 + self.set_moe_parameters()
190 +
191 + def set_moe_parameters(self):
192 + self.expert_weights = []
193 + self.num_moe_layers = self.config.num_nextn_predict_layers
194 + self.num_expert_groups = self.config.n_group
195 +
196 + self.moe_layers = []
197 + self.moe_mlp_layers = []
198 + example_moe = None
199 + for layer in self.model.layers.values():
200 + assert isinstance(layer, DeepSeekMultiTokenPredictorLayer)
201 + layer = layer.mtp_block
202 + assert isinstance(layer, DeepseekV2DecoderLayer)
203 + if isinstance(layer.mlp, DeepseekV2MoE):
204 + example_moe = layer.mlp
205 + self.moe_mlp_layers.append(layer.mlp)
206 + self.moe_layers.append(layer.mlp.experts)
207 + self.extract_moe_parameters(example_moe)
208
209 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
210 return self.model.get_input_embeddings(input_ids)
vllm/model_executor/models/deepseek_v2.py
+57 -49
@@ -166,7 +166,7 @@ class DeepseekV2MoE(nn.Module):
166 self.routed_scaling_factor = config.routed_scaling_factor
167
168 self.ep_group = get_ep_group().device_group
169 - self.ep_rank = self.ep_group.rank()
169 + self.ep_rank = get_ep_group().rank_in_group
170 self.ep_size = self.ep_group.size()
171 self.n_routed_experts: int = config.n_routed_experts
172 self.n_shared_experts: int = config.n_shared_experts
@@ -1122,7 +1122,6 @@ class DeepseekV2Model(nn.Module):
1122 )
1123 else:
1124 self.embed_tokens = PPMissingLayer()
1125 -
1125 self.start_layer, self.end_layer, self.layers = make_layers(
1126 config.num_hidden_layers,
1127 lambda prefix: DeepseekV2DecoderLayer(
@@ -1172,7 +1171,50 @@ class DeepseekV2Model(nn.Module):
1171 return hidden_states
1172
1173
1175 -class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoRA):
1174 +class DeepseekV2MixtureOfExperts(MixtureOfExperts):
1175 + moe_mlp_layers: list[DeepseekV2MoE]
1176 + """
1177 + List of MoE MLP layers in the model.
1178 + """
1179 +
1180 + def extract_moe_parameters(self, example_moe: DeepseekV2MoE | None):
1181 + if example_moe is None:
1182 + self.num_moe_layers = 0
1183 + self.num_expert_groups = 0
1184 + self.num_logical_experts = 0
1185 + self.num_physical_experts = 0
1186 + self.num_local_physical_experts = 0
1187 + self.num_routed_experts = 0
1188 + self.num_shared_experts = 0
1189 + self.num_redundant_experts = 0
1190 + logger.warning("DeepSeekV2: No DeepseekV2MoE layer found in model.layers.")
1191 + else:
1192 + self.num_logical_experts = example_moe.n_logical_experts
1193 + self.num_physical_experts = example_moe.n_physical_experts
1194 + self.num_local_physical_experts = example_moe.n_local_physical_experts
1195 + self.num_routed_experts = example_moe.n_routed_experts
1196 + self.num_shared_experts = example_moe.n_shared_experts
1197 + self.num_redundant_experts = example_moe.n_redundant_experts
1198 +
1199 + def update_physical_experts_metadata(
1200 + self,
1201 + num_physical_experts: int,
1202 + num_local_physical_experts: int,
1203 + ) -> None:
1204 + assert self.num_local_physical_experts == num_local_physical_experts
1205 + self.num_physical_experts = num_physical_experts
1206 + self.num_local_physical_experts = num_local_physical_experts
1207 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
1208 + for moe in self.moe_mlp_layers:
1209 + moe.n_local_physical_experts = num_local_physical_experts
1210 + moe.n_physical_experts = num_physical_experts
1211 + moe.n_redundant_experts = self.num_redundant_experts
1212 + moe.experts.update_expert_map()
1213 +
1214 +
1215 +class DeepseekV2ForCausalLM(
1216 + nn.Module, SupportsPP, DeepseekV2MixtureOfExperts, SupportsLoRA
1217 +):
1218 packed_modules_mapping = {
1219 "gate_up_proj": ["gate_proj", "up_proj"],
1220 }
@@ -1213,13 +1255,19 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoR
1255 self.make_empty_intermediate_tensors = (
1256 self.model.make_empty_intermediate_tensors
1257 )
1258 + # Set MoE hyperparameters
1259 + self.num_moe_layers = (
1260 + self.config.num_hidden_layers - self.config.first_k_dense_replace
1261 + )
1262 + self.set_moe_parameters()
1263 +
1264 + def set_moe_parameters(self):
1265 self.expert_weights = []
1266
1218 - # Set MoE hyperparameters
1219 - self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace
1220 - self.num_expert_groups = config.n_group
1267 + self.num_expert_groups = self.config.n_group
1268
1222 - self.moe_layers: list[SharedFusedMoE] = []
1269 + self.moe_layers = []
1270 + self.moe_mlp_layers = []
1271 example_moe = None
1272 for layer in self.model.layers:
1273 if isinstance(layer, PPMissingLayer):
@@ -1229,50 +1277,10 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts, SupportsLoR
1277 if isinstance(layer.mlp, DeepseekV2MoE):
1278 # Pick last one layer since the first ones may be dense layers.
1279 example_moe = layer.mlp
1280 + self.moe_mlp_layers.append(layer.mlp)
1281 self.moe_layers.append(layer.mlp.experts)
1282
1234 - if example_moe is None:
1235 - raise RuntimeError("No DeepseekV2MoE layer found in model.layers.")
1236 -
1237 - self.num_logical_experts = example_moe.n_logical_experts
1238 - self.num_physical_experts = example_moe.n_physical_experts
1239 - self.num_local_physical_experts = example_moe.n_local_physical_experts
1240 - self.num_routed_experts = example_moe.n_routed_experts
1241 - self.num_shared_experts = example_moe.n_shared_experts
1242 - self.num_redundant_experts = example_moe.n_redundant_experts
1243 -
1244 - def set_eplb_state(
1245 - self,
1246 - expert_load_view: torch.Tensor,
1247 - logical_to_physical_map: torch.Tensor,
1248 - logical_replica_count: torch.Tensor,
1249 - ) -> None:
1250 - for layer_idx, layer in enumerate(self.moe_layers):
1251 - # Register the expert weights.
1252 - self.expert_weights.append(layer.get_expert_weights())
1253 - layer.set_eplb_state(
1254 - moe_layer_idx=layer_idx,
1255 - expert_load_view=expert_load_view,
1256 - logical_to_physical_map=logical_to_physical_map,
1257 - logical_replica_count=logical_replica_count,
1258 - )
1259 -
1260 - def update_physical_experts_metadata(
1261 - self,
1262 - num_physical_experts: int,
1263 - num_local_physical_experts: int,
1264 - ) -> None:
1265 - assert self.num_local_physical_experts == num_local_physical_experts
1266 - self.num_physical_experts = num_physical_experts
1267 - self.num_local_physical_experts = num_local_physical_experts
1268 - self.num_redundant_experts = num_physical_experts - self.num_logical_experts
1269 - for layer in self.model.layers:
1270 - if isinstance(layer.mlp, DeepseekV2MoE):
1271 - moe = layer.mlp
1272 - moe.n_local_physical_experts = num_local_physical_experts
1273 - moe.n_physical_experts = num_physical_experts
1274 - moe.n_redundant_experts = self.num_redundant_experts
1275 - moe.experts.update_expert_map()
1283 + self.extract_moe_parameters(example_moe)
1284
1285 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
1286 return self.model.get_input_embeddings(input_ids)
vllm/model_executor/models/ernie45_moe.py
+1 -17
@@ -133,7 +133,7 @@ class Ernie4_5_MoeMoE(nn.Module):
133
134 self.moe_num_shared_experts = getattr(config, "moe_num_shared_experts", None)
135 self.ep_group = get_ep_group().device_group
136 - self.ep_rank = self.ep_group.rank()
136 + self.ep_rank = get_ep_group().rank_in_group
137 self.ep_size = self.ep_group.size()
138 self.n_routed_experts: int = config.moe_num_experts
139 self.n_shared_experts: int = self.moe_num_shared_experts
@@ -709,22 +709,6 @@ class Ernie4_5_MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, MixtureOfExpe
709 self.num_shared_experts = example_moe.n_shared_experts
710 self.num_redundant_experts = example_moe.n_redundant_experts
711
712 - def set_eplb_state(
713 - self,
714 - expert_load_view: torch.Tensor,
715 - logical_to_physical_map: torch.Tensor,
716 - logical_replica_count: torch.Tensor,
717 - ) -> None:
718 - for layer_idx, layer in enumerate(self.moe_layers):
719 - # Register the expert weights.
720 - self.expert_weights.append(layer.get_expert_weights())
721 - layer.set_eplb_state(
722 - moe_layer_idx=layer_idx,
723 - expert_load_view=expert_load_view,
724 - logical_to_physical_map=logical_to_physical_map,
725 - logical_replica_count=logical_replica_count,
726 - )
727 -
712 def update_physical_experts_metadata(
713 self,
714 num_physical_experts: int,
vllm/model_executor/models/glm4_moe.py
+36 -29
@@ -62,7 +62,7 @@ from vllm.model_executor.model_loader.weight_utils import (
62 )
63 from vllm.sequence import IntermediateTensors
64
65 -from .interfaces import SupportsLoRA, SupportsPP
65 +from .interfaces import MixtureOfExperts, SupportsLoRA, SupportsPP
66 from .utils import (
67 AutoWeightsLoader,
68 PPMissingLayer,
@@ -127,7 +127,7 @@ class Glm4MoE(nn.Module):
127 self.routed_scaling_factor = config.routed_scaling_factor
128
129 self.ep_group = get_ep_group().device_group
130 - self.ep_rank = self.ep_group.rank()
130 + self.ep_rank = get_ep_group().rank_in_group
131 self.ep_size = self.ep_group.size()
132 self.n_routed_experts: int = config.n_routed_experts
133 self.n_shared_experts: int = config.n_shared_experts
@@ -616,7 +616,35 @@ class Glm4MoeModel(nn.Module):
616 return loaded_params
617
618
619 -class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
619 +class Glm4MixtureOfExperts(MixtureOfExperts):
620 + def extract_moe_parameters(self, example_moe: Glm4MoE | None) -> None:
621 + if example_moe is None:
622 + raise RuntimeError("No Glm4MoE layer found in model.layers.")
623 + else:
624 + self.num_logical_experts = example_moe.n_logical_experts
625 + self.num_physical_experts = example_moe.n_physical_experts
626 + self.num_local_physical_experts = example_moe.n_local_physical_experts
627 + self.num_routed_experts = example_moe.n_routed_experts
628 + self.num_shared_experts = example_moe.n_shared_experts
629 + self.num_redundant_experts = example_moe.n_redundant_experts
630 +
631 + def update_physical_experts_metadata(
632 + self,
633 + num_physical_experts: int,
634 + num_local_physical_experts: int,
635 + ) -> None:
636 + assert self.num_local_physical_experts == num_local_physical_experts
637 + self.num_physical_experts = num_physical_experts
638 + self.num_local_physical_experts = num_local_physical_experts
639 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
640 + for moe in self.moe_mlp_layers:
641 + moe.n_local_physical_experts = num_local_physical_experts
642 + moe.n_physical_experts = num_physical_experts
643 + moe.n_redundant_experts = self.num_redundant_experts
644 + moe.experts.update_expert_map()
645 +
646 +
647 +class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA, Glm4MixtureOfExperts):
648 packed_modules_mapping = {
649 "qkv_proj": [
650 "q_proj",
@@ -659,7 +687,9 @@ class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
687 self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace
688 self.num_expert_groups = config.n_group
689
662 - self.moe_layers: list[SharedFusedMoE] = []
690 + self.moe_layers = []
691 + self.moe_mlp_layers: list[Glm4MoE] = []
692 +
693 example_moe = None
694 for layer in self.model.layers:
695 if isinstance(layer, PPMissingLayer):
@@ -669,33 +699,10 @@ class Glm4MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
699 if isinstance(layer.mlp, Glm4MoE):
700 # Pick last one layer since the first ones may be dense layers.
701 example_moe = layer.mlp
702 + self.moe_mlp_layers.append(layer.mlp)
703 self.moe_layers.append(layer.mlp.experts)
704
674 - if example_moe is None:
675 - raise RuntimeError("No Glm4MoE layer found in model.layers.")
676 -
677 - self.num_logical_experts = example_moe.n_logical_experts
678 - self.num_physical_experts = example_moe.n_physical_experts
679 - self.num_local_physical_experts = example_moe.n_local_physical_experts
680 - self.num_routed_experts = example_moe.n_routed_experts
681 - self.num_shared_experts = example_moe.n_shared_experts
682 - self.num_redundant_experts = example_moe.n_redundant_experts
683 -
684 - def set_eplb_state(
685 - self,
686 - expert_load_view: torch.Tensor,
687 - logical_to_physical_map: torch.Tensor,
688 - logical_replica_count: torch.Tensor,
689 - ) -> None:
690 - for layer_idx, layer in enumerate(self.moe_layers):
691 - # Register the expert weights.
692 - self.expert_weights.append(layer.get_expert_weights())
693 - layer.set_eplb_state(
694 - moe_layer_idx=layer_idx,
695 - expert_load_view=expert_load_view,
696 - logical_to_physical_map=logical_to_physical_map,
697 - logical_replica_count=logical_replica_count,
698 - )
705 + self.extract_moe_parameters(example_moe)
706
707 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
708 return self.model.get_input_embeddings(input_ids)
vllm/model_executor/models/glm4_moe_mtp.py
+31 -3
@@ -29,7 +29,7 @@ import torch
29 import torch.nn as nn
30 from transformers import PretrainedConfig
31
32 -from vllm.config import CacheConfig, VllmConfig
32 +from vllm.config import CacheConfig, ParallelConfig, VllmConfig
33 from vllm.model_executor.layers.fused_moe import FusedMoE
34 from vllm.model_executor.layers.layernorm import RMSNorm
35 from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -41,7 +41,12 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
41 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
42 from vllm.sequence import IntermediateTensors
43
44 -from .glm4_moe import Glm4MoeDecoderLayer, get_spec_layer_idx_from_weight_name
44 +from .glm4_moe import (
45 + Glm4MixtureOfExperts,
46 + Glm4MoE,
47 + Glm4MoeDecoderLayer,
48 + get_spec_layer_idx_from_weight_name,
49 +)
50 from .interfaces import SupportsPP
51 from .utils import maybe_prefix
52
@@ -73,6 +78,7 @@ class Glm4MoeMultiTokenPredictorLayer(nn.Module):
78 prefix: str,
79 cache_config: CacheConfig | None = None,
80 quant_config: QuantizationConfig | None = None,
81 + parallel_config: ParallelConfig | None = None,
82 ) -> None:
83 super().__init__()
84 self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -81,11 +87,13 @@ class Glm4MoeMultiTokenPredictorLayer(nn.Module):
87 self.shared_head = SharedHead(
88 config=config, prefix=prefix, quant_config=quant_config
89 )
90 + self.enable_eplb = parallel_config.enable_eplb
91 self.mtp_block = Glm4MoeDecoderLayer(
92 config=config,
93 cache_config=cache_config,
94 quant_config=quant_config,
95 prefix=prefix,
96 + enable_eplb=self.enable_eplb,
97 )
98
99 def forward(
@@ -127,6 +135,7 @@ class Glm4MoeMultiTokenPredictor(nn.Module):
135 f"{prefix}.layers.{idx}",
136 cache_config=vllm_config.cache_config,
137 quant_config=vllm_config.quant_config,
138 + parallel_config=vllm_config.parallel_config,
139 )
140 for idx in range(
141 self.mtp_start_layer_idx,
@@ -175,7 +184,7 @@ class Glm4MoeMultiTokenPredictor(nn.Module):
184 return logits
185
186
178 -class Glm4MoeMTP(nn.Module, SupportsPP):
187 +class Glm4MoeMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):
188 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
189 super().__init__()
190 self.config = vllm_config.model_config.hf_config
@@ -183,6 +192,25 @@ class Glm4MoeMTP(nn.Module, SupportsPP):
192 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
193 )
194
195 + self.expert_weights = []
196 +
197 + # Set MoE hyperparameters
198 + self.num_moe_layers = self.config.num_nextn_predict_layers
199 + self.num_expert_groups = self.config.n_group
200 +
201 + self.moe_layers: list[FusedMoE] = []
202 + self.moe_mlp_layers: list[Glm4MoE] = []
203 + example_moe = None
204 + for layer in self.model.layers.values():
205 + assert isinstance(layer, Glm4MoeMultiTokenPredictorLayer)
206 + layer = layer.mtp_block
207 + assert isinstance(layer, Glm4MoeDecoderLayer)
208 + if isinstance(layer.mlp, Glm4MoE):
209 + example_moe = layer.mlp
210 + self.moe_mlp_layers.append(layer.mlp)
211 + self.moe_layers.append(layer.mlp.experts)
212 + self.extract_moe_parameters(example_moe)
213 +
214 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
215 return self.model.get_input_embeddings(input_ids)
216
vllm/model_executor/models/hunyuan_v1.py
+2 -18
@@ -374,7 +374,7 @@ class HunYuanSparseMoeBlock(nn.Module):
374 self.tp_size = get_tensor_model_parallel_world_size()
375
376 self.ep_group = get_ep_group().device_group
377 - self.ep_rank = self.ep_group.rank()
377 + self.ep_rank = get_ep_group().rank_in_group
378 self.ep_size = self.ep_group.size()
379 self.n_routed_experts = config.num_experts
380
@@ -1007,7 +1007,7 @@ class HunYuanMoEV1Base(HunyuanV1ModelBase, MixtureOfExperts):
1007 # Set MoE hyperparameters
1008 self.expert_weights = []
1009 self.num_expert_groups = 1
1010 - self.moe_layers: list[SharedFusedMoE] = []
1010 + self.moe_layers = []
1011 example_layer = None
1012 for layer in self.model.layers:
1013 if isinstance(layer, PPMissingLayer):
@@ -1028,22 +1028,6 @@ class HunYuanMoEV1Base(HunyuanV1ModelBase, MixtureOfExperts):
1028 self.num_routed_experts = example_layer.n_routed_experts
1029 self.num_redundant_experts = example_layer.n_redundant_experts
1030
1031 - def set_eplb_state(
1032 - self,
1033 - expert_load_view: torch.Tensor,
1034 - logical_to_physical_map: torch.Tensor,
1035 - logical_replica_count: torch.Tensor,
1036 - ) -> None:
1037 - for layer_idx, layer in enumerate(self.moe_layers):
1038 - self.expert_weights.append(layer.get_expert_weights())
1039 - # Register the expert weights.
1040 - layer.set_eplb_state(
1041 - moe_layer_idx=layer_idx,
1042 - expert_load_view=expert_load_view,
1043 - logical_to_physical_map=logical_to_physical_map,
1044 - logical_replica_count=logical_replica_count,
1045 - )
1046 -
1031 def update_physical_experts_metadata(
1032 self,
1033 num_physical_experts: int,
vllm/model_executor/models/interfaces.py
+13 -1
@@ -14,6 +14,7 @@ from typing import (
14
15 import numpy as np
16 import torch
17 +import torch.nn as nn
18 from torch import Tensor
19 from transformers import PretrainedConfig
20 from transformers.models.whisper.tokenization_whisper import LANGUAGES
@@ -641,6 +642,9 @@ class MixtureOfExperts(Protocol):
642 num_redundant_experts: int
643 """Number of redundant experts in this model."""
644
645 + moe_layers: Iterable[nn.Module]
646 + """List of MoE layers in this model."""
647 +
648 def set_eplb_state(
649 self,
650 expert_load_view: Tensor,
@@ -663,7 +667,15 @@ class MixtureOfExperts(Protocol):
667 logical_to_physical_map: Mapping from logical to physical experts.
668 logical_replica_count: Count of replicas for each logical expert.
669 """
666 - ...
670 + for layer_idx, layer in enumerate(self.moe_layers):
671 + # Register the expert weights.
672 + self.expert_weights.append(layer.get_expert_weights())
673 + layer.set_eplb_state(
674 + moe_layer_idx=layer_idx,
675 + expert_load_view=expert_load_view,
676 + logical_to_physical_map=logical_to_physical_map,
677 + logical_replica_count=logical_replica_count,
678 + )
679
680 def update_physical_experts_metadata(
681 self,
vllm/model_executor/models/lfm2_moe.py
+2 -18
@@ -105,7 +105,7 @@ class Lfm2MoeSparseMoeBlock(nn.Module):
105 self.routed_scaling_factor = config.routed_scaling_factor
106
107 self.ep_group = get_ep_group().device_group
108 - self.ep_rank = self.ep_group.rank()
108 + self.ep_rank = get_ep_group().rank_in_group
109 self.ep_size = self.ep_group.size()
110 self.n_routed_experts = config.num_experts
111
@@ -707,7 +707,7 @@ class Lfm2MoeForCausalLM(
707 # Set MoE hyperparameters
708 self.expert_weights = []
709
710 - self.moe_layers: list[FusedMoE] = []
710 + self.moe_layers = []
711 example_layer = None
712 for layer in self.model.layers:
713 if isinstance(layer, PPMissingLayer):
@@ -737,22 +737,6 @@ class Lfm2MoeForCausalLM(
737 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
738 return self.model.get_input_embeddings(input_ids)
739
740 - def set_eplb_state(
741 - self,
742 - expert_load_view: torch.Tensor,
743 - logical_to_physical_map: torch.Tensor,
744 - logical_replica_count: torch.Tensor,
745 - ) -> None:
746 - for layer_idx, layer in enumerate(self.moe_layers):
747 - # Register the expert weights.
748 - self.expert_weights.append(layer.get_expert_weights())
749 - layer.set_eplb_state(
750 - moe_layer_idx=layer_idx,
751 - expert_load_view=expert_load_view,
752 - logical_to_physical_map=logical_to_physical_map,
753 - logical_replica_count=logical_replica_count,
754 - )
755 -
740 def update_physical_experts_metadata(
741 self,
742 num_physical_experts: int,
vllm/model_executor/models/llama4.py
+80 -2
@@ -30,9 +30,11 @@ from vllm.attention.layers.chunked_local_attention import ChunkedLocalAttention
30 from vllm.compilation.decorators import support_torch_compile
31 from vllm.config import CacheConfig, VllmConfig
32 from vllm.distributed import (
33 + get_ep_group,
34 get_tensor_model_parallel_world_size,
35 tensor_model_parallel_all_gather,
36 )
37 +from vllm.logger import init_logger
38 from vllm.model_executor.layers.fused_moe import SharedFusedMoE
39 from vllm.model_executor.layers.layernorm import RMSNorm
40 from vllm.model_executor.layers.linear import (
@@ -46,6 +48,7 @@ from vllm.model_executor.model_loader.weight_utils import (
48 default_weight_loader,
49 maybe_remap_kv_scale_name,
50 )
51 +from vllm.model_executor.models.interfaces import MixtureOfExperts
52 from vllm.model_executor.models.utils import sequence_parallel_chunk
53
54 from .llama import LlamaForCausalLM, LlamaMLP, LlamaModel
@@ -56,6 +59,8 @@ from .utils import (
59 is_pp_missing_parameter,
60 )
61
62 +logger = init_logger(__name__)
63 +
64
65 class Llama4MoE(nn.Module):
66 @staticmethod
@@ -80,6 +85,9 @@ class Llama4MoE(nn.Module):
85 self.tp_size = get_tensor_model_parallel_world_size()
86 self.top_k = config.num_experts_per_tok
87 self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
88 + self.ep_group = get_ep_group().device_group
89 + self.ep_rank = get_ep_group().rank_in_group
90 + self.ep_size = self.ep_group.size()
91
92 intermediate_size_moe = config.intermediate_size
93 self.router = ReplicatedLinear(
@@ -101,6 +109,20 @@ class Llama4MoE(nn.Module):
109 disable_tp=self.is_sequence_parallel,
110 )
111
112 + # Load balancing settings.
113 + eplb_config = parallel_config.eplb_config if parallel_config else None
114 + self.enable_eplb = parallel_config.enable_eplb if parallel_config else False
115 + self.n_redundant_experts = (
116 + eplb_config.num_redundant_experts if eplb_config else 0
117 + )
118 +
119 + self.n_routed_experts: int = config.num_local_experts
120 + self.n_logical_experts = self.n_routed_experts
121 + self.n_shared_experts: int = 1
122 + self.n_local_experts: int = config.num_local_experts
123 + self.n_physical_experts = self.n_local_experts + self.n_redundant_experts
124 + self.n_local_physical_experts = self.n_physical_experts // self.ep_size
125 +
126 self.experts = SharedFusedMoE(
127 shared_experts=self.shared_expert,
128 num_experts=config.num_local_experts,
@@ -114,6 +136,8 @@ class Llama4MoE(nn.Module):
136 quant_config=quant_config,
137 prefix=f"{prefix}.experts",
138 is_sequence_parallel=self.is_sequence_parallel,
139 + enable_eplb=self.enable_eplb,
140 + num_redundant_experts=self.n_redundant_experts,
141 )
142
143 def forward(self, hidden_states):
@@ -378,6 +402,9 @@ class Llama4Model(LlamaModel):
402 layer_type: type[Llama4DecoderLayer] = Llama4DecoderLayer,
403 ):
404 self.num_experts = vllm_config.model_config.hf_config.num_local_experts
405 + self.n_redundant_experts = (
406 + vllm_config.parallel_config.eplb_config.num_redundant_experts
407 + )
408 super().__init__(vllm_config=vllm_config, prefix=prefix, layer_type=layer_type)
409
410 def load_moe_expert_weights(
@@ -499,7 +526,6 @@ class Llama4Model(LlamaModel):
526 shard_id=shard_id,
527 expert_id=expert_id,
528 )
502 -
529 loaded_params.add(full_param_name)
530 expert_param_loaded = True
531
@@ -526,6 +552,7 @@ class Llama4Model(LlamaModel):
552 ckpt_down_proj_name="down_proj",
553 ckpt_up_proj_name="up_proj",
554 num_experts=self.num_experts,
555 + num_redundant_experts=self.n_redundant_experts,
556 )
557 # Expert parameter mapping for the case where the expert weights are
558 # fused into a single weight tensor.
@@ -683,7 +710,7 @@ class Llama4Model(LlamaModel):
710 return loaded_params
711
712
686 -class Llama4ForCausalLM(LlamaForCausalLM):
713 +class Llama4ForCausalLM(LlamaForCausalLM, MixtureOfExperts):
714 packed_modules_mapping = {
715 "qkv_proj": ["q_proj", "k_proj", "v_proj"],
716 "gate_up_proj": ["gate_proj", "up_proj"],
@@ -702,6 +729,57 @@ class Llama4ForCausalLM(LlamaForCausalLM):
729 super().__init__(
730 vllm_config=vllm_config, prefix=prefix, layer_type=Llama4DecoderLayer
731 )
732 + # Set MoE hyperparameters
733 + self.set_moe_parameters()
734 +
735 + def set_moe_parameters(self):
736 + self.expert_weights = []
737 +
738 + self.moe_layers = []
739 + example_moe = None
740 + for layer in self.model.layers:
741 + assert isinstance(layer, Llama4DecoderLayer)
742 + if isinstance(layer.feed_forward, Llama4MoE):
743 + # Pick last one layer since the first ones may be dense layers.
744 + example_moe = layer.feed_forward
745 + self.moe_layers.append(layer.feed_forward.experts)
746 +
747 + if example_moe is None:
748 + self.num_moe_layers = 0
749 + self.num_expert_groups = 0
750 + self.num_logical_experts = 0
751 + self.num_physical_experts = 0
752 + self.num_local_physical_experts = 0
753 + self.num_routed_experts = 0
754 + self.num_shared_experts = 0
755 + self.num_redundant_experts = 0
756 + logger.warning("No Llama4MoE layer found in model.layers.")
757 + else:
758 + self.num_moe_layers = len(self.moe_layers)
759 + self.num_expert_groups = 1
760 + self.num_logical_experts = example_moe.n_logical_experts
761 + self.num_physical_experts = example_moe.n_physical_experts
762 + self.num_local_physical_experts = example_moe.n_local_physical_experts
763 + self.num_routed_experts = example_moe.n_routed_experts
764 + self.num_shared_experts = example_moe.n_shared_experts
765 + self.num_redundant_experts = example_moe.n_redundant_experts
766 +
767 + def update_physical_experts_metadata(
768 + self,
769 + num_physical_experts: int,
770 + num_local_physical_experts: int,
771 + ) -> None:
772 + assert self.num_local_physical_experts == num_local_physical_experts
773 + self.num_physical_experts = num_physical_experts
774 + self.num_local_physical_experts = num_local_physical_experts
775 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
776 + for layer in self.model.layers:
777 + if isinstance(layer.feed_forward, Llama4MoE):
778 + moe = layer.feed_forward
779 + moe.n_local_physical_experts = num_local_physical_experts
780 + moe.n_physical_experts = num_physical_experts
781 + moe.n_redundant_experts = self.num_redundant_experts
782 + moe.experts.update_expert_map()
783
784 def _init_model(
785 self,
vllm/model_executor/models/llama4_eagle.py
+3
@@ -189,6 +189,9 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):
189 self.config.vocab_size, scale=logit_scale
190 )
191
192 + # Set MoE hyperparameters
193 + self.set_moe_parameters()
194 +
195 def get_language_model(self) -> torch.nn.Module:
196 return self.model
197
vllm/model_executor/models/minicpm.py
+3
@@ -578,6 +578,7 @@ class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
578 cache_config = vllm_config.cache_config
579 quant_config = vllm_config.quant_config
580 lora_config = vllm_config.lora_config
581 + parallel_config = vllm_config.parallel_config
582
583 self.prefix = prefix
584 self.vllm_config = vllm_config
@@ -613,6 +614,8 @@ class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
614 self.make_empty_intermediate_tensors = (
615 self.model.make_empty_intermediate_tensors
616 )
617 + if parallel_config.enable_eplb and getattr(config, "num_experts", 0) > 0:
618 + raise NotImplementedError("EPLB is not supported for MiniCPM yet.")
619
620 def _init_model(self, *, vllm_config: VllmConfig, prefix: str = ""):
621 return MiniCPMModel(vllm_config=vllm_config, prefix=prefix)
vllm/model_executor/models/mixtral.py
+2 -18
@@ -98,7 +98,7 @@ class MixtralMoE(nn.Module):
98 self.hidden_size = hidden_size
99
100 self.ep_group = get_ep_group().device_group
101 - self.ep_rank = self.ep_group.rank()
101 + self.ep_rank = get_ep_group().rank_in_group
102 self.ep_size = self.ep_group.size()
103
104 # Expert Parallelism Load balancing settings.
@@ -546,7 +546,7 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
546 )
547
548 self.expert_weights = []
549 - self.moe_layers: list[FusedMoE] = []
549 + self.moe_layers = []
550 example_moe = None
551
552 for layer in self.model.layers:
@@ -572,22 +572,6 @@ class MixtralForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
572 self.num_expert_groups = 1
573 self.num_shared_experts = 0
574
575 - def set_eplb_state(
576 - self,
577 - expert_load_view: torch.Tensor,
578 - logical_to_physical_map: torch.Tensor,
579 - logical_replica_count: torch.Tensor,
580 - ) -> None:
581 - for layer_idx, layer in enumerate(self.moe_layers):
582 - # Register the expert weights.
583 - self.expert_weights.append(layer.get_expert_weights())
584 - layer.set_eplb_state(
585 - moe_layer_idx=layer_idx,
586 - expert_load_view=expert_load_view,
587 - logical_to_physical_map=logical_to_physical_map,
588 - logical_replica_count=logical_replica_count,
589 - )
590 -
575 def update_physical_experts_metadata(
576 self,
577 num_physical_experts: int,
vllm/model_executor/models/mllama4.py
+31 -1
@@ -65,6 +65,7 @@ from vllm.sequence import IntermediateTensors
65 from vllm.utils.tensor_schema import TensorSchema, TensorShape
66
67 from .interfaces import (
68 + MixtureOfExperts,
69 MultiModalEmbeddings,
70 SupportsEagle3,
71 SupportsMultiModal,
@@ -723,7 +724,7 @@ class Mllama4DummyInputsBuilder(BaseDummyInputsBuilder[Mllama4ProcessingInfo]):
724 dummy_inputs=Mllama4DummyInputsBuilder,
725 )
726 class Llama4ForConditionalGeneration(
726 - nn.Module, SupportsMultiModal, SupportsPP, SupportsEagle3
727 + nn.Module, SupportsMultiModal, SupportsPP, MixtureOfExperts, SupportsEagle3
728 ):
729 merge_by_field_config = True
730
@@ -776,6 +777,17 @@ class Llama4ForConditionalGeneration(
777 self.language_model.make_empty_intermediate_tensors
778 )
779
780 + # Set MoE hyperparameters
781 + self.num_expert_groups = 1
782 + self.num_logical_experts = self.language_model.num_logical_experts
783 + self.num_physical_experts = self.language_model.num_physical_experts
784 + self.num_local_physical_experts = self.language_model.num_local_physical_experts
785 + self.num_routed_experts = self.language_model.num_routed_experts
786 + self.num_shared_experts = self.language_model.num_shared_experts
787 + self.num_redundant_experts = self.language_model.num_redundant_experts
788 + self.moe_layers = self.language_model.moe_layers
789 + self.num_moe_layers = len(self.moe_layers)
790 +
791 def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
792 """Set which layers should output auxiliary hidden states for EAGLE3."""
793 # Delegate to underlying language model (Llama4ForCausalLM)
@@ -792,6 +804,24 @@ class Llama4ForConditionalGeneration(
804 assert hasattr(self.language_model, "get_eagle3_aux_hidden_state_layers")
805 return self.language_model.get_eagle3_aux_hidden_state_layers()
806
807 + def set_eplb_state(
808 + self,
809 + expert_load_view: torch.Tensor,
810 + logical_to_physical_map: torch.Tensor,
811 + logical_replica_count: torch.Tensor,
812 + ):
813 + self.language_model.set_eplb_state(
814 + expert_load_view, logical_to_physical_map, logical_replica_count
815 + )
816 + self.expert_weights = self.language_model.expert_weights
817 +
818 + def update_physical_experts_metadata(
819 + self, num_physical_experts: int, num_local_physical_experts: int
820 + ):
821 + self.language_model.update_physical_experts_metadata(
822 + num_physical_experts, num_local_physical_experts
823 + )
824 +
825 def _parse_and_validate_image_input(
826 self, **kwargs: object
827 ) -> Llama4ImagePatchInputs | None:
vllm/model_executor/models/nemotron_h.py
+1 -17
@@ -807,7 +807,7 @@ class NemotronHForCausalLM(
807 self.expert_weights = []
808 self.num_expert_groups = config.n_group
809
810 - self.moe_layers: list[SharedFusedMoE] = []
810 + self.moe_layers = []
811 example_moe = None
812 for layer in self.model.layers:
813 if isinstance(layer, NemotronHMoEDecoderLayer):
@@ -824,22 +824,6 @@ class NemotronHForCausalLM(
824 self.num_shared_experts = example_moe.n_shared_experts
825 self.num_redundant_experts = example_moe.n_redundant_experts
826
827 - def set_eplb_state(
828 - self,
829 - expert_load_view: torch.Tensor,
830 - logical_to_physical_map: torch.Tensor,
831 - logical_replica_count: torch.Tensor,
832 - ) -> None:
833 - for layer_idx, layer in enumerate(self.moe_layers):
834 - # Register the expert weights.
835 - self.expert_weights.append(layer.get_expert_weights())
836 - layer.set_eplb_state(
837 - moe_layer_idx=layer_idx,
838 - expert_load_view=expert_load_view,
839 - logical_to_physical_map=logical_to_physical_map,
840 - logical_replica_count=logical_replica_count,
841 - )
842 -
827 def update_physical_experts_metadata(
828 self,
829 num_physical_experts: int,
vllm/model_executor/models/openpangu.py
+1 -17
@@ -1009,7 +1009,7 @@ class OpenPanguMoEModel(OpenPanguModelBase, MixtureOfExperts):
1009 self.num_moe_layers = config.num_hidden_layers - config.first_k_dense_replace
1010 self.num_expert_groups = 1
1011
1012 - self.moe_layers: list[SharedFusedMoE] = []
1012 + self.moe_layers = []
1013 example_moe = None
1014 for layer in self.model.layers:
1015 if isinstance(layer, PPMissingLayer):
@@ -1031,22 +1031,6 @@ class OpenPanguMoEModel(OpenPanguModelBase, MixtureOfExperts):
1031 self.n_shared_experts = example_moe.n_shared_experts
1032 self.num_redundant_experts = example_moe.n_redundant_experts
1033
1034 - def set_eplb_state(
1035 - self,
1036 - expert_load_view: torch.Tensor,
1037 - logical_to_physical_map: torch.Tensor,
1038 - logical_replica_count: torch.Tensor,
1039 - ) -> None:
1040 - for layer_idx, layer in enumerate(self.moe_layers):
1041 - # Register the expert weights.
1042 - self.expert_weights.append(layer.get_expert_weights())
1043 - layer.set_eplb_state(
1044 - moe_layer_idx=layer_idx,
1045 - expert_load_view=expert_load_view,
1046 - logical_to_physical_map=logical_to_physical_map,
1047 - logical_replica_count=logical_replica_count,
1048 - )
1049 -
1034 def update_physical_experts_metadata(
1035 self,
1036 num_physical_experts: int,
vllm/model_executor/models/qwen3_moe.py
+2 -18
@@ -132,7 +132,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
132 self.tp_size = get_tensor_model_parallel_world_size()
133
134 self.ep_group = get_ep_group().device_group
135 - self.ep_rank = self.ep_group.rank()
135 + self.ep_rank = get_ep_group().rank_in_group
136 self.ep_size = self.ep_group.size()
137 self.n_routed_experts = config.num_experts
138
@@ -665,7 +665,7 @@ class Qwen3MoeForCausalLM(
665 # Set MoE hyperparameters
666 self.expert_weights = []
667
668 - self.moe_layers: list[FusedMoE] = []
668 + self.moe_layers = []
669 example_layer = None
670 for layer in self.model.layers:
671 if isinstance(layer, PPMissingLayer):
@@ -688,22 +688,6 @@ class Qwen3MoeForCausalLM(
688 self.num_routed_experts = example_layer.n_routed_experts
689 self.num_redundant_experts = example_layer.n_redundant_experts
690
691 - def set_eplb_state(
692 - self,
693 - expert_load_view: torch.Tensor,
694 - logical_to_physical_map: torch.Tensor,
695 - logical_replica_count: torch.Tensor,
696 - ) -> None:
697 - for layer_idx, layer in enumerate(self.moe_layers):
698 - # Register the expert weights.
699 - self.expert_weights.append(layer.get_expert_weights())
700 - layer.set_eplb_state(
701 - moe_layer_idx=layer_idx,
702 - expert_load_view=expert_load_view,
703 - logical_to_physical_map=logical_to_physical_map,
704 - logical_replica_count=logical_replica_count,
705 - )
706 -
691 def update_physical_experts_metadata(
692 self,
693 num_physical_experts: int,
vllm/model_executor/models/qwen3_next.py
+52 -59
@@ -107,7 +107,7 @@ class Qwen3NextSparseMoeBlock(nn.Module):
107 self.tp_size = get_tensor_model_parallel_world_size()
108
109 self.ep_group = get_ep_group().device_group
110 - self.ep_rank = self.ep_group.rank()
110 + self.ep_rank = get_ep_group().rank_in_group
111 self.ep_size = self.ep_group.size()
112 self.n_routed_experts = config.num_experts
113
@@ -1095,8 +1095,57 @@ class Qwen3NextModel(nn.Module):
1095 return loaded_params
1096
1097
1098 +class QwenNextMixtureOfExperts(MixtureOfExperts):
1099 + def update_physical_experts_metadata(
1100 + self,
1101 + num_physical_experts: int,
1102 + num_local_physical_experts: int,
1103 + ) -> None:
1104 + assert self.num_local_physical_experts == num_local_physical_experts
1105 + self.num_physical_experts = num_physical_experts
1106 + self.num_local_physical_experts = num_local_physical_experts
1107 + self.num_redundant_experts = num_physical_experts - self.num_logical_experts
1108 + for layer in self.model.layers:
1109 + if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
1110 + moe = layer.mlp
1111 + moe.n_local_physical_experts = num_local_physical_experts
1112 + moe.n_physical_experts = num_physical_experts
1113 + moe.n_redundant_experts = self.num_redundant_experts
1114 + moe.experts.update_expert_map()
1115 +
1116 + def set_moe_parameters(self):
1117 + self.expert_weights = []
1118 +
1119 + self.moe_layers = []
1120 + example_moe = None
1121 + for layer in self.model.layers:
1122 + if isinstance(layer, Qwen3NextDecoderLayer) and isinstance(
1123 + layer.mlp, Qwen3NextSparseMoeBlock
1124 + ):
1125 + example_moe = layer.mlp
1126 + self.moe_layers.append(layer.mlp.experts)
1127 +
1128 + if example_moe is None:
1129 + raise RuntimeError("No Qwen3Next layer found in the model.layers.")
1130 +
1131 + # Set MoE hyperparameters
1132 + self.num_moe_layers = len(self.moe_layers)
1133 + self.num_expert_groups = 1
1134 + self.num_shared_experts = 0
1135 + self.num_logical_experts = example_moe.n_logical_experts
1136 + self.num_physical_experts = example_moe.n_physical_experts
1137 + self.num_local_physical_experts = example_moe.n_local_physical_experts
1138 + self.num_routed_experts = example_moe.n_routed_experts
1139 + self.num_redundant_experts = example_moe.n_redundant_experts
1140 +
1141 +
1142 class Qwen3NextForCausalLM(
1099 - nn.Module, HasInnerState, SupportsLoRA, SupportsPP, MixtureOfExperts, IsHybrid
1143 + nn.Module,
1144 + HasInnerState,
1145 + SupportsLoRA,
1146 + SupportsPP,
1147 + QwenNextMixtureOfExperts,
1148 + IsHybrid,
1149 ):
1150 packed_modules_mapping = {
1151 "qkv_proj": [
@@ -1147,63 +1196,7 @@ class Qwen3NextForCausalLM(
1196 )
1197
1198 # Set MoE hyperparameters
1150 - self.expert_weights = []
1151 -
1152 - self.moe_layers: list[SharedFusedMoE] = []
1153 - example_layer = None
1154 - for layer in self.model.layers:
1155 - if isinstance(layer, PPMissingLayer):
1156 - continue
1157 -
1158 - assert isinstance(layer, Qwen3NextDecoderLayer)
1159 - if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
1160 - example_layer = layer.mlp
1161 - self.moe_layers.append(layer.mlp.experts)
1162 -
1163 - if example_layer is None:
1164 - raise RuntimeError("No Qwen3Next layer found in the model.layers.")
1165 -
1166 - self.num_moe_layers = len(self.moe_layers)
1167 - self.num_expert_groups = 1
1168 - self.num_shared_experts = 0
1169 - self.num_logical_experts = example_layer.n_logical_experts
1170 - self.num_physical_experts = example_layer.n_physical_experts
1171 - self.num_local_physical_experts = example_layer.n_local_physical_experts
1172 - self.num_routed_experts = example_layer.n_routed_experts
1173 - self.num_redundant_experts = example_layer.n_redundant_experts
1174 -
1175 - def set_eplb_state(
1176 - self,
1177 - expert_load_view: torch.Tensor,
1178 - logical_to_physical_map: torch.Tensor,
1179 - logical_replica_count: torch.Tensor,
1180 - ) -> None:
1181 - for layer_idx, layer in enumerate(self.moe_layers):
1182 - # Register the expert weights.
1183 - self.expert_weights.append(layer.get_expert_weights())
1184 - layer.set_eplb_state(
1185 - moe_layer_idx=layer_idx,
1186 - expert_load_view=expert_load_view,
1187 - logical_to_physical_map=logical_to_physical_map,
1188 - logical_replica_count=logical_replica_count,
1189 - )
1190 -
1191 - def update_physical_experts_metadata(
1192 - self,
1193 - num_physical_experts: int,
1194 - num_local_physical_experts: int,
1195 - ) -> None:
1196 - assert self.num_local_physical_experts == num_local_physical_experts
1197 - self.num_physical_experts = num_physical_experts
1198 - self.num_local_physical_experts = num_local_physical_experts
1199 - self.num_redundant_experts = num_physical_experts - self.num_logical_experts
1200 - for layer in self.model.layers:
1201 - if isinstance(layer.mlp, Qwen3NextSparseMoeBlock):
1202 - moe = layer.mlp
1203 - moe.n_local_physical_experts = num_local_physical_experts
1204 - moe.n_physical_experts = num_physical_experts
1205 - moe.n_redundant_experts = self.num_redundant_experts
1206 - moe.experts.update_expert_map()
1199 + self.set_moe_parameters()
1200
1201 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
1202 return self.model.get_input_embeddings(input_ids)
vllm/model_executor/models/qwen3_next_mtp.py
+3 -1
@@ -23,6 +23,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
23 from vllm.model_executor.models.qwen3_next import (
24 Qwen3NextDecoderLayer,
25 Qwen3NextRMSNorm,
26 + QwenNextMixtureOfExperts,
27 )
28 from vllm.sequence import IntermediateTensors
29 from vllm.transformers_utils.configs import Qwen3NextConfig
@@ -226,7 +227,7 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
227
228
229 @support_torch_compile
229 -class Qwen3NextMTP(nn.Module, SupportsPP):
230 +class Qwen3NextMTP(nn.Module, SupportsPP, QwenNextMixtureOfExperts):
231 packed_modules_mapping = {
232 "qkv_proj": [
233 "q_proj",
@@ -265,6 +266,7 @@ class Qwen3NextMTP(nn.Module, SupportsPP):
266 self.make_empty_intermediate_tensors = (
267 self.model.make_empty_intermediate_tensors
268 )
269 + self.set_moe_parameters()
270
271 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
272 return self.model.get_input_embeddings(input_ids)
vllm/model_executor/models/transformers/moe.py
+6 -4
@@ -125,7 +125,7 @@ class MoEMixin(MixtureOfExperts):
125 logical_to_physical_map: torch.Tensor,
126 logical_replica_count: torch.Tensor,
127 ):
128 - for moe_layer_idx, mlp_layer in enumerate(self.mlp_layers):
128 + for moe_layer_idx, mlp_layer in enumerate(self.mlp_moe_layers):
129 mlp_layer.experts.set_eplb_state(
130 moe_layer_idx=moe_layer_idx,
131 expert_load_view=expert_load_view,
@@ -142,7 +142,7 @@ class MoEMixin(MixtureOfExperts):
142 self.num_physical_experts = num_physical_experts
143 self.num_local_physical_experts = num_local_physical_experts
144 self.num_redundant_experts = num_physical_experts - self.num_logical_experts
145 - for mlp in self.mlp_layers:
145 + for mlp in self.mlp_moe_layers:
146 mlp.n_local_physical_experts = num_local_physical_experts
147 mlp.n_physical_experts = num_physical_experts
148 mlp.n_redundant_experts = self.num_redundant_experts
@@ -240,7 +240,8 @@ class MoEMixin(MixtureOfExperts):
240 # MixtureOfExperts mixin settings
241 ep_size = get_ep_group().world_size
242
243 - self.mlp_layers = [] # Used for MixtureOfExperts methods
243 + self.mlp_moe_layers = [] # Used for MixtureOfExperts methods
244 + self.moe_layers = []
245 self.expert_weights = []
246 self.num_moe_layers = 0
247 self.num_expert_groups = 1 if num_expert_group is None else num_expert_group
@@ -298,7 +299,8 @@ class MoEMixin(MixtureOfExperts):
299 mlp.experts = fused_experts
300 log_replacement(qual_name, experts, fused_experts)
301 # Update MixtureOfExperts mixin state
301 - self.mlp_layers.append(mlp)
302 + self.mlp_moe_layers.append(mlp)
303 + self.moe_layers.append(fused_experts)
304 self.expert_weights.append(fused_experts.get_expert_weights())
305 self.num_moe_layers += 1
306 # If results are not all-reduced in FusedMoE, ensure they
vllm/v1/spec_decode/medusa.py
+5
@@ -8,6 +8,7 @@ from vllm.config import VllmConfig
8 from vllm.forward_context import set_forward_context
9 from vllm.logger import init_logger
10 from vllm.model_executor.model_loader import get_model
11 +from vllm.model_executor.models.interfaces import is_mixture_of_experts
12 from vllm.v1.sample.metadata import SamplingMetadata
13
14 # Initialize logger
@@ -56,6 +57,10 @@ class MedusaProposer:
57 vllm_config=self.vllm_config,
58 model_config=self.vllm_config.speculative_config.draft_model_config,
59 )
60 + assert not (
61 + is_mixture_of_experts(self.model)
62 + and self.vllm_config.parallel_config.enable_eplb
63 + ), "EPLB for Medusa is not supported"
64
65 @torch.inference_mode()
66 def dummy_run(self, num_tokens: int) -> None:
vllm/v1/worker/gpu_model_runner.py
+55 -33
@@ -2046,7 +2046,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2046 model = self.get_model()
2047 assert is_mixture_of_experts(model)
2048 self.eplb_state.step(
2049 - model,
2049 is_dummy,
2050 is_profile,
2051 log_stats=self.parallel_config.eplb_config.log_balancedness,
@@ -2803,7 +2802,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2802 else:
2803 indices = []
2804 offset = 0
2806 - assert spec_decode_metadata is not None
2805 + assert spec_decode_metadata is not None, (
2806 + "No spec decode metadata for medusa"
2807 + )
2808 for num_draft, tokens in zip(
2809 spec_decode_metadata.num_draft_tokens, sampled_token_ids
2810 ):
@@ -2934,32 +2935,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2935 self.model_config.model,
2936 scope="global",
2937 )
2937 - if eep_scale_up:
2938 - from vllm.distributed.parallel_state import get_ep_group
2939 -
2940 - num_local_physical_experts = torch.empty(1, dtype=torch.int32, device="cpu")
2941 - torch.distributed.broadcast(
2942 - num_local_physical_experts, group=get_ep_group().cpu_group, group_src=0
2943 - )
2944 - num_local_physical_experts = int(num_local_physical_experts.item())
2945 - new_ep_size = get_ep_group().world_size
2946 - global_expert_load, old_global_expert_indices = EplbState.recv_state()
2947 - num_logical_experts = global_expert_load.shape[1]
2948 - self.parallel_config.eplb_config.num_redundant_experts = (
2949 - num_local_physical_experts * new_ep_size - num_logical_experts
2950 - )
2951 - assert old_global_expert_indices.shape[1] % num_local_physical_experts == 0
2952 - old_ep_size = (
2953 - old_global_expert_indices.shape[1] // num_local_physical_experts
2954 - )
2955 - rank_mapping = {
2956 - old_ep_rank: old_ep_rank for old_ep_rank in range(old_ep_size)
2957 - }
2958 - else:
2959 - global_expert_load = None
2960 - old_global_expert_indices = None
2961 - rank_mapping = None
2938 + global_expert_loads, old_global_expert_indices_per_model, rank_mapping = (
2939 + EplbState.get_eep_state(self.parallel_config)
2940 + if eep_scale_up
2941 + else (None, None, None)
2942 + )
2943
2944 + if self.parallel_config.enable_eplb:
2945 + self.eplb_state = EplbState(self.parallel_config, self.device)
2946 + eplb_models = 0
2947 with DeviceMemoryProfiler() as m:
2948 time_before_load = time.perf_counter()
2949 model_loader = get_model_loader(self.load_config)
@@ -2971,8 +2955,39 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2955 self.model, self.vllm_config, self.device
2956 )
2957 if hasattr(self, "drafter"):
2974 - logger.info("Loading drafter model...")
2958 + logger.info_once("Loading drafter model...")
2959 self.drafter.load_model(self.model)
2960 + if (
2961 + hasattr(self.drafter, "model")
2962 + and is_mixture_of_experts(self.drafter.model)
2963 + and self.parallel_config.enable_eplb
2964 + ):
2965 + logger.info_once(
2966 + "EPLB is enabled for drafter model %s.",
2967 + self.vllm_config.speculative_config.draft_model_config.model,
2968 + )
2969 +
2970 + global_expert_load = (
2971 + global_expert_loads[eplb_models]
2972 + if global_expert_loads
2973 + else None
2974 + )
2975 + old_global_expert_indices = (
2976 + old_global_expert_indices_per_model[eplb_models]
2977 + if old_global_expert_indices_per_model
2978 + else None
2979 + )
2980 + if self.eplb_state is None:
2981 + self.eplb_state = EplbState(self.parallel_config, self.device)
2982 + self.eplb_state.add_model(
2983 + self.drafter.model,
2984 + self.vllm_config.speculative_config.draft_model_config,
2985 + global_expert_load,
2986 + old_global_expert_indices,
2987 + rank_mapping,
2988 + )
2989 + eplb_models += 1
2990 +
2991 if self.use_aux_hidden_state_outputs:
2992 if not supports_eagle3(self.get_model()):
2993 raise RuntimeError(
@@ -3001,18 +3016,25 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
3016 scope="local",
3017 )
3018 prepare_communication_buffer_for_model(self.model)
3004 -
3019 self.is_multimodal_pruning_enabled = (
3020 supports_multimodal_pruning(self.get_model())
3021 and self.model_config.multimodal_config.is_multimodal_pruning_enabled()
3022 )
3023
3024 if is_mixture_of_experts(self.model) and self.parallel_config.enable_eplb:
3011 - logger.info("EPLB is enabled for model %s.", self.model_config.model)
3012 - self.eplb_state = EplbState.build(
3025 + logger.info_once("EPLB is enabled for model %s.", self.model_config.model)
3026 + global_expert_load = (
3027 + global_expert_loads[eplb_models] if global_expert_loads else None
3028 + )
3029 + old_global_expert_indices = (
3030 + old_global_expert_indices_per_model[eplb_models]
3031 + if old_global_expert_indices_per_model
3032 + else None
3033 + )
3034 + assert self.eplb_state is not None
3035 + self.eplb_state.add_model(
3036 self.model,
3014 - self.device,
3015 - self.parallel_config,
3037 + self.model_config,
3038 global_expert_load,
3039 old_global_expert_indices,
3040 rank_mapping,
vllm/v1/worker/gpu_worker.py
+61 -35
@@ -32,6 +32,7 @@ from vllm.distributed.parallel_state import (
32 from vllm.logger import init_logger
33 from vllm.lora.request import LoRARequest
34 from vllm.model_executor import set_random_seed
35 +from vllm.model_executor.models.interfaces import is_mixture_of_experts
36 from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
37 from vllm.platforms import current_platform
38 from vllm.sequence import IntermediateTensors
@@ -613,7 +614,6 @@ class Worker(WorkerBase):
614 }
615 assert self.model_runner.eplb_state is not None
616 self.model_runner.eplb_state.rearrange(
616 - self.model_runner.model,
617 execute_shuffle=True,
618 global_expert_load=None,
619 rank_mapping=rank_mapping,
@@ -626,7 +626,7 @@ class Worker(WorkerBase):
626 self,
627 old_ep_size: int,
628 new_ep_size: int,
629 - global_expert_load: torch.Tensor | None,
629 + global_expert_loads: list[torch.Tensor] | None,
630 ) -> None:
631 from vllm.distributed.parallel_state import get_ep_group
632
@@ -635,9 +635,8 @@ class Worker(WorkerBase):
635 rank_mapping = {old_ep_rank: old_ep_rank for old_ep_rank in range(old_ep_size)}
636 assert self.model_runner.eplb_state is not None
637 self.model_runner.eplb_state.rearrange(
638 - self.model_runner.model,
638 execute_shuffle=True,
640 - global_expert_load=global_expert_load,
639 + global_expert_loads=global_expert_loads,
640 rank_mapping=rank_mapping,
641 )
642 if get_ep_group().rank == 0:
@@ -684,31 +683,56 @@ class Worker(WorkerBase):
683 get_ep_group,
684 prepare_communication_buffer_for_model,
685 )
687 - from vllm.model_executor.layers.fused_moe.layer import FusedMoEParallelConfig
686 + from vllm.model_executor.layers.fused_moe.layer import (
687 + FusedMoE,
688 + FusedMoEParallelConfig,
689 + )
690
691 parallel_config = self.vllm_config.parallel_config
690 - moe_modules = [
691 - module
692 - for module in self.model_runner.model.modules()
693 - if (
694 - module.__class__.__name__ == "FusedMoE"
695 - or module.__class__.__name__ == "SharedFusedMoE"
696 - )
697 - ]
698 - num_local_experts = moe_modules[0].moe_config.num_local_experts
699 - assert all(
700 - module.moe_config.num_local_experts == num_local_experts
701 - for module in moe_modules
702 - ), "All MoE modules must have the same number of experts"
703 - for module in moe_modules:
704 - module.moe_config.num_experts = num_local_experts * new_ep_size
705 - module.global_num_experts = module.moe_config.num_experts
706 - module.moe_parallel_config = FusedMoEParallelConfig.make(
707 - tp_size_=get_tp_group().world_size,
708 - dp_size_=get_dp_group().world_size,
709 - vllm_parallel_config=parallel_config,
710 - )
711 - module.moe_config.moe_parallel_config = module.moe_parallel_config
692 +
693 + def get_moe_modules(model: torch.nn.Module) -> list[FusedMoE]:
694 + return [
695 + module
696 + for module in model.modules()
697 + if (
698 + module.__class__.__name__ == "FusedMoE"
699 + or module.__class__.__name__ == "SharedFusedMoE"
700 + )
701 + ]
702 +
703 + def update_moe_modules(moe_modules: list[FusedMoE], num_local_experts: int):
704 + assert all(
705 + module.moe_config.num_local_experts == num_local_experts
706 + for module in moe_modules
707 + ), "All MoE modules must have the same number of experts"
708 + for module in moe_modules:
709 + module.moe_config.num_experts = num_local_experts * new_ep_size
710 + module.global_num_experts = module.moe_config.num_experts
711 + module.moe_parallel_config = FusedMoEParallelConfig.make(
712 + tp_size_=get_tp_group().world_size,
713 + dp_size_=get_dp_group().world_size,
714 + vllm_parallel_config=parallel_config,
715 + )
716 + module.moe_config.moe_parallel_config = module.moe_parallel_config
717 + return moe_modules
718 +
719 + model_moe_modules = get_moe_modules(self.model_runner.model)
720 + num_local_experts = model_moe_modules[0].moe_config.num_local_experts
721 +
722 + update_moe_modules(model_moe_modules, num_local_experts)
723 + drafter_model = None
724 + if hasattr(self.model_runner, "drafter") and hasattr(
725 + self.model_runner.drafter, "model"
726 + ):
727 + drafter_model = self.model_runner.drafter.model
728 + if drafter_model is not None and is_mixture_of_experts(drafter_model):
729 + drafter_moe_modules = get_moe_modules(drafter_model)
730 + # Check if drafter and model have matching configs
731 + assert (
732 + drafter_moe_modules[0].moe_config.num_local_experts == num_local_experts
733 + ), "Drafter and model configs should be the same"
734 + update_moe_modules(drafter_moe_modules, num_local_experts)
735 +
736 if new_ep_size < old_ep_size:
737 num_local_physical_experts = num_local_experts
738 assert self.model_runner.eplb_state is not None
@@ -719,7 +743,7 @@ class Worker(WorkerBase):
743 new_physical_experts
744 - self.model_runner.eplb_state.logical_replica_count.shape[1]
745 )
722 - global_expert_load = None
746 + global_expert_loads = None
747 else:
748 num_local_physical_experts = torch.tensor(
749 [num_local_experts], dtype=torch.int32, device="cpu"
@@ -730,18 +754,20 @@ class Worker(WorkerBase):
754 num_local_physical_experts = num_local_physical_experts.item()
755 new_physical_experts = num_local_physical_experts * new_ep_size
756 assert self.model_runner.eplb_state is not None
733 - global_expert_load = self.model_runner.eplb_state.rearrange(
734 - self.model_runner.model, execute_shuffle=False
757 + global_expert_loads = self.model_runner.eplb_state.rearrange(
758 + execute_shuffle=False
759 )
760 parallel_config.eplb_config.num_redundant_experts = (
737 - new_physical_experts - global_expert_load.shape[1]
761 + new_physical_experts - global_expert_loads[0].shape[1]
762 )
763 prepare_communication_buffer_for_model(self.model_runner.model)
764 + if drafter_model is not None:
765 + prepare_communication_buffer_for_model(drafter_model)
766 self.model_runner.model.update_physical_experts_metadata(
767 num_physical_experts=new_physical_experts,
768 num_local_physical_experts=num_local_physical_experts,
769 )
744 - return global_expert_load
770 + return global_expert_loads
771
772 def reinitialize_distributed(
773 self, reconfig_request: ReconfigureDistributedRequest
@@ -782,11 +808,11 @@ class Worker(WorkerBase):
808 self.local_rank,
809 )
810
785 - global_expert_load = self._reconfigure_moe(old_ep_size, new_ep_size)
811 + global_expert_loads = self._reconfigure_moe(old_ep_size, new_ep_size)
812
813 if new_ep_size > old_ep_size:
788 - assert global_expert_load is not None
789 - self._eplb_after_scale_up(old_ep_size, new_ep_size, global_expert_load)
814 + assert global_expert_loads is not None
815 + self._eplb_after_scale_up(old_ep_size, new_ep_size, global_expert_loads)
816
817 def save_sharded_state(
818 self,