关于深度学习: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…