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.

Requirements

• Python >= 3.10
• protobuf >= 5.27.2 and < 7
• NVIDIA GPUs with RDMA/InfiniBand support
• NIXL (NVIDIA Interconnect eXchange Library)
• 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 publication

An RL framework creates a weight version through the external Refit API. Each trainer actor then invokes its rank-local client to stage and publish one shard. Worker registration, manifest serving, and internal shard CRUD remain hidden behind the client.

from modelexpress_rl import (
    ModelExpressTrainerClient,
    ModelExpressTrainerConfig,
    WeightVersionRef,
)

trainer = ModelExpressTrainerClient.initialize(ModelExpressTrainerConfig())
trainer.bind_tensors(megatron_tensor_specs)
trainer.publish_version(version=WeightVersionRef(version.uid))

The deployment supplies MODEL_NAME, MX_TRAINER_ENGINE, MX_TRAINER_STAGING_MODE, MX_WEIGHT_PAYLOAD_FORMAT, MX_WORKER_HOST, and the normal ModelExpress server configuration. The Megatron adapter derives its source slot from logical tensor names and shard geometry. DP replicas of the same partition therefore publish redundant workers for one slot, while distinct TP partitions remain separate required slots. The NIXL metadata endpoint is derived from MX_WORKER_HOST and the client-owned NIXL manager's listen port. LOCAL_RANK selects the device unless device_id is passed to initialize().

The client owns the NIXL manager and trainer-side manifest service. server_url selects the central ModelExpress control-plane service and defaults to the normal ModelExpress server configuration. A Megatron worker may initialize the client after selecting its CUDA device but before creating NCCL resources; the engine adapter is created lazily when the worker first requests its source slot or stages a shard after distributed setup.

Initialization fixes the staging mode and payload format. publish() hides manifest publication and the internal CreateWeightVersionShard RPC. The current Megatron adapter registers and exposes its live buffers through IN_PLACE, so callers must keep those tensors immutable while the version is published. The required lifecycle is synchronous: create and publish the version, update every generator, retire and release the version, and only then resume training or begin the next optimizer step.

Version creation and expected-source-slot declaration remain framework-orchestrator responsibilities. Each trainer adapter derives its own source slot from the engine's native topology; the orchestrator declares the expected slots using the same adapter-defined convention. initialize() selects the configured trainer engine and constructs its adapter internally. Megatron and FSDP implementations are available. Megatron-specific APIs live under modelexpress_rl; modelexpress.refit.reshard remains the shared, engine-neutral transfer core.

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

UCX/NIXL Tuning

Package Structure

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