浅看MOE模型
#
依旧给出伟大的csdn参考文献 (opens new window)
# MOE——Mixture of Experts
大概是25年年初,deepseek的爆火把moe架构带入了大众视野;这个用更少的算力来实现比SOTA更强性能的架构确实很有意思;其实博主之前也想过这样的可能,显然用不同的专家来处理多种任务比用一个通用的模型处理多种任务听上去更合理一点对吧x 嗯我真牛()下面来简短介绍一下moe模型(其实和dense区别就仅仅几个点)
MoE模型的核心思想是将输入数据分配给不同的专家子模型,然后将所有子模型的输出进行合并,以生成最终结果。这种分配可以根据输入数据的特征进行动态调整,确保每个专家处理其最擅长的数据类型或任务方面,从而实现更高效、准确的预测。
- 相反,稠密模型就是输入数据会经过模型的所有部分,完整的计算一遍;
# 模型组成
在模型组成上与稠密模型的区别主要是
- 专家:模型中的每个专家都是一个独立的神经网络,专门处理输入数据的特定子集或特定任务;
- 门控网络:决定每个输入样本应该由哪个专家或哪些专家来处理
- 通常是一个简单的神经网络,其输入是原始的输入数据或经过预处理的特征向量。网络通过学习输入数据与各个专家之间的相关性,为每个专家计算一个权重或重要性分数。这些权重反映了输入数据在各个专家擅长的领域内的匹配程度,权重越高表示该专家对当前输入数据的处理能力越强。然后,路由器将输入数据按照计算出的权重分配给相应的专家,完成数据的路由任务。
- 输出前会经过一个softmax函数
- 噪声机制:在路由器的设计中,为了防止某些专家过于频繁地被选择,导致其他专家得不到足够的训练,通常会引入噪声机制。具体方法是在路由器的输出logits上添加一定的噪声,然后应用softmax函数
- Top-K选择策略:除了根据softmax权重进行加权求和外,路由器还可以采用Top-K选择策略来进一步简化计算和提高效率。在这种策略下,路由器只选择权重最高的K个专家来进行后续的计算,而忽略其他权重较低的专家
- 负载均衡与辅助损失函数:为了确保各个专家之间的负载均衡,避免某些专家过度负担而其他专家闲置的情况,通常会在MoE模型的训练过程中引入辅助损失函数;主要是为了专家负载均衡
# 细说专家
(推眼镜)我觉得很重要的是理解这里的专家和稠密模型里的FFN到底有什么区别
MoE(Mixture of Experts)在主流 Transformer 里一般对应的是 Dense 模型中的 FFN/MLP 子层,而不是 Attention 子层。更具体地说:把每个 block 里的“前馈网络(FFN)”从 1 个变成很多个专家 FFN,再用门控路由把每个 token 分配给少数几个专家去算;Attention 通常仍然是共享的 Dense 注意力
以常见的 Pre-LN Transformer block 为例:
- Self-Attention 子层
输入
输出仍是 - FFN/MLP 子层(又叫前馈网络)
典型是两层线性 + 激活:
其中 , (4H 只是常见比例)
Dense 模型里:Attention 负责“token 之间的信息交互”,FFN 负责“对每个 token 做非线性变换/特征提取”(逐 token 计算,不混 token)。
MoE 通常替换的就是 FFN:从 1 个 FFN 变成 EP 个专家 FFN,FFN 位置变成:
- 有
个专家,每个专家本质上就是一个 FFN:
- 有一个门控/路由器(Router/Gate),为每个 token 计算专家得分:
- 只选 Top-
个专家(常见 或 ),把 token 发给它们计算,再按门控权重合并输出:
关键点:每个 token 只激活少数专家,所以计算量接近 Dense FFN,但参数量可以大很多(“稀疏激活,大容量参数”)。
Q:为什么不是 Attention?Attention MoE
Attention 为什么通常保持 Dense
- Attention 里有 QKV 投影、注意力矩阵计算 等,结构更耦合、更敏感;做稀疏路由会引入更复杂的不稳定性与工程成本。
- Attention 是 token 之间交互的“共享通道”,很多经验表明保持其共享有利于稳定与泛化。
- 业界与论文里最常见、最成熟的是 “Sparse MoE in FFN”(如 Switch Transformer / GShard / 多数 MoE LLM 变体)。
# 通信与并行:MoE 引入 All-to-All(dispatch / combine)
MoE-FFN 的关键动作是:每个 token 只去少数几个专家(Top-k)。而在大规模训练里,专家通常被分散放在不同 GPU(Expert Parallel, EP)上,于是每个 GPU 手里的 token 必须“按路由结果重新洗牌”到拥有对应专家的 GPU——这就是 All-to-All
# 1. MoE 一次前向/反向到底通信了什么?
以 Top-1(每 token 选 1 个专家)为例,流程大致是:
- Router 在本地算分配
- 对每个 token 得到 expert id(以及权重α)。
- Dispatch(All-to-All #1):把 token 发到专家所在 GPU
- 发送内容:token 的 hidden states(形如
是本卡 token 数), 其 中 - 发送目标:拥有对应 expert 的 GPU
- 结果:每张卡收到一堆“属于本卡专家”的 token(数量不固定)
- 发送内容:token 的 hidden states(形如
- 专家计算(本地)
对收到的 token,按 expert 分桶(bucket),分别做专家 FFN(常见用 grouped GEMM 提速) - Combine(All-to-All #2):把专家输出送回 token 原位置
- 发送内容:专家输出 hidden states(同样大小
) - 目的:把结果还原回“原 token 的顺序/归属”(并应用α 加权)
- 发送内容:专家输出 hidden states(同样大小
- 反向传播
- 反向会再次经历类似的“按路由路径回传梯度”的通信,因此整体上 MoE 的通信常被粗略看成:前向 2 次 All-to-All + 反向 2 次 All-to-All(实现细节不同会略有变化,但量级直觉正确)。
# 2. 为什么 All-to-All 容易成为瓶颈?
alltoall的痛点
- 小消息 + 高频:MoE 通常发生在每个 block(或若干 block),每次只搬 token 激活,消息粒度可能不大但次数很多,容易被 latency 主导
- 不均匀通信(skew):不同 expert 接收 token 数不同 → 不同 GPU 收发量不同 → All-to-All 完成时间被最慢的那张卡拖住
- 计算与通信难以完全重叠:dispatch 完才能算专家;算完才能 combine,流水线空间有限;
一个很实用的近似:
若每个 token 发送
这里的
是本 step 的 token 总数(或本 EP 组内 token 总数,取决于你怎么划分并行组)。
# 3. Expert Parallel(EP)里“并行组”怎么组织?
工程上一般会把 GPU 划成一个或多个 EP group:
- EP group 内:分布专家,发生 All-to-All
- EP group 外:走数据并行(DP)或张量并行(TP)等
典型组合是 TP(算 Attention/线性)+ EP(算 MoE FFN)+ DP(扩 batch)。
MoE 的 All-to-All 通常发生在 EP group 内,因此 EP group 的大小直接影响:
- 单次 All-to-All 的参与 GPU 数(越大越难做、latency 越高)
- 每卡专家数量与单专家容量
# 负载均衡
路由器如果不受约束,常见现象是:少数专家变成“热门专家”,大部分专家很少被选中
后果有两类:性能和训练质量都下降
# 1. 不均衡带来的直接问题
- 性能:
- 热门专家 token 爆炸 → 该专家所在 GPU 计算排队
- All-to-All 变得极不均匀 → 同步点被最慢 GPU 拖住
- 质量:
- 冷门专家几乎不训练 → 参数浪费
- 热门专家过拟合或成为“单点”,泛化变差
# 2. capacity factor:硬约束“每个专家最多收多少 token”
这是最常见、最工程化的手段:给每个 expert 设置容量上限
设:
- EP group 内总 token 数
- expert 数:
- Top-1 路由(每 token 选 1 expert)
理想均匀时每个 expert 平均收到
- capacity_factor > 1:给波动留余量(更少丢 token,但更多 padding/计算浪费)
- capacity_factor 越小:越省算/省通信,但更容易出现 overflow(超额) 导致 token 被丢弃或改道
# overflow 怎么处理?
常见策略:
- Drop(丢弃):超出容量的 token 不送该 expert(输出置零或走残差),简单但可能伤收敛。
- Second-choice / Top-2 改道:Top-1 满了就尝试 Top-2(或更多备选),更稳但通信与计算会增加。
- Padding 到容量:即使某 expert 实际 token 少,也 pad 到
做计算(实现更简单、吞吐稳定,但浪费 FLOPs)。
这就是你看到的“部分专家很忙/很闲”:capacity 既是“保性能”的保险丝,也是“算力利用率”的权衡点。
# 3. load balancing loss:软约束“鼓励均匀分配”
为了减少热门专家倾斜,训练时会加一个辅助损失(aux loss),鼓励“每个 expert 分到的概率”和“实际分到的 token 数”更均匀。
这里给出一种常见直觉写法:
- 令
表示 expert 实际接收的 token 比例 - 令
表示 router 分配到 expert 的平均概率(或门控权重的平均)
希望
- 让
尽量小(越小越均匀) - 或把
、 的乘积/相关性推向均匀(不同实现细节不同)
关键点:aux loss 是“拉均衡”的软约束,capacity 是“兜底”的硬上限,两者通常一起用
# MOE的优化
# Grouped GEMM
将"每个 expert 各做两次 GEMM(up/down)"合并为一次(或少数几次)分组矩阵乘调度,让 GPU 以更高的计算强度跑起来
工程实现里常配合:按 expert 分桶(bucketize)+ contiguous packing(把同一 expert 的 token 连续存放)
# Fused MoE kernel
很多框架会提供 fused MoE(或 fused dispatch+GEMM+combine,融合算子)以减少:
- 访存次数(读写 hidden states 太贵)
- kernel 启动次数
- 中间 buffer 占用