diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index feacb03d28..d34b844bd3 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -46,6 +46,7 @@ from vllm.distributed.stateless_coordinator import StatelessGroupCoordinator from vllm.distributed.utils import StatelessProcessGroup from vllm.logger import init_logger from vllm.model_executor.models.interfaces import MixtureOfExperts +from vllm.platforms import current_platform from .async_worker import start_async_worker from .eplb_communicator import EplbCommunicator, create_eplb_communicator @@ -825,23 +826,80 @@ class EplbState: eplb_model_state.physical_to_logical_map.cpu(), ) - # Update expert weights - rearrange_expert_weights_inplace( - eplb_model_state.physical_to_logical_map, - new_physical_to_logical_map, - eplb_model_state.model.expert_weights, - eplb_model_state.expert_buffer, - ep_group, - eplb_model_state.communicator, - is_profile, - rank_mapping, - ) + skip_rearrange = False + if ( + current_platform.is_rocm() + and not is_profile + and rank_mapping is None + and bool((eplb_model_state.physical_to_logical_map >= 0).all()) + ): + logical_loads = global_expert_load_window.float() + ep_size = ep_group.size() - if not is_profile: - _commit_eplb_maps( - eplb_model_state, - new_physical_to_logical_map=new_physical_to_logical_map, + def rank_load_imbalance( + mapping: torch.Tensor, + logical_loads: torch.Tensor = logical_loads, + ep_size: int = ep_size, + ) -> float: + mapping = mapping.to( + device=logical_loads.device, + dtype=torch.long, + ) + replica_counts = torch.zeros_like(logical_loads) + replica_counts.scatter_add_( + dim=1, + index=mapping, + src=torch.ones_like(mapping, dtype=logical_loads.dtype), + ) + loads_per_replica = torch.gather( + logical_loads / replica_counts.clamp_min(1), + dim=1, + index=mapping, + ) + loads_per_rank = loads_per_replica.reshape( + logical_loads.shape[0], ep_size, -1 + ).sum(dim=(0, 2)) + mean_load = loads_per_rank.mean() + if mean_load == 0: + return 1.0 + return (loads_per_rank.max() / mean_load).item() + + current_imbalance = rank_load_imbalance( + eplb_model_state.physical_to_logical_map ) + proposed_imbalance = rank_load_imbalance( + new_physical_to_logical_map + ) + relative_improvement = ( + current_imbalance - proposed_imbalance + ) / current_imbalance + skip_rearrange = relative_improvement < 0.05 + if skip_rearrange and is_main_rank: + logger.info( + "[EPLB] Skip rearrange: imbalance %.4f -> " + "%.4f (no material gain)", + current_imbalance, + proposed_imbalance, + ) + + if not skip_rearrange: + # Update expert weights + rearrange_expert_weights_inplace( + eplb_model_state.physical_to_logical_map, + new_physical_to_logical_map, + eplb_model_state.model.expert_weights, + eplb_model_state.expert_buffer, + ep_group, + eplb_model_state.communicator, + is_profile, + rank_mapping, + ) + + if not is_profile: + _commit_eplb_maps( + eplb_model_state, + new_physical_to_logical_map=new_physical_to_logical_map, + ) if is_main_rank: assert start_event is not None diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 8bf0123e57..6f46283829 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -233,17 +233,27 @@ class RoutedExperts(PluggableLayer): # Update local attributes from ExpertMapManager self.local_num_experts = self.expert_map_manager.local_num_experts self.expert_placement_strategy = self.expert_map_manager.placement_strategy - self.register_buffer("_expert_map", self.expert_map_manager.expert_map) - self.register_buffer("expert_mask", self.expert_map_manager.expert_mask) + self.register_buffer( + "_expert_map", self.expert_map_manager.expert_map, persistent=False + ) + self.register_buffer( + "expert_mask", self.expert_map_manager.expert_mask, persistent=False + ) # Get routing tables from ExpertMapManager routing_tables = self.expert_map_manager.routing_tables if routing_tables is not None: # Register routing tables as buffers for this layer global_to_physical, physical_to_global, local_global = routing_tables - self.register_buffer("expert_global_to_physical", global_to_physical) - self.register_buffer("expert_physical_to_global", physical_to_global) - self.register_buffer("expert_local_to_global", local_global) + self.register_buffer( + "expert_global_to_physical", global_to_physical, persistent=False + ) + self.register_buffer( + "expert_physical_to_global", physical_to_global, persistent=False + ) + self.register_buffer( + "expert_local_to_global", local_global, persistent=False + ) def _expert_routing_tables( self,