基于单机最高能效270亿参数GPT模型的文本生成与理解

概述

GPT模型能较好的处理文本生成领域的各种任务，比如文本补全，自由问答，完形填空，写作文，写摘要，写小说，写诗歌等等。最近火爆全网的人工智能产品ChatGPT也是以GPT文本生成模型为底座。虽然GPT大模型作用在这些应用领域的效果很好，但是训练成本非常高。以OpenAI推出的1750亿的GPT-3为例，在1024张A100GPU上预估需要34天，一万亿参数的GPT-3在3072张A100显卡上也至少需要84天；微软/英伟达联合推出的5300亿的NLG模型，在2048张A100显卡上耗时了3个月的训练时间才能达到比较好的收敛效果。

针对GPT基础模型参数量大，训练&推理硬件资源消耗过高等问题，基于MoE的稀疏化训练是目前最具竞争力的降本增效途径。MoE的全称是Mixture of Experts，其中的Expert对应的是Transfomrer模型的MLP层，在训练的时候从多个MLP中选取一个MLP进行激活（如下图所示）。这意味着模型可以在不增加计算强度（FLOPS/Bytes）的情况下，通过增加MLP模块的数量来增加模型参数量级，进而提升模型在下游任务上的泛化性能。采用MoE后的稀疏Transformer模型和同等质量（验证集loss以及zeroshot nlu下游任务性能）的稠密模型相比有将近1.2倍的训练吞吐性能提升，1.3倍的推理吞吐性能提升。我们在稀疏架构总体设计的时候，选择让MoE跟纯Transformer Decoder架构的GPT进行有机结合。原因是MoE跟Decoder结合效果通常会好于跟Encoder的结合效果。具体来讲，Encoder是通过随机masking的方式学习语言模型，而这种被随机masked的token会让expert的路由选择出现不均衡。另一方面，考虑到Decoder类的GPT模型比Encoder类的Bert模型有更广泛使用场景，因此我们采用GPT+MoE的技术架构路线，探索单机最高能效的绿色低碳GPT大模型训练&推理软硬一体化适配技术在中文文本生成场景的落地可行性。

基于当前比较成熟的分布式MoE专家路由选择技术，采用Switch Transformer[2]中的top-1路由机制。每个Expert根据如下的softmax函数被赋予一个概率值，取概率最高（top-1）的那个Expert当作网络的FFN层。其中W_r是做路由选择时需要学习的参数。

GPT-MoE训练&推理能效分析

基础预训练模型训练&推理性能分析

任何一种稠密（Dense）的GPT模型，都有一种效果与之对应的训练&推理速度更快的稀疏（MoE）GPT模型。我们的目标是在受限硬件比如单机条件下找到这种GPT-MoE模型配置，然后通过对MoE算法进行改进来进一步提升它的训练能效。我们通过对比稠密&稀疏模型的训练&推理性能，来发现与稠密模型等价的高能效稀疏模型。

8种GPT模型的参数量，模型结构，训练超参数如下表所示：

GPT模型	参数量	Layers	Heads	hidden size	LR	Batch of Tokens
1.3B Dense	1.3B	24	32	2048	2e-4	1M
2.7B Dense	2.7B	32	32	2560	1.6e-4	1M
3.6B Dense	3.6B	30	32	3072	1.6e-4	1M
0.35B+MoE-64	6.7B	24	16	1024	3e-4	0.5M
1.3B+MoE-32	13B	24	32	2048	2e-4	1M
1.3B+MoE-64	27B	24	32	2048	1.6e-4	1M
2.7B+MoE-64	56B	32	32	2560	1.6e-4	1M
3.6B+MoE-64	75B	30	32	3072	1.6e-4	1M

如下图所示，1.3B+MoE32/64模型在相同的step下对比1.3B dense表现出更低的验证集loss，其中1.3B+MoE-64模型的loss甚至低于2.7B dense模型

5个模型中，0.35B+MoE-64的训练吞吐速度最快，是其他模型的2倍左右。其余四个模型中，吞吐速度较高的是1.3B dense和1.3B+MoE-32，1.3B+MoE-64和2.7B dense的速度相近。如下图所示：

推理吞吐速度方面，1.3B Dense的显存消耗最少，0.35B+MoE64的延迟最低。

input_len = 20

output_len = 128

batch_size = 1

模型	latency （ms）	memory （MB）	num of gpus
1.3B Dense	399.66	9476	1
2.7B Dense	753.37	17340	1
3.6B Dense	777.54	22558	1
0.35B+MoE64	356.22	15772	1
1.3B+MoE32	581.34	33294	1
1.3B+MoE64	586.18	57880