执行摘要

• 一句话：修正 BERT 模型的权重层别名映射
• 推荐动作：该 PR 展示了处理模型权重加载时的通用模式：通过 WeightsMapper 修正前缀并添加 params_dict 不在时的跳过逻辑。值得模型加载相关开发者阅读；deep_gemm_warmup.py 的修复也对跨平台兼容性有参考价值。

功能与动机

当加载 HuggingFace 的 BERT 风格模型 checkpoint 时，ModernBertModel 和 BertModel 的参数名称因别名映射不完整（如 model.layers. 未映射）或 checkpoint 包含任务特定权重（如 cls.predictions.*）而无法正确加载，导致 KeyError。本 PR 修正这些问题以支持更多 BERT 系列模型的加载。

实现拆解

1. 扩展 ModernBertModel 的权重映射：在 modernbert.py 中修改 hf_to_vllm_mapper，增加 "model.layers." 和 "layers." 到 "encoder_layer.layers." 的映射，并增加 "model." 到空字符串的映射，以消除前缀差异。
2. 添加未知权重跳过逻辑：在 load_weights 方法中增加 if name not in params_dict: continue 检查，使加载器能跳过 checkpoint 中包含但模型未定义的权重（如 MLM 头部），避免 KeyError。
3. 修复 DeepGEMM 预热函数调用顺序：在 deep_gemm_warmup.py 中将 get_mk_alignment_for_contiguous_layout() 调用移动到 FP8 线性方法检查之后，以防止在不支持 FP8 DeepGEMM 的平台上（如 TPU）提前触发错误。

关键文件：

• vllm/model_executor/models/modernbert.py（模块 模型执行；类别 source；类型 data-contract；符号 hf_to_vllm_mapper, load_weights）: 核心修复文件，修改权重映射和加载逻辑以支持 ModernBERT 和其他 BERT 模型
• vllm/model_executor/warmup/deep_gemm_warmup.py（模块 模型执行；类别 source；类型 bugfix；符号 _fp8_linear_may_use_deep_gemm）: 修复函数调用顺序以避免在非 FP8 环境下崩溃

关键符号：load_weights, _fp8_linear_may_use_deep_gemm

关键源码片段

vllm/model_executor/models/modernbert.py

核心修复文件，修改权重映射和加载逻辑以支持 ModernBERT 和其他 BERT 模型

class ModernBertModel(nn.Module):
    # 扩展 hf_to_vllm_mapper，覆盖更多 HuggingFace 前缀变体
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            "model.layers.": "encoder_layer.layers.", # 处理 "model.layers.X" 前缀
            "layers.": "encoder_layer.layers.", # 处理 "layers.X" 前缀
            "model.": "", # 去除开头的 "model."
        }
    )
​
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        weights = self.hf_to_vllm_mapper.apply(weights)
        params_dict = dict(self.named_parameters())
        loaded_params: set[str] = set()
        for name, loaded_weight in weights:
            if name.endswith(".bias") and name not in params_dict:
                continue # 跳过缺少的 bias（如 MLM 头部的 bias）
            if name not in params_dict:
                continue # 跳过模型未定义的权重（如 cls.predictions.weight）
            param = params_dict[name]
            weight_loader = getattr(param, "weight_loader", default_weight_loader)
            weight_loader(param, loaded_weight)
            loaded_params.add(name)
        return loaded_params

vllm/model_executor/warmup/deep_gemm_warmup.py

修复函数调用顺序以避免在非 FP8 环境下崩溃

def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool:
    # 首先检查是否为 FP8 线性层（若否，直接返回 False）
    if not (
        isinstance(module, LinearBase)
        and isinstance(module.quant_method, Fp8LinearMethod)
        and not isinstance(module.quant_method, Mxfp8OnlineLinearMethod)
        and getattr(module.quant_method, "block_quant", False)
        and not getattr(module.quant_method, "use_marlin", True)
    ):
        return False
    # 在确认是 FP8 线性层后，再调用 DeepGEMM 相关函数
    # （避免在不支持 DeepGEMM 的环境下提前崩溃）
    block_size = get_mk_alignment_for_contiguous_layout()[0]
    # ... 后续校验

评论区精华

审核者 mgoin 询问 deep_gemm_warmup.py 中移动 block_size 计算的目的。提交者 ap9272 解释该调用在不支持 DeepGEMM FP8 的环境下会失败，但原来的代码顺序在检查之前执行了该调用，导致异常。移动后确保先检查是否 FP8 线性层，再计算 block_size。该澄清获得认可并最终批准。

• deep_gemm_warmup.py 中 block_size 计算位置的改动原因 (question): ap9272 解释：原顺序在不支持 DeepGEMM FP8 时会失败，移动后确保在检查通过后再调用相关函数

风险与影响

• 风险：风险较低。修改仅限于两个文件：modernbert.py 的映射和跳过逻辑增强了对不同 checkpoint 的兼容性，不会破坏现有模型；deep_gemm_warmup.py 的改动只影响函数执行顺序，对 FP8 场景无副作用。但未新增测试用例，可能遗漏边缘情况。
• 影响：直接影响使用 ModernBERT 和 BERT 模型的用户，使其能够加载之前因别名或不匹配权重而失败的 checkpoint。间接影响低，因为改动范围小且属于兼容性增强。团队内其他模型的权重加载模式可能参考此映射模式。
• 风险标记：缺少测试覆盖

关联脉络

• 暂无明显关联 PR