Prhub

#53106 Reduce `AutoWeightsLoader` kwargs

原始 PR 作者 hmellor 合并时间 2026-08-20 23:40 文件变更 84 提交数 7 评论 4 代码增减 +578 / -569

执行摘要

统一权重跳过逻辑至 WeightsMapper,精简 AutoWeightsLoader 参数

PR body 明确指出:向 AutoWeightsLoader 传 kwargs 是移除模型代码中 load_weights 方法的障碍;skip_prefixes 与 skip_substrs 的行为都能用 WeightsMapper 表达,因此把权重跳过逻辑统一到 WeightsMapper,既更接近“不需要 load_weights”的目标,也避免维护两套机制、减少样板代码。

值得精读。核心是 vllm/model_executor/models/utils.py 的 mapper 设计:None 映射的统一语义、get_rename_mapper 对“命名消费者”的过滤,以及 hf_to_vllm_mapper 类属性 + load_weights 只传 mapper 的迁移模板。维护特定模型的团队应检查自身模型文件是否已跟随迁移,避免残留 skip_prefixes 调用导致 API 破坏。

讨论亮点

本 PR 的 review 评论很少,无实质技术争论。claude[bot] 提示这是 fork PR,自动化评审被禁用;DarkLight1337 直接批准。核心设计决策(用 None 映射表达丢弃、get_rename_mapper 过滤 None)均由作者 hmellor 在文档字符串中阐明,未出现未解决的疑虑。

实现拆解

  1. 核心契约变更:在 vllm/model_executor/models/utils.py 中移除 AutoWeightsLoader.__init__ 的 skip_prefixes / skip_substrs 参数,将默认跳过的 rotary 冗余张量改为类属性 REMOVE_UNUSED_ROTARY_EMBEDS_MAPPER(WeightsMapper 的 orig_to_new_substr 映射到 None);同时把 WeightsMapper.get_unstacked_mapper 重构为 get_rename_mapper,除清空 orig_to_new_stacked 外还过滤所有值为 None 的映射,避免“不加载”语义泄漏给 LoRA 名称解析与量化配置层列表等消费者。
  2. 模型侧迁移模式:各模型的 load_weights 不再传 skip_prefixes / skip_substrs,改为在类上声明 hf_to_vllm_mapper(如 WeightsMapper(orig_to_new_prefix={... : None})),再以 loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) 调用。典型示例包括 interns2_preview.py(合并父类 mapper 并补充 time_series 前缀丢弃)、internvl.py 与 skyworkr1v.py(用 orig_to_new_prefix 表达原 skip_prefixes 列表)、aria.py、glm4.py、qwen3_dspark.py、siglip.py、minicpmv.py 等。
  3. Transformers 后端适配:transformers/legacy.py 的 LegacyMixin 不再扩展 skip_prefixes / skip_substrs,改为直接更新 self.hf_to_vllm_mapper.orig_to_new_prefix 与 orig_to_new_substr(如 position_ids、score.bias 映射为 None)。
  4. 测试配套:tests/models/test_utils.py 删除针对 skip_prefixes / skip_substrs 的 test_module_skip_prefix / test_module_skip_substr,新增 test_get_rename_mapper_keeps_only_renames 验证 None 映射被过滤;跳过行为改由 mapper 语义测试覆盖。
  5. 无配置、schema 或部署类配套改动;提交历史显示模型分批迁移(process some models / some more / rest of models),最后合并 main 解决冲突。
文件 模块 状态 重要度
vllm/model_executor/models/utils.py 加载器 modified 7.98
tests/models/test_utils.py 测试 modified 7.15
vllm/model_executor/models/interns2_preview.py 多模态模型 modified 6.62
vllm/model_executor/models/transformers/legacy.py 兼容层 modified 6.1

关键符号

get_rename_mapper AutoWeightsLoader.__init__ InternS2PreviewForConditionalGeneration.load_weights InternVLChatModel.load_weights LegacyMixin.__init__

