David Patterson,Google杰出工程师、UC Berkeley荣誉退休教授、美国国家工程院、科学院院士、文理科学院“三院”院士。他是RISC(精简指令集核算机)、RAID(独立磁盘冗余阵列)和NOW(作业站网络)的缔造者,他与John Hennessy的著作《核算机体系结构:量化研讨方法》在业界久负盛名。

2017年,David Patterson参加Google TPU团队,2018年3月,他与John Hennessy共同取得图灵奖,2008年获ACM/IEEE Eckert-Mauchly 奖(被誉为核算机体系结构最高奖),2000年取得冯诺依曼奖章。

本文是他近期在加州大学伯克利分校的演讲,他共享了Google TPU近十年的发展进程以及心得体会,并论述了进步机器学习硬件能效对碳脚印的影响。OneFlow社区对此进行了编译。

作者|David Patterson

翻译|胡燕君、贾川、程浩源

1

一场由TPU引发的“地震”

2013年,Google AI担任人Jeff Dean通过核算后发现,假如有1亿安卓用户每天运用手机语音转文字服务3分钟,耗费的算力就已是Google一切数据中心总算力的两倍,况且全球安卓用户远不止1亿。

假如仅通过扩展数据中心规划来满足算力需求,不但耗时,并且本钱昂扬。因而,Google决议针对机器学习构建特定范畴核算架构(Domain-specific Architecture),期望将深度神经网络推理的整体具有本钱(TCO)下降至原来的十分之一。

TPU演进十年:Google的十大经验教训

于是,Google在2014年开端研制TPU,项目进展神速,仅15个月后TPU就可在Google数据中心部署运用,并且TPU的功用远超预期,它的每瓦功用是是GPU的30倍、CPU的80倍(数据源自论文:arxiv.org/ftp/arxiv/p… )。

2016年,在Google I/O开发者大会上,Google首席执行官Sundar Pichai对外发布了TPU这一突破性成果,他介绍道:

“通过Google云渠道,用户不但能够接触到Google内部运用的高功用软件,还能够运用Google内部开发的专用硬件。机器学习的核算规划巨大,因而Google研制了机器学习专用硬件,也便是‘张量处理单元(TPU)’。TPU的每瓦功用比市面上一切GPU和FPGA都高出一个数量级。用户能够通过Google云渠道体会TPU的优异功用。DeepMind研制的AlphaGo在与韩国棋手李世石的对战中运用的底层硬件便是TPU。”

希腊神话中,特洛伊战争的原因是两方争夺世界上最美的女人——海伦,后世诗人将海伦的美貌“令成千战舰为之起航”。我以为TPU就像海伦,它的呈现引起了“成千芯片与之竞逐”。

能够说,TPU的面世引发了硅谷的“地震”。TPU宣布诞生后,Intel耗资数十亿美元收购了多家芯片公司,阿里巴巴、Amazon等竞争对手纷繁开端研制相似产品。TPU从头唤起了人们对核算机架构的重视,后来的几年内,呈现了上百家相关草创企业,年均总融资额近20亿美元,各种别致的想法层出不穷。

五年后,Sundar Pichai又在2021年Google I/O开发者大会发布TPU v4:

“AI技能的进步有赖于核算基础设施的支撑,而TPU正是Google核算基础设施的重要部分。新一代TPU v4芯片的速度是v3的两倍多。Google用TPU集群构建出Pod超级核算机,单台TPU v4 Pod包含4096块v4芯片,每台Pod的芯片间互连带宽是其他互连技能的10倍,因而,TPU v4 Pod的算力可达1 ExaFLOP,即每秒执行10的18次方浮点运算,相当于1000万台笔记本电脑的总算力。”

TPU演进十年:Google的十大经验教训

上图展示了TPU的发展前史。其间,Google没有发布TPU v4i(TPU v4 lite)的相关细节。上一年Google宣布TPU v4i已在云服务上可用,也发表了一篇关于TPU v4i的论文(www.gwern.net/docs/ai/sca… )。

2

十年演进,十大经历

过往十年,咱们在ML核算架构的发展中汲取了十大经历。

TPU演进十年:Google的十大经验教训

其间,前五个都和ML模型自身有关,后五个则关乎硬件和架构。这些经历对深度学习以外的范畴也有学习意义。

经历一:DNN所需内存空间和算力迅速增加

TPU演进十年:Google的十大经验教训

