ITPub博客

首页 > 人工智能 > 深度学习 > 大力再出奇迹,1024 张TPU,65536 batch size,仅76分钟训练完BERT!

大力再出奇迹,1024 张TPU,65536 batch size,仅76分钟训练完BERT!

原创 深度学习 作者:大数据文摘 时间:2019-04-04 15:32:46 0 删除 编辑

大数据文摘出品

作者:Andy

BERT 作为目前工业界中训练最耗时的应用,计算量甚至远大于机器视觉中的 ImageNet 训练。在 BERT 原论文中,Jacob Devlin 也是用了 16 台云TPU(64 个 TPU 芯片)花了整整四天,才训练完了 BERT-large 模型。

正因如此难以训练,导致大家也都几乎是直接拿 Google 开源出的模型,用 GPU 在上面 finetune 一下就来用。Github上的 BERT 开源库也因为其在 GPU 机器上预训练的困难和耗时,往往没给出预训练相关脚本。

论文下载地址:

https://arxiv.org/pdf/1904.00962.pdf

而在这篇愚人节当天于arxiv发出的 《Reducing BERT Pre-Training Time from 3 Days to 76 Minutes》论文中,正如标题所说,作者将BERT的训练时间缩减到了仅仅 76 分钟。

而所用的计算机器数也是惊人的,一个 TPUv3 pod (1024张TPU芯片)。不要以为这只是个靠堆积机器就能堆积上去的成果,这还涉及到具体的训练优化问题解决,比如说如何在增大batch size来提高计算通讯比的同时,又能保证其收敛。这篇论文中就是提出了一套训练方案,在新优化器 LAMB 的帮组下来解决这个问题。

这篇论文中就提出了一套训练方案,还有一个新的优化器LAMB来解决这个问题。

这是一作加大伯克利分校博士尤洋作为 Google Brain 实习生时完成的项目,其现在的主攻方向便是“高速并行分布深度学习算法”,这也是与这次论文密切相关的主题。

看他之前的论文也能看到类似的大规模训练项目,比如之前在英伟达实习时的技术报告《LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS》,怎么用大 batch size 训练CNN网络。

他在里面提出的 LARS,也正是这篇论文中提出优化器 LAMB 的最初原型。论文中也有拿 LARS 做了实验。而说到在Google Brain实习也是有迹可循,之前他便已在这实习过一段时间研究关于TPU上使用Tensorflow的课题,也为此次项目打下铺垫。

论文下载地址:

https://arxiv.org/pdf/1708.03888.pdf

看他之前的论文也能看到类似的大规模训练项目,比如他之前在英伟达实习时的技术报告《LARGE BATCH TRAINING OF CONVOLUTIONAL NETWORKS》,怎么用大 batch size 训练CNN网络。

他在里面提出的 LARS,也正是这篇论文中提出优化器 LAMB 的最初原型。论文中也拿LARS做了实验。而说到在Google Brain实习也是有迹可循,之前他便已在这实习过一段时间研究关于TPU上使用Tensorflow的课题。

下面就来介绍论文吧。

首先在导言部分,作者稍稍介绍了大批量训练的困难,还有此次面对的训练对象BERT。为处理大批量 BERT 的训练,作者提出了 LAMB 优化器,通过这个优化器将BERT训练的batch size推到了64k(具体65536)的量级,而同时不损失精度,此外LAMB优化器还有一个优点就是只用调学习率。

最后的预训练包括两个阶段:前九个epoch用128的序列长度和64k的batch size;最后一个epoch用512的序列长度和32k的batch size,只用了惊人的8599个迭代便训练完了BERT。而它和基线模型还有其他batch size训练对比,如下图。

之后,在背景部分,作者给我们分享了关于大批量训练的知识。目前大批量训练常常会遇到的几个问题,还有如何对其中一些进行处理。

  • 大批量会导致测试精度丢失,所以需要调节超参,比如学习率。所以需要随着batch size的增大,线性或平方根级地增大学习率;

  • 但大学习率会导致训练初始不稳定,因此需要使用学习率预热技巧(learning rate warmup),先用一个小学习率然后慢慢增大,到一定点切换到正常的学习率策略;

  • 还有大批量训练里的泛化间距(generalization gap)问题,大批量训练模型会倾向于收敛到比较尖锐的局部最小点,这会导致训练很容易达到一个比较高的训练精度,却很难获得比较好的测试精度,也就是存在比较大的间距(gap)。目前还没找到合适方法解决。

然后再来详细了解一下本文最重要的创新,LAMB (Layer-wise Adaptive Moments optimizer for Batch training)优化器吧。它是基于作者之前提出的 LARS (Layer-wise Adaptive Rate Scaling)优化器,同时又参考了BERT原有优化器进行了改进。

首先为了理解 LAMB 我们需要对 LARS 有个粗略的了解。LARS 提出的背景是,作者发现对于网络各层,其权重和梯度的L2-norm的比值|w|/|g|变化会非常大,比如5.76和1345。

这导致的一个问题就是,一个学习率并不适应于所有层训练,它对于有些层或许适合,但对于有些层却可能太大。因此作者提出应该要按照层,来获得每层的学习率(Local LR),而这个每层学习率的计算则需要之后 LAMB 中多次提到的一个trust ratio,置信比,有多大的可能我们相信当前层会在这次更新中改变它的权重。

于是 LAMB 对 LARS 最大的改进有三点,都是与trust ratio相关。

1.通过在Tensorflow具体的LARS优化器实现中,移除其中一个当某层的|w|和|g|都非零时用于计算 trust ratio 的系数 eeta,从而避免了BERT大批量训练中的发散;

2.LARS 自身的trust ratio在一些自适应优化器,比如BERT里用到的ADAM with weight decay 或 ADAM中,会导致不准确的学习率矫正,因为它们用了元素级的更新策略。于是作者们将trust_ratio做了些修正,保证了它和自适应优化器的结合。

3.最后一点,在计算梯度的L2-norm的时候,还加入了梯度的一阶和二阶惯量的信息。

LAMB 算法具体如下:

最后具体实验部分便不累述,值得一提的一个细节是,在混合batch训练的时候,因为中间有一个阶段转换过程,将batch size从64k降到了32k,这里会给优化过程带来噪音,进而导致训练发散。为解决这个问题,在第二个阶段的时候作者又重新进行了一次学习率预热(re-warm-up)。

来自 “ ITPUB博客 ” ,链接:http://blog.itpub.net/31562039/viewspace-2640433/,如需转载,请注明出处,否则将追究法律责任。

请登录后发表评论 登录
全部评论

注册时间:2018-11-13

  • 博文量
    196
  • 访问量
    146425