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

引言

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

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

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

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

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

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

与文本模型不同，地理空间基础模型（通常实现为视觉 Transformer）不需要标记解码，即它们不需要将输出标记转换为文本。相反，给定一张输入图像，单次推理即可产生原始模型输出，然后将其后处理为输出图像。此外，有时输入图像需要被分割并批量处理成若干子图像或补丁（patches）。然后将这些补丁送入模型进行推理，每个补丁生成的输出图像再被缝合在一起形成最终的输出图像。

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

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

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

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

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

到目前为止一切顺利！不过，这只让我们完成了一半的目标。

通过上述集成，我们确实可以服务地理空间基础模型——但仅限于张量到张量（tensor-to-tensor）的格式。用户在将张量发送到 vLLM 实例之前，仍必须将其图像预处理为张量格式。同样，原始张量输出的后处理也必须在 vLLM 之外进行。其影响是：没有一个端点可以让用户发送一张图片并直接取回一张图片。

这个问题之所以存在，是因为在我们修改之前，vLLM 仅部分支持输入数据的预处理和模型输出的后处理。具体而言，多模态输入数据的预处理只能通过 Transformers 库中可用的处理器进行。然而，Transformers 处理器通常仅支持标准数据类型，无法处理更复杂的数据格式，如 GeoTIFF（带有丰富地理空间元数据的图像文件）。此外，在输出处理方面，vLLM 仅支持反向分词（de-tokenization）为文本或对模型隐藏状态应用池化器——无法进行其他输出处理。

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

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

使用 vLLM IO 处理器插件

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

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

服务实例启动后，在访问 /pooling 端点时，预处理和后处理将自动应用于模型输入和输出。在此阶段，IO 处理器仅适用于池化模型，但未来我们预计也会集成其他端点。

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

可以使用 vLLM 通过 TerraTorch 后端服务的一个模型类示例是用于洪水检测的 Prithvi。针对 Prithvi 地理空间基础模型的完整插件示例可在此处找到。

Prithvi IO 处理器插件

为了说明 IO 处理器插件方法的灵活性，下面的伪代码展示了 Prithvi IO 处理器预处理和后处理的主要步骤。我们想要强调的是数据特定转换与模型推理数据之间的解耦。这为理想情况下任何模型和任何输入/输出数据类型，甚至根据消耗数据的下游任务对同一模型输出应用多个插件留出了空间。

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 处理器插件。

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 处理器插件扩展到更多的 TerraTorch 模型和模态及其他领域，使安装过程更加无缝。从长远来看，我们设想由 IO 处理器驱动的视觉语言系统、结构化推理智能体和多模态流水线，全部由同一个 vLLM 栈提供服务。我们也期待看到社区如何使用 IO 处理器来突破 vLLM 的可能性边界。我们还计划继续与 vLLM 社区合作并做出贡献，以启用更多的多模态模型和端到端用例。

随时欢迎您的贡献、反馈和建议！

要开始使用 IO 处理器插件，请查看文档并探索示例。有关 IBM TerraTorch 的更多信息可在此处获得。

致谢

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