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

[1/2] Add ModelExpress coordination for remote instance weight loading - matching TP (#19920)

Browse Source 
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Ishan Dhanani <ishan@dhanani.dev>
This commit is contained in:
ishandhanani
2026-03-18 15:38:32 -05:00
committed by GitHub
parent 6cca5b9b97
commit 8f0f36c64b
6 changed files with 320 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/configs/load_config.py
View File

@@ -76,6 +76,8 @@ class LoadConfig:
remote_instance_weight_loader_send_weights_group_ports: Optional[List[int]] = None
remote_instance_weight_loader_backend: Optional[str] = None
remote_instance_weight_loader_transfer_engine: Optional[Any] = None
modelexpress_url: Optional[str] = None
modelexpress_model_name: Optional[str] = None
# ModelOpt-specific loading options
modelopt_checkpoint_restore_path: Optional[str] = None

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

@@ -676,6 +676,74 @@ class ModelRunner(ModelRunnerKVCacheMixin):
local_ip, self.remote_instance_transfer_engine.get_rpc_port()
).to_host_port_str()
def _publish_modelexpress_metadata(self):
"""Publish TransferEngine metadata to ModelExpress server (seed mode)."""
try:
from modelexpress import p2p_pb2
from modelexpress.client import MxClient
except ImportError as exc:
raise ImportError(
"ModelExpress support requires the 'modelexpress' package. "
"Install it with: pip install modelexpress"
) from exc
model_name = (
self.server_args.modelexpress_model_name or self.server_args.model_path
)
mx_url = self.server_args.modelexpress_url
session_id = self.remote_instance_transfer_engine_session_id
weight_info = self.remote_instance_transfer_engine_weight_info
if not session_id or weight_info is None:
logger.warning(
"ModelExpress source: skipping publish -- "
"TransferEngine not initialized or no weight info"
)
return
# Build tensor descriptors from weight_info dict
tensors = []
for name, (addr, numel, element_size) in weight_info.items():
tensors.append(
p2p_pb2.TensorDescriptor(
name=name,
addr=addr,
size=numel * element_size,
device_id=self.gpu_id,
)
)
worker = p2p_pb2.WorkerMetadata(
worker_rank=self.tp_rank,
transfer_engine_session_id=session_id,
tensors=tensors,
)
mx_client = MxClient(server_url=mx_url)
try:
logger.info(
"ModelExpress source: publishing metadata for model=%s, "
"tp_rank=%d, session=%s, %d tensors",
model_name,
self.tp_rank,
session_id,
len(tensors),
)
mx_client.publish_metadata(model_name, [worker])
mx_client.publish_ready(
model_name,
worker_id=self.tp_rank,
session_id=mx_client.session_id,
metadata_hash="",
)
logger.info(
"ModelExpress source: published ready for model=%s, tp_rank=%d",
model_name,
self.tp_rank,
)
finally:
mx_client.close()
def model_specific_adjustment(self):
server_args = self.server_args
@@ -963,6 +1031,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
remote_instance_weight_loader_send_weights_group_ports=self.server_args.remote_instance_weight_loader_send_weights_group_ports,
remote_instance_weight_loader_backend=self.server_args.remote_instance_weight_loader_backend,
remote_instance_weight_loader_transfer_engine=self.remote_instance_transfer_engine,
modelexpress_url=self.server_args.modelexpress_url,
modelexpress_model_name=self.server_args.modelexpress_model_name
or self.server_args.model_path,
modelopt_config=modelopt_config,
rl_quant_profile=self.server_args.rl_quant_profile,
draft_model_idx=self.draft_model_idx,
@@ -1015,6 +1086,25 @@ class ModelRunner(ModelRunnerKVCacheMixin):
)
monkey_patch_vllm_parallel_state(reverse=True)
# Publish metadata to ModelExpress if running as seed source
if self.server_args.modelexpress_source:
# Seed loads via DefaultModelLoader (load_format=auto), which doesn't
# call register_memory_region(). Do it here so weight_info is populated.
if (
self.remote_instance_transfer_engine_weight_info is None
and self.remote_instance_transfer_engine is not None
):
from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
register_memory_region,
)
self.remote_instance_transfer_engine_weight_info = (
register_memory_region(
self.model, self.remote_instance_transfer_engine
)
)
self._publish_modelexpress_metadata()
get_offloader().post_init()
# Register model for layerwise NVTX profiling if enabled

138
python/sglang/srt/model_loader/loader.py
View File

