76分钟训练BERT!谷歌大脑新型优化器LAMB加速大批量训练

选自arXiv

作者:Yang You、Jing Li等

机器之心编辑部

去年,谷歌发布了基于双向 Transformer 的大规模预训练语言模型 BERT 并开源。该模型参数量非常大——3 亿,训练时间也很长。近日,来自谷歌大脑的研究者提出了一种新型优化器 LAMB,有效降低 BERT 预训练时间,76 分钟即可完成 BERT 预训练!

尽管 BERT效果惊人,但它所需的计算量非常大,原作者在论文中也表示每次只能预测 15% 的词,因此模型收敛得非常慢。如果我们想保留这种 Mask 机制,那么就需要寻找另一种加速方法了。

当时,BERT 的作者在 Reddit 上也表示预训练的计算量非常大,Jacob 说:「OpenAI 的 Transformer 有 12 层、768 个隐藏单元,他们使用 8 块 P100 在 8 亿词量的数据集上训练 40 个 Epoch 需要一个月,而 BERT-Large 模型有 24 层、2014 个隐藏单元,它们在有 33 亿词量的数据集上需要训练 40 个 Epoch,因此在 8 块 P100 上可能需要 1 年?16 Cloud TPU 已经是非常大的计算力了。」

而在谷歌大脑的这篇新论文中,研究者提出新型 优化器LAMB,通过使用 65536/32768 的批量大小,他们只需要 8599 次迭代、76 分钟就能完成 BERT 预训练。总体而言,相比基线 BERT-Large 用 16 块 TPU 芯片,LAMB 训练 BERT-Large 用了一个 TPU v3 Pod(1024 块 TPU 芯片),因此时间也由 3 天降低为 76 分钟。

不过在 BERT 原论文中,训练 BERT-Large 使用了 64 块 TPU 芯片(16 Cloud TPU),它需要花四天时间完成训练。

论文:Reducing BERT Pre-Training Time from 3 Days to 76 Minutes

论文地址:https://arxiv.org/abs/1904.00962

摘要:大批量训练是加速大型分布式系统中深度神经网络训练的关键。但是,大批量训练难度很大,因为它会产生泛化差距(generalization gap),直接优化通常会导致测试集准确率受损。BERT [4] 是当前最优的深度学习模型,它基于用于语言理解的深度双向 transformer 而构建。当我们扩展批量大小时(比如批量大小超过 8192),之前的大批量训练技术在 BERT 上的执行性能并不好。BERT 预训练需要大量时间(使用 16 个 TPUv3 训练 3 天)。

为了解决这个问题,来自谷歌大脑的研究者提出了一种新型优化器 LAMB,可在不损害准确率的情况下将批量大小扩展至 65536。LAMB 是一款通用优化器,它适用于小批量和大批量,且除了学习率以外其他超参数均无需调整。基线 BERT-Large 模型的预训练需要 100 万次迭代,而 LAMB 使用 65536/32768 的批量大小,仅需 8599 次迭代。研究者将批量大小扩展到 TPUv3 pod 的内存极限,在 76 分钟内完成了 BERT 的训练。

具体来讲,LAMB 优化器支持自适应元素级更新(adaptive element-wise updating)和准确的逐层修正(layer-wise correction)。LAMB 可将 BERT 预训练的批量大小扩展到 64K,且不会造成准确率损失。BERT 预训练包括两个阶段:1)前 9/10 的训练 epoch 使用 128 的序列长度,2)最后 1/10 的训练 epoch 使用 512 的序列长度。

基线 BERT-Large 模型的预训练需要 100 万次迭代,研究者仅使用 8599 次迭代就完成了预训练,将训练时间从 3 天缩短到 76 分钟。该研究使用的训练批量大小接近 TPUv3 pod 的内存极限。LAMB 优化器可以将批量大小扩展到 128k 甚至更大,由于硬件限制,研究者在序列长度为 512 时使用的批量大小为 32768,在序列长度为 128 时使用的批量大小为 65536。该论文中的 BERT 模型均指 BERT-Large。为公平起见,研究中所有实验均运行同样数量的 epoch(即固定数量的浮点运算)。实验结果见下表。

表 1:该研究使用 SQuAD-v1 的 F1 分数作为准确率度量。

