关于深度学习:TPU演进十年Google的十大经验教训

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）升高至原来的十分之一。

于是，Google在2014年开始研发TPU，我的项目停顿神速，仅15个月后TPU就可在Google数据中心部署利用，而且TPU的性能远超预期，它的每瓦性能是是GPU的30倍、CPU的80倍（数据源自论文：https://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 v4i（TPU v4 lite）的相干细节。去年Google发表TPU v4i已在云服务上可用，也发表了一篇对于TPU v4i的论文
（https://www.gwern.net/docs/ai...）。

2

十年演进，十大教训

过往十年，咱们在ML计算架构的倒退中吸取了十大教训。

其中，前五个都和ML模型自身无关，后五个则关乎硬件和架构。这些教训对深度学习以外的畛域也有借鉴意义。

教训一：DNN所需内存空间和算力迅速增长

咱们浏览近几年的论文后发现，推理模型所需的内存空间和算力均匀每年增长50%。因为芯片设计和部署至多各须要1年，投入理论应用并优化须要3年。可见，从一款芯片开始设计到生产周期完结的5年内，模型所需的内存空间和算力已增长到大概8倍。因而，在芯片设计之初就要将这种增长思考在内。

训练模型的增长速度比推理模型更快。依据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模型可优化

通常而言，计算机架构师只需懂硬件、体系结构、编译器，如果还懂操作系统则更好，但他们不须要懂利用。然而，构建针对特定畛域的架构则须要软硬件兼通。

对ML工程师而言，只有能够让模型跑得更好，他们十分违心依据硬件/编译器改良DNN模型。毕竟DNN模型不像GCC编译器，后者已成为被宽泛驳回的编译器规范，不会轻易依据硬件改变。

DNN模型之所以能够优化，局部起因是这些程序自身不算宏大，大概只是成千上万行PyTorch或TensorFlow代码，操作可行性较强。

Google的一篇论文介绍了一种模型优化技术Platform-aware AutoML，AutoML应用的办法称为“神经架构搜寻（Neural Architecture Search）”，即机器主动在搜寻空间中寻找更优的神经网络模型构造。在上述论文的例子中，经机器主动优化后的CNN1模型，在雷同的硬件和编译器上可实现雷同的准确率，而运算性能为原模型的1.6倍。

教训四：影响推理体验的是提早，而非批次规模

一些对于模型推理优化的论文把重点放在数据批次规模（batch size）上，认为要把batch size设置为1能力使提早降到最低。然而，通过MLPerf基准数据可见，Google的生产模型在batch size相当大的状况下也能实现低提早，这可能是因为这些模型是基于TPU开发，因而更加高效。

教训五：生产端推理须要多租户技术

DNN须要应用多租户技术（multi-tenancy）。不少深度学习论文的一个假如是同一时间只需运行一个模型，但在理论利用中，有不少状况都须要在不同模型中切换。

比方，机器翻译波及各种语言对，就须要不同的模型；传统的软件开发须要用到一个主模型和多个试验模型；甚至有时因为对吞吐量和提早有不同的偏重要求，就须要不同的batch size，进而须要不同的模型。

如上图所示，咱们收集了8个模型的基准数据，其中6个模型波及多租户。右方的柱状图展现了模型大小（以MB计算）。红色虚线示意单块芯片的最大SRAM，可见不少模型须要的内存远大于此，这意味着须要有存取速度极快的DRAM。局部芯片的设计思路是利用SRAM解决所有工作，但在多租户利用场景下，咱们认为这很难办到。

教训六：重要的是内存，而非浮点运算数

借用克林顿竞选总统时的口号——“重要的是经济，懂吗？”（OneFlow译注：过后美国正值经济萧条，克林顿将经济作为竞选演说的重要话题，最终博得选举），在此，我想说，“重要的是内存，不是浮点运算数（FLOPs），懂吗？”

古代微处理器最大的瓶颈是能耗，而不是芯片集成度。Yahoo!创始人Mark Horowitz在十多年前就发现，拜访片外DRAM的能耗是拜访片上DRAM的100倍，是算术运算的5000~10,000倍。因而，咱们心愿能够通过减少浮点运算单元（FPU）来摊派内存拜访开销。基于Mark Horowitz的数据，芯片上的FPU数量被设置为10,000个左右。ML模型开发人员经常试图通过缩小浮点运算数来优化模型，但其实缩小内存拜访数才是更无效的方法。

TPU v1有65,000多个乘法单元，比GPU、CPU等硬件高出许多倍。只管它的时钟频率较低，仅为700MHz，但因为其乘法单元数量微小，且每个乘法单元可进行2个运算操作，因而TPU v1每秒可执行65,000×2×700M≈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既要专门优化，也要灵便

作为一种针对特定畛域的架构（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 v1应用DDR3内存，因为它针对的是推理，只需应用已有的权重，不须要生成权重。针对训练的TPU v2则不一样，训练时既要读取权重，也要写入权重，所以在v2中，咱们将本来的DDR3改为与向量存储区相连，这样就既能向其读取数据，又能向其写入数据。

而后，咱们将DDR3改为HBM。因为从DDR3读取参数速度太慢，影响性能，而HBM的读写速度快20倍。

第五步，咱们在HBM和向量存储区之间减少互连（Interconnect），用于TPU之间的连贯，组成咱们之前提到的Pod超级计算机。以上就是从TPU v1到TPU v2的改良。

教训八：半导体技术的倒退速度参差不齐

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

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

因为布线技术的提高绝对滞后，如果仍像TPU v1一样，每块芯片只有一个Tensor Core，就会导致管道更为简短，如果管道出了问题也会更加麻烦。因而，咱们将两个Tensor Core相互连贯，这对编译器而言也更敌对。

