详解谷歌第二代TPU：功耗性能究竟如何？巨头想用它干什么？

安妮 陈桦 编译自 The Next Platform 量子位 报道 | 公众号 QbitAI

在上周召开的Google I/O 大会上，谷歌正式公布了第二代TPU，又称Cloud TPU或TPU 2。但是，谷歌并没有详细介绍自己的新芯片，只展示了一些照片。

The Next Platform今天发布一篇文章，基于谷歌提供的图片和细节，带你深入了解谷歌的TPU2。量子位编译如下：

首先要说明的一点是，谷歌不太可能向公众出售TPU的芯片、主板或是服务器。目前看来，TPU2还是一个只供内部使用的产品。只有极少数人可以通过TensorFlow研究云（TRC）直接访问TPU2的硬件，因为这本身就是为研究人员设计的具有“高度选择性”的项目。研究人员可以分享他们发现的TPU2可加速代码类型。谷歌还推出了Google Compute Engine Cloud TPU Alpha项目，我们可以假定，这也是一个高度选择性的项目。

谷歌设计TPU2的首要目的，肯定是为了加速其面向消费者的核心软件（比如搜索、地图、语音识别和无人车的研究等项目）深度学习的工作量。我们对Google TRC的粗略解读是，Google想借此招募人员去研究适合TPU2的超网格的工作负载。

谷歌表示，TRC项目虽然开始规模较小，但之后会逐步扩大。如果Google的研究推广到一般的应用程序，那么其他人员也可以直接访问TPU2。那时Google将为其谷歌云平台的公共云增添一个TensorFlow硬件实例。

TPU2能有今天，离不开去年Google I/O大会上第一代TPU的贡献。一代TPU也是专为机器学习设计的特定目的芯片，还应用在了AlphaGo、搜索、翻译、相册等背后的机器学习模型中。 TPU通过两个PCI-E 3.0 x8边缘连接器连接协处理器（参见下面两张照片的左下角），总共有16 GB/s的双向带宽。TPU消耗功率高达40瓦，远高于PCI-E 电源规格，可为8位整数运算提供每秒92万亿次的运算，或为16位整数运算提供每秒23万亿次的运算。为了进行比较，Google声称，在半精度浮点数（FP16）情况下，TPU2可以达到每秒45万亿次的浮点运算。

TPU没有内置的调度功能，也不能被虚拟化。它是一个直接连接到服务器主板的简单矩阵乘法协处理器。

△ 谷歌的第一代TPU卡：A图没有散热器；B图有散热器

在主板处理能力或其PCI-E 吞吐量超负载前，Google从不会透露有多少TPU连接到一个服务器主板。协处理器只需要做一件事，它需要以任务设置和拆卸的形式，从主机处理器获取大量信息，并管理每个TPU数据的传输带宽。

Google已将其TPU2设计用于四机架机柜，并将其称为pod。机柜是相对于一组工作负载的标准机架配置（从半机架到多机架）。它为大型数据中心所有者提供更轻松廉价的购买、安装和部署流程。例如，Microsoft的Azure Stack标准半机架就是一个机柜。

四机架机柜大小主要取决与Google正在使用的铜缆类型和全速运行的最大铜线长度。下图显示了机柜的高层次组织。

我们首先注意到，Google通过两根电缆将每个TPU2板连接到一个服务器处理器板上。也可能是谷歌将每个TPU2板连接到两个不同的处理器板，但是，即使是谷歌也不希望混淆该拓扑结构的安装、编程和调度复杂性。如果在服务器主板和TPU2板之间存在一对一的连接，则要简单得多。

△ Google的TPU2机柜：A是CPU机架，B是TPU2机架，C是TPU2机架，D是CPU机架; 固体箱（蓝色）：不断电电源系统（UPS）; 虚线框（红色）是电源; 虚线框（绿色）是机架式网络交换机和机架式交换机顶部

谷歌展示了TPU2机柜的3张不同照片。在这3张照片中，配置和连线方式看起来一致。TPU2连线的颜色编码有助于比较这些照片。

△ 三个谷歌TPU2机柜

谷歌展示了TPU2电路板的顶视图，以及电路板前面板接口的近距离视图。TPU2电路板的所有4个象限共享同一电源分配系统。我们认为，这4个TPU2电路板象限也通过简单的网络开关共享同一网络连接。看起来，每个电路板象限都是一个独立的子系统，而除此以外4个子系统之间并没有相互连接。

△ TPU2板的俯视图:A是四个TPU2芯片和散热片;B是2个BlueLink 25GB / s电缆/ TPU2;C是两种全路径体系结构(OPA)电缆;D是电路板电源连接器，E很可能是一个网络开关

前面板连接看起来像是QSFP网络接口，但我从未在其他地方看到过。IBM BlueLink规范定义，对于最小25GB/s的配置（称作‘子链接’），在上下行每个方向上需要8个200Gb/s信道（总共16个信道）。谷歌是OpenCAPI的成员，同时也是OpenPowerFoundation的创始成员，因此使用BlueLink规范是合理的。

△ TPU2面板连接

前面板中央的两个接口看起来像是QSFP接口，介质为铜双绞线，而不是光纤。这支持两种网络配置，分别为10Gbps以太网和100Gbps英特尔OPA连接。两个100Gbps的OPA链路可以合并提供双向25GB/s的带宽，从而符合BlueLink规范要求的网速。因此我们认为，谷歌采用了100Gbps的OPA连接。

不过为了避免信号衰减带来问题，这些铜缆、BlueLink或OPA的线缆长度不能超过3米。这意味着，CPU和TPU2电路板之间的物理距离不能超过3米。谷歌使用彩色编码的线缆来连接，我猜测这是为了更方便地接线，避免出错。可以看到，在前面板最前方的接口下方，有贴纸与线缆颜色一一对应。我们认为，颜色编码表明，谷歌计划更大规模地部署这些TPU2机柜。

白色线缆最有可能是1Gbps以太网连接，这个网络用于系统管理。在照片中，我们并没有看到，谷歌如何将管理网络连接至TPU2电路板。不过，基于白色线缆的走线方式，我们可以假定，谷歌从机架背面将处理板连接至管理网络。或许，处理板将通过OPA连接管理TPU2板，并评估这些电路板的健康状况。

谷歌的TPU2机柜具有镜像对称性的特点。在下方图片中，我们将处理器机柜D的照片进行镜像翻转，并与处理器机柜A进行比较。这两个机柜看起来一模一样，仅仅只是互为镜像。在再下方的图片中，可以看到机柜B和C也互为镜像。

△ 比较两个TPU2机架

谷歌的照片中并没有展示足够多的连线信息，以判断电路板之间的准确网络拓扑结构。不过这很可能是一种非常复杂的mesh网络。

我们认为，CPU板是标准的英特尔至强双socket主板，符合谷歌的1.5英寸服务器机架单元的尺寸。这是当前一代主板设计。考虑到对OPA的支持，这可能是Skylake主板（参见以下对功耗的探讨）。我们猜测这是双socket主板的原因仅仅在于，我没有听说过，在英特尔供应链中，有哪家厂商曾大量发货单socket主板。不过，随着AMD推出“Naples”Epyc X86服务器芯片，以及高通推出Centriq ARM服务器芯片，突出单socket配置，这样的情况将发生改变。

我们认为，谷歌使用两个OPA线缆将每块CPU板连接至唯一的TPU2板，以实现25GB/s的总带宽。这种一对一连接回答了关于TPU2的一个关键问题：谷歌在设计TPU2机柜时，将TPU2芯片与至强socket数量配比为2:1。这就是说，4颗TPU2芯片对应于一个双socket至强服务器。

在深度学习任务中，GPU加速器采用的配比通常为4:1或6:1，而这种TPU2加速器与处理器之间的紧耦合与此有很大不同。2:1的配比表明，谷歌沿用了第一代TPU的设计哲学：“与GPU相比，TPU与FPU（浮点处理单元）之间关系更紧密。”处理器在谷歌TPU2架构中承担了大量工作，同时把所有矩阵运算任务丢给了TPU2。

