【Hot Chips 29】浅谈 Google 的 TPU

作者:    2020-06-12 01:31:46   835 人阅读  528 条评论
【Hot Chips 29】浅谈 Google 的 TPU

长期关心工人智慧……人工智慧与深度学习的技术宅们,不可能未曾听闻 Google 自行打造、为 TensorFlow 机器学习框架量身订做,让 AlphaGo 在世纪人机围棋大战,打败李世乭九段的祕密武器:人工智慧晶片 TPU(Tensor Processing Unit)。在 Hot Chips 29,Google 也以「比较平易近人」的简报形式(年初发表的是论文),介绍第一世代 TPU 的技术全貌与第二代的概观。

大哉问:Why TPU,因为非「专职专业」不可

开宗明义:过去提升效能的诸多手段,如一次处理更宽的资料、指令管线化、同时执行更多的指令、更高的运行时脉等,在近 40 年来,让处理器变快了 100 万倍。

【Hot Chips 29】浅谈 Google 的 TPU

但包含 CISC 到 RISC 的指令集改革,这些手段和一堆法则,现在都遭遇瓶颈了,现在还想突破性成长?可以说几乎没招了。

【Hot Chips 29】浅谈 Google 的 TPU

最后一句话开宗明义:有鉴于电晶体和功耗预算的限制,现在已经进入专用处理器的时代(有兴趣可以参考 IBM 正在高喊的 Accerlerated Computing)。

【Hot Chips 29】浅谈 Google 的 TPU

专注于「推论」的第一世代 TPU

人工智慧应用可大略区分成两个领域:「训练/学习」和「推论/预测」,在成长(Development)阶段,经过「训练」(Training)的类神经网路,「学习」(Learning)如何执行工作,接着进入生产(Production)阶段,类神经网路会被应用在「推论」(Inference)上,即对资料进行分类并「预测」(Prediction)出结果。

【Hot Chips 29】浅谈 Google 的 TPU

Google 初代 TPU 仅能应用于推论(其实应该可以用来做训练学习,可能是运算精度考量),能够兼备训练类神经网路,则是第二代 TPU 的故事了。Google 在 Hot Chips 29 的主题演讲有公开部分第二代 TPU 的资讯,让兼具 HBM 记忆体与浮点运算能力的 TPU,又变得更像 GPU。

有别于当时多数「训练」平台以 GPU 做为运算底层,而不得不使用浮点运算工作,Google TPU 则是整数运算晶片,以达成更快的速度、更少的功耗,与更小的晶片面积。后来 nVidia 与 AMD 的新型 GPU 也因「英雄所见略同」,扩增整数运算支援性,特别是 8 位元 16 位元短整数。

Google TPU 的技术精髓:古老的脉动阵列(Systolic Array)

Google 打造 TPU,会面对两个主要挑战:

低成本:架构需要尽其所能的简单,不仅要能降低成本,更要缩短产品开发时程,儘快部署到 Google 资料中心。

高效能:尤其致力着墨于免除传统电脑的「范纽曼瓶颈」(von Neumann bottleneck):记忆体跟不上运算。考量到第一个成本因素,像 GPU 如此昂贵暴力的记忆体子系统,是压根儿完全不考虑的选项。换言之,须消耗较少的记忆体频宽,却能实现更高的运算吞吐量。

所以理论存在于古老论文已久,让「运算多跑一会儿」的脉动阵列(或称为脉动管线,Systolic Pipeline),就在相隔 30 余年之后,变身为云端霸主初代人工智慧晶片的灵魂。

其概念也很简单,既然运算动作反覆存取记忆体浪费时间,何不就「串串乐」,让资料经过大量相同的处理单元(Processing Element,PE)?举个例子,我们想要做连续加法运算,要进行 6 次 +1,如其反覆存取记忆体 6 次,不如 6 个 +1 串在一起,不就毕其功于一役了?

【Hot Chips 29】浅谈 Google 的 TPU

但脉动阵列的限制也很明显:只适用于特定的运算,缺乏灵活性与泛用性。很幸运的,像普及于特徵辨识的卷积神经网路(Convolutional Neural Networks,CNN),就是非常适合的应用,在 Google TPU 内转化为一连串的 8 位元整数乘积运算,时脉仅 700MHz,理论尖峰效能就 92T,推论效率高达 GPU 的「10 倍」。

