大语言模型从Scaling Laws到MoE

发布于:2024-05-05 ⋅ 阅读:(34) ⋅ 点赞:(0)

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倍左右

参考

  1. 摩尔定律-维基百科 https://zh.wikipedia.org/zh-hans/%E6%91%A9%E5%B0%94%E5%AE%9A%E5%BE%8B
  2. Scaling Laws for Neural Language Models https://arxiv.org/abs/2001.08361
  3. Training Compute-Optimal Large Language Models https://arxiv.org/abs/2203.15556
  4. GLaM: Efficient Scaling of Language Models with Mixture-of-Experts https://arxiv.org/abs/2112.06905
  5. Mixture-of-Experts with Expert Choice Routing https://arxiv.org/abs/2202.09368
  6. Brainformers: Trading Simplicity for Efficiency https://arxiv.org/abs/2306.00008
  7. Lifelong Language Pretraining with Distribution-Specialized Experts https://arxiv.org/abs/2305.12281