中文稀疏GPT大模型落地 — 通往低成本&高性能多任务通用自然语言理解的关键里程碑

作者：同润，临在

团队：阿里云机器学习平台PAI

1. 概述

GPT模型能较好的处理NLP各个应用领域的任务，比如文本分类，推理，对话，问答，完形填空，阅读理解，摘要，生成等等。百亿/千亿参数量级的GPT大模型作用在这些应用领域虽然效果很好，但是训练成本非常高。以OpenAI推出的1750亿的GPT-3为例，在1024张A100GPU上预估需要34天，一万亿参数的GPT-3在3072张A100显卡上也至少需要84天；微软/英伟达联合推出的5300亿的NLG模型，在2048张A100显卡上耗时了3个月的训练时间才能达到比较好的收敛效果。基于MoE的稀疏化训练是目前最具竞争力的降本途径。另一方面，GPT稀疏大模型在中文领域的效果还没有得到实际的验证，能否有效处理中文超多分类任务，口语化表达的中文推理或者情感分析任务，都是我们亟待探索的课题。综上所述，我们的工作是探索基于稀疏MoE结构的GPT模型在中文任务上的性能上限，同时把大模型训练成本压缩到极致，仅在单机A100硬件支持下就能将百亿参数大模型训练成熟。

继去年使用Transformer Encoder大模型取得中文小样本学习、英文预训练模型知识量度量双冠后，今年阿里云又在中文语言理解权威评测基准CLUE的中文零样本学习权威榜单ZeroCLUE上取得冠军。我们将大模型技术能力向前推进了一大步。基于MoE稀疏结构，仅用一台A100就把160亿参数量级的多任务通用GPT模型训练成熟。此次刷榜的模型是阿里云机器学习PAI团队推出的 GPT-MoE，这也是业界推出的首个稀疏的纯decoder架构模型在zeroclue上登顶。这是通往低成本且高性能多任务通用自然语言理解的重要里程碑。

MoE的全称是Mixture of Experts，其中的Expert对应的是Transfomrer模型的FFN层，在训练的时候从多个FFN中选取一部分FFN进行激活。这意味着模型可以在不增加计算强度（FLOPS/Bytes）的情况下，通过增加FFN的数量进而增加模型参数量级来提升模型在下游任务上的性能。采用MoE后的稀疏Transformer模型和同等质量的稠密模型相比有5倍的吞吐性能提升，如下图所示。我们在稀疏架构总体设计的时候，选择让MoE跟纯Transformer Decoder架构的GPT进行有机结合。原因是MoE跟Decoder结合效果通常会好于跟Encoder的结合效果。具体来讲，Encoder是通过随机masking的方式学习语言模型，而这种被随机masked的token会让expert的路由选择出现不均衡。另一方面，考虑到Decoder类的GPT模型比Encoder类的Bert模型有更广泛使用场景，因此我们采用GPT+MoE的技术架构路线。

2. 算法&工程技术解读

算法方面我们提出了一种更加高效稳定的稀疏路由均衡器，基于Prompt的领域话术再适应的零样本学习以及中文中文复杂任务定向优化。工程方面我们训练了工业级中文BPE分词器，微调过程中的工程优化，transformer训练加速工具rapidformer等，下面为大家一一详细解读。

打造工业级中文BPE分词器 & 高质量预训练数据集

中文每个字都有独立的意思，可以利用复杂多变的自由组合来表达多种含义，如果GPT的词表依旧采用Bert那种按字分词的策略，这无形中极大增加了GPT语言建模的难度。比如下面这个例子，“道”字可以衍生出意义完全不同的三个词：道德，道理，道歉。根据“道”字左边的上下文信息推测其右边的字难度还是比较大的。

