从零理解 MoE

篇外 E：从零理解 MoE（混合专家模型）

这是本系列的又一篇外篇。

如果你关注最近的大模型，一定反复看到这个词：MoE。

• Mixtral 8x7B
• DeepSeek-V2 / V3
• Qwen 的 MoE 版本
• GPT-4 被广泛传为 MoE 架构

它几乎是"大模型越做越大"这条路上最重要的架构创新之一。

这篇我们从直觉讲到底层，再到代码，把 MoE 彻底讲明白。

一、先看矛盾：想让模型更强，但不想让它更慢

大模型的“能力”和“参数量”高度相关：参数越多，通常越强。

但参数越多，问题也越大：

于是产生了一个尖锐的矛盾：

我们想要“大模型的容量”，但不想要“大模型的计算量”。

MoE 就是为解决这个矛盾而生的。

二、核心思想：不要用所有参数，只用“相关”的那部分MoE 的灵感来自一个朴素直觉：

处理不同的问题，其实不需要动用全部“脑细胞”。

比如：

• 问数学题 → 调用“数学专家”
• 问法律问题 → 调用“法律专家”
• 问代码 → 调用“编程专家”

每次都只激活一小部分参数，剩下的“专家”保持待命。

这就是 MoE 的两大关键：

1. 把大模型拆成多个“专家”子网络
2. 每个 token 只激活其中几个专家（这叫稀疏激活）

结果：

总参数量很大（容量大），但每个 token 实际计算量很小（速度快）。

三、对比：Dense vs MoE假设总参数 8×7B = 56B：

Dense（稠密）模型

每个 token 都要过所有参数：

输入 token → [全部 56B 参数] → 输出

计算量对应56B。

MoE（稀疏）模型

假设有 8 个专家，每次只激活Top-2：

输入 token → [选 2 个专家] → 输出 ↑ 只用 ~2×7B = 14B 参数

计算量只对应14B，但总容量是56B。

用一个"看起来像 56B 的大脑"，跑出"14B 的速度"——这就是 MoE 的魔法。

四、MoE 的核心组件一个 MoE 层通常包含三个部分：

输入 token x │ ▼ ┌─────────────────┐ │ Router / Gate │ ← 决定用哪些专家 └─────────────────┘ │ ┌──────┼──────┐ ▼ ▼ ▼ Expert1 Expert2 Expert3 ... ExpertN │ │ │ └──────┴──────┘ │ ▼ 加权求和 → 输出

1. 专家（Experts）

每个专家本质就是一个前馈网络（FFN）：

Experti(x)=W2⋅ReLU(W1x)\text{Expert}_i(x) = W_2 \cdot \text{ReLU}(W_1 x)Experti​(x)=W2​⋅ReLU(W1​x)

它们结构相同，但参数独立——各自学出不同的"专长"。

2. 路由器（Router / Gate）

一个小网络，负责决定“这个 token 该交给哪些专家”：

logits=x⋅Wg\text{logits} = x \cdot W_glogits=x⋅Wg​

然后取Top-K：

topk=TopK(softmax(logits),k)\text{topk} = \text{TopK}(\text{softmax}(\text{logits}), k)topk=TopK(softmax(logits),k)

3. 加权组合

output=∑i∈TopKgi⋅Experti(x)\text{output} = \sum_{i \in \text{TopK}} g_i \cdot \text{Expert}_i(x)output=i∈TopK∑​gi​⋅Experti​(x)

其中gig_igi​是归一化后的门控权重。

五、手写一个最小 MoE 层先把零件写出来：

