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 14, 05:30 UTC
[Core][Model] Terratorch backend integration (#23513)
Signed-off-by: Michele Gazzetti <michele.gazzetti1@ibm.com> Signed-off-by: Christian Pinto <christian.pinto@ibm.com> Co-authored-by: Christian Pinto <christian.pinto@ibm.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
mgazz committed
Sep 4, 2025 at 08:22 UTC
51d5e9be7dbf4d914374447548dd01f9bfb68f89
23 files changed
+305
-208
examples/offline_inference/prithvi_geospatial_mae.py
+5
-1
@@ -45,7 +45,11 @@ datamodule_config = {
45
class PrithviMAE:
46
def __init__(self, model):
47
self.model = LLM(
48
- model=model, skip_tokenizer_init=True, dtype="float16", enforce_eager=True
48
+ model=model,
49
+ skip_tokenizer_init=True,
50
+ dtype="float16",
51
+ enforce_eager=True,
52
+ model_impl="terratorch",
53
)
54
55
def run(self, input_data, location_coords):
examples/offline_inference/prithvi_geospatial_mae_io_processor.py
+1
@@ -37,6 +37,7 @@ def main():
37
# The maximum number depends on the available GPU memory
38
max_num_seqs=32,
39
io_processor_plugin="prithvi_to_tiff_india",
40
+ model_impl="terratorch",
41
)
42
43
pooling_params = PoolingParams(task="encode", softmax=False)
examples/online_serving/prithvi_geospatial_mae.py
+1
@@ -15,6 +15,7 @@ import requests
15
# https://github.com/christian-pinto/prithvi_io_processor_plugin
16
# - start vllm in serving mode with the below args
17
# --model='christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM'
18
+# --model-impl terratorch
19
# --task embed --trust-remote-code
20
# --skip-tokenizer-init --enforce-eager
21
# --io-processor-plugin prithvi_to_tiff_india
requirements/test.in
+1
-1
@@ -53,5 +53,5 @@ runai-model-streamer==0.11.0
53
runai-model-streamer-s3==0.11.0
54
fastsafetensors>=0.1.10
55
pydantic>=2.10 # 2.9 leads to error on python 3.10
56
-terratorch==1.1rc2 # required for PrithviMAE test
56
decord==0.6.0
57
+terratorch==1.1rc3 # required for PrithviMAE test
requirements/test.txt
+1
-1
@@ -1042,7 +1042,7 @@ tensorboardx==2.6.4
1042
# via lightning
1043
tensorizer==2.10.1
1044
# via -r requirements/test.in
1045
-terratorch==1.1rc2
1045
+terratorch==1.1rc3
1046
# via -r requirements/test.in
1047
threadpoolctl==3.5.0
1048
# via scikit-learn
tests/distributed/test_pipeline_parallel.py
+6
@@ -298,6 +298,8 @@ def _compare_tp(
298
tokenizer_mode = model_info.tokenizer_mode
299
hf_overrides = model_info.hf_overrides
300
hf_config = get_config(model_id, trust_remote_code)
301
+ skip_tokenizer_init = model_info.skip_tokenizer_init
302
+ max_num_seqs = model_info.max_num_seqs
303
304
dtype = "float16"
305
if hf_config.model_type in _FLOAT16_NOT_SUPPORTED_MODELS:
@@ -351,6 +353,10 @@ def _compare_tp(
353
common_args.extend(["--load-format", load_format])
354
if hf_overrides:
355
common_args.extend(["--hf-overrides", json.dumps(hf_overrides)])
356
+ if skip_tokenizer_init:
357
+ common_args.append("--skip-tokenizer-init")
358
+ if max_num_seqs:
359
+ common_args.extend(["--max-num-seqs", f"{max_num_seqs}"])
360
361
specific_case = tp_size == 2 and pp_size == 2 and chunked_prefill
362
testing_ray_compiled_graph = False
tests/distributed/test_sequence_parallel.py
+3
@@ -178,6 +178,7 @@ def _compare_sp(
178
trust_remote_code = model_info.trust_remote_code
179
tokenizer_mode = model_info.tokenizer_mode
180
hf_overrides = model_info.hf_overrides
181
+ skip_tokenizer_init = model_info.skip_tokenizer_init
182
183
if load_format == "dummy":
184
# Avoid OOM
@@ -227,6 +228,8 @@ def _compare_sp(
228
common_args.extend(["--load-format", load_format])
229
if hf_overrides:
230
common_args.extend(["--hf-overrides", json.dumps(hf_overrides)])
231
+ if skip_tokenizer_init:
232
+ common_args.append("--skip-tokenizer-init")
233
234
compilation_config = {
235
'level': 3,
tests/entrypoints/openai/test_chat_template.py
+3
-1
@@ -104,7 +104,9 @@ def test_get_gen_prompt(model, template, add_generation_prompt,
104
trust_remote_code=model_info.trust_remote_code,
105
revision=model_info.revision,
106
hf_overrides=model_info.hf_overrides,
107
- )
107
+ skip_tokenizer_init=model_info.skip_tokenizer_init,
108
+ enforce_eager=model_info.enforce_eager,
109
+ dtype=model_info.dtype)
110
111
# Initialize the tokenizer
112
tokenizer = get_tokenizer(
tests/entrypoints/openai/test_skip_tokenizer.py
+4
-2
@@ -11,7 +11,7 @@ import torch
11
12
from ...utils import RemoteOpenAIServer
13
14
-MODEL_NAME = "christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM"
14
+MODEL_NAME = "mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11"
15
DTYPE = "float16"
16
17
@@ -35,7 +35,9 @@ def server():
35
"--trust-remote-code",
36
"--skip-tokenizer-init",
37
"--max-num-seqs",
38
- "32"
38
+ "32",
39
+ "--model-impl",
40
+ "terratorch"
41
]
42
43
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
tests/entrypoints/test_chat_utils.py
+9
-3
@@ -1266,7 +1266,9 @@ def test_resolve_hf_chat_template(sample_json_schema, model, use_tools):
1266
revision=model_info.revision,
1267
trust_remote_code=model_info.trust_remote_code,
1268
hf_overrides=model_info.hf_overrides,
1269
- )
1269
+ skip_tokenizer_init=model_info.skip_tokenizer_init,
1270
+ enforce_eager=model_info.enforce_eager,
1271
+ dtype=model_info.dtype)
1272
1273
# Build the tokenizer group and grab the underlying tokenizer
1274
tokenizer_group = TokenizerGroup(
@@ -1322,7 +1324,9 @@ def test_resolve_content_format_hf_defined(model, expected_format):
1324
revision=model_info.revision,
1325
trust_remote_code=model_info.trust_remote_code,
1326
hf_overrides=model_info.hf_overrides,
1325
- )
1327
+ skip_tokenizer_init=model_info.skip_tokenizer_init,
1328
+ enforce_eager=model_info.enforce_eager,
1329
+ dtype=model_info.dtype)
1330
1331
tokenizer_group = TokenizerGroup(
1332
model,
@@ -1382,7 +1386,9 @@ def test_resolve_content_format_fallbacks(model, expected_format):
1386
revision=model_info.revision,
1387
trust_remote_code=model_info.trust_remote_code,
1388
hf_overrides=model_info.hf_overrides,
1385
- )
1389
+ skip_tokenizer_init=model_info.skip_tokenizer_init,
1390
+ enforce_eager=model_info.enforce_eager,
1391
+ dtype=model_info.dtype)
1392
1393
tokenizer_group = TokenizerGroup(
1394
model_config.tokenizer,
tests/models/multimodal/generation/vlm_utils/core.py
+3
@@ -69,6 +69,9 @@ def run_test(
69
vllm_runner_kwargs_["tokenizer_mode"] = model_info.tokenizer_mode
70
if model_info.hf_overrides:
71
vllm_runner_kwargs_["hf_overrides"] = model_info.hf_overrides
72
+ if model_info.skip_tokenizer_init:
73
+ vllm_runner_kwargs_[
74
+ "skip_tokenizer_init"] = model_info.skip_tokenizer_init
75
76
if vllm_runner_kwargs:
77
vllm_runner_kwargs_.update(vllm_runner_kwargs)
tests/models/multimodal/pooling/test_prithvi_mae.py
+1
-1
@@ -46,7 +46,7 @@ def _run_test(
46
vllm_model.encode(prompt)
47
48
49
-MODELS = ["christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM"]
49
+MODELS = ["mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11"]
50
51
52
@pytest.mark.core_model
tests/models/multimodal/processing/test_common.py
+3
-1
@@ -66,7 +66,9 @@ def _test_processing_correctness(
66
hf_overrides=model_info.hf_overrides,
67
# Ensure that the cache can fit all of the data
68
mm_processor_cache_gb=2048,
69
- )
69
+ skip_tokenizer_init=model_info.skip_tokenizer_init,
70
+ enforce_eager=model_info.enforce_eager,
71
+ dtype=model_info.dtype)
72
73
model_cls = MULTIMODAL_REGISTRY._get_model_cls(model_config)
74
factories = MULTIMODAL_REGISTRY._processor_factories[model_cls]
tests/models/multimodal/processing/test_tensor_schema.py
+3
-1
@@ -196,7 +196,9 @@ def test_model_tensor_schema(model_arch: str, model_id: str):
196
revision=model_info.revision,
197
trust_remote_code=model_info.trust_remote_code,
198
hf_overrides=hf_overrides_fn,
199
- )
199
+ skip_tokenizer_init=model_info.skip_tokenizer_init,
200
+ enforce_eager=model_info.enforce_eager,
201
+ dtype=model_info.dtype)
202
model_cls = MULTIMODAL_REGISTRY._get_model_cls(model_config)
203
factories = MULTIMODAL_REGISTRY._processor_factories[model_cls]
204
tests/models/multimodal/test_mapping.py
+3
-1
@@ -59,7 +59,9 @@ def test_hf_model_weights_mapper(model_arch: str):
59
revision=model_info.revision,
60
trust_remote_code=model_info.trust_remote_code,
61
hf_overrides=model_info.hf_overrides,
62
- )
62
+ skip_tokenizer_init=model_info.skip_tokenizer_init,
63
+ enforce_eager=model_info.enforce_eager,
64
+ dtype=model_info.dtype)
65
model_cls = MULTIMODAL_REGISTRY._get_model_cls(model_config)
66
67
original_weights = create_repo_dummy_weights(model_id)
tests/models/registry.py
+37
-3
@@ -6,10 +6,11 @@ from dataclasses import dataclass, field
6
from typing import Any, Literal, Optional
7
8
import pytest
9
+import torch
10
from packaging.version import Version
11
from transformers import __version__ as TRANSFORMERS_VERSION
12
12
-from vllm.config import TokenizerMode
13
+from vllm.config import ModelDType, TokenizerMode
14
15
16
@dataclass(frozen=True)
@@ -47,6 +48,23 @@ class _HfExamplesInfo:
48
The reason for the minimum/maximum version requirement.
49
"""
50
51
+ skip_tokenizer_init: bool = False
52
+ """
53
+ If true, skip initialization of tokenizer and detokenizer.
54
+ """
55
+
56
+ dtype: ModelDType = "auto"
57
+ """
58
+ The data type for the model weights and activations.
59
+ """
60
+
61
+ enforce_eager: bool = False
62
+ """
63
+ Whether to enforce eager execution. If True, we will
64
+ disable CUDA graph and always execute the model in eager mode.
65
+ If False, we will use CUDA graph and eager execution in hybrid.
66
+ """
67
+
68
is_available_online: bool = True
69
"""
70
Set this to ``False`` if the name of this architecture no longer exists on
@@ -76,6 +94,9 @@ class _HfExamplesInfo:
94
If not specified, the default revision will be used.
95
"""
96
97
+ max_num_seqs: Optional[int] = None
98
+ """Maximum number of sequences to be processed in a single iteration."""
99
+
100
def check_transformers_version(
101
self,
102
*,
@@ -361,8 +382,21 @@ _EMBEDDING_EXAMPLE_MODELS = {
382
"Phi3VForCausalLM": _HfExamplesInfo("TIGER-Lab/VLM2Vec-Full",
383
trust_remote_code=True),
384
"Qwen2VLForConditionalGeneration": _HfExamplesInfo("MrLight/dse-qwen2-2b-mrl-v1"), # noqa: E501
364
- "PrithviGeoSpatialMAE": _HfExamplesInfo("ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11", # noqa: E501
365
- is_available_online=False), # noqa: E501
385
+ "PrithviGeoSpatialMAE": _HfExamplesInfo("mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11", # noqa: E501
386
+ dtype=torch.float16,
387
+ enforce_eager=True,
388
+ skip_tokenizer_init=True,
389
+ # This is to avoid the model
390
+ # going OOM in CI
391
+ max_num_seqs=32,
392
+ ),
393
+ "Terratorch": _HfExamplesInfo("mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
394
+ dtype=torch.float16,
395
+ enforce_eager=True,
396
+ skip_tokenizer_init=True,
397
+ # This is to avoid the model going OOM in CI
398
+ max_num_seqs=32,
399
+ ),
400
}
401
402
_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
tests/models/test_initialization.py
+4
-1
@@ -73,6 +73,9 @@ def can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch,
73
tokenizer=model_info.tokenizer,
74
tokenizer_mode=model_info.tokenizer_mode,
75
revision=model_info.revision,
76
+ enforce_eager=model_info.enforce_eager,
77
+ skip_tokenizer_init=model_info.skip_tokenizer_init,
78
+ dtype=model_info.dtype,
79
speculative_config={
80
"model": model_info.speculative_model,
81
"num_speculative_tokens": 1,
@@ -85,7 +88,7 @@ def can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch,
88
model_impl=ModelImpl.TRANSFORMERS
89
if model_arch in _TRANSFORMERS_BACKEND_MODELS else ModelImpl.VLLM,
90
hf_overrides=hf_overrides_fn,
88
- )
91
+ max_num_seqs=model_info.max_num_seqs)
92
93
94
@pytest.mark.parametrize("model_arch", HF_EXAMPLE_MODELS.get_supported_archs())
tests/models/test_terratorch.py
new
+45
@@ -0,0 +1,45 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from tests.conftest import VllmRunner
8
+from vllm.utils import set_default_torch_num_threads
9
+
10
+
11
+@pytest.mark.parametrize(
12
+ "model",
13
+ [
14
+ "mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
15
+ "mgazz/Prithvi_v2_eo_300_tl_unet_agb"
16
+ ],
17
+)
18
+def test_inference(
19
+ vllm_runner: type[VllmRunner],
20
+ model: str,
21
+) -> None:
22
+
23
+ pixel_values = torch.full((6, 512, 512), 1.0, dtype=torch.float16)
24
+ location_coords = torch.full((1, 2), 1.0, dtype=torch.float16)
25
+ prompt = dict(prompt_token_ids=[1],
26
+ multi_modal_data=dict(pixel_values=pixel_values,
27
+ location_coords=location_coords))
28
+ with (
29
+ set_default_torch_num_threads(1),
30
+ vllm_runner(
31
+ model,
32
+ runner="pooling",
33
+ dtype=torch.float16,
34
+ enforce_eager=True,
35
+ skip_tokenizer_init=True,
36
+ # Limit the maximum number of sequences to avoid the
37
+ # test going OOM during the warmup run
38
+ max_num_seqs=32,
39
+ ) as vllm_model,
40
+ ):
41
+
42
+ vllm_output = vllm_model.llm.encode(prompt)
43
+ assert torch.equal(
44
+ torch.isnan(vllm_output[0].outputs.data).any(),
45
+ torch.tensor(False))
tests/models/utils.py
+2
@@ -294,6 +294,8 @@ def build_model_context(
294
limit_mm_per_prompt=limit_mm_per_prompt,
295
mm_processor_cache_gb=mm_processor_cache_gb,
296
hf_overrides=model_info.hf_overrides,
297
+ skip_tokenizer_init=model_info.skip_tokenizer_init,
298
+ enforce_eager=model_info.enforce_eager,
299
**model_config_kwargs,
300
)
301
return InputContext(model_config)
tests/plugins_tests/test_io_processor_plugins.py
+45
-58
@@ -7,12 +7,11 @@ import requests
7
8
from tests.utils import RemoteOpenAIServer
9
from vllm.config import VllmConfig
10
-from vllm.entrypoints.llm import LLM
10
from vllm.entrypoints.openai.protocol import IOProcessorResponse
11
from vllm.plugins.io_processors import get_io_processor
12
from vllm.pooling_params import PoolingParams
13
15
-MODEL_NAME = "christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM"
14
+MODEL_NAME = "mgazz/Prithvi-EO-2.0-300M-TL-Sen1Floods11"
15
16
image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" # noqa: E501
17
@@ -23,61 +22,7 @@ def test_loading_missing_plugin():
22
get_io_processor(vllm_config, "wrong_plugin")
23
24
26
-def test_loading_engine_with_wrong_plugin():
27
-
28
- with pytest.raises(ValueError):
29
- LLM(
30
- model=MODEL_NAME,
31
- skip_tokenizer_init=True,
32
- trust_remote_code=True,
33
- enforce_eager=True,
34
- # Limit the maximum number of parallel requests
35
- # to avoid the model going OOM in CI.
36
- max_num_seqs=32,
37
- io_processor_plugin="wrong_plugin",
38
- )
39
-
40
-
41
-@pytest.mark.parametrize("model_name", [MODEL_NAME])
42
-def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):
43
-
44
- img_prompt = dict(
45
- data=image_url,
46
- data_format="url",
47
- image_format="tiff",
48
- out_data_format="b64_json",
49
- )
50
-
51
- pooling_params = PoolingParams(task="encode", softmax=False)
52
-
53
- with vllm_runner(
54
- model_name,
55
- runner="pooling",
56
- skip_tokenizer_init=True,
57
- trust_remote_code=True,
58
- enforce_eager=True,
59
- # Limit the maximum number of parallel requests
60
- # to avoid the model going OOM in CI.
61
- max_num_seqs=1,
62
- io_processor_plugin="prithvi_to_tiff_valencia",
63
- ) as llm_runner:
64
- pooler_output = llm_runner.get_llm().encode(
65
- img_prompt,
66
- pooling_params=pooling_params,
67
- )
68
- output = pooler_output[0].outputs
69
-
70
- # verify the output is formatted as expected for this plugin
71
- assert all(
72
- hasattr(output, attr)
73
- for attr in ["type", "format", "data", "request_id"])
74
-
75
- # We just check that the output is a valid base64 string.
76
- # Raises an exception and fails the test if the string is corrupted.
77
- base64.b64decode(output.data)
78
-
79
-
80
-@pytest.fixture(scope="module")
25
+@pytest.fixture(scope="function")
26
def server():
27
args = [
28
"--runner",
@@ -90,7 +35,9 @@ def server():
35
"--max-num-seqs",
36
"32",
37
"--io-processor-plugin",
93
- "prithvi_to_tiff_valencia"
38
+ "prithvi_to_tiff_valencia",
39
+ "--model-impl",
40
+ "terratorch",
41
]
42
43
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
@@ -136,3 +83,43 @@ async def test_prithvi_mae_plugin_online(
83
# We just check that the output is a valid base64 string.
84
# Raises an exception and fails the test if the string is corrupted.
85
base64.b64decode(plugin_data["data"])
86
+
87
+
88
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
89
+def test_prithvi_mae_plugin_offline(vllm_runner, model_name: str):
90
+
91
+ img_prompt = dict(
92
+ data=image_url,
93
+ data_format="url",
94
+ image_format="tiff",
95
+ out_data_format="b64_json",
96
+ )
97
+
98
+ pooling_params = PoolingParams(task="encode", softmax=False)
99
+
100
+ with vllm_runner(
101
+ model_name,
102
+ runner="pooling",
103
+ skip_tokenizer_init=True,
104
+ trust_remote_code=True,
105
+ enforce_eager=True,
106
+ # Limit the maximum number of parallel requests
107
+ # to avoid the model going OOM in CI.
108
+ max_num_seqs=1,
109
+ model_impl="terratorch",
110
+ io_processor_plugin="prithvi_to_tiff_valencia",
111
+ ) as llm_runner:
112
+ pooler_output = llm_runner.get_llm().encode(
113
+ img_prompt,
114
+ pooling_params=pooling_params,
115
+ )
116
+ output = pooler_output[0].outputs
117
+
118
+ # verify the output is formatted as expected for this plugin
119
+ assert all(
120
+ hasattr(output, attr)
121
+ for attr in ["type", "format", "data", "request_id"])
122
+
123
+ # We just check that the output is a valid base64 string.
124
+ # Raises an exception and fails the test if the string is corrupted.
125
+ base64.b64decode(output.data)
vllm/config/__init__.py
+4
-1
@@ -171,6 +171,7 @@ class ModelImpl(str, enum.Enum):
171
AUTO = "auto"
172
VLLM = "vllm"
173
TRANSFORMERS = "transformers"
174
+ TERRATORCH = "terratorch"
175
176
177
def get_attr_docs(cls: type[Any]) -> dict[str, str]:
@@ -496,7 +497,9 @@ class ModelConfig:
497
back to the Transformers implementation if no vLLM implementation is
498
available.\n
499
- "vllm" will use the vLLM model implementation.\n
499
- - "transformers" will use the Transformers model implementation."""
500
+ - "transformers" will use the Transformers model implementation.\n
501
+ - "terratorch" will use the TerraTorch model implementation.
502
+ """
503
override_attention_dtype: Optional[str] = None
504
"""Override dtype for attention"""
505
logits_processors: Optional[list[Union[str, type[LogitsProcessor]]]] = None
vllm/model_executor/models/registry.py
+12
-3
@@ -184,10 +184,11 @@ _EMBEDDING_MODELS = {
184
"LlavaNextForConditionalGeneration": ("llava_next", "LlavaNextForConditionalGeneration"), # noqa: E501
185
"Phi3VForCausalLM": ("phi3v", "Phi3VForCausalLM"),
186
"Qwen2VLForConditionalGeneration": ("qwen2_vl", "Qwen2VLForConditionalGeneration"), # noqa: E501
187
- # Technically PrithviGeoSpatialMAE is a model that works on images, both in
188
- # input and output. I am adding it here because it piggybacks on embedding
187
+ # Technically Terratorch models work on images, both in
188
+ # input and output. I am adding it here because it piggy-backs on embedding
189
# models for the time being.
190
- "PrithviGeoSpatialMAE": ("prithvi_geospatial_mae", "PrithviGeoSpatialMAE"),
190
+ "PrithviGeoSpatialMAE": ("terratorch", "Terratorch"),
191
+ "Terratorch": ("terratorch", "Terratorch"),
192
}
193
194
_CROSS_ENCODER_MODELS = {
@@ -639,6 +640,9 @@ class _ModelRegistry:
640
model_info = self._try_inspect_model_cls(arch)
641
if model_info is not None:
642
return (model_info, arch)
643
+ elif model_config.model_impl == ModelImpl.TERRATORCH:
644
+ model_info = self._try_inspect_model_cls("Terratorch")
645
+ return (model_info, "Terratorch")
646
647
# Fallback to transformers impl (after resolving convert_type)
648
if (all(arch not in self.models for arch in architectures)
@@ -687,6 +691,11 @@ class _ModelRegistry:
691
model_cls = self._try_load_model_cls(arch)
692
if model_cls is not None:
693
return (model_cls, arch)
694
+ elif model_config.model_impl == ModelImpl.TERRATORCH:
695
+ arch = "Terratorch"
696
+ model_cls = self._try_load_model_cls(arch)
697
+ if model_cls is not None:
698
+ return (model_cls, arch)
699
700
# Fallback to transformers impl (after resolving convert_type)
701
if (all(arch not in self.models for arch in architectures)
vllm/model_executor/models/terratorch.py
renamed
+109
-129
@@ -15,13 +15,16 @@
15
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
# See the License for the specific language governing permissions and
17
# limitations under the License.
18
-"""Inference-only IBM/NASA Prithvi Geospatial model."""
18
+"""Wrapper around `Terratorch` models"""
19
20
+from collections import OrderedDict
21
from collections.abc import Iterable, Mapping, Sequence
21
-from typing import Any, Optional, Union
22
+from typing import Any, Callable, Optional, Union
23
24
import torch
25
import torch.nn as nn
26
+from terratorch.vllm import (DummyDataGenerator, InferenceRunner,
27
+ InputDefinition, InputTypeEnum)
28
from transformers import BatchFeature
29
30
from vllm.config import VllmConfig
@@ -29,6 +32,7 @@ from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
32
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
33
from vllm.model_executor.models.utils import AutoWeightsLoader
34
from vllm.multimodal import MULTIMODAL_REGISTRY
35
+from vllm.multimodal.cache import MultiModalProcessorOnlyCache
36
from vllm.multimodal.inputs import (ImageItem, ModalityData,
37
MultiModalDataDict, MultiModalFieldConfig,
38
MultiModalInputs, MultiModalKwargsItems,
@@ -45,52 +49,46 @@ from .interfaces import (IsAttentionFree, MultiModalEmbeddings,
49
from .interfaces_base import default_pooling_type
50
51
48
-def _prithvi_field_config(hf_inputs: Mapping[str, torch.Tensor]):
49
- # This model receives in input a multi-dimensional tensor representing
50
- # a single image patch and therefore it is not to be split
51
- # into multiple elements, but rather to be considered a single one.
52
- # Hence, the decision of using a MultiModalSharedField.
53
- # The expected shape is (num_channels, width, height).
52
+def _terratorch_field_names(pretrained_cfg: dict):
53
+ input_definition = InputDefinition(**pretrained_cfg["input"])
54
+ return set(input_definition.data.keys())
55
55
- # This model however allows the user to also submit multiple image
56
- # patches as a batch, adding a further dimension to the above shape.
57
- # At this stage we only support submitting one patch per request and
58
- # batching is achieved via vLLM batching.
59
- # TODO (christian-pinto): enable support for multi patch requests
60
- # in tandem with vLLM batching.
61
- return dict(
62
- pixel_values=MultiModalFieldConfig.shared(batch_size=1,
63
- modality="image"),
64
- location_coords=MultiModalFieldConfig.shared(batch_size=1,
65
- modality="image"),
66
- )
56
57
+def _terratorch_field_factory(
58
+ pretrained_cfg: dict
59
+) -> Callable[
60
+ [Mapping[str, torch.Tensor]],
61
+ Mapping[str, MultiModalFieldConfig],
62
+]:
63
69
-class PrithviGeoSpatialMAEMultiModalDataParser(MultiModalDataParser):
64
+ def _terratorch_field_config(hf_inputs: Mapping[str, torch.Tensor]):
65
+ input_definition = InputDefinition(**pretrained_cfg["input"])
66
+ fields = {}
67
+ for input_name, input in input_definition.data.items():
68
+ if input.type == InputTypeEnum.tensor:
69
+ fields[input_name] = "image"
70
71
- def _parse_image_data(
72
- self,
73
- data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],
74
- ) -> Optional[ModalityDataItems[Any, Any]]:
75
- if isinstance(data, dict):
76
- return DictEmbeddingItems(
77
- data,
78
- modality="image",
79
- required_fields={"pixel_values", "location_coords"},
80
- fields_factory=_prithvi_field_config,
81
- )
71
+ mm_fields_config = {}
72
+ for field_name, field_modality in fields.items():
73
+ mm_fields_config[field_name] = MultiModalFieldConfig.shared(
74
+ batch_size=1, modality=field_modality)
75
+ return mm_fields_config
76
83
- return super()._parse_image_data(data)
77
+ return _terratorch_field_config
78
79
86
-class PrithviGeoSpatialMAEProcessingInfo(BaseProcessingInfo):
80
+class TerratorchProcessingInfo(BaseProcessingInfo):
81
82
def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:
83
return {"image": None}
84
85
92
-class PrithviGeoSpatialMAEInputBuilder(
93
- BaseDummyInputsBuilder[PrithviGeoSpatialMAEProcessingInfo]):
86
+class TerratorchInputBuilder(BaseDummyInputsBuilder[TerratorchProcessingInfo]):
87
+
88
+ def __init__(self, info: TerratorchProcessingInfo):
89
+ super().__init__(info)
90
+ self.dummy_data_generator = DummyDataGenerator(
91
+ self.info.get_hf_config().to_dict()["pretrained_cfg"])
92
93
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
94
return ""
@@ -100,29 +98,57 @@ class PrithviGeoSpatialMAEInputBuilder(
98
seq_len: int,
99
mm_counts: Mapping[str, int],
100
) -> MultiModalDataDict:
103
- # This model input is fixed and is in the form of a torch Tensor.
104
- # The size of pixel_values might change in the cases where we resize
105
- # the input but never exceeds the dimensions below.
106
- image_data = {
107
- "pixel_values": torch.full((6, 512, 512), 1.0,
108
- dtype=torch.float16),
109
- "location_coords": torch.full((1, 2), 1.0, dtype=torch.float16),
110
- }
101
+ # Dummy data is generated based on the 'input' section
102
+ # defined in the HF configuration file
103
+ return self.dummy_data_generator.get_dummy_mm_data()
104
+
105
+
106
+class TerratorchMultiModalDataParser(MultiModalDataParser):
107
+
108
+ def __init__(self, pretrained_cfg: dict, *args, **kwargs):
109
+ self._pretrained_cfg = pretrained_cfg
110
+ super().__init__(*args, **kwargs)
111
+
112
+ def _parse_image_data(
113
+ self,
114
+ data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],
115
+ ) -> Optional[ModalityDataItems[Any, Any]]:
116
+ if isinstance(data, dict):
117
+
118
+ terratorch_fields = _terratorch_field_names(self._pretrained_cfg)
119
112
- return {"image": image_data}
120
+ return DictEmbeddingItems(
121
+ data,
122
+ modality="image",
123
+ required_fields=terratorch_fields,
124
+ fields_factory=_terratorch_field_factory(self._pretrained_cfg),
125
+ )
126
+
127
+ return super()._parse_image_data(data)
128
129
115
-class PrithviGeoSpatialMAEMultiModalProcessor(BaseMultiModalProcessor):
130
+class TerratorchMultiModalProcessor(BaseMultiModalProcessor):
131
+
132
+ def __init__(
133
+ self,
134
+ info: TerratorchProcessingInfo,
135
+ dummy_inputs: "BaseDummyInputsBuilder[TerratorchProcessingInfo]",
136
+ *,
137
+ cache: Optional[MultiModalProcessorOnlyCache] = None) -> None:
138
+
139
+ self.pretrained_cfg = info.get_hf_config().to_dict()["pretrained_cfg"]
140
+ super().__init__(info=info, dummy_inputs=dummy_inputs, cache=cache)
141
142
def _get_data_parser(self) -> MultiModalDataParser:
118
- return PrithviGeoSpatialMAEMultiModalDataParser()
143
+ return TerratorchMultiModalDataParser(
144
+ pretrained_cfg=self.pretrained_cfg)
145
146
def _get_mm_fields_config(
147
self,
148
hf_inputs: BatchFeature,
149
hf_processor_mm_kwargs: Mapping[str, object],
150
) -> Mapping[str, MultiModalFieldConfig]:
125
- return _prithvi_field_config(hf_inputs)
151
+ return _terratorch_field_factory(self.pretrained_cfg)(hf_inputs)
152
153
def _get_prompt_updates(
154
self,
@@ -173,13 +199,11 @@ class PrithviGeoSpatialMAEMultiModalProcessor(BaseMultiModalProcessor):
199
200
@default_pooling_type("All")
201
@MULTIMODAL_REGISTRY.register_processor(
176
- PrithviGeoSpatialMAEMultiModalProcessor,
177
- info=PrithviGeoSpatialMAEProcessingInfo,
178
- dummy_inputs=PrithviGeoSpatialMAEInputBuilder,
202
+ TerratorchMultiModalProcessor,
203
+ info=TerratorchProcessingInfo,
204
+ dummy_inputs=TerratorchInputBuilder,
205
)
180
-class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal):
181
- """Prithvi Masked Autoencoder"""
182
-
206
+class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):
207
supports_multimodal_raw_input_only = True
208
is_pooling_model = True
209
@@ -190,43 +214,13 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal):
214
215
raise ValueError("Only image modality is supported")
216
193
- def _instantiate_model(self, config: dict) -> Optional[nn.Module]:
194
- # We might be able/need to support different tasks with this same model
195
- if config["task_args"]["task"] == "SemanticSegmentationTask":
196
- from terratorch.cli_tools import SemanticSegmentationTask
197
-
198
- task = SemanticSegmentationTask(
199
- config["model_args"],
200
- config["task_args"]["model_factory"],
201
- loss=config["task_args"]["loss"],
202
- lr=config["task_args"]["lr"],
203
- ignore_index=config["task_args"]["ignore_index"],
204
- optimizer=config["task_args"]["optimizer"],
205
- optimizer_hparams=config["optimizer_params"],
206
- scheduler=config["task_args"]["scheduler"],
207
- scheduler_hparams=config["scheduler_params"],
208
- plot_on_val=config["task_args"]["plot_on_val"],
209
- freeze_decoder=config["task_args"]["freeze_decoder"],
210
- freeze_backbone=config["task_args"]["freeze_backbone"],
211
- )
212
-
213
- return task.model
214
- else:
215
- return None
216
-
217
def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
218
super().__init__()
219
220
- # the actual model is dynamically instantiated using terratorch
221
- # allowing us to perform changes to the model architecture
222
- # at startup time (e.g., change the model decoder class.)
223
- self.model = self._instantiate_model(
224
- vllm_config.model_config.hf_config.to_dict()["pretrained_cfg"])
225
- if self.model is None:
226
- raise ValueError(
227
- "Unsupported task. "
228
- "Only SemanticSegmentationTask is supported for now "
229
- "by PrithviGeospatialMAE.")
220
+ config = vllm_config.model_config.hf_config.to_dict()["pretrained_cfg"]
221
+
222
+ self.inference_runner = InferenceRunner(config)
223
+ self.model = self.inference_runner.model
224
225
pooler_config = vllm_config.model_config.pooler_config
226
assert pooler_config is not None
@@ -234,23 +228,6 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal):
228
self.pooler = DispatchPooler(
229
{"encode": Pooler.for_encode(pooler_config)}, )
230
237
- def _parse_and_validate_multimodal_data(
238
- self, **kwargs) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
239
- pixel_values = kwargs.pop("pixel_values", None)
240
- if not isinstance(pixel_values, torch.Tensor):
241
- raise ValueError(f"Incorrect type of pixel_values. "
242
- f"Got type: {type(pixel_values)}")
243
-
244
- location_coords = kwargs.pop("location_coords", None)
245
- if not isinstance(location_coords, torch.Tensor):
246
- raise ValueError(f"Incorrect type of location_coords. "
247
- f"Got type: {type(location_coords)}")
248
- location_coords = torch.unbind(location_coords, dim=0)[0]
249
- if location_coords.shape == torch.Size([0]):
250
- location_coords = None
251
-
252
- return pixel_values, location_coords
253
-
231
def get_input_embeddings(
232
self,
233
input_ids: torch.Tensor,
@@ -270,10 +247,7 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal):
247
inputs_embeds: Optional[torch.Tensor] = None,
248
**kwargs: object,
249
):
273
- pixel_values, location_coords = (
274
- self._parse_and_validate_multimodal_data(**kwargs))
275
- model_output = self.model(pixel_values,
276
- location_coords=location_coords)
250
+ model_output = self.inference_runner.forward(**kwargs)
251
252
return model_output.output
253
@@ -283,28 +257,34 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal):
257
model_buffers = dict(self.named_buffers())
258
loaded_buffers = []
259
for key, value in weights:
286
- if key == "state_dict":
287
- weights_to_parse = value
288
- for name, weight in weights_to_parse.items():
289
- if "pos_embed" in name:
290
- continue
291
-
292
- if "_timm_module." in name:
293
- name = name.replace("_timm_module.", "")
294
-
295
- # this model requires a couple of buffers to be loaded
296
- # that are not loadable with the AutoWeightsLoader
297
- if name in model_buffers:
260
+ if isinstance(value, (dict, OrderedDict)):
261
+ if key == "state_dict":
262
+ weights_to_parse = value
263
+ for name, weight in weights_to_parse.items():
264
+ name = f"inference_runner.{name}"
265
+
266
+ if "pos_embed" in name:
267
+ continue
268
+
269
if "_timm_module." in name:
270
name = name.replace("_timm_module.", "")
300
- buffer = model_buffers[name]
301
- weight_loader = getattr(buffer, "weight_loader",
302
- default_weight_loader)
303
- weight_loader(buffer, weight)
304
- loaded_buffers.append(name)
305
- else:
306
- params_list.append((name, weight))
307
- break
271
+
272
+ # this model requires a couple of buffers to be loaded
273
+ # that are not loadable with the AutoWeightsLoader
274
+ if name in model_buffers:
275
+ if "_timm_module." in name:
276
+ name = name.replace("_timm_module.", "")
277
+ buffer = model_buffers[name]
278
+ weight_loader = getattr(buffer, "weight_loader",
279
+ default_weight_loader)
280
+ weight_loader(buffer, weight)
281
+ loaded_buffers.append(name)
282
+ else:
283
+ params_list.append((name, weight))
284
+ break
285
+
286
+ elif isinstance(value, torch.Tensor):
287
+ params_list.append((f"inference_runner.model.{key}", value))
288
289
# Load the remaining model parameters
290
loader = AutoWeightsLoader(self)