点击下面卡片，关注我呀，每天给你送来AI技术干货！

来自：AI部落联盟

Part1简介

我们都知道训练BERT是一件很费时费卡的事情。本文研究了Nvidia对BERT模型的实现，看看BERT的训练是如何能够在47分钟内完成的。

本文将重点关注实现的原理，特别是新的LAMB优化器，它允许使用大批量数据进行训练而不会破坏训练的稳定性。此外，本文将介绍一些Nvidia实现时的细节。

在研究解决方案之前，让我们首先系统地研究一下scaling问题。

1异步（Asynchronous）

在神经网络的训练过程中，由于需要不断的迭代更新参数以优化loss函数，该训练方式限制了并行化的可能，所以导致训练过程较慢。

但是有无可能打破这一限制呢？答案是可以通过异步参数更新稍微放松迭代过程的顺序约束，即可以使用稍微过时的模型参数重叠顺序迭代。然而，单独使用异步更新，训练结果的准确性在实验中会受到影响。因此，同步参数更新目前仍然是比较流行的做法。

2更好的精度（Better Accuracy）

由于迭代过程的顺序约束，我们把重点放在提高每次迭代的训练效率上。许多梯度下降方法都采用近似的方法提高每次迭代的训练效率：通过引入二阶导数的同时考虑损失函数的图形曲线特征，从而提高精度。这种方法在强化学习中大有前途。然而该方法也显著增加了计算的复杂度。因此，即使该方法更准确，但它对于深度学习模型的整体训练速度的提升仍然值得怀疑。

3更好的参数更新方法（Better parameter updates）

事实上，大多数深度学习模型的训练过程采用了一种不那么复杂的方法，即专注于调整学习率或引入基于一阶导数的参数更新方案。后一种方法包括基于动量的参数更新和pre-parameter自适应学习率的方法，如RMSprop。流行的Adam优化器将这两个概念结合在了一起。本文将详细介绍并重新应用这些概念，以提高训练的效率。

4更好的并行性（Better parallelism）

为了加快训练速度，我们可以提高每次迭代的并行化。有两种常见的方法：模型并行和数据并行。在模型并行方面，可以将模型划分为多个部分，在不同的GPU上运行不同的部分。然而模型的顺序性质仍然限制了并行性的可能。

5更大的batch_size（Bigger batch size）

在数据并行中，每个GPU会计算不同数据的梯度损失。例如，对于batch_size为1024的数据，可以使用16个gpu，每个gpu负责64个训练样本。因为更大的batch_size更能抵消信息噪声，并有望产生更好的梯度下降。然而，为了显著减少总的训练时间，还需要一个更好的learning schedule。在许多深度学习模型的训练中，可以通过提高学习率来实现的。例如，在batch_size增加了k倍之后，也可以增加gpu的数量k倍来保持每次迭代的训练时间不变(不缺卡的情况)。然后同时将学习率提高k倍，以加快训练速度。但是Krizhevsky在2014年却表明：

非常大的batch_size会对SGD收敛的速度以及最终解决方案的质量产生不利影响。（But very big batch sizes adversely affect the rate at which SGD converges as well as the quality of the final solution.）

6更好的优化器（Better Optimizer）

有时我们会说大batch_size会对训练无益。但明确的一点是，在特定的问题中，一旦我们跨过了一个特定的batch_size规模，必须考虑该学习方案会带来的风险，否则模型将不能很好地收敛。本文将详细介绍将BERT训练缩短到47分钟的方案。其中Nvidia的实现中将Adam优化器变成了LAMB优化器，这一改变是使用更大的batch_size来扩展训练的核心。首先需要了解一些基本概念和一些训练技巧。

---

Part2BERT的训练策略（BERT Training Strategy）

本节将对BERT的训练策略进行一个快速的概述。假设您已经对BERT有基本的了解，如果没有，请参阅本文。

BERT有两种版本：base和large。参数请见下表：

例如，BERT-base模型有12个编码器层，每个编码器的设计如下面的左图所示。

BERT的训练分为两个阶段:

• 预先训练生成输入句子的通用密集向量表示；

• 微调以解决具体的下游问题，如QA。

预训练通过未标记数据对模型进行训练，最后可以得到输入句子的通用向量表示。在微调阶段，我们需要为相应的具体下游任务专门标记数据集。在模型设计方面，我们只是在训练前在模型上添加一到两个全连接层，从而产生微调模型。

由于BERT-large模型具有更好的精度，所以本文将重点讨论这个更大的模型。这个模型非常大，很多公司用它来展示他们是如何进行训练的。

Part3训练BERT的一般技巧

冰冻三尺非一日之寒。深度学习模型的训练过程一般都是分阶段进行的。一些研究人员逐渐增加每个阶段的复杂度，使训练更稳定，更快以及避免陷入局部极小值。对于BERT在预训练过程中可分为两个阶段。第一阶段使用长度为128的较短输入序列。第二阶段使用更少的训练steps，但序列长度更长，为512。下面是关于在Nvidia实现中运行预训练脚本的文档。

run_pretraining_lamb.sh - 使用LAMB优化器使用run_pretraining.py文件分两个阶段运行预训练过程。阶段1在序列长度为128的情况下完成90%的训练。在阶段2中，剩余10%的训练是在序列长度=512的情况下完成的。

下面是谷歌实现的理由:

较长的序列计算代价过高，因为注意力机制的计算复杂度是序列长度的二次方。也就是说，一批长度为512的64个序列比一批长度为128的256个序列的计算代价高得多。全连接/卷积的成本是相同的，但是512长度的序列的注意力成本要大得多。因此，好的方法是在训练前，以128的序列长度完成90,000步，然后再以512的序列长度完成10,000步。非常长的序列主要用于学习位置的embedding，可以快速学习。

7Reuse of pre-trained models

预先训练一个BERT-large模型需要花费大量的GPU或TPU资源。它可以在本地或通过云服务进行训练。幸运的是，有一些预先训练好的模型可以用来启动这个过程。例如，我们可以先用一个预先训练的模型来加载Transformer，然后用一个特定领域的数据集(使用较小的训练率)进一步对它进行预先训练。这种继续训练的方式通常花费的时间更少。

例如，下面段落中的阴影区域是由经过技术期刊语料训练的Transformer自动生成的。

BioBERT是使用PubMed摘要和PMC全文文章语料进一步对BERT模型进行预训练的另一个例子。在原始研究论文中，该模型使用了8个Nvidia V100 gpu，用生物医学语料库经过23天完成。

对于微调部分，使用单个GPU可以在数小时内完成。大部分微调过程只需要训练两个Epoch就能达到不错的效果。

Part4大的Mini-Batch

对模型操作的顺序依赖和迭代优化的方法限制了并行化的程度。如果不能设计一个高度并行的系统，就不能通过添加新的gpu来扩展训练。

我们可以通过增加batch的大小来增加数据的并行性。然而，为了缩短训练时间，我们还需要一个高效的训练策略(更高的学习率)。然而，这种策略很容易适得其反，破坏训练的稳定性。在接下来的几节中，在研究LAMB之前，我们将首先介绍一些基本技巧。

8线性扩展规则

当mini-batch的大小 

 乘以 

 时，我们应该像某些理论所建议的那样，把起始学习速率 




η



 乘以 

 的平方根。然而通过多个研究人员的实验，线性缩放显示出更好的结果，即将开始学习速率乘以 

。

9逐渐预热策略

然而，我们不会一开始就使用这个初始学习速率 

η



 。应该从 




η



 开始，增加一个固定的量，使其在预定义的预热step数之后达到

η



。这种逐步的热身将提供类似于ImageNet训练中的使用的小batch(使用相同数量的训练样本，但更少的训练步骤)的训练精度。

通过应用这两种技术，我们可以在ImageNet训练中将mini-batch的大小推到8K，而不丢失准确性。

Part5Optimizer

为了提高mini-batch的大小，我们需要从Adam优化器切换到LAMB优化器。但首先让我们快速回顾一下一些常用的优化器。

10动量随机梯度下降（Momentum SGD）

从概念上讲，动量SGD可以看作是当前梯度和以前梯度的加权平均值。这种加权方法降低了梯度噪声。为了减少旧更新的影响，我们在每次迭代中衰减 

 值(多项式衰减)。例如，

₃

₂

²

₁

11RMSProp

普通的模型优化器的主要问题是所有模型参数都使用单一的学习率。这种一刀切的解决方案很难满足每一个具有不同曲率的下降方向。但是我们也不可能控制每个模型参数的学习率，因为有数百万个参数，但是可以隐式地实现。

下图的第一个方程记录了在速度 

 下的梯度大小的多项式加权平均。更准确地说，它的平方根值度量了最近参数变化幅度的加权平均值。

在参数更新过程中，可以用它的根来划分全局学习速率 




η



。实际上如果一个参数最近变化太快或太大，该方式就会对其进行抑制。结合全局学习率，控制每个节点的学习率。

12Adam

Adam优化器在更新模型参数时简单地结合了Momentum和RMSProp的思想，如下所示。

13AdamW

L2正则化和权重衰减正则化在普通SGD中是等价的，因为它们在数学上的梯度下降法中是相等的。但对于像Adam这样的自适应方法是不相等的。

AdamW在参数更新中应用权重衰减(下面绿色高亮部分)，并没有在梯度中应用L2正则化，因此它们将产生不同的结果。实验结果表明，训练模型的过拟合程度较低，泛化程度较好。

14Layer-wise Adaptive Rate Scaling (LARS)

在RMSProp中，每个参数独立并以不同的速率学习。但也许学习率在层级别上也是敏感的，那是否可以对其进行实验评价呢？

实验结果表明，在不同的层中，权重和梯度的比值变化幅度有很大的差异。

这为每一层的学习速率应该分别控制的说法提供了支持，这可能有助于慢学习层学习得更快。

梯度下降过程中模型参数的变化与梯度成正比。一个不成比例或不受控制的参数会增加训练过程中的分歧。在许多其他算法中，梯度裁剪和权值归一化被引入，就是为了避免这种训练过程中出现的不可控。

LARS引入以下局部层级学习速率，应用该层的梯度并对梯度进行归一化。

这个新方程允许LARS在最陡的梯度下降方向上移动参数，其幅度与 

ˡ

 成正比。这种归一化有望缓解一些梯度消失和爆炸的问题。下面的信任比作为学习率的一部分，我们可以帮助学习较慢的层使用更高的学习率，同时调节梯度爆炸的变化。

另外，对于模型的正则化，LARS在下面引入 




β



 来进行权重衰减。

这里是LARS的最终优化算法，为全局学习率加上动量SGD加上多项式学习率衰减。

LARS将Resnet-50训练推到32K的batch size，而不损失准确性。然而，该方法在BERT上的表现并不好。

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

LAMB在LARS中使用了相同的分层归一化概念，因此每层的学习率都是不同的。但是对于参数更新，它用AdamW代替了动量概念。

在LAMB中，权重和偏差（weights and biases）被认为是两个独立的层，因为两者有非常不同的信任值，因此应该以不同的学习率来对待。

LAMB将mini-batch大小增加至了32K。这是大批量训练的核心。现在我们准备详细讨论Nvidia gpu的BERT训练。

Part7BERT Datasets

如果你想自己尝试最初的BERT预训练，请仔细阅读这一部分。原始研究论文中的BERT预训练使用了图书语料库中的800万单词和英语维基百科中的2500万单词。Nvidia实现提供了下载预训练数据集的脚本。然而，提供BookCorpus数据的站点在下载了500多篇文章后会阻止您访问IP。因此您可以忽略BookCorpus数据集，或者使用一个小得多的BookCorpus。第一种方法需要对Nvidia脚本做一些简单的更改。如果您想收集BookCorpus，在文章的底部列出了可供选择的方法。

对于我来说，英文维基百科的下载无法用Nvidia脚本完成。因此，我手动下载了zip文件，并使用相同的脚本来处理它(详细信息稍后)。为了存储和准备数据集，可能需要600GB+存储空间。

Nvidia的实现还包含其他脚本。微调数据集和预先训练的谷歌的BERT模型(BERT-base和BERT-large)可以通过这些脚本轻松下载。

Part8显存溢出问题

显存是分级深度学习训练中的一个敏感问题。谷歌和Nvidia的实现都运行在有着顶级的tpu和gpu的强大主机上。有时由于资源限制您可能需要修改代码。Nvidia的实现起源于谷歌实现，然后对其进行了优化以减少GPU内存消耗和速度提升，特别是利用了Nvidia的GPU硬件和AMP。

让我们首先关注谷歌的实现。BERT论文中的所有微调都是在一个具有64GB内存的Cloud TPU上完成的。对于BERT论文中的大多数微调实验，BERT-large需要超过16GB的GPU内存。所有分配给GPU的batch在训练时必须一次性放入GPU内存中。例如，对于大小为64的batch，您不能将其拆分为两个，然后再合并结果。每次小批量训练后，会更新模型参数。听起来还不错，但我接下来会解释这个问题。

如下所示，对于12GB的GPU，批处理的最大大小是12。不幸的是，此批处理规模太小，无法产生BERT论文中相同的性能结果。很可能是因为梯度存在噪声使模型更难收敛。

15梯度累加（Gradient accumulation）

您可能会问，为什么我们不使用更多的小batch，并在模型更新之前结合结果。实际上，它被称为梯度累加，并且实现起来非常简单。下面是一些演示一般概念的示例。

现在，让我们来看看一下Nvidia的实现。它利用梯度累加和AMP(稍后讨论)来降低GPU内存需求。对于一个16GB的GPU，你可以通过运行batch-size为8以及累加步等于32的设置来训练BERT-large，使用长度为128字的序列，有效批处理大小为256。即，32个大小为8的小batch的结果可以被组合成一个256有效batch。

但它仍然需要足够的显存来训练至少一个样本。当运行序列长度设置为为384，对应GPU内存为11GB时，可能会OOM，如果要进一步减少内存需要更先进的技术。如果您对这些技术感兴趣，参考部分提供了部分链接。

在本文的其余部分，将介绍使用24GB的Titan RTX GPU来复制BERT的训练过程。

16自动混合精度 (AMP)

许多深度学习模型都是用32位精度浮点数进行训练的。混合精度在计算节点激活和梯度时使用16位精度。因此可以将内存消耗减少一半。但是实际上并不能那么做，因为我们仍然需要保留32位权重的副本以及其他数据。

在Nvidia的BERT实现中，混合精度可以通过在命令行中使用“use_fp16”标志自动打开，该标志只需在代码中打开一个环境变量。底层引擎将自动使用16位精度进行梯度计算。

if FLAGS.use_fp16:os.environ["TF_ENABLE_AUTO_MIXED_PRECISION_GRAPH_REWRITE"] = "1"

许多算法以计算复杂度为代价来节省内存。不过在Nvidia的演示中，由于AMP可以切换到更简单的数学运算，使用TensorCore可以使BERT速度提高3倍。

但你需要选择使用基于Nvidia架构的TensorCore。

此外，您还需要为BERT安装Nvidia的docker镜像。Nvidia的docker的设置非常简单。此外，本文简化了混合精度的讨论，更多的细节可以在这篇AMP文章中找到。

Part9NVLAMB

Nvidia的BERT实现与LAMB略有不同。下面还有一个额外的步骤(2):使用所有节点的梯度对梯度进行归一化。这种额外的归一化可能有助于训练不那么容易受到梯度的影响。

下图是引入这一额外步骤后loss的曲线。

Part10Nvidia的BERT实现

在接下来的章节中，我们将了解更多关于Nvidia实现的细节。对于BERT，对于输入序列长度分别为128 (phase 1)和512 (phase 2)， LAMB可以实现全局batch大小为64K和32K。对于单个GPU，我们需要64个小batch加上1024个累加steps。这样的设置可能需要几个月的时间来训练BERT。

2019年，Nvidia分别构建了DGX SuperPOD集群，分别拥有92台和64台DGX-2h机器，该集群可以在47分钟和67分钟内完成训练。

DGX-2每台售价约为40万美元。访问这种基础设施的一种可能性是通过云服务，比如微软Azure的NDv2实例，该实例拥有800 Nvidia V100张量核心gpu。然而，这取决于组织的使用和用例。例如，如果训练语料可能随着时间的推移而变化，或者需要多次训练，那么云解决方案的价格可能会增加。

在多节点系统上，与Adam优化器相比，LAMB可以扩展到1024个gpu，将训练速度提高17倍。

Part11软件设置

在使用Nvidia实现之前，您需要设置Nvidia docker环境。最终目标是安装Nvidia gpu加速容器(Docker镜像)。它需要Docker, Nvidia Docker, NGC容器。本文应该包含您可能需要的所有软件设置。或者您可以在这里按照Nvidia的指令操作。

Part12BERT的预训练、微调和推理

这个链接是Nvidia关于预训练、微调和推理BERT的简单指导说明。本文也会在这里介绍其中的重要步骤，并分享一些技巧。另外，本文同时也记录了遇到的一些问题，在文章的最后列出了一些解决方案供参考。

1. 克隆仓库

git clone https://github.com/NVIDIA/DeepLearningExamplescd DeepLearningExamples/TensorFlow/LanguageModeling/BERT

1. 构建BERT TensorFlow NGC容器。(本文在这里遇到了两个问题。如果遇到问题，请参阅问题部分。)

bash scripts/docker/build.sh

1. 下载并预处理数据集

bash scripts/data_download.sh

该脚本最终调用data/create_datasets_from_start.sh。本文建议将create_datasets_from_start.sh中的任务划分为6个单独的步骤，因为其中一些可能会失败或者需要很长时间才能完成。因此，注释掉其他代码并再次运行data_download.sh以执行所需的特定步骤。但是不要跳过此步骤，因为它们可能依赖于前面的步骤。在继续之前，还请检查下一节。

# (Step1)
# Downland bookscorpus
python3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action download --dataset bookscorpus
# (Step2)
# Download English Wikipedia
python3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action download --dataset wikicorpus_en
# (Step3)
# Download pre-trained model and datasets for the fine-tuning
python3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action download --dataset google_pretrained_weights  # Includes vocabpython3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action download --dataset squad
python3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action download --dataset "CoLA"
python3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action download --dataset "MRPC"
python3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action download --dataset "MNLI"
# (Step 4)
# Properly format the text files
python3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action text_formatting --dataset bookscorpus
python3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action text_formatting --dataset wikicorpus_en# (Step 5)
# Shard the text files (group wiki+books then shard)
python3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action sharding --dataset books_wiki_en_corpus# (Step 6)
# Create TFRecord files Phase 1
python3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action create_tfrecord_files --dataset books_wiki_en_corpus --max_seq_length 128 \--max_predictions_per_seq 20 --vocab_file ${BERT_PREP_WORKING_DIR}/download/google_pretrained_weights/uncased_L-24_H-1024_A-16/vocab.txt# Create TFRecord files Phase 2
python3 ${BERT_PREP_WORKING_DIR}/bertPrep.py --action create_tfrecord_files --dataset books_wiki_en_corpus --max_seq_length 512 \--max_predictions_per_seq 80 --vocab_file ${BERT_PREP_WORKING_DIR}/download/google_pretrained_weights/uncased_L-24_H-1024_A-16/vocab.txt

17下载和预处理

以下总结了你可能会遇到的问题:

• BookCorpus: 在下载500篇文章下载会被封禁。请参阅文章末尾的替代方案。

• English Wikipedia: 直接在浏览器下载最新的维基百科。在运行脚本之前，将其重命名为下面的文件。

DeepLearningExamples/TensorFlow/LanguageModeling/BERT/data/download/wikicorpus_en/wikicorpus_en.xml.bz2

• 谷歌预先训练的BERT-large和BERT-base模型:按原样运行脚本。

• 微调数据集:按原样运行脚本。

• 文本格式化:按原样运行脚本。该脚本忽略了下载文件中的头文件等不相关的文本，只提取所需的文本。

• 分片:将数据分成多个GPU训练所需的分片。对于单个CPU，可能需要几个小时才能完成该步骤。即使使用32GB的主机RAM，也会在分片数据集时耗尽内存(大约2800m个单词)。可以通过修改了代码减少TextSharding.py中所需的内存占用。

• 在Tensorflow中创建TFRecord文件以实现更快的数据处理。按原样运行脚本。对于单个CPU，可能需要一天才能完成。

Part13预训练和微调

首先需要在docker中启动NGC映像。

bash scripts/docker/launch.sh

一旦进入了docker，就可以运行训练脚本。

下面的例子将使用LAMB进行预训练，使用8个GPU，每个GPU使用64的mini-batch。此设置基于DGX-1:8张Nvidia V100, 32G内存。

bash scripts/run_pretraining_lamb.sh 64 8 8 7.5e-4 5e-4 fp16 true 8 2000 200 7820 100 128 512 large

但是，如果系统中有多个GPU，建议先运行nvidia-smi，以验证您将在运行应用程序时使用哪个GPU。

以下是训练前脚本的默认设置(针对DGX-1)。

对于本文使用的Titan RTX 24G内存，batch size设置为48。

bash scripts/run_pretraining_lamb.sh
<train_batch_size_phase1> <train_batch_size_phase2> <eval_batch_size>
<learning_rate_phase1> <learning_rate_phase2>
<precision> <use_xla>
<num_gpus>
<warmup_steps_phase1> <warmup_steps_phase2>
<train_steps> <save_checkpoint_steps>
<num_accumulation_phase1> <num_accumulation_steps_phase2> <bert_model>

对于GPU内存小于16G的用户，即使在预训练或微调中batch size设置为1也可能会遇到OOM。因此，大多数人使用更小的BERT-base模型。

对预训练好的BERT-large模型进行微调：

bash scripts/run_squad.sh 10 5e-6 fp16 true 10 384 128 large 1.1 data/download/google_pretrained_weights/uncased_L-24_H-1024_A-16/bert_model.ckpt 1.1

这里是使用SQuAD微调的做法：

bash scripts/run_squad.sh
<batch_size_per_gpu> <learning_rate_per_gpu>
<precision> <use_xla>
<num_gpus>
<seq_length> <doc_stride>
<bert_model> <squad_version> <checkpoint> <epochs>

对于Titan RTX GPU，一个epoch不到2小时，大约需要2个epoch来进行微调

同样，请参考Nvidia的README文件获取其他脚本和命令。

Part14准确率和速度

如下所示，通过增大4倍的batch的大小以及增加了2倍的gpu，获得了2倍的速度提升。

下图展示了预训练在V100 16G gpu上单节点的训练性能。

下图展示了预训练在V100 32G gpu上单节点的训练性能。

18多个节点训练

DGX1每个节点有8个gpu, DGX2H每个节点有16个gpu。

19微调

在DGX-2 32G上使用SQuAD微调的性能。

20推理时的表现

Tesla T4上SQuAD的推理性能(1x T4 16G)。

Part15BERT的实现

Google BERT

Nvidia BERT’s implementation optimized from Google’s implementation

Hugging Face

Part16参考

BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding

BERT For TensorFlow

Pretraining BERT with Layer-wise Adaptive Learning Rates

NVIDIA Container Support Matrix

Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour

Large batch training of convolution networks

Large batch optimization for Deep Learning: Training BERT in 76 minutes

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

Optimization Methods for Large-Scale Machine Learning

One weird trick for parallelizing convolutional neural networks

减少内存: Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups

Part17下载Book Corpus数据集

Gutenberg Dataset

Replicating the Toronto BookCorpus dataset — a write-up

Homemade BookCorpus

---

投稿或交流学习，备注：昵称-学校（公司）-方向，进入DL&NLP交流群。

方向有很多：机器学习、深度学习，python，情感分析、意见挖掘、句法分析、机器翻译、人机对话、知识图谱、语音识别等。

记得备注呦

点击上面卡片，关注我呀，每天推送AI技术干货~

整理不易，还望给个在看！