1、摩尔定律和伸缩法则
摩尔定律(Moore's law)是由英特尔(Intel)创始人之一戈登·摩尔提出的。其内容为:集成电路上可容纳的晶体管数目,约每隔两年便会增加一倍;而经常被引用的“18个月”,则是由英特尔首席执行官大卫·豪斯(David House)提出:预计18个月会将芯片的性能提高一倍(即更多的晶体管使其更快),是一种以倍数增长的观测。[1]
然而,由于受到晶体管的散热问题、内存带宽瓶劲等问题,摩尔定律目前已经走到了物理的极限,限制了我们做出更快的芯片。
另一方面,OpenAI提出了Scaling Laws[2],其中列举了影响模型性能最大的三个因素:计算量、数据集大小、模型参数量。也就是说,当其他因素不成为瓶颈时,计算量、数据集大小、模型参数量这3个因素中的单个因素指数增加时,loss会线性的下降。同时,DeepMind的研究也得出来和OpenAI类似的结论[3]。
Scaling Laws for Neural Language Models
根据Scaling Laws,模型越大、数据量越大、计算量越大,模型效果也越好。因此,目前很多LLM都是朝着更多参数、更多训练数据的方向进行scaling。然而,随着摩尔定律走到了尽头,LLM也不可能做到无限大。那么计算受到限制的时候,该如何进一步提升模型的性能呢?其中一种方法是MoE。
2、使用MoE进行LLM的scaling
MoE(the mixture of experts model)的思想是训练多个神经网络(也就是多个experts),每个神经网络 (expert) 被指定 (specialize) 应用于数据集的不同部分。对于不同来源的数据,有一个managing neural net来判断应该交给哪一个 expert 进行处理。
2.1 GLaM
2022年,Google发布了MoE的模型GLaM[4]。GLaM是一个 decoder only 模型,支持 in-context learning,一共有1.2T的参数量,其中有97B是激活的。
GLaM: Efficient Scaling of Language Models with Mixture-of-Experts
GLaM在训练和推理时所需的计算量,都远低于GPT3。同时,在zero-shot、one-shot 和 few shot 方面,GLaM的效果都优于GPT3。
2.2 Expert Choice Routing
GLaM 虽然效果不错,但是有负载不均衡问题。也就是说,会有一部分 expert 会经常被激活,而有一些 expert 很少被激活。
2022年,在 NeurIPS上,提出了新的 Expert Choice Routing 方法[5],来解决负载不均衡问题。该方法中,每个expert 会独立选择 top-k 的tokens 作为输入。每个token都可能会被不同的 expert 选择。一些比较重要的tokens会得到更多的计算资源,而不重要的 tokens 得到的计算资源会比较少。
Mixture-of-Experts with Expert Choice Routing
采用 Expert Choice Routing 的模型,相比于GLaM,在收敛速度方面可以提升2倍, 在 step time 上提速20%,并且完美解决了负载不均衡问题。8B/64E 的模型(有9.8B激活的参数),在 SuperGLUE 上效果超过了 T5-11B 的模型。
3、MoE进阶技术
3.1 Brainformers
上述MoE模型,在训练速度很慢,会成为进一步 scaling 的瓶颈。在 ICML 2023上,进一步提出了Brainformers[6]。
Brainformers: Trading Simplicity for Efficiency
为了模型计算更快,需要对矩阵乘法进行分解,从而减少计算量。同时,这些矩阵分解,必须不能损害模型的准确性。上图 (a) 中,是两种分解矩阵乘法的主要方法,分别是从横向分解(low-rank)和纵向分解(multi-expert) 。而在图 (b)中,可以对 low-rank 和 multi-expert 进行组合、堆叠,以实现更有趣且计算效率高的模型架构。
如果在 bottleneck 处,插入一个 mixure 层,模型看起来就非常像 transformers。如果在 bottleneck 处,插入一个 attention layer,模型看起来就非常像一个 multi-expert transformers。
作者对不同的参数,进行了搜索,从而找到了最优的网络模型结果。搜索的空间包括:不同的层类型(attn、moe、ffn)、隐层维度、MoE隐层维度、FFN隐层维度、attention 的 head 数、Gating Fuction、Capacity Factor、Activation Function等。
从搜索空间中,采样一组参数,构建一个 100M/32E 的模型,选择top-K的模型,然后进行scaling,如1B/64E 、8B/64E。
Brainformers: Trading Simplicity for Efficiency
相比于GLaM,可以在收敛速度方面,提速2倍;在 step time 上提速5倍。
3.2 Lifelong Language Pretraining
其他问题:如何进行增量训练,同时避免灾难性遗忘。
方法:引入新的数据分布时,引入新的 experts,同时冻结原有的权重,并且加入一个regularization loss 避免灾难性遗忘[7]。
Lifelong Language Pretraining with Distribution-Specialized Experts
相比于baseline,Lifelong-MoE 的灾难性遗忘被显著抑制,并且比 dense oracle 的效果还要好。
4、问答及彩蛋
- MoE model 的训练,有2种方案,一种是 train from scratch,一种是 dense to sparse
- GPT4是一个MoE model,且大概率是一个 train from scratch 的 MoE model
- 如果 finetune 后的 MoE 效果变差,可能是 finetune 没有训好,需要增加新的 experts
- Google 从 NVIDIA 买了 26,000 个 H100 (真有钱)
- H100 针对 transformers 训练做了专门的优化,主要是卡与卡之间的通讯优化,比 A100 提速 10倍左右
参考
- 摩尔定律-维基百科 https://zh.wikipedia.org/zh-hans/%E6%91%A9%E5%B0%94%E5%AE%9A%E5%BE%8B
- Scaling Laws for Neural Language Models https://arxiv.org/abs/2001.08361
- Training Compute-Optimal Large Language Models https://arxiv.org/abs/2203.15556
- GLaM: Efficient Scaling of Language Models with Mixture-of-Experts https://arxiv.org/abs/2112.06905
- Mixture-of-Experts with Expert Choice Routing https://arxiv.org/abs/2202.09368
- Brainformers: Trading Simplicity for Efficiency https://arxiv.org/abs/2306.00008
- Lifelong Language Pretraining with Distribution-Specialized Experts https://arxiv.org/abs/2305.12281