importtorchimporttorch.nnasnnimporttorch.nn.functionalasFclassExpert(nn.Module):"""单个专家：一个普通的前馈网络"""def__init__(self,d_model,d_hidden):super().__init__()self.fc1=nn.Linear(d_model,d_hidden)self.fc2=nn.Linear(d_hidden,d_model)defforward(self,x):returnself.fc2(F.relu(self.fc1(x)))classMoELayer(nn.Module):def__init__(self,d_model,d_hidden,num_experts=8,top_k=2):super().__init__()self.num_experts=num_experts self.top_k=top_k# 每个专家独立参数self.experts=nn.ModuleList([Expert(d_model,d_hidden)for_inrange(num_experts)])# 路由器：把 d_model 映射到 num_experts 个分数self.router=nn.Linear(d_model,num_experts,bias=False)defforward(self,x):""" x: [batch, d_model] 这里按单个 token 处理 """# 1. 路由器打分logits=self.router(x)# [batch, num_experts]# 2. 选择 Top-K 专家topk_vals,topk_idx=torch.topk(logits,self.top_k,dim=-1)# 3. 归一化门控权重gates=F.softmax(topk_vals,dim=-1)# [batch, top_k]# 4. 加权求和参与专家的输出output=torch.zeros_like(x)foriinrange(self.top_k):expert_id=topk_idx[:,i]# 每个样本选了谁gate=gates[:,i].unsqueeze(-1)# [batch, 1]foreinrange(self.num_experts):mask=(expert_id==e)ifmask.any():expert_out=self.experts[e](x[mask])output[mask]+=gate[mask]*expert_outreturnoutput,topk_idx,gates# 测试torch.manual_seed(0)d_model,d_hidden,num_experts,top_k=16,32,8,2moe=MoELayer(d_model,d_hidden,num_experts,top_k)x=torch.randn(4,d_model)# 4 个 tokenout,idx,gates=moe(x)print("输出形状:",out.shape)# [4, 16]print("每个 token 选中的专家:",idx.tolist())print("门控权重:\n",gates.round(decimals=3))

运行后你会看到：

• 每个 token 只选了 2 个专家（top_k=2）
• 不同 token 可能选不同的专家

这正是 MoE 的核心行为。

六、验证稀疏性：算算到底省了多少```python

def compute_flops_compare(num_experts, top_k, params_per_expert, d_model):
“”“粗略对比 Dense 和 MoE 每个 token 的计算量”“”
dense = num_experts * params_per_expert # 全用
moe = top_k * params_per_expert # 只用 top_k

print(f"总专家数: {num_experts}, 激活: {top_k}") print(f"总参数量 (容量): {dense:,}") print(f"每 token 激活参数: {moe:,}") print(f"实际算力占比: {100 * moe / dense:.1f}%")

compute_flops_compare(num_experts=8, top_k=2,
params_per_expert=7_000_000_000, d_model=4096)

输出：

总专家数: 8, 激活: 2
总参数量 (容量): 56,000,000,000
每 token 激活参数: 14,000,000,000
实际算力占比: 25.0%

**容量是 56B，但每个 token 只花 25% 的算力——这就是“Mixtral 8x7B 其实只激活 13B”的由来。** --- ## 七、训练 MoE 的难点：让专家“各司其职” MoE 看似美好，但训练时有个著名的大坑：**负载不均（load imbalance）**。 ### 问题 路由器容易“偷懒”： - 总把 token 都发给**少数几个专家** - 其他专家永远学不到东西 → **“专家崩塌”（expert collapse）** - 被选中的专家累死，没选中的“躺平” ### 解决办法：负载均衡损失 在训练损失里加一项，**惩罚“专家使用不均”**： ```python import torch import torch.nn.functional as F def load_balancing_loss(router_logits, num_experts, top_k): """ router_logits: [num_tokens, num_experts] 目标：让每个专家处理的 token 比例尽量均匀 """ # 每个专家被选中的比例（importance） probs = F.softmax(router_logits, dim=-1) # [T, E] importance = probs.mean(dim=0) # [E] 每个专家的平均概率 # 实际路由到每个专家的 token 比例（load） _, topk_idx = torch.topk(probs, top_k, dim=-1) # [T, k] one_hot = F.one_hot(topk_idx, num_experts).float() # [T, k, E] load = one_hot.sum(dim=(0, 1)) # [E] 每个专家收到的 token 数 load = load / load.sum() # 负载均衡损失 = num_experts * Σ(importance × load) loss = num_experts * torch.sum(importance * load) return loss torch.manual_seed(0) num_tokens, num_experts, top_k = 100, 8, 2 router_logits = torch.randn(num_tokens, num_experts) loss = load_balancing_loss(router_logits, num_experts, top_k) print(f"负载均衡损失: {loss.item():.4f}") print("（越接近 1.0 表示越均衡；偏离越大越不均衡）")

这个损失逼着路由器“雨露均沾”，让所有专家都能训练到。

为什么需要专家容量限制？

工程上还给每个专家设一个capacity（容量上限），超出就丢弃（token dropping），防止某个专家过载拖慢整个 batch：

defexpert_capacity(num_tokens,num_experts,top_k,factor=1.25):"""每个专家最多能处理多少 token"""capacity=factor*num_tokens*top_k/num_expertsreturnint(capacity)print("专家容量:",expert_capacity(2048,8,2))# 约 640

八、MoE 在真实模型里长什么样？

现代大模型里的 MoE，通常是把 Transformer 里的 FFN 换成 MoE 层：

