万字长文,Google TPU往事

万字长文,Google TPU往事
2024年03月29日 13:17 AI科技大本营

【编者按】尽管 Google 的 AGI 之战稍显落后,不敌微软 X OpenAI、Amazon X Claude 以及 META 的 LLaMA,但许多人仍对 Google 怀以希望,毕竟 Google 的 TPU 这张王牌还没有大显身手。本文深入回顾了 TPU 的起源以及架构的细节,还提到了 AlexNet 作者 Alex Krizhevsky 是如何让模型训练走上了 GPU 之路。故事与技术相交织,可以一窥 AI 的演进。

作者 | Babbage

责编 | 沭七

出品丨AI 科技大本营(ID:rgznaI100)

计算历史中被遗忘或过时的想法仍具有研究价值,并且有时极具实用性。使用神经网络的例子,经历了几次“AI 寒冬”与热潮交替的历程,几乎七十年来时冷时热,无疑是其中最具影响力的案例。

若说对神经网络的兴趣历经起伏,那么为其加速的专用硬件的发展同样经历了多次浮沉。然而,这类硬件面临着一个典型的“鸡生蛋蛋生鸡”的问题。CPU 并非执行神经网络的最高效方式。但是,没有神经网络的实用证明,如何为专用硬件的持续投入提供理由呢?同时,神经网络的有效运作又需要更强大的硬件。随着炒作周期的变化,一次又一次地,建造专用硬件的项目被搁置。最终,原本为其他目的而构建的硬件——GPU,打破了这一僵局。

DALL·E 生成

随后,专用硬件的出现似乎是必然的。我们现在正处于一个“炒作期”,这一炒作导致许多公司,从初创公司到超大型企业,纷纷打造自己的加速器。尽管大多数设计最终可能会失败,专用神经网络加速器肯定会在未来的计算领域占据一席之地。谷歌的张量处理单元(TPU),作为首批大规模部署的加速器,无疑将成为其中的幸存者。

为何谷歌的 TPU 值得关注并具有重要性?它们并非可以从谷歌或任何其他方购买的产品。你可以在谷歌云平台上“租用”它们,但继续使用 TPU 依赖于谷歌继续提供支持的承诺。这似乎是一个不稳定的基础。

但这毕竟是谷歌,仅仅因为谷歌在其服务中使用 TPU,就足以证明它们的重要性。有时人们容易忘记谷歌服务的规模之大。该公司有九项服务(包括搜索、Android、Chrome、Gmail、地图、Play 商店、YouTube、照片)拥有超过十亿用户。此外,谷歌还拥有世界第三大的“公共云”平台。

在一个由 Nvidia 在数据中心机器学习硬件领域所主导的时代,谷歌的 TPU 提供了一个已被大规模部署的、证明有效的 Nvidia GPU 的替代方案的实例。谷歌 TPU 项目背后的核心理念——创建专门用于加速机器学习任务的数据中心硬件,已经促使许多其他公司,从初创企业(如 Cerebras 和 Tenstorrent)到大型企业(如 AWS 和 Microsoft),构建自己的硬件。其中一个风头无两的初创公司 Groq,由原 TPU 团队成员 Jonathan Ross 领导。

第一款 TPU 的开发历史同样引人入胜。它展示了即使是像谷歌这样的公司,也能在迫切需要时迅速且灵活地行动。一个拥有必要资源的公司可以迅速开发出创新的硬件。它还展示了数十年前的想法如何重新浮现,成为现代重要设计的核心。

首先,我们澄清一下,当我们讨论“张量处理单元”时,我们讨论的究竟是什么。

张量硬件与软件

近年来,“张量”一词开始频繁出现在一些计算机硬件和软件的命名中。以下是一个简短的列表:

1. TensorFlow:谷歌的机器学习框架;

2. 谷歌张量处理单元(Tensor Processing Unit,TPU):谷歌的自定义数据中心加速器;

3. Tensor Core:Nvidia 最新 GPU 中的一个执行单元;

4. 谷歌 Tensor:最新的 Pixel 智能手机中的基于 ARM 的系统级芯片。

本文将重点讨论第二项,即谷歌的自定义数据中心加速器,涵盖第一代谷歌 TPU 的历史和架构。如有必要,我将使用 TPU v1 来区分其后续版本。

这些“张量”的突然出现,当然是因为机器学习,特别是深度学习的重要性迅速增加。张量是深度学习公式化和实现的核心。

