【Bug已解决】[Bug]: key error: ‘Layer.34.mlp.experts.gate_up

【Bug已解决】[Bug]: key error: 'Layer.34.mlp.experts.gate_up_proj' 解决方案

一、现象长什么样

在用 vLLM 加载一个 MoE（混合专家）模型的权重时，启动阶段直接抛出一个KeyError，程序在load_state_dict阶段崩溃，日志里反复出现下面这一行：

KeyError: 'Layer.34.mlp.experts.gate_up_proj'

如果你把调试级别调高，会看到更完整的一段调用链，大致长这样：

File ".../vllm/model_executor/models/.../model.py", line N, in load_weights param = params[name] File ".../torch/nn/modules/module.py", line M, in _load_from_state_dict ... KeyError: 'Layer.34.mlp.experts.gate_up_proj'

最让人迷惑的地方在于：报错只说「找不到这个 key」，但你自己print(state_dict.keys())时，明明能看到名为Layer.34.mlp.experts.gate_up_proj的键就在字典里。换句话说，模型定义期望的权重名，和 checkpoint 实际保存的权重名，在字面形式上并不一致——它们描述的其实是同一个张量，只是命名约定不同。

这类问题有几个典型特征，可以帮助你快速判断自己是不是踩中了同一个坑：

报错永远落在 loaded checkpoint 与模型state_dict对齐的那一步，而不是在 forward、也不是在 tokenizer。
报错的 key 往往带有experts、.gate_up_proj、.down_proj、.w1、.w2、w3这类 MoE 专属字段。
改一个num_experts、换一个 checkpoint 来源（比如从 HF 原始权重换成某框架导出的权重），报错里 Layer 的序号、专家层名字可能变，但「key error + experts」的模式不变。
用strict=False强行跳过，进程不崩了，但推理结果明显不对（输出乱码、困惑度爆表），说明权重根本没被真正加载进去。

二、背景

要理解这个误报，先得搞清楚 vLLM 在 MoE 模型里是怎么组织专家权重的。

一个典型的 MoE MLP 层，结构是这样的：

Layer.34.mlp.gate # 路由门，决定每个 token 走哪几个专家 Layer.34.mlp.experts # 一组专家，数量 = num_experts Layer.34.mlp.experts.gate_up_proj # 专家的前半投影（某些格式） Layer.34.mlp.experts.down_proj # 专家的后半投影（某些格式）

问题在于：不同框架对「一组专家」的存储方式，约定并不统一。

约定 A——「逐专家拆分」（per-expert）：每个专家是一个独立子模块，名字里带专家下标i：

Layer.34.mlp.experts.0.gate_up_proj Layer.34.mlp.experts.1.gate_up_proj Layer.34.mlp.experts.2.gate_up_proj ... Layer.34.mlp.experts.7.gate_up_proj

约定 B——「融合专家」（fused）：所有专家在某一投影上被沿第 0 维堆叠成一个大张量，名字里不带下标，直接挂在experts下面：

Layer.34.mlp.experts.gate_up_proj # shape = [num_experts, ...] Layer.34.mlp.experts.down_proj # shape = [num_experts, ...]

约定 A 常见于 PyTorchnn.ModuleList直接堆出来的模型定义；约定 B 常见于为了推理效率把专家权重融合成一个张量、用一次大矩阵乘代替循环小矩阵乘的推理引擎。

两件事同时成立时，就会触发我们看到的KeyError：

模型类在load_weights里，按约定 A 去params["Layer.34.mlp.experts.0.gate_up_proj"]这样取权重；
而你手里的 checkpoint 是按约定 B 存的，字典里只有Layer.34.mlp.experts.gate_up_proj，没有带下标的版本。

于是params[逐专家key]找不到，抛出KeyError: 'Layer.34.mlp.experts.gate_up_proj'（注意：报错打印的，往往是它「期望但没拿到」的那个 key，所以字面看起来像是有这个 key，其实字典里那个没下标）。

还有第三、第四种可能，实践中同样常见：

