深入了解Google的第一个Tensor Processing Unit（TPU）

作者：

Kaz Sato（谷歌云Staff Developer Advocate）

Cliff Young（谷歌大脑软件工程师）

David Patterson（谷歌大脑杰出工程师）

谷歌搜索，街景，谷歌照片和谷歌翻译等Google服务都有一个共同点：他们都使用谷歌的张量处理单元（TPU）来加速他们在幕后的神经网络计算。

Google在印刷电路板上的第一张张力处理单元（TPU）（左）; 部署在Google数据中心的TPU（右侧）

我们去年宣布了TPU，并且最近对其性能和结构进行了详细研究。简而言之，我们发现TPU的性能比现在的CPU和GPU高出15-30倍，每瓦特性能高出30-80倍。这些优势帮助Google的许多服务以规模和成本实现了最先进的神经网络。在这篇文章中，我们将深入研究Google TPU内部的技术，并讨论它如何实现如此出色的性能。

TPU之路

尽管早在2006年，Google就开始考虑为神经网络构建专用集成电路（ASIC），但在2013年，情况变得非常紧迫。当时，我们意识到神经网络的快速增长的计算需求可能要求我们将我们运营的数据中心。

通常，ASIC开发需要几年的时间。然而，就TPU而言，我们在15个月内就设计，验证，构建并将处理器部署到我们的数据中心。TPU项目（也是MIPS处理器的主要架构师之一）的技术主管Norm Jouppi 以这种方式描述了冲刺：

“我们做了一个非常快的芯片设计。这真的非常了不起。我们开始发货第一块硅，没有修正错误或修改掩膜。考虑到我们在筹建芯片的时候雇佣了这个团队，然后聘请了RTL（电路设计）人员，并且急于聘请设计验证人员，这是忙碌的。

（从深入研究Google的TPU架构，下一个平台）

TPU ASIC采用28nm工艺制造，工作频率为700MHz，运行时功耗为40W。由于我们需要尽快将TPU部署到Google现有的服务器上，因此我们选择将处理器作为外置加速卡进行封装，以便插入SATA硬盘插槽进行安装。TPU通过PCIe Gen3 x16总线连接到主机，提供12.5GB / s的有效带宽。

用神经网络预测

为了理解我们为什么设计TPU的原因，让我们看看运行一个简单的神经网络的计算。

TensorFlow游乐场的双螺旋问题

TensorFlow Playground上的这个例子训练神经网络，根据训练数据集将数据点分类为蓝色或橙色。（参见这篇文章以学习更多关于这个例子。）运行一个训练的神经网络的过程，以标签分类数据或估计一些缺失或未来的价值被称为推断。为了推断，神经网络中的每个神经元进行以下计算：

• 将输入数据（x）与权重（w）相乘以表示信号强度
• 添加结果以将神经元的状态聚合为单个值
• 应用激活函数（f）（如ReLU，Sigmoid，tanh或其他）调节人造神经元的活动。

神经网络获取输入数据，将它们与权重矩阵相乘并应用激活函数

例如，如果有三个输入和两个具有完全连接的单层神经网络的神经元，则必须在权重和输入之间执行六次乘法运算，并将两组乘法运算加起来。这个乘法和加法的序列可以写成一个矩阵乘法。这个矩阵乘法的输出然后被激活函数进一步处理。即使在处理复杂得多的神经网络模型体系结构时，乘法矩阵通常是运行经过训练的模型中计算量最大的部分。

在生产规模上需要多少次乘法运算？2016年7月，我们调查了谷歌生产服务中的六个具有代表性的神经网络应用程序，并总结了每个神经网络架构中的权重总数。您可以在下表中看到结果。

(MLP: Multi Layer Perceptron, LSTM: Long Short-Term Memory, CNN: Convolutional Neural Network)

从表中可以看出，每个神经网络中权重的数量从500万到1亿不等。每个预测都需要很多步骤，将处理后的输入数据乘以权重矩阵并应用激活函数。

总的来说，这是一个大量的计算。作为第一次优化，我们不是使用普通的32位或16位浮点运算在CPU或GPU上执行所有这些数学运算，而是应用称为量化的技术，以允许我们使用整数运算。这使我们能够减少使用我们的神经网络模型进行有用预测所需的内存和计算资源总量。

量化神经网络

如果外面正在下雨，你可能不需要知道每秒钟有多少滴水滴 - 你只是想知道下雨还是下雨。同样，神经网络预测通常不需要使用32位或甚至16位数字进行浮点计算的精度。通过一些努力，您可以使用8位整数来计算神经网络预测，并保持适当的准确度。

量化是一种优化技术，它使用8位整数来近似预设的最小值和最大值之间的任意值。有关更多详细信息，请参阅如何使用TensorFlow量化神经网络。

在张量流中的量化

