A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 03:00 UTC
[Core] Add native ModelExpress load format (#43105)
Signed-off-by: Zheng Luo <zheluo@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Zheng Luo committed
May 21, 2026 at 13:05 UTC
17b69828a013acb7af0cd1d16d24ecc8d7582094
5 files changed
+212
-1
tests/model_executor/model_loader/test_modelexpress_loader.py
new
+133
@@ -0,0 +1,133 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import sys
5
+from types import ModuleType, SimpleNamespace
6
+
7
+import pytest
8
+from torch import nn
9
+
10
+from vllm.config import VllmConfig
11
+from vllm.config.load import LoadConfig
12
+from vllm.model_executor.model_loader import get_model_loader
13
+from vllm.model_executor.model_loader.modelexpress_loader import (
14
+ ModelExpressModelLoader,
15
+)
16
+
17
+
18
+class FakeModelexpressLoader:
19
+ calls: list[tuple[str, tuple, dict]] = []
20
+ loaded_model: nn.Module
21
+
22
+ def __init__(self, load_config: LoadConfig):
23
+ self.load_config = load_config
24
+
25
+ def download_model(self, *args, **kwargs):
26
+ self.calls.append(("download_model", args, kwargs))
27
+
28
+ def load_weights(self, *args, **kwargs):
29
+ self.calls.append(("load_weights", args, kwargs))
30
+
31
+ def load_model(self, *args, **kwargs):
32
+ self.calls.append(("load_model", args, kwargs))
33
+ return self.loaded_model
34
+
35
+
36
+def _install_fake_modelexpress(monkeypatch):
37
+ FakeModelexpressLoader.calls = []
38
+ FakeModelexpressLoader.loaded_model = nn.Module()
39
+
40
+ for name in [
41
+ "modelexpress",
42
+ "modelexpress.engines",
43
+ "modelexpress.engines.vllm",
44
+ ]:
45
+ monkeypatch.setitem(sys.modules, name, ModuleType(name))
46
+
47
+ module = ModuleType("modelexpress.engines.vllm.loader")
48
+ module.__dict__["MxModelLoader"] = FakeModelexpressLoader
49
+ monkeypatch.setitem(sys.modules, module.__name__, module)
50
+
51
+
52
+def test_modelexpress_load_format_resolves_to_modelexpress_loader(monkeypatch):
53
+ _install_fake_modelexpress(monkeypatch)
54
+
55
+ loader = get_model_loader(LoadConfig(load_format="modelexpress"))
56
+
57
+ assert isinstance(loader, ModelExpressModelLoader)
58
+
59
+
60
+def test_modelexpress_loader_delegates_to_modelexpress(monkeypatch):
61
+ _install_fake_modelexpress(monkeypatch)
62
+ loader = ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
63
+ model = nn.Module()
64
+ model_config = SimpleNamespace()
65
+ vllm_config = SimpleNamespace()
66
+
67
+ loader.download_model(model_config)
68
+ loader.load_weights(model, model_config)
69
+ FakeModelexpressLoader.loaded_model.train()
70
+ result = loader.load_model(
71
+ vllm_config=vllm_config,
72
+ model_config=model_config,
73
+ prefix="model",
74
+ )
75
+
76
+ assert result is FakeModelexpressLoader.loaded_model
77
+ assert not result.training
78
+ assert FakeModelexpressLoader.calls == [
79
+ ("download_model", (model_config,), {}),
80
+ ("load_weights", (model, model_config), {}),
81
+ (
82
+ "load_model",
83
+ (),
84
+ {
85
+ "vllm_config": vllm_config,
86
+ "model_config": model_config,
87
+ "prefix": "model",
88
+ },
89
+ ),
90
+ ]
91
+
92
+
93
+def test_modelexpress_loader_missing_modelexpress_error(monkeypatch):
94
+ import importlib
95
+
96
+ def missing_modelexpress(name):
97
+ raise ModuleNotFoundError(name=name)
98
+
99
+ monkeypatch.setattr(importlib, "import_module", missing_modelexpress)
100
+
101
+ with pytest.raises(ImportError, match="requires the ModelExpress Python package"):
102
+ ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
103
+
104
+
105
+def test_modelexpress_loader_preserves_internal_import_errors(monkeypatch):
106
+ import importlib
107
+
108
+ def missing_dependency(name):
109
+ raise ModuleNotFoundError(name="not_modelexpress_dependency")
110
+
111
+ monkeypatch.setattr(importlib, "import_module", missing_dependency)
112
+
113
+ with pytest.raises(ModuleNotFoundError) as exc_info:
114
+ ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
115
+ assert exc_info.value.name == "not_modelexpress_dependency"
116
+
117
+
118
+def test_modelexpress_load_format_allows_object_storage_model_weights():
119
+ model_config = SimpleNamespace(
120
+ architecture="UnknownForTest",
121
+ config_updated=False,
122
+ convert_type=None,
123
+ is_hybrid=False,
124
+ model="test-model",
125
+ model_weights="s3://bucket/model",
126
+ )
127
+ vllm_config = object.__new__(VllmConfig)
128
+ vllm_config.model_config = model_config
129
+ vllm_config.load_config = LoadConfig(load_format="modelexpress")
130
+
131
+ vllm_config.try_verify_and_update_config()
132
+
133
+ assert vllm_config.load_config.load_format == "modelexpress"
vllm/config/load.py
+1
@@ -55,6 +55,7 @@ class LoadConfig:
55
https://github.com/ggml-org/ggml/blob/master/docs/gguf.md).
56
- "mistral" will load weights from consolidated safetensors files used by
57
Mistral models.
58
+ - "modelexpress" will load weights using ModelExpress.
59
- Other custom values can be supported via plugins.
60
"""
61
download_dir: str | None = None
vllm/config/vllm.py
+2
-1
@@ -1901,12 +1901,13 @@ class VllmConfig:
1901
)
1902
self.load_config.load_format = "runai_streamer"
1903
elif self.load_config.load_format not in (
1904
+ "modelexpress",
1905
"runai_streamer",
1906
"runai_streamer_sharded",
1907
):
1908
raise ValueError(
1909
f"To load a model from object storage (S3/GCS/Azure), "
1909
- f"'load_format' must be 'runai_streamer' or "
1910
+ f"'load_format' must be 'modelexpress', 'runai_streamer' or "
1911
f"'runai_streamer_sharded', "
1912
f"but got '{self.load_config.load_format}'. "
1913
f"Model: {self.model_config.model}"
vllm/model_executor/model_loader/__init__.py
+6
@@ -13,6 +13,9 @@ from vllm.model_executor.model_loader.bitsandbytes_loader import BitsAndBytesMod
13
from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
14
from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader
15
from vllm.model_executor.model_loader.gguf_loader import GGUFModelLoader
16
+from vllm.model_executor.model_loader.modelexpress_loader import (
17
+ ModelExpressModelLoader,
18
+)
19
from vllm.model_executor.model_loader.runai_streamer_loader import (
20
RunaiModelStreamerLoader,
21
)
@@ -37,6 +40,7 @@ LoadFormats = Literal[
40
"gguf",
41
"instanttensor",
42
"mistral",
43
+ "modelexpress",
44
"npcache",
45
"pt",
46
"runai_streamer",
@@ -54,6 +58,7 @@ _LOAD_FORMAT_TO_MODEL_LOADER: dict[str, type[BaseModelLoader]] = {
58
"gguf": GGUFModelLoader,
59
"instanttensor": DefaultModelLoader,
60
"mistral": DefaultModelLoader,
61
+ "modelexpress": ModelExpressModelLoader,
62
"npcache": DefaultModelLoader,
63
"pt": DefaultModelLoader,
64
"runai_streamer": RunaiModelStreamerLoader,
@@ -150,6 +155,7 @@ __all__ = [
155
"BaseModelLoader",
156
"BitsAndBytesModelLoader",
157
"GGUFModelLoader",
158
+ "ModelExpressModelLoader",
159
"DefaultModelLoader",
160
"DummyModelLoader",
161
"RunaiModelStreamerLoader",
vllm/model_executor/model_loader/modelexpress_loader.py
new
+70
@@ -0,0 +1,70 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+import importlib
7
+
8
+from torch import nn
9
+
10
+from vllm.config import ModelConfig, VllmConfig
11
+from vllm.config.load import LoadConfig
12
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
13
+from vllm.tracing import instrument
14
+
15
+_MODELEXPRESS_LOADER_MODULE = "modelexpress.engines.vllm.loader"
16
+_MISSING_MODELEXPRESS_MODULES = frozenset(
17
+ {
18
+ "modelexpress",
19
+ "modelexpress.engines",
20
+ "modelexpress.engines.vllm",
21
+ _MODELEXPRESS_LOADER_MODULE,
22
+ }
23
+)
24
+
25
+
26
+def _missing_modelexpress_error() -> ImportError:
27
+ return ImportError(
28
+ "The 'modelexpress' load format requires the ModelExpress Python package. "
29
+ "Install it with `pip install modelexpress`."
30
+ )
31
+
32
+
33
+class ModelExpressModelLoader(BaseModelLoader):
34
+ """Thin vLLM loader wrapper for ModelExpress."""
35
+
36
+ def __init__(self, load_config: LoadConfig):
37
+ super().__init__(load_config)
38
+ self._loader = self._load_modelexpress_loader(load_config)
39
+
40
+ @staticmethod
41
+ def _load_modelexpress_loader(load_config: LoadConfig) -> BaseModelLoader:
42
+ try:
43
+ module = importlib.import_module(_MODELEXPRESS_LOADER_MODULE)
44
+ except ModuleNotFoundError as exc:
45
+ if exc.name not in _MISSING_MODELEXPRESS_MODULES:
46
+ raise
47
+ raise _missing_modelexpress_error() from exc
48
+
49
+ ModelExpressVllmLoader = module.MxModelLoader
50
+ return ModelExpressVllmLoader(load_config)
51
+
52
+ def download_model(self, model_config: ModelConfig) -> None:
53
+ self._loader.download_model(model_config)
54
+
55
+ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
56
+ self._loader.load_weights(model, model_config)
57
+
58
+ @instrument(span_name="Load model")
59
+ def load_model(
60
+ self,
61
+ vllm_config: VllmConfig,
62
+ model_config: ModelConfig,
63
+ prefix: str = "",
64
+ ) -> nn.Module:
65
+ model = self._loader.load_model(
66
+ vllm_config=vllm_config,
67
+ model_config=model_config,
68
+ prefix=prefix,
69
+ )
70
+ return model.eval()