Prhub

#50411 [Model] Fused mm preprocess normalisation on the Device

原始 PR 作者 noooop 合并时间 2026-08-06 18:11 文件变更 13 提交数 43 评论 47 代码增减 +359 / -11

执行摘要

图像归一化搬上 GPU,Qwen2VL 预处理提速近一倍

PR body 中作者 noooop 给出两条动机:normalize 在 GPU 上计算时可融合进后续算子,即使不融合也几乎不耗时;图像数据从入口到引擎再到 GPU 全程以 uint8 传输,直到输入归一化前才转成 bf16,把归一化搬上设备与现有传输路径天然契合,且可推广到所有视觉模型。实测数据:Qwen2VLImageProcessor 对 1024x1024 图像预处理速度接近翻倍(12 ms → 6 ms);vllm bench 平均 TTFT 从 3809.70 ms 降至 3459.52 ms(Rust 前端基线 3078.26 ms);PPL 测试开/关该特性与 Transformers 的差异完全一致(约 0.036%)。

值得精读。设计上有三个可借鉴点:用 eval 模式 batch_norm 折叠 rescale + normalize 为一次仿射变换并自动检测恒等路径;归一化参数「配置文件 → HF 实例 → 开关默认值 → 快速失败」的分级解析回退链;模型能力标记 + 前端能力约束的开关解析模式(_resolve_mm_device_do_normalize)。review 讨论中关于「读处理器类默认值不严谨」与「fail-fast vs 静默回退」的取舍对后续多模态预处理开发有参考价值。

讨论亮点

review 中围绕参数读取来源、fail-fast 与静默回退、模块形态、Rust 前端兼容性展开了多轮交锋。DarkLight1337 质疑直接读 processor 类默认值不严谨,Isotr0py 建议改读 preprocessor_config.json 并指出自定义 remote code 难以导入;noooop 主张参数缺失时快速失败以避免用错 mean/std(Qwen2-VL 用 OPENAI_CLIP、Qwen3.5 用 IMAGENET_STANDARD)。DarkLight1337 还要求对被强制关闭的场景输出告警、构造器增加 dtype 参数,并在 issue 中质疑端到端收益、指出部分收益来自引擎批处理 MM 输入。所有讨论最终落地:三级参数解析回退链、保留 assert、_resolve_mm_device_do_normalize 中打 warning、FusedInputNorm 增加 dtype 与 identity 快速路径,两位 reviewer 最终均 APPROVED。

实现拆解

  1. 配置契约扩展:MultiModalConfig 新增 mm_device_do_normalize 字段(默认 True)并纳入 compute_hash;ModelConfig 新增同名 InitVar 与 _resolve_mm_device_do_normalize 决策方法——未显式指定时跟随模型能力标记(Rust 前端强制 False),用户显式开启但不被支持时回退并输出 warning;EngineArgs 暴露 --mm-device-do-normalize CLI 参数(默认 None 走模型能力推断)。
  2. 处理器参数读取重构:vllm/transformers_utils/processor.py 将 get_processor_cls_name_from_config 拆分为返回完整配置 dict 的 get_processor_config 与原有包装函数,供 FusedInputNorm.from_model_config 读取 preprocessor_config.json 中的 image_mean、image_std、rescale_factor、do_rescale、do_normalize;任一字段缺失时回退到 HF image_processor 实例属性。
  3. 核心归一化模块:vllm/model_executor/models/vision.py 新增 FusedInputNorm,将 (input * rescale_factor - mean) / std 折叠为一次 F.batch_norm(weight = 1/std、bias = -mean/std、running_var = 1、eps = 0),并自动检测恒等场景跳过计算。
  4. 模型接入:Qwen2VLForConditionalGeneration 与 Qwen2_5_VLForConditionalGeneration 声明 supports_mm_device_do_normalize = True,在 _mark_tower_model 上下文内构造 self.input_norm,并在 _process_image_input / _process_video_input 中对图像与视频帧 pixel_values 做设备端归一化,image_embeds 分支统一 .type(self.visual.dtype)。
  5. 配套改动:MultiModalConfig.merge_mm_processor_kwargs 在开启时强制注入 do_normalize=False、do_rescale=False(推理期 inference_kwargs 仍可覆盖);测试新增 test_mm_device_do_normalize 数值一致性用例与 Qwen 系列 PPL 双模式参数化;文档更新 docs/design/mm_processing.md。