咱们阅览近几年的论文后发现,推理模型所需的内存空间和算力均匀每年增加50%。因为芯片规划和部署至少各需求1年,投入实践运用并优化需求3年。可见,从一款芯片开端规划到生产周期结束的5年内,模型所需的内存空间和算力已增加到大约8倍。因而,在芯片规划之初就要将这种增加考虑在内。

TPU演进十年:Google的十大经验教训

练习模型的增加速度比推理模型更快。依据OpenAI的统计,2012-2019年,SOTA练习模型的算力需求年均增加10倍。备受重视的GPT-3模型的参数量更是从15亿(GPT-2)增加到1750亿,进步了100倍。

经历二:DNN作业负载跟着DNN突破不断演变

深度学习是一个日新月异的范畴。2016年,MLP(多层感知器)模型仍是干流,但到2020年,CNN、RNN和BERT等不同模型百家争鸣。BERT是一种全新的Transformer模型,诞生于2018年,短短两年后,四分之一以上的Google内部运用都在运用BERT模型,可见深度学习发展变化之快。因而,ML核算架构需求能够支撑多种模型。

经历三:DNN模型可优化

一般而言,核算机架构师只需懂硬件、体系结构、编译器,假如还懂操作系统则更好,但他们不需求懂运用。但是,构建针对特定范畴的架构则需求软硬件兼通。

TPU演进十年:Google的十大经验教训

对ML工程师而言,只要能够让模型跑得更好,他们十分愿意依据硬件/编译器改善DNN模型。毕竟DNN模型不像GCC编译器,后者已成为被广泛采纳的编译器规范,不会轻易依据硬件改动。

DNN模型之所以能够优化,部分原因是这些程序自身不算巨大,大约仅仅不计其数行PyTorch或TensorFlow代码,操作可行性较强。

Google的一篇论文介绍了一种模型优化技能Platform-aware AutoML,AutoML运用的方法称为“神经架构搜索(Neural Architecture Search)”,即机器主动在搜索空间中寻找更优的神经网络模型结构。在上述论文的比方中,经机器主动优化后的CNN1模型,在相同的硬件和编译器上可完成相同的准确率,而运算功用为原模型的1.6倍。

经历四:影响推理体会的是推迟,而非批次规划

TPU演进十年:Google的十大经验教训

一些关于模型推理优化的论文把要点放在数据批次规划(batch size)上,以为要把batch size设置为1才能使推迟降到最低。但是,通过MLPerf基准数据可见,Google的生产模型在batch size相当大的状况下也能完成低推迟,这或许是因为这些模型是根据TPU开发,因而愈加高效。

经历五:生产端推理需求多租户技能

TPU演进十年:Google的十大经验教训

DNN需求运用多租户技能(multi-tenancy)。不少深度学习论文的一个假设是同一时刻只需运转一个模型,但在实践运用中,有不少状况都需求在不同模型中切换。

比方,机器翻译触及各种语言对,就需求不同的模型;传统的软件开发需求用到一个主模型和多个实验模型;乃至有时因为对吞吐量和推迟有不同的侧重要求,就需求不同的batch size,然后需求不同的模型。

TPU演进十年:Google的十大经验教训

如上图所示,咱们收集了8个模型的基准数据,其间6个模型触及多租户。右方的柱状图展示了模型巨细(以MB核算)。赤色虚线表明单块芯片的最大SRAM,可见不少模型需求的内存远大于此,这意味着需求有存取速度极快的DRAM。部分芯片的规划思路是利用SRAM解决一切使命,但在多租户运用场景下,咱们以为这很难办到。

经历六:重要的是内存,而非浮点运管用

TPU演进十年:Google的十大经验教训

借用克林顿竞选总统时的标语——“重要的是经济,懂吗?”(OneFlow译注:其时美国正值经济萧条,克林顿将经济作为竞选讲演的重要话题,最终赢得推举),在此,我想说,“重要的是内存,不是浮点运管用(FLOPs),懂吗?”

现代微处理器最大的瓶颈是能耗,而不是芯片集成度。Yahoo!创始人Mark Horowitz在十多年前就发现,拜访片外DRAM的能耗是拜访片上DRAM的100倍,是算术运算的5000~10,000倍。因而,咱们期望能够通过增加浮点运算单元(FPU)来分摊内存拜访开支。根据Mark Horowitz的数据,芯片上的FPU数量被设置为10,000个左右。ML模型开发人员常常试图通过削减浮点运管用来优化模型,但其实削减内存拜访数才是更有效的方法。

TPU演进十年:Google的十大经验教训

