MoE(Mixture of Experts)的相关问题

陈锦龙，12/14/2023

MoE的发明的动机是什么？

MoE提出自hinton 1991年的论文Adaptive mixtures of local experts，主要目的是加快模型训练的速度。因为正常训练过程中，如果使用同一个网络去拟合全部的数据，因为数据是多样性的，所以拟合的过程会比较慢。这时候如果有多个专家网络，每个专家网络只拟合一部分的数据，那么模型整体的学习速度以及泛化的能力都会增强。

第一次将MoE应用到transformer中的工作是Google 2021年的GShard，并且确定了最近几年MoE工作的主要动机：保持相同训练和推理资源的同时，通过增加模型的体积代价来提升模型学习效果。

为什么在1991年提出直到最近才重新进入视野？

1991年还处在BP算法刚刚提出来的阶段，最优的模型也就是多层感知机。当模型本身的容量较低的时候，在复杂的场景下，用一个网络去拟合所有的数据，会因为数据的多样性，所以拟合的过程会比较慢。所以MoE被提出用来增加模型在复杂场景下学习的效果，虽然在LLM时代，只要有足够的算力和数据，模型规模扩大一定能带来更好的效果。但在当时的算力稀缺并且缺少模型scaleup需要的技术时，这种方法还是可以提高参数利用率的。

而当LLM发展到GPT3的规模，推理和训练对应的优化方法也趋近于完善，继续scale up更多依赖于硬件的提升。那么当算力发展变缓或者获取成本变高的时候，就需要另外一种可以继续scale up但不那么依赖于硬件的方式，MoE开始进入人们的视野。

MoE为什么会起作用？

scaling law层面的解释

Scaling laws for neural language models揭示了模型规模、数据大小、以及算力大小之间的规律，并且建议对算力最佳利用方式是：在固定数据集大小下尽量训练更大规模的模型。这里的规模一般指的是模型参数数量以及需要的计算量，参数量增加同时计算量也会增加。

MoE相当于一次解耦，只增加模型参数数量、同时保持需求计算量相对恒定，所以效果提升符合scaling law的规律。

典型的MoE在transformer结构中应用如下（switch-transformer, Google, Jul 2022）：

模型结构层面的猜想

观察1：FFN可以理解为KV对，K代表了文本模式 、V代表了文本分布^[1]^，越靠后的层学习到越复杂的模式。

观察2：模型越靠后的层学习越不够充分。

基于这2个观察，可以做出一个假设：学习不充分的原因是容量不够，越靠后的层对容量需求越大。支撑这个假设的一个实验观察是：如果只有一个MoE层，放在越靠后的位置，最终效果越好^[3],[4]^。另外，从直觉上复杂的模式对应的数量是简单模式的指数倍，需要更多的参数去拟合。

所以MoE通过增加FFN的数量，增加了模型的容量，可以学习到更多的文本模式，所以得到更好的效果。

[1] Transformer Feed-Forward Layers Are Key-Value Memories [2] Deepspeed-MoE#PR-MoE [3] Hash Layers For Large Sparse Models#Figure 3

Expert真的是专家吗？

考虑下面3个MoE路由配置的实验：

将语义相似的token路由到相同的expert上；
将语义相似的token路由到不同的expert上；
随机将token指定到一个固定的expert上；

如果把expert理解成不同领域的专家，那么应该是1>2>3，但真实的实验结论是2>3>>1（23.22>23.27>>23.99）。

对此一个合理的解释是：experts并不是理解中的"领域专家" 分别学习不同的领域知识，而是增加对相似文本模式之间的区分度；相似的文本模式更可能发生在相似的token上面，所以相似的token应该路由到不同的expert上。

如何计算训练和推理成本？

结论1 ：MoE训练所需的显存与基底模型并没有明显差别；

以Mixtral-7B*8 MoE为例：

dim	n_layers	hidden_dim	n_heads	n_kv_heads	vocab_size	num_experts_per_tok	num_experts
4096	32	14336	32	8	32000	2	8

整体参数量：46b
活跃参数量：12b
对应基底模型参数量：7.2b

计算过程如下：

transformers:

FFNs: 176,160,768 * 8 = 45,097,156,608

gate: 4096 * 8 = 1,048,576

MA: 4096 * (128 * 48) + 4096 * 4096= 41,943,040

