diff --git a/docs/training/weight_transfer/README.md b/docs/training/weight_transfer/README.md index 77d2728345..c8382e2a06 100644 --- a/docs/training/weight_transfer/README.md +++ b/docs/training/weight_transfer/README.md @@ -36,6 +36,7 @@ engine drives on your behalf: | [sparse_nccl](nccl.md#sparse-nccl) | NCCL broadcast | Checkpoint-coordinate sparse weight patches | | [sharded_rdt](sharded_rdt.md) | NIXL / Ray Direct Transport (pull-based) | Very large models where each worker needs only its own slice (MoE with expert parallelism) | | [nccl_m2n](m2n.md) | NCCL M2N reshard | Trainer and inference use different sharding layouts | +| [ModelExpress](modelexpress.md) | ModelExpress version-based transfer | Install immutable published weight versions, including object-storage checkpoints and deltas | ## Quickstart diff --git a/docs/training/weight_transfer/modelexpress.md b/docs/training/weight_transfer/modelexpress.md new file mode 100644 index 0000000000..fd47e09d22 --- /dev/null +++ b/docs/training/weight_transfer/modelexpress.md @@ -0,0 +1,156 @@ +# ModelExpress Engine + +The `modelexpress` weight transfer engine installs immutable weight versions +published through [ModelExpress](https://github.com/ai-dynamo/modelexpress). +Each update identifies the exact version to install, while ModelExpress handles +source selection, transfer, and installation. + +## When to Use ModelExpress + +- Your workflow manages immutable weight versions through the ModelExpress + Refit API. +- Inference workers need to pull a specific published version through + ModelExpress, using GPU peer transfer or object storage. +- You want to load full checkpoints or replay checkpoint deltas from S3. + +## How It Works + +1. A publisher creates a weight-version resource through the ModelExpress Refit + API, publishes its payload, and marks the version READY. +2. The orchestrator pauses generation and calls `start_weight_update` on the + inference engine. +3. The orchestrator calls `update_weights` with the target `version_id`. + Each worker's `ModelExpressGeneratorClient` resolves and stages that version, + then applies it to the live vLLM model. +4. After a successful update, `finish_weight_update` releases the staged handle. + The orchestrator can then resume inference. + +Initialize the backend before starting an update session, and send exactly one +version per session. Invalid transitions must raise; an orchestrator must keep +inference paused after an update error and retry with a new session only after +cleanup. An error does not roll back weights that were already applied. + +`ModelExpressTrainerClient` is an optional publisher. Custom integrations can +manage weight-version resources directly through the ModelExpress server's +Refit API. For S3, create a STAGING version with the correct model, payload +format, lineage, and index URI; upload the index and referenced shards in the +format expected by `ModelExpressGeneratorClient`; then mark the version READY. +See the [S3 Delta Weight Refit guide](https://github.com/ai-dynamo/modelexpress/blob/main/docs/S3_DELTA_WEIGHT_REFIT.md) +for the publication sequence and generator artifact requirements. + +vLLM's `modelexpress_engine.py` re-exports the engine and its init/update +dataclasses from +`modelexpress_rl.inference.engines.vllm.weight_transfer_engine`. +The implementation and configuration schema live in ModelExpress. Upgrading +the installed ModelExpress package and restarting workers picks up compatible +changes without a vLLM code change. The package must continue to implement +vLLM's `WeightTransferEngine` interface. + +## Inference Side + +!!! warning "Development version" + + This backend currently uses unreleased ModelExpress features from `main`, + which may introduce breaking changes. These features will be included in + the official ModelExpress v0.8.0 release. + +Install ModelExpress from the `main` branch on every worker: + +```bash +uv pip install "git+https://github.com/ai-dynamo/modelexpress@main#subdirectory=modelexpress_client/python" +``` + +The backend is registered natively and loaded only when selected; +`VLLM_PLUGINS=modelexpress` is not required for weight transfer. If the plugin +is enabled, it preserves the native registration and uses the same engine +implementation. + +```python +from vllm import LLM +from vllm.config import WeightTransferConfig + +llm = LLM(model="my-model", weight_transfer_config=WeightTransferConfig(backend="modelexpress")) +``` + +```bash +vllm serve /models/launch \ + --load-format safetensors \ + --weight-transfer-config '{"backend":"modelexpress"}' +``` + +### Initialize the engine + +Initialize once by passing `init_info` to `init_weight_transfer_engine`. +For HTTP administration, set `VLLM_SERVER_DEV_MODE=1` before starting vLLM and +restrict these endpoints to trusted orchestrators. + +For S3 updates, send the following body to `POST /init_weight_transfer_engine`: + +```json +{ + "init_info": { + "model_name": "policy", + "server_url": "modelexpress:8001", + "object_storage_type": "S3", + "initial_base_version_id": "policy-v0", + "seed_checkpoint_path": "/models/launch", + "refit_checkpoint_dir": "/mxdelta/receiver", + "refit_checkpoint_max_size_gb": 200, + "object_storage_region_name": "us-west-2" + } +} +``` + +### `ModelExpressWeightTransferInitInfo` + +These fields belong inside `init_info`, rather than +`--weight-transfer-config`. The installed ModelExpress package defines the +schema and defaults. + +| Field | Default | Description | +| ----- | ------- | ----------- | +| `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. | +| `server_url` | ModelExpress configuration | ModelExpress server address, such as `modelexpress:8001`. If omitted, uses ModelExpress's environment settings, then `localhost:8001`. | +| `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. | +| `registration_ttl_seconds` | Three MX heartbeat intervals | Worker registration lifetime in seconds, renewed by ModelExpress. | +| `lease_ttl_seconds` | Registration lifetime | Weight-version lease lifetime in seconds, renewed while held. | +| `max_transfer_attempts` | `3` | Maximum source-discovery and transfer attempts for one staged update. | +| `max_replay_chain_length` | `64` | Maximum number of payload revisions replayed to reach the requested version. | +| `rpc_timeout_seconds` | `30.0` | Deadline in seconds for each control-plane or manifest RPC. | + +### Object storage settings + +To enable object storage, provide `object_storage_type`, +`initial_base_version_id`, `seed_checkpoint_path`, and `refit_checkpoint_dir` +inside `init_info`. Endpoint and region settings also require these four fields. + +| Field | Default | Description | +| ----- | ------- | ----------- | +| `object_storage_type` | `None` | Object storage provider. Use `"S3"`; the generator currently supports S3. | +| `initial_base_version_id` | `None` | Registered base version matching the seed checkpoint. Required for object storage. | +| `seed_checkpoint_path` | `None` | Local seed checkpoint path accessible to each worker. Required for object storage. | +| `refit_checkpoint_dir` | `None` | Local directory for cached checkpoints and delta replay. Required for object storage. | +| `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. | +| `object_storage_endpoint_url` | `None` | Optional endpoint override for an S3-compatible service. | +| `object_storage_region_name` | `None` | Optional S3 region, such as `us-west-2`. | + +The seed checkpoint must match `initial_base_version_id`. Credentials use +ModelExpress's normal credential chain. Omit the object storage settings when +using only GPU peer transfer. + +### Apply a weight update + +After the target version is READY, pause generation and call +`POST /start_weight_update`. Then send the exact immutable version ID to +`POST /update_weights`: + +```json +{"update_info": {"version_id": "policy-v1"}} +``` + +Call `POST /finish_weight_update` only after the update succeeds, then resume +inference. The staged handle is retained until finish or shutdown. Staging and +installation failures propagate; do not resume a worker with uncertain weights. + +See the ModelExpress documentation for publication, source selection, object +storage, and deployment requirements. diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index 1353def738..3af7b7aeb9 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/tests/distributed/test_weight_transfer.py @@ -6,10 +6,14 @@ Unit tests for engine classes (parsing, validation, registry). Integration tests for NCCL and IPC weight transfer between processes using Ray. """ +import builtins import pickle +import runpy import threading import time from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock import pybase64 as base64 @@ -297,6 +301,36 @@ class TestEngineRegistry: ) assert isinstance(engine, SparseNCCLWeightTransferEngine) + def test_modelexpress_registration_is_native_and_lazy(self, monkeypatch): + from vllm.distributed.weight_transfer import factory + + imported = [] + module = MagicMock() + module.ModelExpressWeightTransferEngine.__name__ = ( + "ModelExpressWeightTransferEngine" + ) + + def import_module(name): + imported.append(name) + return module + + monkeypatch.setattr(factory.importlib, "import_module", import_module) + config = WeightTransferConfig(backend="modelexpress") + vllm_config = create_mock_vllm_config() + device = torch.device("cpu") + model = MagicMock(spec=torch.nn.Module) + + assert imported == [] + engine = WeightTransferEngineFactory.create_engine( + config, vllm_config, device, model + ) + + assert imported == ["vllm.distributed.weight_transfer.modelexpress_engine"] + module.ModelExpressWeightTransferEngine.assert_called_once_with( + config, vllm_config, device, model + ) + assert engine is module.ModelExpressWeightTransferEngine.return_value + def test_create_engine_invalid_backend(self): config = WeightTransferConfig(backend="invalid") with pytest.raises(ValueError, match="Invalid weight transfer backend"): @@ -1976,3 +2010,322 @@ def test_sparse_nccl_trainer_non_sender_skips_client(): engine.send_weights([_sparse_patch()]) engine.send_weight_chunk([_sparse_patch()]) assert client.order == [] + + +# --- Optional ModelExpress Client Lifecycle --- + + +def _import_modelexpress_shim(): + path = ( + Path(__file__).parents[2] + / "vllm/distributed/weight_transfer/modelexpress_engine.py" + ) + runpy.run_path(str(path), run_name="_test_modelexpress_shim") + + +@pytest.mark.parametrize( + "missing_module", + [ + "modelexpress", + "modelexpress_rl", + "modelexpress_rl.inference", + "modelexpress_rl.inference.engines", + "modelexpress_rl.inference.engines.vllm", + "modelexpress_rl.inference.engines.vllm.weight_transfer_engine", + ], +) +def test_modelexpress_missing_backend_explains_source_install( + monkeypatch, missing_module +): + original_import = builtins.__import__ + error = ModuleNotFoundError(name=missing_module) + + def import_without_backend(name, *args, **kwargs): + if name == "modelexpress": + raise error + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", import_without_backend) + with pytest.raises(ImportError, match="modelexpress_client/python") as exc_info: + _import_modelexpress_shim() + assert exc_info.value.__cause__ is error + assert "uv pip install" in str(exc_info.value) + + +@pytest.mark.parametrize("missing_module", ["grpc", "modelexpress.internal_dependency"]) +def test_modelexpress_preserves_transitive_import_errors(monkeypatch, missing_module): + original_import = builtins.__import__ + error = ModuleNotFoundError(name=missing_module) + + def import_without_dependency(name, *args, **kwargs): + if name == "modelexpress": + raise error + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", import_without_dependency) + with pytest.raises(ModuleNotFoundError) as exc_info: + _import_modelexpress_shim() + assert exc_info.value is error + + +@pytest.fixture +def mx_client(monkeypatch): + client_module = pytest.importorskip("modelexpress_rl.inference.client") + client = MagicMock() + monkeypatch.setattr( + client_module.ModelExpressGeneratorClient, + "initialize", + MagicMock(return_value=client), + ) + monkeypatch.setattr(torch.accelerator, "synchronize", MagicMock()) + return client + + +def _make_modelexpress_engine(): + return WeightTransferEngineFactory.create_engine( + WeightTransferConfig(backend="modelexpress"), + SimpleNamespace( + parallel_config=SimpleNamespace(), + model_config=SimpleNamespace(model="test/model"), + ), + torch.device("cpu"), + torch.nn.Linear(2, 2), + ) + + +@pytest.fixture +def mx_worker(mx_client, monkeypatch): + from vllm.v1.worker import gpu_worker + + monkeypatch.setattr(gpu_worker, "set_current_vllm_config", lambda _: nullcontext()) + worker = object.__new__(gpu_worker.Worker) + worker.weight_transfer_engine = _make_modelexpress_engine() + worker.vllm_config = worker.weight_transfer_engine.vllm_config + worker._weight_update_active = False + worker._weight_update_is_draft = False + worker.model_runner = MagicMock() + return worker + + +def test_modelexpress_worker_rejects_updates_before_initialization( + mx_worker, mx_client +): + with pytest.raises(RuntimeError, match="not initialized"): + mx_worker.start_weight_update() + with pytest.raises(RuntimeError, match="start_weight_update must be called"): + mx_worker.update_weights({"version_id": "version-a"}) + with pytest.raises(RuntimeError, match="without a matching"): + mx_worker.finish_weight_update() + mx_client.stage_weight.assert_not_called() + + +@pytest.mark.parametrize( + "second_update, error_type", + [ + ({"version_id": "version-a"}, RuntimeError), + ({"version_id": "version-b"}, RuntimeError), + ({"version_id": ""}, ValueError), + ({"version_id": None}, ValueError), + ({}, ValueError), + ], +) +def test_modelexpress_worker_recovers_after_rejected_update( + mx_worker, mx_client, second_update, error_type +): + """A rejected payload must release A so a fresh session can actually install B.""" + mx_worker.weight_transfer_engine.init_transfer_engine( + mx_worker.weight_transfer_engine.parse_init_info({}) + ) + staged_a = SimpleNamespace(version_id="version-a", metrics={}, release=MagicMock()) + staged_b = SimpleNamespace(version_id="version-b", metrics={}, release=MagicMock()) + mx_client.stage_weight.side_effect = [staged_a, staged_b] + mx_client.apply_weight.return_value = {} + mx_worker.start_weight_update() + mx_worker.update_weights({"version_id": "version-a"}) + + with pytest.raises(error_type): + mx_worker.update_weights(second_update) + staged_a.release.assert_called_once_with() + mx_client.apply_weight.assert_called_once_with(staged_a) + + mx_worker.start_weight_update() + mx_worker.update_weights({"version_id": "version-b"}) + mx_worker.finish_weight_update() + + assert [ + call.kwargs["version"].version_id + for call in mx_client.stage_weight.call_args_list + ] == ["version-a", "version-b"] + mx_client.apply_weight.assert_called_with(staged_b) + staged_b.release.assert_called_once_with() + + +def test_modelexpress_worker_can_retry_failed_finish(mx_worker, mx_client): + mx_worker.weight_transfer_engine.init_transfer_engine( + mx_worker.weight_transfer_engine.parse_init_info({}) + ) + staged = mx_client.stage_weight.return_value + staged.release.side_effect = [RuntimeError("release failed"), None] + mx_worker.start_weight_update() + mx_worker.update_weights({"version_id": "version-a"}) + + with pytest.raises(RuntimeError, match="release failed"): + mx_worker.finish_weight_update() + mx_worker.finish_weight_update() + + assert staged.release.call_count == 2 + assert not mx_worker._weight_update_active + mx_worker.start_weight_update() + + +def test_native_engine_applies_and_releases_exact_versions(mx_client): + from modelexpress_rl.inference.engines.vllm import ( + weight_transfer_engine as mx_engine, + ) + + from vllm.distributed.weight_transfer import modelexpress_engine + + engine = _make_modelexpress_engine() + assert type(engine) is mx_engine.ModelExpressWeightTransferEngine + assert ( + modelexpress_engine.ModelExpressWeightTransferInitInfo + is mx_engine.ModelExpressWeightTransferInitInfo + ) + assert ( + modelexpress_engine.ModelExpressWeightTransferUpdateInfo + is mx_engine.ModelExpressWeightTransferUpdateInfo + ) + assert not engine.supports_draft_weight_update + engine.init_transfer_engine(engine.parse_init_info({})) + + for version_id in ("version-a", "version-b"): + staged = SimpleNamespace(version_id=version_id, metrics={}, release=MagicMock()) + mx_client.stage_weight.return_value = staged + mx_client.apply_weight.return_value = {} + engine.start_weight_update() + engine.update_weights({"version_id": version_id}) + + assert ( + mx_client.stage_weight.call_args.kwargs["version"].version_id == version_id + ) + mx_client.apply_weight.assert_called_with(staged) + torch.accelerator.synchronize.assert_called() + staged.release.assert_not_called() + engine.finish_weight_update() + staged.release.assert_called_once_with() + + engine.shutdown() + engine.shutdown() + mx_client.close.assert_called_once_with() + + +def test_vime_init_config_reaches_mx_client(mx_client): + from modelexpress_rl.inference.client import ModelExpressGeneratorClient + + engine = _make_modelexpress_engine() + ModelExpressGeneratorClient.initialize.assert_not_called() + engine.init_transfer_engine( + engine.parse_init_info( + { + "model_name": "policy", + "server_url": "mx:8001", + "initial_serving_version_id": "serving-a", + "object_storage_type": "S3", + "initial_base_version_id": "base-a", + "seed_checkpoint_path": "/models/launch", + "refit_checkpoint_dir": "/mxdelta/receiver", + "refit_checkpoint_max_size_gb": 200, + "object_storage_endpoint_url": "http://minio:9000", + "object_storage_region_name": "us-west-2", + "registration_ttl_seconds": 90, + "lease_ttl_seconds": 60, + "max_transfer_attempts": 4, + "max_replay_chain_length": 17, + "rpc_timeout_seconds": 12.5, + } + ) + ) + config = ModelExpressGeneratorClient.initialize.call_args.args[0] + assert config.engine_context.model is engine.model + assert config.engine_context.vllm_config is engine.vllm_config + assert config.model_name == "policy" + assert config.server_url == "mx:8001" + assert config.initial_serving_version_id == "serving-a" + assert config.registration_ttl_seconds == 90 + assert config.lease_ttl_seconds == 60 + assert config.max_transfer_attempts == 4 + assert config.max_replay_chain_length == 17 + assert config.rpc_timeout_seconds == 12.5 + storage = config.object_storage + assert storage.storage_type.value == "S3" + assert storage.initial_base_version_id == "base-a" + assert storage.seed_checkpoint_path == "/models/launch" + assert storage.refit_checkpoint_dir == "/mxdelta/receiver" + assert storage.refit_checkpoint_max_size_gb == 200 + assert storage.endpoint_url == "http://minio:9000" + assert storage.region_name == "us-west-2" + + +@pytest.mark.parametrize( + "key", + [ + "object_storage_type", + "initial_base_version_id", + "seed_checkpoint_path", + "refit_checkpoint_dir", + "object_storage_endpoint_url", + "object_storage_region_name", + ], +) +def test_incomplete_storage_config_rejected_before_init(mx_client, key): + from modelexpress_rl.inference.client import ModelExpressGeneratorClient + + engine = _make_modelexpress_engine() + with pytest.raises(ValueError, match="object storage requires"): + engine.init_transfer_engine(engine.parse_init_info({key: "value"})) + ModelExpressGeneratorClient.initialize.assert_not_called() + + +@pytest.mark.parametrize("version_id", ["", " ", None, 1]) +def test_invalid_version_rejected_before_staging(mx_client, version_id): + engine = _make_modelexpress_engine() + engine.init_transfer_engine(engine.parse_init_info({})) + engine.start_weight_update() + with pytest.raises(ValueError, match="version_id is required"): + engine.update_weights({"version_id": version_id}) + mx_client.stage_weight.assert_not_called() + + +@pytest.mark.parametrize("failure", ["stage", "apply", "release"]) +def test_update_failure_preserves_error_and_releases_handle(mx_client, failure): + engine = _make_modelexpress_engine() + engine.init_transfer_engine(engine.parse_init_info({})) + staged = SimpleNamespace(version_id="version-a", metrics={}, release=MagicMock()) + mx_client.stage_weight.return_value = staged + if failure == "stage": + mx_client.stage_weight.side_effect = RuntimeError("stage failed") + else: + mx_client.apply_weight.side_effect = RuntimeError("apply failed") + if failure == "release": + staged.release.side_effect = RuntimeError("release failed") + engine.start_weight_update() + with pytest.raises(RuntimeError, match="stage failed|apply failed"): + engine.update_weights({"version_id": "version-a"}) + assert staged.release.call_count == (0 if failure == "stage" else 1) + torch.accelerator.synchronize.assert_not_called() + engine.shutdown() + mx_client.close.assert_called_once_with() + + +def test_plugin_does_not_replace_native_backend(mx_client): + from modelexpress.engines.vllm.registration import ( + register_plugin_weight_transfer_engine, + ) + from modelexpress_rl.inference.engines.vllm.weight_transfer_engine import ( + ModelExpressWeightTransferEngine, + ) + + native_loader = WeightTransferEngineFactory._registry["modelexpress"] + register_plugin_weight_transfer_engine() + assert WeightTransferEngineFactory._registry["modelexpress"] is native_loader + assert type(_make_modelexpress_engine()) is ModelExpressWeightTransferEngine diff --git a/vllm/distributed/weight_transfer/factory.py b/vllm/distributed/weight_transfer/factory.py index 82228bef34..c0e88e03bf 100644 --- a/vllm/distributed/weight_transfer/factory.py +++ b/vllm/distributed/weight_transfer/factory.py @@ -250,6 +250,13 @@ WeightTransferEngineFactory.register_engine( ) +WeightTransferEngineFactory.register_engine( + "modelexpress", + "vllm.distributed.weight_transfer.modelexpress_engine", + "ModelExpressWeightTransferEngine", +) + + # Trainer-side engines, parallel to the worker registry above. WeightTransferTrainerFactory.register_engine( "nccl", diff --git a/vllm/distributed/weight_transfer/modelexpress_engine.py b/vllm/distributed/weight_transfer/modelexpress_engine.py new file mode 100644 index 0000000000..1172542e85 --- /dev/null +++ b/vllm/distributed/weight_transfer/modelexpress_engine.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Re-export the weight-transfer backend maintained by ModelExpress. + +Install ModelExpress from the main branch's modelexpress_client/python +subdirectory. +""" + +try: + from modelexpress import configure_vllm_logging + from modelexpress_rl.inference.engines.vllm.weight_transfer_engine import ( + ModelExpressWeightTransferEngine, + ModelExpressWeightTransferInitInfo, + ModelExpressWeightTransferUpdateInfo, + ) +except ModuleNotFoundError as exc: + if exc.name not in { + "modelexpress", + "modelexpress_rl", + "modelexpress_rl.inference", + "modelexpress_rl.inference.engines", + "modelexpress_rl.inference.engines.vllm", + "modelexpress_rl.inference.engines.vllm.weight_transfer_engine", + }: + raise + raise ImportError( + "The 'modelexpress' weight transfer backend requires ModelExpress. " + "Install it with `uv pip install " + "'git+https://github.com/ai-dynamo/modelexpress@main" + "#subdirectory=modelexpress_client/python'`." + ) from exc + +configure_vllm_logging() + +__all__ = [ + "ModelExpressWeightTransferEngine", + "ModelExpressWeightTransferInitInfo", + "ModelExpressWeightTransferUpdateInfo", +]