【Bug已解决】GPT-OSS fails to load with FSPD2 解决方案

【Bug已解决】GPT-OSS fails to load with FSPD2 解决方案一、现象长什么样在把 OpenAI 开源权重模型GPT-OSS通过accelerate的 FSDP2 路径fully_shard加载到多卡上时很多人会卡在一个非常早期、且错误信息并不直观的阶段。最常见的几种报错形态如下KeyError: lm_head.weight或者RuntimeError: shape mismatch: tied lm_head expects (32000, 2880), got (2880, 32000)又或者在启用了cpu_ram_efficient_loadingTrue、先在meta设备上构建模型再逐张量落盘加载时程序直接停在 materialization 阶段RuntimeError: Materialization failed: parameter lm_head.weight was already materialized by tie最隐蔽的一种是加载看似成功但训练第一步 forward 报RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and meta这些现象的共同点是模型里lm_head与embed_tokens是权重共享tied的而 FSDP2 的切分逻辑在处理同一个nn.Parameter被两个模块引用时和 GPT-OSS 这种先 meta 初始化、再按 shard 落盘的加载流程发生了冲突。下面把根因讲清楚。二、背景GPT-OSS 是一个 decoder-only 的 MoE 模型结构上有两个关键特征恰恰都和 FSDP2 的加载假设打架embedding 与 lm_head 权重共享。即lm_head.weight is embed_tokens.weight。这是标准做法可以少一份输出层参数。大模型普遍采用 meta 设备初始化 分片落盘加载。也就是先在devicemeta上把模型结构搭起来不占显存再用load_checkpoint_and_dispatch把 safetensors 里对应 shard 的张量逐块 materialize 到目标设备。FSDP2torch.distributed.fsdp的fully_shard的工作方式是遍历module.parameters()对每个参数调用Mesh上的分片。问题在于——当lm_head.weight与embed_tokens.weight是同一个 Python 对象时module.parameters()在nn.Module层面只会枚举一次因为lm_head持有的是embed_tokens.weight这个引用不是独立参数这本该是好事但load_checkpoint_and_dispatchaccelerate 的实现在按模块名路径加载时会同时尝试为embed_tokens.weight和lm_head.weight两个 key 各自 materialize 一份张量如果 materialize 逻辑没有识别出二者是 tie就会对同一个 meta 张量触发两次materialize()或者第二次发现它已经不是 meta 而报错更进一步当lm_head被注册为Linear(..., biasFalse)时如果权重以转置形式存放GPT-OSS 的lm_head实际上直接复用embed_tokens.weight形状是[vocab, hidden]而Linear内部做x weight.T逻辑上没问题一旦切分策略把lm_head当成需要独立切分的层就会出现形状错位。下面用一段可以独立运行的最小复现把tie 被重复 materialize这一最关键的根因钉死。三、根因把问题抽象成一个最小模型一个共享 embedding 的小 Transformer。我们模拟 accelerate 的按模块名加载 按分片 materialize流程故意对 tie 的两个名字都调一次materialize观察冲突。关键根因只有一句话FSDP2 / accelerate 的加载器把tie当成两个独立张量来处理而nn.Module只把它当成一个参数——两者对需要 materialize 几份的认知不一致。具体有三处失配枚举视角不一致module.parameters()只看对象引用tie 只算一份加载器按 state_dict key 枚举tie 算两份。materialize 幂等性缺失同一个 meta 张量被materialize()两次第二次要么报错要么产生两份物理张量导致后续fully_shard在 hook 里拿到错误的张量。切分策略视角不一致fully_shard想对lm_head单独切分但 tie 不允许把同一份物理内存切成两份embed 按 hidden 维切、lm_head 按 vocab 维切方向不同。四、最小可运行复现下面这段代码不依赖任何真实大模型权重纯torch即可复现tie 被重复 materialize的冲突。import torch import torch.nn as nn from dataclasses import dataclass, field from typing import Dict, List class TinyTiedLM(nn.Module): 模拟 GPT-OSSembed_tokens 与 lm_head 共享权重tie。 def __init__(self, vocab: int, hidden: int): super().__init__() self.embed_tokens nn.Parameter(torch.empty(vocab, hidden)) # [vocab, hidden] # lm_head 直接复用 embed_tokens不新建参数 self.lm_head nn.Linear(hidden, vocab, biasFalse) self.lm_head.weight self.embed_tokens # 关键tie def forward(self, x): h self.embed_tokens[x] # [B, T, hidden] return self.lm_head(h) # [B, T, vocab] dataclass class _Meta: materialized: int 0 def buggy_loader(model: nn.Module, keys: List[str], dev: str, meta: Dict[str, _Meta]): 模拟 accelerate 的 load_checkpoint_and_dispatch按 key 逐个 materialize。 for name, p in model.named_parameters(): if name not in keys: continue # 错误点对每个 key 都 materialize没有识别 tie if p.is_meta: torch.nn.init.normal_(p) # 模拟从 safetensors 落盘 meta[name].materialized 1 else: # 已经被另一个 key 当做 tie materialize 过了 raise RuntimeError( fMaterialization failed: parameter {name} was already fmaterialized by tie (count{meta[name].materialized}) ) def main(): model TinyTiedLM(vocab200, hidden16).to(meta) meta { embed_tokens: _Meta(), lm_head.weight: _Meta(), } # 模拟按 state_dict key 同时加载两个名字 try: buggy_loader( model, keys[embed_tokens, lm_head.weight], devcpu, metameta, ) print(加载成功但这是 bug 路径没触发检查复现) except RuntimeError as e: print(复现到根因错误, e) if __name__ __main__: main()运行后你会看到复现到根因错误: Materialization failed: parameter lm_head.weight was already materialized by tie (count1)这就复现了 GPT-OSS 在 FSDP2 加载流程里最典型的失败——lm_head.weight和embed_tokens共享同一份 meta 张量加载器却对两个 key 各 materialize 一次。五、解决方案第一层最小直接修复最立竿见影的修复在加载阶段做一次 tie 归一化canonicalize把lm_head.weight映射回embed_tokens全程只对 canonic 名字 materialize 一次。同时在fully_shard之前显式告诉 FSDP2 哪些参数不要单独切分。import torch import torch.nn as nn from typing import Dict, List, Set def build_tie_map(model: nn.Module) - Dict[str, str]: 返回 {alias_name: canonical_name}例如 {lm_head.weight: embed_tokens}。 seen: Dict[int, str] {} tie: Dict[str, str] {} for name, p in model.named_parameters(): pid id(p) if pid in seen: tie[name] seen[pid] # 这是别名指向已经见过的参数 else: seen[pid] name return tie def fixed_loader(model: nn.Module, keys: List[str], meta_count: Dict[str, int]): 修复版先归一化 tie再只按 canonic 名字 materialize。 tie_map build_tie_map(model) canonic_keys [] for k in keys: canonic tie_map.get(k, k) if canonic not in canonic_keys: canonic_keys.append(canonic) for name, p in model.named_parameters(): if name in canonic_keys and p.is_meta: torch.nn.init.normal_(p) meta_count[name] meta_count.get(name, 0) 1 def main(): model TinyTiedLM(vocab200, hidden16).to(meta) meta_count: Dict[str, int] {} fixed_loader(model, [embed_tokens, lm_head.weight], meta_count) print(materialize 次数:, meta_count) # 期望只 materialize 一次{embed_tokens: 1} out model(torch.randint(0, 200, (2, 5))) print(forward 通过输出形状:, tuple(out.shape)) if __name__ __main__: main()这一层修复直接消除了Materialization failed错误且lm_head因为复用embed_tokens自动拿到了同一份张量。六、解决方案第二层结构性改进仅仅在加载时归一化还不够因为后面fully_shard仍可能把 tie 当成两个切分单元或把lm_head沿错误维度切分。第二层做三件事建立全局TieRegistry作为模型加载与切分的唯一事实来源。在fully_shard之前把 tie 别名从待切分参数集合里剔除避免重复切分。显式声明切分维度embed/lm_head 这类 tie 只在 hidden 维最后一维切分不要按 vocab 维切。import torch import torch.nn as nn from dataclasses import dataclass, field from typing import Dict, List, Set dataclass class TieRegistry: 把 tie 关系收口到一个地方加载器和切分器都读它。 aliases: Dict[str, str] field(default_factorydict) def canonical(self, name: str) - str: return self.aliases.get(name, name) def is_alias(self, name: str) - bool: return name in self.aliases def seen_params(self, names: List[str]) - List[str]: out, done [], set() for n in names: c self.canonical(n) if c not in done: done.add(c) out.append(c) return out def make_registry(model: nn.Module) - TieRegistry: reg TieRegistry() seen: Dict[int, str] {} for name, p in model.named_parameters(): pid id(p) if pid in seen: reg.aliases[name] seen[pid] else: seen[pid] name return reg dataclass class ShardPlan: 声明哪些参数参与切分、沿哪个维度切。tie 只在 hidden 维切。 shardable: Set[str] field(default_factoryset) dim: Dict[str, int] field(default_factorydict) def build_plan(model: nn.Module, reg: TieRegistry) - ShardPlan: plan ShardPlan() for name, p in model.named_parameters(): if reg.is_alias(name): continue # 别名不参与切分 plan.shardable.add(name) plan.dim[name] p.dim() - 1 # 默认沿最后一维hidden切 return plan def main(): model TinyTiedLM(vocab200, hidden16).to(meta) # 第一步元设备下先 materialize用第五节的 fixed_loader 思路按 canonic reg make_registry(model) keys reg.seen_params([embed_tokens, lm_head.weight]) for name, p in model.named_parameters(): if name in keys and p.is_meta: torch.nn.init.normal_(p) plan build_plan(model, reg) print(切分单元:, sorted(plan.shardable)) # 期望只有 {embed_tokens}lm_head.weight 已被剔除 print(tie 别名:, reg.aliases) out model(torch.randint(0, 200, (2, 5))) print(forward 通过输出形状:, tuple(out.shape)) if __name__ __main__: main()这一层的价值在于加载和切分不再各自为政都从TieRegistry读同一份真相后续再换模型结构也不会重现 tie 双切分问题。七、解决方案第三层断言 / CI 守护为防止回归比如有人改了lm_head初始化又悄悄打破了 tie加一组pytest断言作为 CI 守护import torch import torch.nn as nn import pytest class TinyTiedLM(nn.Module): def __init__(self, vocab, hidden): super().__init__() self.embed_tokens nn.Parameter(torch.empty(vocab, hidden)) self.lm_head nn.Linear(hidden, vocab, biasFalse) self.lm_head.weight self.embed_tokens def forward(self, x): return self.lm_head(self.embed_tokens[x]) def build_tie_map(model): seen, tie {}, {} for name, p in model.named_parameters(): if id(p) in seen: tie[name] seen[id(p)] else: seen[id(p)] name return tie def test_tie_is_single_parameter(): 核心断言tie 必须是同一个 Python 对象加载后仍是。 model TinyTiedLM(200, 16) assert model.lm_head.weight is model.embed_tokens tie build_tie_map(model) assert lm_head.weight in tie assert tie[lm_head.weight] embed_tokens def test_no_duplicate_materialize(): 修复后的加载每个 canonic 名字只 materialize 一次。 model TinyTiedLM(200, 16).to(meta) tie build_tie_map(model) count {} for name, p in model.named_parameters(): if tie.get(name, name) not in count and p.is_meta: torch.nn.init.normal_(p) count[tie.get(name, name)] 1 # embed_tokens 与 lm_head.weight 归一后是同一个所以 count 只有 1 项 assert len(count) 1 assert model.lm_head.weight is not None assert not model.embed_tokens.is_meta def test_forward_shape(): model TinyTiedLM(200, 16) out model(torch.randint(0, 200, (2, 5))) assert tuple(out.shape) (2, 5, 200) if __name__ __main__: pytest.main([__file__, -q])CI 里只要test_no_duplicate_materialize通过就能保证tie 被重复 materialize这一类 GPT-OSS FSDP2 加载失败不会再回来。八、排查清单当你在 GPT-OSS FSDP2 加载路径上遇到类似问题时按以下顺序排查先确认是不是 tie 问题打印id(model.embed_tokens) id(model.lm_head.weight)相等即为 tie。看报错发生在哪一步是 materialization 阶段meta→真实设备还是fully_shard阶段前者是加载器把 tie 当两份后者是切分把 tie 当两份。检查加载日志如果出现already materialized说明加载器对同一个 meta 张量调了两次materialize按第五节加 tie 归一化。检查named_parameters与 state_dict key 的差异nn.Module只枚举一份 tiestate_dict/checkpoint 里却常有两个 key二者必须做 canonic 映射。检查切分维度tie 的lm_head/embed_tokens只能在 hidden 维切不能按 vocab 维切若报形状错位优先怀疑切分维度配置。逐层缩小先用单卡、非 meta、直接.cuda()验证模型本身能 forward再逐步打开cpu_ram_efficient_loading、meta init、fully_shard定位是哪一层引入冲突。断言守护在模型构建函数末尾加一句assert model.lm_head.weight is model.embed_tokens把 tie 关系变成硬约束。九、小结GPT-OSS 在 FSDP2 加载路径上失败根子不在 FSDP2 本身而在于**tie 共享权重这件事在nn.Module视角和state_dict/加载器视角下被数了两次**前者只算一份参数、后者按两个 key 各 materialize 一次于是冲突爆发为Materialization failed、形状错位或 device 不一致。修复分三层第一层在加载阶段做 tie 归一化只按 canonic 名字 materialize第二层用TieRegistry收口 tie 真相让加载与切分都从同一处读取并明确 tie 只在 hidden 维切分第三层用 pytest 断言test_no_duplicate_materialize等把tie 必须是单一参数、只能 materialize 一次变成 CI 不可逾越的红线。三层叠加后GPT-OSS 这类共享权重 meta 初始化 分片加载的组合就能稳定落到多卡上。