约定 C：checkpoint 用w1/w2/w3命名（Megatron-LM 风格），而 loader 找的是gate_proj/up_proj/down_proj（HuggingFace 风格），纯属名字别名不同。
约定 D：checkpoint 顶部多了一层前缀，比如model.或module.，而 loader 用的是去掉前缀的裸名，导致整体对不齐。

本文聚焦最棘手的「fused vs per-expert」命名错位，其余几种在第八节的排查清单里一并给出处理办法。

三、根因

根因只有一句话：模型定义里「专家权重」的命名约定，与 checkpoint 实际落盘的命名约定不一致；loader 在state_dict对齐阶段按自己的约定去做键查找，找不到就抛KeyError。

展开来说，这是两个相互独立演进的事实撞在了一起：

模型侧：vllm的某个 MoE 模型实现，用nn.ModuleList([Expert(...) for _ in range(num_experts)])把专家组织成带下标的子模块。当它在load_weights里遍历expected_names时，天然期望experts.{i}.{proj}这种逐专家键。
权重侧：你加载的 checkpoint，可能来自：
- 用融合内核训练/导出得到的权重（约定 B）；
- 某个转换脚本只做了张量堆叠，没做键名改写；
- 或者来源模型本身就把experts当「一个融合张量」存。

loader 默认假设「checkpoint 的键名 == 我期望的键名」。一旦这个假设被打破，且 loader 没有做键名兼容/重映射，就会直接KeyError。这不是张量形状问题、不是 dtype 问题、也不是显存问题，纯粹是字符串键对不上。

一个很容易误导人的细节：torch.load之后如果你print(state_dict.keys())，确实能看到Layer.34.mlp.experts.gate_up_proj。你会以为「键在啊，为什么还报找不到」。原因是 loader 内部查找的 key 是带下标的Layer.34.mlp.experts.0.gate_up_proj，而你眼睛看到的、checkpoint 里实际有的，是没下标的那个。打印出来「看起来一样」，其实差了一个.0、.1…前缀。

四、最小可运行复现

下面用一个纯 Python 字典，复现「checkpoint 用 fused 命名，loader 用 per-expert 命名」导致的KeyError，不依赖任何深度学习框架就能跑：

# reproduce_keyerror.py # 复现：checkpoint 是 fused 命名，loader 期望 per-expert 命名 NUM_EXPERTS = 8 # 1) checkpoint 实际落盘的权重（融合专家，单键不带下标） checkpoint = { "Layer.34.mlp.gate.weight": "<gate tensor>", "Layer.34.mlp.experts.gate_up_proj": f"<fused tensor, shape=[{NUM_EXPERTS}, ...]>", "Layer.34.mlp.experts.down_proj": f"<fused tensor, shape=[{NUM_EXPERTS}, ...]>", } # 2) loader 期望的键（逐专家） def expected_keys(num_experts: int): keys = [f"Layer.34.mlp.gate.weight"] for i in range(num_experts): keys.append(f"Layer.34.mlp.experts.{i}.gate_up_proj") keys.append(f"Layer.34.mlp.experts.{i}.down_proj") return keys # 3) 模拟 loader 取权重——这一步会抛 KeyError try: for name in expected_keys(NUM_EXPERTS): _ = checkpoint[name] # 期望逐专家 key，但 checkpoint 只有 fused key except KeyError as e: print("复现成功，抛出 KeyError:", e) print("原因：checkpoint 有 'Layer.34.mlp.experts.gate_up_proj'") print(" loader 找的是 'Layer.34.mlp.experts.0.gate_up_proj'")

运行python reproduce_keyerror.py，你会精确看到和线上一致的KeyError。这证明问题出在「键名约定不一致」，与具体框架无关。

五、解决方案（第一层：最小直接修复）

最小修复思路：在 loader 取权重之前，先把 fused 命名的 checkpoint，重映射成 loader 期望的 per-expert 命名。也就是把experts.gate_up_proj（一个[num_experts, ...]的张量）「展开」成experts.0.gate_up_proj…experts.{N-1}.gate_up_proj，并对每个专家做tensor[i]切片。

