@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8560a5b258
Fork 0

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, 23:45 UTC

[Core][Model] PrithviMAE Enablement on vLLM v1 engine (#20577)

[Core][Model] PrithviMAE Enablement on vLLM v1 engine (#20577) Signed-off-by: Christian Pinto <christian.pinto@ibm.com>

Christian Pinto committed Jul 23, 2025 at 19:00 UTC 8560a5b258946354b180e7fc694061d5296ec24f
15 files changed +704 -238
examples/offline_inference/prithvi_geospatial_mae.py
+70 -175
@@ -1,122 +1,27 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""
4 -This is a demo script showing how to use the
5 -PrithviGeospatialMAE model with vLLM
6 -This script is based on: https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11/blob/main/inference.py # noqa
7 -
8 -Target model weights: https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11/resolve/main/Prithvi-EO-V2-300M-TL-Sen1Floods11.pt # noqa
9 -
10 -The requirements for running this script are:
11 -- Installing [terratorch, albumentations, rasterio] in your python environment
12 -- downloading the model weights in a 'model' folder local to the script
13 - (temporary measure until the proper config.json file is uploaded to HF)
14 -- download an input example image (India_900498_S2Hand.tif) and place it in
15 - the same folder with the script (or specify with the --data_file argument)
16 -
17 -Run the example:
18 -python prithvi_geospatial_mae.py
19 -
20 -""" # noqa: E501
21 -
3 import argparse
4 import datetime
5 import os
6 +import re
7 from typing import Union
8
9 import albumentations
10 import numpy as np
11 import rasterio
30 -import regex as re
12 import torch
13 from einops import rearrange
14 from terratorch.datamodules import Sen1Floods11NonGeoDataModule
15
16 from vllm import LLM
17
18 +torch.set_default_dtype(torch.float16)
19 +
20 NO_DATA = -9999
21 NO_DATA_FLOAT = 0.0001
22 OFFSET = 0
23 PERCENTILE = 99
24
42 -model_config = """{
43 - "architectures": ["PrithviGeoSpatialMAE"],
44 - "num_classes": 0,
45 - "pretrained_cfg": {
46 - "task_args": {
47 - "task": "SemanticSegmentationTask",
48 - "model_factory": "EncoderDecoderFactory",
49 - "loss": "ce",
50 - "ignore_index": -1,
51 - "lr": 0.001,
52 - "freeze_backbone": false,
53 - "freeze_decoder": false,
54 - "plot_on_val": 10,
55 - "optimizer": "AdamW",
56 - "scheduler": "CosineAnnealingLR"
57 - },
58 - "model_args": {
59 - "backbone_pretrained": false,
60 - "backbone": "prithvi_eo_v2_300_tl",
61 - "decoder": "UperNetDecoder",
62 - "decoder_channels": 256,
63 - "decoder_scale_modules": true,
64 - "num_classes": 2,
65 - "rescale": true,
66 - "backbone_bands": [
67 - "BLUE",
68 - "GREEN",
69 - "RED",
70 - "NIR_NARROW",
71 - "SWIR_1",
72 - "SWIR_2"
73 - ],
74 - "head_dropout": 0.1,
75 - "necks": [
76 - {
77 - "name": "SelectIndices",
78 - "indices": [
79 - 5,
80 - 11,
81 - 17,
82 - 23
83 - ]
84 - },
85 - {
86 - "name": "ReshapeTokensToImage"
87 - }
88 - ]
89 - },
90 - "optimizer_params" : {
91 - "lr": 5.0e-05,
92 - "betas": [0.9, 0.999],
93 - "eps": [1.0e-08],
94 - "weight_decay": 0.05,
95 - "amsgrad": false,
96 - "maximize": false,
97 - "capturable": false,
98 - "differentiable": false
99 - },
100 - "scheduler_params" : {
101 - "T_max": 50,
102 - "eta_min": 0,
103 - "last_epoch": -1,
104 - "verbose": "deprecated"
105 - }
106 - },
107 -
108 -
109 - "torch_dtype": "float32"
110 -}
111 -"""
112 -
113 -# Temporarily creating the "config.json" for the model.
114 -# This is going to disappear once the correct config.json is available on HF
115 -with open(
116 - os.path.join(os.path.dirname(__file__), "./model/config.json"), "w"
117 -) as config_file:
118 - config_file.write(model_config)
119 -
25 datamodule_config = {
26 "bands": ["BLUE", "GREEN", "RED", "NIR_NARROW", "SWIR_1", "SWIR_2"],
27 "batch_size": 16,
@@ -138,28 +43,24 @@ datamodule_config = {
43
44
45 class PrithviMAE:
141 - def __init__(self):
142 - print("Initializing PrithviMAE model")
143 - self.llm = LLM(
144 - model=os.path.join(os.path.dirname(__file__), "./model"),
145 - skip_tokenizer_init=True,
146 - dtype="float32",
46 + def __init__(self, model):
47 + self.model = LLM(
48 + model=model, skip_tokenizer_init=True, dtype="float16", enforce_eager=True
49 )
50
51 def run(self, input_data, location_coords):
150 - print("################ Running inference on vLLM ##############")
52 # merge the inputs into one data structure
53 + if input_data is not None and input_data.dtype == torch.float32:
54 + input_data = input_data.to(torch.float16)
55 + input_data = input_data[0]
56 +
57 mm_data = {
153 - "pixel_values": torch.empty(0) if input_data is None else input_data,
154 - "location_coords": torch.empty(0)
155 - if location_coords is None
156 - else location_coords,
58 + "pixel_values": input_data,
59 + "location_coords": location_coords,
60 }
61
62 prompt = {"prompt_token_ids": [1], "multi_modal_data": mm_data}
160 -
161 - outputs = self.llm.encode(prompt, use_tqdm=False)
162 - print("################ Inference done (it took seconds) ##############")
63 + outputs = self.model.encode(prompt, use_tqdm=False)
64
65 return outputs[0].outputs.data
66
@@ -181,11 +82,12 @@ def process_channel_group(orig_img, channels):
82 """
83 Args:
84 orig_img: torch.Tensor representing original image (reference)
184 - with shape = (bands, H, W).
85 + with shape = (bands, H, W).
86 channels: list of indices representing RGB channels.
87
88 Returns:
188 - torch.Tensor with shape (num_channels, height, width) for original image
89 + torch.Tensor with shape (num_channels, height, width)
90 + for original image
91 """
92
93 orig_img = orig_img[channels, ...]
@@ -260,10 +162,10 @@ def load_example(
162
163 Args:
164 file_paths: list of file paths .
263 - mean: list containing mean values for each band in the images
264 - in *file_paths*.
265 - std: list containing std values for each band in the images
266 - in *file_paths*.
165 + mean: list containing mean values for each band in the
166 + images in *file_paths*.
167 + std: list containing std values for each band in the
168 + images in *file_paths*.
169
170 Returns:
171 np.array containing created example
@@ -308,7 +210,7 @@ def load_example(
210 print(f"Could not extract timestamp for {file} ({e})")
211
212 imgs = np.stack(imgs, axis=0) # num_frames, H, W, C
311 - imgs = np.moveaxis(imgs, -1, 0).astype("float32")
213 + imgs = np.moveaxis(imgs, -1, 0).astype("float32") # C, num_frames, H, W
214 imgs = np.expand_dims(imgs, axis=0) # add batch di
215
216 return imgs, temporal_coords, location_coords, metas
@@ -332,8 +234,10 @@ def run_model(
234 )
235
236 # Build sliding window
237 +
238 batch_size = 1
336 - batch = torch.tensor(input_data, device="cpu")
239 + # batch = torch.tensor(input_data, device="cpu")
240 + batch = torch.tensor(input_data)
241 windows = batch.unfold(3, img_size, img_size).unfold(4, img_size, img_size)
242 h1, w1 = windows.shape[3:5]
243 windows = rearrange(
@@ -344,18 +248,16 @@ def run_model(
248 num_batches = windows.shape[0] // batch_size if windows.shape[0] > batch_size else 1
249 windows = torch.tensor_split(windows, num_batches, dim=0)
250
347 - device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
348 -
251 if temporal_coords:
350 - temporal_coords = torch.tensor(temporal_coords, device=device).unsqueeze(0)
252 + temporal_coords = torch.tensor(temporal_coords).unsqueeze(0)
253 else:
254 temporal_coords = None
255 if location_coords:
354 - location_coords = torch.tensor(location_coords[0], device=device).unsqueeze(0)
256 + location_coords = torch.tensor(location_coords[0]).unsqueeze(0)
257 else:
258 location_coords = None
259
358 - # Run model
260 + # Run Prithvi-EO-V2-300M-TL-Sen1Floods11
261 pred_imgs = []
262 for x in windows:
263 # Apply standardization
@@ -363,15 +265,7 @@ def run_model(
265 x = datamodule.aug(x)["image"]
266
267 with torch.no_grad():
366 - x = x.to(device)
268 pred = model.run(x, location_coords=location_coords)
368 - if lightning_model:
369 - pred_lightning = lightning_model(
370 - x, temporal_coords=temporal_coords, location_coords=location_coords
371 - )
372 - pred_lightning = pred_lightning.output.detach().cpu()
373 - if not torch.equal(pred, pred_lightning):
374 - print("Inference output is not equal")
269 y_hat = pred.argmax(dim=1)
270
271 y_hat = torch.nn.functional.interpolate(
@@ -403,52 +297,18 @@ def run_model(
297 return pred_imgs
298
299
406 -def parse_args():
407 - parser = argparse.ArgumentParser("MAE run inference", add_help=False)
408 -
409 - parser.add_argument(
410 - "--data_file",
411 - type=str,
412 - default="./India_900498_S2Hand.tif",
413 - help="Path to the file.",
414 - )
415 - parser.add_argument(
416 - "--output_dir",
417 - type=str,
418 - default="output",
419 - help="Path to the directory where to save outputs.",
420 - )
421 - parser.add_argument(
422 - "--input_indices",
423 - default=[1, 2, 3, 8, 11, 12],
424 - type=int,
425 - nargs="+",
426 - help="0-based indices of the six Prithvi channels to be selected from the "
427 - "input. By default selects [1,2,3,8,11,12] for S2L1C data.",
428 - )
429 - parser.add_argument(
430 - "--rgb_outputs",
431 - action="store_true",
432 - help="If present, output files will only contain RGB channels. "
433 - "Otherwise, all bands will be saved.",
434 - )
435 -
436 -
300 def main(
301 data_file: str,
302 + model: str,
303 output_dir: str,
304 rgb_outputs: bool,
305 input_indices: list[int] = None,
306 ):
307 os.makedirs(output_dir, exist_ok=True)
308
445 - # Load model ---------------------------------------------------------------
446 -
447 - model_obj = PrithviMAE()
309 + model_obj = PrithviMAE(model=model)
310 datamodule = generate_datamodule()
449 - img_size = 256 # Size of Sen1Floods11
450 -
451 - # Loading data -------------------------------------------------------------
311 + img_size = 512 # Size of Sen1Floods11
312
313 input_data, temporal_coords, location_coords, meta_data = load_example(
314 file_paths=[data_file],
@@ -460,8 +320,6 @@ def main(
320 if input_data.mean() > 1:
321 input_data = input_data / 10000 # Convert to range 0-1
322
463 - # Running model ------------------------------------------------------------
464 -
323 channels = [
324 datamodule_config["bands"].index(b) for b in ["RED", "GREEN", "BLUE"]
325 ] # BGR -> RGB
@@ -469,7 +327,6 @@ def main(
327 pred = run_model(
328 input_data, temporal_coords, location_coords, model_obj, datamodule, img_size
329 )
472 -
330 # Save pred
331 meta_data.update(count=1, dtype="uint8", compress="lzw", nodata=0)
332 pred_file = os.path.join(
@@ -487,6 +344,7 @@ def main(
344 orig_img=torch.Tensor(input_data[0, :, 0, ...]),
345 channels=channels,
346 )
347 + rgb_orig = rgb_orig.to(torch.float32)
348
349 pred[pred == 0.0] = np.nan
350 img_pred = rgb_orig * 0.7 + pred * 0.3
@@ -503,9 +361,10 @@ def main(
361
362 # Save image rgb
363 if rgb_outputs:
364 + name_suffix = os.path.splitext(os.path.basename(data_file))[0]
365 rgb_file = os.path.join(
366 output_dir,
508 - f"original_rgb_{os.path.splitext(os.path.basename(data_file))[0]}.tiff",
367 + f"original_rgb_{name_suffix}.tiff",
368 )
369 save_geotiff(
370 image=_convert_np_uint8(rgb_orig),
@@ -515,6 +374,42 @@ def main(
374
375
376 if __name__ == "__main__":
518 - args = parse_args()
377 + parser = argparse.ArgumentParser("MAE run inference", add_help=False)
378 +
379 + parser.add_argument(
380 + "--data_file",
381 + type=str,
382 + default="./India_900498_S2Hand.tif",
383 + help="Path to the file.",
384 + )
385 + parser.add_argument(
386 + "--model",
387 + type=str,
388 + default="christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM",
389 + help="Path to a checkpoint file to load from.",
390 + )
391 + parser.add_argument(
392 + "--output_dir",
393 + type=str,
394 + default="output",
395 + help="Path to the directory where to save outputs.",
396 + )
397 + parser.add_argument(
398 + "--input_indices",
399 + default=[1, 2, 3, 8, 11, 12],
400 + type=int,
401 + nargs="+",
402 + help="""
403 + 0-based indices of the six Prithvi channels to be selected from the input.
404 + By default selects [1,2,3,8,11,12] for S2L1C data.
405 + """,
406 + )
407 + parser.add_argument(
408 + "--rgb_outputs",
409 + action="store_true",
410 + help="If present, output files will only contain RGB channels. "
411 + "Otherwise, all bands will be saved.",
412 + )
413 + args = parser.parse_args()
414
415 main(**vars(args))
requirements/test.in
+1
@@ -54,3 +54,4 @@ runai-model-streamer==0.11.0
54 runai-model-streamer-s3==0.11.0
55 fastsafetensors>=0.1.10
56 pydantic>=2.10 # 2.9 leads to error on python 3.10
57 +terratorch==1.1rc2 # required for PrithviMAE test
\ No newline at end of file
requirements/test.txt
+360 -14
@@ -6,6 +6,10 @@ accelerate==1.0.1
6 # via
7 # lm-eval
8 # peft
9 +aenum==3.1.16
10 + # via lightly
11 +affine==2.4.0
12 + # via rasterio
13 aiohappyeyeballs==2.4.3
14 # via aiohttp
15 aiohttp==3.10.11
@@ -21,8 +25,18 @@ aiosignal==1.3.1
25 # via
26 # aiohttp
27 # ray
28 +albucore==0.0.16
29 + # via terratorch
30 +albumentations==1.4.6
31 + # via terratorch
32 +alembic==1.16.4
33 + # via mlflow
34 annotated-types==0.7.0
35 # via pydantic
36 +antlr4-python3-runtime==4.9.3
37 + # via
38 + # hydra-core
39 + # omegaconf
40 anyio==4.6.2.post1
41 # via
42 # httpx
@@ -34,10 +48,12 @@ arrow==1.3.0
48 attrs==24.2.0
49 # via
50 # aiohttp
51 + # fiona
52 # hypothesis
53 # jsonlines
54 # jsonschema
55 # pytest-subtests
56 + # rasterio
57 # referencing
58 audioread==3.0.1
59 # via librosa
@@ -46,9 +62,13 @@ backoff==2.2.1
62 # -r requirements/test.in
63 # schemathesis
64 bitsandbytes==0.46.1
49 - # via -r requirements/test.in
65 + # via
66 + # -r requirements/test.in
67 + # lightning
68 black==24.10.0
69 # via datamodel-code-generator
70 +blinker==1.9.0
71 + # via flask
72 blobfile==3.0.0
73 # via -r requirements/test.in
74 bm25s==0.2.13
@@ -64,11 +84,18 @@ bounded-pool-executor==0.0.3
84 buildkite-test-collector==0.1.9
85 # via -r requirements/test.in
86 cachetools==5.5.2
67 - # via google-auth
87 + # via
88 + # google-auth
89 + # mlflow-skinny
90 certifi==2024.8.30
91 # via
92 + # fiona
93 # httpcore
94 # httpx
95 + # lightly
96 + # pyogrio
97 + # pyproj
98 + # rasterio
99 # requests
100 cffi==1.17.1
101 # via soundfile
@@ -79,11 +106,28 @@ charset-normalizer==3.4.0
106 click==8.1.7
107 # via
108 # black
109 + # click-plugins
110 + # cligj
111 + # fiona
112 + # flask
113 # jiwer
114 + # mlflow-skinny
115 # nltk
116 + # rasterio
117 # ray
118 # schemathesis
119 # typer
120 + # uvicorn
121 +click-plugins==1.1.1.2
122 + # via
123 + # fiona
124 + # rasterio
125 +cligj==0.7.2
126 + # via
127 + # fiona
128 + # rasterio
129 +cloudpickle==3.1.1
130 + # via mlflow-skinny
131 colorama==0.4.6
132 # via
133 # sacrebleu
@@ -99,6 +143,8 @@ cupy-cuda12x==13.3.0
143 # via ray
144 cycler==0.12.1
145 # via matplotlib
146 +databricks-sdk==0.59.0
147 + # via mlflow-skinny
148 datamodel-code-generator==0.26.3
149 # via -r requirements/test.in
150 dataproperty==1.0.1
@@ -122,13 +168,21 @@ distlib==0.3.9
168 # via virtualenv
169 dnspython==2.7.0
170 # via email-validator
171 +docker==7.1.0
172 + # via mlflow
173 docopt==0.6.2
174 # via num2words
127 -einops==0.8.0
175 +docstring-parser==0.17.0
176 + # via jsonargparse
177 +efficientnet-pytorch==0.7.1
178 + # via segmentation-models-pytorch
179 +einops==0.8.1
180 # via
181 # -r requirements/test.in
182 # encodec
183 # mamba-ssm
184 + # terratorch
185 + # torchgeo
186 # vector-quantize-pytorch
187 # vocos
188 einx==0.3.0
@@ -141,6 +195,8 @@ eval-type-backport==0.2.2
195 # via mteb
196 evaluate==0.4.3
197 # via lm-eval
198 +fastapi==0.116.1
199 + # via mlflow-skinny
200 fastparquet==2024.11.0
201 # via genai-perf
202 fastrlock==0.8.2
@@ -156,6 +212,10 @@ filelock==3.16.1
212 # torch
213 # transformers
214 # virtualenv
215 +fiona==1.10.1
216 + # via torchgeo
217 +flask==3.1.1
218 + # via mlflow
219 fonttools==4.54.1
220 # via matplotlib
221 fqdn==1.5.1
@@ -173,6 +233,8 @@ fsspec==2024.9.0
233 # evaluate
234 # fastparquet
235 # huggingface-hub
236 + # lightning
237 + # pytorch-lightning
238 # torch
239 ftfy==6.3.1
240 # via open-clip-torch
@@ -180,18 +242,41 @@ genai-perf==0.0.8
242 # via -r requirements/test.in
243 genson==1.3.0
244 # via datamodel-code-generator
245 +geopandas==1.0.1
246 + # via terratorch
247 +gitdb==4.0.12
248 + # via gitpython
249 +gitpython==3.1.44
250 + # via mlflow-skinny
251 google-api-core==2.24.2
252 # via opencensus
253 google-auth==2.40.2
186 - # via google-api-core
254 + # via
255 + # databricks-sdk
256 + # google-api-core
257 googleapis-common-protos==1.70.0
258 # via google-api-core
259 +graphene==3.4.3
260 + # via mlflow
261 graphql-core==3.2.6
190 - # via hypothesis-graphql
262 + # via
263 + # graphene
264 + # graphql-relay
265 + # hypothesis-graphql
266 +graphql-relay==3.2.0
267 + # via graphene
268 +greenlet==3.2.3
269 + # via sqlalchemy
270 grpcio==1.71.0
271 # via ray
272 +gunicorn==23.0.0
273 + # via mlflow
274 h11==0.14.0
194 - # via httpcore
275 + # via
276 + # httpcore
277 + # uvicorn
278 +h5py==3.13.0
279 + # via terratorch
280 harfile==0.3.0
281 # via schemathesis
282 hf-xet==1.1.3
@@ -204,7 +289,7 @@ httpx==0.27.2
289 # via
290 # -r requirements/test.in
291 # schemathesis
207 -huggingface-hub==0.33.0
292 +huggingface-hub==0.33.1
293 # via
294 # -r requirements/test.in
295 # accelerate
@@ -212,13 +297,19 @@ huggingface-hub==0.33.0
297 # evaluate
298 # open-clip-torch
299 # peft
300 + # segmentation-models-pytorch
301 # sentence-transformers
302 + # terratorch
303 # timm
304 # tokenizers
305 # transformers
306 # vocos
307 humanize==4.11.0
308 # via runai-model-streamer
309 +hydra-core==1.3.2
310 + # via
311 + # lightly
312 + # lightning
313 hypothesis==6.131.0
314 # via
315 # hypothesis-graphql
@@ -236,6 +327,14 @@ idna==3.10
327 # jsonschema
328 # requests
329 # yarl
330 +imageio==2.37.0
331 + # via scikit-image
332 +importlib-metadata==8.7.0
333 + # via
334 + # mlflow-skinny
335 + # opentelemetry-api
336 +importlib-resources==6.5.2
337 + # via typeshed-client
338 inflect==5.6.2
339 # via datamodel-code-generator
340 iniconfig==2.0.0
@@ -244,9 +343,13 @@ isoduration==20.11.0
343 # via jsonschema
344 isort==5.13.2
345 # via datamodel-code-generator
346 +itsdangerous==2.2.0
347 + # via flask
348 jinja2==3.1.6
349 # via
350 # datamodel-code-generator
351 + # flask
352 + # mlflow
353 # torch
354 jiwer==3.0.5
355 # via -r requirements/test.in
@@ -259,6 +362,10 @@ joblib==1.4.2
362 # librosa
363 # nltk
364 # scikit-learn
365 +jsonargparse==4.35.0
366 + # via
367 + # lightning
368 + # terratorch
369 jsonlines==4.0.0
370 # via lm-eval
371 jsonpointer==3.0.0
@@ -277,12 +384,33 @@ kaleido==0.2.1
384 # via genai-perf
385 kiwisolver==1.4.7
386 # via matplotlib
387 +kornia==0.8.1
388 + # via torchgeo
389 +kornia-rs==0.1.9
390 + # via kornia
391 lazy-loader==0.4
281 - # via librosa
392 + # via
393 + # librosa
394 + # scikit-image
395 libnacl==2.1.0
396 # via tensorizer
397 librosa==0.10.2.post1
398 # via -r requirements/test.in
399 +lightly==1.5.20
400 + # via
401 + # terratorch
402 + # torchgeo
403 +lightly-utils==0.0.2
404 + # via lightly
405 +lightning==2.5.1.post0
406 + # via
407 + # terratorch
408 + # torchgeo
409 +lightning-utilities==0.14.3
410 + # via
411 + # lightning
412 + # pytorch-lightning
413 + # torchmetrics
414 llvmlite==0.44.0
415 # via numba
416 lm-eval==0.4.8
@@ -291,16 +419,27 @@ lxml==5.3.0
419 # via
420 # blobfile
421 # sacrebleu
422 +mako==1.3.10
423 + # via alembic
424 mamba-ssm==2.2.4
425 # via -r requirements/test.in
426 +markdown==3.8.2
427 + # via mlflow
428 markdown-it-py==3.0.0
429 # via rich
430 markupsafe==3.0.1
431 # via
432 + # flask
433 # jinja2
434 + # mako
435 # werkzeug
436 matplotlib==3.9.2
303 - # via -r requirements/test.in
437 + # via
438 + # -r requirements/test.in
439 + # lightning
440 + # mlflow
441 + # pycocotools
442 + # torchgeo
443 mbstrdecoder==1.1.3
444 # via
445 # dataproperty
@@ -310,6 +449,10 @@ mdurl==0.1.2
449 # via markdown-it-py
450 mistral-common==1.8.0
451 # via -r requirements/test.in
452 +mlflow==2.22.0
453 + # via terratorch
454 +mlflow-skinny==2.22.0
455 + # via mlflow
456 more-itertools==10.5.0
457 # via lm-eval
458 mpmath==1.3.0
@@ -328,10 +471,14 @@ multiprocess==0.70.16
471 # via
472 # datasets
473 # evaluate
474 +munch==4.0.0
475 + # via pretrainedmodels
476 mypy-extensions==1.0.0
477 # via black
478 networkx==3.2.1
334 - # via torch
479 + # via
480 + # scikit-image
481 + # torch
482 ninja==1.11.1.3
483 # via mamba-ssm
484 nltk==3.9.1
@@ -348,6 +495,8 @@ numpy==1.26.4
495 # via
496 # -r requirements/test.in
497 # accelerate
498 + # albucore
499 + # albumentations
500 # bitsandbytes
501 # bm25s
502 # contourpy
@@ -358,9 +507,15 @@ numpy==1.26.4
507 # evaluate
508 # fastparquet
509 # genai-perf
510 + # geopandas
511 + # h5py
512 + # imageio
513 # librosa
514 + # lightly
515 + # lightly-utils
516 # matplotlib
517 # mistral-common
518 + # mlflow
519 # mteb
520 # numba
521 # numexpr
@@ -368,18 +523,30 @@ numpy==1.26.4
523 # pandas
524 # patsy
525 # peft
526 + # pycocotools
527 + # pyogrio
528 + # rasterio
529 + # rioxarray
530 # rouge-score
531 # runai-model-streamer
532 # sacrebleu
533 + # scikit-image
534 # scikit-learn
535 # scipy
536 + # segmentation-models-pytorch
537 + # shapely
538 # soxr
539 # statsmodels
540 + # tensorboardx
541 # tensorizer
542 + # tifffile
543 + # torchgeo
544 + # torchmetrics
545 # torchvision
546 # transformers
547 # tritonclient
548 # vocos
549 + # xarray
550 nvidia-cublas-cu12==12.8.3.14
551 # via
552 # nvidia-cudnn-cu12
@@ -417,6 +584,10 @@ nvidia-nvjitlink-cu12==12.8.61
584 # torch
585 nvidia-nvtx-cu12==12.8.55
586 # via torch
587 +omegaconf==2.3.0
588 + # via
589 + # hydra-core
590 + # lightning
591 open-clip-torch==2.32.0
592 # via -r requirements/test.in
593 opencensus==0.11.4
@@ -426,7 +597,18 @@ opencensus-context==0.1.3
597 opencv-python-headless==4.11.0.86
598 # via
599 # -r requirements/test.in
600 + # albucore
601 + # albumentations
602 # mistral-common
603 +opentelemetry-api==1.35.0
604 + # via
605 + # mlflow-skinny
606 + # opentelemetry-sdk
607 + # opentelemetry-semantic-conventions
608 +opentelemetry-sdk==1.35.0
609 + # via mlflow-skinny
610 +opentelemetry-semantic-conventions==0.56b0
611 + # via opentelemetry-sdk
612 packaging==24.2
613 # via
614 # accelerate
@@ -435,26 +617,44 @@ packaging==24.2
617 # datasets
618 # evaluate
619 # fastparquet
620 + # geopandas
621 + # gunicorn
622 # huggingface-hub
623 + # hydra-core
624 + # kornia
625 # lazy-loader
626 + # lightning
627 + # lightning-utilities
628 # mamba-ssm
629 # matplotlib
630 + # mlflow-skinny
631 # peft
632 # plotly
633 # pooch
634 + # pyogrio
635 # pytest
636 # pytest-rerunfailures
637 + # pytorch-lightning
638 # ray
639 + # rioxarray
640 + # scikit-image
641 # statsmodels
642 + # tensorboardx
643 + # torchmetrics
644 # transformers
645 # typepy
646 + # xarray
647 pandas==2.2.3
648 # via
649 # datasets
650 # evaluate
651 # fastparquet
652 # genai-perf
653 + # geopandas
654 + # mlflow
655 # statsmodels
656 + # torchgeo
657 + # xarray
658 pathspec==0.12.1
659 # via black
660 pathvalidate==3.2.1
@@ -468,9 +668,14 @@ peft==0.13.2
668 pillow==10.4.0
669 # via
670 # genai-perf
671 + # imageio
672 + # lightly-utils
673 # matplotlib
674 # mistral-common
675 + # scikit-image
676 + # segmentation-models-pytorch
677 # sentence-transformers
678 + # torchgeo
679 # torchvision
680 platformdirs==4.3.6
681 # via
@@ -489,6 +694,8 @@ portalocker==2.10.1
694 # via sacrebleu
695 pqdm==0.2.0
696 # via -r requirements/test.in
697 +pretrainedmodels==0.7.4
698 + # via segmentation-models-pytorch
699 prometheus-client==0.22.0
700 # via ray
701 propcache==0.2.0
@@ -499,8 +706,10 @@ protobuf==5.28.3
706 # via
707 # google-api-core
708 # googleapis-common-protos
709 + # mlflow-skinny
710 # proto-plus
711 # ray
712 + # tensorboardx
713 # tensorizer
714 psutil==6.1.0
715 # via
@@ -515,6 +724,7 @@ pyarrow==18.0.0
724 # via
725 # datasets
726 # genai-perf
727 + # mlflow
728 pyasn1==0.6.1
729 # via
730 # pyasn1-modules
@@ -523,6 +733,8 @@ pyasn1-modules==0.4.2
733 # via google-auth
734 pybind11==2.13.6
735 # via lm-eval
736 +pycocotools==2.0.8
737 + # via terratorch
738 pycountry==24.6.1
739 # via pydantic-extra-types
740 pycparser==2.22
@@ -532,8 +744,12 @@ pycryptodomex==3.22.0
744 pydantic==2.11.5
745 # via
746 # -r requirements/test.in
747 + # albumentations
748 # datamodel-code-generator
749 + # fastapi
750 + # lightly
751 # mistral-common
752 + # mlflow-skinny
753 # mteb
754 # pydantic-extra-types
755 # ray
@@ -543,15 +759,24 @@ pydantic-extra-types==2.10.5
759 # via mistral-common
760 pygments==2.18.0
761 # via rich
762 +pyogrio==0.11.0
763 + # via geopandas
764 pyparsing==3.2.0
547 - # via matplotlib
765 + # via
766 + # matplotlib
767 + # rasterio
768 +pyproj==3.7.1
769 + # via
770 + # geopandas
771 + # rioxarray
772 + # torchgeo
773 pyrate-limiter==3.7.0
774 # via schemathesis
775 pystemmer==3.0.0
776 # via mteb
777 pytablewriter==1.2.0
778 # via lm-eval
554 -pytest==8.3.3
779 +pytest==8.3.5
780 # via
781 # -r requirements/test.in
782 # buildkite-test-collector
@@ -564,6 +789,7 @@ pytest==8.3.3
789 # pytest-subtests
790 # pytest-timeout
791 # schemathesis
792 + # terratorch
793 pytest-asyncio==0.24.0
794 # via -r requirements/test.in
795 pytest-forked==1.6.0
@@ -578,15 +804,23 @@ pytest-subtests==0.14.1
804 # via schemathesis
805 pytest-timeout==2.3.1
806 # via -r requirements/test.in
807 +python-box==7.3.2
808 + # via terratorch
809 python-dateutil==2.9.0.post0
810 # via
811 # arrow
812 # botocore
813 + # graphene
814 + # lightly
815 # matplotlib
816 # pandas
817 # typepy
818 python-rapidjson==1.20
819 # via tritonclient
820 +pytorch-lightning==2.5.2
821 + # via
822 + # lightly
823 + # lightning
824 pytrec-eval-terrier==0.5.7
825 # via mteb
826 pytz==2024.2
@@ -596,11 +830,17 @@ pytz==2024.2
830 pyyaml==6.0.2
831 # via
832 # accelerate
833 + # albumentations
834 # datamodel-code-generator
835 # datasets
836 # genai-perf
837 # huggingface-hub
838 + # jsonargparse
839 + # lightning
840 + # mlflow-skinny
841 + # omegaconf
842 # peft
843 + # pytorch-lightning
844 # ray
845 # responses
846 # schemathesis
@@ -609,6 +849,11 @@ pyyaml==6.0.2
849 # vocos
850 rapidfuzz==3.12.1
851 # via jiwer
852 +rasterio==1.4.3
853 + # via
854 + # rioxarray
855 + # terratorch
856 + # torchgeo
857 ray==2.43.0
858 # via -r requirements/test.in
859 redis==5.2.0
@@ -627,12 +872,16 @@ regex==2024.9.11
872 requests==2.32.3
873 # via
874 # buildkite-test-collector
875 + # databricks-sdk
876 # datasets
877 + # docker
878 # evaluate
879 # google-api-core
880 # huggingface-hub
881 + # lightly
882 # lm-eval
883 # mistral-common
884 + # mlflow-skinny
885 # mteb
886 # pooch
887 # ray
@@ -650,8 +899,11 @@ rfc3987==1.3.8
899 rich==13.9.4
900 # via
901 # genai-perf
902 + # lightning
903 # mteb
904 # typer
905 +rioxarray==0.19.0
906 + # via terratorch
907 rouge-score==0.1.2
908 # via lm-eval
909 rpds-py==0.20.1
@@ -660,6 +912,8 @@ rpds-py==0.20.1
912 # referencing
913 rsa==4.9.1
914 # via google-auth
915 +rtree==1.4.0
916 + # via torchgeo
917 runai-model-streamer==0.11.0
918 # via -r requirements/test.in
919 runai-model-streamer-s3==0.11.0
@@ -677,21 +931,32 @@ safetensors==0.4.5
931 # transformers
932 schemathesis==3.39.15
933 # via -r requirements/test.in
934 +scikit-image==0.25.2
935 + # via albumentations
936 scikit-learn==1.5.2
937 # via
938 + # albumentations
939 # librosa
940 # lm-eval
941 + # mlflow
942 # mteb
943 # sentence-transformers
944 scipy==1.13.1
945 # via
946 + # albumentations
947 # bm25s
948 # librosa
949 + # mlflow
950 # mteb
951 + # scikit-image
952 # scikit-learn
953 # sentence-transformers
954 # statsmodels
955 # vocos
956 +segmentation-models-pytorch==0.4.0
957 + # via
958 + # terratorch
959 + # torchgeo
960 sentence-transformers==3.2.1
961 # via
962 # -r requirements/test.in
@@ -700,21 +965,30 @@ sentencepiece==0.2.0
965 # via mistral-common
966 setuptools==77.0.3
967 # via
968 + # lightning-utilities
969 # mamba-ssm
970 # pytablewriter
971 # torch
972 # triton
973 +shapely==2.1.1
974 + # via
975 + # geopandas
976 + # torchgeo
977 shellingham==1.5.4
978 # via typer
979 six==1.16.0
980 # via
981 # junit-xml
982 + # lightly
983 # opencensus
984 # python-dateutil
985 # rfc3339-validator
986 # rouge-score
987 + # segmentation-models-pytorch
988 smart-open==7.1.0
989 # via ray
990 +smmap==5.0.2
991 + # via gitdb
992 sniffio==1.3.1
993 # via
994 # anyio
@@ -727,10 +1001,17 @@ soundfile==0.12.1
1001 # librosa
1002 soxr==0.5.0.post1
1003 # via librosa
1004 +sqlalchemy==2.0.41
1005 + # via
1006 + # alembic
1007 + # mlflow
1008 sqlitedict==2.1.0
1009 # via lm-eval
1010 +sqlparse==0.5.3
1011 + # via mlflow-skinny
1012 starlette==0.46.2
1013 # via
1014 + # fastapi
1015 # schemathesis
1016 # starlette-testclient
1017 starlette-testclient==0.4.1
@@ -751,18 +1032,29 @@ tenacity==9.0.0
1032 # via
1033 # lm-eval
1034 # plotly
1035 +tensorboardx==2.6.4
1036 + # via lightning
1037 tensorizer==2.10.1
1038 # via -r requirements/test.in
1039 +terratorch==1.1rc2
1040 + # via -r requirements/test.in
1041 threadpoolctl==3.5.0
1042 # via scikit-learn
1043 +tifffile==2025.3.30
1044 + # via
1045 + # scikit-image
1046 + # terratorch
1047 tiktoken==0.7.0
1048 # via
1049 # lm-eval
1050 # mistral-common
762 -timm==1.0.11
1051 +timm==1.0.15
1052 # via
1053 # -r requirements/test.in
1054 # open-clip-torch
1055 + # segmentation-models-pytorch
1056 + # terratorch
1057 + # torchgeo
1058 tokenizers==0.21.1
1059 # via
1060 # -r requirements/test.in
@@ -776,18 +1068,28 @@ torch==2.7.1+cu128
1068 # -r requirements/test.in
1069 # accelerate
1070 # bitsandbytes
1071 + # efficientnet-pytorch
1072 # encodec
1073 # fastsafetensors
1074 + # kornia
1075 + # lightly
1076 + # lightning
1077 # lm-eval
1078 # mamba-ssm
1079 # mteb
1080 # open-clip-torch
1081 # peft
1082 + # pretrainedmodels
1083 + # pytorch-lightning
1084 # runai-model-streamer
1085 + # segmentation-models-pytorch
1086 # sentence-transformers
1087 # tensorizer
1088 + # terratorch
1089 # timm
1090 # torchaudio
1091 + # torchgeo
1092 + # torchmetrics
1093 # torchvision
1094 # vector-quantize-pytorch
1095 # vocos
@@ -796,22 +1098,40 @@ torchaudio==2.7.1+cu128
1098 # -r requirements/test.in
1099 # encodec
1100 # vocos
1101 +torchgeo==0.7.0
1102 + # via terratorch
1103 +torchmetrics==1.7.4
1104 + # via
1105 + # lightning
1106 + # pytorch-lightning
1107 + # terratorch
1108 + # torchgeo
1109 torchvision==0.22.1+cu128
1110 # via
1111 # -r requirements/test.in
1112 + # lightly
1113 # open-clip-torch
1114 + # pretrainedmodels
1115 + # segmentation-models-pytorch
1116 + # terratorch
1117 # timm
1118 + # torchgeo
1119 tqdm==4.66.6
1120 # via
1121 # datasets
1122 # evaluate
1123 # huggingface-hub
1124 + # lightly
1125 + # lightning
1126 # lm-eval
1127 # mteb
1128 # nltk
1129 # open-clip-torch
1130 # peft
1131 # pqdm
1132 + # pretrainedmodels
1133 + # pytorch-lightning
1134 + # segmentation-models-pytorch
1135 # sentence-transformers
1136 # tqdm-multiprocess
1137 # transformers
@@ -843,18 +1163,34 @@ typer==0.15.2
1163 # via fastsafetensors
1164 types-python-dateutil==2.9.0.20241206
1165 # via arrow
1166 +typeshed-client==2.8.2
1167 + # via jsonargparse
1168 typing-extensions==4.12.2
1169 # via
1170 + # albumentations
1171 + # alembic
1172 + # fastapi
1173 + # graphene
1174 # huggingface-hub
1175 # librosa
1176 + # lightning
1177 + # lightning-utilities
1178 # mistral-common
1179 + # mlflow-skinny
1180 # mteb
1181 + # opentelemetry-api
1182 + # opentelemetry-sdk
1183 + # opentelemetry-semantic-conventions
1184 # pqdm
1185 # pydantic
1186 # pydantic-core
1187 # pydantic-extra-types
1188 + # pytorch-lightning
1189 + # sqlalchemy
1190 # torch
1191 + # torchgeo
1192 # typer
1193 + # typeshed-client
1194 # typing-inspection
1195 typing-inspection==0.4.1
1196 # via pydantic
@@ -866,9 +1202,13 @@ urllib3==2.2.3
1202 # via
1203 # blobfile
1204 # botocore
1205 + # docker
1206 + # lightly
1207 # requests
1208 # responses
1209 # tritonclient
1210 +uvicorn==0.35.0
1211 + # via mlflow-skinny
1212 vector-quantize-pytorch==1.21.2
1213 # via -r requirements/test.in
1214 virtualenv==20.31.2
@@ -880,11 +1220,15 @@ wcwidth==0.2.13
1220 webcolors==24.11.1
1221 # via jsonschema
1222 werkzeug==3.1.3
883 - # via schemathesis
1223 + # via
1224 + # flask
1225 + # schemathesis
1226 word2number==1.1
1227 # via lm-eval
1228 wrapt==1.17.2
1229 # via smart-open
1230 +xarray==2025.7.1
1231 + # via rioxarray
1232 xxhash==3.5.0
1233 # via
1234 # datasets
@@ -893,5 +1237,7 @@ yarl==1.17.1
1237 # via
1238 # aiohttp
1239 # schemathesis
1240 +zipp==3.23.0
1241 + # via importlib-metadata
1242 zstandard==0.23.0
1243 # via lm-eval
tests/models/multimodal/pooling/test_prithvi_mae.py new
+63
@@ -0,0 +1,63 @@
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 vllm.utils import set_default_torch_num_threads
8 +
9 +from ....conftest import VllmRunner
10 +
11 +
12 +def generate_test_mm_data():
13 + mm_data = {
14 + "pixel_values": torch.full((6, 512, 512), 1.0, dtype=torch.float16),
15 + "location_coords": torch.full((1, 2), 1.0, dtype=torch.float16),
16 + }
17 + return mm_data
18 +
19 +
20 +def _run_test(
21 + vllm_runner: type[VllmRunner],
22 + model: str,
23 +) -> None:
24 +
25 + prompt = [
26 + {
27 + # This model deals with no text input
28 + "prompt_token_ids": [1],
29 + "multi_modal_data": generate_test_mm_data(),
30 + } for _ in range(10)
31 + ]
32 +
33 + with (
34 + set_default_torch_num_threads(1),
35 + vllm_runner(
36 + model,
37 + task="embed",
38 + dtype=torch.float16,
39 + enforce_eager=True,
40 + skip_tokenizer_init=True,
41 + # Limit the maximum number of sequences to avoid the
42 + # test going OOM during the warmup run
43 + max_num_seqs=32,
44 + ) as vllm_model,
45 + ):
46 + vllm_model.encode(prompt)
47 +
48 +
49 +MODELS = ["christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM"]
50 +
51 +
52 +@pytest.mark.core_model
53 +@pytest.mark.parametrize("model", MODELS)
54 +def test_models_image(
55 + hf_runner,
56 + vllm_runner,
57 + image_assets,
58 + model: str,
59 +) -> None:
60 + _run_test(
61 + vllm_runner,
62 + model,
63 + )
vllm/config.py
+4 -2
@@ -651,6 +651,8 @@ class ModelConfig:
651 self.original_max_model_len = self.max_model_len
652 self.max_model_len = self.get_and_verify_max_len(self.max_model_len)
653 self.multimodal_config = self._init_multimodal_config()
654 + self.model_supports_multimodal_raw_input = (
655 + self.registry.supports_multimodal_raw_input(self.architectures))
656 if not self.skip_tokenizer_init:
657 self._verify_tokenizer_mode()
658
@@ -1243,10 +1245,10 @@ class ModelConfig:
1245 return self.get_hf_config_sliding_window()
1246
1247 def get_vocab_size(self) -> int:
1246 - return self.hf_text_config.vocab_size
1248 + return getattr(self.hf_text_config, "vocab_size", 0)
1249
1250 def get_hidden_size(self) -> int:
1249 - return self.hf_text_config.hidden_size
1251 + return getattr(self.hf_text_config, "hidden_size", 0)
1252
1253 @property
1254 def is_deepseek_mla(self) -> bool:
vllm/engine/llm_engine.py
+5 -5
@@ -238,14 +238,14 @@ class LLMEngine:
238 self.log_stats = log_stats
239 self.use_cached_outputs = use_cached_outputs
240
241 - if not self.model_config.skip_tokenizer_init:
242 - self.tokenizer = self._init_tokenizer()
243 - self.detokenizer = Detokenizer(self.tokenizer)
244 - tokenizer_group = self.get_tokenizer_group()
245 - else:
241 + if self.model_config.skip_tokenizer_init:
242 self.tokenizer = None
243 self.detokenizer = None
244 tokenizer_group = None
245 + else:
246 + self.tokenizer = self._init_tokenizer()
247 + self.detokenizer = Detokenizer(self.tokenizer)
248 + tokenizer_group = self.get_tokenizer_group()
249
250 # Ensure that the function doesn't contain a reference to self,
251 # to avoid engine GC issues
vllm/model_executor/models/interfaces.py
+34
@@ -136,6 +136,40 @@ def supports_multimodal(
136 return getattr(model, "supports_multimodal", False)
137
138
139 +@runtime_checkable
140 +class SupportsMultiModalWithRawInput(SupportsMultiModal, Protocol):
141 + """The interface required for all multi-modal models."""
142 +
143 + supports_multimodal_raw_input: ClassVar[Literal[True]] = True
144 + """
145 + A flag that indicates this model supports multi-modal inputs and processes
146 + them in their raw form and not embeddings.
147 +
148 + Note:
149 + There is no need to redefine this flag if this class is in the
150 + MRO of your model class.
151 + """
152 +
153 +
154 +@overload
155 +def supports_multimodal_raw_input(
156 + model: object) -> TypeIs[SupportsMultiModalWithRawInput]:
157 + ...
158 +
159 +
160 +@overload
161 +def supports_multimodal_raw_input(
162 + model: type[object]) -> TypeIs[type[SupportsMultiModalWithRawInput]]:
163 + ...
164 +
165 +
166 +def supports_multimodal_raw_input(
167 + model: Union[type[object], object]
168 +) -> Union[TypeIs[type[SupportsMultiModalWithRawInput]],
169 + TypeIs[SupportsMultiModalWithRawInput]]:
170 + return getattr(model, "supports_multimodal_raw_input", False)
171 +
172 +
173 @runtime_checkable
174 class SupportsScoreTemplate(Protocol):
175 """The interface required for all models that support score template."""
vllm/model_executor/models/prithvi_geospatial_mae.py
+57 -17
@@ -16,6 +16,7 @@
16 # See the License for the specific language governing permissions and
17 # limitations under the License.
18 """Inference-only IBM/NASA Prithvi Geospatial model."""
19 +
20 from collections.abc import Iterable, Mapping, Sequence
21 from typing import Optional, Union
22
@@ -27,13 +28,14 @@ from vllm.config import VllmConfig
28 from vllm.model_executor.layers.pooler import (AllPool, PoolerHead,
29 PoolerIdentity, SimplePooler)
30 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
30 -from vllm.model_executor.models.interfaces import (IsAttentionFree,
31 - SupportsMultiModal,
32 - SupportsV0Only)
31 +from vllm.model_executor.models.interfaces import (
32 + IsAttentionFree, MultiModalEmbeddings, SupportsMultiModalWithRawInput)
33 from vllm.model_executor.models.utils import AutoWeightsLoader
34 from vllm.multimodal import MULTIMODAL_REGISTRY
35 from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,
36 - MultiModalInputs, MultiModalKwargs)
36 + MultiModalFieldElem, MultiModalInputs,
37 + MultiModalKwargs, MultiModalKwargsItem,
38 + MultiModalSharedField, PlaceholderRange)
39 from vllm.multimodal.parse import MultiModalDataItems
40 from vllm.multimodal.processing import (BaseMultiModalProcessor,
41 BaseProcessingInfo, PromptUpdate)
@@ -62,8 +64,9 @@ class PrithviGeoSpatialMAEInputBuilder(
64 # The size of pixel_values might change in the cases where we resize
65 # the input but never exceeds the dimensions below.
66 return {
65 - "pixel_values": torch.full((1, 6, 512, 512), 1.0),
66 - "location_coords": torch.full((1, 2), 1.0),
67 + "pixel_values": torch.full((6, 512, 512), 1.0,
68 + dtype=torch.float16),
69 + "location_coords": torch.full((1, 2), 1.0, dtype=torch.float16),
70 }
71
72
@@ -75,8 +78,10 @@ class PrithviGeoSpatialMAEMultiModalProcessor(BaseMultiModalProcessor):
78 hf_processor_mm_kwargs: Mapping[str, object],
79 ) -> Mapping[str, MultiModalFieldConfig]:
80 return dict(
78 - pixel_values=MultiModalFieldConfig.batched("image"),
79 - location_coords=MultiModalFieldConfig.batched("image"),
81 + pixel_values=MultiModalFieldConfig.shared(batch_size=1,
82 + modality="image"),
83 + location_coords=MultiModalFieldConfig.shared(batch_size=1,
84 + modality="image"),
85 )
86
87 def _get_prompt_updates(
@@ -99,23 +104,48 @@ class PrithviGeoSpatialMAEMultiModalProcessor(BaseMultiModalProcessor):
104
105 for k, v in mm_data.items():
106 mm_kwargs[k] = v
107 + mm_placeholders = {"image": [PlaceholderRange(offset=0, length=0)]}
108 +
109 + # This model receives in input a multi-dimensional tensor representing
110 + # a single image patch and therefore it is not to be split
111 + # into multiple elements, but rather to be considered a single one.
112 + # Hence, the decision of using a MultiModalSharedField.
113 + # The expected shape is (num_channels, width, height).
114 +
115 + # This model however allows the user to also submit multiple image
116 + # patches as a batch, adding a further dimension to the above shape.
117 + # At this stage we only support submitting one patch per request and
118 + # batching is achieved via vLLM batching.
119 + # TODO (christian-pinto): enable support for multi patch requests
120 + # in tandem with vLLM batching.
121 + multimodal_kwargs_items = [
122 + MultiModalKwargsItem.from_elems([
123 + MultiModalFieldElem(
124 + modality="image",
125 + key=key,
126 + data=data,
127 + field=MultiModalSharedField(1),
128 + ) for key, data in mm_kwargs.items()
129 + ])
130 + ]
131
132 return MultiModalInputs(
133 type="multimodal",
134 prompt=prompt,
135 prompt_token_ids=[1],
107 - mm_kwargs=MultiModalKwargs(mm_kwargs),
136 + mm_kwargs=MultiModalKwargs.from_items(multimodal_kwargs_items),
137 mm_hashes=None,
109 - mm_placeholders={},
138 + mm_placeholders=mm_placeholders,
139 )
140
141
142 @MULTIMODAL_REGISTRY.register_processor(
143 PrithviGeoSpatialMAEMultiModalProcessor,
144 info=PrithviGeoSpatialMAEProcessingInfo,
116 - dummy_inputs=PrithviGeoSpatialMAEInputBuilder)
117 -class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,
118 - SupportsV0Only):
145 + dummy_inputs=PrithviGeoSpatialMAEInputBuilder,
146 +)
147 +class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree,
148 + SupportsMultiModalWithRawInput):
149 """Prithvi Masked Autoencoder"""
150
151 is_pooling_model = True
@@ -128,10 +158,10 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,
158 raise ValueError("Only image modality is supported")
159
160 def _instantiate_model(self, config: dict) -> Optional[nn.Module]:
131 -
161 # We might be able/need to support different tasks with this same model
162 if config["task_args"]["task"] == "SemanticSegmentationTask":
163 from terratorch.cli_tools import SemanticSegmentationTask
164 +
165 task = SemanticSegmentationTask(
166 config["model_args"],
167 config["task_args"]["model_factory"],
@@ -144,7 +174,8 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,
174 scheduler_hparams=config["scheduler_params"],
175 plot_on_val=config["task_args"]["plot_on_val"],
176 freeze_decoder=config["task_args"]["freeze_decoder"],
147 - freeze_backbone=config["task_args"]["freeze_backbone"])
177 + freeze_backbone=config["task_args"]["freeze_backbone"],
178 + )
179
180 return task.model
181 else:
@@ -168,12 +199,10 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,
199
200 def _parse_and_validate_multimodal_data(
201 self, **kwargs) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
171 -
202 pixel_values = kwargs.pop("pixel_values", None)
203 if not isinstance(pixel_values, torch.Tensor):
204 raise ValueError(f"Incorrect type of pixel_values. "
205 f"Got type: {type(pixel_values)}")
176 - pixel_values = torch.unbind(pixel_values, dim=0)[0]
206
207 location_coords = kwargs.pop("location_coords", None)
208 if not isinstance(location_coords, torch.Tensor):
@@ -185,6 +214,17 @@ class PrithviGeoSpatialMAE(nn.Module, IsAttentionFree, SupportsMultiModal,
214
215 return pixel_values, location_coords
216
217 + def get_input_embeddings(
218 + self,
219 + input_ids: torch.Tensor,
220 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
221 + ) -> torch.Tensor:
222 + # We do not really use any input tokens and therefore no embeddings
223 + # to be calculated. However, due to the mandatory token ids in
224 + # the input prompt we pass one token and the size of the dummy
225 + # embedding tensors must reflect that.
226 + return torch.empty((input_ids.shape[0], 0))
227 +
228 def forward(
229 self,
230 input_ids: Optional[torch.Tensor],
vllm/model_executor/models/registry.py
+11 -2
@@ -22,8 +22,8 @@ from vllm.logger import init_logger
22
23 from .interfaces import (has_inner_state, has_noops, is_attention_free,
24 is_hybrid, supports_cross_encoding,
25 - supports_multimodal, supports_pp,
26 - supports_transcription, supports_v0_only)
25 + supports_multimodal, supports_multimodal_raw_input,
26 + supports_pp, supports_transcription, supports_v0_only)
27 from .interfaces_base import is_text_generation_model
28
29 logger = init_logger(__name__)
@@ -287,6 +287,7 @@ class _ModelInfo:
287 is_pooling_model: bool
288 supports_cross_encoding: bool
289 supports_multimodal: bool
290 + supports_multimodal_raw_input: bool
291 supports_pp: bool
292 has_inner_state: bool
293 is_attention_free: bool
@@ -304,6 +305,7 @@ class _ModelInfo:
305 is_pooling_model=True, # Can convert any model into a pooling model
306 supports_cross_encoding=supports_cross_encoding(model),
307 supports_multimodal=supports_multimodal(model),
308 + supports_multimodal_raw_input=supports_multimodal_raw_input(model),
309 supports_pp=supports_pp(model),
310 has_inner_state=has_inner_state(model),
311 is_attention_free=is_attention_free(model),
@@ -573,6 +575,13 @@ class _ModelRegistry:
575 model_cls, _ = self.inspect_model_cls(architectures)
576 return model_cls.supports_multimodal
577
578 + def supports_multimodal_raw_input(
579 + self,
580 + architectures: Union[str, list[str]],
581 + ) -> bool:
582 + model_cls, _ = self.inspect_model_cls(architectures)
583 + return model_cls.supports_multimodal_raw_input
584 +
585 def is_pp_supported_model(
586 self,
587 architectures: Union[str, list[str]],
vllm/multimodal/registry.py
+1 -1
@@ -266,7 +266,7 @@ class MultiModalRegistry:
266 if not model_config.is_multimodal_model:
267 raise ValueError(f"{model_config.model} is not a multimodal model")
268
269 - if tokenizer is None:
269 + if tokenizer is None and not model_config.skip_tokenizer_init:
270 tokenizer = cached_tokenizer_from_config(model_config)
271 if disable_cache is None:
272 mm_config = model_config.get_multimodal_config()
vllm/v1/engine/async_llm.py
+12 -5
@@ -94,11 +94,14 @@ class AsyncLLM(EngineClient):
94 self.log_requests = log_requests
95 self.log_stats = log_stats
96
97 - # Tokenizer (+ ensure liveness if running in another process).
98 - self.tokenizer = init_tokenizer_from_configs(
99 - model_config=vllm_config.model_config,
100 - scheduler_config=vllm_config.scheduler_config,
101 - lora_config=vllm_config.lora_config)
97 + if self.model_config.skip_tokenizer_init:
98 + self.tokenizer = None
99 + else:
100 + # Tokenizer (+ ensure liveness if running in another process).
101 + self.tokenizer = init_tokenizer_from_configs(
102 + model_config=vllm_config.model_config,
103 + scheduler_config=vllm_config.scheduler_config,
104 + lora_config=vllm_config.lora_config)
105
106 # Processor (converts Inputs --> EngineCoreRequests).
107 self.processor = Processor(
@@ -525,6 +528,10 @@ class AsyncLLM(EngineClient):
528 self,
529 lora_request: Optional[LoRARequest] = None,
530 ) -> AnyTokenizer:
531 + if self.tokenizer is None:
532 + raise ValueError("Unable to get tokenizer because "
533 + "skip_tokenizer_init is True")
534 +
535 return self.tokenizer.get_lora_tokenizer(lora_request)
536
537 async def is_tracing_enabled(self) -> bool:
vllm/v1/engine/llm_engine.py
+8 -5
@@ -82,11 +82,14 @@ class LLMEngine:
82 self.dp_group = None
83 self.should_execute_dummy_batch = False
84
85 - # Tokenizer (+ ensure liveness if running in another process).
86 - self.tokenizer = init_tokenizer_from_configs(
87 - model_config=vllm_config.model_config,
88 - scheduler_config=vllm_config.scheduler_config,
89 - lora_config=vllm_config.lora_config)
85 + if self.model_config.skip_tokenizer_init:
86 + self.tokenizer = None
87 + else:
88 + # Tokenizer (+ ensure liveness if running in another process).
89 + self.tokenizer = init_tokenizer_from_configs(
90 + model_config=vllm_config.model_config,
91 + scheduler_config=vllm_config.scheduler_config,
92 + lora_config=vllm_config.lora_config)
93
94 # Processor (convert Inputs --> EngineCoreRequests)
95 self.processor = Processor(vllm_config=vllm_config,
vllm/v1/engine/output_processor.py
+10 -8
@@ -327,14 +327,16 @@ class OutputProcessor:
327 if request_id in self.request_states:
328 raise ValueError(f"Request id {request_id} already running.")
329
330 - req_state = RequestState.from_new_request(
331 - tokenizer=self.tokenizer.get_lora_tokenizer(request.lora_request),
332 - request=request,
333 - prompt=prompt,
334 - parent_req=parent_req,
335 - request_index=request_index,
336 - queue=queue,
337 - log_stats=self.log_stats)
330 + tokenizer = None if not self.tokenizer else \
331 + self.tokenizer.get_lora_tokenizer(request.lora_request)
332 +
333 + req_state = RequestState.from_new_request(tokenizer=tokenizer,
334 + request=request,
335 + prompt=prompt,
336 + parent_req=parent_req,
337 + request_index=request_index,
338 + queue=queue,
339 + log_stats=self.log_stats)
340 self.request_states[request_id] = req_state
341 self.lora_states.add_request(req_state)
342 if parent_req:
vllm/v1/engine/processor.py
+8 -4
@@ -380,7 +380,6 @@ class Processor:
380 prompt_type: Literal["encoder", "decoder"],
381 ):
382 model_config = self.model_config
383 - tokenizer = self.tokenizer.get_lora_tokenizer(lora_request)
383
384 prompt_ids = prompt_inputs["prompt_token_ids"]
385 if not prompt_ids:
@@ -389,9 +388,14 @@ class Processor:
388 else:
389 raise ValueError(f"The {prompt_type} prompt cannot be empty")
390
392 - max_input_id = max(prompt_ids, default=0)
393 - if max_input_id > tokenizer.max_token_id:
394 - raise ValueError(f"Token id {max_input_id} is out of vocabulary")
391 + if self.model_config.skip_tokenizer_init:
392 + tokenizer = None
393 + else:
394 + tokenizer = self.tokenizer.get_lora_tokenizer(lora_request)
395 + max_input_id = max(prompt_ids, default=0)
396 + if max_input_id > tokenizer.max_token_id:
397 + raise ValueError(
398 + f"Token id {max_input_id} is out of vocabulary")
399
400 max_prompt_len = self.model_config.max_model_len
401 if len(prompt_ids) > max_prompt_len:
vllm/v1/worker/gpu_model_runner.py
+60
@@ -126,6 +126,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
126
127 self.is_multimodal_model = model_config.is_multimodal_model
128 self.is_pooling_model = model_config.pooler_config is not None
129 + self.model_supports_multimodal_raw_input = (
130 + model_config.model_supports_multimodal_raw_input)
131 self.max_model_len = model_config.max_model_len
132 self.max_num_tokens = scheduler_config.max_num_batched_tokens
133 self.max_num_reqs = scheduler_config.max_num_seqs
@@ -328,6 +330,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
330 Args:
331 scheduler_output: The scheduler output.
332 """
333 + # Attention free models have zero kv_cache_goups, however models
334 + # like Mamba are also attention free but use the kv_cache for
335 + # keeping its internal state. This is why we check the number
336 + # of kv_cache groups instead of solely checking
337 + # for self.model_config.is_attention_free.
338 + if len(self.kv_cache_config.kv_cache_groups) == 0:
339 + return
340 +
341 self.attn_metadata_builders[0].reorder_batch(self.input_batch,
342 scheduler_output)
343
@@ -565,6 +575,38 @@ class GPUModelRunner(LoRAModelRunnerMixin):
575 # Refresh batch metadata with any pending updates.
576 self.input_batch.refresh_metadata()
577
578 + def _init_model_kwargs_for_multimodal_model(
579 + self,
580 + scheduler_output: Optional["SchedulerOutput"] = None,
581 + num_reqs: int = -1,
582 + ) -> dict[str, Any]:
583 +
584 + model_kwargs: dict[str, Any] = {}
585 + if self.model_supports_multimodal_raw_input:
586 + # This model requires the raw multimodal data in input.
587 + if scheduler_output:
588 + multi_modal_kwargs_list = []
589 + for req in scheduler_output.scheduled_new_reqs:
590 + req_mm_inputs = req.mm_inputs
591 + if not isinstance(req_mm_inputs, list):
592 + req_mm_inputs = list(req_mm_inputs)
593 + multi_modal_kwargs_list.extend(req_mm_inputs)
594 + multi_modal_kwargs = MultiModalKwargs.batch(
595 + multi_modal_kwargs_list)
596 + else:
597 + # The only case where SchedulerOutput is None is for
598 + # a dummy run let's get some dummy data.
599 + dummy_data = [
600 + self.mm_registry.get_decoder_dummy_data(
601 + model_config=self.model_config,
602 + seq_len=1).multi_modal_data for i in range(num_reqs)
603 + ]
604 + multi_modal_kwargs = MultiModalKwargs.batch(dummy_data)
605 +
606 + model_kwargs.update(multi_modal_kwargs)
607 +
608 + return model_kwargs
609 +
610 def _get_cumsum_and_arange(
611 self,
612 num_tokens: np.ndarray,
@@ -1359,10 +1401,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1401 # embeddings), we always use embeddings (rather than token ids)
1402 # as input to the multimodal model, even when the input is text.
1403 input_ids = self.input_ids[:num_scheduled_tokens]
1404 +
1405 + model_kwargs = self._init_model_kwargs_for_multimodal_model(
1406 + scheduler_output=scheduler_output)
1407 inputs_embeds = self.model.get_input_embeddings(
1408 input_ids=input_ids,
1409 multimodal_embeddings=mm_embeds or None,
1410 )
1411 +
1412 # TODO(woosuk): Avoid the copy. Optimize.
1413 self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds)
1414 inputs_embeds = self.inputs_embeds[:num_input_tokens]
@@ -1374,6 +1420,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1420 # then the embedding layer is not included in the CUDA graph.
1421 input_ids = self.input_ids[:num_input_tokens]
1422 inputs_embeds = None
1423 + model_kwargs = {}
1424 if self.uses_mrope:
1425 positions = self.mrope_positions[:, :num_input_tokens]
1426 else:
@@ -1406,6 +1453,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1453 positions=positions,
1454 intermediate_tensors=intermediate_tensors,
1455 inputs_embeds=inputs_embeds,
1456 + **MultiModalKwargs.as_kwargs(
1457 + model_kwargs,
1458 + device=self.device,
1459 + ),
1460 )
1461
1462 self.maybe_wait_for_kv_save()
@@ -2084,11 +2135,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2135 num_scheduled_tokens):
2136 model = self.model
2137 if self.is_multimodal_model:
2138 + model_kwargs = self._init_model_kwargs_for_multimodal_model(
2139 + num_reqs=num_reqs)
2140 input_ids = None
2141 inputs_embeds = self.inputs_embeds[:num_tokens]
2142 else:
2143 input_ids = self.input_ids[:num_tokens]
2144 inputs_embeds = None
2145 + model_kwargs = {}
2146 +
2147 if self.uses_mrope:
2148 positions = self.mrope_positions[:, :num_tokens]
2149 else:
@@ -2117,7 +2172,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2172 positions=positions,
2173 intermediate_tensors=intermediate_tensors,
2174 inputs_embeds=inputs_embeds,
2175 + **MultiModalKwargs.as_kwargs(
2176 + model_kwargs,
2177 + device=self.device,
2178 + ),
2179 )
2180 +
2181 if self.use_aux_hidden_state_outputs:
2182 hidden_states, _ = outputs
2183 else: