\

在IPU上训练BERT-Large详述

作者:

分享  

Share on weixin
Share on weibo
Share on linkedin

订阅

凭借其准确性和灵活性,BERT是当今使用最广泛的自然语言处理模型之一。它也是Graphcore客户需求最多的模型之一。

我们的工程师已为IPU系统实施并优化了BERT,以支持各种基于语言的应用程序,并且使用行业标准的机器学习训练方案展示了出色的吞吐量。

在本文中,我们将详细说明Graphcore的应用程序部门、软件部门和中国工程团队的工程师所采用的优化技术,他们使用这些技术对BERT-Large实施方案进行预训练和微调。

BERT的兴起

Transformer的双向编码器表示(BERT)是谷歌创建的基于transformer的语言表示模型。自2018年底发布以来,它的受欢迎程度有了大幅提高。

BERT最初在谷歌约10%的以英语进行的美国搜索查询中试用,如今已用于几乎所有的谷歌英语搜索

与搜索引擎一样,BERT及其衍生工具也已应用于其他领域,例如问题解答、基于内容的推荐、视频理解和蛋白质特征提取,这证明了它的多功能性。

与生成式预训练Transformer(Generative Pre-trained Transformer,GPT)等以前的transformer模型不同,BERT的设计目的是在所有transformer层中利用每个未标记的训练示例的左右语境,并在两个方向上有效地检索其表示形式。这种双向关注为各种下游任务提供了更大的灵活性。

BERT能够在预训练阶段使用庞大的、未标记的数据集(并且仅需少量标记数据即可在微调阶段获得最先进的准确性),从而使类似于BERT、基于transformer的大型语言模型很有吸引力。有鉴于此,对训练和微调这些大型神经网络语言模型的需求正在急剧增长。

除了诸如语言建模之类的大规模任务之外,超小型的下游任务也可以从这些方法中受益。例如,Microsoft Research Paraphrase Corpus(MRPC)仅有3600个训练样本,但是通过将模型大小从BERT-Base扩大到BERT-Large,其准确性从84.4%提高到了86.6%。BERT变体保持了相对于MRPC和许多类似的benchmark的领先地位

随着机器学习研究人员和工程师希望继续改善这些模型的任务性能,模型尺寸正变得越来越大。在BERT模型变体和其他一些最新模型(例如具有1750亿个参数的GPT-3)中都可以看到这一点。

在IPU-POD上预训练BERT

使大规模语言建模成为可能的并不仅仅是建模方面的进步。开发能够提供更高效率的新AI硬件和系统也使我们能够在合理的时间范围内训练这些大规模模型,从而有可能在几天甚至几小时内完成对数十亿个样本的预训练。Graphcore的IPU‑POD系统解决了这一性能挑战,并大大提高了研究人员和工程师的工作效率。IPU-POD系统利用极高性能的处理器内存储提供出色的计算性能,并通过最大程度地减少数据移动实现更优的能效。高速横向扩展互连和智能存储管理使应用程序可以有效地扩展到数百个IPU。

为了在IPU‑POD上高效运行BERT,我们将整个模型的参数加载到IPU上。为此,我们在四个IPU之间拆分或“切分”了BERT模型,并在训练过程中将该模型作为流水线执行。

下图中您可以看到有关如何划分BERT-Large的示例。IPU 0包含嵌入层和投影/损失层以及三个编码器层,其余的21个层均匀分布在其他三个IPU上。由于嵌入层和投影层共享参数,因此我们可以将投影层、遮掩语言模型(Masked Language Mode,MLM)层和下一句预测(Next Sentence Prediction,NSP)层放回到IPU 0上。

在IPU‑POD4上的BERT-Large的模型并行性

为了帮助减少片上的存储占用,我们使用重新计算。这意味着我们无需为了在计算后向传递时使用而存储中间层激活。重新计算是一种可以在训练模型时采用的策略,并且在实施流水线策略时特别有用。由于在流水线中我们随时有多个批次处于“进行中”,如果不使用重新计算,存储的激活量的规模可能会很大。

预训练系统的优化程序状态存储在流存储中,并在优化程序步骤中按需加载。

流存储是我们在Graphcore上描述片外存储的方式。IPU-POD64系统中的每个核心“构建块” IPU-M2000平台均具有多达450GB的存储,可通过其四个IPU进行寻址。它分为每个IPU芯片中包含的900MB的处理器内存储,以及每个IPU片外流存储中的112GB。IPU-Machine上的DDR4 DIMM支持该流存储。

TensorFlow、PyTorch和PopART的Graphcore BERT实施

在我们的GitHub应用示例上的Graphcore示例中,您可以找到在TensorFlow、PyTorch和PopART中针对IPU的BERT实现。

我们的TensorFlow实施与原始的Google BERT实施共享模型代码,并为了利用Graphcore的TensorFlow流水线API,进行了自定义和扩展。

我们的PyTorch实施基于Hugging Face transformer库中的模型描述和实用程序。面向IPU,我们使用了Graphcore PopTorch库,包括流水线执行、重新计算和多副本/数据并行性。

我们的Github上也提供了使用Poplar高级运行时(Poplar Advanced Runtime,PopART)的BERT实施。PopART让您可以从ONNX模型描述中导入或创建模型,以进行训练和推理,并且同时包含C ++ API和Python API。PopART支持本文所述的优化程序、梯度和参数划分。我们将它们统称为复制张量切分(Replicated Tensor Sharding,RTS)。对于PopART中复制的流水线模型并行BERT系统,我们使用优化程序和梯度划分。

将BERT扩展为IPU-POD

无监督的预训练使BERT可以利用Wikipedia、BookCorpus和其他来源中数十亿的训练样本。即使使用IPU-POD4,在这些大型数据集上进行多次传递也要花费大量时间。为了减少训练时间,我们使用数据并行模型训练,将预训练扩展到IPU‑POD16IPU-POD64以及更多。

数据并行训练涉及将训练数据集分成多个部分,每个部分都由一个模型副本使用。在每个优化步骤中,所有副本之间的梯度都是均值缩减的,因此权重更新和模型状态在所有副本之间都是相同的。

集体梯度缩减

梯度缩减使用了Graphcore通信库(GCL),它是Poplar软件栈的扩展,旨在实现高性能的横向扩展。GCL包括高效的、基于环的全缩减,以及许多其他常见的通信原语。如下图所示,梯度在所有副本之间进行平均,并且在所有的多副本梯度完全缩减后应用权重更新。

可以使用优化程序wrapper功能或自定义优化程序定义中的更精细控制来自动添加数据并行梯度缩减。这种缩减可以在一个IPU-POD内部或不同IPU-POD之间无缝进行。

有关API在TensorFlow中工作方式的更多详细信息,您可以在TensorFlow用户指南中阅读有关复制计算图的信息。

对于PyTorch,我们在PopTorch IPU模型选项中设置了一个复制因数。

针对多副本BERT的数据并行性

训练大批次BERT

该模型的并行缩放效率的主要决定因素之一是梯度传递的频率。我们使用梯度累积(GA)将计算或微型批尺寸与全局批尺寸分离开来。这也使得我们能够在通过调整梯度累积因数来更改IPU和/或副本的数量时,得以保持恒定的全局批尺寸:全局批尺寸=副本批尺寸×副本数量,其中副本批尺寸=计算批尺寸×梯度累积因数

因此,对于固定数量的副本,较大的全局批尺寸可实现较高的梯度累积因数,并减少优化程序和通信步骤。

但是,如果GA太大(例如几千),则可能导致FP16下溢问题。在GA较小的情况下,这会由于气泡开销而导致流水线效率低下,如此处所述。可能需要进行一些实验才能找到最佳值。

在以下各节中,我们将回顾在训练BERT时所利用的学习率、预热和优化程序方案。

线性缩放规则

在关于使用SGD微型批次训练ImageNet的论文中,研究人员在训练ResNet-50和Mask R-CNN的全局批次时,使用了线性缩放规则。该规则确定当副本批尺寸η乘以k(其中k,通常是模型副本的数量)时,基本学习率η,应乘以k并设置为kη。

将全局批尺寸从n增加到nk,同时使用相同数量的训练时期并保持测试准确性,可以将总训练时间减少k倍,并显著缩短从模型到生产的时间。但是,随着全局批次的增加,任务性能显示出了下降。

逐步预热策略

缓解这种情况的一种方法是预热,其中学习率不会立即初始化为kη。相反,训练过程从零或任意小的学习率开始,并在预定义数量的预热步骤中线性增加,直至达到kη。这种逐步的预热可以减少大型全局批次训练的许多步骤,从而达到与较小批尺寸相似的训练准确性。在使用SGD微型批次训练ImageNet的论文中,这种做法使得全局批尺寸为8,000左右的训练成为可能。

在BERT-Large预训练情况下,阶段1和阶段2的预定义预热步骤不同。与BERT论文所述相同,阶段1使用训练数据的最大序列长度为128,阶段2的最大序列长度为384。阶段1的预热为2000个步骤,约占阶段1中整个训练步骤的30%。相比之下,阶段2中总共有2100个步骤,其中约13%的步骤是预热。

预热步骤可能需要针对不同的预训练数据集进行调整。

AdamW优化器

标准随机梯度下降算法对所有权重更新使用单一学习率,并在训练过程中保持恒定。相反,Adam(Adaptive Moment Estimation,自适应矩估计)利用梯度的第一矩和第二矩的移动指数平均值,并基于这些矩来调整学习率参数。

研究人员Loshchilov和Hutter在其ICLR 2019论文中发现L2正则化在Adam中无效。相反,他们提出了AdamW,它在权重更新步骤中应用了权重衰减正则化,而不是在损失函数采用L2正则化。这使权重可以成倍地衰减,而不是以加性常数衰减。他们证明了,对于CIFAR-10和ResNet32x32,具有热重启功能的AdamW在训练损失和泛化误差上的表现更优。

在IPU‑POD16 BERT预训练中,我们使用预先训练的全局批尺寸(范围从512到2560)对AdamW进行了实验,当在SQuAD下游任务上进行微调时,所有批尺寸都收敛到了参考准确性。

LAMB优化器

LAMB优化器(在此有更详细的描述)旨在克服由于批尺寸增长而引起的梯度不稳定性和损失差异,从而实现更大的批尺寸。LAMB使用与逐层自适应率缩放(Layer-wise Adaptive Rate Scaling,LARS)相同的逐层归一化概念,因此学习率对层保持敏感。但是,对于参数更新,它改而使用AdamW的动量和方差概念。
每一层的学习率按照以下方法计算:

其η是整体学习率,‖x‖是该层参数的规范,‖g‖是同一个AdamW优化程序更新的规范。

这意味着LAMB对更新进行了归一化,然后乘以‖x‖,以使其与每个层的参数具有相同的数量级,从而确保更新对每个层都进行了一些实质性的更改。然后将结果乘以总体学习率η。

在LAMB中,权重和偏差被认为是两个单独的层,因为它们的信任值非常不同,因此应以不同的学习率来对待。偏差和伽马以及批规范或组规范的beta经常被排除在层适应之外。

对于BERT,LAMB在阶段1启用的全局批尺寸最多为65,536,在阶段2启用的最大批尺寸为32,768。

低精度训练

在深度学习的早期发展中,许多模型是使用32位精度浮点算法训练的。使用较低精度很有吸引力,因为它可以同时实现更高的计算吞吐量:对于IPU而言,FP16的峰值性能是FP32的四倍。在较低精度下,张量的尺寸减小了两倍,从而降低了存储压力和通信成本。

但是,FP16精度浮点比FP32具有更低的精度和更低的动态范围。有些观点建议使用FP16来计算前向传递(激活)和后向传递(梯度),并使用损失定标来管理激活和梯度的动态范围,同时保留具有FP32权重梯度累积的主权重的FP32副本。

在FP16中使用激活和梯度进行训练时,我们以类似的方式实现损失定标。

Graphcore IPU具有使用随机舍入以及传统IEEE舍入模式的功能。随机舍入基于与上下舍入边界的近似值成正比的概率。在大量样本中,这会产生无偏差的四舍五入结果。

使用随机舍入可以使我们在整个训练过程中保持FP16的权重,而不会出现端到端训练或下游任务性能的准确性的可见损失。

计算优化程序中的一阶矩和二阶矩并将其存储在FP32中,归一化也是在FP32中进行。训练过程中的其余操作在FP16中进行计算。

训练结果

Graphcore最新的横向扩展系统显示出前所未有的BERT-Large训练效率。与基于DGX A100的同类系统相比,其训练时间缩短了2.6倍。

IPU-POD64[30]包含16个最新的IPU-M2000加速器,它利用计算、通信和存储技术的创新,在短得多的时间框架内提供与BERT-Large上领先的AI平台相同的准确性。下图中,我们使用TensorFlow和PyTorch的标准高级框架以及我们自己的基于PopART的实现来提供结果。这些是与NVIDIA已公布的最优PyTorch结果进行比较,并使用类似的方法论得出可比的训练时间结果。