TPU v1有65,000多个乘法单元,比GPU、CPU等硬件高出许多倍。尽管它的时钟频率较低,仅为700MHz,但因为其乘法单元数量巨大,且每个乘法单元可进行2个运算操作,因而TPU v1每秒可执行65,0002700M≈90 TeraOPS次操作。

上图右侧展示了运算时的首要作业循环。65,000多个乘法器组成矩阵乘法单元(Matrix Multiply Unit)。核算时,首要启动累加器(Accumulator),然后通过激活函数管道(Activation Pipeline)进行非线性函数运算。累加器和激活函数输出存储(Activation Storage)是两个首要功用单元之间的缓冲区。内存(DDR3)向矩阵乘法单元输入参数;最终,核算成果通过PCIe队列返回服务器。

因而,TPU v1中首要的数据活动如下图赤色箭头所示,此外的数据活动还包含DDR3向其间输入权重,以及核算输出成果发送至主机。

TPU v1运用了脉动阵列(systolic array),这一概念早在40年前就被提出,做法是以固定的时刻间隔使数据从不同方向流入阵列中的处理单元(cell),最终将数据累积,以完成大型矩阵乘法运算。因为70年代的芯片只要一个金属层,不能很好地完成互连,所以Kung和Leiserson提出“脉动阵列“以削减布线,简化衔接。

现代芯片有多达10个金属层,不存在这方面的问题,其最大难点是能耗,而脉动阵列的能效极高,运用脉动阵列能够使芯片容纳更多乘法单元,然后分摊内存拜访开支。

经历七:DSA既要专门优化,也要灵活

TPU演进十年:Google的十大经验教训

作为一种针对特定范畴的架构(DSA),TPU的难点在于既要进行针对性的优化,一起还须坚持一定的灵活性。Google在推出用于推理的TPU v1之后,决议霸占更难的问题——练习。

练习之所以比推理愈加杂乱,是因为练习的核算量更大,包含反向传达、转置和求导等运算。并且练习时需求将很多运算成果储存起来用于反向传达的核算,因而也需求更大的内存空间。

TPU v1只支撑INT8核算,对练习而言动态规模不够大,因而Google在TPU v2引入了一种的新的浮点格式BFloat16,用于机器学习核算。练习的并行化比推理的并行化更难。因为针对的是练习而非推理,所以TPU v2的可编程性也比TPU v1更高。

与TPU v1比较,TPU v2的改善分为5步。第一步,TPU v1有两个存储区域:Accumulator和Activation Storage,前者担任储存矩阵相乘成果,后者担任储存激活函数输出。

为了进步灵活性,TPU v2将上述两个相互独立的缓冲区调整位置后合并为向量存储区(Vector Memory),然后进步可编程性,这也更相似传统的内存区。

第二步改善针对的是激活函数管道(Activation Pipeline),TPU v1的管道内包含一组担任非线性激活函数运算的固定功用单元。TPU v2则将其改为可编程性更高的向量单元(Vector Unit),使其对编译器和编程人员而言更易用。

第三步,将矩阵乘法单元直接与向量存储区衔接,如此一来,矩阵乘法单元就成为向量单元的协处理器。这种结构对编译器和编程人员而言更友爱。

TPU演进十年:Google的十大经验教训

第四步,TPU v1运用DDR3内存,因为它针对的是推理,只需运用已有的权重,不需求生成权重。针对练习的TPU v2则不相同,练习时既要读取权重,也要写入权重,所以在v2中,咱们将本来的DDR3改为与向量存储区相连,这样就既能向其读取数据,又能向其写入数据。

TPU演进十年:Google的十大经验教训

然后,咱们将DDR3改为HBM。因为从DDR3读取参数速度太慢,影响功用,而HBM的读写速度快20倍。

第五步,咱们在HBM和向量存储区之间增加互连(Interconnect),用于TPU之间的衔接,组成咱们之前说到的Pod超级核算机。以上便是从TPU v1到TPU v2的改善。

经历八:半导体技能的发展速度良莠不齐

TPU演进十年:Google的十大经验教训

回顾过去能够发现,各类技能的发展速度并不同步。核算逻辑的进步速度很快,芯片布线的发展速度则较慢,而SRAM和HBM比DDR4和GDDR6的速度更快,能效更高。

TPU演进十年:Google的十大经验教训

上图虚线框内展示了单个Tensor Core运算单元。TPU v2中有两个互连的Tensor Core。

