【Bug已解决】Nemotron-3-Nano NemotronHTopkRouter 解决方案一、现象长什么样你在用 PEFT 给Nemotron-3-Nano一个 MoE 模型做 LoRA 微调模型的专家路由用NemotronHTopkRouter一种 top-k 路由把 token 分给前 k 个专家。但你发现不知道该把 LoRA 挂到哪是挂 router 还是挂每个 expert 的线性层挂 expert 时target_modules怎么写才能命中 MoE 块里的专家层命名常是block.self_attn...之外还有block.mlp.experts.X.xxx路由权重router 的 gate 参数要训吗训了会不会破坏负载均衡训练时 top-k 路由的dispatch/combine逻辑和 LoRA 的前向是否兼容。本文把NemotronHTopkRouter这类 top-k 路由在 PEFT/LoRA 下的正确用法讲清。二、背景MoEMixture of Experts层结构一个 token 进来先过一个router门控router 算每个专家的得分取top-k个最高分专家把 token 送到这些专家算再按得分加权 combine 回。Nemotron 的NemotronHTopkRouter就是这套逻辑H 可能指某个变体/层级。PEFT/LoRA 在 MoE 上微调的常见做法挂 LoRA 到 expert 的线性层每个 expert 内部是标准 FFN两个 Linear 激活。对 expert 的w1/w2或gate_proj/up_proj/down_proj挂 LoRA最常用也最安全。是否训 router一般不训 router 的 gate保持负载均衡除非专门做路由调优。LoRA 默认也不碰 router除非target_modules明确包含 router 的层。top-k 路由与 LoRA 兼容LoRA 是“在 expert 线性层上加低秩增量”路由的 dispatch/combine 逻辑不变二者天然兼容——LoRA 只改 expert 内部计算不改 token 分发。难点在 PEFT 的自动 target 识别MoE 块的命名和稠密块不同PEFT 默认的 target 映射可能漏掉 expert 层需要显式写出target_modules或扩展映射。三、根因为什么容易做错根因 Atarget_modules没命中 MoE expert 层最常见。默认 targetq_proj/v_proj等只命中 attention没命中experts.*.w1/w2LoRA 只挂在 attention漏了 MoE 主体。根因 B误训 router gate破坏负载均衡若target_modules误包含 router 的 gate 参数LoRA 会改路由权重可能让 top-k 选择失衡某些专家被过度/过少选中。根因 C把 LoRA 挂到 router 而非 expertrouter 通常很小一个线性映射 token→expert 分挂 LoRA 收益低且危险应挂 expert 内部。根因 DPEFT 自动映射不含 Nemotron MoE 命名PEFT 默认TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES可能没 Nemotron-3-Nano 的 expert 命名需手动写 target。根因小结MoE 上 LoRA 应挂 expert 内部线性层不挂/慎挂 routertarget_modules必须显式命中experts.*层默认映射常漏修复手动写 target 不训 router除非专门调路由 确认 top-k 兼容。四、最小可运行复现下面脚本演示“MoE top-k 路由 给 expert 挂 LoRA”的结构与正确 targetimport torch import torch.nn as nn import torch.nn.functional as F class Expert(nn.Module): def __init__(self, d): super().__init__() self.w1 nn.Linear(d, d * 4) # up self.w2 nn.Linear(d * 4, d) # down self.w3 nn.Linear(d, d * 4) # gate def forward(self, x): return self.w2(F.silu(self.w3(x)) * self.w1(x)) class NemotronHTopkRouter(nn.Module): top-k 路由token - 选前 k 个专家 - 加权 combine def __init__(self, d, n_experts, k2): super().__init__() self.gate nn.Linear(d, n_experts, biasFalse) # 路由 gate通常不训 LoRA self.experts nn.ModuleList([Expert(d) for _ in range(n_experts)]) self.k k def forward(self, x): scores self.gate(x) # [B, T, E] topk scores.topk(self.k, dim-1) # 前 k 专家 idx, w topk.indices, topk.values.softmax(-1) out torch.zeros_like(x) for i in range(self.k): e topk.indices[..., i] expert_out torch.stack([self.experts[e[b, t]](x[b, t:b1])[0] for b in range(x.shape[0]) for t in range(x.shape[1])]) # 简化 combine示意 out out w[..., i:i1] * expert_out.reshape_as(out) return out def demo(): torch.manual_seed(0) model NemotronHTopkRouter(16, n_experts4, k2) # 正确 target命中 expert 的 w1/w2/w3不命中 gate target [] for n, m in model.named_modules(): if isinstance(m, nn.Linear) and experts in n and gate not in n: target.append(n) print(应挂 LoRA 的 expert 层:, target[:4], ...) x torch.randn(2, 5, 16) out model(x) print(MoE top-k 路由前向形状:, tuple(out.shape)) if __name__ __main__: demo()运行后target 正确命中 expert 的w1/w2/w3不命中 routergate前向正常证明 MoE LoRA 的正确挂法。五、解决方案第一层最小直接修复显式写target_modules命中 expert 内部层排除 router gatefrom peft import get_peft_model, LoraConfig # 用子串命中所有 expert 的线性层但用 name_filter 排除 router gate config LoraConfig( r8, target_modules[w1, w2, w3], # 命中 expert 内部 # 或用更精确的正则式PEFT 支持列表子串匹配 ) # 关键不要把 router 的 gate 列入 target # 若要完全排除可后处理 def exclude_router(names): return [n for n in names if gate not in n and router not in n] model get_peft_model(base_nemotron, config)若 PEFT 默认映射不含 Nemotron手动列出 expert 路径target_modules [ model.layers.*.mlp.experts.*.w1, model.layers.*.mlp.experts.*.w2, model.layers.*.mlp.experts.*.w3, ]六、解决方案第二层结构性改进6.1 不训 router除非专门调路由# router gate 保持冻结避免破坏 top-k 负载均衡 for n, p in model.named_parameters(): if router in n or gate in n: p.requires_grad False6.2 确认 top-k 与 LoRA 兼容LoRA 只改 expert 内部线性计算router 的 dispatch/combine 不变天然兼容。无需改路由逻辑。6.3 给 PEFT 注册 Nemotron MoE target 映射# 扩展 PEFT 的 transformers 映射让 Nemotron-3-Nano 自动识别 expert 层 NEMOTRON_TARGET [w1, w2, w3, q_proj, v_proj]七、解决方案第三层断言 / CI 守护import torch import pytest def test_lora_targets_experts(model_nemotron): cfg LoraConfig(r8, target_modules[w1, w2, w3]) m get_peft_model(model_nemotron, cfg) lora_names [n for n, _ in m.named_parameters() if lora in n.lower()] assert any(experts in n for n in lora_names), LoRA 应挂到 expert 层 def test_router_not_trained(model_nemotron): cfg LoraConfig(r8, target_modules[w1, w2, w3]) m get_peft_model(model_nemotron, cfg) for n, p in m.named_parameters(): if gate in n or router in n: assert not p.requires_grad, router gate 不应被 LoRA 训练 def test_topk_forward_runs(model_nemotron): m get_peft_model(model_nemotron, LoraConfig(r8, target_modules[w1, w2, w3])) out m(torch.randn(2, 5, 16)) assert out.shape (2, 5, 16) def test_expert_lora_shape(model_nemotron): m get_peft_model(model_nemotron, LoraConfig(r8, target_modules[w1])) # 每个 expert 的 w1 应有 lora_A/lora_B assert any(experts in n and lora_A in n for n, _ in m.named_parameters())CI 跑这四条MoE LoRA 的 target/router/forward 被正确守护。八、排查清单Nemotron MoE LoRA 时查LoRA 挂到 expert 内部了吗target_modules命中experts.*.w1/w2/w3。router gate 被训了吗不该训避免破坏 top-k 负载均衡。默认 target 映射含 Nemotron MoE 吗不含就手动写 expert 路径。top-k 路由与 LoRA 兼容吗LoRA 只改 expert 内部路由逻辑不变天然兼容。每个 expert 都有 LoRA 吗断言 expert 层出现 lora_A/lora_B。前向形状对吗[B,T,d]不变。CI 测了 MoE LoRA 吗加 target/router/forward 断言防回归。九、小结“Nemotron-3-Nano NemotronHTopkRouter” 下的 PEFT 用法MoE 模型用 top-k 路由分发 token 到专家LoRA 应挂到 expert 内部线性层不挂/慎挂 router gate根因默认target_modules只命中 attention漏掉 MoE expert 层误训 router 会破坏 top-k 负载均衡第一层显式写target_modules[w1,w2,w3]命中 expert排除 router gate第二层冻结 router、确认 top-k 与 LoRA 兼容LoRA 只改 expert 内部、注册 Nemotron MoE target 映射第三层pytest 守护“LoRA 挂 expert router 不训 前向通过 每 expert 有 LoRA”。一句话Nemotron MoE LoRA 要把 adapter 挂到experts.*内部线性层w1/w2/w3刻意排除 router gate保持 top-k 负载均衡target_modules需手动命中 MoE 命名。