文件 模块 状态 重要度
vllm/model_executor/models/vision.py 视觉模型 modified 8.67
vllm/config/model.py 配置层 modified 7.34
vllm/config/multimodal.py 配置层 modified 5.88
vllm/model_executor/models/qwen2_vl.py 视觉模型 modified 5.98
vllm/model_executor/models/qwen2_5_vl.py 视觉模型 modified 5.52
vllm/transformers_utils/processor.py 处理器工具 modified 6.36
vllm/engine/arg_utils.py 启动参数 modified 5.52
vllm/model_executor/models/interfaces.py 模型接口 modified 5.0
tests/models/multimodal/processing/test_qwen2_vl.py 测试 modified 5.82
tests/models/multimodal/generation_ppl_test/test_qwen.py 测试 modified 5.01
docs/design/mm_processing.md 设计文档 modified 3.05

关键符号

FusedInputNorm.from_model_config FusedInputNorm.forward FusedInputNorm.identity _resolve_mm_device_do_normalize MultiModalConfig.merge_mm_processor_kwargs MultiModalConfig.compute_hash get_processor_config get_processor_cls_name_from_config Qwen2VLForConditionalGeneration._process_image_input Qwen2VLForConditionalGeneration._process_video_input Qwen2_5_VLForConditionalGeneration._process_image_input Qwen2_5_VLForConditionalGeneration._process_video_input test_mm_device_do_normalize

关键源码片段

vllm/config/model.py configuration

新增 mm_device_do_normalize InitVar 与 _resolve_mm_device_do_normalize 决策方法,处理默认值推断、Rust 前端强制关闭与模型能力校验,是整套开关逻辑的入口。

def _resolve_mm_device_do_normalize(
    self, mm_device_do_normalize: bool | None
) -> bool:
    if mm_device_do_normalize is None:
        # 未显式指定时:Rust 前端暂不支持设备端归一化,强制关闭;
        # 否则跟随模型能力标记 supports_mm_device_do_normalize 决定默认值。
        if envs.VLLM_USE_RUST_FRONTEND:
            logger.debug(
                "VLLM_USE_RUST_FRONTEND is set. "
                "Rust frontend does not currently support mm_device_do_normalize, "
                "forcing mm_device_do_normalize = False."
            )
            mm_device_do_normalize = False
        else:
            mm_device_do_normalize = self._model_info.supports_mm_device_do_normalize
            logger.debug(
                "mm_device_do_normalize is %s by default.",
                "enabled" if mm_device_do_normalize else "disabled",
            )
    else:
        # 用户显式开启但 Rust 前端或模型不支持时,强制回退并告警,
        # 避免用户困惑为什么开关不生效。
        if mm_device_do_normalize and envs.VLLM_USE_RUST_FRONTEND:
            logger.warning(
                "VLLM_USE_RUST_FRONTEND is set. "
                "Rust frontend does not currently support mm_device_do_normalize, "
                "forcing mm_device_do_normalize = False."
            )
            mm_device_do_normalize = False
​
        if (
            mm_device_do_normalize
            and not self._model_info.supports_mm_device_do_normalize
        ):
            logger.warning(
                "Model does not support mm_device_do_normalize, "
                "forcing mm_device_do_normalize = False."
            )
            mm_device_do_normalize = False
​
        logger.debug(
            "mm_device_do_normalize is %s.",
            "enabled" if mm_device_do_normalize else "disabled",
        )
​
    return mm_device_do_normalize
vllm/model_executor/models/qwen2_vl.py core-logic