在上表中,基线 F1 分数来自于 BERT 开源 GitHub 中 BERT-Large 预训练模型的对应分数。实验所用硬件为 TPUv3,实验设置与基线 BERT-Large 相同:前 9/10 的训练 epoch 使用 128 的序列长度,最后 1/10 的训练 epoch 使用 512 的序列长度。所有实验均运行同样数量的 epoch。

LAMB (Layer-wise Adaptive Moments optimizer for Batch training)

BERT 基线模型的训练使用 Adam with weight decay(Adam优化器的变体)作为优化器 [15]。另一个成功用于大批量卷积神经网络训练的自适应优化器是 LARS [21]。这些优化器启发了研究者提出用于大批量 BERT 训练的新型优化器 LAMB。LAMB 优化器详见算法 1。

实验

常规训练

TPU 是强大的浮点运算计算硬件。研究者在所有实验中均使用 TPUv3。每个 TPUv3 pod 拥有 1024 个芯片,可提供超过 100 petaflops 的混合精度计算。实验结果见表 1。基线模型预训练过程中使用了 Wikipedia 和 BooksCorpus 数据集。研究者使用与开源 BERT 模型相同的数据集进行预训练,即包含 2.5B 单词的 Wikipedia 数据集和包含 800M 单词的 BooksCorpus 数据集。

BERT 作者首先以 128 的序列长度进行 900k 次迭代,然后以 512 的序列长度进行 100k 次迭代。在 16 块 TPUv3 上训练的总时间大约为 3 天。该研究使用 SQuAD-v1 的 F1 分数作为准确率度量。较高的 F1 分数意味着较好的准确率。研究者下载了 BERT 开源项目中提供的预训练模型。使用 BERT 作者提供的脚本,基线模型得到了 90.395 的 F1 分数。

该研究中,研究者使用 BERT 作者提供的数据集和基线模型,仅改变了优化器。使用新型优化器 LAMB 后,研究者以 32768 的批量大小进行了 15625 次迭代,得到了 91.460 的 F1 分数(用于序列长度 128 的迭代有 14063 次,用于序列长度 512 的迭代有 1562 次)。研究者将 BERT 训练时间从 3 天缩短到大约 100 分钟。

该研究取得了 76.7% 的弱可扩展性效率(weak scaling efficiency)。研究者在 TPU Pod 上使用了分布式训练的同步数据并行化,因此梯度迁移会产生通信开销。这些梯度与训练后的模型大小一样。在 ImageNet 数据集上训练 ResNet-50 时的弱可扩展性效率可达到 90+%,因为 ResNet-50 比 BERT 的参数少得多(25 million vs 300 million)。LAMB 优化器无需用户调整超参数,用户只需输入学习率即可。

混合批次训练(Mixed-Batch Training)

如前所述,BERT 预训练主要分为两部分,1)前面 9/10 的 Epoch 使用 128 的序列长度;2)最后 1/10 的 Epoch 使用 512 的序列长度进行训练。对于第二阶段而言,因为内存限制,TPUv3 Pod 上最大的批量大小为 32768,因此第二阶段使用的批大小为 32768。对于第一阶段,受限于内存,TPUv3 Pod 上最大的批量大小为 131072。然而,研究者将批大小从 65536 增加到 131072 时,并没有发现明显的加速,因此研究者将阶段 1 的批大小定为 65536。

此前,也有其它研究者讨论了混合批次训练,然而他们是在训练中增大批量大小;而本研究是降低批大小,这样他们能从开始到结束充分利用硬件资源。增大批大小能起到 warm-up 初始训练和稳定最优化过程的效果,但是降低批大小可能会带来最优化过程的紊乱,并导致训练的不收敛。

在研究者的实验中,他们发现有一些技术能稳定第二阶段的最优化过程。因为这两阶段会切换到不同的最优化问题,因此有必要重新 warm-up 最优化过程。研究者没有在第二阶段进行学习率衰减,而是将学习率从零开始增加(re-warm-up)。和第一阶段一样,研究者在 re-warm-up 之后执行学习率衰减。因此,研究者能够以 8599 次迭代完成 BERT 训练,且仅用时 76 分钟,达到了 101.8% 的弱可扩展性效率。

原文发布于微信公众号 - 机器之心(almosthuman2014)

原文发表时间:2019-04-03

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

编辑于

我来说两句

0 条评论
登录 后参与评论

扫码关注云+社区

领取腾讯云代金券