Google做出TPU v2之后，心愿再花一年工夫欠缺v2，所以TPU v3没有引进新技术，只是v2的改进版。

与v2相比，TPU v3有以下特点：

• 体积只大了不到10%；
• 矩阵乘法单元（MXU）的数量翻倍，因而峰值性能也翻倍；
• 时钟频率放慢了30%，进一步放慢计算速度；
• 内存带宽扩充了30%；
• 容量翻倍，可使多种利用更加不便；
• 芯片间带宽扩充30%；
• 可连贯的节点数是之前的4倍。

上图左上角即为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 v3到TPU v4i，矩阵乘法单元的数量再次翻倍，但芯片面积却没有扩充。如前所述，计算逻辑的倒退速度是最快的。

如果想理解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兼容性非常重要

XLA（减速线性代数）编译器可对全程序进行剖析和优化，优化分为与机器无关的高级操作和与机器相干的低级操作，高级优化操作将影响TPU和GPU，所以通常咱们不会改变高级优化操作，免得导致失灵，但咱们能够改变低级优化操作。

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

最初的问题是，与CPU相比，DSA的软件栈还不够成熟。那么编译器优化最终可能提速多少？

实际上的提速相当可观。其中重要的优化之一称为算子交融（Operator Fusion），如将矩阵乘法与激活函数进行交融，省略将两头后果写入HBM再读取进去的步骤。上图是咱们的MLPerf基准测试后果，可见，应用算子交融均匀能够带来超两倍的性能晋升。

上图显示了编译器优化的提速成果。蓝色示意应用GPU，红色示意应用TPU，通过短短十几个月的优化后，不少模型的性能都晋升到了两倍。要晓得，对C++编译器而言，如果能在一年内把性能晋升5%-10%就曾经十分了不起了。

此外，编译器的后向ML兼容性十分重要。我的共事Luiz Barroso主管Google的一个与计算机架构无关的部门，他示意不心愿在训练中花太多工夫，心愿一早晨就能够训练好模型，第二天能够间接部署。咱们心愿训练和推理时后果统一，这就是咱们说的后向ML兼容性，咱们还心愿它能在所有TPU上运行，而不是每次更改TPU时都要从新训练。

为什么放弃后向ML兼容性如此艰难？因为浮点加法不合乎结合律，所以运算程序可能会影响运算后果。而TPU的工作就是让所有机器对编译器而言都没有区别，以便能够在重组代码的同时取得雷同的高质量后果，以实现后向ML兼容性。

教训十：优化的指标是Perf/TCO还是Perf/CapEx

在将研究成果利用到理论生产时，咱们优化的指标是什么？Google构建硬件是为了用在本人的数据中心，所以咱们所要管制的老本是指总体领有老本（TCO），包含资本老本（洽购老本）和运行老本（电力、冷却、空间老本）。资金老本是一次性的，而运行老本须要继续收入3~4年。

因而，芯片和主板生产商只须要思考产品性能/资本老本的比率；而Google却要思考整个硬件生命周期的老本，关注性能/总体领有老本之间的比率。如下面的饼状图所示，电力可占总体领有老本的一半。所以，如果把眼光扩充到总体领有老本上，在零碎设计时就很可能会做出不同的取舍。

之前提到，TPU v1有一个Tensor Core，v2和v3有两个。到了v4时，基于对总体领有老本的思考，Google决定离开设计：用于训练的TPU v4有两个Tensor Core，用于推理的TPU v4i只有一个。这样就大大晋升了性能和总体领有老本之间的比率。

上图是TPU v4i和TPU v3的每瓦性能比照，红色是TPU v4i，蓝色是TPU v3，前者的每瓦性能是后者的两倍以上。

3

晋升机器学习能效，缩小碳脚印

2021年10月的IEEE Spectrum杂志有一篇文章提到，训练某一模型须要数年工夫，破费1000亿美元，总碳排放量相当于纽约一个月的排碳量，如果还要进一步晋升模型准确度，这些数字还会更夸大。

2022年1月，又有文章示意，依据以后的算力需要增长曲线预计，到2026年，训练最大AI模型的老本将相当于美国的GDP，大略是20万亿美元。

Google钻研了ML硬件的能源消耗。不同于全生命周期耗费的能源（包含从芯片制作到数据中心构建的所有间接碳排放），咱们只关注硬件运行时的能源消耗。

展望未来，咱们有方法让机器学习的能耗升高到原来的100倍，碳排放量升高1000倍。咱们能够从四方面协同着手，极大地促成机器学习在更多畛域的可继续倒退：

第一个因素是模型。Google在2017年颁布Transformer模型，四年后，又开发了Primer模型，其计算品质雷同，但能效更高。Primer的能耗和碳排放量相比Transformer升高了4倍。

第二个因素是硬件。2017年所应用的P100 GPU和以后最新TPU的性能相差了14倍。所以，前两个因素联合，能够将能耗和碳排放量升高60倍。

第三个因素是数据中心的能效。Google的PUE大概是其余数据中心的1.4倍。所以，前三个因素累积，能够将能耗和碳排放量升高80倍。

第四个因素是数据中心的地理位置。即便在同一个国家，不同地区应用无碳能源的比例也可能大不相同。在Google所有数据中心所在地中，俄克拉荷马州的无碳能源占比最高，Primer模型就是在此处训练的，这可将碳排放量在上述根底上升高9倍。

综合上述四个因素，咱们可将机器学习的能耗升高80倍，碳排放量升高700倍。这十分了不起。

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既缩小了碳排放量，而且计算品质更好。

（原视频链接：https://www.youtube.com/watch...）

欢送下载体验 OneFlow v0.8.0 最新版本：https://github.com/Oneflow-In...