在TPU2机柜中，我们看不到任何存储模块。或许这正是下图中机柜上方大量蓝色光纤存在的原因。数据中心网络连接至CPU板，同时没有任何光纤连接至机柜B和C，而TPU2板上也没有任何网络连接。

△ 很多光纤带宽连接到谷歌数据中心的其余部分

不管是TPU2还是CPU，每个机架上都有32个计算单位。每个机柜中有64个CPU板和64个TPU板，共有128个CPU芯片和256个TPU2芯片。

谷歌表示，其TRC包含1000个TPU2芯片，不过这个数字是去掉了零头的。四个机柜包含1024个TPU2芯片。因此，四个机柜是Google已经部署了多少TPU2芯片的下限。在Google I/O上公布的照片中，可以看到三个机柜，也可能是四个。

现在我们还不清楚一个机柜中的CPU和TPU2芯片如何关联，让TPU2芯片可以通过超网格中的连接有效地共享数据。我们几乎可以肯定，TRC不能跨机柜(256个TPU2芯片)处理单个任务。第一代TPU是一个简单的协处理器，因此CPU负责处理所有数据流量。在这种架构中，CPU通过数据中心网络访问远程存储器数据。

谷歌没有描述机柜的内存模型。TPU2芯片可以在OPA上使用远程直接存储器访问（RDMA）从处理器板上的内存中加载自己的数据吗？大概可以。

CPU板似乎也可能在机柜上执行相同操作，创建了大型共享内存池。该共享内存池不会像惠普企业版机器共享内存系统原型中的内存池那么快，但是有着25 GB/s的带宽，它速度不会太慢，而是在两位数太字节范围内（每个DIMM 16GB，每个处理器有8个DIMM，每个板有两个处理器，64个板产生16TB的内存）。

我们推测，在一个机柜上安排一个需要多个TPU2的任务看起来是这样：

处理器池应该有一个机柜的超网格拓扑图，哪些TPU2芯片可用于运行任务。

处理器组可能相互关联，对每个TPU2进行编程，以明确地链接TPU2芯片之间的网格。

每个处理器板将数据和指令，加载到其配对的TPU2板上的四个TPU2芯片上，包括网状互连的流量控制。

处理器在互连的TPU2芯片之间同步引导任务。

当任务完成时，处理器从TPU2芯片收集结果数据（该数据可能已经通过RDMA传输到全局存储器池中），并将TPU2芯片标记为可用于另一任务。

这种方法的优点是TPU2芯片不需要理解多任务，虚拟化或多租户，机柜上的所有这类运算都由CPU来处理。

这也意味着，如果Google想将云端TPU实例作为其谷歌云自定义机器类型IaaS的一种来提供，该实例将必须包括处理器和TPU2芯片。

目前我们还不清楚工作负载是否可以跨邮票进行缩放，并保留超级网格的低延迟和高吞吐量。虽然研究人员可以通过TRC访问1,024个TPU2芯片中的一些，但将计算量扩展到整个机柜看起来是一个挑战。研究人员或许能连接到多达256个TPU2芯片的集群，这足以令人印象深刻，因为云GPU连接目前才扩展到32个互连设备。

谷歌的第一代TPU运行时功耗40瓦，能以23 TOPS的速率执行16位整数矩阵乘法。TPU2的运行速度提高到45 TFLOPS，是上代的两倍，同时通过升级到16位浮点运算，提高了计算复杂度。一个粗略的经验法则表明，这样功耗至少翻两番：只是将运行速率提高一倍并升级到16位浮点运算，TPU2的功耗至少要提高到160瓦。

从散热器的尺寸来看，TPU2的功耗可能更高，甚至高于200瓦。

TPU2板上，在TPU2芯片顶部有巨大的散热片，它们是多年来我看到的最高的风冷散热片。同时，它们还具有内部密封循环的水冷系统。在下图中，我们将TPU2散热片与过去几个月看到的最大的散热片相比较：