关键源码片段

vllm/model_executor/models/utils.py data-contract

核心变更文件:移除 AutoWeightsLoader 的 skip_prefixes / skip_substrs 参数,新增 REMOVE_UNUSED_ROTARY_EMBEDS_MAPPER,并将 get_unstacked_mapper 重构为过滤 None 映射的 get_rename_mapper,是整个 PR 的契约中心。

class AutoWeightsLoader:
    """自动权重加载器:遍历 checkpoint 权重一次,自动分发给子模块或参数。
​
    子模块可通过定义 `load_weights` 覆盖加载逻辑,参数可通过定义
    `weight_loader` 覆盖。详细加载信息可用环境变量
    `VLLM_LOGGING_LEVEL=DEBUG` 查看。
    """
​
    # ColossalAI 早期版本或 GPTQModel 量化产生的 checkpoint 可能包含
    # rotary 相关的冗余张量,vLLM 不会使用。统一用 mapper 表达 " 丢弃 ",
    # 而不再给 AutoWeightsLoader 增加参数。
    REMOVE_UNUSED_ROTARY_EMBEDS_MAPPER = WeightsMapper(
        orig_to_new_substr={
            "rotary_pos_emb.inv_freq": None,
            "rotary_emb.inv_freq": None,
            "rotary_emb.cos_cached": None,
            "rotary_emb.sin_cached": None,
        }
    )
​
    def __init__(
        self,
        module: nn.Module,
        *,
        ignore_unexpected_prefixes: list[str] | None = None,
        ignore_unexpected_suffixes: list[str] | None = None,
    ) -> None:
        super().__init__()
        self.module = module
        self.ignore_unexpected_prefixes = ignore_unexpected_prefixes or []
        self.ignore_unexpected_suffixes = ignore_unexpected_suffixes or []
​
        # Weight tying 会让两个 qualname 指向同一个 nn.Parameter
        # (例如 lm_head.weight 与 model.embed_tokens.weight)。
        # 两次加载会写入同一块存储,因此只加载遍历中先遇到的名称,
        # 其余按别名关系跳过。
        self.aliased_params = _get_tied_embedding_params(module)
        self._skipped_aliases = dict[str, str]()
        self._loaded_params_are_complete = True
​
    def get_rename_mapper(self) -> "WeightsMapper":
        """只保留"改名"语义的 mapper 变体。
​
        只按名字消费权重、并不实际加载的调用方(LoRA 名称解析、
        量化配置的层列表)需要保留原始名字,例如 `q_proj`
        而不是被改写成堆叠后的 `qkv_proj`。因此这里清空
        堆叠映射;同时把值为 `None` 的映射过滤掉,因为
        "不加载该权重"对这类消费者没有意义,保留会悄悄缩短
        量化配置的忽略列表,或让 LoRA 名称解析失败。
        """
        remove_none = lambda d: {k: v for k, v in d.items() if v is not None}
        return replace(
            self,
            orig_to_new_regex=remove_none(self.orig_to_new_regex),
            orig_to_new_substr=remove_none(self.orig_to_new_substr),
            orig_to_new_stacked={},
            orig_to_new_prefix=remove_none(self.orig_to_new_prefix),
            orig_to_new_suffix=remove_none(self.orig_to_new_suffix),
        )
vllm/model_executor/models/interns2_preview.py data-contract

展示了最有代表性的迁移模式:合并父类 mapper 并用 orig_to_new_prefix 表达原 skip_prefixes,是其他模型迁移的模板。

class InternS2PreviewForConditionalGeneration(Qwen3_5MoeForConditionalGeneration):
    # `mtp.` 前缀已由父类 `Qwen3_5ForConditionalGeneration` 的 mapper 丢弃,
    # 这里只补充时间序列相关前缀的丢弃规则,通过 `|` 合并父类 mapper。
    hf_to_vllm_mapper = Qwen3_5MoeForConditionalGeneration.hf_to_vllm_mapper | (
        WeightsMapper(
            orig_to_new_prefix={"model.time_series.": None, "time_series.": None}
        )
    )
