【Bug已解决】FSDP2 fails due to KeyError： ‘lm_head.weight‘ 解决方案

【Bug已解决】FSDP2 fails due to KeyError： 'lm_head.weight' 解决方案

一、现象长什么样

用 FSDP2 训练一个tie_word_embeddings=True的模型（即lm_head.weight与model.embed_tokens.weight共享），在fully_shard或加载 / 分片阶段直接抛：

KeyError： 'lm_head.weight'

最小判据：

触发：FSDP2 + 模型 tie_word_embeddings=True（lm_head 与 embed 共享） 现象：fully_shard / 分片计划找不到 'lm_head.weight' 根因：tied 后 lm_head.weight 不是独立参数（是 embed 的别名）， 分片代码按名字查 'lm_head.weight' 却查不到 影响：带 tied embedding 的模型无法用 FSDP2

最迷惑的是：不 tie 的模型（lm_head 独立）用 FSDP2 正常，一旦 tie 就 KeyError。因为 tied 后lm_head.weight从"独立参数"变成了"embed_tokens 的别名"，名字还在但参数对象不在参数列表里，按名查找即失败。

二、背景

tie_word_embeddings=True时，模型的lm_head不持有独立nn.Parameter，而是直接复用embed_tokens.weight。在 PyTorch 里这通常通过：self.lm_head.weight = self.model.embed_tokens.weight（别名赋值）实现，于是model.parameters()里只有一份embed_tokens.weight，没有独立的lm_head.weight。

FSDP2 的fully_shard（或infer_auto_device_map/ 分片计划）在某些实现里会：

1. 遍历模块的命名参数或按预定义的权重名清单去规划分片；
2. 清单里包含'lm_head.weight'（因为很多模型 lm_head 独立）；
3. 用'lm_head.weight'去named_parameters/state_dict里查；
4. tied 模型里这个键不存在（被别名合并了）->KeyError。

更深层：FSDP2 在分片时需要知道"每组参数属于哪个 module、怎么切"。如果它用固定名字'lm_head.weight'去找 lm_head 模块的参数，而 tied 后该模块没有独立参数，查找失败。

根因是"分片代码按固定权重名查找，未处理 tied embedding 的别名合并"。

三、根因

抽象成代码（示意）：

def shard_plan(model): names = [n for n, _ in model.named_parameters()] # BUG：假设 'lm_head.weight' 一定存在 plan = build_plan_for(names + ["lm_head.weight"]) # tied 时 'lm_head.weight' 不在 names if "lm_head.weight" not in names: raise KeyError("lm_head.weight") # 这里炸

根因链条：

1. tied 模型里lm_head.weight是embed_tokens.weight的别名，参数列表只有一份；
2. 分片代码按固定名'lm_head.weight'查找；
3. 该键在 tied 模型里不存在 ->KeyError；
4. 不 tie 模型有独立lm_head.weight，正常；
5. 根因是"分片代码假设 lm_head 总是独立参数"。

一句话：tied embedding 把lm_head.weight合并成 embed 别名，FSDP2 按固定名查找即 KeyError。

四、最小可运行复现

用纯 Python 模拟"按固定名查找 tied 缺失参数"：

# repro_tied_lmhead.py def shard_plan_buggy(param_names, expected): if expected not in param_names: raise KeyError(expected) return "planned" def main(): # tied 模型：只有 embed_tokens.weight，没有独立 lm_head.weight tied_params = ["model.embed_tokens.weight", "model.layers.0.weight"] try: shard_plan_buggy(tied_params, "lm_head.weight") except KeyError as e: print("复现成功 ->", e) if __name__ == "__main__": main()

运行输出：

复现成功 -> lm_head.weight

tied 模型参数列表里没有lm_head.weight，按名查找即 KeyError，正是真实 bug 的抽象。

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

最小且必须的一步：分片前识别 tied，把lm_head.weight解析回它实际指向的embed_tokens.weight，不再按不存在的名字查找：

# fix_layer1.py def resolve_lm_head(params, tie_map): # tie_map: {"lm_head.weight": "model.embed_tokens.weight"} resolved = [] for name in params: resolved.append(name) for alias, target in tie_map.items(): if alias not in params and target in params: resolved.append(alias) # 补上别名入口，指向同一份 return resolved # 用法 tie_map = {"lm_head.weight": "model.embed_tokens.weight"} plan_names = resolve_lm_head(list(model.state_dict().keys()), tie_map)

要点：

• resolve_lm_head把别名lm_head.weight补回查找清单，指向真实embed_tokens.weight；
• 分片时两者共享同一份参数，不重复分片；
• 不再KeyError。

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

把"tied 权重解析"做成模型结构感知的预处理：在分片前统一把 tied 别名映射好，并校验别名与目标确实指向同一对象（防误配）：

# fix_layer2.py from dataclasses import dataclass, field from typing import Dict @dataclass class TieResolver: tie_map: Dict[str, str] = field(default_factory=dict) def expand(self, param_names): out = list(param_names) for alias, target in self.tie_map.items(): if alias not in param_names and target in param_names: out.append(alias) # 别名补入口 return out def validate_shared_object(self, model): # 校验别名和目标确实指向同一 Parameter 对象（tied 的本质） for alias, target in self.tie_map.items(): a = dict(model.named_parameters()).get(alias) t = dict(model.named_parameters()).get(target) if a is not None and t is not None: assert a is t, f"{alias} 与 {target} 未共享同一对象，tie 配置错误" # 用法 resolver = TieResolver({"lm_head.weight": "model.embed_tokens.weight"}) plan_names = resolver.expand(list(model.state_dict().keys())) resolver.validate_shared_object(model)

要点：

• TieResolver.expand把 tied 别名补回分片清单；
• validate_shared_object校验别名和目标确实is同一对象，防 tie 配置错；
• 分片逻辑只面对"展开后的名字"，不再 KeyError。

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

写 pytest 验证"tied 模型分片不 KeyError、别名指向同一对象"：

# test_tied_lmhead.py import pytest def expand(tie_map, names): out = list(names) for alias, target in tie_map.items(): if alias not in names and target in names: out.append(alias) return out def test_tied_no_keyerror(): tied = ["model.embed_tokens.weight", "model.layers.0.weight"] tie_map = {"lm_head.weight": "model.embed_tokens.weight"} out = expand(tie_map, tied) assert "lm_head.weight" in out, "tied 别名应被补回，不 KeyError" def test_alias_points_to_target(): tie_map = {"lm_head.weight": "model.embed_tokens.weight"} assert tie_map["lm_head.weight"] == "model.embed_tokens.weight" def test_untied_untouched(): untied = ["lm_head.weight", "model.embed_tokens.weight"] out = expand({}, untied) assert out == untied, "不 tie 时不应改动"

CI 一旦有人把 tied 处理删掉，test_tied_no_keyerror立刻变红。

八、排查清单

FSDP2 报KeyError: 'lm_head.weight'时：

1. 确认模型是否tie_word_embeddings=True；
2. 检查分片代码是否按固定名'lm_head.weight'查找；
3. 打印list(model.state_dict().keys())，看是否真的没有lm_head.weight；
4. 按第五 / 六节把 tied 别名解析回embed_tokens.weight；
5. 不 tie 正常、tie 异常，几乎可断定是 tied 别名缺失；
6. 校验别名与目标指向同一对象，防 tie 配置错；
7. 把第七节的 pytest 接进 CI，守护 tied 模型可分片。

九、小结

FSDP2 在tie_word_embeddings=True的模型上报KeyError: 'lm_head.weight'，根因是 tied 后lm_head.weight是embed_tokens.weight的别名、不再是独立参数，而分片代码按固定名'lm_head.weight'查找，查不到即 KeyError。不 tie 模型有独立 lm_head，正常。

三层层级：

• 第一层：分片前把lm_head.weight解析回embed_tokens.weight，补回别名入口；
• 第二层：用TieResolver统一展开 tied 别名并校验共享对象；
• 第三层：pytest 验证 tied 模型不 KeyError、别名指向同一对象，锁进 CI。

核心教训：任何"按固定权重名查找"的分片 / 加载代码，都必须先处理tied embedding 的别名合并。把 tie 关系做成显式映射并校验"别名与目标共享同一对象"，是支持各类模型结构的前提——本系列第 531 篇（GPT-OSS）从 dispatcher 角度、本篇从 FSDP2 分片角度覆盖了同一类问题。