LN: 4096 * 2 = 8,192

total: (176,160,768 * 8 + 41,943,040 + 8,192 + 4096 * 8) * 32 = 46,440,644,608

others:

embed & output_w: 262,144,000

total:

46,440,644,608 + 262,144,000 = 46,702,788,608 = 46B

active params:

(176,160,768 * 2 + 41,943,040 + 8,192 + 4096 * 8) * 32 + 262,144,000 = 12,879,921,152 = 12.8B

模型训练的并行方式分为3种，DP(data parallel) / TP(tensor parallel) / PP(pipline parallel)，MoE模型在训练时可以同时使用这3种模型外，还可以加入EP(expert parallel)方式。EP的精髓就是多个模型中共享expert，2点理解：

因为MoE每次前向中只用到一小部分expert，如果每个模型保留完整的expert，一定会导致大多数expert空闲的情况；
如果DP是8，EP是2，那么2个模型共用一套完整的experts；

训练并行设定：TP2 DP8 EP8 （megatron方案），需要显存如下：

module	MoE参数/单卡	Dense参数/单卡
Emb and Output	<math xmlns="http://www.w3.org/1998/Math/MathML"> h ∗ v o c a b ∗ 2 / T P h * vocab *2 /TP </math>h∗vocab∗2/TP=262 144	<math xmlns="http://www.w3.org/1998/Math/MathML"> h ∗ v o c a b ∗ 2 / T P h * vocab *2 /TP </math>h∗vocab∗2/TP=262 144
experts	<math xmlns="http://www.w3.org/1998/Math/MathML"> h ∗ f f n ∗ 3 ∗ n u m _ e x p e r t s ∗ n _ l a y e r s / T P / E P hffn3 * num\_experts * n\_layers / TP / EP </math>h∗ffn∗3∗num_experts∗n_layers/TP/EP=88 080 384	<math xmlns="http://www.w3.org/1998/Math/MathML"> h ∗ f f n ∗ 3 ∗ n _ l a y e r s / T P hffn3 * n\_layers / TP </math>h∗ffn∗3∗n_layers/TP=88 080 384
gate	<math xmlns="http://www.w3.org/1998/Math/MathML"> h ∗ n u m _ e x p e r t s ∗ n _ l a y e r s hnum\_experts n\_layers </math>h∗num_experts∗n_layers=1 048 576	/
GQA	<math xmlns="http://www.w3.org/1998/Math/MathML"> h ∗ h d i m ∗ ( n h e a d + n _ k v _ h e a d s ) ∗ 2 ∗ n _ l a y e r s / T P hhdim(nhead+n\_kv\_heads) * 2 * n\_layers / TP </math>h∗hdim∗(nhead+n_kv_heads)∗2∗n_layers/TP=671 088 640	<math xmlns="http://www.w3.org/1998/Math/MathML"> h ∗ h d i m ∗ ( n h e a d + n _ k v _ h e a d s ) ∗ 2 ∗ n _ l a y e r s / T P hhdim(nhead+n\_kv\_heads) * 2 * n\_layers / TP </math>h∗hdim∗(nhead+n_kv_heads)∗2∗n_layers/TP=671 088 640
LN	<math xmlns="http://www.w3.org/1998/Math/MathML"> h ∗ 2 ∗ n _ l a y e r s h2 n\_layers </math>h∗2∗n_layers=262 144	<math xmlns="http://www.w3.org/1998/Math/MathML"> h ∗ 2 ∗ n _ l a y e r s h2 n\_layers </math>h∗2∗n_layers=262 144
total	3,622,043,648	3,620,995,072
推理所需显存/单卡	7,244,087,296	7,241,990,144
训练所需显存/单卡	57,952,698,368	57,935,921,152

通过加入EP的方式，在7B的模型大小下，MoE训练所需的显存于正常7B相差不大，结论成立。

注：实际计算中，MoE 的激活值会相比原有增大 EP 倍，和训练长度有关，以实际为准。

结论2 ：expert数量越多，训练时所需的最小设备数越多；

不超过7B的情况下，最小设备数的计算的逻辑如下：

