服务于地理空间、视觉及更多领域：在 vLLM 中启用多模态输出处理

简介

直到最近，生成式 AI 基础设施还主要与自回归文本生成模型紧密耦合，这些模型按标记（token）生成输出，通常以自然语言形式呈现。vLLM 最初顺应这一趋势，支持处理文本输入和输出的传统大型语言模型（LLM）。随着多模态大语言模型（MLLM）的引入，趋势开始转向多模态数据，这些模型能够处理文本以及各种模态的数据（如图像、音频、视频等）。vLLM 也随之跟进，支持 LLaVA 风格的 MLLM，能够对多模态输入数据进行推理并生成文本。

我们现在正见证一个新的趋势转变，即越来越多的非自回归模型能够在单次推理中生成多模态输出，从而在广泛的模态中实现更快、更高效的生成。从推理的角度来看，这些模型可以被视为汇聚（pooling）模型，但需要额外的输入和输出处理支持。这类模型的应用领域不仅限于文本：从图像分类和分割，到音频合成和结构化数据生成。我们在 vLLM 中迈出了下一步，增加了对此类模型的支持。

我们最初的集成重点是地理空间基础模型。这类卷积模型或视觉 Transformer 模型需要除 RGB 通道之外的数据（如多光谱或雷达数据）以及元数据（如地理位置、图像获取日期），常用于卫星影像的灾害响应或土地利用分类等任务，但不仅限于此。然而，这些更改是通用的，为服务各种非文本生成模型铺平了道路。

作为一个具体的例子，我们通过通用后端将 TerraTorch 框架中的所有地理空间模型（其中一些是与 NASA 和 ESA 合作开发的）集成到了 vLLM 中，使其成为 vLLM 生态系统中的一等公民。

在接下来的章节中，我们将描述对 vLLM 进行的技术改进，从服务地理空间基础模型的要求和挑战开始。

在 vLLM 中集成地理空间基础模型

与文本模型不同，地理空间基础模型（通常实现为视觉 Transformer）不需要标记解码，即不需要将输出标记转换为文本。相反，给定一张输入图像，单次推理即可生成原始模型输出，随后将其后处理为输出图像。此外，有时需要将输入图像分区并批处理为多个子图像或补丁（patches）。这些补丁被输入到模型进行推理，来自每个补丁的结果图像被拼接在一起，形成最终的输出图像。

基于这些需求，显而易见的选择是将地理空间基础模型作为汇聚（pooling）模型集成到 vLLM 中。汇聚是深度学习中常用的技术，用于减少特征图的空间维度。常见的类型包括最大汇聚、平均汇聚和全局汇聚，每种都使用不同的策略来聚合信息。在 vLLM 中，汇聚可以应用于诸如嵌入向量计算和分类之类的任务。此外，vLLM 支持恒等汇聚（identity pooler），即在不应用任何转换的情况下返回模型的隐藏状态——这正是我们所需要的。对于输入，我们预处理图像为张量，然后将其馈送到模型进行推理，利用 vLLM 现有的多模态输入能力。

由于我们希望在 vLLM 中开箱即用地支持多个地理空间基础模型，我们还添加了一个用于 TerraTorch 模型的模型实现后端，遵循与 HuggingFace Transformers 库后端相同的模式。

然而，实现这一点并非易事。启用这些模型类需要对 vLLM 的多个部分进行更改，例如：

• 添加对无注意力机制（attention-free）模型的支持
• 改进对不需要分词器（tokenizer）的模型的支持
• 启用对原始输入数据的处理，而不是默认的多模态输入嵌入
• 扩展 vLLM 服务 API。

认识 IO Processor：适用于任何模型的灵活输入/输出处理

目前进展顺利！但这仅仅是我们目标的一半。

通过上述集成，我们确实可以服务地理空间基础模型——尽管仅限于张量到张量的格式。用户仍然必须在将张量发送到 vLLM 实例之前将图像预处理为张量格式。同样，原始张量输出的后处理必须在 vLLM 之外进行。结果是：没有用户可以发送图像并收回图像的端点。

