GPT-MoE中文67亿诗歌生成模型

概述

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模型和同等质量的稠密模型相比有将近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模型的参数量，模型结构，训练超参数如下表所示：

如下图所示，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

通过以上的图表分析，我们能大致判断出2.7B-Dense模型对应的高能效的稀疏模型是以1.3B的dense模型为底座，同时配置32或者64个专家的MoE模型。下面我们重点分析下1.3B+MoE-32/64和2.7B dense模型的性价比。在单机A100上预训练200个小时后，借助Tensorboard画出预训练验证集loss曲线。我们发现在验证集loss到达2.16时，1.3B+MoE-64模型的收敛速度是2.7B dense的1.17，1.3B+MoE-32的收敛速度落后于2.7B dense模型15%，如下图所示：

从下图中的基于Faster Transformer的单级单卡推理性能上看。1.3B+MoE-32和1.3B+MoE64的吞吐速度差不多且都高于2.6B dense模型，这是符合预期的，因为他们的底座大小都只有1.3B。

中文ZeroShot-NLU效果评测

中文文本生成效果评测

文本补全

诗歌生成

广告文案生成

作文生成

诗歌生成训练

下面是基于GPT-MoE中文0.35B+MoE64模型在诗词生成数据集上二次开发训练

from modelscope.metainfo import Trainers
from modelscope.msdatasets import MsDataset
from modelscope.trainers import build_trainer

if __name__ == '__main__':
    dataset_dict = MsDataset.load('chinese-poetry-collection')
    train_dataset = dataset_dict['train'].to_hf_dataset().rename_columns(
        {'text1': 'src_txt'})
    eval_dataset = dataset_dict['test'].to_hf_dataset().rename_columns(
        {'text1': 'src_txt'})
    max_epochs = 10
    tmp_dir = './gpt_moe_poetry'
    num_warmup_steps = 100

    def noam_lambda(current_step: int):
        current_step += 1
        return min(current_step**(-0.5),
                   current_step * num_warmup_steps**(-1.5))

    def cfg_modify_fn(cfg):
        cfg.train.lr_scheduler = {
            'type': 'LambdaLR',
            'lr_lambda': noam_lambda,
            'options': {
                'by_epoch': False
            }
        }
        cfg.train.optimizer = {'type': 'AdamW', 'lr': 3e-4}
        cfg.train.dataloader = {'batch_size_per_gpu': 1, 'workers_per_gpu': 1}
        return cfg

    kwargs = dict(model='PAI/nlp_gpt3_text-generation_0.35B_MoE-64',
                  train_dataset=train_dataset,
                  eval_dataset=eval_dataset,
                  max_epochs=max_epochs,
                  work_dir=tmp_dir,
                  cfg_modify_fn=cfg_modify_fn)

    # Construct trainer and train
    trainer = build_trainer(name=Trainers.gpt_moe_trainer, default_args=kwargs)

    trainer.train()