确定了训练使用设备的GPU显存大小，以及基底模型规模；比如：80G，7B
确定可以使用的最小TP(tensor parallel)数；因为每张卡上可以训练的参数量差不多是3-4B，因此TP=2的时候满足要求，每张卡3.5B；
EP(expert parallel)使用最大，以节约内存； EP = num_experts
因为expert在DP中共享，所以DP必须是EP的整数倍； DP=EP
最终，最小设备数为：DP * TP。

在7B情况下，num_experts=8 最小设备数为16、num_experts=16 最小设备数为32。

在Mixtral-7Bx8 MoE的例子中，训练时可以跑起来的最小硬件要求：A100 80G x 8张 x 2台。

注1：超过7B时，必须要使用PP才能训练。

注2：模型较大的情况下，还需要考虑稳定性其他因素，最优的并行组合需要实测后才能得出；

训练成本

刨除显存以及设备数量这2点因素，训练成本以及速度计算如下：

差异主要来自推理过程中活跃参数的数量差异；
MoE 7b*8 的活跃参数是12b，所需算力相当于训练12b模型所需的算力；

推理成本

以Mixtral-7Bx8为例，整体参数量为46b，活跃参数是12b，但是在并行条件下推理速度接近于7.2b基底模型。因为虽然模型计算量翻倍了，但是模型容量主要是由宽度带来的，wide model更利于并行计算。

但是由于计算量增加，所以在固定设备数量的前提下，对外服务的最大吞吐量会减半。也就是说：

MoE模型会降低推理系统的最大吞吐量 ，但是时延变化不大；
吞吐量降低多少主要与活跃参数相关。

Mixtral-7Bx8MoE的收益和成本？

成本

成本主要取决于推理过程种的活跃参数，Mistral 7Bx8的MoE模型每次推理使用2个expert（top2），所以训练成本大概相当于7B模型的2倍，既训练一个12B模型所需的算力资源。

收益

Model	active params	MMLU	HellaS	WinoG	PIQA	Arc-e	Arc-c	NQ	TriQA	HumanE	MBPP	Math	GSM8K
LLaMA 2 7B	7b	44.4	77.1	69.5	77.9	68.7	43.2	17.5	56.6	11.6	26.1	3.9	16.0
LLaMA 2 13B	13b	55.6	80.7	72.9	80.8	75.2	48.8	16.7	64.0	18.9	58.4	6.0	34.3
LLaMA 2 33B	33b	56.8	83.7	76.2	82.2	79.6	54.4	24.1	68.5	25.0	40.9	8.4	44.1
LLaMA 2 70B	70B	69.9	85.4	80.4	82.6	79.9	56.5	25.4	73.0	29.3	49.8	13.8	69.6
Mixtral 7B	7B	62.5	81.0	81.0	82.2	80.5	54.9	23.2	62.5	26.2	50.2	12.7	50.0
Mistral 8x7B	12B	70.6	84.4	84.4	83.6	83.1	59.7	30.6	71.5	40.2	60.7	28.4	74.4

Mixtral 7B -> MoE 7Bx8，上涨15.6%；
LLaMA 7B -> 13B，上涨19.4%；

结论：考虑到分数越高越难提升，mixtral的MoE 7Bx8的效果收益可以认为与直接训练13B的模型相仿。

数据来自：Mixtral of experts | Mistral AI | Open source models

MoE应用场景与对应的收益成本

*基础要求：建议的模型方案，所消耗的算力与基底模型保持不变。

注：Switch-transformers, HashMoE, DeepspeedMoE等论文中的方案都可以满足该要求。

场景一：在模型规模扩大遇到瓶颈时，通过MoE继续提升模型效果，并且使用保持相同的算力资源。

收益：

10%以上的效果提升；类比mixtral 7bx8 MoE提升15%，GPT4比GPT3提升 20+%

成本：

没有明显的成本增加；

场景二：对于线上模型，可以替换为更小基底的MoE模型。

收益：

在相同的推理资源下，维持相同的最大吞吐量，并大大降低模型的时延；
小基底的MoE模型使用的训练资源更少。

成本：

没有明显的成本增加，除了开发的复杂度变高；

MoE有哪些重要的工作

时间	单位	模型名称	gating	sparsity	Arch
Jun 2020	Google	GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding 引用：543	learnable token-based	top2	MA>MoE>MA>FFN
Jun 2021	Google	Switch-transformers引用：1024	learnable token-based	top1	MA>MoE
Jul 2021	FAIR	Hash Layers For Large Sparse Models引用：95	determinative hash-based	*top1（hash决定）	MA>MoE
Apr 2022	Google Brain	ST-MoE: DesigningStable and Transferable Sparse Expert Models引用：28	learnable token-based	top2	[MA>FFN]*3 >MA>MoE
Aug 2022	Google	GLaM: Efficient Scaling of Language Models with Mixture-of-Experts引用：240	learnable token-based	top2	MA>MoE>MA>FFN
Oct 2022	Google	Mixture-of-Experts with Expert Choice Routing引用：56	learnable expert-based	top2	MA>MoE
Oct 2022	University of Texas	Residual Mixture of Experts引用：19	core + residual	Top2	MA>RMoE

MoE的scaling law：

来自 Deepmind, 2022, Unified Scaling Laws for Routed Language Models

计算公式如下：

需要注意的有以下几点：

log的基底为10；
默认MoE的位置是 MA->FFN->MA->MoE；

也就是隔一层放置一个，如果是每层放置，对应的expert数量要*2；
对应的参数有3组，分别对应了3种路由算法：S-BASE、RL-R、HASH；

其中Hash效果最差，与self-learning参数的路由最接近；

mega-blocks对应的方案是S-BASE的方案，效果最好；所以如果只是使用的话注意选择合适的参数；

一个scaling law的使用场景是计算有效参数，能够把MoE的模型参数对应到对应的Dense网络参数，计算方式如下：

实测，Mixtral 7bx8对应的有效参数大小是12.2B，与前面从效果和算力层面的推测一致。

难收敛如何解决

原因1：MoE刚开始训练时expert是随机初始化的，并且训练过程中expert的能力在不断变化，所以初期gating网络一直在拟合变化中的mapping关系，导致收敛时间边长；

解决方法1：分阶段训练；Evo-MoE

a. 首先训练一个正常的网络；b. 通过随机mask将expert分化成不同的expert；c. 加入gating网络，从topk逐渐降低到top1；

解决方法2：使用固定的映射作为route，hash-based routing。

这种方式可以提前生成token对应的expert映射，所以不存在波动的问题。

原因2：gating网络数值不稳定性；

gating网络在更新时，输入的logits会增大。因为softmax的导数与输入成正比，所以输入越大会导致梯度越大，容易产生不稳定。

解决方法：加入z-loss，限制输入gating网络的logits的大小。

其他原因：

乘性算子：RMS LN、GEGLU使用会增加不稳定性；
越大的模型越不稳定；

相关解决方法：

float32去计算gating 网络中的softmax；
使用更小的初始化参数；

expert负载均衡如何解决？

问题提出：Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

We have observed that the gating network tends to converge to a state where it always produces large weights for the same few experts；

如果不做任何控制，gating网络会倾向于给某几个expert更高的权重。

少量的负载不均衡不会影响模型效果，但是会影响推理和训练的速度。

解决方法：

附加损失auxiliary loss
hash-based route

微调效果不稳定如何解决？

问题原因：

MoE模型比Dense模型更容易在少量数据上过拟合；
MoE参数微调容易带来基础性能的丢失；

解决方法：

expert dropout调整到比正常参数高一些；正常0.1，expert 0.4
只更新非expert参数往往能够得到和全量更新同样的效果。

附录

Adaptive mixtures of local experts论文阅读

BP提出自1986年，这篇文章写自1991年。

主要想解决的问题是：当BP用来训练一个解决多个不同问题的网络时，因为任务之间的干涉导致模型难以收敛和泛化。虽然在LLM时代，只要有足够的算力和数据，模型规模扩大一定能带来更好的效果。但在当时的算力稀缺并且缺少模型scaleup需要的技术时，这种方法还是可以提高参数利用率的。

原来的公式：
<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"> E c = ∥ d c ⃗ − ∑ i p i c o i c ⃗ ∥ 2 \mathbb{E}^c = {\parallel \vec{d^c}-\sum_i{p_i^c \vec{o_i^c}} \parallel}^2 </math>Ec=∥dc −i∑picoic ∥2

其中， <math xmlns="http://www.w3.org/1998/Math/MathML"> c c </math>c 是case， <math xmlns="http://www.w3.org/1998/Math/MathML"> d d </math>d是目标输出， <math xmlns="http://www.w3.org/1998/Math/MathML"> p p </math>p 是expert对应的gate weight， <math xmlns="http://www.w3.org/1998/Math/MathML"> o o </math>o 是export对应的输出。

