ModelExpress Python Client

Python client for ModelExpress -- high-performance GPU-to-GPU model weight transfers using NVIDIA NIXL over RDMA/InfiniBand.

Instead of each inference engine instance loading model weights from storage, one instance loads the model and transfers weights directly to later instances via GPUDirect RDMA, bypassing the CPU entirely.

Installation

# From PyPI (coming soon)
pip install modelexpress

# Editable install from source
pip install -e .

# With test dependencies
pip install -e ".[dev]"

# Additionally install the pinned protobuf code generator when changing protobuf APIs
pip install -e ".[codegen]"

NIXL is expected to be supplied by the runtime environment (TRT-LLM, SGLang, Dynamo, and NemoRL runtime images all ship nixl-cu12 or nixl-cu13). For a bare-environment install, run pip install nixl-cu12 or pip install nixl-cu13 separately, matching your host CUDA toolkit.

The collective refit path in modelexpress_rl.collective needs NCCL M2N, which ships separately from nccl4py as the nccl-extensions distribution under the same nccl namespace. Run pip install "nccl-extensions[cu12]" or pip install "nccl-extensions[cu13]", matching your host CUDA toolkit. The CUDA major has to be named: a bare pip install nccl-extensions resolves and then fails to import, because cuda-bindings only arrives through one of those two extras. Nothing else in the client imports it, so an installation that never uses collective refit does not need it.

Reshard needs NCCL 2.30.7 or newer, and an image that ships its own older libnccl will win the load over the one pip installed. That failure surfaces inside the reshard call rather than at import, so check which library the process resolves rather than which one is on disk. The NGC PyTorch 26.06 image ships 2.30.5 and needs the newer library preloaded.

Requirements

• Python >= 3.10
• protobuf >= 5.27.2 and < 7
• NVIDIA GPUs with RDMA/InfiniBand support
• NIXL (NVIDIA Interconnect eXchange Library)
• NCCL M2N (nccl-extensions, collective refit only)
• A running ModelExpress server (Rust gRPC service backed by Redis)

Quick Start with vLLM

vLLM 0.23.0 and newer recognize --load-format modelexpress natively. Install the ModelExpress Python package in the vLLM image; no VLLM_PLUGINS setting or manual loader registration is required. For older vLLM versions, set VLLM_PLUGINS=modelexpress or call register_modelexpress_loaders() manually.

export MX_SERVER_ADDRESS="modelexpress-server:8001"

vllm serve deepseek-ai/DeepSeek-V4-Pro \
    --load-format modelexpress \
    --tensor-parallel-size 8 \
    --trust-remote-code

Starting the vLLM engine with the modelexpress load format on the source worker will load the weights from disk and register/publish the NIXL and tensor metadata to the MX server. The mx load format is kept as a backward-compatible alias. On the target worker, it retrieves metadata from the MX server and streams weights over RDMA from GPU to GPU. Set MX_ARTIFACT_TRANSFER=1 to also reuse compatible vLLM JIT caches from a ready source.

Quick Start with SGLang

SGLang integrates through its remote_instance loader with the modelexpress backend. Use an SGLang image that includes upstream sgl-project/sglang#24723, such as the known-good release image lmsysorg/sglang:v0.5.13.post1, and install the ModelExpress package into that image.

export MX_SERVER_ADDRESS="modelexpress-server:8001"

python -m sglang.launch_server \
    --model-path deepseek-ai/DeepSeek-V3 \
    --tp 8 \
    --load-format remote_instance \
    --remote-instance-weight-loader-backend modelexpress \
    --modelexpress-config '{"transport": "nixl"}'

Quick Start with TensorRT-LLM

TensorRT-LLM integrates through its native checkpoint_format="MX" interface. Install ModelExpress in a qualified TensorRT-LLM image, then construct the PyTorch backend with the ModelExpress server configuration:

from tensorrt_llm.llmapi import LLM

llm = LLM(
    model="/model",
    checkpoint_format="MX",
    mx_config={
        "server_url": "modelexpress-server:8001",
    },
    tensor_parallel_size=4,
    backend="pytorch",
)

The first replica falls back to the Hugging Face checkpoint and publishes its post-transform weights; later compatible replicas receive them through ModelExpress. The current qualified scope is the LlamaForCausalLM family. See the TensorRT-LLM P2P example for the qualified-image requirement and production-style Kubernetes deployment.