量化是降低神经网络预测成本的有力工具，相应减少内存使用量也非常重要，特别是对于移动和嵌入式部署。例如，当您将量化应用到流行的图像识别模型Inception时，它会从91MB压缩到23MB，大约是原始尺寸的四分之一。

能够使用整数而不是浮点运算大大减少了我们TPU的硬件占地面积和能耗。TPU包含65,536个8位整数乘法器。在云环境中广泛使用的流行的GPU包含数千个32位浮点乘法器。只要您能够满足8位应用的精度要求，即可以达到25倍或更多乘数。

RISC，CISC和TPU指令集

可编程性是TPU的另一个重要设计目标。TPU不是设计用于只运行一种类型的神经网络模型。相反，它的设计足够灵活，可以加速运行许多不同种类的神经网络模型所需的计算。

大多数现代CPU受到精简指令集计算机（RISC）设计风格的严重影响。利用RISC，重点是定义大多数应用程序常用的简单指令（例如，加载，存储，添加和乘法），然后尽可能快地执行这些指令。我们选择了复杂指令集计算机（CISC）风格作为TPU指令集的基础。一个CISC设计侧重于实现用每个指令执行更复杂任务的高级指令（例如多次计算乘加）。我们来看看TPU的框图。

TPU框图

TPU包含以下计算资源：

• 矩阵乘法单元（MXU）：矩阵运算的65,536个8位乘加单元
• 统一缓冲区（UB）：24MB的SRAM作为寄存器
• 激活单元（AU）：硬连线激活功能

为了控制MXU，UB和AU如何进行操作，我们定义了十几个专门为神经网络推理设计的高级指令。其中五个操作在下面突出显示。

该指令集侧重于前面提到的神经网络推断所需的主要数学运算：在输入数据和权重之间执行矩阵乘法并应用激活函数。

标准说：

“神经网络模型由各种大小的矩阵乘法组成 - 这就是形成一个完全连接的层，或者在CNN中，它往往是较小的矩阵乘法。这种架构就是要做这些事情 - 当你累积了所有的部分和并从累加器输出，一切都通过这个激活管道，非线性使得它成为一个神经网络，即使它主要是线性代数。 （从深入研究Google的TPU架构，下一个平台）

总之，TPU设计封装了神经网络计算的本质，并且可以为各种各样的神经网络模型编程。为了对其进行编程，我们创建了一个编译器和软件堆栈，将来自TensorFlow图形的API调用转换为TPU指令。

从TensorFlow到TPU：软件堆栈

矩阵乘子单元的并行处理

典型的RISC处理器提供简单计算的指令，例如乘法或增加数字。这些是所谓的标量处理器，因为它们用每条指令处理单个操作（=标量操作）。

即使CPU以千兆赫范围内的时钟速度运行，但仍然需要很长时间才能通过一系列标量操作来执行大型矩阵运算。改善这种大型矩阵运算的性能的一种有效且众所周知的方法是通过向量处理，其中同时在大量数据元素上同时执行相同的操作。CPU包含指令集扩展，如SSE和AVX，可以表达这种矢量操作。GPU的流式多处理器（SM）是有效的向量处理器，在单个GPU上有许多这样的SM死亡。具有矢量处理支持的机器可以在单个时钟周期内处理数百至数千个操作。

在TPU的情况下，Google将其MXU设计为矩阵处理器，在单个时钟周期内处理数十万个操作（=矩阵操作）。把它看作像一次打印一个字符的文件，一次打印一行，一次打印一页。

使用矩阵处理器就像一次打印整页文档，而不是逐个字符或逐行地打印文档

TPU的心脏：一个收缩阵列

为了实现这样的大规模矩阵处理器，MXU具有与典型的CPU和GPU（被称为收缩阵列）完全不同的架构。CPU被设计为几乎可以运行任何计算; 他们是通用电脑。为了实现这种通用性，CPU将值存储在寄存器中，并且程序告诉算术逻辑单元（ALU）哪些寄存器要读取，要执行的操作（例如加法，乘法或逻辑AND）以及将寄存器放入结果。一个程序由这些读取/操作/写入操作的序列组成。所有这些支持通用性（寄存器，ALU和程序控制）的功能在功耗和芯片面积上都有成本。

CPU和GPU经常花费精力在每个操作中访问多个寄存器。一个收缩阵列将多个ALU链接在一起，重新使用读取单个寄存器的结果。

但是，对于MXU，矩阵乘法将重复使用这两个输入作为产生输出的一部分。我们可以读取每个输入值一次，但将其用于许多不同的操作，而不将其存储回寄存器。电线仅连接空间上相邻的ALU，这使得它们更短，更节能。ALU只执行固定模式下的乘法和加法，这简化了它们的设计。

将权重矩阵与输入向量相乘得到一个收缩阵列