分析：如果其中1个expert的输出（ <math xmlns="http://www.w3.org/1998/Math/MathML"> o i o_i </math>oi）改变，导致expert所有输出的调和超过了最终的输出（ <math xmlns="http://www.w3.org/1998/Math/MathML"> d c ⃗ \vec{d^c} </math>dc ），那么所有expert梯度都会变为反方向。造成的结果就是，expert之间会互相干扰，导致收敛速度变慢，并且学习到的expert会倾向于合作的关系，共同作用于1个case并得到最终结果。

那么有没有方法expert之间更倾向于竞争关系呢。

这个问题有2篇相关的论文：

Learning piecewise control strategies inamo dular connectionist architecture

Task decomposition through competition in a modular connectionist architecture: The what and where vision tasks

==TODO==：一种方法是给目标函数添加惩罚项，来奖励expert之间更多的竞争。

如果把loss改成如下：
<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"> E c = ∑ i p i c ∥ d c ⃗ − o i c ⃗ ∥ 2 \mathbb{E}^c = \sum_i p_i^c{\parallel \vec{d^c}-\vec{o_i^c} \parallel}^2 </math>Ec=i∑pic∥dc −oic ∥2

最直观的变化是，每个expert的输出单独拿出来和目标进行比较，所以不管gate网络和其他网络的输出是什么，每个expert得到的梯度方向只取决于自身的输出和目标的差值。但如果其中某个expert输出对应的error变小，那么gate网络对应该expert的概率就会增加。这点可以通过梯度 <math xmlns="http://www.w3.org/1998/Math/MathML"> ∂ E c p i = ∥ d c ⃗ − o i c ⃗ ∥ 2 \frac{\partial \mathbb{E}^c}{p_i}={\parallel \vec{d^c}-\vec{o_i^c} \parallel}^2 </math>pi∂Ec=∥dc −oic ∥2很明显的看到， <math xmlns="http://www.w3.org/1998/Math/MathML"> p i p_i </math>pi梯度大小取决于error平方的大小，梯度越大下降的越快。

新的目标函数对应每个expert的梯度如下：
<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"> ∂ E c ∂ o i c ⃗ = − 2 p i c ( d c ⃗ − o i c ⃗ ) \frac{\partial \mathbb{E}^c}{\partial \vec{o_i^c}}=-2p_i^c(\vec{d^c}-\vec{o_i^c}) </math>∂oic ∂Ec=−2pic(dc −oic )

此时会出现一个问题，当初期网络随机初始化后，可以认为每个expert对应的权重 <math xmlns="http://www.w3.org/1998/Math/MathML"> p i c p_i^c </math>pic是相同的，那么每个expert的梯度取决于error的大小，所以拟合最好的expert得到的梯度最小。为了解决这个问题，提出了新的目标函数如下：
<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"> E c = − l o g ∑ i p i c e − 1 2 ∥ d c ⃗ − o i c ⃗ ∥ 2 \mathbb{E}^c=-log\sum_i{p_i^c e^{-\frac12\parallel \vec{d^c}-\vec{o_i^c} \parallel^2}} </math>Ec=−logi∑pice−21∥dc −oic ∥2

此时每个expert输出对应的梯度为：
<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"> ∂ E c ∂ o i c ⃗ = − [ p i c e − 1 2 ∥ d c ⃗ − o i c ⃗ ∥ 2 ∑ j p j c e − 1 2 ∥ d c ⃗ − o j c ⃗ ∥ 2 ] ( d c ⃗ − o j c ⃗ ) \frac{\partial \mathbb{E}^c}{\partial \vec{o_i^c}}=-\Big[ \frac{p_i^c e ^{-\frac12\parallel \vec{d^c}-\vec{o_i^c} \parallel^2}}{\sum_j{p_j^c e ^{-\frac12\parallel \vec{d^c}-\vec{o_j^c} \parallel^2}}} \Big](\vec{d^c}-\vec{o_j^c}) </math>∂oic ∂Ec=−[∑jpjce−21∥dc −ojc ∥2pice−21∥dc −oic ∥2](dc −ojc )

梯度的前部分确保了表现好的expert能够得到更快的训练。