那么,什么是张量?已经熟悉张量的读者可以跳过这一节。

维基百科给出了一个不那么直观的定义:

在数学中,张量是一种代数对象,它描述了与向量空间相关的一组代数对象之间的多线性关系。张量可以映射不同对象之间的关系,如向量、标量,甚至其他张量。

不过,不要因此而关闭页面!实际上,我们可以简化这一定义,几乎完全依赖于张量的一种表现形式:

张量可以表示为一个(可能是多维的)数组。

因此,根据张量的性质,它可以被表示为一个 n 维数组,其中 n 是 0、1、2、3 等等。这些表示形式中的一些有更熟悉的名字:

  • 维度 0 - 标量

  • 维度 1 - 向量

  • 维度 2 - 矩阵

为什么称之为张量处理单元(TPU)?因为它设计来加速涉及张量的操作。具体来说,是哪些操作呢?我们原始的维基百科定义中提到的操作,将张量描述为“在不同对象之间的映射(多线性关系),如向量、标量,甚至其他张量”。

让我们举个简单的例子。一个二维数组可以描述两个一维数组之间的多线性关系。数学倾向的人会认识到,这个过程是通过将一个向量乘以一个矩阵来得到另一个向量。

这可以推广到表示更高维数组之间关系的张量。然而,尽管张量描述了任意高维数组之间的关系,实际上我们将考虑的 TPU 硬件是设计来执行与一维和二维数组相关的计算。或者,更具体地说,向量和矩阵操作。

谷歌 TPU 的起源故事

要了解 TPU 的起源,我们需要回顾谷歌与机器学习和深度学习长期且不断发展的历史。谷歌对机器学习的兴趣是显而易见的。谷歌的使命是“组织世界上的信息,使之普遍可访问且有用。”机器学习提供了一套强大的工具,帮助谷歌组织那些信息并使其更有用。从图像和语音识别到语言翻译,再到“大型语言模型”,当然,还有在谷歌每年赚取数十亿美元的广告业务中的众多应用。对谷歌来说,机器学习至关重要。

因此,谷歌对机器学习的兴趣始于 2000 年代初。随着时间的推移,公司的关注点开始转向深度学习。一些关键的里程碑展示了谷歌对深度学习的快速增长的关注和承诺,包括:

  • 2011 年:由 Jeff Dean、Greg Corrado 和 Andrew Ng 发起的深度学习研究项目 Google Brain 作为 Google X 的一部分启动。

  • 2013 年:在 2012 年底成功的 AlexNet 图像识别项目之后,经过与 Microsoft、百度和 DeepMind 的竞标过程,谷歌以 4400 万美元的价格收购了由 Geoffrey Hinton、Alex Krizhevsky 和 Ilya Sutskever 组成的团队。

  • 2014 年:谷歌以高达 6.5 亿美元的价格收购了由 Demis Hassabis、Shane Legg 和 Mustafa Suleyman 创立的 DeepMind。

2012 年是一个关键时刻,当时 Google Brain 团队能够利用深度学习在图像识别性能上实现突破性进展。

使用这个大规模神经网络,我们还显著提高了标准图像分类测试的状态——事实上,我们看到了 70% 的相对准确性改进。我们通过利用网络上大量的未标记数据来增强有限的标记数据集来实现这一点。这是我们非常关注的一个方向——如何开发能够很好地扩展的机器学习系统,以便我们可以利用大量的未标记训练数据。

早在 2006 年,就有关于开发机器学习定制硬件的讨论。当时的结论是,设想的应用所需的计算能力可以使用谷歌庞大数据中心中已有的服务器容量以低成本提供。记录这一点的论文“利用大规模无监督学习构建高级特征”展示了他们使用“一个拥有 1000 台机器(16,000 核心)的集群训练了三天”来训练论文中描述的模型。那是 1000 个 CPU,每个 CPU 有 16 个核心。

Cade Metz 的书《天才制造者》回顾了当 Alex Krizhevsky(AlexNet 作者)在 2013 年来到谷歌时,他发现他们现有的模型都在 CPU 上运行。Krizhevsky 需要 GPU。所以他亲自解决了这个问题!

在公司的最初几天,他从当地的电子商店购买了一台 GPU 机器,把它放在离他的办公桌不远的走廊的衣橱里,将其接入网络,并开始在这台独立的硬件上训练他的神经网络。

最终,Krizhevsky 的新同事们意识到,他们需要 GPU,而且是很多 GPU。2014 年,谷歌决定购买大约 40,000 个 Nvidia GPU,花费约 1.3 亿美元。这些 GPU 被用于训练谷歌业务中使用的深度学习模型。

顺便说一句,像谷歌这样的公司购买 40,000 个 GPU 并不会被忽视。如果 Nvidia 需要一个信号,说明将 GPU 应用于深度学习可能是一个规模可观的业务,那么谷歌几乎可以肯定在 2014 年提供了这个信号。

但这些 GPU 并不一定解决了谷歌最大的潜在问题。它们非常适合训练谷歌正在开发的模型。但是,部署这些模型在谷歌的规模上会如何呢?深度学习的应用,如语音识别,可能会非常受欢迎。这种潜在的受欢迎程度给谷歌带来了一个问题。让我们考虑一下谷歌与之同义的产品,搜索。

谷歌搜索对任何拥有网络浏览器和互联网连接的人来说都是免费的。这种搜索的质量以及它的免费性,是推动其爆炸性增长和迅速占主导地位的关键因素。据估计,到 2023 年,谷歌提供了大约两万亿次搜索结果。要能够以这种规模免费提供搜索,谷歌需要能够非常便宜地提供每一组搜索结果。

谷歌团队可以预见到使用这些新的深度学习技术构建的一系列服务将会非常受欢迎。谷歌可以利用其品牌,以及将这些服务与其现有广受欢迎的产品如搜索、Gmail 和 Android 集成,再次使这些服务占据主导地位。

但有一个重大问题。在这种规模下,这些服务将需要投资更多的硬件。2013 年,谷歌团队设计了一个使用语音识别的示例,说明了规模。如果用户每天只对他们的 Android 手机说三分钟话,并且谷歌使用深度学习将这些语音转换为文本使用 CPU,那么公司将需要增加 2 倍或 3 倍的服务器数量,而这些服务器已经在使用中。这将是极其昂贵的。

一个替代方案是扩大使用 GPU,Nvidia GPU 是显而易见的首选。到 2013 年,Nvidia 的 GPU 及其相关的软件生态系统已经非常成熟,被用作机器学习研究的关键工具,用于训练和推理。Nvidia 的第一款针对通用计算任务的 GPU 出现在 2006 年,Nvidia 的 CUDA 框架,用于在 GPU 上编程通用计算任务,出现在 2007 年。关键的是,Nvidia GPU 可以立即购买。

正如我们所看到的,谷歌确实购买了 Nvidia GPU 并将它们安装在其数据中心。然而,依赖 Nvidia GPU 并不一定是技术上或战略上的最佳解决方案。GPU 比 CPU 更适合深度学习所需的计算,但它们并不完全专门用于这些计算。依赖 GPU 会让潜在的效率损失,而在谷歌的规模下,这些效率损失代表了重大的成本节省损失。此外,依赖单一供应商提供对谷歌战略至关重要的硬件会增加该战略的显著风险。

有两个其他选择。第一个是谷歌可以使用 FPGA,它们可以编程执行深度学习所需的专用计算。

第二个是谷歌可以设计和构建自己的定制硬件,形式为特定应用集成电路(ASIC)。使用这种专门用于深度学习的定制硬件,将开放更多的效率潜力,并将减少谷歌对单一供应商的依赖。

与已经安装的 GPU 相比,FPGA 的性能很快就显示出来会被击败。这些 FPGA 并没有完全被放弃。

FPGA 仍然在数据中心中使用,比 ASIC 可用的时间早六个月;它充当了一个“清洁管道”,用于所有支持生产中新加速器的谷歌部署过程。

但重点将是 ASIC。目标将是开发一个 ASIC,与 GPU 相比,在推理上产生 10 倍的成本性能优势。

但开发定制硬件需要时间。具体需要多长时间将是至关重要的。对于该项目来说,它不能是一个拖延多年的“研究”项目。它需要快速地向谷歌的数据中心交付硬件,并且要以大规模交付。为此,该项目将需要利用谷歌能够快速获取的任何资源和现有知识。

TPU 团队和现有研究

如何在没有现有内部团队的情况下快速设计复杂和创新的硬件。令人惊讶的是,谷歌很快就能组建一个非常有效的团队。它是如何做到的?

谷歌已经在为其数据中心创建定制硬件,并且可以召集该团队的成员来帮助新项目。然而,该公司以前没有创建过自己的处理器芯片,所以它开始引入具有相关经验的新团队成员。他们是如何做到这一点的?

