[AMD][Bugfix][EPLB] Fix elastic EP scaling accuracy on ROCm (#47206)
Signed-off-by: okorzh <okorzh@amd.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
This commit is contained in:
Oxana Korzh
co-authored by
Claude Opus 4.8
Andreas Karatzas
Claude Opus 4.8
Andreas Karatzas
parent
5d8e90a966
commit
2279575cd9
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user