这种设计被称为收缩因为数据流经芯片在波浪中，让人联想到心脏泵血的方式。在执行矩阵乘法时，MXU中特定类型的脉动阵列针对功率和面积效率进行了优化，并且不适合通用计算。它做了一个工程权衡：限制寄存器，控制和操作的灵活性，以换取效率和更高的操作密度。

将输入矩阵乘以权重矩阵与收缩阵列

TPU矩阵乘法单元具有一个包含256×256 =总共65,536个ALU的并发阵列机制。这意味着TPU每个周期可以处理65,536个8位整数的乘加运算。由于TPU运行在700MHz，TPU可以在矩阵单元中计算65,536×700,000,000 = 46×1012的乘加运算或每秒92 Teraops（92×1012）。

我们来比较CPU，GPU和TPU之间每个周期的操作数量。

相比之下，没有向量扩展的典型RISC CPU只能对每条指令执行一次或两次算术运算，GPU每条指令可执行数千次运算。使用TPU，MatrixMultiply指令的单个周期可以调用数十万次操作。

在执行这种大规模矩阵乘法的过程中，所有中间结果都直接在64K ALU之间传递，没有任何内存访问，显着降低了功耗并提高了吞吐量。其结果是，基于CISC矩阵处理器设计提供杰出的性能的每瓦比：TPU提供了83X与当代的CPU相比更好比和29X比当代的GPU更好的比率。

性能/瓦特，相对于当代CPU和GPU（以对数为单位）（增量，加权平均值）

最小和确定性的设计

设计一个针对神经网络推断进行优化的新处理器的另一个显着优点是您可以成为您设计中的最终极简主义者。正如我们的TPU文件所述：

“与CPU和GPU相比，单线程TPU没有任何复杂的微架构功能，消耗晶体管和能量来改善平均情况，但不是99％的情况：没有高速缓存，分支预测，无序执行，多处理，推测预取，地址合并，多线程，上下文切换等等，极简主义是领域特定处理器的优点。“ （P8）

由于通用处理器（如CPU和GPU）必须在广泛的应用程序中提供良好的性能，因此已经发展了无数复杂的，面向性能的机制。作为副作用，这些处理器的行为可能难以预测，这使得难以保证对神经网络推断有一定的延迟限制。相比之下，TPU设计是严格的最小和确定性的，因为它一次只能运行一个任务：神经网络预测。您可以在TPU模具的平面图中看到其简单性。

TPU模具的平面图（黄色 =计算，蓝色 =数据，绿色 = I / O，红色 =控制）

如果您将其与CPU和GPU的平面图进行比较，您会注意到CPU和GPU的红色部分（控制逻辑）要大得多（因此更难设计），因为他们需要实现上述复杂的构造和机制。在TPU中，控制逻辑非常小，占据了芯片的2％以下。

更重要的是，尽管有更多的算术单元和大的片上存储器，TPU芯片是其他芯片尺寸的一半。由于芯片的成本是面积3的函数- 每个硅晶圆的芯片更小，而小芯片的产量更高，因为它们不太可能存在制造缺陷*芯片尺寸减半可将芯片成本降低大约8倍（2 3）。

使用TPU，我们可以很容易地估计出需要多少时间运行一个神经网络并进行预测。这使我们能够在芯片吞吐量接近峰值的情况下运行，同时对几乎所有的预测都保持严格的延迟限制。例如，尽管上述MLP0应用有7ms的限制，但TPU的吞吐量比当前的CPU和GPU高15-30倍。

吞吐量低于7毫秒的延迟限制（以对数为单位）（MLP0的响应为99％：CPU = 7.2毫秒，GPU = 6.7毫秒，TPU = 7.0毫秒）

我们使用神经网络预测来支持面向终端用户的产品和服务，并且每个人都知道，如果服务需要很长的响应时间，用户就会变得不耐烦。因此，对于MLP0应用程序，我们将99百分位预测延迟限制在7毫秒左右，以便基于TPU的Google服务提供始终如一的快速用户体验。以下是在六个神经网络应用程序的时延限制下，TPU与当代CPU和GPU之间的整体性能（每秒预测）比较。在最令人惊叹的情况下，TPU与CNN1应用的CPU相比提供了71倍的性能。

6个参考工作负载的CPU，GPU和TPU性能（对数级）

神经特定的体系结构

在这篇文章中，我们已经看到TPU出色表现的秘诀就是对神经网络推理的奉献。当我们决定把重点放在神经网络推理上时，量化选择，CISC指令集，矩阵处理器和最小设计都成为可能。Google对于TPU的投资充满信心，因为我们看到神经网络正在推动计算模式的转变，我们预计TPU将在未来几年成为快速，智能和负担得起的服务的重要组成部分。

Acknowledgement

Thanks to Zak Stone, Brennan Saeta, Minoru Natsutani and Alexandra Barrett for their insight and feedback on earlier drafts of this article.

* Hennessy, J.L. and Patterson, D.A., 2017. Computer architecture: a quantitative approach, 6th edition, Elsevier.