因为布线技能的进步相对滞后,假如仍像TPU v1相同,每块芯片只要一个Tensor Core,就会导致管道更为冗长,假如管道出了问题也会愈加费事。因而,咱们将两个Tensor Core相互衔接,这对编译器而言也更友爱。

Google做出TPU v2之后,期望再花一年时刻完善v2,所以TPU v3没有引进新技能,仅仅v2的改善版。

TPU演进十年:Google的十大经验教训

与v2比较,TPU v3有以下特点:

  • 体积只大了不到10%;
  • 矩阵乘法单元(MXU)的数量翻倍,因而峰值功用也翻倍;
  • 时钟频率加快了30%,进一步加快核算速度;
  • 内存带宽扩展了30%;
  • 容量翻倍,可使多种运用愈加便利;
  • 芯片间带宽扩展30%;
  • 可衔接的节点数是之前的4倍。

TPU演进十年:Google的十大经验教训

上图左上角即为TPU v1的主板。中心是v2,v2的散热方法是风冷,所以图中可见高高突起的风冷散热器。右上角是v3,v3的运转温度太高,所以只能选用液冷。左下角是TPU v2组成的Pod超级核算机,共有256张TPU,峰值功用为11 PFLOP/s;右侧的TPU v3 Pod有1024张TPU,峰值功用可达100 PFLOP/s(1 PFLOP/s即每秒1015次浮点运算)。

TPU演进十年:Google的十大经验教训

从TPU v3到TPU v4i,矩阵乘法单元的数量再次翻倍,但芯片面积却没有扩展。如前所述,核算逻辑的发展速度是最快的。

TPU演进十年:Google的十大经验教训

假如想了解TPU v4i,能够阅览论文《Ten Lessons From Three Generations Shaped Google’s TPUv4i》。TPU v4i中,单个Tensor Core有4个矩阵乘法单元,是v3的两倍,且v4i的片上内存更大。此外,TPU v4i的时钟频率加快了10%,矩阵乘法单元中运用四位加法器(4-input adder),能够大幅节约芯片面积和功耗。

功用计数器(Performance counter)的重要性不言而喻,Google在v4i的很多当地都放置了功用计数器,能够更好地帮忙编译器,并能更清楚地掌握运转状况。

功用计数器没有缓存,它们都在编译器操控的4D DMA (直接存储器拜访)之下,并且能够进行自定义互连。最终,为了操控更多的MPU(微处理器)和CMEM,VLIW(超长指令字)指令拓宽到将近400位。

经历九:编译器优化和ML兼容性十分重要

TPU演进十年:Google的十大经验教训

XLA(加速线性代数)编译器可对全程序进行剖析和优化,优化分为与机器无关的高档操作和与机器相关的低级操作,高档优化操作将影响TPU和GPU,所以一般咱们不会改动高档优化操作,以免导致失灵,但咱们能够改动低级优化操作。

XLA编译器能够处理多达4096个芯片的多核并行,2D向量和矩阵功用单元的数据级并行性,以及322~400位VLIW指令集的指令级并行。因为向量寄存器和核算单元是2D,这就要求功用单元和内存中有良好的数据布局。此外,因为没有缓存,所以编译器需求管理一切的内存传输。

最终的问题是,与CPU比较,DSA的软件栈还不够成熟。那么编译器优化最终能够提速多少?

TPU演进十年:Google的十大经验教训

实践上的提速相当可观。其间重要的优化之一称为算子交融(Operator Fusion),如将矩阵乘法与激活函数进行交融,省略将中心成果写入HBM再读取出来的过程。上图是咱们的MLPerf基准测试成果,可见,运用算子交融均匀能够带来超两倍的功用进步。

TPU演进十年:Google的十大经验教训

上图显示了编译器优化的提速效果。蓝色表明运用GPU,赤色表明运用TPU,通过短短十几个月的优化后,不少模型的功用都进步到了两倍。要知道,对C++编译器而言,假如能在一年内把功用进步5%-10%就现已十分了不得了。

此外,编译器的后向ML兼容性十分重要。我的搭档Luiz Barroso主管Google的一个与核算机架构无关的部分,他表明不期望在练习中花太多时刻,期望一晚上就能够练习好模型,第二天能够直接部署。咱们期望练习和推理时成果共同,这便是咱们说的后向ML兼容性,咱们还期望它能在一切TPU上运转,而不是每次更改TPU时都要从头练习。

TPU演进十年:Google的十大经验教训