其中一位新成员 Norm Jouppi 回忆说,他是如何在一次会议上与 Jeff Dean 交谈后被招募到 TPU 团队的,以及如何被一个创新的“green field”项目所吸引。加入团队的其他知名人物包括 David Patterson(图灵奖得主),原始 Berkeley RISC 设计的开发者和 RISC-V 指令集架构开发的关键人物。

谷歌很快就拥有了一个才华横溢且经验丰富的团队。但即使如此,如果他们要从头开始开发新系统的架构,他们可能还是难以满足为他们设定的目标。幸运的是,对于谷歌来说,有一个现有的架构方法可以很容易地应用。这种方法是在开始工作于 TPU 之前三十五年多的时间里首次开发的。

在他们 1978 年的论文 Systolic Arrays (for VLSI)中,卡内基梅隆大学的 H.T Kung 和 Charles E. Leiserson 提出了他们所称的“systolic system”(脉动系统)的建议。

systolic system 是一个处理器网络,它有节奏地计算并通过系统传递数据....在脉动计算机系统中,处理器的功能类似于心脏。每个处理器定期地输入和输出数据,每次进行一些简短的计算,以便在网络中保持数据的正常流动。

然后他们确定了 systolic system 的一个应用,矩阵计算。

许多基本的矩阵计算可以在具有数组结构的脉动网络上优雅而有效地进行流水线处理。作为一个例子,六边形连接的处理器可以最优地执行矩阵计算...这些脉动阵列享有简单和规则的通信路径,而且网络中使用的几乎所有处理器都是相同的。因此,基于脉动阵列的专用硬件设备可以使用 VLSI 技术廉价构建。

作为 Kung 和 Leiserson 所处理的技术的年代的一个插图,他们说明了如何使用“脉动系统”作为 DEC PDP-11 小型计算机的附件。

该论文描述了脉动阵列的处理元素可能连接的替代方式。

其中一个提议的安排 - 正交连接 - 与 ILLIAC IV 超级计算机中处理元素的安排产生了共鸣,我们在 ILLIAC IV 超级计算机:DARPA、SIMD、Fairchild 和 Stanley Kubrick 的'2001'中讨论了这一点)。

Kung 和 Leiserson 继续描述了如何使用六边形连接的脉动阵列来执行矩阵乘法。

Kung 和 Lieberson 的这项原始工作在接下来的十年中得到了更多的研究和更多的论文,展示了如何使用脉动阵列解决各种问题。

脉动阵列是如何工作的?我们将在下一篇帖子中更详细地查看 Google 的 TPU 团队如何有效地实现矩阵运算的概念。基本思想是数据通过数组的一侧或多侧输入系统,并且数据和中间结果随着每个“脉冲”的流动通过系统。所需的结果在经过足够的脉冲导致所需计算完成后,从数组的一侧或多侧出现。

到 2013 年,Kung 和 Lieberson 想法背后的一些原始动机,特别是处理 1970 年代制造技术的限制,已经消失。然而,这种方法对于像矩阵乘法这样的任务的固有效率,特别是在 2013 年相关的,相对较低的功耗,仍然存在。因此,TPU 将使用脉动阵列。

谷歌拥有工程师和架构方法,但在将想法转化为工作硅片的能力上仍然存在差距。为此,它将转向 LSI 公司(现为 Broadcom 的一部分)。乍一看,现在的 Broadcom 可能看起来像一个令人惊讶的合作伙伴。它不是一个通常与机器学习硬件相关联的公司。但它们可以与像 TSMC 这样的制造合作伙伴合作,使谷歌的设计成为现实,并且以大规模实现。

最后,让 TPU 能够用于实际工作负载意味着的不仅仅是制造大量芯片。这意味着构建软件,以便谷歌的深度学习工具可以在新设计上运行。新设计必然意味着一个全新的指令集架构,这意味着编译器必须适应新架构,这本身就是一个不小的壮举。

部署和使用

TPU 团队的成员总结了他们的目标如下:

快速构建

实现高性能

......在大规模......

...对新工作负载开箱即用......

同时具有成本效益

他们确实快速构建了它。项目启动仅 15 个月后,第一款 TPU 就于 2015 年初部署在谷歌的数据中心。

他们是如何做到这一点的?论文《“数据中心性能分析的回顾:张量处理单元”》总结了一些因素:

我们有时会吹嘘 TPU 从项目启动到数据中心部署的 15 个月时间,远远短于标准的生产芯片。确实,谷歌依赖于生产的后续 TPUs 具有多年的设计周期。快速上市得益于单一的日程焦点,不仅在架构上——其中 700MHz 的时钟速率使容易实现时间闭环,老化的 28nm 工艺已经完全调试——还在我们的数据中心部署团队的英勇工作上。

“老化的 28nm 工艺”和相对较低的时钟速率是有帮助的因素。然而,这忽略了谷歌能够调动必要资源来支持这样的项目的能力。

TPU 很快就会在谷歌内部用于广泛的应用。很快迎来了一个有趣且高调的例子:谷歌的 DeepMind 使用 TPU 用于其 AlphaGo Zero 项目。AlphaGo Zero 是在 2016 年 3 月击败了世界围棋冠军李世石的模型。

AlphaGo Zero 是本文所描述的程序。它从自我对弈的强化学习中学习,从随机的初始权重开始,不使用展开,没有人类的监督,仅使用原始棋盘历史作为输入特征。它仅使用谷歌云中的一台机器和 4 个 TPU。

2016 年 AlphaGo Zero 的成功似乎是一个重要时刻。以如此少的硬件实现这一成就是非凡的。

TPU 的宣布

像谷歌这样的公司不需要分享其数据中心幕后发生的细节。因此,TPU 在 2015 年首次引入后一年多保持秘密。

然后谷歌的 CEO Sundar Pichai 在 2016 年 5 月 18 日谷歌 IO 大会的主题演讲中宣布:

我们在我们的数据中心内运行 TPUs 超过一年了,并且发现它们在机器学习任务中提供了每瓦性能的数量级提升。

这个发布伴随着一个简短的博文,谷歌使用 TPU 定制芯片为机器学习任务增加动力。除了一些非常简短的技术细节,该博文还分享了一些 TPU 如何帮助支持谷歌的服务。

TPU 已经为谷歌的许多应用提供动力,包括用于提高搜索结果相关性的 RankBrain,以及用于提高我们的地图和导航的准确性和质量的街景查看。AlphaGo 在与围棋世界冠军李世石的比赛中由 TPU 提供动力,使其能够“更快地思考”并在移动之间更远地看到。

深入 TPU 架构

正如我们所见,TPU v1 团队采用的方法是 H.T Kung 和 Charles E. Leiserson 在 1978 年的论文《脉动阵列(针对 VLSI)》中首次提出的架构。

脉动系统是一个由处理器组成的网络,这些处理器有节奏地计算并传递数据……在脉动计算机系统中,处理器的功能类似于心脏。每个处理器定期地输入和输出数据,每次进行一些简短的计算,以此保持网络中数据的正常流动。

那么,在 TPU v1 中如何使用脉动方法有效执行矩阵乘法呢?让我们回到我们的 2x2 矩阵乘法示例。

如果我们有一个 2x2 的乘法单元数组,这些单元在一个简单的网格中连接,并且我们以正确的顺序将我们要相乘的矩阵的元素输入到网格中,那么矩阵乘法的结果自然会从数组中显现。

计算可以用以下图表示。每个角落的方块代表一个乘法/累加单元(MAC),它可以执行乘法和加法操作。

在这张图中,黄色的值是从顶部和左侧输入矩阵的输入值。浅蓝色的值是存储的部分和。深蓝色的值是最终结果。

让我们一步步来分析。

第 1 步:

值 a11 和 b11 被加载进左上角的乘法/累加单元(MAC)。它们被相乘,结果被存储。

第 2 步:

值 a12 和 b21 被加载进左上角的 MAC。它们被相乘,并加到之前计算的结果上。这给出了结果矩阵的左上角值。

与此同时,b11 被传输到右上角的 MAC,在那里它与新加载的值 a21 相乘,结果被存储。同时,a11 被传输到左下角的 MAC,在那里它与新加载的值 b12 相乘,结果被存储。

第 3 步:

b21 被传输到右上角的 MAC,在那里它与新加载的值 a22 相乘,结果被加到之前存储的结果上。同时,a12 被传输到左下角的 MAC,在那里它与新加载的值 b22 相乘,结果被加到之前存储的结果上。在这一步中,我们计算了结果矩阵的右上角和左下角值。

与此同时,a12 和 b21 被传输到右下角的 MAC,在那里它们被相乘,结果被存储。