这个问题存在是因为，在我们进行更改之前，vLLM 对输入数据的预处理和模型输出的后处理仅得到了部分支持。具体来说，多模态输入数据的预处理只能通过 Transformers 库中可用的处理器实现。然而，transformers 处理器通常只支持标准数据类型，无法处理更复杂的数据格式（如 GeoTIFF，即带有丰富地理空间元数据的图像文件）。此外，在输出处理方面，vLLM 仅支持去标记化为文本或对模型隐藏状态应用汇聚器——无法进行其他输出处理。

这就是我们引入的新 IO Processor 插件框架的作用所在。IO Processor 框架允许开发者在同一个 vLLM 服务实例中自定义模型输入和输出的预处理与后处理方式。无论您的模型返回字符串、JSON 对象、图像张量还是自定义数据结构，IO Processor 都能在将其返回给客户端之前将其转换为所需的格式。

IO Processor 框架为 vLLM 用户开启了新的灵活性。这意味着非文本模型（例如，图像生成器、图像到分割掩码、表格到分类等）可以使用标准的 vLLM 基础设施进行服务。通过 IO Processors，用户可以插入自定义逻辑来转换或丰富输出，例如将模型输出解码为图像，或为下游系统格式化响应。这保持了统一的服务栈，降低了操作复杂性并提高了可维护性。

使用 vLLM IO Processor 插件

每个 IO Processor 插件都实现了一个预定义的 IO Processor 接口，并位于 vLLM 源代码树之外。在安装时，每个插件在 vllm.io_processor_plugins 组中注册一个或多个入口点。这允许 vLLM 在引擎初始化时自动发现并加载插件。

使用 IO Processor 插件非常简单，只需将其与 vLLM 安装在同一个 Python 环境中，并在启动服务实例时添加 --io-processor-plugin <plugin_name> 参数即可。目前，每个 vLLM 实例可以加载一个 IO Processor 插件。

服务实例启动后，当服务 /pooling 端点时，预处理和后处理会自动应用于模型输入和输出。现阶段，IO Processors 仅适用于汇聚模型，但我们预计未来其他端点也会进行集成。

分步指南：在 vLLM 中服务 Prithvi 模型

使用 TerraTorch 后端可以在 vLLM 中服务的模型类的一个例子是 用于洪水检测的 Prithvi。Prithvi 地理空间基础模型的完整插件示例可在此处获取。

Prithvi IO Processor 插件

为了说明 IO Processor 插件方法的灵活性，下面的伪代码展示了 Prithvi IO Processor 预处理和后处理的主要步骤。我们要强调的是数据特定转换与模型推理数据之间的解耦。这为理想中的任何模型和任何输入/输出数据类型，甚至应用于相同模型输出的多个插件提供了空间，具体取决于消费数据的下游任务。

def pre_process(request_data: dict):
    # Downloads geotiff
    # In this example the input image has 7 bands
    image_url = request_data["url"]
    image_obj = download_image(image_url)
 
    # Extract image data:
    # - pixel_values([n, 6, 512, 512])
    #   - 6 input bands R, G, B, +3 multispectral wavelengths
    #   - n > 1 if the size of the input image is > [512, 512]
    # - metadata
    #   - GPS coordinates
    #   - date
    pixel_values, metadata = process_image(image_obj)
 
    # Process the image data into n vLLM prompts
    model_prompts = pixels_to_prompts(pixel_values)
 
    return model_prompts
 
 
def post_process(model_outputs: list[PoolingRequestOutput]):
    # Uses the previously extracted metadata to guarantee the output
    # contains the same georeferences and date.
    return image_object(model_outputs, metadata)

安装 Python 依赖

在您的 Python 环境中安装 terratorch (>=1.1rc3) 和 vllm 包。在撰写本文时，复制此示例所需的更改尚未成为 vLLM 发布版本的一部分（目前最新版本为 v0.10.1.1），我们建议用户安装最新代码。