下面是一段可以直接落到项目里的工具函数，纯标准库 + 框架无关的命名处理（真正的张量切片留给框架，这里用占位对象演示逻辑）：

# fix_layer1_remap.py import re from typing import Dict, Mapping # 匹配 "....experts.<proj>" 形式的 fused 键 FUSED_EXPERT_RE = re.compile(r"^(?P<prefix>.*?\.experts)\.(?P<proj>gate_up_proj|up_proj|gate_proj|down_proj|w1|w2|w3)$") # 匹配 "....experts.<i>.<proj>" 形式的 per-expert 键 PER_EXPERT_RE = re.compile(r"^(?P<prefix>.*?\.experts)\.(?P<idx>\d+)\.(?P<proj>gate_up_proj|up_proj|gate_proj|down_proj|w1|w2|w3)$") def expand_fused_experts(state: Mapping[str, object], num_experts: int) -> Dict[str, object]: """把 fused 命名展开为 per-expert 命名。 约定：fused 键对应的张量在第 0 维 = num_experts， 展开后 experts.{i}.<proj> 取该张量的第 i 片。 """ out: Dict[str, object] = {} for key, value in state.items(): m = FUSED_EXPERT_RE.match(key) if not m: out[key] = value continue prefix, proj = m.group("prefix"), m.group("proj") # 真实场景：out[f"{prefix}.{i}.{proj}"] = value[i] # 这里用占位字符串表达「第 i 个专家切片」的语义 for i in range(num_experts): out[f"{prefix}.{i}.{proj}"] = f"{value}#slice[{i}]" return out # 反向：把 per-expert 合并回 fused（用于某些 loader 期望 fused 的情况） def fuse_per_experts(state: Mapping[str, object], num_experts: int) -> Dict[str, object]: grouped: Dict[str, list] = {} rest: Dict[str, object] = {} for key, value in state.items(): m = PER_EXPERT_RE.match(key) if not m: rest[key] = value continue fused_key = f"{m.group('prefix')}.{m.group('proj')}" grouped.setdefault(fused_key, [None] * num_experts)[int(m.group("idx"))] = value out = dict(rest) for fused_key, slices in grouped.items(): # 真实场景：out[fused_key] = torch.stack(slices, dim=0) out[fused_key] = f"stacked[{num_experts}] from {slices}" return out if __name__ == "__main__": ckpt = { "Layer.34.mlp.gate.weight": "g", "Layer.34.mlp.experts.gate_up_proj": "fused_gu", "Layer.34.mlp.experts.down_proj": "fused_down", } remapped = expand_fused_experts(ckpt, num_experts=8) print("展开后是否还有 fused 键:", any(FUSED_EXPERT_RE.match(k) for k in remapped)) print("Layer.34.mlp.experts.0.gate_up_proj =", remapped["Layer.34.mlp.experts.0.gate_up_proj"])

这一层修复的特点是：不动模型定义、不动 checkpoint 文件，只在内存里做一次键名 + 张量维度的对齐。改动面最小，风险最低，适合先让服务跑起来。

注意一个坑：展开时num_experts必须和 checkpoint 里 fused 张量的第 0 维严格相等。如果 fused 张量的第 0 维是 8，你却按 64 展开，会越界；反过来按 4 展开，会漏掉一半专家。这个数值应当从模型配置（config.json 里的num_experts/n_routed_experts）读取，不要写死。

六、解决方案（第二层：结构性改进）

第一层是「看到 fused 就展开」。但真实世界更复杂：有的 checkpoint 是 fused，有的是 per-expert，有的甚至混着来（同一份权重里，gate_up 是 fused、down 却是 per-expert）。更稳妥的做法是让 loader 在加载前先探测 checkpoint 的命名风格，再自动选择对齐策略。