第 4 步:

最后,a22 和 b22 被传输到右下角的 MAC,在那里它们被相乘,结果被加到之前存储的值上,给出了结果矩阵的右下角值。

因此,矩阵乘法的结果沿着 MAC 矩阵中移动的“对角线”显现出来。

在我们的示例中,进行 2x2 矩阵乘法需要 4 个步骤,但这仅是因为在计算开始和结束时一些 MAC 没有被利用。实际上,一旦 MAC 变为空闲,新的矩阵乘法就会从左上角开始。因此,该单元每两个周期就能完成一次新的矩阵乘法。

这是对脉动阵列工作方式的简化表示,我们简要地介绍了 TPU v1 中脉动阵列实现的一些细节。我希望这种架构工作的原理已经清楚。

这是最简单可能的矩阵乘法,但可以扩展到更大的矩阵和更大的乘法单元阵列。

关键点是,如果数据以正确的顺序输入到脉动阵列中,则通过系统的值和结果的流动将确保所需结果随时间从阵列中出现。没有必要将中间结果存储并从“主存储”区域提取。由于矩阵乘法单元的结构和输入被输入到单元的顺序,中间结果在需要时自动可用。

当然,矩阵乘法单元不是孤立存在的,整个系统的最简单呈现如下:

首先要注意的是,TPUv1 依赖于通过 PCIe(高速串行总线)接口与主机计算机的通信。它还可以直接访问自己的 DDR3 动态 RAM 存储。

我们可以将其扩展为设计的更详细展示:

让我们从这个设计的展示中挑选一些关键元素,从顶部开始,大致顺时针移动:

  • DDR3 DRAM / Weight FIFO:权重存储在通过 DDR3-2133 接口连接到 TPU v1 的 DDR3 RAM 芯片中。权重是从主计算机的内存通过 PCIe 预加载到这些芯片上的,然后可以转移到“Weight FIFO”内存中,为矩阵乘法单元的使用做好准备。

  • 矩阵乘法单元:这是一个“脉动”阵列,拥有 256 x 256 的矩阵乘法/累加单元,由顶部的 256 个“权重”值和左侧的 256 个数据输入供给。

  • 累加器:结果从脉动矩阵单元的底部出现,并存储在“累加器”内存存储中。

  • 激活:在此应用上述神经网络中描述的激活函数。

  • 统一缓冲区/脉动数据设置:应用激活函数的结果存储在“统一缓冲区”内存中,准备作为输入反馈到矩阵乘法单元以计算下一层所需的值。

乘法/累加单元(MACs)

到目前为止,我们还没有具体说明矩阵乘法单元执行的乘法的性质。TPU v1 执行 8 位 x8 位整数乘法,利用量化来避免对更占用更多硅片面积的浮点计算的需求。

指令集

TPU v1 采用 CISC(复杂指令集计算机)设计,大约只有 20 条指令。值得注意的是,这些指令是由主计算机通过 PCIe 接口发送给它的,而不是从内存中获取的。

五个关键指令如下:

Read_Host_Memory

通过 PCIe 从主计算机的内存读取输入值到统一缓冲区。

Read_Weights

从权重内存读取权重到 Weight FIFO。注意,权重内存将已经通过 PCIe 从计算机的主内存中读取了权重。

Matrix_Multiply / Convolve

根据论文,这条指令:

…使矩阵单元执行矩阵乘法或卷积,从统一缓冲区到累加器。矩阵操作接受一个变量大小的 B*256 输入,将其与一个 256x256 的恒定权重输入相乘,产生一个 B*256 的输出,完成 B 流水线周期。

这是实现脉动阵列矩阵乘法的指令。它还可以执行卷积计算,这对卷积神经网络是必需的。

Activate

根据论文,这条指令:

执行人工神经元的非线性函数,有 ReLU、Sigmoid 等选项。它的输入是累加器,输出是统一缓冲区。

如果我们回到我们简单的神经网络模型,隐藏层中的值是将“激活函数”应用于输入的权重乘以输入的总和的结果。ReLU 和 Sigmoid 是两种最受欢迎的激活函数。在硬件中实现这些将提供应用激活函数的有用加速。

Write_Host_Memory

通过 PCIe 将结果从统一缓冲区写入主计算机的内存。

值得暂停一刻,反思这五条指令在提供几乎完整的 TPU v1 推理实现方面的优雅。在伪代码中,我们可以大致描述 TPU v1 的操作如下:

Read_Host_MemoryRead_WeightsLoop_Start    Matrix_Multiply    ActivateLoop_EndWrite_Host_Memory

强调脉动单元在使这成为可能和高效中的重要性也是有用的。正如 TPU v1 团队所描述:

..矩阵单元使用脉动执行以节省能量,通过减少对统一缓冲区的读写....它依靠来自不同方向的数据在数组中的单元以规律间隔到达时被组合。...数据从左侧流入,权重从顶部加载。给定的 256 元素乘法-累加操作以对角线波前的形式通过矩阵移动。

软件

没有软件支持的 TPU v1 硬件将毫无用处。谷歌开发并使用了 Tensorflow,因此创建“驱动程序”以使 Tensorflow 能够与 TPU v1 协同工作是主要步骤。

TPU 软件栈必须与为 CPU 和 GPU 开发的软件栈兼容,以便应用程序可以快速移植到 TPU。在 TPU 上运行的应用程序的部分通常用 TensorFlow 编写,并编译成可以在 GPU 或 TPUs 上运行的 API。

与 GPUs 一样,TPU 栈分为用户空间驱动程序和内核驱动程序。内核驱动程序很轻量级,仅处理内存管理和中断。它设计用于长期稳定。用户空间驱动程序变化频繁。它设置和控制 TPU 执行,将数据重新格式化为 TPU 顺序,将 API 调用转换为 TPU 指令,并将它们转换为应用程序二进制文件。

制造和模具

TPU v1 是由 TSMC 使用相对“成熟”的 28nm TSMC 工艺制造的。谷歌表示,与谷歌此时在其数据中心使用的更先进工艺制造的 Intel Haswell CPU 和 Nvidia 的 K80 GPU 芯片相比,模具面积不到一半。

我们已经看到了 TPU v1 的指令集有多简单,只有 20 条 CISC 指令。ISA 的简单性导致 TPU v1 的模具中用于解码和相关活动的“开销”非常低,仅占模具面积的 2%,标记为“控制”。

相比之下,24%的模具面积专用于矩阵乘法单元,29%专用于存储输入和中间结果的“统一缓冲区”内存。

性能

此时,回顾一下 TPU v1 旨在使推理——即在谷歌规模的真实世界服务中使用已经训练好的模型——更高效是有用的。它不旨在提高训练的速度或效率。尽管推理和训练有一些共同的特征,但在开发专用硬件时,推理和训练提供了完全不同的挑战。

那么 TPU v1 的表现如何呢?

2013 年,TPU v1 的关键比较是与 Intel 的 Haswell CPU 和 Nvidia 的 K80 GPU。

  • TPU v1 拥有 K80 GPU 的 25 倍 MAC 和 3.5 倍的片上内存。

  • TPU v1 在推理方面比 K80 GPU 和 Haswell CPU 快大约 15X - 30X。

并且关键的是 TPU v1 在能源效率方面远超 GPU:

  • TPU v1 的相对增量性能/瓦特是 GPU 的 25 到 29 倍。

教训

在本文中,我们关注了像谷歌这样的组织如何能够快速构建 TPU v1。同时,也看到了 TPU v1 的定制架构对于使其能够以比当代 CPU 和 GPU 更低的能耗产生更好的性能至关重要。

TPU v1 只是故事的开始。TPU v1 设计得很快,唯一的目标是使推理更快和更节能。它有一些明显的限制,不是为训练设计的。不仅谷歌内部,外部公司很快就开始考虑如何改进 TPU v1。

本文使用 GPT-4 生成,原文参见:

https://thechipletter.substack.com/p/googles-first-tpu-architecture

https://thechipletter.substack.com/p/googles-first-tensor-processing-unit

4 月 25 ~ 26 日,由 CSDN 和高端 IT 咨询和教育平台 Boolan 联合主办的「全球机器学习技术大会」将在上海环球港凯悦酒店举行,特邀近 50 位技术领袖和行业应用专家,与 1000+ 来自电商、金融、汽车、智能制造、通信、工业互联网、医疗、教育等众多行业的精英参会听众,共同探讨人工智能领域的前沿发展和行业最佳实践。欢迎所有开发者朋友访问官网 http://ml-summit.org、点击「阅读原文」或扫码进一步了解详情。

财经自媒体联盟更多自媒体作者

新浪首页 语音播报 相关新闻 返回顶部