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

本篇blog是UCberkely的尤洋等基于BERT模型提出的新的pre-training方法。

---

Abstract

大型batch的训练是加速分布式Distributed DNN网络的关键。然而，large-batch训练是很困难的，他产生了很大的泛化误差generalization gap。直接的优化经常会导致accuracy在测试集上的降低（generalization）.BERT[4]作为一个SOTA几倍的DeepLearning模型，为语言理解模型构建了一个顶级的Deep bidirectional transformers。当我们使用一个很大的batch size=8292，之前的large-batch训练技巧在BERT上表现很差。BERT pre-training也会需要耗费大量时间（在16个TPUv3芯片上耗费3天）。为了解决这个问题，我们提倡使用LAMB优化器，这帮助我们固定batch size=65536，并且我们不会损失accuracy。LAMB是一个通用优化器，在很小或者很大的batch size上工作都很有效，并且也不会需要超参数的调试（除了learning rate以外）。baseline BERT-Large模型需要1m次循环来完成预训练，等同于LAMB以batch size=65536/32768，循环次数8599。我们让这个batch size限制在使用一个TPUv3，并且在76分钟内完成BERT训练(如表1)。

---

 

1 Introduction

训练DNN是很耗费时间的。目前，最有效减少时间的方法是使用多个并行的芯片（例如CPUs，GPUs，TPUs）来并行优化SGD的过程。因为数据依赖在多个不同的层之间进行前向传播和BP，在不同的层之间并行计算就不适用了。代替的是，研究者在每一次mini-batc的epoch中并行计算数据点，线性增加batch size，减少循环次数（例如权重更新的数量）。为了最小化wall clock时间，最大化batch size会是最有效的。

然而，large-batch训练是很困难的。例如，使用batch size大小为512的AlexNet[13]能在ImageNet上测试达到80%的top5的准确率。一个直接的前向方法只需要达到50%-60%的top5准确率在不同的结果上，当人们使用的batch size超过了4096。Keskar et al. [10]提出这里有generalization gap。Hoffer et al. [6] 提出训练更长时间可能帮助减小generalization gap。然而，训练更长时间意味着不会对large-batch训练有帮助。尽管，large-batch训练的膜目标是实现在固定的epochs下能够达到一个相当好水平的accuracy[5]。通过设计一系列的学习率[5，12，21]，研究者有能力将训练ImageNet的batch size达到32k，并且还有很小的accuracy loss。为了更好展现我们的想法，最快的结果是达到了76%的top1准确率Ying et al.[20]。通过使用LARS优化器[21]来调整batch size = 32k，Ying et al.[20]能使用ResNet-59在一块TPUv3上训练ImageNet，这个过程只花费了2.2分钟。

BERT[4]是一个SOAT级别DeepLearning模型。它为语言理解模型构建了最好的deep bidirectional transformers。先前的large-batch技术在这样大batch size（超过8192）中表现的十分差。BERT预训练需要很长时间来完成（在16块TPUv3上需要3天时间）。调整BERT batch size的大小，在这篇论文中我们提倡使用LAMB optimizer。LAMB提供adaptive element-wise 更新以及准确的layer-wise。LAMB是一个通用的优化器，不管batch size有多大。使用者仅仅需要调整学习率，并且不需要其他的参数。通过使用LAMB，我们只能调整BERT预训练的batch size到64K，并且不会有accuracy loss。BERT pre-training包含两个步骤：

1. 在9/10次epoch，使用序列长度为128
2. 最后1/10次epoch中，使用序列长度为512

baseline需要1m次循环来完成BERT预训练，我们只需要8599次循环，这样我们能减少BERT预训练时间从3天到76分钟。我们把batch size塞到一个TPU硬件中。这个batch size比32768大，序列长度为512，完全占据了内存。而batch size大于65536，序列长度为128不会带来任何加速。我们的优化器可能调整batch size到128k，甚至更大。因为硬件的限制，我们停止在序列长度为512，batch size=32768/序列长度为128，batch size=65536。所有的BERT模型都指出了BERT-Large模型。为了构建这个公平的比较，所有的实验都使用了相同大小的epochs。我们的结果在表1中。

---

2 Background and Related Work 

Deep神经网络是SGD。 作为batch size。对于每个步骤t，从训练集X选出 个样本。损失函数由数据的子集，权重w基于SGD更新：

SGD优化器的计算能通过数据并行计算每个机器的processor，每个处理器或机器计算它的小批量部分。增加小批处理大小允许扩展到更多的机器，而不会减少每台机器上的工作负载。然而，据观察，训练非常大的批次是困难的[6, 10]。研究人员需要仔细调整训练的超参数(如学习速度、动力等)，以避免失去测试的准确性[5, 14, 23]。

