diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 19fc15b275..2602edf2df 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -550,10 +550,7 @@ class ElasticEPScalingExecutor: is_async_enabled = eplb_state.is_async run_async = async_op and is_async_enabled eplb_state.is_async = run_async - if rank_mapping is None: - eplb_state.rearrange() - else: - eplb_state.rearrange(rank_mapping=rank_mapping) + eplb_state.rearrange(rank_mapping=rank_mapping, use_last_expert_load=True) if not run_async: # Wait for non-blocking expert resharding copies before continuing # the Elastic EP reconfiguration. diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index e4d805b9fe..205226cde4 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -183,6 +183,8 @@ class EplbModelState: See: https://github.com/vllm-project/vllm/pull/22167#pullrequestreview-3086143856 """ + last_expert_load: torch.Tensor + """Global logical-expert totals from the last reshuffle; equal demand initially.""" model_name: str model: MixtureOfExperts expert_buffer: list[torch.Tensor] @@ -499,6 +501,11 @@ class EplbState: expert_load_pass=expert_load_pass, expert_load_pass_buffer=expert_load_pass_buffer, expert_load_window=expert_load_window, + last_expert_load=torch.ones( + (model.num_moe_layers, model.num_logical_experts), + dtype=torch.int64, + device=self.device, + ), model_name=model_config.model, model=model, expert_buffer=expert_buffer, @@ -744,6 +751,7 @@ class EplbState: self, is_profile: bool = False, rank_mapping: dict[int, int] | None = None, + use_last_expert_load: bool = False, ) -> torch.Tensor | None: """Rearrange the experts according to the current load. @@ -753,6 +761,7 @@ class EplbState: no memory movement will be performed. Default is False. rank_mapping (dict[int, int] | None): The rank mapping when scaling is done in EEP. + use_last_expert_load: Use the last reshuffle's logical-expert totals. """ ep_group = get_ep_group().device_group @@ -775,6 +784,13 @@ class EplbState: # Map the physical expert load to global logical experts global_expert_load_windows = [] for eplb_model_state in self.model_states.values(): + if use_last_expert_load: + # Distribute rank 0's saved counts through the existing all-reduce. + saved_load = eplb_model_state.last_expert_load + global_expert_load_windows.append( + saved_load.clone() if is_main_rank else torch.zeros_like(saved_load) + ) + continue expert_load_window = eplb_model_state.expert_load_window physical_to_logical = eplb_model_state.physical_to_logical_map invalid_idx = eplb_model_state.model.num_logical_experts @@ -835,6 +851,8 @@ class EplbState: for eplb_model_state, global_expert_load_window in zip( self.model_states.values(), global_expert_load_windows ): + if not is_profile: + eplb_model_state.last_expert_load = global_expert_load_window if not self.is_async or is_profile: # Get new expert mappings for the model. The policy runs on the # host, so the load window and current map have to come back.