我们需要构建一种性能鲁棒的中文分词工具，它能够从大规模语料中按照上下文统计特征提炼出中文单个汉字构成的大颗粒度词组。我们决定采用先进的BPE算法来尝试做到这点。我们使用的Sentencepiece工具在5亿条高质量中文语料上训练了一个星期产出了一个5w词表量级的中文BPE分词器。另外中文超大规模语料的选择也有讲究，去年我们使用PAI数据集就能把13亿的Bert模型训练成熟，今年我们尝试了更大的WuDao2.0数据集，大概是PAI数据集的十倍。原始WuDoa2.0数据集质量还是不太符合要求的，我们按照去年对PAI数据集的清洗过程，重新对原始的WuDao2.0数据集做了清洗，生成了一份相对高质量的WuDao2.0数据集。我们在数据集和分词器上做了一些消融实验，结论如下：

1. 精选后的wudao数据集比精选的PAI数据集（只有wudao量级的十分之一）在验证集loss上要小，同时在TNEWS和EPRSTMT等下游任务指标要高。
2. 5亿条句子训练出来的BPE分词器比1亿条句子训练出来的BPE分词器在验证集loss上有0.02的降低，同时在TNEWS和EPRSTMT下游任务指标有小幅提升。

强化型稀疏路由均衡器

MoE概念最早出现在2017年的ICLR[1]上。它的基本概念十分简单，在前向传播的时候动态地从多个FFN中随机选择一些执行运算，从而达到参数扩充的效果。后续一系列论文[4][5][6][7][8][9]也都是围绕设计稀疏路由均衡器展开。

真正将MoE带到工业级发扬光大的是谷歌公司的GShard[2]和Switch Transformer[3]。其采用top-1路由机制。每个Expert根据如下的softmax函数被赋予一个概率值，取概率最高（top-1）的那个Expert当作网络的FFN层。其中W_r是做路由选择时需要学习的参数。

我们作出的改进是对Expert路由选择也就是到token到FFN的映射模块进行强化。让该模块在做映射路径选择的时候更加保持平衡。我们通过大量的消融实验发现了一种性能更好更稳定的稀疏路由均衡器Large-Margin[10] Router。其在Transformer中的功能位置如下图所示：

采用Large-Margin Router后可以看出validation loss的下降十分明显，同时下游任务的提升的效果也提升不少。

刷榜采用的模型参数配置如下所示：

领域话术再适应的零样本学习

我们仅使用一台A100机型就能拉起160亿量级的稀疏GPT大模型的训练。这种低成本的训练环境，允许我们通过继续预训练或者微调来能进一步释放GPT大模型的能力，让其在各个NLP任务中发挥更大的价值。

继续稀疏预训练的再适应损失函数

训练GPT大模型使用的超大规模中文数据集WuDao2.0，其话术风格偏向正式的用语风格。而ZeroCLUE9个任务中有六个任务是口语化的风格，两个任务是是科技论文风格。这两种风格是大模型没有见过的，因此需要通过学习更多特定领域的陈述表达方式让大模型去适应新的话术风格。具体来讲，当基础大模型训练成熟后，我们对每个任务或者多个相似任务进行继续预训练，最后对每个任务执行无监督零样本学习。如下图所示：

我们使用CLUE官方提供的未标注的任务语料来学习风格话术。对loss函数设计是总的loss= 上下文 loss + 提示语 loss + 伪标签loss+MoEloss。示例如下：

其中上下文和提示语的loss使用GPT训练语言模型标准的负对数最大似然，伪标签loss使用Masked LM，MoE loss和预训练保持一致。同时通过相关系数调节各个loss之间的贡献比例。

寻找最优中文提示语及其答案

提示语和对应答案的设计对zeroshot的性能影响非常大。以EPRSTMT情感分析任务为例，我们发现使用“满意/不满意”比使用“负面/正面”效果好不少。提示语中加上评论两个字毫无作用，这跟我们的直觉貌似有些不一致。需采用更为通用的提示语：“这句话表达的是”。下面是我们通过实践总结出来的比较好的提示语和答案。

备注：CHID任务不需要提示语和答案

单机环境优化训练显存 & 超大体积大模型加载
我们在使用160亿参数大模型做继续预训练的时候面临两个非常大的工程挑战

• 超大容量ckpt导致单机无法加载
• 原始数据on the fly继续预训练显存消耗更大导致无法继续预训练

