跳到内容

torch.compile 与多模态编码器

现在可以将 torch.compile 应用于 vLLM 中的多模态编码器和其他 nn 模块，包括像 LLaMA 4、Qwen-VL 等类似的基于编码器架构的视觉语言模型。

本文档介绍了 torch.compile 集成在 vLLM 中如何作用于多模态编码器，以及如何将该装饰器应用于新模型以提高性能。

注意

有关 vLLM 中 torch.compile 集成的常规信息，请参阅 torch.compile 设计文档。

概述

我们最近启用了 @support_torch_compile 装饰器，使其能够支持模型类型内的多个 nn 模块组件；这使得可以为多模态编码器开启编译，从而为堆栈的其他组件带来性能提升。

当应用于 Qwen2_5_vl 的视觉模块时，我们观察到端到端性能提升了约 4.5%，但编译时间有所增加。

此功能默认关闭，但当模型具有 @support_torch_compile 装饰器时，可以通过在编译配置中设置 compile_mm_encoder: true 来启用。

编译对多模态组件的工作原理

启用 API

为了编译像编码器这样的多模态组件，我们遵循与 LLM 文本骨干网络相同的机制，并增加了一些额外的框架支撑：

  1. @support_torch_compile 装饰器应包含 enable_if=should_torch_compile_mm_encoder。这将通过我们的 compile_mm_encoder 配置来控制编译的开启。

  2. 对于编码器组件，@support_torch_compile 装饰器应包含 is_encoder=True。这对于编译范围集成（参见“编译范围集成”）是必需的。该装饰器会自动使用类名作为缓存目录的前缀，从而避免独立编译的子模块（例如视觉编码器组件与文本骨干网络）之间的冲突。

CompilationConfig

除了 compile_mm_encoder: true 外，多模态编码器将继承与文本 LLM 相同的编译配置。未来我们可能会扩展此配置以支持更多功能。

将 torch.compile 应用于新的多模态模型/组件

要将 support_torch_compile 应用于新的通用 nn.Module，建议遵循 debug_vllm_compile 中的相同步骤；这包括：

  1. 先对较小的模块（如基本 MLP 层）应用 support_torch_compile，然后逐渐扩展到更通用的模块，直到达到良好的性能平衡点。

  2. 利用 tlparse 来识别并消除重新编译和图中断的源头。

  3. 使用 dynamic_arg_dims 和适当的 dynamic_shapes_config 来处理动态性。

常见陷阱

VllmBackend 功能支持

编译范围

torch.compile 集成将尝试依赖 max_batch_size 来推断动态形状的编译范围；然而，对于编码器中使用的模块，由于编码器可能看到的输入形状范围未指定，该形状可能难以推断。因此，我们依靠 @support_torch_compile 装饰器中的 is_encoder=True 来提醒 torch.compile 该范围无法推断，并默认设置为范围 (1, MAX_INT)。

注意

未来我们可能会尝试收窄此范围以获得更好的性能。

Cudagraphs

我们尚未探索针对多模态编码器集成 CUDAGraphs 的编译；目前其行为未定义。

故障排除

视觉编码器中的图中断

某些视觉编码器操作可能会导致图中断。识别它们的方法如下：

TORCH_LOGS="+dynamo" vllm serve <MODEL>

多模态模型中图中断的常见原因：

  • 动态图像尺寸：使用 dynamic_shapes_config 来处理可变分辨率。
  • 无法追踪的操作：某些操作（如 to_list）可能不受 Dynamo 支持。
  • 条件处理：基于图像属性的依赖于数据的分支。

编译错误

如果多模态模型的编译失败：

  1. 禁用并测试：首先验证模型在不开启编译的情况下是否能正常工作。

    VLLM_TORCH_COMPILE_LEVEL=0 vllm serve <model> --compilation-config='{"compile_mm_encoder":"false"}'

  2. 检查日志：启用调试日志以查看编译详情。

    VLLM_LOGGING_LEVEL=DEBUG vllm serve <model> --compilation-config='{"compile_mm_encoder":"true"}'

  3. 报告问题：如果您发现 Bug，请在 GitHub 上提交 issue。

参见