Krizhevsky[12]介绍了批量培训的一些实用方案。一个重要的规则是根据批量大小调整LR(学习率)。当批量|B|增加k倍时，应将LR增加k倍，使梯度估计方差保持在同一水平。在实践中，他发现线性缩放效果更好：当我们将|B|增加k时，我们应该将LR增加k，同时保持其他超参数(动量、重量衰减等)不变。使用这种方案，Krizhevsky扩大了批量AlexNet的大小为1024，与基线|B|=128相比，精度有轻微(1%)的损失。

通过线性或平方根LR缩放，大批量训练需要很大的学习率。然而，在训练的初始阶段，高的初始学习率往往会导致不稳定。为了克服这一问题，Goyal等人(5)提出使用学习率热身：训练从较小的LR开始，然后逐渐增加LR到较大的值。在热身阶段(通常是几个阶段)之后，我们切换到常规LR策略(多步骤、指数衰减或多项式衰减等)。Goyal(5)采用LR预热和线性标度法，成功地用批量|B|=8K训练了ResNet-50，测试精度没有损失。

大规模批量训练的另一个问题是泛化缺口，Keskar等人(10)在CIFAR-10和CIFAR-100上进行实验时发现了这一点。他们得出的结论是，泛化差距是由于大批量方法往往收敛到训练函数的急剧极小值。大批量训练容易达到较高的训练精度，但模型很难达到较高的测试精度。他们尝试了一些方法来提高泛化使用数据增强，热身开始小批量，但没有找到一个有效的解决方案。

有几组研究人员(1 3、7、8、16 19、22、23)成功地将批处理大小扩展到较大的值，并在几分钟内完成了ImageNet培训。据我们所知，Ying等人(20)实现了76+% top-1准确率的最快训练结果。通过使用LARS optimizer(21)将批处理大小扩展到32K, Ying等人(20)能够在TPUv3 Pod上用ResNet-50在2.2分钟内完成ImageNet培训。在本文中，我们的主要任务是BERT (4)， BERT是一项重要的技术，在GLUE、multiple li、SQuAD等几个NLP任务中都取得了最先进的成果。本文是第一个将BERT训练的批处理规模扩展到极大(即大于2K)的报告。

---

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

BERT训练的基线使用体重衰减的Adam作为优化器(15)，这是Adam优化器(11)的变体。另一个成功应用于大批量卷积神经网络训练的自适应优化器是LARS(21)。这些优化器启发我们提出了新的优化器，用于批量BERT培训。算法1显示了我们所提议的LAMB优化器的概述。

3.1 Correcting the LARS trust ratio

在LARS实现1中，作者定义了一个系数eeta来控制信任比。作者使用0.001作为默认系数eeta。系数eeta仅用于具有非零|g|(层梯度的L2范数)和非零|w|(层权重的L2范数)的层。对于零|w|或零|g|的层，LARS使用的信任比为1.0。然而，在某些情况下，这种方案可能会导致严重的问题或分歧。例如，我们假设使用默认系数eeta(0.001)。0 |w|和非0 |g|的层可能会得到非常大的信任比(如1000倍大)，这可能会导致网络发散。在信度计算中去掉系数eeta。对于具有非零|w|和非零|g|的层，我们使用不乘以系数eeta的规则LARS比。对于零|w|或零|g|的层，我们使用的信任比为1.0。这种修正避免了大批量BERT训练的发散。

3.2 Fixing weight decay in trust ratio

在LARS实现中，作者使用权值衰减来计算信任比，以精确地证明学习率。具体地说，他们使用公式(2)中的规则信任比在动量和SGD等优化器中工作得很好，因为它们对所有权重都有一个统一的学习率。

然而，对于像AdamW和Adam这样的自适应优化器，LARS比率会导致不准确的学习率调整，因为这些自适应优化器使用元素明智的更新方案。因此，我们使用了一种计算信任比的新方法来确保权重衰减项以元素的方式工作。LAMB的信任比定义在式(3)中，式(3)计算分母时采用的是元素方式，而不是分层方式。

3.3 Using more accurate information in gradient norm

在LARS优化器中，作者使用梯度范数和权值衰减作为信任比的分母。具体来说，他们用。然而，在实践中，我们发现在BERT训练中这个信息并不准确。LARS在ImageNet培训(例如通过ResNet-50或AlexNet)中工作得很好。这些模型的参数远远少于BERT(6100万对3亿)。对于BERT training，作者使用了具有重量衰减的Adam，这是一个明智的元素更新优化器。我们在LAMB优化器中保留了元素更新的特性。元素方面的信息来自对梯度的第一和第二时刻的估计。因此，我们使用梯度的第一阶矩和第二阶矩的估计值来进行分层学习速率证明。在这里，我们使用。这些符号在算法1中定义。为了便于理解，我们使用与传统Adam算法描述相同的符号。我们使用信任比来更新全局学习率(详见算法1)，不同的层会有不同的信任比。我们的信任比LARS优化器更精确。这样，我们不仅有了元素自适应学习率，而且有了分层校正的方案。我们通过LAMB优化器成功地将批处理大小扩展到32K，用于BERT培训。在我们的实验中，LARS optimizer能够将BERT的批处理大小扩展到8K。

