【Bug已解决】FSDP + torch.nn.Parameter (MoE layer) lora fine-tuning doesn‘t work 解决方案-港品优选

【Bug已解决】FSDP + torch.nn.Parameter (MoE layer) lora fine-tuning doesn't 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），于是专家部分原样前向、无法微调，梯度也传不进去；
用 FSDP1（FullyShardedDataParallel）包裹时，因为专家是nn.Parameter而不是nn.Module子模块，FSDP 无法把它作为独立单元分片，门控选择专家时的切片操作与分片状态冲突，报full tensor required一类错误；
用summon_full_params想临时聚合参数做评估时，发现nn.Parameter专家张量没有被正确还原成完整形态。

一句话：MoE 层如果用torch.nn.Parameter（而不是nn.Module）承载专家权重，LoRA 注入和 FSDP 分片都会“看不见”它。

二、背景

要搞清楚为什么这套组合会坏，得分别看 FSDP 和 PEFT 各自依赖什么：

FSDP（FullyShardedDataParallel）的分片单位是“参数”，但它通过遍历模块树来发现参数。FSDP1 以nn.Module为包裹单元，对每个被wrap的模块，把它下面的nn.Parameter收集起来做分片。nn.ParameterList虽然也是参数，但如果它直接挂在一个非nn.Module容器里（比如你用一个普通 Pythonlist存nn.Parameter），FSDP 的递归遍历会漏掉它们。即便用nn.ParameterList，FSDP1 也倾向于把整个模块当成一个分片单元，无法对“列表里单个专家”做细粒度分片，门控按索引切片就会和分片后的形状打架。

PEFT（LoRA）的注入逻辑是：遍历模块的named_modules()，凡是命中target_modules的nn.Linear/nn.Conv1d等模块，就包一层lora.Linear。如果你的专家权重是nn.Parameter（一个裸张量），那它既不是nn.Linear也不是任何 Module，PEFT 根本不会去碰它，LoRA 自然挂不上。

于是矛盾就来了：你想用nn.Parameter存专家（省内存、方便切片），但 FSDP 和 LoRA 都假设权重是“挂在nn.Module下的nn.Parameter”，二者视角不一致。

下面用最小可运行代码复现“专家用nn.Parameter时 LoRA 挂不上、FSDP 分片错位”。

三、根因

根因有三条：

专家不是nn.Module：torch.nn.Parameter（尤其放在普通list或单个大张量里）不会出现在model.named_modules()的模块遍历中，PEFT 遍历模块时看不到“可注入点”，LoRA 挂不上。
FSDP1 以模块为分片单元：当专家是nn.ParameterList里的一项，FSDP1 无法对单个专家做独立分片；门控用index_select取专家时，张量已被分片成[shard, ...]，索引语义失效。
requires_grad设置错位：nn.Parameter默认requires_grad=True，但 LoRA 只训练lora_A/lora_B；若专家是nn.Parameter且你忘记requires_grad_(False)，会和 LoRA 一起更新，既浪费显存也容易和分片冲突。

正确的方向是：把专家做成真正的nn.Module子模块（每个专家是一个nn.Linear），这样 FSDP 和 PEFT 都能正确发现它们；并优先使用 FSDP2（torch.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, dim=16, n_experts=4): 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, bias=False) def forward(self, x, gate=None): 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, dim=16): super().__init__() self.linear = nn.Linear(dim, dim) def forward(self, x): return self.linear(x) class GoodMoE(nn.Module): def __init__(self, dim=16, n_experts=4): 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, bias=False) def forward(self, x, gate=None): 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(r=4, lora_alpha=8, target_modules=["linear"], lora_dropout=0.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… 并正常注入 LoRA；
FSDP 能对每个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()

修复 3：FSDP1 包裹时按专家模块 wrap

from 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_policy=auto_wrap, device_id=torch.cuda.current_device())

因为专家是nn.Module，FSDP1 能正确把它们作为分片单元。

六、解决方案（第二层：结构性改进）

改进 1：优先用 FSDP2（`fully_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, gate=None): 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:e+1, 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, d=16): super().__init__() self.linear = nn.Linear(d, d) def forward(self, x): return self.linear(x) class GoodMoE(nn.Module): def __init__(self, d=16, e=4): super().__init__() self.experts = nn.ModuleList([Expert(d) for _ in range(e)]) self.router = nn.Linear(d, e, bias=False) def forward(self, x): g = torch.softmax(self.router(x.mean(dim=-2)), -1) out = sum(g[:, i:i+1, 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(r=4, lora_alpha=8, target_modules=["linear"], lora_dropout=0.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 版本：优先 FSDP2（fully_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 doesn't work的根因是：当专家权重以torch.nn.Parameter（尤其nn.ParameterList/ 裸张量）形式存在时，它既不出现在named_modules()的模块遍历里（LoRA 注入不到），也不能被 FSDP1 作为独立单元分片（门控切片与分片形状冲突）。

最小修复是把专家改成nn.Module子模块（每个专家一个nn.Linear），LoRA 和 FSDP 立刻都能看见；结构性改进是优先用 FSDP2 的参数级分片、把门控从逐样本index_select改成软路由加权、写校验器强制专家必须是 Module；最后用 CI 测试守护“LoRA 命中专家层、专家是 Module、禁止 ParameterList”。这样 MoE 微调才能在多卡分片下既正确又可训练。

企业官网建设流程全解析