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.

When creating a WeightVersion, the orchestrator may supply its UID or let MX generate one. A caller-supplied UID already assigned to another request returns ALREADY_EXISTS; an identical request retried with the same idempotency key returns the existing version.

from modelexpress_rl import (
    MegatronTrainerContext,
    ModelExpressTrainerClient,
    ModelExpressTrainerConfig,
    WeightVersionRef,
)

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

The deployment supplies MODEL_NAME, 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().

Canonical S3 staging consumes Hugging Face tensor buckets produced by the training framework. Framework-native bucket settings remain the default. The public trainer API accepts ModelExpressTrainerConfig(object_storage=ObjectStorageConfig(...)); its storage_type selects the provider. Weight versions use the corresponding typed ObjectStorageSource envelope. Generator clients likewise accept ModelExpressGeneratorConfig(object_storage=ObjectStorageGeneratorConfig(...)). The current trainer and generator clients support only ObjectStorageType.S3. Integrations may use MX_REFIT_DELTA_BUCKET_BYTES as an explicit override, or its 512 MiB default when they have no native setting. CPU workers are configured by MX_REFIT_DELTA_WORKERS (default min(32, CPU count)), while MX_S3_UPLOAD_WORKERS controls concurrent full-checkpoint batch uploads. MX_REFIT_CHECKSUM_FORMAT selects the checksum algorithm and defaults to adler32. The framework integration reads the bucket-size setting while constructing the stream; ModelExpress processes each supplied bucket without splitting or merging it. Before training begins, the framework calls prepare_delta_base() with one bucket stream. ModelExpress submits each framework bucket directly for concurrent rank-local seed-checkpoint reads. Real delta staging therefore performs no seed-checkpoint reads. A FULL_HF_CHECKPOINT version serializes the current buckets as native HF safetensor shards, omits base_version_id, and replaces the retained snapshot so the next XOR_DELTA uses it as its exact base. A bounded worker pool updates the rank-local snapshot with immutable CPU tensors. Each publishing rank groups its snapshot into concurrently uploaded safetensors objects with at most MX_REFIT_FULL_CHECKPOINT_BATCH_BYTES tensor bytes (4 GiB by default); an oversized tensor occupies its own object. The objects are sent directly to S3 without a trainer-side temporary checkpoint. Framework integrations own the optional full-checkpoint period; it is disabled by default. Generators use the ModelExpress S3 client to download full-checkpoint batches concurrently. Each worker validates one downloaded batch and copies its tensors into their existing local mmap destinations, without materializing a second full checkpoint. ModelStreamer integration remains a future optimization. The local checkpoint state changes from READY to UPDATING before mutation and returns to READY only after success. An interrupted update must be reseeded from seed_checkpoint_path during initialization. The framework supplies each version's exact object_storage.uri under the configured uri_prefix. That URI names the global safetensors index; its objects are stored beside it. A delta index records the target WeightVersion.uid and its base_version_id as metadata.version and metadata.base_version. After upload, the orchestrator changes the version from STAGING to READY. S3 versions remain READY for rollout recovery; their immutable objects are governed by the bucket's external lifecycle policy.

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 before its distributed process group is ready; the explicitly selected engine_context is constructed lazily on the first tensor operation. Deployment environment variables do not select Python implementations.

Initialization fixes the staging mode. NIXL also fixes its payload format; canonical S3 publication follows each target WeightVersion. On NIXL, 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() constructs the adapter selected by engine_context 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

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

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