From c9b4081016037ca493c8a99e3d1afcc60e6aaa5c Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Tue, 14 Jul 2026 15:54:40 +0800 Subject: [PATCH] Extract expert location updating into EPLBManager (#31149) --- python/sglang/srt/eplb/eplb_manager.py | 67 ++++++++++++++++++- .../sglang/srt/model_executor/model_runner.py | 49 +------------- 2 files changed, 66 insertions(+), 50 deletions(-) diff --git a/python/sglang/srt/eplb/eplb_manager.py b/python/sglang/srt/eplb/eplb_manager.py index 27e7e1bc6..2302984ee 100644 --- a/python/sglang/srt/eplb/eplb_manager.py +++ b/python/sglang/srt/eplb/eplb_manager.py @@ -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] diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 87652128e..a6b95349b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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.