​
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        # 不再向 AutoWeightsLoader 传 skip_prefixes,
        # 跳过逻辑全部收敛到 mapper。
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/transformers/legacy.py data-contract

Transformers 后端的 LegacyMixin 从扩展 skip 列表改为直接修改 hf_to_vllm_mapper 字典,体现了对类属性 mapper 的可变更新模式,值得关注共享状态风险。

class LegacyMixin:
    def __init__(self, *, vllm_config: "VllmConfig", prefix: str = ""):
        super().__init__(vllm_config=vllm_config, prefix=prefix)
​
        # 丢弃不支持的输出 embedding 层,直接写入 mapper 的前缀映射。
        # 注意:hf_to_vllm_mapper 是类属性,这里的 update 会修改类级
        # 字典,当前每个实例更新内容一致,因此行为保持不变;若未来
        # 需要按实例区分,需改为实例级副本。
        self.hf_to_vllm_mapper.orig_to_new_prefix.update(
            {
                "model.lm_head.": None,
                "model.predictions.": None,
                "model.qa_outputs.": None,
                "model.embeddings_project.": None,
                "model.discriminator_predictions.": None,
            }
        )
        # 部分 encoder 模型的 checkpoint 带 position_ids buffer,
        # vLLM 总是显式传入 position_ids,因此直接丢弃该 buffer。
        self.hf_to_vllm_mapper.orig_to_new_substr["position_ids"] = None
        # 分类器偏置同理,vLLM 不使用该偏置。
        self.hf_to_vllm_mapper.orig_to_new_substr["score.bias"] = None

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

  1. 语义等价性风险:原 skip_prefixes 是前缀匹配、skip_substrs 是子串匹配,迁移后分别对应 orig_to_new_prefix / orig_to_new_substr 映射到 None。84 个模型文件中任一处混淆前缀/子串语义都会静默改变权重加载行为,例如 InternS2Preview 的 model.time_series. 与 time_series. 前缀、legacy.py 的 position_ids 子串。
  2. 共享可变类属性风险:hf_to_vllm_mapper 是类属性,legacy.py 通过 self.hf_to_vllm_mapper.orig_to_new_prefix.update(...) 修改类级字典,所有实例共享同一 mapper;当前每个实例更新内容一致,行为不变,但后续若需按实例差异化配置则会产生污染。REMOVE_UNUSED_ROTARY_EMBEDS_MAPPER 同理(通常只读)。
  3. 回归面:改动横跨多模态、encoder、MoE 等大量模型,且新增 mapper 逻辑仅有 test_utils.py 单测覆盖,模型级行为等价性依赖 CI 全量验证;同时 84 个文件的 import 变化会抬高后续合入其他 PR 的冲突概率。
  4. 消费者行为变化:get_unstacked_mapper 升级为 get_rename_mapper 后会额外过滤 None 映射,LoRA 名称解析与量化配置层列表的行为随之改变,属于有意的语义修正,但依赖下游测试确认。

用户侧:无可感知的变化,权重加载结果应保持一致。系统侧:权重跳过收敛为 WeightsMapper 单一机制,AutoWeightsLoader 仅保留 ignore_unexpected_* 参数,加载路径更简洁。团队侧:模型开发者新增模型时只需声明 hf_to_vllm_mapper 即可表达跳过规则,无需在 load_weights 内维护 skip 列表,向“无需 load_weights”目标迈出实质一步;影响面覆盖绝大多数模型文件,后续相关 PR 需注意合并冲突与行为回归。

核心权重加载路径变更 涉及 84 个文件大范围迁移 类属性 mapper 可变共享风险 行为等价性依赖 CI 验证 skip 语义需与 mapper 语义对齐

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论