# fix_layer2_autodetect.py import re from typing import List FUSED_EXPERT_RE = re.compile(r"\.experts\.(gate_up_proj|up_proj|gate_proj|down_proj|w1|w2|w3)$") PER_EXPERT_RE = re.compile(r"\.experts\.\d+\.(gate_up_proj|up_proj|gate_proj|down_proj|w1|w2|w3)$") PREFIX_RE = re.compile(r"^(model\.|module\.|state_dict\.|model\.)") def detect_expert_layout(state_keys: List[str]) -> str: has_fused = any(FUSED_EXPERT_RE.search(k) for k in state_keys) has_per = any(PER_EXPERT_RE.search(k) for k in state_keys) if has_fused and has_per: return "mixed" if has_fused: return "fused" if has_per: return "per_expert" return "none" def strip_common_prefix(state_keys: List[str]) -> str: """探测是否存在统一前缀（如 'model.'），返回最常见的候选前缀。""" from collections import Counter candidates = [m.group(1) for k in state_keys if (m := PREFIX_RE.match(k))] if not candidates: return "" return Counter(candidates).most_common(1)[0][0] def normalize_keys(raw_keys: List[str], num_experts: int) -> dict: """返回 {normalized_key: original_key} 的映射，loader 据此取张量。""" layout = detect_expert_layout(raw_keys) prefix = strip_common_prefix(raw_keys) mapping = {} for k in raw_keys: nk = k[len(prefix):] if prefix and k.startswith(prefix) else k if layout in ("fused", "mixed") and FUSED_EXPERT_RE.search(k): # 展开：生成 num_experts 个 per-expert 标准化键 proj = FUSED_EXPERT_RE.search(k).group(1) base = nk[: nk.rfind(proj) - 1] # ...experts for i in range(num_experts): mapping[f"{base}.{i}.{proj}"] = k else: mapping[nk] = k return mapping if __name__ == "__main__": sample = [ "model.Layer.34.mlp.gate.weight", "model.Layer.34.mlp.experts.gate_up_proj", "model.Layer.34.mlp.experts.down_proj", ] print("detect layout:", detect_expert_layout(sample)) # fused print("common prefix:", strip_common_prefix(sample)) # model. mp = normalize_keys(sample, num_experts=8) print("标准化后含 per-expert 键:", any(".experts.0." in k for k in mp))

这一层把「命名风格探测」做成了 loader 的前置步骤：

detect_expert_layout区分 fused / per_expert / mixed / none 四种情况；
strip_common_prefix处理约定 D 的顶部前缀问题；
normalize_keys输出一份「标准化键 → 原始键」的映射表，loader 在取权重时统一用标准化键去找，内部再回查原始张量。

这样无论 checkpoint 来自哪个框架、怎么命名，loader 只认「标准化键」，把所有兼容逻辑收敛到一处，模型定义本身不需要为不同来源反复改。

七、解决方案（第三层：断言 / CI 守护）

前两层的修复都依赖「我们能正确猜出 checkpoint 的命名风格」。如果猜错（比如把 per-expert 误判成 fused，或反过来），后果是静默加载了错误的权重——进程不崩，但推理结果错得离谱，这种 bug 比KeyError更危险。因此需要用断言和 CI 把「对齐完整性」钉死。

# fix_layer3_guard.py # 放在 loader 取完权重之后、真正 forward 之前执行 def assert_no_unresolved_expert_keys(state_keys, normalized_keys, num_experts: int): """断言：所有专家投影键，最终都能在归一化映射里被解释。""" expected_count = {} for k in normalized_keys: m = __import__("re").search(r"\.experts\.(\d+)\.(gate_up_proj|up_proj|gate_proj|down_proj|w1|w2|w3)$", k) if m: proj = m.group(2) expected_count[proj] = expected_count.get(proj, 0) + 1 for proj, cnt in expected_count.items(): assert cnt == num_experts, ( f"专家投影 '{proj}' 解析出 {cnt} 个键，但 num_experts={num_experts}，" f"说明 fused 展开时专家数不匹配或键名误判" ) def assert_prefix_consistent(state_keys): """断言：要么全部带同一前缀，要么全部不带，避免半前缀导致部分漏加载。""" import re has_prefix = [bool(re.match(r"^(model\.|module\.)", k)) for k in state_keys] assert all(has_prefix) or not any(has_prefix), ( "checkpoint 键名前缀不一致：部分带 'model.' 前缀、部分不带，" "说明 normalize_keys 的前缀剥离逻辑漏掉了某些键" ) # ---- 可写成 pytest 用例，进 CI ---- def test_fused_expands_to_num_experts(): from fix_layer2_autodetect import normalize_keys ckpt = [ "model.Layer.34.mlp.experts.gate_up_proj", "model.Layer.34.mlp.experts.down_proj", ] mp = normalize_keys(ckpt, num_experts=8) gate_keys = [k for k in mp if k.endswith(".gate_up_proj")] down_keys = [k for k in mp if k.endswith(".down_proj")] assert len(gate_keys) == 8, gate_keys assert len(down_keys) == 8, down_keys # 下标必须连续 0..7 idxs = sorted(int(k.split(".experts.")[1].split(".")[0]) for k in gate_keys) assert idxs == list(range(8)), idxs def test_mixed_layout_does_not_drop_keys(): from fix_layer2_autodetect import detect_expert_layout mixed = [ "Layer.34.mlp.experts.gate_up_proj", # fused "Layer.34.mlp.experts.0.down_proj", # per-expert ] assert detect_expert_layout(mixed) == "mixed"