3.3.1 Upper bound of trust ratio (a variant of LAMB)

即使我们使用基于梯度的一阶矩和二阶矩估计的元素明智更新，我们仍然可以在信度比计算中使用。然而，由于不准确的信息，我们观察到一些LARS的信任比率是非常大的。由于最初的LARS优化器使用动量SGD作为基础优化器，较大的信任比不会产生显著的负面影响，因为某个层中的所有元素使用相同的学习率。然而，对于Adam和Adagrad这样的自适应优化器，不同的元素具有不同的元素明智学习率。在这种情况下，较大的信度比可能会导致学习率较大的权重出现发散。一个实际的解决方案是设置信任比率的上限(例如将上限设置为10)。通过这种方式，我们仍然可以成功地将BERT训练的批处理大小扩展到16K，并且不会给LAMB优化器增加任何计算和内存开销。

---

4 Experimental Results

4.1 Regular Training

张量处理单元(9)是浮点运算的强大计算硬件。我们在所有的实验中都使用了TPUv3。TPUv3 Pod有1024个芯片，可以为混合精度计算提供超过100 petaflops的性能。我们的结果如表1所示。基线在培训前使用Wikipedia和BooksCorpus数据集(24)。我们使用与发布的BERT模型相同的数据集，即Wikipedia和BooksCorpus的拼接，分别为2.5B和8亿字。作者首先进行900k迭代对序列长度128进行训练，然后对序列长度512进行100k迭代训练。16个TPUv3芯片的总训练时间约为3天。我们使用SQuAD-v1的F1分数作为精度指标。F1得分越高，准确度越高。斯坦福问答数据集(SQuAD)是一个阅读理解数据集。它包含众包工作者在维基百科的一组文章中提出的问题，每一个问题的答案都是对应阅读文章的一段文字，或者是一个无法回答的问题。我们从公共BERT github3下载了预先训练好的模型。通过使用作者提供的脚本，基线的F1得分为90.395。在我们的代码中，我们使用了BERT的作者提供的数据集和基线模型，只修改了优化器。通过使用LAMB优化器，我们能够在批大小为32768的15625次迭代中获得91.460的F1分数(序列长度为128的14063次迭代和序列长度为512的1562次迭代)。我们把训练时间从3天减少到100分钟左右。我们将批量大小推到TPU Pod的硬件极限。批处理大小大于32768(序列长度为512)将导致TPU Pod耗尽内存。

我们实现了76.7%的弱缩放效率(49.1倍的加速，64倍的计算资源)。由于我们在TPU Pod上使用同步数据并行性进行分布式训练，因此在互连上传输梯度会带来通信开销。梯度的大小与训练后的模型相同。对于使用ResNet-50进行ImageNet训练，研究人员能够实现90%以上的弱标度效率，因为ResNet-50的参数比BERT少得多(2500万对3亿)。LAMB优化器不期望用户调优超参数。像Adam和AdamW一样，用户只需要输入学习率。β1用于衰减的运行平均梯度。β2用于衰减的运行平均梯度的平方。其他参数的默认设置算法1:权重衰变率。

4.2 Mixed-Batch Training

如前所述，BERT pre-training包括两个阶段:(1)所有纪元的前9/10使用序列长度为128，(2)所有纪元的后1/10使用序列长度为512。对于第二阶段，由于内存限制，TPUv3 Pod上的最大批处理大小为32768，因此我们将在第2阶段的32768处停止。对于第一个阶段，由于内存限制，TPUv3 Pod上的最大批处理大小是131072。但是，当我们将批大小从65536增加到131072时，并没有看到加速，因此我们在阶段1的65536处停止。在此之前，Smith等人(19)也研究了混合批处理训练。然而，他们在训练中增加了批次大小，而我们减少了批次大小。我们能够从头到尾充分利用硬件资源。它们只在最后阶段充分利用硬件资源。增加批大小可以预热和稳定优化过程，但是减少批大小会给优化过程带来混乱，导致发散。在实验中，我们发现了一种稳定第二阶段优化的方法。由于我们切换到一个不同的优化问题，有必要重新预热优化。在第二阶段，我们没有降低学习速度，而是在第二阶段(重新热身)将学习速度从零再次提高。和第一个阶段一样，我们在重新热身之后会降低学习速度。通过这些改变，我们只需要8599次迭代，可以在76分钟左右完成BERT训练，实现101.8%的弱缩放效率(65.2倍的加速，64倍的计算资源)。

 

5 Conclusion and Future Work

大批量技术是加快神经网络深度训练的关键。在本文中，我们提出了支持自适应元素更新和分层校正的LAMB优化器。LAMB是一个通用的优化器，适用于小批和大批。通过使用LAMB，我们可以将BERT pre-training的批量大小扩展到64K，而不会丢失准确性。我们将训练时间从3天减少到76分钟左右，并将批量大小推到TPU Pod的硬件极限。我们正在研究LAMB优化器的理论分析