【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_loading=True、先在 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 的加载假设打架：

1. embedding 与 lm_head 权重共享。即 lm_head.weight is embed_tokens.weight。这是标准做法，可以少一份输出层参数。
2. 大模型普遍采用 meta 设备初始化 + 分片落盘加载。也就是先在 device="meta" 上把模型结构搭起来（不占显存），再用 load_checkpoint_and_dispatch 把 safetensors 里对应 shard 的张量逐块 materialize 到目标设备。

FSDP2（torch.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_dispatch（accelerate 的实现）在按"模块名路径"加载时，会同时尝试为 embed_tokens.weight 和 lm_head.weight 两个 key 各自 materialize 一份张量；
• 如果 materialize 逻辑没有识别出二者是 tie，就会对同一个 meta 张量触发两次 materialize()，或者第二次发现它已经不是 meta 而报错；
• 更进一步，当 lm_head 被注册为 Linear(..., bias=False) 时，如果权重以转置形式存放（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 几份"的认知不一致。

具体有三处失配：

1. 枚举视角不一致：module.parameters() 只看对象引用，tie 只算一份；加载器按 state_dict key 枚举，tie 算两份。
2. materialize 幂等性缺失：同一个 meta 张量被 materialize() 两次，第二次要么报错，要么产生两份物理张量，导致后续 fully_shard 在 hook 里拿到错误的张量。
3. 切分策略视角不一致：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-OSS：embed_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, bias=False)
        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(
                f"Materialization failed: parameter '{name}' was already "
                f"materialized by tie (count={meta[name].materialized})"
            )


def main():
    model = TinyTiedLM(vocab=200, hidden=16).to("meta")
    meta = {
        "embed_tokens": _Meta(),
        "lm_head.weight": _Meta(),
    }
    # 模拟按 state_dict key 同时加载两个名字
    try:
        buggy_loader(
            model,
            keys=["embed_tokens", "lm_head.weight"],
            dev="cpu",
            meta=meta,
        )
        print("加载成功（但这是 bug 路径没触发，检查复现）")
    except RuntimeError as e:
        print("复现到根因错误：", e)


if __name__ == "__main__":
    main()

运行后你会看到：

复现到根因错误: Materialization failed: parameter 'lm_head.weight' was already materialized by tie (count=1)

这就复现了 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(vocab=200, hidden=16).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 沿错误维度切分。第二层做三件事：

1. 建立全局 TieRegistry，作为模型加载与切分的唯一事实来源。
2. 在 fully_shard 之前，把 tie 别名从"待切分参数集合"里剔除，避免重复切分。
3. 显式声明切分维度：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_factory=dict)

    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_factory=set)
    dim: Dict[str, int] = field(default_factory=dict)


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(vocab=200, hidden=16).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, bias=False)
        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 加载路径上遇到类似问题时，按以下顺序排查：

1. 先确认是不是 tie 问题：打印 id(model.embed_tokens) == id(model.lm_head.weight)，相等即为 tie。
2. 看报错发生在哪一步：是 materialization 阶段（meta→真实设备）还是 fully_shard 阶段？前者是加载器把 tie 当两份，后者是切分把 tie 当两份。
3. 检查加载日志：如果出现 already materialized，说明加载器对同一个 meta 张量调了两次 materialize，按第五节加 tie 归一化。
4. 检查 named_parameters 与 state_dict key 的差异：nn.Module 只枚举一份 tie，state_dict/checkpoint 里却常有两个 key，二者必须做 canonic 映射。
5. 检查切分维度：tie 的 lm_head/embed_tokens 只能在 hidden 维切，不能按 vocab 维切；若报形状错位，优先怀疑切分维度配置。
6. 逐层缩小：先用单卡、非 meta、直接 .cuda() 验证模型本身能 forward，再逐步打开 cpu_ram_efficient_loading、meta init、fully_shard，定位是哪一层引入冲突。
7. 断言守护：在模型构建函数末尾加一句 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 初始化 + 分片加载"的组合就能稳定落到多卡上。