@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 156 lines 7.29 KB
Rendered Raw
1 # ModelExpress Engine
2
3 The `modelexpress` weight transfer engine installs immutable weight versions
4 published through [ModelExpress](https://github.com/ai-dynamo/modelexpress).
5 Each update identifies the exact version to install, while ModelExpress handles
6 source selection, transfer, and installation.
7
8 ## When to Use ModelExpress
9
10 - Your workflow manages immutable weight versions through the ModelExpress
11 Refit API.
12 - Inference workers need to pull a specific published version through
13 ModelExpress, using GPU peer transfer or object storage.
14 - You want to load full checkpoints or replay checkpoint deltas from S3.
15
16 ## How It Works
17
18 1. A publisher creates a weight-version resource through the ModelExpress Refit
19 API, publishes its payload, and marks the version READY.
20 2. The orchestrator pauses generation and calls `start_weight_update` on the
21 inference engine.
22 3. The orchestrator calls `update_weights` with the target `version_id`.
23 Each worker's `ModelExpressGeneratorClient` resolves and stages that version,
24 then applies it to the live vLLM model.
25 4. After a successful update, `finish_weight_update` releases the staged handle.
26 The orchestrator can then resume inference.
27
28 Initialize the backend before starting an update session, and send exactly one
29 version per session. Invalid transitions must raise; an orchestrator must keep
30 inference paused after an update error and retry with a new session only after
31 cleanup. An error does not roll back weights that were already applied.
32
33 `ModelExpressTrainerClient` is an optional publisher. Custom integrations can
34 manage weight-version resources directly through the ModelExpress server's
35 Refit API. For S3, create a STAGING version with the correct model, payload
36 format, lineage, and index URI; upload the index and referenced shards in the
37 format expected by `ModelExpressGeneratorClient`; then mark the version READY.
38 See the [S3 Delta Weight Refit guide](https://github.com/ai-dynamo/modelexpress/blob/main/docs/S3_DELTA_WEIGHT_REFIT.md)
39 for the publication sequence and generator artifact requirements.
40
41 vLLM's `modelexpress_engine.py` re-exports the engine and its init/update
42 dataclasses from
43 `modelexpress_rl.inference.engines.vllm.weight_transfer_engine`.
44 The implementation and configuration schema live in ModelExpress. Upgrading
45 the installed ModelExpress package and restarting workers picks up compatible
46 changes without a vLLM code change. The package must continue to implement
47 vLLM's `WeightTransferEngine` interface.
48
49 ## Inference Side
50
51 !!! warning "Development version"
52
53 This backend currently uses unreleased ModelExpress features from `main`,
54 which may introduce breaking changes. These features will be included in
55 the official ModelExpress v0.8.0 release.
56
57 Install ModelExpress from the `main` branch on every worker:
58
59 ```bash
60 uv pip install "git+https://github.com/ai-dynamo/modelexpress@main#subdirectory=modelexpress_client/python"
61 ```
62
63 The backend is registered natively and loaded only when selected;
64 `VLLM_PLUGINS=modelexpress` is not required for weight transfer. If the plugin
65 is enabled, it preserves the native registration and uses the same engine
66 implementation.
67
68 ```python
69 from vllm import LLM
70 from vllm.config import WeightTransferConfig
71
72 llm = LLM(model="my-model", weight_transfer_config=WeightTransferConfig(backend="modelexpress"))
73 ```
74
75 ```bash
76 vllm serve /models/launch \
77 --load-format safetensors \
78 --weight-transfer-config '{"backend":"modelexpress"}'
79 ```
80
81 ### Initialize the engine
82
83 Initialize once by passing `init_info` to `init_weight_transfer_engine`.
84 For HTTP administration, set `VLLM_SERVER_DEV_MODE=1` before starting vLLM and
85 restrict these endpoints to trusted orchestrators.
86
87 For S3 updates, send the following body to `POST /init_weight_transfer_engine`:
88
89 ```json
90 {
91 "init_info": {
92 "model_name": "policy",
93 "server_url": "modelexpress:8001",
94 "object_storage_type": "S3",
95 "initial_base_version_id": "policy-v0",
96 "seed_checkpoint_path": "/models/launch",
97 "refit_checkpoint_dir": "/mxdelta/receiver",
98 "refit_checkpoint_max_size_gb": 200,
99 "object_storage_region_name": "us-west-2"
100 }
101 }
102 ```
103
104 ### `ModelExpressWeightTransferInitInfo`
105
106 These fields belong inside `init_info`, rather than
107 `--weight-transfer-config`. The installed ModelExpress package defines the
108 schema and defaults.
109
110 | Field | Default | Description |
111 | ----- | ------- | ----------- |
112 | `model_name` | vLLM model name | Logical model identity used for published ModelExpress versions. Set this when it differs from the vLLM model path or name. |
113 | `server_url` | ModelExpress configuration | ModelExpress server address, such as `modelexpress:8001`. If omitted, uses ModelExpress's environment settings, then `localhost:8001`. |
114 | `initial_serving_version_id` | `None` | Version already loaded by the inference worker. If omitted, ModelExpress uses the initial version supplied by its runtime, when available. |
115 | `registration_ttl_seconds` | Three MX heartbeat intervals | Worker registration lifetime in seconds, renewed by ModelExpress. |
116 | `lease_ttl_seconds` | Registration lifetime | Weight-version lease lifetime in seconds, renewed while held. |
117 | `max_transfer_attempts` | `3` | Maximum source-discovery and transfer attempts for one staged update. |
118 | `max_replay_chain_length` | `64` | Maximum number of payload revisions replayed to reach the requested version. |
119 | `rpc_timeout_seconds` | `30.0` | Deadline in seconds for each control-plane or manifest RPC. |
120
121 ### Object storage settings
122
123 To enable object storage, provide `object_storage_type`,
124 `initial_base_version_id`, `seed_checkpoint_path`, and `refit_checkpoint_dir`
125 inside `init_info`. Endpoint and region settings also require these four fields.
126
127 | Field | Default | Description |
128 | ----- | ------- | ----------- |
129 | `object_storage_type` | `None` | Object storage provider. Use `"S3"`; the generator currently supports S3. |
130 | `initial_base_version_id` | `None` | Registered base version matching the seed checkpoint. Required for object storage. |
131 | `seed_checkpoint_path` | `None` | Local seed checkpoint path accessible to each worker. Required for object storage. |
132 | `refit_checkpoint_dir` | `None` | Local directory for cached checkpoints and delta replay. Required for object storage. |
133 | `refit_checkpoint_max_size_gb` | ModelExpress default | Checkpoint cache quota in decimal GB. The example sets an explicit 200 GB quota; it does not reserve disk space. |
134 | `object_storage_endpoint_url` | `None` | Optional endpoint override for an S3-compatible service. |
135 | `object_storage_region_name` | `None` | Optional S3 region, such as `us-west-2`. |
136
137 The seed checkpoint must match `initial_base_version_id`. Credentials use
138 ModelExpress's normal credential chain. Omit the object storage settings when
139 using only GPU peer transfer.
140
141 ### Apply a weight update
142
143 After the target version is READY, pause generation and call
144 `POST /start_weight_update`. Then send the exact immutable version ID to
145 `POST /update_weights`:
146
147 ```json
148 {"update_info": {"version_id": "policy-v1"}}
149 ```
150
151 Call `POST /finish_weight_update` only after the update succeeds, then resume
152 inference. The staged handle is retained until finish or shutdown. Staging and
153 installation failures propagate; do not resume a worker with uncertain weights.
154
155 See the ModelExpress documentation for publication, source selection, object
156 storage, and deployment requirements.