【Bug已解决】Add support for GPT_OSS with tp_plan or enable native tensor parallelism 解决方案一、现象长什么样GPT-OSS 是 OpenAI 开源的 MoE 模型带权重共享的lm_head/embed_tokenstie weights和大量专家experts。当你尝试用 Transformers 的tp_plan给它开原生张量并行TP时报这类错# 现象 AMoE 专家层没被切分OOM torch.OutOfMemoryError: CUDA out of memory. # tp_plan 只列了 self_attn/mlp没列 experts专家权重全量复制 # 现象 B共享 lm_head 在 TP 下形状错 RuntimeError: mat1 and mat2 shapes cannot be multiplied (1024x4096 and 8192x4096) # lm_head 与 embed_tokens 共享权重却被按 col_parallel 切了输出维 # 但 embed 没切二者维度契约破坏 # 现象 Ctp_plan 缺少 gpt-oss 特有的层parallelize_module 跳过导致不对称 AssertionError: expert gate weight not sharded; collective mismatch. # 部分 rank 的专家 gate 未被切all-to-all 集合通信不对称 - 卡死/断言 # 典型触发 model AutoModelForCausalLM.from_pretrained(gpt-oss-20b, torch_dtypeauto) parallelize_module(model, mesh, model.base_model_tp_plan) # plan 没覆盖 experts最典型的指纹单卡/TP1 正常TP2 必崩崩在 MoE 专家或共享 lm_head——因为base_model_tp_plan是按普通 Transformer模板写的没覆盖 gpt-oss 的 MoE 专家层和共享权重。二、背景GPT-OSS 相比普通 decoder-only 模型有两个特殊点让它的tp_plan不能套用通用模板MoE混合专家结构。 每个 decoder block 的 FFN 被换成若干 expertmlp.experts.{i}.w1/w2/w3或gate/up/down外加一个 routergate。这些专家权重要沿tp切分通常按专家数或按专家内部维度否则每个 rank 全量持有所有专家 → OOM。而且 router 的 all-to-all 集合通信要求每个 rank 对专家权重的切分完全一致否则通信不对称卡死。权重共享tie weights。lm_head.weight is embed_tokens.weight。在 TP 下如果embed_tokens被col_parallel输出维切分那lm_head必须也按同样方式切并且因为是共享同一份数据不能一个切一个不切。通用tp_plan往往只把lm_head标成col_parallel却没保证和embed_tokens同步于是形状错。三、根因根因有三类tp_plan漏登 MoE 专家层。 gpt-oss 的base_model_tp_plan沿用了普通模型的模板只列self_attn.*/mlp.gate_proj之类没列mlp.experts.*与 routergate。parallelize_module遍历时跳过这些层 → 专家权重全量复制OOM router 集合通信不对称卡死/断言。共享 lm_head 与 embed_tokens 切分不同步。 二者共享权重但tp_plan若只把lm_head标col_parallel而embed_tokens标local或不一致前向embed_tokens输出完整hiddenlm_head却用切过的权重乘 → 形状错1024x4096vs8192x4096。专家切分方式与 router all-to-all 不匹配。 MoE 的 TP 通常按专家维度切每 rank 持有部分专家或按专家内部隐藏维切。若tp_plan把专家内部w1/w2/w3按列并行切了隐藏维但 router 的 all-to-all 假设按专家数切二者不匹配 → 集合通信形状错。四、最小可运行复现下面用纯 Python 模拟tp_plan 漏登 MoE 专家 → 全量复制OOM与 router 切分不对称from typing import Dict, List HIDDEN 4096 NUM_EXPERTS 8 TP 2 # gpt-oss 实际存在的可切分层 REAL [ model.layers.0.self_attn.q_proj, model.layers.0.mlp.gate, # router model.layers.0.mlp.experts.0.w1, model.layers.0.mlp.experts.0.w2, model.layers.0.mlp.experts.0.w3, model.layers.0.mlp.experts.7.w1, model.layers.0.mlp.experts.7.w3, lm_head, embed_tokens, ] # 通用模板漏了 experts 与 router GENERIC_PLAN: Dict[str, str] { model.layers.*.self_attn.q_proj: col_parallel, model.layers.*.self_attn.o_proj: row_parallel, lm_head: col_parallel, embed_tokens: col_parallel, } def covered(mod: str, plan: Dict[str, str]) - bool: for k in plan: if k.replace(.*, .0) mod: return True return False def audit(plan): unsharded, asym 0, False for m in REAL: if not covered(m, plan): unsharded 1 # router 未切 - all-to-all 不对称 if not covered(model.layers.0.mlp.gate, plan): asym True return unsharded, asym n_unsharded, asym audit(GENERIC_PLAN) print(漏切分层数:, n_unsharded, router 不对称:, asym) # 专家全部漏切 router 不对称 assert n_unsharded 0 and asym, 复现失败应当漏切专家且 router 不对称运行后GENERIC_PLAN漏掉了所有mlp.experts.*与mlp.gate导致专家全量复制OOM且 router 集合通信不对称卡死正好对应现象 A/C。五、解决方案第一层最小直接修复最快的止血在parallelize_module前把 MoE 专家、router、以及共享 lm_head 的同步切分补齐进tp_plandef patch_gptoss_tp_plan(model, num_experts: int): 第一层修复把 gpt-oss 的 MoE 专家、router、共享 lm_head 补全进 tp_plan。 plan dict(model.base_model_tp_plan) # 1) routergate走列并行保证 all-to-all 对称 plan[model.layers.*.mlp.gate] col_parallel # 2) 每个专家的 w1/w3 列并行、w2 行并行 for i in range(num_experts): plan[fmodel.layers.*.mlp.experts.{i}.w1] col_parallel plan[fmodel.layers.*.mlp.experts.{i}.w3] col_parallel plan[fmodel.layers.*.mlp.experts.{i}.w2] row_parallel # 3) 共享权重lm_head 与 embed_tokens 必须同步切分都 col_parallel plan[lm_head] col_parallel plan[embed_tokens] col_parallel # 告诉 TP 引擎二者共享避免重复切分/不一致 if hasattr(model, _tied_weights_keys): model._tied_weights_keys [lm_head.weight, model.embed_tokens.weight] model.base_model_tp_plan plan return model # 使用 from transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained(gpt-oss-20b, torch_dtypeauto) model patch_gptoss_tp_plan(model, num_experts8) from torch.distributed.tensor.parallel import parallelize_module from torch.distributed.device_mesh import init_device_mesh mesh init_device_mesh(cuda, (TP,), mesh_dim_names(tp,)) parallelize_module(model, mesh, model.base_model_tp_plan)第一层让 gpt-oss 在 TP2 下正常加载与运行不再 OOM 或形状错。六、解决方案第二层结构性改进用GptOssTpAuditor自动推断 MoE 专家的切分计划避免手写num_experts漏配from dataclasses import dataclass from typing import Dict, List dataclass class GptOssTpAuditor: 自动为 gpt-oss 补全 MoE 共享权重的 tp_plan。 def discover_experts(self, model) - int: # 数一下模型里实际有几个专家 max_idx -1 for name, _ in model.named_modules(): if .experts. in name: try: idx int(name.split(.experts.)[1].split(.)[0]) max_idx max(max_idx, idx) except ValueError: pass return max_idx 1 def build_plan(self, model, base_plan: Dict[str, str]) - Dict[str, str]: plan dict(base_plan) n self.discover_experts(model) for i in range(n): plan[fmodel.layers.*.mlp.experts.{i}.w1] col_parallel plan[fmodel.layers.*.mlp.experts.{i}.w3] col_parallel plan[fmodel.layers.*.mlp.experts.{i}.w2] row_parallel plan[model.layers.*.mlp.gate] col_parallel # 共享权重同步 plan[lm_head] col_parallel plan[model.embed_tokens] col_parallel return plan # 使用 auditor GptOssTpAuditor() model.base_model_tp_plan auditor.build_plan(model, model.base_model_tp_plan)GptOssTpAuditor的语义是凡是模型里真实存在的 MoE 专家层自动按w1/w3 列并行、w2 行并行登记并同步共享权重切分——即使上游 transformers 的base_model_tp_plan漏了 gpt-oss 特有层也会被自动补全。七、解决方案第三层断言 / CI 守护用 pytest 固化gpt-oss tp_plan 必须覆盖专家/router、共享权重同步切分import pytest def test_tp_plan_covers_all_experts(): from gptoss_tp import GptOssTpAuditor # 构造一个最小 gpt-oss 替身 class FakeExpert(torch.nn.Module): def __init__(self): super().__init__() self.w1 torch.nn.Linear(64, 64) self.w2 torch.nn.Linear(64, 64) self.w3 torch.nn.Linear(64, 64) class FakeModel(torch.nn.Module): def __init__(self): super().__init__() self.experts torch.nn.ModuleList([FakeExpert() for _ in range(4)]) def named_modules(self): for i, e in enumerate(self.experts): yield fmodel.layers.0.mlp.experts.{i}.w1, e.w1 yield fmodel.layers.0.mlp.experts.{i}.w2, e.w2 yield fmodel.layers.0.mlp.experts.{i}.w3, e.w3 auditor GptOssTpAuditor() plan auditor.build_plan(FakeModel(), {}) for i in range(4): assert fmodel.layers.*.mlp.experts.{i}.w1 in plan assert plan[fmodel.layers.*.mlp.experts.{i}.w2] row_parallel def test_router_and_tied_weights_in_plan(): from gptoss_tp import GptOssTpAuditor plan GptOssTpAuditor().build_plan(FakeModel(), {}) assert plan.get(model.layers.*.mlp.gate) col_parallel assert plan.get(lm_head) col_parallel assert plan.get(model.embed_tokens) col_parallel def test_no_unsharded_expert(): from gptoss_tp import audit plan GptOssTpAuditor().build_plan(FakeModel(), {}) unsharded, asym audit(plan) assert unsharded 0, 仍有专家层未切分 - OOM 风险 assert asym is False, router 切分不对称 - 集合通信卡死CI 跑pytest tests/test_gptoss_tp.py以后只要base_model_tp_plan又漏了 gpt-oss 的专家/router测试立刻红灯。八、排查清单当 gpt-oss 开 TP 失败按顺序查OOM 且 TP2 →tp_plan漏登mlp.experts.*用GptOssTpAuditor补全。形状错Xx4096 vs 8192x4096→ 共享 lm_head 与 embed_tokens 切分不同步二者都标col_parallel并声明_tied_weights_keys。卡死/断言在 router all-to-all →mlp.gate没切或切分方式与专家不匹配补col_parallel。确认num_experts与实际一致用discover_experts自动数避免手填错。长期方案把 gpt-oss 的 MoE/共享权重 tp_plan 收进模型类的base_model_tp_plan而不是每次手动 patch。九、小结Add support for GPT_OSS with tp_plan 的根因是gpt-oss 的base_model_tp_plan套用了普通 Transformer 模板漏掉了 MoE 专家层与 router且没同步共享 lm_head/embed_tokens 的切分导致专家全量复制OOM、router 集合通信不对称卡死、共享权重形状错。第一层手动把mlp.experts.*、mlp.gate、lm_head/embed_tokens同步切分补齐进tp_plan立即能跑。第二层用GptOssTpAuditor自动发现专家数并补全切分计划杜绝手填遗漏。第三层pytest 断言专家全部覆盖、router 与共享权重同步切分、无未切分专家防止回归。记住MoE 模型开 TP专家层与 router 必须进tp_plan且权重共享的两个模块要同步切分——否则要么 OOM要么集合通信不对称卡死。