160亿稀疏GPT大模型需要占用2T的内存空间才能加载进来，经过探索尝试，我们最终用NVME挂载来扩充虚拟内存的方式解决了该问题，幸运的是我们使用的A100机型NVME空间很大。命令如下所示：

sudo mkfs -t ext4 /dev/nvme0n1
sudo mkfs -t ext4 /dev/nvme3n1
sudo mkdir /mnt /mnt3
sudo mount /dev/nvme0n1 /mnt
sudo mount /dev/nvme3n1 /mnt3
sudo dd if=/dev/zero of=/mnt/swapfile bs=1G count=1024
sudo dd if=/dev/zero of=/mnt3/swapfile bs=1G count=1024
sudo chmod 600 /mnt/swapfile
sudo chmod 600 /mnt3/swapfile
sudo swapoff -a
sudo mkswap /mnt/swapfile
sudo mkswap /mnt3/swapfile
sudo swapon /mnt/swapfile /mnt3/swapfile

更为致命的是我们发现160亿参数的大模型按照和基础预训练同样的训练加速&降显存配置加载进来显存就OOM了。OOM的主要原因是由于采用了on the fly的数据加载方式以及更为复杂的反向传播。得益于Rapidformer加速库良好的用户接口设计，我们发现打开zero-stage-2开关同时限制number workers个数可以执行继续预训练。这样做存在一个负效应就是基础预训练和继续预训练的优化器存在不一致，不过影响在可控范围内。

攻坚中文超多分类任务&复杂推理任务&关键字识别任务

中文超多分类任务

榜单中有二个很困难的超多分类任务，是英文任务中从来没有遇到过的。众所周知，GPT模型擅长做生成类的任务，BERT模型擅长做识别类的任务。IFLYTEK和CSLDCP是识别类的任务且类别超多，类别数过多导致likelihood-ranking的性能下降，直接使用ZeroShot的精度很低。

为了提升超多分类问题的精度，我们首先对这些叶子类目进行归纳整理，抽象出一级类目出来，比如：

接着在执行loglikehood-ranking的时候先在一级类目范围内进行粗排，接着在叶子泪目范围内进行精排。采用这样一级泪目辅助的精细化排序后，这三个多分类任务平均有4个点左右的提升。

复杂推理任务

这两个任务的难点在于口语化，句子短。即使进行领域话术再适应，也很难在几个epoch内让160亿大模型的风格迅速适配过来。

我们将CMNLI/OCNLI/BUSTM的无标签数据混合在一起进行联合领域自适应，采用同样的Prompt设计，进行联合推理。执行多任务数据混合后在OCNLI和BUSTM的平均分有5个点的提升。
关键字识别任务

起初我们认为这个任务不行是prompt的设计有问题，后来通过大量的实验我们发现使用log-likelihood-ranking推理准则没有使用perplexity推理准则有用，部署perplexity-ranking后有将近7个点的提升。

大模型训练加速工具Rapidformer

Rapidformer的定位是为Transformer模型库提供大模型大数据量的训练加速能力。这是通过有机整合微软的DeepSpeed，英伟达的Megatron以及Meta的FairScale来做到的。如下图所示：

下图对比了Rapidformer，Megatron，DeepSpeed三者之间的加速能力差异，Rapidformer融合了Megatron和DeepSpeed的各自的优点，提供了一个统一的训练环境。

基于Rapidformer，我们可以通过加速开关的组合，来找到最优的吞吐。下面我们用130亿参数的GPT-MoE大模型为例，介绍下如何寻找最优的吞吐加速策略组合。130亿稀疏模型的参数配置如下：

我们用到的训练加速核心技术包括：

激活检查点（Activation Checkpoint） 在神经网络中间设置若干个检查点(checkpoint)，检查点以外的中间结果全部舍弃，反向传播求导数的时间，需要某个中间结果就从最近的检查点开始计算，这样既节省了显存，又避免了从头计算的繁琐过程。

