chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Hank Han <hanhan.hank@bytedance.com>
Co-authored-by: Hank Han <hanhan7630@outlook.com>
This commit is contained in:
Xun Sun
2026-03-16 18:40:58 +08:00
committed by GitHub
parent 3055b6906d
commit 549fbcc864
4 changed files with 110 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/elastic_ep/expert_backup_client.py
View File

@@ -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

47
python/sglang/srt/eplb/expert_location_updater.py
View File

@@ -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]

56
python/sglang/srt/model_executor/model_runner.py
View File

@@ -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,

31
python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py
View File

@@ -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,