From 549fbcc864c24bc7d98b6ac33981967b79a67c3c Mon Sep 17 00:00:00 2001 From: Xun Sun Date: Mon, 16 Mar 2026 18:40:58 +0800 Subject: [PATCH] [5/N] (Elastic EP) Use GPU P2P to exchange expert weights during EPLB as much as possible (#12068) Co-authored-by: Hank Han Co-authored-by: Hank Han --- .../srt/elastic_ep/expert_backup_client.py | 4 +- .../srt/eplb/expert_location_updater.py | 47 +++++++++++++++- .../sglang/srt/model_executor/model_runner.py | 56 +++++++++++-------- .../deepseek_common/deepseek_weight_loader.py | 31 +++++++++- 4 files changed, 110 insertions(+), 28 deletions(-) diff --git a/python/sglang/srt/elastic_ep/expert_backup_client.py b/python/sglang/srt/elastic_ep/expert_backup_client.py index 9bd61b6f5..1f4481348 100644 --- a/python/sglang/srt/elastic_ep/expert_backup_client.py +++ b/python/sglang/srt/elastic_ep/expert_backup_client.py @@ -98,7 +98,7 @@ class ExpertBackupClient: logger.warning("Register fails. Stop using expert weight backup!") break - def update_weights(self): + def update_weights(self, weight_name_filter=None): global_expert_location_metadata = get_global_expert_location_metadata() num_experts = ( self.model_config.hf_config.n_routed_experts @@ -111,6 +111,8 @@ class ExpertBackupClient: weight_size_list = [] for name, weight_info in self.dram_map_list[i].items(): + if weight_name_filter is not None and not weight_name_filter(name): + continue layer_id, expert_id, weight_name = extract_layer_and_expert_id(name) if layer_id >= self.model_config.hf_config.num_hidden_layers: continue diff --git a/python/sglang/srt/eplb/expert_location_updater.py b/python/sglang/srt/eplb/expert_location_updater.py index 286f1d0e3..d1d6387fc 100644 --- a/python/sglang/srt/eplb/expert_location_updater.py +++ b/python/sglang/srt/eplb/expert_location_updater.py @@ -20,6 +20,7 @@ import torch import torch.distributed from torch.distributed import P2POp +from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager from sglang.srt.eplb.expert_location import ( ExpertLocationMetadata, get_global_expert_location_metadata, @@ -45,6 +46,12 @@ class ExpertLocationUpdater: nnodes: int, rank: int, ): + """ + Update experts' physical location after EPLB. + + Returns a map of layer_id to expert_ids that are missing due to rank + failures during fault conditions when elastic EP is enabled. + """ if self._first_execution: self._first_execution = False torch.get_device_module().empty_cache() @@ -52,7 +59,7 @@ class ExpertLocationUpdater: old_expert_location_metadata = get_global_expert_location_metadata() assert old_expert_location_metadata is not None - _update_expert_weights( + missing_logical_experts_by_layers = _update_expert_weights( routed_experts_weights_of_layer=routed_experts_weights_of_layer, old_expert_location_metadata=old_expert_location_metadata, new_expert_location_metadata=new_expert_location_metadata, @@ -65,6 +72,8 @@ class ExpertLocationUpdater: update_layer_ids=update_layer_ids, ) + return missing_logical_experts_by_layers + def _update_expert_weights(**kwargs): if get_bool_env_var("SGLANG_EXPERT_LOCATION_UPDATER_CANARY"): @@ -101,7 +110,7 @@ def _update_expert_weights_with_canary( ) routed_experts_weights_of_layer[layer_id].append(canary_tensor) - _update_expert_weights_raw( + missing_logical_experts_by_layers = _update_expert_weights_raw( routed_experts_weights_of_layer=routed_experts_weights_of_layer, old_expert_location_metadata=old_expert_location_metadata, new_expert_location_metadata=new_expert_location_metadata, @@ -120,6 +129,8 @@ def _update_expert_weights_with_canary( f"{new_expert_location_metadata.physical_to_logical_map_cpu.tolist()=} " ) + return missing_logical_experts_by_layers + def _update_expert_weights_raw( routed_experts_weights_of_layer: Dict[int, List[torch.Tensor]], @@ -139,7 +150,10 @@ def _update_expert_weights_raw( num_local_physical_experts = old_expert_location_metadata.num_local_physical_experts num_gpu_per_node = world_size // nnodes + missing_logical_experts_by_layers: Dict[int, List[int]] = {} + for layer_id in update_layer_ids: + missing_logical_experts_info: List[int] = [] update_expert_weights_single_layer( routed_experts_weights=routed_experts_weights_of_layer[layer_id], temp_buffers=temp_buffers, @@ -153,8 +167,12 @@ def _update_expert_weights_raw( num_gpu_per_node=num_gpu_per_node, rank=rank, world_size=world_size, + missing_logical_experts_info=missing_logical_experts_info, log_metrics=log_metrics, ) + if len(missing_logical_experts_info) > 0: + missing_logical_experts_by_layers[layer_id] = missing_logical_experts_info + return missing_logical_experts_by_layers def create_temp_buffers(sample_tensors): @@ -170,6 +188,7 @@ def update_expert_weights_single_layer( num_gpu_per_node: int, rank: int, world_size: Optional[int] = None, + missing_logical_experts_info: Optional[List[int]] = None, debug: bool = False, log_metrics: bool = False, ): @@ -213,6 +232,7 @@ def update_expert_weights_single_layer( _handle_recv(buffer2weight_copy_infos, p2p_op_infos) _create_isend_ops(p2p_op_infos) + _filter_p2p_ops(p2p_op_infos) _execute_p2p_ops(p2p_op_infos) _execute_buffer2weight_copies(buffer2weight_copy_infos) @@ -434,6 +454,29 @@ def update_expert_weights_single_layer( return same_node_mapping, cross_node_mapping, need_comm_self_node_dst_ranks + def _filter_p2p_ops(p2p_op_infos): + elastic_ep_state = ElasticEPStateManager.instance() + if elastic_ep_state is not None and missing_logical_experts_info is not None: + # Filter out inactive P2P ops and record missing expert IDs in missing_logical_experts_info + is_active = elastic_ep_state.active_ranks_cpu + for i, (logical_expert_id, ops) in enumerate(p2p_op_infos): + has_isend = any(op.op == torch.distributed.isend for op in ops) + has_irecv = any(op.op == torch.distributed.irecv for op in ops) + assert not (has_isend and has_irecv), ( + "Each p2p_op_infos entry is expected to contain only send " + "or only recv ops." + ) + + if has_isend: + p2p_op_infos[i] = ( + logical_expert_id, + [op for op in ops if is_active[op.peer]], + ) + elif has_irecv: + if any(not is_active[op.peer] for op in ops): + missing_logical_experts_info.append(logical_expert_id) + p2p_op_infos[i] = (logical_expert_id, []) + def _execute_p2p_ops(p2p_op_infos): sorted_infos = sorted(p2p_op_infos, key=lambda info: info[0]) p2p_ops = [op for _, ops in sorted_infos for op in ops] diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 689afe18f..d85a7d1bf 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1100,34 +1100,42 @@ class ModelRunner(ModelRunnerKVCacheMixin): new_expert_location_metadata: ExpertLocationMetadata, update_layer_ids: List[int], ): - if ElasticEPStateManager.instance() is not None: - # TODO: refactor the weights update when elastic ep - old_expert_location_metadata = get_global_expert_location_metadata() - assert old_expert_location_metadata is not None - old_expert_location_metadata.update( - new_expert_location_metadata, - update_layer_ids=update_layer_ids, - ) + p2p_missing_logical_experts = self.expert_location_updater.update( + self.model.routed_experts_weights_of_layer, + new_expert_location_metadata, + update_layer_ids=update_layer_ids, + nnodes=self.server_args.nnodes, + rank=self.tp_rank, + ) + + if len(p2p_missing_logical_experts) > 0: + # Load the missing expert weights from disk + if callable(getattr(self.model, "generate_weight_name_filter", None)): + # Filter and load only missing expert weights + weight_name_filter = self.model.generate_weight_name_filter( + p2p_missing_logical_experts + ) + else: + # Do a full reload from disk/DRAM + logger.info( + "[Elastic EP] Model does not implement generate_weight_name_filter. " + "Performing full weight reload." + ) + weight_name_filter = None + if ( self.expert_backup_client is not None and self.expert_backup_client.use_backup ): - self.expert_backup_client.update_weights() - return - - self.update_weights_from_disk( - self.server_args.model_path, - self.server_args.load_format, - lambda name: "mlp.experts" in name and "mlp.shared_experts" not in name, - ) - else: - self.expert_location_updater.update( - self.model.routed_experts_weights_of_layer, - new_expert_location_metadata, - update_layer_ids=update_layer_ids, - nnodes=self.server_args.nnodes, - rank=self.tp_rank, - ) + # Load the missing weights from the DRAM backup + self.expert_backup_client.update_weights(weight_name_filter) + else: + # Load the missing weights from disk + self.update_weights_from_disk( + get_global_server_args().model_path, + get_global_server_args().load_format, + weight_name_filter=weight_name_filter, + ) def update_weights_from_disk( self, diff --git a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py index deda4f97e..c54ab358d 100644 --- a/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py +++ b/python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py @@ -15,7 +15,7 @@ import concurrent.futures import logging from dataclasses import dataclass -from typing import Iterable, List, Optional, Tuple +from typing import Dict, Iterable, List, Optional, Tuple import torch import torch.nn as nn @@ -609,6 +609,35 @@ class DeepseekV2WeightLoaderMixin: self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc.contiguous()) self_attn.use_deep_gemm_bmm = True + @classmethod + def generate_weight_name_filter(cls, logical_experts_map: Dict[int, List[int]]): + """ + Generates a filter function that tests whether the (layer_id, expert_id) + indicated by a param name lies in the `logical_experts` map + Args: + logical_experts_map: a map of layer_id to expert_ids, specifying a list of expert_ids by a specific layer_id. + + Returns: + A function (name: str) -> bool + """ + import re + + # Regex pattern to extract layer_id and expert_id from weight name + pattern = re.compile(r"layers\.(\d+)\.mlp\.experts\.(\d+)\.") + + def weight_name_filter(name: str) -> bool: + match = pattern.search(name) + if match: + layer_id, expert = int(match.group(1)), int(match.group(2)) + # First check if layer_id exists, then check if expert is in the list + return ( + layer_id in logical_experts_map + and expert in logical_experts_map[layer_id] + ) + return False + + return weight_name_filter + def _maybe_quant_weights_to_fp8_ue8m0( self, weights,