diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index d6b7eafc3..a36a97e50 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -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 diff --git a/python/sglang/srt/elastic_ep/elastic_ep.py b/python/sglang/srt/elastic_ep/elastic_ep.py index f3367980c..8f31fe4c7 100644 --- a/python/sglang/srt/elastic_ep/elastic_ep.py +++ b/python/sglang/srt/elastic_ep/elastic_ep.py @@ -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, diff --git a/python/sglang/srt/elastic_ep/expert_backup_client.py b/python/sglang/srt/elastic_ep/expert_backup_client.py new file mode 100644 index 000000000..51e9b2912 --- /dev/null +++ b/python/sglang/srt/elastic_ep/expert_backup_client.py @@ -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 diff --git a/python/sglang/srt/elastic_ep/expert_backup_manager.py b/python/sglang/srt/elastic_ep/expert_backup_manager.py new file mode 100644 index 000000000..47f26852f --- /dev/null +++ b/python/sglang/srt/elastic_ep/expert_backup_manager.py @@ -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 diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 9c460f559..8246c849e 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -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. diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 2e909d07c..57e279589 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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 diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 1252f9a81..d5a8338b9 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index b885729c4..8f1dc0db4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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() diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index d3ff835bb..4fe3be6d4 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c350bb661..a91e17d80 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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, diff --git a/test/manual/ep/test_mooncake_expert_backup.py b/test/manual/ep/test_mooncake_expert_backup.py new file mode 100644 index 000000000..fc3089d88 --- /dev/null +++ b/test/manual/ep/test_mooncake_expert_backup.py @@ -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()