百分点认知智能实验室出品：TPU与BERT的邂逅

编者按

预训练语言模型BERT是业界NLP领域的主流技术，百分点认知智能实验室在各种NLP任务中也广泛使用BERT技术，并进行持续的优化和改进。随着实验室队伍的不断壮大，NLP任务数量的快速增加，我们在提高实验效率并降低算力成本上开始遇到挑战。

经过探索，我们发现使用谷歌的TPU是一种可行的解决方案。通过实验室的研究和实践，相比GPU，TPU能显著提高BERT模型运算的效率，且大幅度降低所需要的成本。

1

实验室为什么要采用TPU

百分点认知智能实验室应用预训练语言模型的方法主要有两种。一种方式是通过互联网上开源的预训练语言模型，如BERT和XLNet等，并根据上游任务建立相应的模型进行finetune。另外一种方式，在遇到对效果要求很高的NLP任务上，实验室会根据目标语料进行预训练语言模型的增量或重新训练，通过新生成的预训练模型，实验室在文本分类、信息抽取、问题等价性等方面均取得了重要的进展。

然而，预训练的语言模型也带来了一定的问题，在传统的word embedding加LSTM或CNN为主的模型训练当中，模型的参数很少会超过千万级，而这样的模型只用一块或几块GPU即可在短时间内完成训练。相比之下，base版本的BERT模型有1.5亿个参数，而large版本的则超过3亿，这样的模型即使仅仅根据任务finetune也可能需要数日，而重新训练则更是需要以月为单位计算。

计算成本开销的增大，不但增加了应用该模型的困难，同时也为NLP算法工程师的工作方式带来了一定的挑战。由于使用预训练语言模型需要经过大量的调参，所以这使得算法工程师的工作面临着两难的决策。

一方面，如果采用串行的方法进行调参，则NLP算法工程师需要花费大量的时间进行等待，从而严重降低了其单位时间的产出。另一方面，如果采用并行的方式进行调参，则消耗的算力会显著的增加。因此，在预训练语言模型已经改变NLP领域的时代，如何探索出一种高效省钱的NLP算法工作流程，将是NLP算法研究员和工程师面临的普遍问题。

值得注意的是，这个问题并非是单纯的增加GPU数量就可以解决的，其原因在于在一些训练任务中，并非增加batchsize也可以达到同样的效果。因此在理想状况下，NLPer需要的是在短期内可以显著减少预训练语言模型计算时间的高效芯片，而不仅仅是大量的简单的堆叠现有的芯片。

为了尝试解决这个问题，百分点认知智能实验室在国内首先与谷歌合作，对其提供的TPU进行比较全面的测评，并取得了积极的效果。由于TPU的中文文档和用例相对匮乏，我们首先对TPU进行简单的介绍。

2

关于TPU的简单介绍

TPU 即谷歌推出的 Tensor ProcessingUnit，一个只专注于神经网络计算的处理器。TPU的主要功能是矩阵运算，是一个单线程芯片，不需要考虑缓存、分支预测、多道处理等问题的处理器，其可以在单个时钟周期内处理数十万次矩阵运算，而GPU在单个时钟周期内只可以处理数百到数千的计算。

此外，GPU作为“图像处理器”，其基本的硬件架构仍然要能够支持图像相关运算中所存在的各种需求，这意味着GPU多少受到其固有架构的桎梏，这些机制的存在使得预测GPU的处理性能和利用率变得比较困难，故在现实使用中，使GPU达到理论计算峰值是比较困难的。

而相比于GPU，TPU的实际运算表现要比GPU更加稳定，除此之外，TPU对于数据的读取也进行了大量的优化，所以当所涉及的运算包含大量的矩阵计算时，TPU在训练的优势就显现出来了。

因为BERT模型是multi-head dot-product attention，所以从数学上而言，BERT不论是前向传播还以后向传播计算都包含了大量矩阵运算，这使得TPU天生就非常适合计算BERT等神经网络。

当然任何理论分析都是没有Benchmark有说服力的。那么TPU效果究竟如何呢？下面我们来展现一下百分点认知智能实验室对TPU的测评结果。

3

实验室使用TPU的结果

如上文所述，TPU在训练BERT 模型中到底有怎样的优势呢？我们分别就重新预训练语言模型和finetune语言模型的任务进行了比对。

实验室采用了500M的文本语料，并根据主流的BERT BASE版模型的参数要求：首先将文本数据mask十遍产生预训练数据，然后采用了序列长度为128、12层、768维等参数进行500k步训练。

本实验室使用TPUV2-8进行计算并与主流GPU从运算时间和花费上进行对比，结果如下：

实验室分别使用了TPU和GPU对BERT模型进行500K步训练，通过与主流的GPU Tesla V100*8 对比发现，使用GPU进行训练花费了大约7天时间，而是用TPU进行训练仅需要了1.2天即可完成。同时，在总费用成本上也是大量的缩减。

TPU在BERT预训练模型的运算时间和总成本上完胜了当前的主流的GPU。这里我们采用了已经训练好的BERT模型在10万条数据上进行十轮的finetune再次进行对比，结果如下：

实验室现在使用TPU作为加速硬件的方法对BERT模型进行十轮的fintune，如上表所示，通过和现在主流的GPUTesla V100*8进行对比发现：TPU完成fintune仅仅需要约10min的时间，而GPU完成同样的finetune 需要超过一个小时的时间，这项技术大大提升了实验室在NLP领域进行神经网络计算的效率，而且从总花费的角度上来看，使用TPU的成本只有GPU的3.5%左右。

