Perseus-BERT——业内性能极致优化的BERT训练方案

【做者】 笋江（林立翔）   驭策（龚志刚）   蜚廉（王志明）   昀龙（游亮）

一，背景——横空出世的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资源，这是专门作4x4矩阵乘法的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。

原文连接 本文为云栖社区原创内容，未经容许不得转载。