Qwen2VLForConditionalGeneration 声明 supports_mm_device_do_normalize = True,在 _mark_tower_model 内构造 input_norm,并在图像 / 视频处理路径接入设备端归一化,是本特性第一个落地模型。

# __init__ 中,在 _mark_tower_model 上下文内构造:
with self._mark_tower_model(vllm_config, {"image", "video"}):
    self.visual = Qwen2VisionTransformer(
        config.vision_config,
        norm_eps=getattr(config, "rms_norm_eps", 1e-6),
        quant_config=quant_config,
        prefix=maybe_prefix(prefix, "visual"),
    )
    self.input_norm = FusedInputNorm.from_model_config(self.model_config)
​
​
def _process_image_input(
    self, image_input: Qwen2VLImageInputs
) -> tuple[torch.Tensor, ...]:
    grid_thw = image_input["image_grid_thw"]
    assert grid_thw.ndim == 2
​
    if image_input["type"] == "image_embeds":
        image_embeds = image_input["image_embeds"].type(self.visual.dtype)
    else:
        pixel_values = image_input["pixel_values"]
        # 此前在 CPU 上完成的 rescale + normalize 被搬到此处,
        # 输入仍是 uint8 语义的原始像素,FusedInputNorm 在 GPU 上完成等价变换。
        pixel_values = self.input_norm(pixel_values, self.visual.dtype)
​
        if self.use_data_parallel:
            return run_dp_sharded_mrope_vision_model(
                self.visual, pixel_values, grid_thw.tolist(), rope_type="rope_3d"
            )
        else:
            image_embeds = self.visual(pixel_values, grid_thw=grid_thw)
​
    # 按每张图的 patch 数切分拼接后的 embedding
    merge_size = self.visual.spatial_merge_size
    sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()
    return image_embeds.split(sizes)

视频路径 _process_video_input 采用相同模式,对 pixel_values_videos 同样调用 self.input_norm。

评论区精华

归一化参数读取来源:processor 类 vs preprocessor_config.json vs HF 实例 设计

DarkLight1337 指出最初实现直接读 processor 类默认值不严谨:“you're just reading the default values from the processor class, instead of the actual HF processor instance that's being loaded from the config”。Isotr0py 补充自定义 remote code 的 processor 难以导入,建议改读 preprocessor_config.json;noooop 承认同架构不同模型可能配置不同。

结论:最终实现为分三级解析:优先 get_processor_config 读取配置文件,字段缺失再回退到 HF image_processor 实例属性,仍缺失则按 do_rescale/do_normalize 开关套用恒等默认值,最后 assert 快速失败。 · 已解决

参数缺失时快速失败还是静默回退 设计

DarkLight1337 质疑 assert 的必要性:“Can't we just return nn.Identity or default to a rescale/normalize factor of 1”。noooop 列举 OPENAI_CLIP_MEAN、IMAGENET_STANDARD_MEAN 等常用常量,主张 fail fast 以免用错 mean/std,并指出 Qwen2-VL 用 OPENAI_CLIP 而 Qwen3.5 用 IMAGENET_STANDARD。

结论:保留 assert,但同时扩展默认值兜底逻辑(do_rescale=False 时 factor 置 1、do_normalize=False 时 mean/std 取恒等值),只在参数完全无法解析时才失败。 · 已解决

模块形态与放置位置(nn.Identity 与 _mark_tower_model) 设计

Isotr0py 建议 input_norm 用 nn.Identity() 而非 None,并移到 _mark_tower_model 上下文内;DarkLight1337 进一步建议构造器增加 dtype 参数、forward 内使用 grid_thw.to(self.dtype) 并正确注册缓冲。

结论:全部采纳,最终演化为 FusedInputNorm nn.Module,自动检测恒等场景跳过计算。 · 已解决

Rust 前端兼容性与端到端收益验证 设计

