基于单机最高能效 270 亿参数 GPT 模型的文本生成与理解
作者:李鹏,王玮,陈嘉乐,黄松芳,黄俊
单位:阿里云智能机器学习平台 PAI & 达摩院自然语言基础技术
概述
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 模型的参数量,模型结构,训练超参数如下表所示:
如下图所示,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 效果评测
中文文本生成效果评测
文本补全
诗歌生成
在线体验地址:https://www.modelscope.cn/models/PAI/nlp_gpt3_text-generation_0.35B_MoE-64/summary
广告文案生成
在线体验地址:https://www.modelscope.cn/models/PAI/nlp_gpt3_text-generation_1.3B_MoE-32/summary
作文生成
在线体验地址:https://www.modelscope.cn/models/PAI/nlp_gpt3_text-generation_1.3B_MoE-64/summary
自研 GPT-MoE 算法创新 &实验分析
背景
Top-1 Gating 是目前最主流也最有效的 Routing 算法,但是也有着明显的缺点。例如,在 Top-1 Gating 中,每一个 Token 仅会被交给一个 expert 处理,因此,时常会出现某些 expert 需要处理很多 token,而有些 expert 仅需处理极少数量的 token 的情况,这导致处理极少 token 的 expert 无法获得足够多的信息,无法得到充分的利用。
高能效专家路由选择
因此,我们自研了新的路由算法,如下图所示。我们的算法让 expert 主动选择固定数量(capacity)的 token,同一个 token 可以同时被不同的 expert 处理,从而使每一个 expert 都能得到充分的训练。
最后生成的 token 的表示,则采用 weighted sum 的方式,将来自不同 expert 生成的表示加权求和并经过 Expert Residual 模块,以获得最终的 token 表示。这样的表示,由于有多个 expert 的共同作用,因此更加鲁棒。
计算 expert 对 token 的偏好:
其中是输入的 tokens 的表示,是 experts 的权重,表示输入 token 的数量, 表示 expert 的数量, 表示隐藏特征的维度, 为每一个 expert 对输入的每一个 token 的偏好程度;
2. L-Softmax 算法训练 expert 权重
训练过程中利用 L-Softmax loss 优化 experts 的权重 ,实现每一个 expert 对 token 都具有可区分的的偏好。
3. 每个 expert 选取固定数量的 token:
其中, 就是我们预先确定的,每个 expert 可以处理的最多 token 的数量,记录了每个 expert 需要处理的 tokens 的索引;
4. 计算最后的输出:
每一个 expert 根据索引计算出对应 tokens 的表示,并将不同 expert 对同一 token 生成的表示加权求和,最后通过 Expert Residual 模块生成最终的输出 。
实验分析
下图是采用自研算法与 Top-1 Gating、s-BASE 算法的验证集 loss 随训练 step 变化的曲线图,我们可以发现,采用自研算法的验证集 loss 始终低于 top-1 gating 和 s-BASE 算法的验证集 loss,证明了我们自研算法的有效性。
同时,我们观察到,验证集 loss 首次低于 2.7 的时间,自研算法的速度是 s-BASE 的 1.48 倍,极大的减少了模型训练的开销。
此外,我们也分析了自研算法和 Top-1 Gating、s-BASE 的训练吞吐,如下图所示,采用自研算法相比于 s-BASE 的训练吞吐提升了 1.17 倍。
基于 PAI DLC 的 GPT-MoE 预训练
Rapidformer 为 EasyNLP 中各种 Transformer 模型提供训练加速能力,这是通过有机整合微软的 DeepSpeed,英伟达的 Megatron 来做到的,如下图所示:
在 GPT-MoE 大模型的预训练中,我们用到的主要训练加速核心技术包括:
混合精度训练(Mixed Precision Training)采用混合精度训练的好处主要有以下两点:1. 减少显存占用,由于 FP16 的内存占用只有 FP32 的一半,自然地就可以帮助训练过程节省一半的显存空间。2. 加快训练和推断的计算,FP16 除了能节约内存,还能同时节省模型的训练时间。具体原理如下图所示,核心是在反向传播参数更新的时候需要维护一个 FP32 的备份来避免舍入误差,另外会通过 Loss Scaling 来缓解溢出错误。
选择性激活重算(Selective Activation Recomputation)在神经网络中间设置若干个检查点(checkpoint),检查点以外的中间结果全部舍弃,反向传播求导数的时间,需要某个中间结果就从最近的检查点开始计算,这样既节省了显存,又避免了从头计算的繁琐过程。实际使用时,有些 layers 产生的激活值大但是计算量小,需要选择性的过滤掉这一部分的激活值,保留重要的激活值,以节省重计算量。
Zero 优化器状态切分 (The Zero Redundancy Optimizer)是一种用于大规模分布式深度学习的新型内存优化技术。ZeRO 具有三个主要的优化阶段分别对应于优化器状态,梯度和参数的划分。我们这里使用的是 Zero-1 优化器状态分区。
序列并行 (Sequence Parallelism)是一种对长序列进行切分从而加速训练的技术,在 Tensor Parallelism 的基础上,将 Transformer 核的 LayerNorm 以及 Dropout 层的输入按 Sequence Length 维度进行了切分,使得各个设备上面只需要做一部分的 Dropout 和 LayerNorm 即可。这样做的好处有两个:1. LayerNorm 和 Dropout 的计算被平摊到了各个设备上,减少了计算资源的浪费;2. LayerNorm 和 Dropout 所产生的激活值也被平摊到了各个设备上,进一步降低了显存开销。
接下来我们通过PAI-DLC产品来演示如何执行 GPT-MoE 的基础预训练。
环境准备
首先通过阿里云产品机器学习平台 PAI 页进入容器训练 DLC 创建新的训练任务,如下图所示。
有三个关键的地方需要配置,分别是专属镜像,数据集以及执行命令。节点镜像地址配置为:pai-image-manage-registry.cn-shanghai.cr.aliyuncs.com/pai/pai-rf-image-poc-moe:0.2。执行命令配置为:
其中 dlc_run_pretrain_megatron_gpt.sh 可以从 EasyNLP 的 github 中获取:https://github.com/alibaba/EasyNLP/blob/master/examples/rapidformer/gpt_moe/run_pretrain_megatron_gpt.sh 传递给 dlc_run_pretrain_megatron_gpt.sh 的主要有以下 11 个参数,分别是:
资源 &数据准备
由于预训练数据集体积比较大,采用下载的方式不是很方便。借助于 DLC 提供的 OSS 路径的挂载功能,可以很方便的直接使用存储在 OSS 上的大规模预训练数据集。
任务资源配置建议使用单机八卡 A100 来训练 1.3B 以上的 GPT-MoE 模型。
任务创建
任务创建完成后,进入任务监控页面观察任务执行状态,如下所示:
点击日志查看运行状态,如下所示:
基于 PAI DSW 的 GPT-MoE 微调诗歌生成
使用 DLC 完成 GPT-MoE 的基础预训练后,还需要对模型进行微调才能在下游任务生成领域比如诗歌,小说,文等领域获得比较好的生成效果。接下来我们通过 PAI-DSW 产品来演示如何执行 GPT-MoE 的下游任务微调。我们以 0.35B+MoE64 的 GPT 大模型为例,尝试对诗歌生成任务进行微调来获得比较好的生成效果。
环境准备
首先创建 DSW 任务,资源类型选择单机八卡 V100-32G,如下所示:
由于预训练 checkpoint 体积比较大,采用下载的方式不是很方便。借助于 DSW 提供的 OSS 路径的挂载功能,可以很方便的直接使用存储在 OSS 上的预训练模型来进行下游任务微调,配置方法如下图所示。
接着配置镜像,采用和 DLC 中同样的镜像地址:pai-image-manage-registry.cn-shanghai.cr.aliyuncs.com/pai/pai-rf-image-poc-moe:0.2
创建好后,点击打开进入到 DSW 交互式开发环境
数据准备
首先,您需要下载用于本示例的训练和测试集,并创建保存模型的文件夹,命令如下:
数据下载完成后,可以通过以下代码查看第一条数据。在训练集验证集中,每一行为一条诗歌数据。
初始化
在 EasyNLP 专属镜像环境下,我们首先引入模型运行需要的各种库,并做一些初始化。在本教程中,我们使用 gpt-moe-0.35B+MoE64 作为预训练模型底座。
模型训练 &评估
一步执行
值得一提的是,上述所有训练/评估,都已经被集成到 EasyNLP/examples/rapidformer/gpt_moe/finetune_megatron_gpt.py 中,可采用直接执行脚本文件 EasyNLP/examples/rapidformer/gpt_moe/run_finetune_text_generation_gpt.sh 的方式,一步执行上述所有训练/评估。根据脚本 run_finetune_text_generation_gpt.sh,需要设置的具体参数如下:
分析 Tensorboard
分析 Tensorboard,对比不同实验参数设置,比如我们选择调试训练轮次分别是 2 和 5,观察收敛曲线的变化,如下所示:
从上图可以看出,训练轮次越大,验证集 loss 越低,收敛效果越好。于是我们选择训练 5 个 epoch 后的 ckpt 来做诗歌生成的预测
模型预测
文本生成效果预测已经被集成到 EasyNLP/examples/rapidformer/gpt_moe/generate_text_gpt.py
然后根据脚本 run_finetune_text_generation_gpt.sh 里面的参数设置,运行命令:sh run_text_generation_gpt.sh
生成结果如下:
基于 PAI EAS 的在线推理部署
使用 DSW 完成 GPT-MoE 在下游任务上的微调后,就可以在线部署下游文本生成任务的服务了。接下来我们通过 PAI-EAS 产品来演示如何执行 GPT-MoE 的下游任务服务的在线部署。
开发基于 FasterTransformer 的 Processor
首先使用 Faster Transformer Converter 将微调后的诗歌模型进行格式转换
cd ~/RapidformerPro/examples/fastertransformer
sh run_convert.sh
然后基于 EAS 提供的开发参考https://help.aliyun.com/document_detail/130248.html,开发文本生成 processor。核心 process 方法参考如下
搭建在线服务
基于 ModelScope 的在线文本生成演示
我们已经搭建了诗歌,广告文案,作文三个文本生成服务,具体的在线体验地址是:https://www.modelscope.cn/organization/PAI
我们可以修改 modelscope 页面上的 README.md 信息来设计应用样式
在 modelscope 上呈现出来的效果如下所示:
参考文献
[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
相关阅读
版权声明: 本文为 InfoQ 作者【阿里云大数据AI技术】的原创文章。
原文链接:【http://xie.infoq.cn/article/dd6c3d0f80653ffc35135ca1f】。文章转载请联系作者。
评论