预训练

下面的图表显示了我们的TensorFlow、PyTorch和PopART实施的预训练损失曲线,显示了等效的最终训练损失和非常相似的训练曲线的收敛性。我们还纳入了IPU-POD16和IPU-POD64上所有三个模型实现的训练吞吐量。

BERT-Large的预训练损失

表:预训练吞吐量

微调

在我们对BERT进行了充分的训练之后(可能需要数亿个训练样本),我们可以将这些预训练的权重用作初始权重,以用于具有较少标记数据的、特定于任务的微调训练过程。

BERT的预训练和微调

这种两阶段的设置在实践中得到了广泛使用,因为它具有以下优点:

  • 对于所有微调任务,我们只需对BERT模型进行一次预训练
  • 对这些经过预训练的模型进行微调,可以以较少的标记数据提供良好的任务性能,从而节省了大量的标记特定于任务的数据的工作

在IPU‑POD4或IPU‑POD16上,微调可以在几分钟或几小时内完成,具体取决于训练数据集的大小。许多微调训练可以在训练数据集上的少量传递后被终止。

SQuAD v1.1

斯坦福大学问答数据集(SQuAD[31])v1.1是一个大型阅读理解数据集,其中包含500多篇文章中的100,000多对问答。

下表显示了使用参考预训练权重和IPU预训练权重在IPU上针对SQuAD v1.1任务对BERT-Large进行微调时的准确性。如图所示,IPU的表现可以优于此任务的参考精度。

表:IPU上的SQuAD 1.1任务准确性

BERT-LargeExact match F1
Reference (Devlin et al. 2018)84.190.9
IPU-POD16 with IPU pre-trained weights84.590.9

中文理解评估

我们将要查看的下一组数据显示了使用谷歌预训练权重对IPU上的中文理解评估(Chinese Language Understanding Evaluation,CLUE)任务进行BERT-Base微调时的准确性。

CLUE分数是所有CLUE任务的测试准确性的平均值。每个任务的测试准确度是五个实验结果的平均值。如下表所示,IPU可以实现与其他AI平台(例如DGX-1 V100)相同的精度。

表:在GPU和IPU上的CLUE任务准确性

只需进行一些优化,AI从业者就可以使用IPU-POD系统大幅减少BERT-Large的训练时间,并且不会影响准确性。如前所见,在我们的新硬件架构上运行像BERT这样的大型模型比您想像的要简单得多。我们在这里描述的许多优化技术,您可能都不陌生。我们还提供了TensorFlowPyTorch实现,以便任何机器学习从业者都能对应该如何使用这些框架在IPU上运行BERT有所了解。

随着研究人员和工程师越来越多地寻求使用诸如BERT之类的更大型的模型来实现最佳的任务性能,硬件效率将发挥关键作用。在本文中我们已经看到,通过使用IPU-POD系统,创新者不仅可以获得BERT-Large和类似模型在准确性上的优势,而且还可以从显著的性能提升中受益。IPU处理器架构的独特特性使加速当今最先进的模型成为可能,同时使我们能够重新想象如何在熟悉的软件环境中开发未来的最先进模型。

向引领本研究工作的工程师们致敬:

James Briggs、Oskar Bunyan、Lorenzo Cevolani、Arjun Chandra、Nic Couronneau、Lakshmi Krishnan、李国颖、Visu Loganathan、Sam Maddrell-Mander、田至伟、Sylvain Viguier、王献、温锡怀、夏承舜、赵晗。

除了上述诸位,还要感谢Graphcore的BERT应用程序、PopART、TensorFlow、PyTorch、GCL和Poplar团队为IPU系统实施和优化BERT所提供的支持。

More Posts

Graphcore开源Argon Streams AV1视频解码器验证工具

没见过雪,神经网络如何模拟降雪?

Paperspace和Graphcore携手为开发者提供免费IPU

深入探讨:Hugging Face Optimum Graphcore上的视觉transformer

Paperspace上的IPU入门

Pienso和Graphcore携手,以更深入、更快速的人工智能洞见为企业赋能

获取最新的GRAPHCORE资讯

在下方注册以获取最新的资讯和更新:




    获取最新的GRAPHCORE资讯

    在下方注册以获取最新的资讯和更新: