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, and larger downloads use ranged GETs through one persistent s3transfer.TransferManager per S3Client. The receiver's file-level pool and the manager's global request concurrency use the same worker setting, so all whole-object and ranged data GETs share one 16-request budget. HEAD requests use the manager's separate submission executor. Downloads target a seekable BytesIO, so the complete downloaded object remains resident; the I/O settings below bound queued chunks, not the final object size.

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