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

[4/N] (Elastic EP) Back up Expert Weights in DRAM (#17374)

Browse Source 
Co-authored-by: UNIDY2002 <unidy2002@outlook.com>
This commit is contained in:
ympcMark
2026-02-27 15:59:13 +08:00
committed by GitHub
parent eef44ec916
commit 43fade5f69
11 changed files with 553 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/elastic_ep/elastic_ep.py
View File

@@ -56,7 +56,6 @@ class ElasticEPStateManager:
def _build_state(
cls, *, ep_size: Optional[int] = None, device: Optional[torch.device] = None
) -> ElasticEPState:
active = cls.healthy_rank_state(ep_size=ep_size, device=device)
return ElasticEPState(
active_ranks=active,

171
python/sglang/srt/elastic_ep/expert_backup_client.py Normal file
View File

@@ -0,0 +1,171 @@
import logging
import re
import threading
import time
import torch
import zmq
from sglang.srt.distributed.parallel_state import (
get_world_group,
get_world_size,
)
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
from sglang.srt.managers.io_struct import UpdateExpertBackupReq
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import get_local_ip_auto
PORT_BASE = envs.SGLANG_BACKUP_PORT_BASE.get()
logger = logging.getLogger(__name__)
def extract_layer_and_expert_id(param_name):
pattern = r"layers\.(\d+)\.mlp\.experts\.(\d+)\.(.+?)\."
match = re.search(pattern, param_name)
if match:
return int(match.group(1)), int(match.group(2)), match.group(3)
return -1, -1, ""
class ExpertBackupClient:
def __init__(self, server_args: ServerArgs, model_runner):
context = zmq.Context(2)
self.server_args = server_args
self.engine_num = server_args.nnodes
self.engine_rank = server_args.node_rank
self.recv_list = [None] * self.engine_num
self.ready_sockets = [None] * self.engine_num
self.model_runner = model_runner
self.moe_ep_size = model_runner.moe_ep_size
self.model_config = model_runner.model_config
self.moe_ep_rank = model_runner.moe_ep_rank
self.dram_map_list = [None] * self.engine_num
self.session_id_list = [None] * self.engine_num
self.transfer_engine = None
self.gpu_buffer = None
self.buffer_size = 0
self.use_backup = False
local_ip = get_local_ip_auto()
all_ips = [None] * get_world_size()
torch.distributed.all_gather_object(
all_ips, local_ip, group=get_world_group().cpu_group
)
logger.info(f"all_ips: {all_ips}")
for i in range(self.engine_num):
self.recv_list[i] = context.socket(zmq.SUB)
self.recv_list[i].connect(
f"tcp://{all_ips[i * get_world_size() // server_args.nnodes]}:{PORT_BASE + i * 2 + 1}"
)
self.recv_list[i].setsockopt(zmq.SUBSCRIBE, b"")
# Synchronization channel to notify the manager when this client is ready.
self.ready_sockets[i] = context.socket(zmq.PUSH)
self.ready_sockets[i].connect(
f"tcp://{all_ips[i * get_world_size() // server_args.nnodes]}:{PORT_BASE + i * 2}"
)
self.ready_sockets[i].send_pyobj(UpdateExpertBackupReq())
self._receive_thread = threading.Thread(target=self._receive_loop, daemon=True)
self._receive_thread.start()
def _receive_loop(self):
cnt = 0
while cnt < self.engine_num:
response = self.recv_list[cnt].recv_pyobj()
self.dram_map_list[response.rank] = response.weight_pointer_map
self.session_id_list[response.rank] = response.session_id
self.buffer_size = max(self.buffer_size, response.buffer_size)
cnt += 1
self.use_backup = True
self.start_transfer_client()
def start_transfer_client(self):
from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
self.transfer_engine = get_mooncake_transfer_engine()
self.params_dict = dict(self.model_runner.model.named_parameters())
for name, param in self.params_dict.items():
param_data = param.data
ret_value = self.transfer_engine.engine.register_memory(
param_data.data_ptr(), param_data.numel() * param_data.element_size()
)
if ret_value != 0:
self.use_backup = False
logger.warning("Register fails. Stop using expert weight backup!")
break
def update_weights(self):
global_expert_location_metadata = get_global_expert_location_metadata()
num_experts = (
self.model_config.hf_config.n_routed_experts
+ self.server_args.ep_num_redundant_experts
)
num_local_experts = num_experts // self.moe_ep_size
for i in range(self.engine_num):
server_ptr_list = []
local_ptr_list = []
weight_size_list = []
for name, weight_info in self.dram_map_list[i].items():
layer_id, expert_id, weight_name = extract_layer_and_expert_id(name)
if layer_id >= self.model_config.hf_config.num_hidden_layers:
continue
if weight_name == "gate_proj":
shard_id = "w1"
param_name = "experts.w13_"
elif weight_name == "down_proj":
shard_id = "w2"
param_name = "experts.w2_"
elif weight_name == "up_proj":
shard_id = "w3"
param_name = "experts.w13_"
else:
raise RuntimeError(f"Unknown weight name {weight_name}")
name = name.replace(f"experts.{expert_id}.{weight_name}.", param_name)
weight_param = self.params_dict[name]
physical_expert_ids = (
global_expert_location_metadata.logical_to_all_physical(
layer_id, expert_id
)
)
for physical_expert_id in physical_expert_ids:
if physical_expert_id not in range(
num_local_experts * self.moe_ep_rank,
num_local_experts * (self.moe_ep_rank + 1),
):
continue
param = weight_param[physical_expert_id % num_local_experts]
if shard_id == "w1":
param = param.narrow(0, 0, param.shape[0] // 2)
elif shard_id == "w3":
param = param.narrow(
0, param.shape[0] // 2, param.shape[0] // 2
)
server_ptr_list.append(weight_info["weight_ptr"])
local_ptr_list.append(param.data_ptr())
assert (
param.numel() * param.element_size() == weight_info["byte_size"]
)
weight_size_list.append(weight_info["byte_size"])
before_transfer = time.time()
ret = self.transfer_engine.engine.batch_transfer_sync_read(
self.session_id_list[i],
local_ptr_list,
server_ptr_list,
weight_size_list,
)
after_transfer = time.time()
logger.info(f"transfer time = {after_transfer - before_transfer} s")
if ret != 0:
raise RuntimeError(
f"Failed to read weights from backup, error code: {ret}"
)
return

186
python/sglang/srt/elastic_ep/expert_backup_manager.py Normal file
View File

@@ -0,0 +1,186 @@
import logging
import multiprocessing as mp
import re
import signal
import torch
import zmq
from sglang.srt.configs.load_config import LoadConfig
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.environ import envs
from sglang.srt.managers.io_struct import BackupDramReq
from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader
from sglang.srt.model_loader.utils import set_default_torch_dtype
from sglang.srt.server_args import (
PortArgs,
ServerArgs,
set_global_server_args_for_scheduler,
)
from sglang.srt.utils import get_local_ip_auto
PORT_BASE = envs.SGLANG_BACKUP_PORT_BASE.get()
logger = logging.getLogger(__name__)
def extract_expert_id(param_name):
pattern = r"\.experts\.(\d+)\."
match = re.search(pattern, param_name)
if match:
return int(match.group(1))
return -1
class ExpertBackupManager:
def __init__(self, server_args: ServerArgs, port_args: PortArgs):
self.load_format = server_args.load_format
self.model_config = ModelConfig.from_server_args(server_args)
self.continuous_buffer = None
self.weight_pointer_map = {}
self.transfer_engine = None
self.session_id = None
self.engine_num = server_args.nnodes
self.engine_rank = server_args.node_rank
self.expert_num = self.model_config.hf_config.n_routed_experts
self.idmn = (self.expert_num // self.engine_num) * self.engine_rank
self.idmx = (self.expert_num // self.engine_num) * (self.engine_rank + 1)
context = zmq.Context(2)
# Synchronization socket to avoid PUB/SUB slow joiner issues.
self.recv_from_expert_backup_client = context.socket(zmq.PULL)
self.recv_from_expert_backup_client.bind(
f"tcp://{get_local_ip_auto()}:{PORT_BASE + server_args.node_rank * 2}"
)
self.send_to_expert_backup_client = context.socket(zmq.PUB)
self.send_to_expert_backup_client.bind(
f"tcp://{get_local_ip_auto()}:{PORT_BASE + server_args.node_rank * 2 + 1}"
)
self.backup_weights_from_disk()
self.start_transfer_server()
# Block until all expert backup clients have reported readiness, to avoid
# losing the initial PUB message due to slow joiners.
num_ready_clients = 0
while num_ready_clients < server_args.tp_size:
self.recv_from_expert_backup_client.recv_pyobj()
num_ready_clients += 1
back_req = BackupDramReq(
rank=self.engine_rank,
weight_pointer_map=self.weight_pointer_map,
session_id=self.session_id,
buffer_size=self.continuous_buffer.numel()
* self.continuous_buffer.element_size(),
)
self.send_to_expert_backup_client.send_pyobj(back_req)
# Keep the manager subprocess alive until signals
signal.pause()
def backup_weights_from_disk(self):
load_config = LoadConfig(load_format=self.load_format)
loader = get_model_loader(load_config, self.model_config)
with set_default_torch_dtype(self.model_config.dtype):
iter = loader._get_weights_iterator(
DefaultModelLoader.Source.init_new(self.model_config, None)
)
total_bytes = 0
weight_info_dict = {}
for name, weight in iter:
expert_id = extract_expert_id(name)
if expert_id < self.idmx and expert_id >= self.idmn:
numel = weight.numel()
element_size = weight.element_size()
byte_size = numel * element_size
weight_info_dict[name] = {
"name": name,
"weight": weight,
"numel": numel,
"shape": weight.shape,
"dtype": weight.dtype,
"element_size": element_size,
"byte_size": byte_size,
}
total_bytes += byte_size
if total_bytes == 0:
self.continuous_buffer = None
self.weight_pointer_map = {}
return
self.continuous_buffer = torch.empty(
total_bytes, dtype=torch.uint8, device="cpu"
)
buffer_base_ptr = self.continuous_buffer.data_ptr()
self.weight_pointer_map = {}
current_byte_offset = 0
for name in sorted(weight_info_dict.keys()):
weight_info = weight_info_dict[name]
weight = weight_info["weight"]
byte_size = weight_info["byte_size"]
weight_flat = weight.flatten().contiguous()
weight_bytes = weight_flat.view(torch.uint8)
start_byte = current_byte_offset
end_byte = current_byte_offset + byte_size
weight_ptr = buffer_base_ptr + current_byte_offset
self.continuous_buffer[start_byte:end_byte].copy_(weight_bytes)
self.weight_pointer_map[name] = {
"name": name,
"weight_ptr": weight_ptr,
"shape": weight_info["shape"],
"numel": weight_info["numel"],
"dtype": weight_info["dtype"],
"element_size": weight_info["element_size"],
"byte_size": byte_size,
}
current_byte_offset = end_byte
def start_transfer_server(self):
from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
self.transfer_engine = get_mooncake_transfer_engine()
self.session_id = self.transfer_engine.session_id
server_ptr = self.continuous_buffer.data_ptr()
server_len = (
self.continuous_buffer.numel() * self.continuous_buffer.element_size()
)
ret_value = self.transfer_engine.engine.register_memory(server_ptr, server_len)
if ret_value != 0:
raise RuntimeError("Mooncake memory registration failed.")
def run_expert_backup_manager_process(
server_args: ServerArgs,
port_args: PortArgs,
):
set_global_server_args_for_scheduler(server_args)
from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
init_mooncake_transfer_engine,
)
init_mooncake_transfer_engine(
hostname=get_local_ip_auto(),
gpu_id=0,
ib_device=(
server_args.disaggregation_ib_device or server_args.mooncake_ib_device
),
)
manager = ExpertBackupManager(server_args, port_args)
def run_expert_backup_manager(
server_args: ServerArgs,
port_args: PortArgs,
):
proc = mp.Process(
target=run_expert_backup_manager_process,
args=(server_args, port_args),
)
proc.start()
return proc

7
python/sglang/srt/entrypoints/engine.py
View File

@@ -36,6 +36,7 @@ import torch
import uvloop
import zmq
from sglang.srt.elastic_ep.expert_backup_manager import run_expert_backup_manager
from sglang.srt.entrypoints.EngineBase import EngineBase
from sglang.srt.managers.data_parallel_controller import (
run_data_parallel_controller_process,
@@ -1068,6 +1069,12 @@ def _launch_subprocesses(
run_scheduler_process_func=run_scheduler_process_func,
)
if (
server_args.enable_elastic_expert_backup
and server_args.elastic_ep_backend is not None
):
run_expert_backup_manager(server_args, port_args)
if server_args.node_rank >= 1:
# In multi-node cases, non-zero rank nodes do not need to run tokenizer or detokenizer,
# so they can just wait here.

3
python/sglang/srt/environ.py
View File

@@ -489,6 +489,9 @@ class Envs:
SGLANG_ENCODER_RECV_TIMEOUT = EnvFloat(180.0)
SGLANG_ENCODER_SEND_TIMEOUT = EnvFloat(180.0)
# Elastic EP Backup Port
SGLANG_BACKUP_PORT_BASE = EnvInt(10000)
envs = Envs()
EnvField._allow_set_name = False

13
python/sglang/srt/managers/io_struct.py
View File

@@ -1401,6 +1401,19 @@ class SendWeightsToRemoteInstanceReqOutput(BaseReq):
message: str
@dataclass
class UpdateExpertBackupReq(BaseReq):
pass
@dataclass
class BackupDramReq(BaseReq):
rank: int
weight_pointer_map: Dict[str, Any]
session_id: str
buffer_size: int
@dataclass
class InitWeightsUpdateGroupReqInput(BaseReq):
# The master address

2
python/sglang/srt/managers/scheduler.py
View File

@@ -2403,7 +2403,7 @@ class Scheduler(
if (
self.server_args.enable_dp_attention
and self.server_args.elastic_ep_backend == "mooncake"
and self.server_args.elastic_ep_backend is not None
):
# Get the tensors indicating rank activeness
tp_active_ranks = self.tp_group.active_ranks.detach().cpu().numpy()

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

@@ -70,6 +70,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
)
from sglang.srt.distributed.parallel_state import monkey_patch_vllm_parallel_state
from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
from sglang.srt.elastic_ep.expert_backup_client import ExpertBackupClient
from sglang.srt.environ import envs
from sglang.srt.eplb.eplb_manager import EPLBManager
from sglang.srt.eplb.expert_distribution import (
@@ -494,6 +495,16 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.sampler = create_sampler()
self.load_model()
# Load the expert backup client
self.expert_backup_client = (
ExpertBackupClient(self.server_args, self)
if (
self.server_args.enable_elastic_expert_backup
and self.server_args.elastic_ep_backend is not None
)
else None
)
if (
self.server_args.remote_instance_weight_loader_use_transfer_engine()
and self.remote_instance_transfer_engine is not None
@@ -866,6 +877,10 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.server_args.language_only
and self.server_args.encoder_transfer_backend == "mooncake"
)
or (
self.server_args.enable_elastic_expert_backup
and self.server_args.elastic_ep_backend is not None
)
)
if use_mooncake_te:
@@ -1087,6 +1102,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
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,

7
python/sglang/srt/server_args.py
View File

@@ -517,6 +517,7 @@ class ServerArgs:
deepep_config: Optional[str] = None
moe_dense_tp_size: Optional[int] = None
elastic_ep_backend: Literal[None, "mooncake"] = None
enable_elastic_expert_backup: bool = False
mooncake_ib_device: Optional[str] = None
# Mamba cache
@@ -4255,6 +4256,12 @@ class ServerArgs:
choices=["none", "mooncake"],
help="Specify the collective communication backend for elastic EP. Currently supports 'mooncake'.",
)
parser.add_argument(
"--enable-elastic-expert-backup",
action="store_true",
default=ServerArgs.enable_elastic_expert_backup,
help="Enable elastic expert backup feature.",
)
parser.add_argument(
"--mooncake-ib-device",
type=str,