因此，TPU的超高效率和低廉价格将神经网络计算变得更加“亲民”了，TPU可以从根本解决了中小公司算力要求高但经费不足的顾虑，曾经那种需要几十台GPU几天时间的BERT预训练由一个TPU一天可以轻松解决。这让所有的中小型企业也可以拥有之前所缺少的强大算力。

4

使用TPU的经验总结

下面是实验室对使用TPU的一个经验总结，希望大家可以借此排坑。

我们以租用TPU V2八核心为例，系统的说明一下创建虚拟机实例和TPU实例的方法。进入到谷歌云的首页，其页面应类似于如下：

图一：谷歌云初始页面

首先需要创建一个VM实例，在选项中进行现存、内存数量、系统镜像等配置。

二：创建VM实例的界面

在该页面的选项中，有几个是比较重要的。

接下来创建TPU界面：

图三：创建TPU界面

在创建TPU的页面，有几个选项值得说明：

1.TPUtype一项中，会出现v2-8，v3-8，v3-32等选项的说明（注意不同的区域提供不同型号的TPU）。其中v2或v3为tpu的型号，-8或-32则为核心数量。最小的核心数量为8核心。在该模式下，我们可以选用抢占式的模式。而对于大于8核心的选项，则意味着TPUpod。该模式尚不支持抢占式（但是抢占式正在谷歌内部进行内测）。

2.IPaddress range一项中，如不涉及到部署，则可以填写10.1.x.0，其中x为大于101的数字（如102，103等）。值得注意的是，如果之前已有TPU填写了某范畴，而新创建的TPU的IP地址范畴和之也有重叠，则新创建的TPU会覆盖掉原先的实例。

3.Preemptibility一项为是否采用抢占式实例的选项。如前文所述，只有V2-8和V3-8两种型号支持创建抢占式实例。

如果以上选项均已设定完毕，则可以点击CREATE按钮创建TPU实例，然后就可以顺利运行TPU程序了。

5

实验室使用TPU的心得

我们在上文中介绍的VM实例和TPU实例的管理方式为众多方法中的一种，除去以上办法外，还可以通过命令行模型ctpuup等创建TPU实例，对此我们不再做详细介绍。下面我们重点结合实验室在此期间使用TPU的经验，和大家分享心得，帮助大家排雷。

首先大家需要注意的是TPU创建完毕并开始运行，即使没有实际的程序运行也会发生扣费。我们实验室建议在开启TPU之前，就先将代码在本地环境中调通，避免没有必要的费用流失。从实验室运行TPU的实战经历来看，我们建议使用TensorFlow的Estimator框架，因为我们只需在创建Estimator时将普通的Estimator改为TPU Estimator，即可使用TPU进行BERT神经网络预训练。这样大大减少了实验室的工作量。

在进行BERT模型训练过程中，batch_size的大小直接影响模型训练的性能，通过对google TPU了解得知，每个TPU包含了8个core，建议设置batch_size大小为8的倍数。同时google 推出了TPU 的pod模式以满足用户对更大算力的追求，实际上就是将多个TPU进行封装成一个整体供用户使用，比如V3-32，就是用4个V3的TPU，因此我们建议batch_size的大小能被整体core的数目整除即可，这样可以最大效率的利用TPU。

当然除此之外，在实战中我们还需要处理较大的文件在VM中的问题，因为这样会消耗大量硬盘资源以及增加运算成本，实验室用到了Buckets—一个价格相对亲民的存储方式来提升资源运用的效率，我们建议将较大的文件（如BERT文件的初始权重）存储在bucket当中，该步骤较为简单，故我们省略详细介绍，有需要可以进一步阅读https://cloud.google.com/storage/docs/creating-buckets。

6

实验室对TPU的建议

实验室在经过一段时间的使用后，TPU虽然在多个方面上完胜GPU，但是我们认为TPU还是有很多可改进的地方：

1. TPU的使用门槛很高，TPU自开发以来，拥有较少的代码示例和文档，官方提供的实例也不够完善，对于初学者不够友好。尤其由于TensorFlow静态图的本质，这使得基于TPU的Debug比较困难；

2. 究其根本TPU是围绕TensorFlow框架设计的硬件，实际使用过程中TPU硬件和TensorFlow版本具有较大的相互依赖性，大大减少了其可兼容性，使得使用其他人工智能框架的项目很难高效低成本地运用TPU进行运算；

3. 由于TensorFlow本身是基于静态图的，而TPU从本质上也只能支持静态图，这使得需要依赖于动态图的应用难以在TPU上运行。这类应用包括semantic parsing和图网络等。

当然在TPU新技术的引进下，实验室会迎来新的转机和工作模式。随着计算效率的提高，实验室的算法研究员可以将冗长的计算等待时间极大缩短，并且提升整体实验室的研究效率，增强实验项目的可控性。之前，研究员因为考虑到时间成本问题，会采用多个假设并行验证的工作方式。因为每个实验都有风险出现问题，并行实验会使得研究员无法估测每个验证的具体成功率，很有可能耗费大量算力后空手而归。如今，研究员可以将所有假设串联并一一快速验证，显著提高实验的效率，大大降低了项目的成本风险，增加可预测性。

整体上来说，实验室认为TPU结合BERT模型是一个不错的开始，大大减少了我们预训练模型的时间，显著提升了BERT模型整体运算的效率，大幅度降低了硬件资源的算力成本。在这方面实验室还会进一步研究和探索，让我们期待未来更多的实践成果。