Perseus-BERT——业内性能极致优化的BERT训练方案【阿里云弹性人工智能】

38次阅读

共计 5138 个字符,预计需要花费 13 分钟才能阅读完成。

一,背景——横空出世的 BERT 全面超越人类
2018 年在自然语言处理(NLP)领域最具爆炸性的一朵“蘑菇云”莫过于 Google Research 提出的 BERT(Bidirectional Encoder Representations from Transformers)模型。作为一种新型的语言表示模型,BERT 以“摧枯拉朽”之势横扫包括语言问答、理解、预测等各项 NLP 锦标的桂冠,见图 1 和图 2。

【图 1】SQuAD 是基于 Wikipedia 文章的标准问答数据库的 NLP 锦标。目前 SQuAD2.0 排名前十名均为基于 BERT 的模型(图中列出前五名),前 20 名有 16 席均是出自 BERT *

【图 2】GLUE 是一项通用语言理解评估的 benchmark,包含 11 项 NLP 任务。BERT 自诞生日起长期压倒性霸占榜首(目前 BERT 排名第二,第一为 Microsoft 提交的 BIGBIRD 模型,由于没有 URL 链接无从知晓模型细节,网传 BIGBIRD 的名称上有借鉴 BERT BIG 模型之嫌)
 业内将 BERT 在自然语言处理的地位比作 ResNet 之于计算机视觉领域的里程碑地位。在 BERT 横空出世之后,所有的自然语言处理任务都可以基于 BERT 模型为基础展开。
一言以蔽之,现如今,作为 NLP 的研究者,如果不了解 BERT,那就是落后的科技工作者;作为以自然语言处理为重要依托的科技公司,如果不落地 BERT,那就是落后生产力的代表。
二,痛点——算力成为 BERT 落地的拦路虎
BERT 强大的原因在哪里?让我们拂去云霭,窥探下硝烟下的奥秘。
BERT 模型分为预训练模型(Pretrain)和精调模型(Finetune)。Pretrain 模型为通用的语言模型。Finetune 只需要在 Pretrain 的基础上增加一层适配层就可以服务于从问答到语言推理等各类任务,无需为具体任务修改整体模型架构,如图 3 所示。这种设计方便 BERT 预处理模型适配于各类具体 NLP 模型(类似于 CV 领域基于 ImageNet 训练的各种 Backbone 模型)。

【图 3】左图基于 BERT pretrain 的模型用于语句问答任务(SQuAD)的 finetune 模型,右图为用于句对分类(Sentence Pair Classification Tasks)的 finetune 模型。他们均是在 BERT Pretrain 模型的基础上增加了一层具体任务的适配层
因此,BERT 的强大主要归功于精确度和鲁棒性俱佳的 Pretrain 语言模型。大部分的计算量也出自 Pretrain 模型。其主要运用了以下两项技术,都是极其耗费计算资源的模块。
1. 双向 Transformer 架构
图 4 可见,与其他 pre-training 的模型架构不同,BERT 从左到右和从右到左地同时对语料进行 transformer 处理。这种双向技术能充分提取语料的时域相关性,但同时也大大增加了计算资源的负担。【关于 Transformer 是 Google 17 年在 NLP 上的大作,其用全 Attention 机制取代 NLP 常用的 RNN 及其变体 LSTM 等的常用架构,大大改善了 NLP 的预测准确度。本文不展开,该兴趣的同学可以自行搜索一下】。

【图 4】Pretrain 架构对比。其中 OpenAI GPT 采用从左到右的 Transformer 架构,ELMo 采用部分从左到右和部分从右到左的 LSTM 的级联方式。BERT 采用同时从左到右和从右到左的双向 Transformer 架构。
1. 词 / 句双任务随机预测
BERT 预训练模型在迭代计算中会同时进行单词预测和语句预测两项非监督预测任务。
其一,单词预测任务对语料进行随机 MASK 操作(Masked LM)。在所有语料中随机选取 15% 的单词作为 Mask 数据。被选中 Mask 的语料单词在迭代计算过程中 80% 时间会被掩码覆盖用于预测、10% 时间保持不变、10% 时间随机替换为其他单词,如图 5 所示。
其二,语句预测任务(Next Sentence Prediction)。对选中的前后句 A 和 B,在整个迭代预测过程中,50% 的时间 B 作为 A 的真实后续语句(Label=IsNext),另外 50% 的时间则从语料库里随机选取其他语句作为 A 的后续语句(Label=NotNext),如图 5 所示

【图 5】词 / 句双任务随机预测输入语料实例。蓝框和红框为同一个语料输入在不同时刻的随机状态。对单词预测任务,蓝框中的“went”为真实数据,到了红框则被 [MASK],红框中的“the”则相反;对于语句预测任务,蓝框中的句组为真实的前后句,而红框中的句组则为随机的组合。
这种随机选取的单词 / 语句预测方式在功能上实现了非监督数据的输入的功能,有效防止模型的过拟合。但是按比例随机选取需要大大增加对语料库的迭代次数才能消化所有的语料数据,这给计算资源带来了极大的压力。
综上,BERT 预处理模型功能需要建立在极强的计算力基础之上。BERT 论文显示,训练 BERT BASE 预训练模型(L=12, H=768, A=12, Total Parameters=110M, 1000,000 次迭代)需要 1 台 Cloud TPU 工作 16 天;而作为目前深度学习主流的 Nvidia GPU 加速卡面对如此海量的计算量更是力不从心。即使是目前主流最强劲的 Nvidia V100 加速卡,训练一个 BERT-Base Pretrain 模型需要一两个月的时间。而训练 Large 模型,需要花至少四五个月的时间。
花几个月训练一个模型,对于绝大部分在 GPU 上训练 BERT 的用户来说真是伤不起。
三,救星——擎天云加速框架为 BERT 披荆斩棘
阿里云弹性人工智能团队依托阿里云强大的基础设施资源打磨业内极具竞争力的人工智能创新方案。基于 BERT 的训练痛点,团队打造了擎天优化版的 Perseus-BERT, 极大地提升了 BERT pretrain 模型的训练速度。在云上一台 V100 8 卡实例上,只需 4 天不到即可训练一份 BERT 模型。
Perseus-BERT 是如何打造云上最佳的 BERT 训练实践?以下干货为您揭秘 Perseus-BERT 的独门绝技。
1.  Perseus 统一分布式通信框架 —— 赋予 BERT 分布式训练的轻功
Perseus(擎天)统一分布式通信框架是团队针对人工智能云端训练的痛点,针对阿里云基础设施极致优化的分布式训练框架。其可轻便地嵌入主流人工智能框架的单机训练代码,在保证训练精度的同时高效地提升训练的多机扩展性。擎天分布式框架的干货介绍详见团队另一篇文章《Perseus(擎天):统一深度学习分布式通信框架》。
针对 tensorflow 代码的 BERT,Perseus 提供 horovod 的 python api 方便嵌入 BERT 预训练代码。基本流程如下:

让每块 GPU 对应一个 Perseus rank 进程;
对 global step 和 warmup step 做基于 rank 数的校准;
对训练数据根据 rank-id 做划分;
给 Optimizer 增加 DistributeOptimizer 的 wrapper。

值得注意的是,BERT 源码用的自定义的 Optimizer,在计算梯度时采用了以下 api
grads = tf.gradients(loss, tvars)
Perseus 的 DistributeOptimizer 继承标准的 Optimizer 实现,并在 compute_gradients api 上实现分布式的梯度更新计算。因此对 grads 获取做了如下微调
grads_and_vars = optimizer.compute_gradients(loss, tvars)

grads = list()

for grad, var in grads_and_vars:

grads.append(grad)
2.  混合精度训练和 XLA 编译优化——提升 BERT 单机性能的内功
混合精度
在深度学习中,混合精度训练指的是 float32 和 float16 混合的训练方式,一般的混合精度模式如图 6 所示

