TPU和GPU的带宽模型

前几天，谷歌AI团队发布了一款新的语言表征模型——BERT，是来自Transformer的双向编码器表征。它的效果很强大，在11项NLP任务中都刷新了最佳成绩。

https://arxiv.org/abs/1810.04805

在计算力方面，关于BERT选择TPU还是GPU仍然存在争议。BERT用了四天时间，在4个TPU pod上完成的（共有256个TPU芯片），这是否意味着只有谷歌才能训练像BERT这样的模型呢？难道GPU已经走到尽头了吗？这里需要明确两点基础知识：

一台TPU是一个矩阵乘法单元，它仅可以进行矩阵乘法和矩阵操作。在计算矩阵乘法时，它的速度很快。

进行矩阵乘法的过程中，最慢的部分就是从主记忆体中得到元素，并将其载入处理器中。

换句话说，矩阵相乘中，最烧钱的部分是内存负载。对BERT来说，矩阵相乘应该占计算负载的90%。了解了这些背景，我们可以对这一问题进行小小的技术分析。

TPU和GPU的带宽模型

TPU上的Transformer

BERT中常见的操作是矩阵乘法：AB=C，其中A的尺寸为256×256，B为1024×1024。TPU在对矩阵执行相乘的过程中，会将矩阵分解成更小的128×128矩阵。这就意味着我们需要对A加载16个128×128的tile，从B中加载64个tile。总共就是1664=1024个128×128的负载量。在16位的条件下，这就是32MB的数据。

现在我们进一步简化它。我们假设，在进行两个记忆负载时没有延迟时间，这也是有可能的，因为通常你可以在线程并行下隐藏内存访问延迟。简单地说，这意味着，当我们等待一个128×128的矩阵副本完成时，已经完成了下一个。这样一来，我们只需要等待第一个内存副本的完成，不用等待其他的。这就是GPU速度快的核心原因，以及为什么我们要在GPU中使用多个线程，无延迟的重叠内存传输与实际情况相差无几。使用了这种简化，我们现在可以直接使用内存带宽计算为矩阵乘法加载内存所需要的时间。如果我们查看TPU的带宽，就会放发现有600GB/s，所以我们需要5.2e-05秒来传输32MB的数据。

GPU上的Transformer

对于GPU，过程相同，只不过使用更小的tile和更多处理器。和TPU相似，我们同时使用两个负载来隐藏内存延迟。对GPU来说，16位数据的tile尺寸是96×96的。如果我们用一个V100 Tesla GPU，那么需要同时运行160个tile，还会有稍许延迟。与TPU相比，这意味着，和两个能处理128×128的矩阵的单元不同，GPU有160个单元（80个SM，160个线程块，每个线程块有两个96×96的矩阵）。这也能保证我们可以通过并行隐藏内存延迟。

重复以上计算过程，可以得到下面的结果：

对矩阵A，我们有33个96×96的tile；对矩阵B，我们有121个96×96的tile。总共需要33*121=3993次负载，数据总量为70MB。V100每秒运行速度为900GB，所以内存负载可能会花7.6r-05秒。所以，我们的模型判断，在这一场景下，一台GPU比一台TPU慢32% 。注意，对一台RTX 2080 Ti GPU来说，矩阵tile是一样的，但是内存带宽减少到了616GB/s，说明RTX 2080 Ti比TPU慢了54%。

注意，可用Tensor Core的TPU和GPU都能在一次运行中分别计算矩阵乘法tile，所以，就速度来说二者是差不多的，区别就在于内存是如何被载入的。

在GPU上BERT的训练时间

利用这一数据，用V100和RTX 2080 Ti构成的GPU群组，以及高速网络和好的并行算法（例如用微软的CNTK），我们能在64台GPU上（相当于四个TPU pod）、用5天多的时间或8天半的时间训练出BERT。在有八台GPU的设备上，使用任意软件和并行算法（PyTorch或者TensorFlow），我们训练BERT需要42天或者68天。对于标准的四个GPU的台式机，我们需要99天。

带宽模型的限制

带宽模型最大的限制就是这些计算是针对特定矩阵大小的，计算的难度在各种尺寸之间都不同。例如，如果你的batch size是128，那么GPU的速度会比TPU稍快一点。如果batch size小于128，GPU的速度就会快很多。增加矩阵B的尺寸会让TPU比GPU快得多。在BERT的原始论文中，研究人员是针对TPU进行矩阵尺寸的调整，如果你用GPU训练的话就不要参考了。

未来可能遇到的限制包括融合运算，TPU可以计算额外的操作，例如非线性激活函数或矩阵乘法中的偏差。这意味着，TPU不需要从较慢的全局内存中进行加载。另外，GPU也支持这种操作，但是英伟达还未将它们实现，所以GPU用户可能无法使用这一方法。所以，用户可能会遇到1.6%的速度变慢，例如，如果你应用一个非线性函数和一个偏见，那么TPU可能会比GPU快3.2%。

32位、16位和8位的区别

如果在32位情况下重复上面的计算过程，那么TPU会比之前快5.3倍。所以数据类型的尺寸在TPU和GPU的选择上还是很重要的。

TPU不支持8位训练，但是图灵GPU可以。所以，我们也可以看看8位矩阵乘法会有怎样的表现。我此前总结了8位模型的表现，发现并不难训练。如果我们重复上面的训练，用8位GPU训练，就会发现GPU比TPU快了3倍。在四台RTX 2080 Ti上的8位训练需要21天。

结语

TPU在训练类似BERT的模型时，比GPU快32%到54%。你也可以用40到70天的时间，在八台GPU上复现BERT。而用四台普通GPU，在16位的情况下训练出BERT需要99天，在8位情况下则需要21天。