【Hot Chips 29】浅谈 Google 的 TPU

TPU 内的 256×256「Matrix」部分就是很标準的脉动阵列,而 TPU 内的区块名称,也溢出了满满的「人工智慧味」。

在脉动阵列内,「用来决定输入资讯重要性」的「权重」(Weight)由上而下流动,而「替类神经网路加入非线性因素,以解决更複杂问题」的「激活值」(Activation)由左向右传递,如此简单的结构,就足以同时达成指令流与资料流的管线化,并减轻记忆体子系统的负担。当然,天底下没有白吃的午餐,简单的硬体意味着複杂的软体,但这对 Google 绝对不是问题。

【Hot Chips 29】浅谈 Google 的 TPU

自己做想要的晶片总是有好处的

Google 在一台伺服器安装 4 张 TPU 运算卡。

【Hot Chips 29】浅谈 Google 的 TPU

与 Intel 18 核 Haswell 和 nVidia K80 规格比一比,无论晶粒面积和耗电量都少了一大截,而初代 TPU 还是比较落后的 28nm 製程。

【Hot Chips 29】浅谈 Google 的 TPU

Google 是依据以下组态,部署并比较同时期 CPU、GPU 与自家 TPU 的差异性。

【Hot Chips 29】浅谈 Google 的 TPU

经过人体实验证实,Google TPU 享有数十倍于同期 CPU 与 GPU 的效能功耗比。

更重要的是,此计画从开案、研发、晶片 Tape Out、量产,一路到上线部署,仅仅只花了「15 个月」(反观动不动就好几年的 x86 微架构),做到完美的 Time To Market,世界上恐怕没有其他企业,能比 Google 能透彻理解资料中心使用深度类神经网路进行推论的需求,採用简单硬体设计,充分发挥软体能力的策略,更是功不可没。

【Hot Chips 29】浅谈 Google 的 TPU

像 GPU 的第二代 TPU

但一个便当不够,你可以吃第二个,Google 第一代 TPU 只能做推论,所以就有可以做深度学习的第二代。我们就来瞧瞧乍看之下与 GPU「殊途同归」的第二代 TPU,有什幺有趣的地方。

【Hot Chips 29】浅谈 Google 的 TPU

机械学习究竟需要何种运算精度,一直是没有标準答案的问题,Google 基于自身的需求,让第二代 TPU 具备 32 位元单精度浮点运算的能力。原先第一代 TPU 的 256×256 Systolic Matrix,演化成两个 128×128 的运算核心(但牺牲掉部分精确度),并新增泛用性较高的纯量运算单元。

此外,运算精度激增也加重记忆体频宽需求,两个运算核心个别加挂 8GB HBM,提供较前代多出 20 倍的 600GB/s 理论频宽,让第二代 TPU 更近似新型 GPU。因 HBM 採用多晶片堆叠封装,更能精简电路板的设计与面积,提高布署密度。

至于第二代 TPU 是否保留大型化的晶片内(on-chip)记忆体,如前代多达 24MB 用来存放激活值的缓冲区,与 4MB 用于暂存 Matrix 计算结果,Google 尚未公布细节,就不得而知了。

【Hot Chips 29】浅谈 Google 的 TPU

总之,第二代 TPU 在 Google 资料中心的布署样貌就长这样,64 个 TPU Pod(四颗TPU)即可提供 11.5P 的浮点运算能量。

【Hot Chips 29】浅谈 Google 的 TPU

最后加映工商服务时间:第二代 TPU 即将上线 Google Cloud 替各位服务,而且宣称使用 TensorFlow 机械学习框架,你也不需要大改程式码,就可享受 TPU 的惊人效能。

【Hot Chips 29】浅谈 Google 的 TPU

在 2013 年底,Google 爆炸性公开其「规划部署已达 3 年」、开大规模商业化导入软体定义网路(SDN)之先河的 B4 资料中心广域网路,到现在不知不觉中,在众多新创企业与学术机关四处宣传其研发成果的当下,也一路遥遥领先,早已让自家打造的人工智慧晶片迈进到第二代并投入商业运转,让人不得不佩服云端服务霸主的研发能量与执行效率。微软、苹果、IBM、Amazon 等竞争对手,会不会尽速跟进如法炮製?我们可以等着,但唯一可以确定的是,以后销售「人工智慧晶片」可能不会是好做的生意。