Programmatic Usage

RL trainer and generator APIs

See the RL package quickstart and support guide for trainer publication, generator initialization, and import examples.

MxClient

MxClient is a lightweight gRPC client for communicating with the ModelExpress server:

from modelexpress import MxClient

client = MxClient(server_url="modelexpress-server:8001")

# Query for a source model
response = client.get_metadata("deepseek-ai/DeepSeek-V4-Pro")
if response.found:
    for worker in response.workers:
        print(f"Worker rank {worker.worker_rank}: {len(worker.tensors)} tensors")

# Wait for source readiness (blocks until ready or timeout)
success, session_id, metadata_hash = client.wait_for_ready(
    model_name="deepseek-ai/DeepSeek-V4-Pro",
    worker_id=0,
    timeout_seconds=7200,
)

client.close()

Registering Loaders Manually

Manual registration is only needed for integrations that construct vLLM loaders outside vLLM 0.23.0's native load-format path.

from modelexpress import register_modelexpress_loaders

register_modelexpress_loaders()
# Now vLLM recognizes --load-format modelexpress and mx

Environment Variables

Shared client settings are registered in modelexpress.envs. RL settings are listed in the RL configuration reference.

Canonical S3 Transfer Tuning

Objects below the configured thresholds use one PUT or GET. Larger uploads use multipart parts. Downloads default to one persistent s3transfer.TransferManager per S3Client. Set MX_S3_DOWNLOAD_BACKEND=preallocated to opt in to parallel GETs into a preallocated buffer. This affects canonical S3 artifact downloads for refit and recovery; the normal ModelStreamer loading path has its own downloader.

The preallocated backend starts with HEAD for size and ETag. Below the range threshold, it performs one whole-object GET; otherwise it issues parallel ranges. All data GETs, including the first range and body-read retries, share one per-client pool limited by MX_REFIT_DOWNLOAD_WORKERS (falling back to MX_S3_DOWNLOAD_WORKERS, then 16). HEAD requests are outside this data-GET budget. Every GET is pinned to the HEAD ETag with IfMatch; an object change or an invalid range response fails without retrying. Endpoints must support HEAD, conditional GET, and byte ranges. Transient body-read failures use capped exponential backoff with full jitter (0–100 ms before the first retry, doubling to a 2 s cap); request retries remain managed by botocore. Each backend owns its download resources behind a shared get/close interface. Select transfer_manager explicitly to return to the default implementation; there is no automatic backend switch on an error.

Ranges write directly into a preallocated BytesIO buffer. Exported views are released before returning immutable bytes, avoiding a second full-object allocation on CPython. Older SDK bodies without readinto use bounded read chunks. The complete object remains resident, and concurrent file downloads and retained delta shards add to peak memory; this is not a bounded-memory model streamer. Validate peak RSS and throughput on your target endpoint and runtime before enabling it for large models.

To compare the SDK writer with direct parallel writes, run the three-arm benchmark from modelexpress_client/python against an immutable S3 object:

PYTHONPATH=. python benchmarks/bench_s3_download.py s3://bucket/checkpoint \
  --iterations 3 --warmup 1

It compares default TransferManager, TransferManager with preallocated output, and the preallocated backend using the same transfer settings. It rotates arm order, checks payload hashes and object ETags, and emits per-download JSON timing and throughput. Hashing is outside the timed interval; HEAD and buffer setup needed by each arm are included. Use --endpoint-url and --region-name for compatible endpoints. Run on the target host and endpoint: the reported upstream GB300 result does not establish that buffer growth causes the performance gap.

UCX/NIXL Tuning

Package Structure

The optional RL APIs are installed with this distribution under the modelexpress_rl namespace. See the RL package guide.

How It Works

1. Source loads weights from disk, registers raw tensors with NIXL before FP8 processing, and publishes metadata to the ModelExpress server.
2. Target creates dummy weights, waits for the source ready flag, then pulls raw tensors via RDMA read.
3. Both source and target run process_weights_after_loading() independently, producing identical FP8-transformed weights.
4. When artifact transfer is enabled, a healthy source publishes its pod-scoped JIT caches and later pods install compatible caches before model initialization.

This pre-processing transfer strategy is critical for FP8 models (e.g., DeepSeek-V4-Pro) where tensors are renamed and transformed during processing.

License

Apache-2.0