把test_fused_expands_to_num_experts和test_mixed_layout_does_not_drop_keys接进 CI 后，任何一次「误判命名风格」或「num_experts读错」都会在合并前就被拦下，而不是等到线上推理出错才回头查。

八、排查清单

当你再看到KeyError: 'Layer.xx.mlp.experts.***'这类报错，按下面顺序排查，通常五分钟内能定位：

先打印两边键名：print(sorted(checkpoint.keys()))和 loader 里expected_names，用眼睛（或diff）对比experts那一段到底带不带.<i>下标。
判断 fused 还是 per-expert：fused 键形如experts.gate_up_proj（无下标，张量第 0 维 = 专家数）；per-expert 键形如experts.0.gate_up_proj。
检查顶部前缀：checkpoint 键是否统一带了model./module./state_dict.之类前缀，而 loader 用裸名——是的话做前缀剥离。
检查别名差异：w1/w2/w3（Megatron 风格）vsgate_proj/up_proj/down_proj（HF 风格）；gate_up_proj（融合）vsgate_proj+up_proj（拆分）。必要时做别名映射表。
核对num_experts来源：展开 fused 时用的专家数，必须等于 checkpoint 张量第 0 维，也必须等于 config 里的num_experts。三处不一致必出问题。
不要直接用strict=False蒙混：strict=False只会让缺失键「静默跳过」，问题从「崩」变成「结果错」，反而更难查。除非你确信那些 missing key 本来就该没有。
看错误里的 Layer 序号：报错停在Layer.34不代表只有 34 层有问题，往往是 loader 从某一层开始第一次遇到 experts 权重，前面的层还没走到 experts 那段。
混合布局要小心：同一份权重可能 gate_up 是 fused、down 是 per-expert，探测逻辑必须支持 mixed，不能一刀切。
用断言兜底：加载完做一次「专家键计数 == num_experts」的断言，把它写进 CI，防止回归。
最后才考虑改模型定义：优先在 loader 侧做键名兼容（内存重映射），不要为了迁就某一个 checkpoint 去改nn.ModuleList结构，否则会影响其它所有权重来源。

九、小结

KeyError: 'Layer.34.mlp.experts.gate_up_proj'看着吓人，本质却很朴素：模型期望「逐专家命名」，checkpoint 却是「融合命名」，loader 按自己的约定去查键，查不到就报KeyError。修复分三层——第一层在内存里把 fused 键展开成 per-expert 键并做张量切片，让服务先跑起来；第二层加一个命名风格探测器，自动识别 fused / per-expert / mixed / 前缀差异，把所有兼容逻辑收敛到一处；第三层用「专家键计数 == num_experts」的断言和 pytest 用例钉进 CI，杜绝静默加载错误权重。核心原则只有两条：加载侧的键名兼容优先于改动模型定义；任何命名风格的猜测都必须有断言兜底，不能让对齐失败演变成「进程不死、结果全错」的隐蔽故障。