Isotr0py 要求补充 Rust 前端对比的 e2e TTFT 数据,因为 Rust 前端也有 normalize+to_tensor 融合;noooop 提供 VLLM_USE_RUST_FRONTEND=1 基线 3078.26 ms,并说明该优化与 Rust 前端正交、后续可加支持。

结论:Rust 前端暂不支持,VLLM_USE_RUST_FRONTEND 下强制 mm_device_do_normalize=False 并打 warning,形成已知限制。 · 已解决(已知限制)

被强制关闭时需提示用户 documentation

DarkLight1337 要求:“add a warning if the user enables this for a configuration that's not supported, so that the user doesn't get confused about why the setting isn't working”。

结论:采纳,_resolve_mm_device_do_normalize 对不支持场景输出 logger.warning。 · 已解决

merge_mm_processor_kwargs 强制覆写 do_normalize/do_rescale 的边界情况 正确性

DarkLight1337 提出边缘场景:processor 类默认 do_rescale=False 或 do_normalize=False、或这些值在 HF 配置或原始 mm_processor_kwargs 中设置时应如何处理。

结论:merge_mm_processor_kwargs 开启时强制注入 do_normalize=False、do_rescale=False,但推理期 inference_kwargs 仍有更高优先级可覆盖;from_model_config 按开关应用默认值兜底。 · 已解决

端到端收益质疑与批处理 MM 输入的方向 question

DarkLight1337 在 issue 中质疑:“How much does this actually improve E2E though?”,并指出部分收益来自引擎在 forward 前批处理 MM 输入而非单纯 GPU 计算;他提到一直想做 batched MM processing 但被 Rust 前端等工作占用时间。

结论:noooop 用 TTFT 数据回应;批处理 MM 输入被确认为后续优化方向,与本 PR 正交可叠加。 · 已解决(讨论性)

风险与影响

  1. 行为变更范围:mm_device_do_normalize 默认 True,凡声明 supports_mm_device_do_normalize = True 的模型(当前为 Qwen2-VL、Qwen2.5-VL)默认启用设备端归一化,数值等价但执行路径有差异,其他架构默认关闭,回归面可控。
  2. Rust 前端不兼容:VLLM_USE_RUST_FRONTEND 下强制关闭并告警,是明确的已知限制;双前端行为不一致可能造成用户困惑。
  3. 配置契约扩展:新字段纳入 MultiModalConfig.compute_hash,影响计算图 hash;EngineArgs、ModelConfig、MultiModalConfig 三处需同步维护。
  4. 参数解析失败风险:from_model_config 中的 assert 对个别缺失 image_mean/image_std 的模型或自定义 remote code processor 可能导致启动失败,虽有「配置文件 → HF 实例 → 开关默认值」三级兜底但覆盖面未完全验证。
  5. 数值一致性依赖:等价性依赖 F.batch_norm 的 eps=0 语义与 float32 计算顺序,仅 Qwen 两代模型有 assert_close 与 PPL 双模式测试,其他视觉模型缺测试覆盖。
  6. GPU 负载转移:预处理从 CPU 移到 GPU,GPU 侧增加一次轻量 batch_norm(可与后续算子融合),CPU 释放约一半预处理耗时,对 CPU 受限的视觉部署有利。

对用户:Qwen2-VL / Qwen2.5-VL 用户默认获得更快的图像预处理与更低平均 TTFT(约 9% 改善),PPL 数值不变;可通过 --mm-device-do-normalize 显式关闭。使用 Rust 前端的用户该开关自动失效。对系统:缓解 CPU 预处理瓶颈,将归一化计算转移至 GPU(开销可融合,几乎为 0)。对团队:多模态配置契约扩展(EngineArgs → ModelConfig → MultiModalConfig + hash),新增模型若想启用该优化需声明 supports_mm_device_do_normalize = True,并保持 processor 配置可解析。

默认开启改变视觉模型行为 Rust 前端暂不支持 配置 hash 契约扩展 归一化参数断言可能误伤 缺少其他视觉模型覆盖

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论