Google硬件工程师揭密TPU为何比CPU、GPU快30倍

在Google 发表 TPU 一年后，这款机器学习定制晶片的神秘面纱终于被揭开了。

5 日，Google 资深硬体工程师Norman Jouppi 刊文表示，Google 的专用机器学习晶片TPU 处理速度要比GPU 和CPU 快15-30 倍（和TPU 对比的是英特尔Haswell CPU 以及Nvidia Tesla K80 GPU），而在能效上，TPU 更是提升了30 到80 倍。

从这次发表的测试结果来看，TPU 似乎已经超出业界的预期，但是藏在这晶片背后的内部架构究竟有什么秘密？我们从Jouppi 之前发表的论文当中，可以找到答案。

据了解，早在4 年前，Google 内部就开始使用消耗大量计算资源的深度学习模型，这对CPU、GPU 组合而言是一个巨大的挑战，Google 深知如果利用现有硬体，他们将不得不将资料中心数量翻倍来支援这些复杂的运算任务。

所以Google 开始研发一种新的架构，Jouppi 称之为「下一个平台」。Jouppi 曾是MIPS 处理器的首席架构师之一，他开创了内存系统中的新技术。3 年前他加入Google 的时候，公司上下正在用CPU、GPU 混合架构上来进行深度学习的训练。

Jouppi 表示，Google 的硬体工程团队在转向定制ASIC 之前，早期还曾用FPGA 来解决廉价、高效和高性能推理的问题。但他指出，FPGA 的性能和每瓦性能相比ASIC 都有很大的差距。他解释，「TPU 可以像CPU 或GPU 一样可编译，它可以在不同的网路（卷积神经网路、LSTM 模型和大规模完全连接的模型）上执行CISC 指令，而不是为某个专用的神经网路模型设计的。」一言以蔽之，TPU 兼具了CPU 和ASIC 的优点，它不仅是可编译的，而且比CPU、GPU 和FPGA 拥有更高的效率和更低的能耗。

TPU 的内部架构

▲ 该图显示了TPU 上的内部结构，除了外挂的DDR3 记忆体，左侧是主机介面。指令从主机发送到队列中（没有循环）。这些触发控制逻辑可以根据指令多次运行相同的指令。

TPU 并非一款复杂的硬体，它看起来像是雷达应用的讯号处理引擎，而不是标准的X86 衍生架构。Jouppi 说，尽管它有众多的矩阵乘法单元，但是GPU 更精于浮点单元的协处理。另外，需要注意的是，TPU 没有任何储存的程序，它可以直接从主机发送指令。

TPU 上的DRAM 做为一个单元并行运行，因为需要获取更多的权重以馈送到矩阵乘法单元（算下来，吞吐量达到了64,000）。Jouppi 并没有提到是他们是如何缩放（systolic）数据流的，但他表示，使用主机体件加速器都将成为瓶颈。

▲ 256×256 阵列缩放数据流引擎，经过矩阵乘法积累后实现非线性输出。

从第二张图片可以看出，TPU 有两个记忆体单元，以及一个用于模型中参数的外部DDR3 DRAM。参数进来后，可从顶部加载到矩阵乘法单元中。同时，可以从左边加载触发（或从「神经元」输出）。那些以收缩的方式进入矩阵单元以产生矩阵乘法，它可以在每个周期中进行64,000 次累加。

毋庸置疑，Google 可能使用了一些新的技巧和技术来加快TPU 的性能和效率。例如，使用高频宽内存或混合3D 记忆体。然而，Google 的问题在于保持分布式硬体的一致性。

TPU 对比Haswell 处理器

在和英特尔「Haswell」Xeon E5 v3 处理器来的对比中，我们可以看到，TPU 各方面的表现都要强于前者。

在Google 的测试中，使用64 位浮点数学运算器的18 核心运行在2.3 GHz 的Haswell Xeon E5-2699 v3 处理器能够处理每秒1.3 TOPS 的运算，并提供51GB／秒的记忆体频宽；Haswell 晶片功耗为145 瓦，其系统（拥有256GB 记忆体）满载时消耗455 瓦特。

相比之下，TPU 使用8 位元整数数学运算器，拥有256GB 的主机记忆体以及32GB 的记忆体，能够达成34G B／秒的记忆体频宽，处理速度高达92 TOPS ，这比Haswell 提升了71 倍，此外，TPU 伺服器的热功率只有384 瓦。

除此之外，Google 还测试了CPU、GPU 和TPU 处理不同批量大小的每秒推断的吞吐量。

如上图所示，在小批量任务中（16），Haswell CPU 的响应时间接近7 毫秒，其每秒提供5,482 次推断（IPS），其可以实现的最大批量任务（64）每秒则可以完成13,194次推断，但其响应时间为21.3 毫秒。相比之下，TPU 可以做到批量大小为200，而响应时间低于7 毫秒，并提供225,000 个IPS 运行推理基准，是其峰值性能的80%，当批量大小为250，响应时间为10 毫秒。

不过需要注意的是，Google 所测试的Haswell Xeon 处理器似乎也不能完全说明问题，英特尔Broadwell Xeon E5 v4 处理器和最新的「Skylake」Xeon E5，每核心时钟（IPC）的指令比这款处理器提升了约5%。在Skylake 是28 核，而Haswell 为18 核，所以Xeon 的总体吞吐量可能会上升80%。当然，这样的提升与TPU 相比仍有差距。

最后需要强调的是，TPU 是一个推理晶片，它并非是要取代GPU，可以确定的是，TPU 与CPU 一起使用对训练分析更加有益。但对于CPU 制造商而言，如何研发出像ASIC 一样兼顾性能和能效的晶片是现在以及未来要做的。

可以确定的是，Google 已经逐渐在自己的数据中心部署TPU，但是上述测试数据只是理论结果，实际应用表现如何？我们还不得而知。