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

【Bug已解决】FSDP2 fails due to KeyError lm_head.weight 解决方案一、现象长什么样用 FSDP2 训练一个tie_word_embeddingsTrue的模型即lm_head.weight与model.embed_tokens.weight共享在fully_shard或加载 / 分片阶段直接抛KeyError lm_head.weight最小判据触发FSDP2 模型 tie_word_embeddingsTruelm_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_embeddingsTrue时模型的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/ 分片计划在某些实现里会遍历模块的命名参数或按预定义的权重名清单去规划分片清单里包含lm_head.weight因为很多模型 lm_head 独立用lm_head.weight去named_parameters/state_dict里查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) # 这里炸根因链条tied 模型里lm_head.weight是embed_tokens.weight的别名参数列表只有一份分片代码按固定名lm_head.weight查找该键在 tied 模型里不存在 -KeyError不 tie 模型有独立lm_head.weight正常根因是分片代码假设 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.weighttied 模型参数列表里没有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_factorydict) 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时确认模型是否tie_word_embeddingsTrue检查分片代码是否按固定名lm_head.weight查找打印list(model.state_dict().keys())看是否真的没有lm_head.weight按第五 / 六节把 tied 别名解析回embed_tokens.weight不 tie 正常、tie 异常几乎可断定是 tied 别名缺失校验别名与目标指向同一对象防 tie 配置错把第七节的 pytest 接进 CI守护 tied 模型可分片。九、小结FSDP2 在tie_word_embeddingsTrue的模型上报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 分片角度覆盖了同一类问题。