minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Delegate ModelExpress loading to package (#24723)

Browse Source 
Signed-off-by: Zheng Luo <zheluo@nvidia.com>
This commit is contained in:
Zheng Luo
2026-05-16 11:16:44 -07:00
committed by GitHub
parent 0be539024f
commit 435ea41cf0
7 changed files with 60 additions and 503 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/docs/advanced_features/rfork.mdx
+5 -11
View File
@@ -47,7 +47,7 @@ To learn more details about R-Fork, please check **<a href="https://lmsys.org/bl
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>modelexpress-config</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>JSON config for <code>modelexpress</code> backend. Keys: <code>"url"</code> (required, gRPC host:port of ModelExpress server), <code>"model_name"</code> (optional, defaults to <code>--model-path</code>), <code>"source"</code> (optional bool, <code>true</code> for seed mode).</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>JSON config for <code>modelexpress</code> backend. Keys: <code>"url"</code> (optional gRPC host:port override) and <code>"transport"</code> (<code>"nixl"</code> or <code>"transfer_engine"</code>, defaults to <code>"nixl"</code>).</td>
</tr>
</tbody>
</table>
@@ -87,22 +87,16 @@ python -m sglang.launch_server [args] \
### ModelExpress as backend
[ModelExpress](https://github.com/ai-dynamo/modelexpress) is a coordination service that manages P2P weight transfer metadata. It removes the need for direct seed IP/port configuration by providing a centralized registry that seeds publish to and clients discover from. Under the hood it uses TransferEngine (Mooncake) for the actual RDMA data transfer.
[ModelExpress](https://github.com/ai-dynamo/modelexpress) is a coordination service that manages P2P weight transfer metadata. It removes the need for direct seed IP/port configuration by providing a centralized registry that instances publish to and discover from. The ModelExpress Python package must be installed in the SGLang image.
A running ModelExpress server is required. See the [ModelExpress documentation](https://github.com/ai-dynamo/modelexpress) for setup instructions.
seed instance:
```bash Command
python -m sglang.launch_server [args] \
--modelexpress-config '{"url": "[modelexpress_grpc_host:port]", "model_name": "[model_name]", "source": true}'
```
client instance:
server instance:
```bash Command
python -m sglang.launch_server [args] \
--load-format remote_instance \
--remote-instance-weight-loader-backend modelexpress \
--modelexpress-config '{"url": "[modelexpress_grpc_host:port]", "model_name": "[model_name]"}'
--modelexpress-config '{"url": "[modelexpress_grpc_host:port]", "transport": "nixl"}'
```
The seed publishes its TransferEngine session ID and tensor layout to ModelExpress. The client queries ModelExpress to discover the seed, then pulls weights directly via RDMA. This enables dynamic seed discovery without hardcoding IPs, and supports multiple models through a single ModelExpress instance.
All SGLang instances use the same command shape. If no ready source exists, the instance loads weights natively and publishes metadata to ModelExpress. If a compatible source exists, it loads weights through ModelExpress P2P transfer. Set <code>"transport": "transfer_engine"</code> to use Mooncake TransferEngine instead of the default NIXL transport.
docs_new/docs/advanced_features/server_arguments.mdx
+2 -2
View File
@@ -2547,9 +2547,9 @@ Please consult the documentation below and [server_args.py](https://github.com/s
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--remote-instance-weight-loader-backend</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>The backend for loading weights from remote instance. Can be 'transfer_engine' or 'nccl'. Default is 'nccl'.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>The backend for loading weights from remote instance. Can be 'transfer_engine', 'nccl', or 'modelexpress'. Default is 'nccl'.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>nccl</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>transfer_engine</code>, <code>nccl</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>transfer_engine</code>, <code>nccl</code>, <code>modelexpress</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>--remote-instance-weight-loader-start-seed-via-transfer-engine</code></td>
python/sglang/srt/configs/load_config.py
+2 -8
View File
@@ -79,15 +79,9 @@ 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
remote_instance_weight_loader_transfer_engine_session_id: Optional[str] = None
modelexpress_url: Optional[str] = None
modelexpress_model_name: Optional[str] = None
# Fields for building SourceIdentity (needed by both seed and client)
modelexpress_tp_size: Optional[int] = None
modelexpress_pp_size: Optional[int] = None
modelexpress_ep_size: Optional[int] = None
modelexpress_dtype: Optional[str] = None
modelexpress_quantization: Optional[str] = None
modelexpress_transport: str = "transfer_engine"
modelexpress_transport: str = "nixl"
# ModelOpt-specific loading options
modelopt_checkpoint_restore_path: Optional[str] = None
python/sglang/srt/model_executor/model_runner.py
+6 -176
View File
@@ -25,7 +25,6 @@ import os
import socket
import threading
import time
import uuid
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
@@ -655,6 +654,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
if (
self.server_args.remote_instance_weight_loader_use_transfer_engine()
# ModelExpress owns TransferEngine memory registration and metadata
# publishing for backend=modelexpress. Re-registering here would
# overlap the same weight buffers.
and self.server_args.remote_instance_weight_loader_backend
!= RemoteInstanceWeightLoaderBackend.MODELEXPRESS
and self.remote_instance_transfer_engine is not None
and self.remote_instance_transfer_engine_weight_info is None
):
@@ -963,155 +967,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
f"Failed to register transfer engine info for tp_rank={self.tp_rank}: {e}"
)
def _publish_modelexpress_metadata(self):
"""Publish metadata to ModelExpress server (seed mode).
Supports two transport backends:
- transfer_engine: publishes TransferEngine session_id (Mooncake)
- nixl: creates NIXL agent, registers tensors, publishes nixl_metadata
"""
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
transport = self.server_args.modelexpress_transport
# Build SourceIdentity for this instance
identity = p2p_pb2.SourceIdentity(
model_name=model_name,
backend_framework=p2p_pb2.BACKEND_FRAMEWORK_SGLANG,
tensor_parallel_size=self.server_args.tp_size,
pipeline_parallel_size=self.server_args.pp_size,
expert_parallel_size=self.server_args.ep_size,
dtype=self.server_args.dtype or "",
quantization=self.server_args.quantization or "",
)
if transport == "nixl":
worker, tensor_count = self._build_nixl_worker_metadata(p2p_pb2)
else:
worker, tensor_count = self._build_transfer_engine_worker_metadata(p2p_pb2)
if worker is None:
return
# Generate a unique worker_id for this running instance
worker_id = str(uuid.uuid4())
mx_client = MxClient(server_url=mx_url)
try:
logger.info(
"ModelExpress source [%s]: publishing metadata for model=%s, "
"tp_rank=%d, %d tensors, worker_id=%s",
transport,
model_name,
self.tp_rank,
tensor_count,
worker_id,
)
mx_source_id = mx_client.publish_metadata(identity, worker, worker_id)
mx_client.update_status(
mx_source_id=mx_source_id,
worker_id=worker_id,
worker_rank=self.tp_rank,
status=p2p_pb2.SOURCE_STATUS_READY,
)
logger.info(
"ModelExpress source: published ready for model=%s, "
"tp_rank=%d, mx_source_id=%s",
model_name,
self.tp_rank,
mx_source_id,
)
finally:
mx_client.close()
def _build_transfer_engine_worker_metadata(self, p2p_pb2):
"""Build WorkerMetadata using TransferEngine session_id."""
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 None, 0
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,
)
return worker, len(tensors)
def _build_nixl_worker_metadata(self, p2p_pb2):
"""Build WorkerMetadata using NIXL agent for RDMA transfers."""
from modelexpress.nixl_transfer import NixlTransferManager
agent_name = f"sglang-seed-rank{self.tp_rank}-{uuid.uuid4().hex[:8]}"
nixl_mgr = NixlTransferManager(agent_name, self.gpu_id)
nixl_mgr.initialize()
# Collect model tensors for NIXL registration
model_tensors = {}
for name, param in self.model.named_parameters():
t = param.data
if t.is_contiguous():
model_tensors[name] = t
else:
# Non-contiguous tensors: register underlying storage as byte view
sv = torch.empty(0, dtype=torch.uint8, device=t.device).set_(
t.untyped_storage()
)
if sv.data_ptr() not in {v.data_ptr() for v in model_tensors.values()}:
model_tensors[f"{name}.__storage"] = sv
nixl_metadata = nixl_mgr.register_tensors(model_tensors)
# Build tensor descriptors from registered tensors
tensors = []
for td in nixl_mgr.tensor_descriptors:
tensors.append(
p2p_pb2.TensorDescriptor(
name=td.name,
addr=td.addr,
size=td.size,
device_id=td.device_id,
dtype=td.dtype,
)
)
worker = p2p_pb2.WorkerMetadata(
worker_rank=self.tp_rank,
nixl_metadata=nixl_metadata,
tensors=tensors,
)
# Keep reference alive so NIXL agent isn't garbage collected
self._nixl_manager = nixl_mgr
return worker, len(tensors)
def model_specific_adjustment(self):
server_args = self.server_args
@@ -1399,14 +1254,8 @@ 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,
remote_instance_weight_loader_transfer_engine_session_id=self.remote_instance_transfer_engine_session_id,
modelexpress_url=self.server_args.modelexpress_url,
modelexpress_model_name=self.server_args.modelexpress_model_name
or self.server_args.model_path,
modelexpress_tp_size=self.server_args.tp_size,
modelexpress_pp_size=self.server_args.pp_size,
modelexpress_ep_size=self.server_args.ep_size,
modelexpress_dtype=self.server_args.dtype,
modelexpress_quantization=self.server_args.quantization or "",
modelexpress_transport=self.server_args.modelexpress_transport,
modelopt_config=modelopt_config,
rl_quant_profile=self.server_args.rl_quant_profile,
@@ -1464,25 +1313,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
torch.npu.empty_cache()
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()
if not self.is_draft_worker:
get_offloader().post_init()
python/sglang/srt/model_loader/loader.py
+17 -271
View File
@@ -2200,10 +2200,18 @@ class RemoteInstanceModelLoader(BaseModelLoader):
load_config.remote_instance_weight_loader_backend
== RemoteInstanceWeightLoaderBackend.MODELEXPRESS
):
self.load_model_from_modelexpress(
model,
load_config,
device_config,
try:
from modelexpress.engines.sglang.loader import MxModelLoader
except ImportError as exc:
raise ImportError(
"ModelExpress support requires the 'modelexpress' "
"package. Install it in the SGLang image."
) from exc
model = MxModelLoader(load_config).load_model(
model=model,
model_config=model_config,
device_config=device_config,
)
else:
raise ValueError("Invalid remote instance weight loader backend.")
@@ -2316,267 +2324,6 @@ class RemoteInstanceModelLoader(BaseModelLoader):
return True
def load_model_from_modelexpress(
self,
model,
load_config: LoadConfig,
device_config: DeviceConfig,
):
"""Load weights via ModelExpress coordination + RDMA transfer.
Supports two transport backends:
- transfer_engine: Mooncake TransferEngine (default)
- nixl: NIXL UCX-based RDMA
"""
try:
import grpc
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
tp_rank = load_config.tp_rank
model_name = load_config.modelexpress_model_name
transport = load_config.modelexpress_transport
# Process quantized weights to establish final tensor layout
target_device = torch.device(device_config.device)
for _, module in model.named_modules():
quant_method = getattr(module, "quant_method", None)
if quant_method is not None:
with device_loading_context(module, target_device):
quant_method.process_weights_after_loading(module)
# Register local memory for the chosen transport
if transport == "nixl":
nixl_mgr = self._init_nixl_for_target(model, load_config, device_config)
else:
transfer_engine = load_config.remote_instance_weight_loader_transfer_engine
if transfer_engine is None:
raise RuntimeError(
"TransferEngine is not initialized for modelexpress backend."
)
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
)
# --- Shared MX discovery logic ---
identity = p2p_pb2.SourceIdentity(
model_name=model_name,
backend_framework=p2p_pb2.BACKEND_FRAMEWORK_SGLANG,
tensor_parallel_size=load_config.modelexpress_tp_size or 1,
pipeline_parallel_size=load_config.modelexpress_pp_size or 1,
expert_parallel_size=load_config.modelexpress_ep_size or 1,
dtype=load_config.modelexpress_dtype or "",
quantization=load_config.modelexpress_quantization or "",
)
mx_client = MxClient(server_url=load_config.modelexpress_url)
try:
logger.info(
"ModelExpress [%s]: looking for seed (model=%s, rank=%d)...",
transport,
model_name,
tp_rank,
)
try:
resp = mx_client.list_sources(
identity=identity,
status_filter=p2p_pb2.SOURCE_STATUS_READY,
)
except grpc.RpcError as e:
raise RuntimeError(
f"ModelExpress: cannot reach server at "
f"{load_config.modelexpress_url}: "
f"{e.code()}: {e.details()}"
) from e
source_ref = None
for inst in resp.instances:
if inst.worker_rank == tp_rank:
source_ref = inst
break
if source_ref is None:
raise RuntimeError(
f"ModelExpress: no READY source found for "
f"model={model_name}, rank={tp_rank}. "
f"Ensure the seed instance is running and has published metadata."
)
response = mx_client.get_metadata(
mx_source_id=source_ref.mx_source_id,
worker_id=source_ref.worker_id,
)
if not response.found:
raise RuntimeError(
f"ModelExpress: no metadata found for "
f"source_id={source_ref.mx_source_id}, "
f"worker_id={source_ref.worker_id}"
)
source_worker = response.worker
finally:
mx_client.close()
# --- Transport-specific transfer ---
if transport == "nixl":
self._transfer_via_nixl(model, nixl_mgr, source_worker, tp_rank)
else:
self._transfer_via_transfer_engine(
model, transfer_engine, source_worker, tp_rank
)
_post_load_weights(model)
logger.info("ModelExpress: weight transfer complete for tp_rank=%d", tp_rank)
def _transfer_via_transfer_engine(
self, model, transfer_engine, source_worker, tp_rank
):
"""Execute weight transfer using Mooncake TransferEngine."""
backend_field = source_worker.WhichOneof("backend_metadata")
if backend_field != "transfer_engine_session_id":
raise RuntimeError(
f"ModelExpress: expected transfer_engine_session_id, "
f"got backend_metadata={backend_field}"
)
seed_session_id = source_worker.transfer_engine_session_id
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,
)
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 TransferEngine RDMA 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}"
)
def _init_nixl_for_target(self, model, load_config, device_config):
"""Initialize NIXL agent and register local tensors for the target."""
import uuid
from modelexpress.nixl_transfer import NixlTransferManager
tp_rank = load_config.tp_rank
device_id = device_config.gpu_id
agent_name = f"sglang-target-rank{tp_rank}-{uuid.uuid4().hex[:8]}"
nixl_mgr = NixlTransferManager(agent_name, device_id)
nixl_mgr.initialize()
# Collect local tensors, handling non-contiguous via storage views
local_tensors = {}
seen_ptrs = set()
for name, param in model.named_parameters():
t = param.data
if t.is_contiguous():
ptr = t.data_ptr()
if ptr in seen_ptrs:
continue
seen_ptrs.add(ptr)
local_tensors[name] = t
else:
sv = torch.empty(0, dtype=torch.uint8, device=t.device).set_(
t.untyped_storage()
)
ptr = sv.data_ptr()
if ptr in seen_ptrs:
continue
seen_ptrs.add(ptr)
local_tensors[f"{name}.__storage"] = sv
nixl_mgr.register_tensors(local_tensors)
logger.info(
"ModelExpress [nixl]: registered %d tensors for tp_rank=%d",
len(local_tensors),
tp_rank,
)
return nixl_mgr
def _transfer_via_nixl(self, model, nixl_mgr, source_worker, tp_rank):
"""Execute weight transfer using NIXL RDMA."""
from modelexpress.types import TensorDescriptor
backend_field = source_worker.WhichOneof("backend_metadata")
if backend_field != "nixl_metadata":
raise RuntimeError(
f"ModelExpress: expected nixl_metadata, "
f"got backend_metadata={backend_field}"
)
source_tensors = [
TensorDescriptor(
name=td.name,
addr=td.addr,
size=td.size,
device_id=td.device_id,
dtype=td.dtype,
)
for td in source_worker.tensors
]
logger.info(
"ModelExpress [nixl]: starting RDMA transfer of %d tensors...",
len(source_tensors),
)
total_bytes, matched, duration = nixl_mgr.receive_from_source(
source_metadata=source_worker.nixl_metadata,
source_tensors=source_tensors,
coalesce_transfers=False,
)
logger.info(
"ModelExpress [nixl]: transferred %d tensors, " "%.2f GB in %.2fs",
matched,
total_bytes / 1e9,
duration,
)
class RemoteModelLoader(BaseModelLoader):
"""Model loader that can load Tensors from remote database."""
@@ -3281,12 +3028,11 @@ def get_model_loader(
return DummyModelLoader(load_config)
# ModelOptModelLoader's local-copy quantize-and-export workflow doesn't apply
# to RUNAI_STREAMER, which streams weights directly from object storage.
# RUNAI_STREAMER loads always fall through to the unconditional branch at
# the bottom of this function. This also avoids calling _is_already_quantized()
# on RunAI streamer cache paths, where huggingface_hub raises HFValidationError.
model_optloader_allowed = (
model_config and load_config.load_format != LoadFormat.RUNAI_STREAMER
# to non-local loaders. These loaders own their weight transport path and still
# initialize the model with ModelOpt quantization config where applicable.
model_optloader_allowed = model_config and load_config.load_format not in (
LoadFormat.RUNAI_STREAMER,
LoadFormat.REMOTE_INSTANCE,
)
if model_optloader_allowed and (
python/sglang/srt/server_args.py
+10 -35
View File
@@ -3539,19 +3539,7 @@ class ServerArgs:
self.custom_weight_loader = []
if self.load_format == "remote_instance":
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 (
if self.remote_instance_weight_loader_backend != "modelexpress" and (
self.remote_instance_weight_loader_seed_instance_ip is None
or self.remote_instance_weight_loader_seed_instance_service_port is None
):
@@ -6623,7 +6611,7 @@ class ServerArgs:
"--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}\'',
help='JSON config for ModelExpress P2P weight loading. Keys: "url" (optional gRPC host:port override), "transport" ("nixl" or "transfer_engine"). Example: \'{"url": "localhost:8001", "transport": "nixl"}\'',
)
# For PD-Multiplexing
@@ -7267,35 +7255,22 @@ class ServerArgs:
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)
@property
def modelexpress_transport(self) -> str:
"""Transport backend for modelexpress: 'transfer_engine' (default) or 'nixl'."""
return self._parsed_modelexpress_config.get("transport", "transfer_engine")
"""Transport backend for modelexpress."""
return self._parsed_modelexpress_config.get("transport", "nixl")
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 needs TransferEngine init only if transport is transfer_engine.
if (
self.modelexpress_source
and self.modelexpress_transport == "transfer_engine"
):
return True
# Use TransferEngine as client backend.
elif (
self.load_format == "remote_instance"
and self.remote_instance_weight_loader_backend
in ("transfer_engine", "modelexpress")
and self.modelexpress_transport == "transfer_engine"
if self.load_format == "remote_instance" and (
self.remote_instance_weight_loader_backend == "transfer_engine"
or (
self.remote_instance_weight_loader_backend == "modelexpress"
and self.modelexpress_transport == "transfer_engine"
)
):
return True
else:
test/registered/unit/model_loader/test_runai_model_streamer_loader.py
+18
View File
@@ -123,6 +123,24 @@ class TestRunaiModelStreamerLoader(CustomTestCase):
self.assertIsInstance(model_loader, loader_mod.RunaiModelStreamerLoader)
def test_get_model_loader_uses_remote_instance_for_prequantized_modelopt(self):
load_config = LoadConfig(
load_format=LoadFormat.REMOTE_INSTANCE,
model_loader_extra_config={},
)
model_config = cast(
ModelConfig,
SimpleNamespace(
quantization="modelopt_fp4",
modelopt_quant=False,
_is_already_quantized=lambda: True,
),
)
model_loader = loader_mod.get_model_loader(load_config, model_config)
self.assertIsInstance(model_loader, loader_mod.RemoteInstanceModelLoader)
if __name__ == "__main__":
unittest.main()