下载并安装用于 Prithvi 洪水检测的 IO Processor 插件。

git clone git@github.com:christian-pinto/prithvi_io_processor_plugin.git
cd prithvi_io_processor_plugin
pip install .

这将安装 prithvi_to_tiff 插件。

启动 vLLM 服务实例

启动一个加载 prithvi_to_tiff 插件和用于洪水检测的 Prithvi 模型的 vLLM 服务实例。

vllm serve \
    --model=ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11 \
    --model-impl terratorch \
    --task embed --trust-remote-code \
    --skip-tokenizer-init --enforce-eager \
    --io-processor-plugin prithvi_to_tiff

实例运行后，即可准备使用所选插件处理请求。下面的日志条目确认您的 vLLM 实例已启动并正在运行，且正在监听 8000 端口。

INFO: Starting vLLM API server 0 on http://0.0.0.0:8000
...
...
INFO: Started server process [409128]
INFO: Waiting for application startup.
INFO: Application startup complete.

向模型发送请求

下面的 Python 脚本向 vLLM /pooling 端点发送请求，并附带一个特定的 JSON 有效负载，其中 model 和 softmax 参数是预定义的，而 data 字段由用户定义，并取决于所使用的插件。

**注意：** 必须将 softmax 字段设置为 False，以确保插件接收到原始模型输出。在这种情况下，我们将输入图像作为 URL 发送到 vLLM，并要求响应为 base64 编码的 GeoTIFF 图像。脚本会对图像进行解码并将其作为 tiff (GeoTIFF) 文件写入磁盘。

import base64
import os
import requests
 
def main():
  image_url = "https://hugging-face.cn/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff"
  server_endpoint = "https://:8000/pooling"
 
  request_payload = {
      "data": {
          "data": image_url,
          "data_format": "url",
          "image_format": "tiff",
          "out_data_format": "b64_json",
      },
      "model": "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11",
      "softmax": False,
  }
 
  ret = requests.post(server_endpoint, json=request_payload)
 
  if ret.status_code == 200:
    response = ret.json()
 
    decoded_image = base64.b64decode(response["data"]["data"])
 
    out_path = os.path.join(os.getcwd(), "online_prediction.tiff")
 
    with open(out_path, "wb") as f:
        f.write(decoded_image)
  else:
    print(f"Response status_code: {ret.status_code}")
    print(f"Response reason:{ret.reason}")
 
 
if __name__ == "__main__":
    main()

下面是输入和预期输出的一个示例。输入图像（左）是 2024 年洪水期间西班牙瓦伦西亚的卫星照片。输出图像（右）显示了 Prithvi 模型预测为洪水淹没的区域（白色）。

未来展望

这仅仅是一个开始。我们计划将 IO Processor 插件扩展到更多的 TerraTorch 模型、模态及其他领域，使安装变得无缝。从长远来看，我们设想由 IO Processors 驱动的视觉语言系统、结构化推理代理和多模态管道，所有这些都从同一个 vLLM 栈中提供服务。我们也期待看到社区如何使用 IO Processors 来突破 vLLM 可能性的边界。我们还计划继续与 vLLM 社区合作并做出贡献，以支持更多的多模态模型和端到端用例。

欢迎随时提供贡献、反馈和想法！

要开始使用 IO Processor 插件，请查看文档并探索示例。有关 IBM TerraTorch 的更多信息，请访问此处。

致谢

我们感谢 vLLM 社区成员帮助改进我们的贡献。特别感谢 Cyrus Leung 在塑造将 vLLM 扩展到文本生成之外的整体概念方面给予的支持。最后，我们感谢 IBM 的 TerraTorch 团队，特别是 Paolo Fraccaro 和 Joao Lucas de Sousa Almeida，感谢他们帮助在 vLLM 中集成了通用的 TerraTorch 后端。