△ A是4个芯片构成的TPU2板侧面图；B是搭载双IBM Power9的Zaius主板；C是双IBM Power8的Minsky主板；D是双英特尔至强的Facebook Yosemite主板；E是带有散热片的英伟达P100 SMX2模块和Facebook Big Basin主板。

这些散热器的尺寸都在高喊着“个个超过200瓦”。很容易看出，它们比上一代TPU上的40瓦散热器大得多。这些散热器的高度约能填满两个机架单元，接近3英寸。 （Google机架单元高度为1.5英寸，比行业标准1.75英寸U型稍矮）。

增加的功耗用在哪里了呢？

因此我们可以推测，TPU2芯片的内存容量也有所扩大，这有助于提高吞吐量，但同时也增加了功耗。

此外，Google从PCI-E插槽驱动单TPU（PCI-Express插槽向TPU卡供电）发展到单芯片TPU2板设计共享双OPA端口和交换机，以及每个TPU2芯片两个专用的BlueLink端口。OPA和BlueLink都增加了TPU2板级功耗。

Google的开放计算项目机架规格展示了功率为6千瓦，12千瓦和20千瓦的电力输送配置文件；20千瓦的功率分配可以实现带动90瓦的CPU。我们猜测，使用Skylake架构的Xeon处理器和处理大部分计算负载的TPU2芯片，机架A和D可能使用20千瓦电源。

机架B和C则不同。功率输送为30千瓦，能够为每个TPU2插槽提供200瓦的功率输送，每个机架36千瓦将为每个TPU2插座提供250瓦的功率输送。36千瓦是一种常见的高性能计算能力传输规范。我们相信，每芯片250瓦功耗也是Google为TPU2配置巨大散热器的唯一原因。因此，单个TPU2机柜的功率传输可能在100千瓦至112千瓦范围内，并且可能更接近较高数量。

这意味着TRC在满负荷运行时消耗将近50万瓦的功率。虽然四个机柜部署成本高昂，但却是一次性的资本费用，并不占用大量的数据中心空间。然而，用50万瓦的电力来持续资助学术研究，就算对Google这个规模的公司来说，也不是一笔小费用。如果TRC在一年内仍然运行，将表明Google正在认真为TPU2研究新用例。

TPU2机柜包含256个TPU2芯片。按每个TPU2芯片45 TFLOPS计算，每个机柜产生总共11.5 petaflops深度学习加速器的性能，计算它是16位浮点运算的峰值性能，也足以让人印象深刻。深度学习训练通常需要更高的精度，因此FP32矩阵乘法性能可能是FP16性能的四分之一，或者每个机柜约为2.9 petaflop，整个TRC是11.5 FP32 petaflops。

在峰值性能方面，这意味着在整个机柜上的FP16运算（不包括CPU性能贡献或位于机柜之外的存储），在每瓦100 - 115 gigaflops之间。

英特尔公布了双插槽Skylake生成Xeon核心计数和功耗配置后，可以计算Xeon处理器的FP16和FP32性能，并将其增加到每瓦特的总体性能。

关于Google的TPU2机柜，还没有足够的信息让我们能将其与英伟达Volta等商用产品进行比较。它们的架构差别太大，如果没有基准，根本无法进行比较。仅仅是比较FP16峰值性能，就像对比两台处理器、内存、硬盘、显卡都不一样的PC，却只考虑CPU主频一样。

也就是说，我们认为真正的比赛不在芯片级别。真正的挑战在于扩展这些加速器的使用范围。英伟达靠NVLink迈出了第一步，追求将自己的芯片独立于CPU。 英伟达正在将其软件基础架构和工作负载从单一GPU扩展到GPU集群。

在谷歌推出第一代TPU时，选择将其作为CPU的协处理器，到推出TPU2时，也只是将其扩展为处理器的2：1加速器。然而，TPU2超网格编程模型似乎还没有可以扩展的工作负载类型。 Google正在寻求第三方帮助，来寻找可使用TPU2架构扩展的工作负载。