【Bug已解决】FSDP torch.nn.Parameter (MoE layer) lora fine-tuning doesnt work 解决方案一、现象长什么样在做一个 MoE混合专家模型微调时很多人会把多个专家存成一个nn.ParameterList或一个大的torch.nn.Parameter张量然后用门控router在前向时按索引取出对应专家的权重做计算。问题出现在给这个 MoE 层叠 LoRA并用 FSDP 做分片训练时出现以下几种情况报错RuntimeError: ... must be a Parameter or a Tensor, got ...或者FSDP cannot shard parameter ... because it is not found in the module tree训练能跑但 LoRA 的lora_A/lora_B根本没被 FSDP 纳入分片导致某个 rank 上显存爆掉而其他 rank 几乎空载更隐蔽LoRA 只挂到了外层线性层nn.Parameter形式的专家权重没被 PEFT 识别因为里面不是nn.Linear于是专家部分原样前向、无法微调梯度也传不进去用 FSDP1FullyShardedDataParallel包裹时因为专家是nn.Parameter而不是nn.Module子模块FSDP 无法把它作为独立单元分片门控选择专家时的切片操作与分片状态冲突报full tensor required一类错误用summon_full_params想临时聚合参数做评估时发现nn.Parameter专家张量没有被正确还原成完整形态。一句话MoE 层如果用torch.nn.Parameter而不是nn.Module承载专家权重LoRA 注入和 FSDP 分片都会“看不见”它。二、背景要搞清楚为什么这套组合会坏得分别看 FSDP 和 PEFT 各自依赖什么FSDPFullyShardedDataParallel的分片单位是“参数”但它通过遍历模块树来发现参数。FSDP1 以nn.Module为包裹单元对每个被wrap的模块把它下面的nn.Parameter收集起来做分片。nn.ParameterList虽然也是参数但如果它直接挂在一个非nn.Module容器里比如你用一个普通 Pythonlist存nn.ParameterFSDP 的递归遍历会漏掉它们。即便用nn.ParameterListFSDP1 也倾向于把整个模块当成一个分片单元无法对“列表里单个专家”做细粒度分片门控按索引切片就会和分片后的形状打架。PEFTLoRA的注入逻辑是遍历模块的named_modules()凡是命中target_modules的nn.Linear/nn.Conv1d等模块就包一层lora.Linear。如果你的专家权重是nn.Parameter一个裸张量那它既不是nn.Linear也不是任何 ModulePEFT 根本不会去碰它LoRA 自然挂不上。于是矛盾就来了你想用nn.Parameter存专家省内存、方便切片但 FSDP 和 LoRA 都假设权重是“挂在nn.Module下的nn.Parameter”二者视角不一致。下面用最小可运行代码复现“专家用nn.Parameter时 LoRA 挂不上、FSDP 分片错位”。三、根因根因有三条专家不是nn.Moduletorch.nn.Parameter尤其放在普通list或单个大张量里不会出现在model.named_modules()的模块遍历中PEFT 遍历模块时看不到“可注入点”LoRA 挂不上。FSDP1 以模块为分片单元当专家是nn.ParameterList里的一项FSDP1 无法对单个专家做独立分片门控用index_select取专家时张量已被分片成[shard, ...]索引语义失效。requires_grad设置错位nn.Parameter默认requires_gradTrue但 LoRA 只训练lora_A/lora_B若专家是nn.Parameter且你忘记requires_grad_(False)会和 LoRA 一起更新既浪费显存也容易和分片冲突。正确的方向是把专家做成真正的nn.Module子模块每个专家是一个nn.Linear这样 FSDP 和 PEFT 都能正确发现它们并优先使用 FSDP2torch.distributed.fsdp.fully_shard它对参数级分片友好能和 LoRA 干净组合。四、最小可运行复现下面用单进程把 world_size 模拟为 1演示两种结构的差异。需要torch多卡部分用注释标注单卡也能跑通结构验证。import torch import torch.nn as nn from peft import get_peft_model, LoraConfig # 反例专家用 nn.ParameterList 存裸张量 - LoRA 挂不上、FSDP 难分片 class BadMoE(nn.Module): def __init__(self, dim16, n_experts4): super().__init__() self.n_experts n_experts # 4 个专家 4 个 [dim, dim] 权重存成 ParameterList self.expert_w nn.ParameterList( [nn.Parameter(torch.randn(dim, dim)) for _ in range(n_experts)] ) self.router nn.Linear(dim, n_experts, biasFalse) def forward(self, x, gateNone): if gate is None: gate self.router(x.mean(dim-2)) # [B, n_experts] idx gate.argmax(dim-1) # [B] # 按样本选专家切片与分片冲突点 out torch.stack([x[i] self.expert_w[idx[i]] for i in range(x.shape[0])]) return out # 正例专家是真正的 nn.Module 子模块 - LoRA / FSDP 都看得见 class Expert(nn.Module): def __init__(self, dim16): super().__init__() self.linear nn.Linear(dim, dim) def forward(self, x): return self.linear(x) class GoodMoE(nn.Module): def __init__(self, dim16, n_experts4): super().__init__() self.n_experts n_experts self.experts nn.ModuleList( [Expert(dim) for _ in range(n_experts)] # 每个都是 nn.Module ) self.router nn.Linear(dim, n_experts, biasFalse) def forward(self, x, gateNone): if gate is None: gate self.router(x.mean(dim-2)) idx gate.argmax(dim-1) out torch.zeros_like(x) for e in range(self.n_experts): mask (idx e) if mask.any(): out[mask] self.experts[e](x[mask]) return out # 演示 LoRA 能否命中专家 cfg LoraConfig(r4, lora_alpha8, target_modules[linear], lora_dropout0.0) bad BadMoE() try: pb get_peft_model(bad, cfg) print(BadMoE 挂 LoRA 后可训练:, sum(p.numel() for p in pb.parameters() if p.requires_grad)) except Exception as e: print(BadMoE LoRA 失败:, repr(e)) good GoodMoE() pg get_peft_model(good, cfg) print(GoodMoE 挂 LoRA 后可训练:, sum(p.numel() for p in pg.parameters() if p.requires_grad))运行后会看到BadMoE要么挂 LoRA 时找不到linear因为专家不是 Module要么trainable为 0而GoodMoE的每个Expert.linear都被lora.Linear包住可训练参数正常。五、解决方案第一层最小直接修复修复 1把专家改成nn.Module子模块如上GoodMoE用nn.ModuleList([Expert(dim) ...])代替nn.ParameterList。这样PEFT 能在named_modules()里看到experts.0.linear、experts.1.linear… 并正常注入 LoRAFSDP 能对每个Expert模块独立分片。target_modules直接写[linear]或精确写[experts.0.linear, ...]。修复 2冻结基础专家权重只训练 LoRA# 注入后确认只有 lora_A / lora_B 可训练 for name, p in pg.named_parameters(): if lora_ not in name: p.requires_grad_(False) pg.print_trainable_parameters()修复 3FSDP1 包裹时按专家模块 wrapfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy # 每个 Expert 是一个 wrap 单元 auto_wrap size_based_auto_wrap_policy model FSDP(pg, auto_wrap_policyauto_wrap, device_idtorch.cuda.current_device())因为专家是nn.ModuleFSDP1 能正确把它们作为分片单元。六、解决方案第二层结构性改进改进 1优先用 FSDP2fully_shard与 LoRA 干净组合FSDP2 在参数级别做分片不需要把模块当成“分片单元”和 LoRA 注入后的lora.Linear本身也是 Module天然兼容from torch.distributed.fsdp import fully_shard def apply_fsdp2(module): # 对每个子模块单独分片门控切片不再与“模块级分片”冲突 for child in module.children(): apply_fsdp2(child) fully_shard(module) model GoodMoE() model get_peft_model(model, cfg) model model.cuda() apply_fsdp2(model) # 参数级分片LoRA 参数也一起被管理FSDP2 的好处是它分片的是nn.Parameter本身而不是“模块”所以即便你仍有少量nn.Parameter形式的权重如 router也能正确分片不会和门控索引打架。改进 2门控用torch.index_select而非 Python 循环保证可微分与分片友好def forward(self, x, gateNone): if gate is None: gate self.router(x.mean(dim-2)) # [B, E] weights torch.softmax(gate, dim-1) # [B, E] # 加权混合所有专家输出避免按样本切分对分片友好 out torch.zeros_like(x) for e in range(self.n_experts): out out weights[:, e:e1, None] * self.experts[e](x) return out这种“软路由加权”不依赖逐样本index_select对 FSDP 分片后的张量是安全的所有 rank 都持有完整 batch 的逻辑切片参数被分片但计算图一致。改进 3统一“专家必须是 Module”的约定写个校验器def assert_experts_are_modules(moe): for name, mod in moe.named_modules(): if name.endswith(experts) or expert in name: assert isinstance(mod, nn.Module), \ f{name} 必须是 nn.Module不能是裸 nn.Parameter return True七、解决方案第三层断言 / CI 守护import torch import torch.nn as nn import pytest from peft import get_peft_model, LoraConfig class Expert(nn.Module): def __init__(self, d16): super().__init__() self.linear nn.Linear(d, d) def forward(self, x): return self.linear(x) class GoodMoE(nn.Module): def __init__(self, d16, e4): super().__init__() self.experts nn.ModuleList([Expert(d) for _ in range(e)]) self.router nn.Linear(d, e, biasFalse) def forward(self, x): g torch.softmax(self.router(x.mean(dim-2)), -1) out sum(g[:, i:i1, None] * self.experts[i](x) for i in range(len(self.experts))) return out def _trainable(model): return sum(p.numel() for p in model.parameters() if p.requires_grad) def test_lora_hits_expert_linear(): cfg LoraConfig(r4, lora_alpha8, target_modules[linear], lora_dropout0.0) peft get_peft_model(GoodMoE(), cfg) assert _trainable(peft) 0 # 确认专家层真的被包成了 lora.Linear from peft.tuners.lora import LoraLinearLayer hit any(lora in n for n, _ in peft.named_modules() if experts in n) assert hit def test_experts_are_modules(): m GoodMoE() for name, mod in m.named_modules(): if experts in name: assert isinstance(mod, nn.Module) def test_no_bare_parameterlist_for_experts(): m GoodMoE() for name, mod in m.named_modules(): assert not isinstance(mod, nn.ParameterList), \ f{name} 不应是 ParameterList这三个测试守护“LoRA 能命中专家线性层”“专家是 Module”“专家不用 ParameterList”。八、排查清单MoE LoRA FSDP 组合出问题时按序查专家是不是nn.Module用named_modules()确认能看到experts.i.linear若是nn.Parameter/裸张量PeFT 挂不上。target_modules是否命中打印专家模块真实名确认linear等子串匹配。requires_grad是否只留给 LoRA注入后冻结非lora_参数。FSDP 版本优先 FSDP2fully_shard参数级分片对 LoRA 友好FSDP1 要按专家模块wrap。门控是否依赖逐样本index_select分片后张量形状变化逐样本循环会失效改用软路由加权。summon_full_params是否必要评估时才聚合训练时保持分片。router 参数router 是nn.Linear也可单独决定是否加 LoRA通常不加避免破坏负载均衡。显存核对分片后每个 rank 显存应均衡若某 rank 爆显存说明有参数漏分片通常是裸nn.Parameter。九、小结FSDP torch.nn.Parameter (MoE layer) lora fine-tuning doesnt work的根因是当专家权重以torch.nn.Parameter尤其nn.ParameterList/ 裸张量形式存在时它既不出现在named_modules()的模块遍历里LoRA 注入不到也不能被 FSDP1 作为独立单元分片门控切片与分片形状冲突。最小修复是把专家改成nn.Module子模块每个专家一个nn.LinearLoRA 和 FSDP 立刻都能看见结构性改进是优先用 FSDP2 的参数级分片、把门控从逐样本index_select改成软路由加权、写校验器强制专家必须是 Module最后用 CI 测试守护“LoRA 命中专家层、专家是 Module、禁止 ParameterList”。这样 MoE 微调才能在多卡分片下既正确又可训练。