【图 6】混合精度训练示例。在 Forward+Backward 计算过程中用 float16 做计算,在梯度更新时转换为 float32 做梯度更新。
混合梯度对 Bert 训练带来如下好处,
增大训练时的 batch size 和 sequence_size 以保证模型训练的精度。
      目前阿里云上提供的主流的 Nvidia 显卡的显存最大为 16GB,对一个 BERT-Base 模型在 float32 模式只能最高设置为 sequence_size=256,batch_size=26。BERT 的随机预测模型设计对 sequence_size 和 batch_size 的大小有一定要求。为保证匹配 BERT 的原生训练精度,需要保证 sequece_size=512 的情况下 batch_size 不小于 16。Float16 的混合精度可以保证如上需求。
混合精度能充分利用硬件的加速资源。
      NVidia 从 Volta 架构开始增加了 Tensor Core 资源,这是专门做 4 ×4 矩阵乘法的 fp16/fp32 混合精度的 ASIC 加速器,一块 V100 能提供 125T 的 Tensor Core 计算能力,只有在混合精度下计算才能利用上这一块强大的算力。
   受限于 float16 的表示精度,混合精度训练的代码需要额外的编写,NVidia 提供了在 Tensorflow 下做混合精度训练的教程。其主要思路是通过 tf.variable_scope 的 custom_getter 参数保证存储的参数为 float32 并用 float16 做计算。
   在 BERT 预训练模型中,为了保证训练的精度,Perseus-BERT 没有简单的利用 custom_getter 参数,而是显式指定训地参数中哪些可以利用 float16 不会影响精度,哪些必须用 float32 已保证精度。我们的经验如下:

Embedding 部分要保证 float32 精度;
Attetion 部分可以利用 float16 加速;
Gradients 相关的更新和验证需要保证 float32 精度;
非线性激活等模块需要保证 float32 精度。

XLA 编译器优化
XLA 是 Tensorflow 新近提出的模型编译器,其可以将 Graph 编译成 IR 表示,Fuse 冗余 Ops,并对 Ops 做了性能优化、适配硬件资源。然而官方的 Tensorflow release 并不支持 xla 的分布式训练,为了保证分布式训练可以正常进行和精度,我们自己编译了带有额外 patch 的 tensorflow 来支持分布式训练,Perseus-BERT 通过启用 XLA 编译优化加速训练过程并增加了 Batch size 大小。
3.  数据集预处理的加速
Perseus BERT 同时对文本预处理做的 word embedding 和语句划分做了并行化的优化。这里就不展开说明。
四,性能——计算时间单位从月降低到天
图 7 展示了 Perseus BERT 在 P100 实例上的性能,与开源主流的 horovod 相比,Peseus-BERT 双机 16 卡的分布式性能是前者的 5 倍之多。
目前某大客户已在阿里云 P100 集群上大规模上线了 Perseus BERT,用 10 台 4 卡 P100 只需要 2.5 天即可训练完成业务模型,如果用开源的 horovod(Tensorflow 分布式性能优化版)大概需要 1 个月的时间。

【图 7】Bert 在阿里云上 P100 实例的对比(实验环境 Bert on P100;Batch size: 22;Max seq length: 256;Data type:float32;Tensorflow 1.12;Perseus: 0.9.1;Horovod: 0.15.2)
为了和 Google TPU 做对比,我们量化了 TPU 的性能,性能依据如图 8。一个 Cloud TPU 可计算的 BERT-Base 性能 256(1000000/4/4/24/60/60)= 185 exmaples/s。而一台阿里云上的 V100 单机八卡实例在相同的 sequence_size=512 下, 通过 Perseus-BERT 优化的 Base 模型训练可以做到 680 examples/s,接近一台 Cloud TPU 的 4 倍性能。对一台 Cloud TPU 花费 16 天才能训练完的 BERT 模型,一台阿里云的 V100 8 卡实例只需要 4 天不到便可训练完毕。

【图 8】BERT Pretain 在 Google Cloud TPU 上的性能依据 *
五,总结——基于阿里云基础设施的 AI 极致性能优化
弹性人工智能团队一直致力基于阿里云基础设施的 AI 极致性能优化的创新方案。Perseus-BERT 就是一个非常典型的案例,我们在框架层面上基于阿里云的基础设施做深度优化,充分释放阿里云上基础资源的计算能力,让阿里云的客户充分享受云上的 AI 计算优势,让天下没有难算的 AI。

本文作者:笋江阅读原文
本文为云栖社区原创内容,未经允许不得转载。

正文完
 0