标准 Transformer Block: MoE Transformer Block: ┌────────────────────┐ ┌────────────────────┐ │ Multi-Head Attn │ │ Multi-Head Attn │ └────────────────────┘ └────────────────────┘ ┌────────────────────┐ ┌────────────────────┐ │ FFN (一个) │ ───────▶ │ MoE FFN (N 个专家)│ └────────────────────┘ └────────────────────┘

注意：注意力层通常不 MoE 化，只替换 FFN。

真实案例对比

DeepSeek 的两个创新值得一提

1. 细粒度专家（Fine-Grained）
不只用 8 个“大专家”，而是拆成几百个“小专家”，组合更灵活。

2. 共享专家（Shared Expert）
保留少数所有 token 都会经过的专家，负责捕捉通用知识；其余专家负责差异化专长。

classMoEWithShared(nn.Module):"""共享专家 + 路由专家的示意结构"""def__init__(self,d_model,d_hidden,n_routed=8,top_k=2):super().__init__()self.shared=Expert(d_model,d_hidden)# 永远的公共专家self.routed=MoELayer(d_model,d_hidden,n_routed,top_k)defforward(self,x):shared_out=self.shared(x)# 通用知识routed_out,idx,gates=self.routed(x)# 专长知识returnshared_out+routed_out,idx,gates

共享专家保证通用能力，路由专家提供专业分工。

九、MoE 的优缺点一览

十、常见问题 FAQ

1. MoE 和集成学习（Ensemble）有什么区别？

集成学习是训练多个完整模型，推理时把它们的输出全部加权平均；而 MoE 是一个模型内部拆成多个专家，每个 token 只激活其中Top-K个（稀疏激活），其余专家不参与计算。前者算力随模型数线性增长，后者容量大但算力几乎不变。

2. 为什么注意力层不 MoE 化？

注意力层负责 token 之间的全局交互，每个 token 都需要看到其他 token 的信息，无法按“专长”拆分；而 FFN 是逐 token 独立计算的，天然适合“按需激活”。所以 MoE 只替换 FFN，注意力保持 Dense。

3. Top-K 的 K 值如何选择？

K 越大，激活的专家越多，容量利用越充分，但算力也越高。常见取K=2（如 Mixtral 8x7B），兼顾容量与速度。K 过大会失去稀疏性优势，K=1 则退化为“单专家”，路由抖动风险更高。

4. 专家数量越多越好吗？

不是。专家多 → 容量大、组合更灵活（如 DeepSeek-V2 的 160 个细粒度专家），但显存需求和通信开销也同步上涨，且更容易出现负载不均、专家崩塌。实际要在容量、算力、显存、训练稳定性之间权衡。

5. MoE 在推理时如何做 batch 并行？

batch 内不同 token 会路由到不同专家，工程上先把 token 按所选专家分组，再让每个专家并行处理自己分到的 token（如第五节代码里的mask分组），最后按门控权重加权合并。专家分散在不同设备时，还需跨设备通信汇总结果。

✅ 优点

❌ 缺点 / 挑战

特别注意第一个：MoE 省的是计算，不是显存。你可以理解为——一个"参数很多但每次只用一部分"的模型。

defmoe_memory_truth(total_params_b,active_params_b):"""澄清常见误解"""print(f"显存需求（要装下所有参数）: ≈{total_params_b}B 参数")print(f"计算量（每 token 激活）: ≈{active_params_b}B 参数")print("\n→ MoE 省算力，不省显存！")moe_memory_truth(total_params_b=47,active_params_b=13)# Mixtral 8x7B

十、总结

MoE 是什么

MoE 把模型的 FFN 层拆成多个“专家”，用路由器为每个 token 只挑少量专家（Top-K）计算，实现“大容量、低算力”。

三句话概括

1. 稀疏激活：总参数大，但每 token 只用一小部分
2. 路由器决定分工：Top-K 选择 + 门控加权
3. 负载均衡是关键：训练必须防止专家崩塌

核心公式

MoE(x)=∑i∈TopK(softmax(xWg))gi⋅Experti(x)\text{MoE}(x) = \sum_{i \in \text{TopK}(\text{softmax}(xW_g))} g_i \cdot \text{Expert}_i(x)MoE(x)=i∈TopK(softmax(xWg​))∑​gi​⋅Experti​(x)

一句话总结

MoE 用"专家分工 + 按需激活"，让模型在参数量暴涨的同时，计算量几乎不变——这是当前大模型"又大又快"的核心秘密之一。