概述
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 |