|
| 1 | +--- |
| 2 | +layout: post |
| 3 | +title: "Serving Geospatial, Vision, and Beyond: Enabling Any-Modality Models in vLLM" |
| 4 | +author: Christian Pinto (IBM Research Europe - Dublin), Michele Gazzetti (IBM Research Europe - Dublin), Michael Johnston (IBM Research Europe - Dublin) |
| 5 | +image: /assets/logos/vllm-logo-text-light.png |
| 6 | +--- |
| 7 | +## Introduction |
| 8 | + |
| 9 | +Until recently, generative AI infrastructure has been tightly coupled with autoregressive text generation models that produce output token-by-token, typically in the form of natural language. |
| 10 | +But not all models work this way. |
| 11 | +A growing class of non-autoregressive models generate their outputs in a single inference pass, enabling faster and more efficient generation across a wide range of modalities. |
| 12 | + |
| 13 | +These models are increasingly used in domains beyond text: from image classification and segmentation, to audio synthesis and structured data generation. |
| 14 | +Supporting such any-modality models where inputs and outputs may be images, audio, tabular data, or combinations thereof, requires a fundamental shift in how inference frameworks operate. |
| 15 | + |
| 16 | +We've made that shift in vLLM. |
| 17 | + |
| 18 | +In this article, we introduce a set of enhancements to vLLM that enable non-autoregressive, any-modality model serving. |
| 19 | +While our initial integration focuses on geospatial vision transformers, used for tasks like flood detection, burn scar identification, and land use classification from satellite imagery, the changes are generic and pave the way for serving a wide variety of non text-generating models. |
| 20 | + |
| 21 | +As a concrete example, we've integrated all geospatial models from the [Terratorch](https://github.com/IBM/terratorch) framework (developed in collaboration with NASA) into vLLM via a generic backend, making them first-class citizens in the vLLM ecosystem. |
| 22 | + |
| 23 | +In the sections that follow, we describe the technical changes made to vLLM, starting with the requirements and challenges of serving geospatial vision transformer models. |
| 24 | + |
| 25 | +## Integrating geospatial vision transformer models in vLLM |
| 26 | + |
| 27 | +Unlike text models, vision transformers don’t need token decoding i.e. they do not need output tokens to be transformed into text. |
| 28 | +Instead, given one input image, a single inference generates the raw model output, and then this is post-processed into the output image. |
| 29 | +In addition, sometimes the input image needs to be partitioned and batched into a number of sub-images, or patches. |
| 30 | +These patches are then fed to the model for inference, with the resulting output images from each patch being stitched together to form the final output image. |
| 31 | + |
| 32 | +<p align="center"> |
| 33 | +<picture> |
| 34 | +<img src="/assets/figures/beyond-text/models-diff.png" width="80%"> |
| 35 | +</picture> |
| 36 | +</p> |
| 37 | + |
| 38 | +Given these requirements, the obvious choice was to integrate vision transformers in vLLM as pooling models. |
| 39 | +In vLLM pooling models allow extracting the raw model output via an identity pooler. |
| 40 | +Identity poolers do not apply any transformation to the data and return it as is - exactly what we need. |
| 41 | +For the input, we exploit the existing multimodal input capabilities of vLLM to pre-proces images into tensors that are then fed to vLLM for inference. |
| 42 | + |
| 43 | +Since we wanted to support multiple geospatial vision transformers out-of-the-box in vLLM we have also added a model implementation backend for TerraTorch models, following the same pattern as the backend for the HuggingFace Transformers library. |
| 44 | + |
| 45 | +Getting this to work was no easy task, though. |
| 46 | +Enabling these model classes required changes to various parts of vLLM such as: |
| 47 | + |
| 48 | +* adding support for attention free models |
| 49 | +* improving support for models that do not require a tokenizer |
| 50 | +* enabling processing of raw input data as opposed to the default multimodal input embeddings |
| 51 | +* extending the vLLM serving API. |
| 52 | + |
| 53 | +## Meet IO Processor: Flexible Input/Output Handling for Any Model |
| 54 | + |
| 55 | +So far so good! Well, this brings us only halfway towards our goal. |
| 56 | + |
| 57 | +With the above integration we can indeed serve geospatial vision transformer models but only in tensor-to-tensor format. |
| 58 | +Users must pre-process their image to a tensor format, before sending these tensors to the vLLM instance. |
| 59 | +Similarly post-processing of the raw tensor output has to happen outside vLLM. |
| 60 | +The impact: there is no endpoint that users can send an image to and get an image back. |
| 61 | + |
| 62 | +This problem existed because, before our changes, pre-processing of input data and post-processing of the model output was only partially supported in vLLM. |
| 63 | +Specifically, pre-processing of multi-modal input data was only possible via the processors available in the transformers library. |
| 64 | +However, the transformers processors usually support only default data types and do not deal with complex formats such as `geotiff` that enrich tiff files with georeferenced metadata. |
| 65 | +Also, on the output processing side vLLM only supported de-tokenization into text or the application of poolers to the model hidden states - no other output processing was possible. |
| 66 | + |
| 67 | +This is where the new IO Processor plugin framework we introduced comes in. |
| 68 | +The IO Processor framework allows developers to customize how model inputs and outputs are pre and post processed, all within the same vLLM serving instance. |
| 69 | +Whether your model returns a string, a JSON object, an image tensor, or a custom data structure, an IO Processor can translate it into the desired format before returning it to the client. |
| 70 | + |
| 71 | +<p align="center"> |
| 72 | +<picture> |
| 73 | +<img src="/assets/figures/beyond-text/io-plugins-flow.png" width="70%"> |
| 74 | +</picture> |
| 75 | +</p> |
| 76 | + |
| 77 | +The IO Processor framework unlocks a new level of flexibility for vLLM users. |
| 78 | +It means non-text models (e.g., image generators, image to segmentation mask, tabular to classification, etc.) can be served using standard vLLM infrastructure. |
| 79 | +Via IO Processors users can plug in custom logic to transform or enrich outputs such as decoding model outputs into images, or formatting responses for downstream systems. |
| 80 | +This maintains a unified serving stack, reducing operational complexity and improving maintainability. |
| 81 | + |
| 82 | +### Using vLLM IO Processor plugins |
| 83 | + |
| 84 | +Each IO Processor plugin implements a pre-defined [IO Processor interface](https://github.com/vllm-project/vllm/blob/main/vllm/plugins/io_processors/interface.py) and resides outside of the vLLM source code tree. |
| 85 | +At installation time each plugin registers one or more entrypoints in the `vllm.io_processor_plugins` group. |
| 86 | +This allows vLLM to automatically discover and load plugins at engine initialization time. |
| 87 | +A full plugin example for geospatial vision transformers is available [here](https://github.com/christian-pinto/prithvi_io_processor_plugin). |
| 88 | + |
| 89 | +Using an IO Processor plugin is as easy as just installing it in the same python environment with vLLM and adding the `--io-processor-plugin <plugin_name>` when starting the serving instance. |
| 90 | +Currently, one IO Processor plugin can be loaded for each vLLM instance. |
| 91 | + |
| 92 | +Once the serving instance is started, pre and post processing is automatically applied to the model input and output when serving the `pooling` endpoint. |
| 93 | +At this stage IO Processors are only available for pooling models, but in the future we expect other endpoints to be integrated too. |
| 94 | + |
| 95 | +## Step-by-Step: Serving the Prithvi Model in vLLM |
| 96 | + |
| 97 | +One example model class that can be served with vLLM using the Terratorch backend is [Prithvi for flood detection](https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11). |
| 98 | + |
| 99 | +### Install the python requirements |
| 100 | + |
| 101 | +Install terratorch (>=1.1rc3) and vLLM in your python environment. |
| 102 | +At the time of writing this article the changes required for replicating this example are not yet part of a vLLM release (current latest is v0.10.1.1) and we advise users to install the [latest code](https://docs.vllm.ai/en/latest/getting_started/installation/gpu.html#install-the-latest-code_1). |
| 103 | + |
| 104 | +Download and install the IO Processor plugin for flood detection with Prithvi. |
| 105 | + |
| 106 | +```bash |
| 107 | +git clone [email protected]:christian-pinto/prithvi_io_processor_plugin.git |
| 108 | +cd prithvi_io_processor_plugin |
| 109 | +pip install . |
| 110 | +``` |
| 111 | + |
| 112 | +This installs the `prithvi_to_tiff` plugin. |
| 113 | + |
| 114 | +### Start a vLLM serving instance |
| 115 | + |
| 116 | +Start a vLLM serving instance that loads the `prithvi_to_tiff` plugin and the Prithvi model for flood detection. |
| 117 | + |
| 118 | +```bash |
| 119 | +vllm serve \ |
| 120 | + --model=ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11 \ |
| 121 | + --model-impl terratorch \ |
| 122 | + --task embed --trust-remote-code \ |
| 123 | + --skip-tokenizer-init --enforce-eager \ |
| 124 | + --io-processor-plugin prithvi_to_tiff |
| 125 | +``` |
| 126 | + |
| 127 | +Once the serving instance is fully up and running it is ready to serve requests with the selected plugin. |
| 128 | +The below log entries confirm that your vLLM instance is up and running and that it is listening on port `8000`. |
| 129 | + |
| 130 | +```bash |
| 131 | +INFO: Starting vLLM API server 0 on http://0.0.0.0:8000 |
| 132 | +... |
| 133 | +... |
| 134 | +INFO: Started server process [409128] |
| 135 | +INFO: Waiting for application startup. |
| 136 | +INFO: Application startup complete. |
| 137 | +``` |
| 138 | + |
| 139 | +### Send requests to the model |
| 140 | +The below python script sends a request to the vLLM `pooling` endpoint with a specific JSON payload where the `model` and `softmax` arguments are pre-defined, while the `data` field is defined by the user and depends on the plugin in use. |
| 141 | +Note setting the `softmax` field to `False`, required to ensure the plugin receives the raw model output. |
| 142 | +In this case we send the input image to vLLM as a URL and we request the response to be a geotiff image in base64 encoding. |
| 143 | +The script decodes the image and writes it to disk as a tiff(geotiff) file. |
| 144 | + |
| 145 | +```python |
| 146 | +import base64 |
| 147 | +import os |
| 148 | +import requests |
| 149 | + |
| 150 | +def main(): |
| 151 | + image_url = "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff" |
| 152 | + server_endpoint = "http://localhost:8000/pooling" |
| 153 | + |
| 154 | + request_payload = { |
| 155 | + "data": { |
| 156 | + "data": image_url, |
| 157 | + "data_format": "url", |
| 158 | + "image_format": "tiff", |
| 159 | + "out_data_format": "b64_json", |
| 160 | + }, |
| 161 | + "model": "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11", |
| 162 | + "softmax": False, |
| 163 | + } |
| 164 | + |
| 165 | + ret = requests.post(server_endpoint, json=request_payload) |
| 166 | + |
| 167 | + if ret.status_code == 200: |
| 168 | + response = ret.json() |
| 169 | + |
| 170 | + decoded_image = base64.b64decode(response["data"]["data"]) |
| 171 | + |
| 172 | + out_path = os.path.join(os.getcwd(), "online_prediction.tiff") |
| 173 | + |
| 174 | + with open(out_path, "wb") as f: |
| 175 | + f.write(decoded_image) |
| 176 | + else: |
| 177 | + print(f"Response status_code: {ret.status_code}") |
| 178 | + print(f"Response reason:{ret.reason}") |
| 179 | + |
| 180 | + |
| 181 | +if __name__ == "__main__": |
| 182 | + main() |
| 183 | +``` |
| 184 | + |
| 185 | +Below is an example of input and the output you should obtain. |
| 186 | +The input image (left) is a satellite picture of Valencia, Spain during the 2024 flood. |
| 187 | +The output image (right) shows the ares predicted as flooded (in white) by the Prithvi model. |
| 188 | + |
| 189 | +<p align="center"> |
| 190 | +<picture> |
| 191 | +<img src="/assets/figures/beyond-text/prithvi-prediction.png" width="100%"> |
| 192 | +</picture> |
| 193 | +</p> |
| 194 | + |
| 195 | +## What’s Next |
| 196 | + |
| 197 | +This is just the beginning. |
| 198 | +We plan to expand IO Processor plugins across more Terratorch models and modalities and beyond, making installation seamless. |
| 199 | +Longer-term, we envision IO Processors powered vision-language systems, structured reasoning agents, and multi-modal pipelines, all served from the same vLLM stack. |
| 200 | + |
| 201 | +To get started, check out the IO Processor [documentation](https://docs.vllm.ai/en/latest/design/io_processor_plugins.html) and explore the [examples](https://github.com/vllm-project/vllm/tree/main/examples). |
| 202 | +More information on IBM's Terratorch are available [here](https://github.com/IBM/terratorch). |
| 203 | + |
| 204 | +Contributions, feedback, and ideas are always welcome! |
| 205 | + |
| 206 | +## Acknowledgement |
| 207 | +We would like to thank the members of the vLLM community who helped improving our contribution, and [Maximilien Philippe Marie de Bayser](https://github.com/maxdebayser) (IBM Research Brazil) for his contributions to the IO Processor plugins framework and [Cyrus Leung](https://github.com/DarkLight1337) (HKUST) for his support in shaping up the overall concept of extending vLLM beyond text generation. |
0 commit comments