7.ST-MoE
2022年2月，Google发表了《ST-MOE: DESIGNING STABLE AND TRANSFERABLE SPARSE EXPERT MODELS》。ST-MoE可以说不仅仅是一个MoE工作，对于模型结构、工程实现、训练策略等都做了很多分析，可以说是MoE的必读论文。
ST-MoE最大模型包含269B总参数量，和与32B dense模型相当的激活计算量。论文中把模型称为称为Stable Transferable Mixture-of-Experts，或者ST-MoE-32B。
在MoE层的使用上，ST-MoE比Switch Transformer更“节省”一点，每四层才替换1个MoE层。
论文中主要训了两个规模的ST-MoE模型，分别有4B和269B的总参数量。ST-MoE以及其他用于对比的模型参数如下表
7.1.稳定性与效果分析
论文通过对乘性操作、噪音和裁剪这几个内容进行探索，来指导模型的设计。
论文首先研究了乘性操作对模型的训练稳定性和最终效果的影响。
之前已经有一些工作表明更多的乘法对模型效果有收益。
Some architectural improvements involve more multiplications than additions or do not sum many items at once
（1）GELU Gated Linear Units (GEGLU)
第一个例子是关于激活函数的。GLU是一个对两个输入向量进行component-wise相乘的操作，之后被扩展成GELU-Linear FFN变体，用于替换transformer中的ReLU FFN变体，其计算如下
这样在一些其他工作里已经被证明了对模型效果有提升。
（2）RMSNorm
第二个例子是RMSNorm中的缩放参数，也就是下面公式的 。
ST-MoE针对GEGLU和RMSNorm这两个乘性操作，做了实验，结果如下表。
发现移除乘性操作可以使模型稳定性更好（训练中发散的情况减少），但是最终效果变差了。
（3）增加dense层
ST-MoE还验证了在expert层增加更多dense层的效果。结果发现增加更多的乘法交互（增加dense层），可以在带来效果收益的同时，基本不影响推理速度，如下表所示。
（4）增加一个bias
在FFN层的第一个矩阵乘法后面增加一个可学习的bias B，分别通过加法和乘法加入
乘法的收敛速度更快，效果也更好。
上面这些实验显示，后续在模型效果的探索方向可以往多使用乘性操作去考虑。
接下来ST-MoE探索了“噪音可以提升模型稳定性”的假设。
通过input-jitter，给router的输入logits乘以一个在[1e-2, 1e2]之间的均匀随机变量来添加噪音。
结果是增加noise之后，有助于让模型的收敛更加稳定，但是对模型最终效果有负面影响。
这里论文还提到，小模型上的结果不一定能直接推广到更大的模型上，比如在小模型上稳定的配置，在大模型就可能就不稳定了。因此还是需要在大模型上也进行充分实验。
对activation和gradient进行限制是目前广泛应用的提升模型训练稳定性的手段。在反向传播过程中，通过裁剪梯度的范数来缓解梯度爆炸，就是一种常用的限制手段。
但是在ST-MoE训练269B的大规模模型时，发现裁剪会使得模型收敛的效果很差。
为了解决这个问题，ST-MoE在训练中引入了router z-loss，形式如下。
是token的数量， 是专家数， 是router的输入。
z-loss会对进入router的较大的logits值进行惩罚，以达到尽量减少进入指数函数的较大误差的目的。什么意思呢？后面来解释，先看下使用z-loss的效果。
ST-MoE认为，在模型训练过程中，由于精度不足或者其他问题，会产生很大的值，从而引入误差。而对梯度进行裁剪是在误差发生之后，并且裁剪本身也造成了数据的不连续性，某种程度上，裁剪本身也是一种误差。相反地，z-loss自然地鼓励模型产生较小的对数值，因此可以更精确地建模。
z-loss乘以一个权重超参 加入到模型训练的总损失中，如下式所示。
ST-MoE经过实验，选择了。
是 auxiliary load balance loss负载均衡损失，ST-MoE这里使用了和GShard/Switch Transformer所用的相同的损失计算，这里回顾一下：
是专家数， 是包含 个token的batch。 表示被分配到第 个expert的token数，这个不可导； 表示整个batch每个token分配给第 个expert的概率的总和，这个可导。
目前大部分的大模型训练都使用混合精度训练：模型权重以float32格式存储以进行梯度更新，然后在正向和反向传播的矩阵乘法中转换为bfloat16；此外，所有激活值都以bfloat16存储和操作，而allreduce通信可以在bfloat16或float32数值精度中进行。
对于ST-MoE-32B的训练，allreduce的数值使用半精度可以加速训练，然而这也会使训练变得不稳定，因此ST-MoE保持allreduce的数值精度为float32。
bfloat16和float32在不同范围的舍入误差如下表所示
可以看到，表达的数值越大，舍入误差越大。而z-loss限制了数值大小，也就将误差值限制在比较小的范围。
MoE模型天生对舍入误差敏感，因为它们由于router的使用而有更多的指数函数，而指数函数会将小的输入误差放大很多，这就加剧舍入误差所导致的训练不稳定。
另外，ST-MoE有一个策略：只有当排第二的专家的权重大于等于第一的专家的1/5时，token才会被路由到其第二位专家，否则第二个专家就会被忽略。
因此虽然舍入误差不会改变softmax运算中各个概率的排序，但它确实会影响MoE中第二个专家的激活。
7.2.模型设计
dense模型的设计有scaling law进行指导，但是MoE模型的设计比dense模型多出几个要考虑的点：
（1）使用多少个expert
（2）怎么routing
（3）专家容量系数怎么定
（4）硬件的影响
（这里提到MoE模型的scaling law工作：《Unified scaling laws for routed language models》，可以了解一下）
ST-MoE认为，从以往的经验来看，在总专家数量较少的情况下（如8/16/32），提升专家数量，能有收益。但是在特别稀疏的情况下（如激活专家数量<1%），或者总专家数较大（比如>256）之后，提升专家数量收益就很小了。
从另一个角度来看，如果一个计算核心使用>1个专家，那么就会出现比较大的加载参数张量的成本，因此建议每个计算核心使用<=1个专家。
论文做了一系列实验来探索capacity factor的选择，如下表所示
从这些实验中得到几个结论：
（1）训练和推理的capacity factor增大都会有收益
（2）如果硬件资源足够，推理的capacity facotr可以设得比训练的时候大，会有进一步提升
（3）激活的expert数量提升会有收益，但是收益随着capacity factor提升而越来越小
当然，选择capacity factor还要看硬件的特性，如果通讯很快，可以适当增大capacity factor，否则就不能选择太大的。
下表展示了不同capacity factor对推理速度的影响
7.3.实验
ST-MoE-32B在下游任务上和以往最佳结果对比如下表，ST-MoE-32B刷新了超过一半任务的最佳效果
论文还对各个专家的专业化进行了追踪，发现decoder中几乎没有专业化的迹象，各种类型的token近乎随机分配给不同的专家。而在encoder中则表现出了高度专业化的特征，如下表
此外，还发现在多语言的模型的encoder中，专业化的情况并不想原先预想那样，按不同语言划分，而是每个专家都会处理一种语言的一部分token，如下表
8.DeepseekMoE
2024年1月，幻方量化（旗下的深度求索）开源了DeepseekMoE，是国内首个开源的MoE大模型。幻方还发布了论文《DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models》，给出了一些DeepSeekMoE的细节内容，颇为实在了。
DeepSeekMoE在其他MoE工作的基础上，进一步给出了2个模型设计的主要思路：
（1）对expert的粒度进行细分，以提供更多样的expert激活组合；
（2）对expert的类型进行区分，从所有expert中保留一部分作为shared expert共享专家，这部分专家对所有输入都保持激活。
这样的做法可以帮助每个expert达到更高程度的专业化(specialization)的水平，更好地学习不同的专业知识。
DeepSeekMoE先在2B的较小MoE模型上进行了充分的实验，然后把方案应用到16B参数的MoE模型上，并获得了较好的效果。其中DeepSeekMoE-16B不需要量化就可以在40GB显存的设备上运行。
DeepSeekMoE-2B模型具有和稠密2B模型相当的性能，而DeepSeekMoE-16B则具有和7B稠密模型相当的性能，且计算量仅为稠密模型的40%。
DeepSeekMoE-16B的参数效率相比稠密模型有明显的优势，如下图所示
并且DeepSeekMoE-2B和16B模型都开源了。
在前面实验的基础上，幻方还训练了DeepSeekMoE-145B的超大MoE模型，具有和稠密的DeepSeek-67B模型相当的表现，但计算量更小。这个后续也有机会放出来。
8.1.模型设计
MoE，mixture of expert，顾名思义，一个最初始的motivation就是让不同expert学习不同的内容，然后再混合起来。
比如最上面提到的1991年的工作里，就是让不同的expert学习不同的元音特征，以此提升特征提取的准确率。
但是当前大部分的MoE架构都会遇到“knowledge hybridity”和“knowledge redundancy”的问题，即知识的杂糅和冗余：
（1）知识冗余
有些基础的常识在不同的领域都需要用到，每个expert就都会学一点，这样这些常识就被多个expert重复学习了。
（2）知识杂糅
在expert数量不够多的情况下，一个expert就可能要负责学习多个领域的内容。以学习高中知识为例，在只有两个expert的时候，只能一个expert学习理科知识，另一个学习文科知识；当我们有8个expert的时候，不同expert就可以分别学习语文、英语、历史、地理、物理、生物、化学、数学知识。显然后者所学知识的专业化程度更高。
知识的杂糅和冗余阻碍了专家专业化(expert specialization)的程度，也就阻碍了模型达到MoE结构理论上限性能。
我们期望每个expert能够学习到non-overlap & foucusd knowledge的知识。
针对上面的问题，DeepSeekMoE的架构设计有2个主要策略：
（1）Fine-Grained Expert Segmentation
参数总量不变的情况下，将expert分成更细的粒度（每个expert更小）。这样可以带来更灵活的激活组合，让每个expert可以有更强的specialization。比如原本是16个expert选择激活2个，那么总的组合数是120种；如果把每个expert缩小为原来的1/4，那在总参数量和激活数量不变的情况下，是64个expert选择激活8个，那么总的排列组合数就是 ，排列组合数比原来多了很多。
（2）Shared Expert Isolation
把部分expert分离出来，保持永远激活。我们期望这部分专家能够学到在多个领域间都通用的common knowledge。这样的策略同样可以使得其他expert能够提高专业化的程度，并且减少不同expert间的知识冗余。还是以学习高中知识为例，数学、物理和化学都需要算术能力，如果让学这三个领域的expert都学习算术技能，就会有冗余；我们可以把通用算术的技能剥离出来，由一个助手专门负责算术任务，相当于给他们发了一个计算器，这样学习数学、物理和化学的expert就能把更多的精力放在专业知识上，也就能达到更好的专业化效果。
下图展示了在传统MoE结构上增加Fine-Grained Expert Segmentation和Shared Expert Isolation策略的设计
（expert isolation的思路最早可以追溯到2022年1月发表的《DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale》，这里就不展开了。）
假设传统的MoE模型每层的expert数量为 ，激活expert数为 ，DeepSeekMoE使用的细粒度expert大小为原来的 ，那DeepSeekMoE每层就有 个expert，激活的expert数量为 个。假设 为输入长度， 为模型层数， 表示第 个expert，DeepSeekMoE可以公式化为以下表示（忽略了layernorm）
8.2.负载均衡
如之前工作反复提及的，如果任由MoE模型自主学习gating，可能会遇到两个问题
（1）routing collapse：专家分配的不均衡，也就是gating倾向于总是选择特定的少量expert，并且这种情况还会自我增强。
（2）计算效率问题：多设备间，不平衡的负载可能会成为计算效率的瓶颈。
针对routing collapse的问题，DeepSeekMoE引入一个expert-level balance loss，如下所示
叫做expert-level balance factor，是人工设定的超参。
而 和 和Switch Transformer里的设定基本一样。
在Switch Transformer里， 表示分配到第 个expert的token数量。在DeepSeekMoE这里也是一样的含义，只是多乘了一个系数 ，其中 ，， 是划分出来的共享expert的数量。这个系数是个常数，可以拿到求和符号外面，这样DeepSeekMoE里的 就和Switch Transformer里的完全一样了。
这个系数可以使得在使用不同的数量的expert时，在完美平均分配的情况下，负载均衡loss都是相同的常数。
表示所有每个token分配给第 个expert的权重的总和，和Switch Transformer里的含义一样。
注意这里 是不可导的， 是可导的。
针对多设备间负载均衡的问题，DeepSeekMoE引入一个device-level balance loss，如下所示
叫做device-level balance factor，是人工设定的超参。
指第 个设备。
device-level balance loss形式上和expert-level balance loss一样，只是 和 对应的对象从单个expert变成单个设备了。
当我们的目标是缓解计算瓶颈时，我们不需要强制执行expert间的均匀分配，而只需确保设备之间计算量的平衡。比如我们每层有64个expert，均匀分布在8个设备上，我们只需要每个设备处理的token数平衡即可，在设备内部即使所有token都是同一个expert处理的，依然能满足设备间负载平衡的要求。
相比expert间严格的负载平衡，只要求设备间的平衡是更松的限制条件，这样缓解了因为过度的负载平衡而损害模型性能的问题。
8.3.实验
为了验证以上策略的有效性，先拿100B token的语料数据在DeepSeekMoE-2B模型做实验。词表也是通过BPE在语料上训练的8k词表，后面训练更大规模模型的时候再扩大词表。
DeepSeekMoE-2B模型参数初始化方差为0.006，使用multi-head attention，前向激活参数量约0.3B，具体参数如下表
relative expert size指的是DeepSeekMoE所用的细粒度expert的大小和正常FFN层大小的比值。
训练的具体参数设置如下
| 属性 | 数值 |
|---|
| optimizer | AdamW |
| adam_beta_1 | 0.9 |
| adam_beta_2 | 0.95 |
| adam_weight_decay | 0.1 |
| warmup schedule | linear |
| warmup step | 2000 |
| max lr | 1.08e-3 |
| dropout | 0 |
| sequence length | 2k |
| batch size | 2k |
| total step | 25,000 |
其他训练细节：
- 所有expert放在单个GPU上，没有使用device-level balance loss
- expert-level balance factor设为0.01
- 训练到80%的时候，学习率乘以0.316，训练到90%的时候，再乘以0.316
使用相同的100B训练数据，训了deepseek moe 2b，在包含语言模型和下游任务的benchmark上和其他4个模型做对比：dense，hash layer（也是一种moe，《Hash layers for large sparse models》），switch transformer，gshard
可以得到几个结论：
- 更大的模型参数量和稀疏的架构，使得Hash Layer和Switch Transformer和具有同样激活参数的dense模型相比，有明显的优势。
- 同样的模型参数下，GSshard比Hash Layer和Switch Transformer有更多激活参数，效果也更好
- 同样的模型参数和激活参数下，DeepSeekMoE效果比GShard有明显优势。
为了进一步探索DeepSeekMoE架构带来的收益，提升了dense模型和GShard模型的激活参数，直到效果和DeepSeekMoE-2B差不多。
结果dense模型和GShard模型需要分别扩大到16倍和1.5倍的参数量，才能达到DeepSeekMoE-2B相近的效果，如下表所示
DeepSeekMoE的优势在更大规模的情况下，依然成立。训了DeepSeekMoE-13B, 对比参数量提升至1.2和1.5倍的GShard，DeepSeekMoE-13B依然能match，具体如下表
针对DeepSeekMoE架构的两个主要设计，shared expert和fine-grained expert进行消融实验。使用不同数量的共享专家和不同粒度的expert进行效果对比，结果如下图。
（1）对比蓝色和橙色，可以看到增加共享专家带来了收益
（2）绿色和红色在橙色的基础上进一步把专家颗粒分得更细，效果进一步提升
（3）共享专家和路由专家的比例：在总共64个expert的情况下，对比了1/2/4个共享专家的情况，结果并没有显著差别，在pile上的loss分别是1.808,1.806,1.811。最终选择了共享专家和激活路由专家1:3（2+6）的比例。
通过实验来验证DeepSeekMoE中expert specialization的优化。
（1）前面实验看到DeepSeekMoE-2B和1.5倍参数量的GShard模型效果相当。在这个基础上，通过禁用不同数量的top专家，而只能从次优的专家中选择进行回答。
实验结果如下
发现deepseek moe损失更大，说明每个专家的专门化程度更好，必要性更高，这说明deepseek的专家专门化程度更高。
（2）另外，通过禁用deepseek moe的共享专家，而额外激活一个专家，发现loss也大大提升。这个结果突出了共享专家的关键功能，并表明共享专家捕捉到了与路由专家不共享的基本且重要的知识，使得它无法被路由专家替代。
（3）只激活更少专家，也能和gshard达到相同水平，这一观察结果支持了DeepSeekMoE可以更准确和高效地获取所需知识的观点。
此外还从零训了一个只用1个共享专家和3个激活专家的2b模型（正常是2个共享专家+6个激活专家），也比gshard好，说明deepseek moe的有效参数效率更高
DeepSeekMoE-16B模型使用了2T数据训练（和LLAMA2-7B对齐）训练，并使用了100k的词表。其他参数如下表所示
论文中提到，除了第一层以外，其他层都使用了MoE层。
第一层不使用MoE是因为观察到第一层的负载均衡loss在训练中收敛得特别慢。
DeepSeekMoE-16B每层有64个专家，其中有2个作为共享专家保持永远激活，加上6个通过gating function选择激活的，每个token共使用8个专家。每个token会激活16.4B中的2.8B参数。
这里没有把专家的dimension再减小，是因为如果专家太小，计算效率就下降得太厉害。
训练中使用的其他设置：
- 训练进行到80%和90%的时候，lr都会缩小到0.316倍
- batch size = 4.5k，训练窗口长度是4k，因此每个batch有18M token，2T数据差不多是10.6w步
expert level balance loss的系数设得比较小，0.001，因为实验中发现设得再大并不能进一步优化负载平衡，反而会损害模型效果。
DeepSeekMoE-16B和DeepSeek-7B模型的对比如下
DeepSeekMoE-16B和LLAMA2-7B模型的对比如下
幻方还用245B的token训练了DeepSeekMoE-145B，模型效果上达到DeepSeek-67B的同等水平
9.DBRX
2024年3月27日，Databricks开源了DBRX，一个拥有有132B参数，激活参数为36B的MoE模型。
结构上，DBRX使用了RoPE、GLU、GQA，采用了fine-grained expert的设计，每层有16个专家，每个token激活其中4个。相比Mixtral和Grok-1在8个专家中激活2个，DBRX有更多的专家组合方式。
DBRX训练的上下文长度为32k，并使用了12T文本和代码token进行训练。DBRX在3072个H100上完成预训练，加上post-training、效果评估、red-team优化，整个过程耗费3个月时间。
DBRX整体效果超过GPT-3.5，与Gemini 1.0 Pro相当，并且具有比较强的代码能力，甚至超过了在代码上专门优化过的模型，如CodeLLaMA-70B，如下图所示。
推理效率效率上，DBRX也领先于其他模型。
10.Qwen1.5-MoE
2024年3月28日，阿里放出了Qwen1.5-MoE-A2.7B，以2.7B的模型参数，达到了Qwen1.5-7B模型的相近效果。
Qwen1.5-MoE-A2.7B参考了DeepSeekMoE和DBRX的工作，采用了fine-grained expert的做法，总共有64个专家，每个token激活8个专家，其中有4个为共享专家。
Qwen1.5-MoE-A2.7B使用Qwen-1.8B进行初始化，并在初始化阶段引入随机性，这样可以显著加快收敛速度，并得到更好的收敛结果。
Qwen1.5-MoE-A2.7B和其他模型效果对比如下
虽然Qwen1.5-MoE-A2.7B总参数量较大，但激活的non-embedding参数量远小于7B模型，如下表所示
实践中，Qwen1.5-MoE-A2.7B相比于Qwen1.5-7B，训练成本降低了75%。
推理性能上，在A100-80G用vLLM部署Qwen1.5-7B和Qwen1.5-MoE-A2.7B模型进行了性能测试。
输入/输出token数都设置为1000，输出token数设置为1000，TPS和throughput如下
虽然MoE模型对内存需求更大，但是由于稀疏激活以及共享专家的设计，但是在速度和吞吐量上都比dense模型更好。Qwen1.5-MoE-A2.7B与Qwen1.5-7B相比，速度提高了约1.74倍。
11.Mistral
11.1.Mistral 8x7B
2023年12月11日，Mistral AI开源Mistral-8x7B，每个token激活8个专家中的2个。
Mistral-8x7B支持32k推理窗口和多语言，并且代码能力较好。和LLAM2-70B以及GPT-3.5的对比如下。
Mistral-8x7B在大多数任务表现优于LLAM2-70B，且推理速度提高了6倍。
而和激活参数量相近的LLAM2-13B比，优势更为明显
11.2.Mistral 8x22B
2024年4月17日，Mistral AI开源Mistral-8x22B模型，一个总参数为141B，激活参数为39B的超大MoE模型。
Mistral-8x22B支持多语言，并且具有较强的数学和代码能力。此外，推理窗口长度也从Mistral-8x7B的32k增加到64k。Mistral-8x22B还具备function call的能力。
在各个维度的评测结果如下
12.小结
- 现有的工作都表明，MoE模型相比dense模型具有更高的参数效率，即同样的计算量下，MoE模型普遍能有更优的效果
- 因此MoE不仅能支持更大规模模型的训练，在较小规模模型上使用MoE架构也有很大收益
- 但是相比dense模型，MoE模型的训练也需要考虑更多内容，包括专家数量、激活数量和专家容量的设计，负载均衡的问题，如何在多设备上的并行等，训练难度更大
- 负载均衡上，GShard和Switch Transformer的负载均衡损失被广泛采用
- 推理时需要对底层框架进行优化以适配MoE机制，否则难以发挥MoE的性能优势
