MoE介绍
MoE(Mixture of Experts)和标准Dense模型的整体区别如下:
- 同参数下与Dense模型相比,预训练速度更快、推理速度更快。
- 需要更高的内存来装载MoE模型。
- 更大的训练挑战。
和标准Transformer模型的结构区别如下如所示[5],主要是将FFN层进行了更改,添加一个gate层,并将FFN变成n_expert个,由gate来对每一个token进行路由,选择不同的expert,最终再将这些expert的输出使用某种方式(大多数为相加)合并成output。
直觉上,由于使用了多个expert来对不同的token进行操作,使得每一个expert可以更加专注于某一项任务。类似于多领域推荐任务,每一个领域都有一个expert,由一个前置路由(或者请求来源)来决定具体执行的expert,并进行汇总返回。
但是这种结构的训练相对较难,由于使用门控来控制token进入的expert,因此直觉上需要保证每一个expert都会得到较好的学习,否则模型就会坍缩到一个expert上,变成了一个Dense模型,失去了MoE结构的意义。因此有些工作会在gate上添加z-loss来保证选择的expert的logit相对平衡。
而对于专家的数量和选择也是一个挑战,可能存在知识混合和知识冗余等问题,这在很大程度上限制了专家的专业性,影响了模型的整体性能。
MoE训练
现在主要的涉及到MoE的资料主要有Mixtral[4]、DeepSeekMoE[3]和上海人工智能研究院的Llama-MoE[2]。他们采用了不同的训练和初始化策略。
from scratch训练
DeepSeekMoE采用了从头训练,进行了多尺寸实验。大部分认知和结构验证都是在2B(等效0.3B)上完成,随后扩展到了16B(等效7B)、145B(等效22B)。
DeepSeekMoE采用了更加细致的expert划分,并使用了一个共享专家。主要动机是:
- 需要一个共享专家来提供综合性的知识。
- 使用更多的expert来将知识进行分解,同时提供更多样性的专家组合方式,更加有针对性的对知识进行获取。
模型整体性能在2B规模时,达到了Dense模型相当的性能。在16B规模下,使用40%的计算量达到了Llama2-7B的性能。在145B的规模性,使用28.5%的计算来那个达到了DeepSeek-67B的性能。
训练细节
对于2B的MoE模型,使用了9层Transformer,1280的hidden size,10个注意力头,维度为128。
参数初始化标准差为0.006,64个expert(1+63),每一个expert的大小为1/4个FFN大小。也就是整个MoE层为16*FFN。
AdamW+0.9&0.95。weight decay 0.1,batch size token 4M。学习率调度没有使用cosine及其变种。而是:
- warmup 2000个step到1.08e-3。
- 在80%step时,学习率缩小到0.316。
- 在90%step时,学习率继续缩小0.316(0.316^^2)。
论文中并没有交代为什么是1.08这种精确到-2级别的学习率以及0.316这种数,猜测是做了一些相关的scaling law实验。
通过对几个结构的验证,在小模型上验证了当前结构的有效性。性能超过同尺寸的Dense模型。
即使是和同参数量的Dense模型相比,该结构也具备一定的优势,说明了共享专家的有效性。
在16B的试验下,仅用40%的计算量就达到了DeepSeek-7B的性能,同时MoE结构在知识密集型任务上的表现更加出色。
从现有模型扩展
Mixtral和Llama-MoE均使用了从现有的模型进行初始化随后微调的策略。但是他们却选择了不同的FFN初始化方向。
向上扩展
Mixtral使用了向上扩展的初始化策略,既用mistral-7B”扩展”出一个47B等效与12B模型的MoE模型。
Mixtral(2/8)选择复用Mistral,并将FFN复制8份,其他参数完全复用。并在32k的上下文长度上进行了微调。整体性能基本接近Llama2-70B和GPT-3.5。整体结构类似于GShard。
由于Mis/xtral使用的滑动窗口对长上下文具有独特的优势(因此Qwen2也是用这一结构),Mixtral还对长文本能力进行了测试,在秘钥检索这个任务重,获得了100%的准确率。
向下扩展
Llama-MoE使用了向下扩展的初始化策略,既用Llama2-7B“缩小”出一个7B等效3~3.5B的MoE模型。
Llama-MoE-v1的性能显著低于Llama2-7B,因为采用向下扩展,因此整个模型的激活参数显著减少(根据不同的expert配置,在3~3.5B间)。不过Llama-MoE-v1既然在同激活量的Dense模型上获得了较为出色的表现。
训练细节
在Llama2-7B的基础上,使用112块A100 80G,global batch token 15M。
学习率使用cosine调度,从2e-4 ~ 2e-5。总计训练200B token(13.6k step)。
测评如下:
整体看,多expert似乎是一个正确的路线(和DeepSeekMoE一致)
认知
这几篇文章提到了一些涉及到expert和数据上的一些认知。不过比较粗,可以当做一个坑留下,后面填。
expert认知
expert初始化
对于expert的初始化,对于Mixtral和DeepSeekMoE相对简单。对于Llama-MoE这类向下拆分的结构,会存在如何拆分expert的问题。均等拆分?聚类拆分?
论文对拆分方式做了验证。
- Independent Random: 这种方法涉及随机地将原始FFNs的参数分割成多个专家网络。每个专家网络都是独立且不重叠的,这意味着它们不共享任何中间神经元。这种方法简单直接,但可能不会考虑到专家网络之间的潜在结构和相互关系。
- Independent Clustering: 这种方法基于聚类来构建专家网络。它首先对FFN的权重进行聚类,然后将聚类结果用作专家网络的划分依据。这种方法试图找到权重之间的自然分组,以便每个专家网络都能专注于处理特定的输入特征。
- Sharing Inter: 这种方法涉及到在专家网络之间共享中间神经元。具体来说,不是将所有中间神经元随机分配给不同的专家网络,而是将一些神经元设置为共享的,而其他神经元则根据重要性分配给各个专家。这种方法试图在保持模型结构的同时,通过共享来减少模型的复杂性。
- Sharing Inner: 这种方法类似于Sharing Inter,但它在专家网络之间共享更多的中间神经元。这种方法可能会保留更多的模型结构,因为它允许专家网络有更多的重叠,这可能有助于模型更好地捕捉输入数据的复杂性。
Independent可能不会充分利用共享神经元的优势,而Sharing可能会导致专家网络之间的功能重叠,但也可能有助于捕捉输入数据的复杂性。
整体看,独立随机拆分的效果最好,不过这个效果好需要再15B token之后才能显现出来,聚类的效果一开始很好,长期效果不太行。猜测是降低了expert的表达性。
expert区分
几个模型也都验证了不同的expert是否有效,究竟在捕捉什么信息?
Mixtral对几个expert的选择进行了分析,通过对同一数据集的连续2个token的expert进行分析,发现浅层网络基本接近于随机分配(1/8=12.5%,1-6/8 * 5/7=46%),但是随着层数的增高,他们选择专门的expert的概率显著提升。
Llama-MoE也对expert进行了定向分析,发现不同的任务/数据类型选择的expert有明显的区分。
Mixtral还做了更详细的token级别分布。整体看,一些模版字符、符号、数字等 更容易被分配到一个expert中,说明了expert的有效性。也说明了expert学习到底并不像多expert的推荐系统在“领域”上学习,而是个更细致的语法级别。
数据/组合认识
Llama-MoE还实验了一些数据配比和清洗策略。
Sheared-Llama的配比似乎更优秀一点。考虑到这个MoE是基于Llama2-7B的,因此对于Llama2的后续工作似乎有指导意义。对于其他模型的后续工作,意义尚不明确。
同时还对数据集中的广告和不流利文本进行了过滤,结论是不流利文本可以显著提升性能,但是广告文本性能有所下降,猜测是降低了数据多样性。
实现
Mixtral是使用megablocks来实现的MoE,Llama-MoE基于transformers进行了开发,在github进行了开源。
不过需要更新到Deepspeed 0.13.1以上,并且可能还有一些其他问题(2024.01.26),目前对于x/y架构的MoE模型可能训练还有一些问题(x!=y),当x=y时应该没有问题。主要原因是Deepspeed在数据并行时,只对一个进程进行监控,当该进程没有使用某些expert时,其他使用这些expert的进程的更新可能会产生错误。
https://github.com/microsoft/DeepSpeed/pull/5008
https://github.com/microsoft/DeepSpeed/pull/4966
对于是否需要对初始化的gate进行冻结,询问了Llama-MoE的人,结论是不需要。他们确实做过相关冻结实验,不过在大概20B参数的时候两者loss就没有明显区别(大概2.1,区别全参数微调),因此他们的训练并未使用任何“gate相关的魔法”。
另外使用类似于Mixtral方式初始化的模型,初始阶段梯度过大,不过观察到了明显的下降。考虑到大尺寸的梯度确实会大一点,同时Deepspeed的一些问题以及参数融合等代码可能存在的bug,可能需要后续观察。
引用
1.https://huggingface.co/blog/moe
2.Llama-Moe https://github.com/pjlab-sys4nlp/llama-moe/blob/main/docs/LLaMA_MoE.pdf
3.DeepSeekMoE http://arxiv.org/abs/2401.06066
4.Mixtral http://arxiv.org/abs/2401.04088
5.Switch Transformers https://arxiv.org/abs/2101.03961
0 条评论