梯度累积 (Gradient Accumulation)  以batch_size=16为例，可以每次算16个样本的平均梯度，然后缓存累加起来，算够了4次之后，然后把总梯度除以4，然后才执行参数更新，这个效果等价于batch_size=64。这是一种有效的增加Batch Size的方法。通过该策略可以将每个step的batch size扩充到很大，结合LAMB优化器会提升收敛速度。

混合精度训练（Mixed Precision Training） 采用混合精度训练的好处主要有以下两点：1. 减少显存占用，由于FP16的内存占用只有FP32的一半，自然地就可以帮助训练过程节省一半的显存空间。2. 加快训练和推断的计算，FP16除了能节约内存，还能同时节省模型的训练时间。具体原理如下图所示，核心是在反向传播参数更新的时候需要维护一个FP32的备份来避免舍入误差，另外会通过Loss Scaling来缓解溢出错误。

Zero显存优化器 ZeRO（The Zero Redundancy Optimizer）是一种用于大规模分布式深度学习的新型内存优化技术。ZeRO具有三个主要的优化阶段分别对应于优化器状态，梯度和参数的划分。

• 优化器状态分区(Pos) ：减少了4倍的内存，通信容量与数据并行性相同
• 增加梯度分区(Pos+g) ： 8x内存减少，通信容量与数据并行性相同
• 增加参数分区(Pos+g+p) ：内存减少与数据并行度和复杂度成线性关系。

所有吞吐实验都采用sequence length=2048，Global Batch Size=256，通过下表我们发现当在local DDP环境下，打开Activation Checkpoint，关闭Zero-stage-2，8步梯度累积每步消耗的时间是最少的，吞吐也是最高的。同时显存也得到了充分的利用。

3. 使用示例

本次刷榜使用的160亿参数的中文GPT-MOE稀疏大模型及其ZeroShot应用已在EasyNLP（https://github.com/alibaba/EasyNLP）中开源，详情请见链接。具体使用方式如下所示：

首先通过如下的命令拉起EasyNLP专属镜像并创建容器，里面已经安装了Rapidformer及其全部的依赖库。

docker pull pai-image-manage-registry.cn-shanghai.cr.aliyuncs.com/pai/easy_nlp:0.0.7

然后下载中文版的GPT-MOE模型，注意需要在硬盘上预留2T的存储空间同时需要通过swap设置2T的内存空间。

./ossutil64 cp -r oss://atp-modelzoo-sh/tutorial/rapidformer/zeroclue/ckpts ckpts

接下来就可以开始针对每个任务执行继续预训练了，在run_finetune_gpt_moe.sh脚本中已经通过开关打开了必要的训练降显存加速技术比如：Activation Checkpoint和Zero-stage-2。从CLUE官网下载无标签的训练数据，然后打上伪标签，注入进大模型中开始进行任务话术领域再适应。

sh run_finetune_gpt_moe.sh

训练完成后保存ckpt就可以执行针对特定任务的Zeroshot推理了，使用如下的推理脚本。

sh run_predict_gpt_moe.sh

4. 总结

在这次中文ZeroCLUE刷榜评测中，我们围绕中文百亿稀疏GPT大模型落地挖掘了以下核心技术：

• 基于large-margin的稀疏路由均衡器比传统Top-1均衡器效果更好更稳定。
• 针对160亿参数稀疏大模型的领域话术再适应算法以及工程底座经刷榜检验是可靠的。
• 提升了中文超多分类问题和复杂推理等问题的零样本学习任务效果。
• 大模型训练加速工具的可靠性进一步得到了验证。

后续在EasyNLP中还会陆续放出更多高质量的大模型，敬请期待。

参考文件

[1] . Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

[2]. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding

[3]. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity

[4]. BASE Layers: Simplifying Training of Large, Sparse Models

[5]. Hash Layers For Large Sparse Models

[6]. TAMING SPARSELY ACTIVATED TRANSFORMER WITH STOCHASTIC EXPERTS

[7]. GLaM: Efficient Scaling of Language Models with Mixture-of-Experts

[8]. Unified Scaling Laws for Routed Language Models

[9]. Designing Effective Sparse Expert Models

[10]. Large Margin Deep Networks for Classification