为什么坚持后向ML兼容性如此困难?因为浮点加法不符合结合律,所以运算顺序或许会影响运算成果。而TPU的使命便是让一切机器对编译器而言都没有区别,以便能够在重组代码的一起取得相同的高质量成果,以完成后向ML兼容性。

经历十:优化的方针是Perf/TCO仍是Perf/CapEx

TPU演进十年:Google的十大经验教训

在将研讨成果运用到实践生产时,咱们优化的方针是什么?Google构建硬件是为了用在自己的数据中心,所以咱们所要操控的本钱是指整体具有本钱(TCO),包含本钱本钱(收购本钱)和运转本钱(电力、冷却、空间本钱)。资金本钱是一次性的,而运转本钱需求持续开销3~4年。

因而,芯片和主板生产商只需求考虑产品功用/本钱本钱的比率;而Google却要考虑整个硬件生命周期的本钱,重视功用/整体具有本钱之间的比率。如上面的饼状图所示,电力可占整体具有本钱的一半。所以,假如把眼光扩展到整体具有本钱上,在系统规划时就很或许会做出不同的取舍。

TPU演进十年:Google的十大经验教训

之前说到,TPU v1有一个Tensor Core,v2和v3有两个。到了v4时,根据对整体具有本钱的考虑,Google决议分开规划:用于练习的TPU v4有两个Tensor Core,用于推理的TPU v4i只要一个。这样就大大进步了功用和整体具有本钱之间的比率。

TPU演进十年:Google的十大经验教训

上图是TPU v4i和TPU v3的每瓦功用比照,赤色是TPU v4i,蓝色是TPU v3,前者的每瓦功用是后者的两倍以上。

3

进步机器学习能效,削减碳脚印

2021年10月的IEEE Spectrum杂志有一篇文章说到,练习某一模型需求数年时刻,花费1000亿美元,总碳排放量相当于纽约一个月的排碳量,假如还要进一步进步模型准确度,这些数字还会更夸张。

2022年1月,又有文章表明,依据当时的算力需求增加曲线估计,到2026年,练习最大AI模型的本钱将相当于美国的GDP,大概是20万亿美元。

TPU演进十年:Google的十大经验教训

Google研讨了ML硬件的动力耗费。不同于全生命周期耗费的动力(包含从芯片制造到数据中心构建的一切直接碳排放),咱们只重视硬件运转时的动力耗费。

展望未来,咱们有方法让机器学习的能耗下降到原来的100倍,碳排放量下降1000倍。咱们能够从四方面协同着手,极大地促进机器学习在更多范畴的可持续发展:

TPU演进十年:Google的十大经验教训

第一个要素是模型。Google在2017年发布Transformer模型,四年后,又开发了Primer模型,其核算质量相同,但能效更高。Primer的能耗和碳排放量比较Transformer下降了4倍。

第二个要素是硬件。2017年所运用的P100 GPU和当时最新TPU的功用相差了14倍。所以,前两个要素结合,能够将能耗和碳排放量下降60倍。

第三个要素是数据中心的能效。Google的PUE大约是其他数据中心的1.4倍。所以,前三个要素累积,能够将能耗和碳排放量下降80倍。

第四个要素是数据中心的地理位置。即便在同一个国家,不同区域运用无碳动力的份额也或许大不相同。在Google一切数据中心地点地中,俄克拉荷马州的无碳动力占比最高,Primer模型便是在此处练习的,这可将碳排放量在上述基础上下降9倍。

归纳上述四个要素,咱们可将机器学习的能耗下降80倍,碳排放量下降700倍。这十分了不得。

TPU演进十年:Google的十大经验教训

OpenAI的GPT-3面世后引起了Google一切机器学习工程师的注意,他们都卯足了劲想做得更好。18个月后,Google推出了GlaM模型,在相同的基准测试中它的体现比GPT-3更好。GlaM的参数是GPT-3的七倍,达到1.2万亿,但它的能耗却并不大,因为它利用了稀少性。GlaM是一个MoE模型(Mixture of Experts,专家混合模型),它平时只调用每个token中的8%的参数,而密集型模型会运用100%的参数。因而,GlaM中加速器的作业时长和能耗都下降了3倍。

最终,与GPT-3不同的是,GlaM在俄克拉荷马州运用清洁动力进行练习,因而累计下来,其碳排放量下降了14倍。所以GlaM的比方表明,比较V100 GPU,运用TPU v4既削减了碳排放量,并且核算质量更好。

(原视频链接:www.youtube.com/watch?v=PLK…

欢迎下载体会 OneFlow v0.8.0 最新版别: github.com/Oneflow-Inc…