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
docs/advanced_features/server_arguments.md
View File

@@ -331,6 +331,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
| `--deepep-config` | Tuned DeepEP config suitable for your own cluster. It can be either a string with JSON content or a file path. | `None` | Type: str |
| `--moe-dense-tp-size` | TP size for MoE dense MLP layers. This flag is useful when, with large TP size, there are errors caused by weights in MLP layers having dimension smaller than the min dimension GEMM supports. | `None` | Type: int |
| `--elastic-ep-backend` | Specify the collective communication backend for elastic EP. Currently supports 'mooncake'. | `none` | `none`, `mooncake` |
| `--enable-elastic-expert-backup` | Enable elastic EP backend to backup expert weights in DRAM feature. Currently supports 'mooncake'.| `False` | bool flag (set to enable) |
| `--mooncake-ib-device` | The InfiniBand devices for Mooncake Backend transfer, accepts multiple comma-separated devices (e.g., --mooncake-ib-device mlx5_0,mlx5_1). Default is None, which triggers automatic device detection when Mooncake Backend is enabled. | `None` | Type: str |
## Mamba Cache

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,

142
test/manual/ep/test_mooncake_expert_backup.py Normal file
View File

@@ -0,0 +1,142 @@
import time
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.server_fixtures.disaggregation_fixture import get_rdma_devices_args
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
CustomTestCase,
popen_launch_pd_server,
)
ib_devices = get_rdma_devices_args()
class TestBackup(CustomTestCase):
extra_args = []
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
cls.base_port = 20000
cls.num_processes = 2
# TODO (stage 100): in the future, implement a specified multiprocess launcher
cls.processes = [
popen_launch_pd_server(
cls.model,
f"http://127.0.0.1:{cls.base_port + i}",
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--enable-dp-attention",
"--dp",
"4",
"--elastic-ep-backend",
"mooncake",
"--mooncake-ib-device",
ib_devices,
"--moe-a2a-backend",
"mooncake",
"--deepep-mode",
"low_latency",
"--moe-dense-tp-size",
"1",
"--enable-dp-lm-head",
"--enable-two-batch-overlap",
"--disable-custom-all-reduce",
"--enable-elastic-expert-backup",
"--enable-eplb",
"--eplb-rebalance-num-iterations",
"50",
"--chunked-prefill-size",
"512",
"--cuda-graph-max-bs",
"128",
"--max-running-requests",
"512",
"--mem-fraction-static",
"0.5",
"--dist-init-addr",
"127.0.0.1:5000",
"--nnodes",
f"{cls.num_processes}",
"--node-rank",
f"{i}",
"--base-gpu-id",
f"{i * 2}",
],
)
for i in range(cls.num_processes)
]
server_ready = [False] * cls.num_processes
start_time = time.perf_counter()
with requests.Session() as session:
while (
time.perf_counter() - start_time < DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
and not all(server_ready)
):
for i, process in enumerate(cls.processes):
return_code = process.poll()
if return_code is not None:
# Server failed to start (non-zero exit code) or crashed
raise Exception(
f"Server process exited with code {return_code}. "
"Check server logs for errors."
)
try:
headers = {
"Content-Type": "application/json; charset=utf-8",
}
response = session.get(
f"http://127.0.0.1:{cls.base_port + i}/health_generate",
headers=headers,
)
if response.status_code == 200:
server_ready[i] = True
except requests.RequestException:
pass
return_code = process.poll()
if return_code is not None:
raise Exception(
f"Server unexpectedly exits ({return_code=}). Usually there will be error logs describing the cause far above this line."
)
time.sleep(10)
if not all(server_ready):
for process in cls.processes:
kill_process_tree(process.pid)
raise TimeoutError("Server failed to start within the timeout period.")
@classmethod
def tearDownClass(cls):
for process in cls.processes:
kill_process_tree(process.pid)
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host="http://127.0.0.1",
port=self.base_port,
)
metrics = run_eval_few_shot_gsm8k(args)
print(metrics)
self.assertGreater(metrics["accuracy"], 0.60)
if __name__ == "__main__":
unittest.main()