B体育【新智元导读】在英伟达统治AI时代硬件的当下,谷歌的TPU却开辟了另一条道路。今天,小编带你了解第一块TPU的前世今生,同时揭开TPU架构的神秘面纱。
多数人可能都了解,在神经网络近70年的历史中,寒冬和泡沫交替出现,——事实上,藏在神经网络背后的专用硬件加速器(ASIC)也是如此。
神经网络和ASIC属于是「先有鸡还是先有蛋」的关系,大家都知道神经网络在CPU上效率不高B体育,但是,如果没有证据证明神经网络是有用的,公司凭什么要为神经网络投资开发专门的硬件?
在历史的周期演进中,构建专用硬件的项目一次又一次地被放弃,最终,作为外来者的GPU打破了僵局。
毫无疑问,我们现在正处于泡沫周期,在大模型的激励下,从初创公司到超大规模企业,都在构建建立自己的硬件加速器。
他们之中的大多数都可能会失败,但也肯定会有成功的设计成为未来计算领域的重要组成部分。
——而谷歌的张量处理单元(TPU)作为第一批大规模部署的加速器,肯定会成为幸存者之一。
为什么谷歌的TPU既有趣又重要?因为这是谷歌,TPU被切实应用于谷歌庞大的服务(搜索、Android、Chrome、Gmail、地图、Play 商店、YouTube、照片),用户超过10亿。此外,谷歌还拥有第三大公共云。
在英伟达的GPU主导了这个AI时代硬件供应的当下,谷歌的TPU是一个真正经历了时间和规模考验的竞品。
以下的内容,小编分成两部分:第一部分讲故事,关于第一个谷歌TPU的前世今生;第二部分讲技术,揭秘TPU的架构细节和性能。
机器学习帮助谷歌发掘信息的价值,从图像和语音识别到语言翻译,再到大型语言模型,当然也少不了谷歌的「摇钱树」——年入数十亿美元的广告业务。
于是,在公司工作的第一天,他出去从当地的一家电子商店买了一台GPU机器,把它插入网络,然后开始在GPU上训练神经网络。
最终,大家意识到需要自己需要更多的GPU,2014年,谷歌决定以约13000万美元的价格购买40,000个英伟达GPU。
深度学习所提供的能力可以大规模应用于谷歌的各种业务,然而,无论是在技术上还是在战略上,依赖英伟达GPU都不一定是最佳解决方案。
GPU不是ASIC,它不是为神经网络量身打造的,对于需要大规模应用的谷歌来说,相当于要付出很大的额外代价;同时,依赖单一供应商显然也是战略上的重大风险。
谷歌当时有两种选择:现场可编程门阵列(FPGA)和专用集成电路(ASIC)。
当时的实验证明,FPGA的性能打不过GPU,而定制ASIC可能在推理方面产生10倍的性价比优势。
开发ASIC的时间成本是一个重要的考量,如果落地周期太长,那么硬件本身也就失去了意义B体育。为此谷歌动用了一切可以快速访问的资源和现有知识。
谷歌迅速招募了一支才华横溢、经验丰富的团队,其中包括David Patterson,——伯克利RISC原始设计的开发者,也是RISC-V指令集架构开发的关键人物。
为了赶时间,谷歌没有去从头开发新的架构。幸运的是,在35年前,就有人为他们准备好了合适的架构。
systolic system是一个处理器网络,它有节奏地计算并通过系统传递数据......处理器的功能类似于心脏,每个处理器都会定期将数据泵入和泵出,每次都执行一些简短的计算,以便在网络中保持常规的数据流。
许多基本的矩阵计算可以在具有数组结构的脉动网络上优雅而有效地执行流水线。例如,六边形连接的处理器可以最佳地执行矩阵计算......这些脉动阵列享有简单而规则的通信路径,并且网络中使用的几乎所有处理器都是相同的。因此,使用VLSI技术可以廉价地构建基于脉动阵列的专用硬件设备。
有了工程师和架构之后,谷歌还与当时的LSI公司(现在是Broadcom的一部分)合作,帮助芯片落地。
另外,全新的芯片意味着需要全新的指令集架构、全新的编译器以及上层软件,这是一个庞大的工程。
2015年初,第一批TPU正式部署在谷歌的数据中心,此时,距离项目启动仅仅过去了15个月,看一下这个庞大的研发团队:
为了这个进度,谷歌也不得不做了很多取舍,包括使用较老的工艺(28nm),以及较低的时钟频率(700MHz).
2016年3月,使用了TPU的AlphaGo Zero击败了当时的世界围棋冠军李世石,举世震惊。
在上面的故事中,我们提到了TPU的原理:systolic arrays(脉动阵列),假设有如下的2*2矩阵乘法:
计算结果如下图所示。每个角上的方块代表一个乘法/累加单位 (MAC),可以执行乘法和加法运算。
在此图中,黄色值是从顶部和左侧输入矩阵的输入。浅蓝色值是存储的部分总和。深蓝色值是最终结果。
第二步,a12和b21加载到左上角的MAC中,相乘并添加到先前计算的结果中B体育。这一步得到了结果矩阵的左上角值。同时,b11被传输到右上角的MAC,乘以新加载的a21,并存储结果;a11被传输到左下角的MAC,乘以新加载的b12,并存储结果;
第三步,b21被传输到右上角的MAC,乘以新加载的值a22,结果被添加到以前存储的结果中;a12被传输到左下角的MAC,乘以新加载的b22,并将结果添加到先前存储的结果中。此时得到了结果矩阵的右上角和左下角值B体育。同时,a12和b21被传输到右下角的MAC,相乘并存储结果。
第四步,将a22和b22传输到右下角的MAC,相乘并将结果添加到先前存储的值中,从而得到结果矩阵的右下角值。
首先要注意的是,TPUv1依赖于通过PCIe(高速串行总线)接口与主机进行通信。它还可以直接访问自己的DDR3存储,
DDR3 DRAM / Weight FIFO:权重存储在通过DDR3-2133接口连接到TPU v1的DDR3 RAM芯片中。权重通过PCIe从主机的内存预加载,然后可以传输到权重FIFO存储器中,供矩阵乘法单元使用。
Accumulators:运算结果从脉动阵列的底部汇总到累加器中(内存)。
Unified Buffer / Systolic Data Setup:应用激活函数的结果存储在统一缓冲区存储器中,然后可以作为输入反馈到矩阵乘法单元,以计算下一层所需的值。
TPU v1使用CISC(复杂指令集)设计,只有大约20条指令。这些指令是由主机通过PCIe接口发送给TPU的,而不是从内存中获取的。
矩阵运算采用可变大小的B*256输入,将其乘以256x256恒定权重输入,生成B*256输出,所以需要B次流水线循环才能完成。
由于TPU v1的指令集非常简单,可以保证芯片用于解码和相关活动的开销非常低,只有2%的芯片区域专门用于控制。
而24%的芯片面积专用于矩阵乘法单元,29%用于存储输入和中间结果的统一缓冲区存储器B体育。
时至今日,这个「临时赶工」的ASIC,已经不断完善和壮大,在英伟达的统治之下,开辟了另一条道路。