@@ -2151,6 +2151,15 @@ class RemoteInstanceModelLoader(BaseModelLoader):
raise RuntimeError(
"Failed to load weights from remote instance via transfer engine."
)
elif (
load_config.remote_instance_weight_loader_backend
== RemoteInstanceWeightLoaderBackend.MODELEXPRESS
):
self.load_model_from_modelexpress(
model,
load_config,
device_config,
)
else:
raise ValueError("Invalid remote instance weight loader backend.")
@@ -2264,6 +2273,135 @@ class RemoteInstanceModelLoader(BaseModelLoader):
return True
def load_model_from_modelexpress(
self,
model,
load_config: LoadConfig,
device_config: DeviceConfig,
):
"""Load weights via ModelExpress coordination + TransferEngine RDMA."""
try:
from modelexpress.client import MxClient
except ImportError as exc:
raise ImportError(
"ModelExpress support requires the 'modelexpress' package. "
"Install it with: pip install modelexpress"
) from exc
transfer_engine = load_config.remote_instance_weight_loader_transfer_engine
if transfer_engine is None:
raise RuntimeError(
"TransferEngine is not initialized for modelexpress backend."
)
tp_rank = load_config.tp_rank
model_name = load_config.modelexpress_model_name
logger.info(
"ModelExpress: registering memory regions for tp_rank=%d...", tp_rank
)
self.remote_instance_transfer_engine_weight_info = register_memory_region(
model, transfer_engine
)
# Wait for seed to be ready via ModelExpress
mx_client = MxClient(server_url=load_config.modelexpress_url)
try:
logger.info(
"ModelExpress: waiting for seed ready (model=%s)...",
model_name,
)
ready, session_id, metadata_hash = mx_client.wait_for_ready(
model_name,
worker_id=tp_rank,
)
if not ready:
raise RuntimeError(
f"ModelExpress: timed out waiting for seed ready "
f"(model={model_name}, worker={tp_rank})"
)
response = mx_client.get_metadata(model_name)
if not response.found:
raise RuntimeError(
f"ModelExpress: no metadata found for model={model_name}"
)
# Find the worker matching our tp_rank
source_worker = None
for w in response.workers:
if w.worker_rank == tp_rank:
source_worker = w
break
if source_worker is None:
raise RuntimeError(
f"ModelExpress: no worker metadata for rank={tp_rank}"
)
# Extract session_id from oneof backend_metadata
backend_field = source_worker.WhichOneof("backend_metadata")
if backend_field == "transfer_engine_session_id":
seed_session_id = source_worker.transfer_engine_session_id
else:
raise RuntimeError(
f"ModelExpress: expected transfer_engine_session_id, "
f"got backend_metadata={backend_field}"
)
# Build {name: (addr, size_bytes)} from seed tensor descriptors
seed_weight_info = {}
for td in source_worker.tensors:
seed_weight_info[td.name] = (td.addr, td.size)
logger.info(
"ModelExpress: got %d tensor descriptors from seed (session=%s)",
len(seed_weight_info),
seed_session_id,
)
finally:
mx_client.close()
# Transfer weights via TransferEngine RDMA
seed_ptr_list = []
client_ptr_list = []
client_len_list = []
for name, tensor in model.named_parameters():
weight_info = seed_weight_info.get(name, None)
if weight_info is None:
raise RuntimeError(
f"ModelExpress: cannot find weight info for {name} "
f"in seed metadata"
)
seed_ptr, seed_size = weight_info
local_size = tensor.numel() * tensor.element_size()
if seed_size != local_size:
raise RuntimeError(
f"ModelExpress: size mismatch for {name}: "
f"seed={seed_size} bytes, local={local_size} bytes"
)
seed_ptr_list.append(seed_ptr)
client_ptr_list.append(tensor.data_ptr())
client_len_list.append(local_size)
logger.info(
"ModelExpress: starting RDMA transfer of %d tensors...",
len(seed_ptr_list),
)
ret = transfer_engine.batch_transfer_sync_read(
seed_session_id,
client_ptr_list,
seed_ptr_list,
client_len_list,
)
if ret < 0:
raise RuntimeError(
f"ModelExpress: batch_transfer_sync_read failed, error={ret}"
)
if hasattr(model, "post_load_weights"):
model.post_load_weights()
logger.info("ModelExpress: weight transfer complete for tp_rank=%d", tp_rank)
class RemoteModelLoader(BaseModelLoader):
"""Model loader that can load Tensors from remote database."""

1
python/sglang/srt/model_loader/remote_instance_weight_loader_utils.py
View File

@@ -15,6 +15,7 @@ logger = logging.getLogger(__name__)
class RemoteInstanceWeightLoaderBackend(str, enum.Enum):
NCCL = "nccl"
TRANSFER_ENGINE = "transfer_engine"
MODELEXPRESS = "modelexpress"
def trigger_init_weights_send_group_for_remote_instance_request(

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

@@ -700,8 +700,11 @@ class ServerArgs:
remote_instance_weight_loader_seed_instance_ip: Optional[str] = None
remote_instance_weight_loader_seed_instance_service_port: Optional[int] = None
remote_instance_weight_loader_send_weights_group_ports: Optional[List[int]] = None
remote_instance_weight_loader_backend: Literal["transfer_engine", "nccl"] = "nccl"
remote_instance_weight_loader_backend: Literal[
"transfer_engine", "nccl", "modelexpress"
] = "nccl"
remote_instance_weight_loader_start_seed_via_transfer_engine: bool = False
modelexpress_config: Optional[str] = None
# For PD-Multiplexing
enable_pdmux: bool = False
@@ -2967,7 +2970,19 @@ class ServerArgs:
self.custom_weight_loader = []
if self.load_format == "remote_instance":
if (
if self.remote_instance_weight_loader_backend == "modelexpress":
# ModelExpress backend: requires url in --modelexpress-config
if self.modelexpress_url is None:
logger.warning(
"Fallback load_format to 'auto' due to missing 'url' in --modelexpress-config."
)
self.load_format = "auto"
elif not self.validate_transfer_engine():
logger.warning(
"Fallback load_format to 'auto' due to 'transfer_engine' (required by modelexpress) not being supported."
)
self.load_format = "auto"
elif (
self.remote_instance_weight_loader_seed_instance_ip is None
or self.remote_instance_weight_loader_seed_instance_service_port is None
):
@@ -5567,15 +5582,21 @@ class ServerArgs:
parser.add_argument(
"--remote-instance-weight-loader-backend",
type=str,
choices=["transfer_engine", "nccl"],
choices=["transfer_engine", "nccl", "modelexpress"],
default=ServerArgs.remote_instance_weight_loader_backend,
help="The backend for loading weights from remote instance. Can be 'transfer_engine' or 'nccl'. Default is 'nccl'.",
help="The backend for loading weights from remote instance. Can be 'transfer_engine', 'nccl', or 'modelexpress'. Default is 'nccl'.",
)
parser.add_argument(
"--remote-instance-weight-loader-start-seed-via-transfer-engine",
action="store_true",
help="Start seed server via transfer engine backend for remote instance weight loader.",
)
parser.add_argument(
"--modelexpress-config",
type=str,
default=ServerArgs.modelexpress_config,
help='JSON config for ModelExpress P2P weight loading. Keys: "url" (required, gRPC host:port), "model_name" (optional, defaults to --model-path), "source" (optional bool, true for seed mode). Example: \'{"url": "localhost:8001", "model_name": "my-model", "source": true}\'',
)
# For PD-Multiplexing
parser.add_argument(
@@ -6109,7 +6130,11 @@ class ServerArgs:
)
def validate_transfer_engine(self):
if importlib.util.find_spec("mooncake.engine") is None:
try:
mooncake_available = importlib.util.find_spec("mooncake.engine") is not None
except (ModuleNotFoundError, ValueError):
mooncake_available = False
if not mooncake_available:
logger.warning(
"Failed to import mooncake.engine. Does not support using TransferEngine as remote instance weight loader backend."
)
@@ -6122,14 +6147,44 @@ class ServerArgs:
else:
return True
@property
def _parsed_modelexpress_config(self) -> dict:
cache = getattr(self, "_mx_config_cache", None)
if cache is not None:
return cache
if self.modelexpress_config is None:
result = {}
elif isinstance(self.modelexpress_config, str):
result = json.loads(self.modelexpress_config)
else:
result = self.modelexpress_config
object.__setattr__(self, "_mx_config_cache", result)
return result
@property
def modelexpress_url(self) -> Optional[str]:
return self._parsed_modelexpress_config.get("url")
@property
def modelexpress_model_name(self) -> Optional[str]:
return self._parsed_modelexpress_config.get("model_name")
@property
def modelexpress_source(self) -> bool:
return self._parsed_modelexpress_config.get("source", False)
def remote_instance_weight_loader_use_transfer_engine(self):
# Use TransferEngine as seed backend.
if self.remote_instance_weight_loader_start_seed_via_transfer_engine:
return True
# ModelExpress source mode also needs TransferEngine init.
if self.modelexpress_source:
return True
# Use TransferEngine as client backend.
elif (
self.load_format == "remote_instance"
and self.remote_instance_weight_loader_backend == "transfer_engine"
and self.remote_instance_weight_loader_backend
in ("transfer_engine", "modelexpress")
):
return True
else: