[5/N] (Elastic EP) Use GPU P2P to exchange expert weights during EPLB as much as possible (#12068)

Co-authored-by: Hank Han <hanhan.hank@bytedance.com>
Co-authored-by: Hank Han <hanhan7630@outlook.com>

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

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]

python/sglang/srt/model_executor/model_runner.py

@@ -1100,28 +1100,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         new_expert_location_metadata: ExpertLocationMetadata,
         update_layer_ids: List[int],
     ):
-        if ElasticEPStateManager.instance() is not None:
+        p2p_missing_logical_experts = self.expert_location_updater.update(
-            # 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,
-            )
-            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,
@@ -1129,6 +1108,35 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             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
+            ):
+                # 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,
         model_path: str,

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,