minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Extract expert location updating into EPLBManager (#31149)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-14 15:54:40 +08:00
committed by GitHub
parent 6507d4a090
commit c9b4081016
2 changed files with 66 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/eplb/eplb_manager.py
+65 -2
View File
@@ -3,6 +3,7 @@ import time
from typing import TYPE_CHECKING, List
import torch.cuda
from torch import nn
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -12,6 +13,8 @@ from sglang.srt.eplb.expert_location import (
format_expert_location_layout_diff,
get_global_expert_location_metadata,
)
from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater
from sglang.srt.runtime_context import get_server_args
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
@@ -94,9 +97,17 @@ class EPLBManager:
for chunk_layer_ids in update_layer_ids_chunks:
if len(update_layer_ids_chunks) > 1:
yield
self._model_runner.update_expert_location(
expert_location_metadata,
update_expert_location_with_recovery(
expert_location_updater=self._model_runner.expert_location_updater,
model=self._model_runner.model,
new_expert_location_metadata=expert_location_metadata,
update_layer_ids=chunk_layer_ids,
nnodes=self._model_runner.server_args.nnodes,
tp_rank=self._model_runner.tp_rank,
expert_backup_client=self._model_runner.expert_backup_client,
update_weights_from_disk_callable=self._model_runner.weight_updater.update_weights_from_disk,
ep_dispatch_algorithm=self._model_runner.server_args.ep_dispatch_algorithm,
init_lplb_solvers_callable=self._model_runner._init_lplb_solvers,
)
self._log_rebalance_layout_after_update(update_layer_ids=all_update_layer_ids)
@@ -181,6 +192,58 @@ class EPLBManager:
)
def update_expert_location_with_recovery(
*,
expert_location_updater: ExpertLocationUpdater,
model: nn.Module,
new_expert_location_metadata: ExpertLocationMetadata,
update_layer_ids: List[int],
nnodes: int,
tp_rank: int,
expert_backup_client,
update_weights_from_disk_callable,
ep_dispatch_algorithm: str,
init_lplb_solvers_callable,
):
p2p_missing_logical_experts = expert_location_updater.update(
model.routed_experts_weights_of_layer,
new_expert_location_metadata,
update_layer_ids=update_layer_ids,
nnodes=nnodes,
rank=tp_rank,
)
if len(p2p_missing_logical_experts) > 0:
# Load the missing expert weights from disk
if callable(getattr(model, "generate_weight_name_filter", None)):
# Filter and load only missing expert weights
weight_name_filter = 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 expert_backup_client is not None and expert_backup_client.use_backup:
# Load the missing weights from the DRAM backup
expert_backup_client.update_weights(weight_name_filter)
else:
# Load the missing weights from disk
update_weights_from_disk_callable(
get_server_args().model_path,
get_server_args().load_format,
weight_name_filter=weight_name_filter,
)
# Re-init LPLB solvers after expert location update
if ep_dispatch_algorithm == "lp":
init_lplb_solvers_callable()
def _chunk_list(items: List, chunk_size):
for start_index in range(0, len(items), chunk_size):
yield items[start_index : start_index + chunk_size]
python/sglang/srt/model_executor/model_runner.py
+1 -48
View File
@@ -25,7 +25,7 @@ import threading
import time
from collections import defaultdict
from dataclasses import dataclass
from typing import Any, List, Optional, Union
from typing import Any, Optional, Union
import torch
import torch.distributed as dist
@@ -96,7 +96,6 @@ from sglang.srt.eplb.expert_distribution import (
set_global_expert_distribution_recorder,
)
from sglang.srt.eplb.expert_location import (
ExpertLocationMetadata,
broadcast_global_expert_location_metadata,
compute_initial_expert_location_metadata,
format_expert_location_layout,
@@ -1649,52 +1648,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
set_global_lplb_solver(lid, solver)
logger.info(f"Initialized LPLB solvers for {metadata.num_layers} layers")
def update_expert_location(
self,
new_expert_location_metadata: ExpertLocationMetadata,
update_layer_ids: List[int],
):
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
):
# 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.weight_updater.update_weights_from_disk(
get_server_args().model_path,
get_server_args().load_format,
weight_name_filter=weight_name_filter,
)
# Re-init LPLB solvers after expert location update
if self.server_args.ep_dispatch_algorithm == "lp":
self._init_lplb_solvers()
def maybe_recover_ep_ranks(self):
# TODO(perf): `active_ranks.all()` on a CUDA tensor triggers host-device
# synchronization, and this function is on the forward-path.