谷歌的张量处理单元（TPU）是如何工作的？

Posted on 2026-01-06

扩展LLM应用规模的4个关键见解（赞助内容）

LLM 工作流可能复杂、不透明且难以保障安全。Datadog 最新电子书为您提供实用策略，助您在生产环境中监控、排查故障并保护 LLM 应用。您将深入了解如何克服大规模安全部署 LLM 的挑战，从调试多步骤工作流到检测提示注入攻击，应有尽有。

下载电子书

免责声明：本文内容均来自谷歌工程团队在线分享的信息。所有技术细节的版权均归谷歌工程团队所有。原文链接和参考文献位于文末的参考文献部分。我们已尽力分析这些信息并提供我们的观点。如果您发现任何不准确或遗漏之处，请留言，我们将尽力修正。

2016年3月，DeepMind的AlphaGo击败围棋世界冠军李世石，这在人工智能领域具有里程碑式的意义。而这场比赛所使用的硬件，谷歌已经投入生产运行一年多，但从未公开承认过。

张量处理单元（TPU）的意义远不止于一款高速芯片。它标志着计算理念的根本性转变：有时少即是多。

自 2015 年以来，谷歌的 TPU 系列芯片已发展了七代，从最初处理图像识别查询的单芯片扩展到如今拥有 9216 个芯片的超级计算机，用于训练目前规模最大的语言模型。本文将探讨谷歌为何要打造定制芯片，以及其工作原理，并揭示他们必须面对的物理限制和工程权衡。

对TPU的需求

2013年，谷歌的基础设施团队进行了一项计算。如果安卓用户像谷歌预期的那样大规模使用语音搜索，即使每天只使用三分钟，其计算需求也需要谷歌将其全球数据中心的规模扩大一倍。

当时这个问题似乎没有明显的解决方案。建造更多配备传统处理器的数据中心在经济上不可行。更重要的是，摩尔定律多年来一直在放缓。几十年来，半导体行业一直依赖于晶体管密度大约每两年翻一番的观察结果，从而在不改变架构的情况下实现性能的持续提升。然而，到了2013年，这种趋势开始减弱。谷歌不能仅仅等待英特尔的下一代CPU来解决它的问题。

造成这种情况的根本原因在于架构。传统计算机遵循冯·诺依曼架构，其中处理器和内存通过共享总线进行通信。要执行任何计算，CPU 必须先获取指令，然后从内存中检索数据，执行运算，最后将结果写回内存。处理器和内存之间这种持续不断的信息传输，造成了计算机科学家所说的冯·诺依曼瓶颈。

在总线上传输数据所需的能量通常超过计算本身所需的能量。例如，想象一下一位厨师正在准备一道菜，但他必须走到很远的储藏室去取每种食材。烹饪只需几秒钟，但走动却要花费数小时。对于文字处理或网页浏览等通用计算任务来说，这种设计是合理的，因为工作负载是不可预测的。然而，神经网络的情况则不同。

深度学习模型执行的操作几乎全部都是矩阵乘法。神经网络通过将输入数据与学习到的权重矩阵相乘、加上偏置值并应用激活函数来处理信息。对于单个预测，这个过程会重复数十亿次。拥有数千亿个参数的现代语言模型，每次查询都需要数千亿次乘加运算。关键在于，这些运算是可预测的、并行的和确定性的。

CPU 将大量处理能力用于分支预测和乱序执行等功能，旨在处理不可预测的代码。图形处理器 (GPU) 通过数千个核心并行工作改善了这种情况，但它们仍然继承了图形处理方面的架构开销。谷歌的远见卓识在于制造只执行神经网络所需功能的芯片，并去除所有其他功能。

脉动阵列：一种不同的计算方法

TPU 的核心是一种名为脉动阵列的架构。这个名称源自希腊语中“心跳”一词，指的是数据在芯片内有节奏地脉动传输。为了理解这一点的重要性，我们可以思考不同的处理器是如何处理同一任务的。

  • CPU 的运行就像一个工人来回奔波于水井和火堆之间，一次只装满一个桶。

  • GPU 部署数千个工作进程同时执行相同的任务。吞吐量增加，但油井和火源之间的流量变得混乱且耗能。

  • 循环式灌溉系统采用了一种截然不同的方法。工人们排成一列，手拉手传递水桶。水在链条中流动，直到工作完成，期间无人返回水源。

在 TPU 中，工作单元是排列成密集网格的简单乘加运算单元。第一代 TPU 使用 256×256 的阵​​列，这意味着 65,536 个计算器同时运行。计算过程如下：

  • 神经网络权重从上方加载到每个计算器中，并保持不变。

  • 输入数据从左侧逐行流入。

  • 数据经过每个计算器时，都会乘以居民体重。

  • 该乘积会累加到累计总和中，然后向右传递给下一个计算器。

  • 局部结果会累积并向下传递。

  • 所有计算完成后，最终结果将从底部得出。

请看下图：

这种设计意味着数据只需从内存中读取一次，但在遍历阵列的过程中会被使用数千次。传统处理器几乎每次操作都需要访问内存。而脉动阵列消除了这一瓶颈。数据仅在空间上相邻的计算器之间通过短导线传输，从而显著降低了能耗。

数据有力地证明了这种方法的合理性。

  • TPU v1 的 256×256 阵列每个时钟周期可以执行 65536 次乘加运算。在 700 MHz 的运行频率下，它每秒可完成 92 万亿次 8 位运算，而功耗仅为 40 瓦。

  • 现代 GPU 每个周期可能执行数万次运算，而 TPU 则执行数十万次运算。

  • 超过 90% 的硅片执行了有效的计算，而 GPU 中只有大约 30% 的硅片执行了有效的计算。

这里的权衡之处在于绝对的专业化。脉动阵列只能高效地执行矩阵乘法，而无法渲染图形、浏览网页或运行电子表格。谷歌接受了这一限制，因为神经网络推理本质上就是多次重复的矩阵乘法。

支持架构

脉动阵列需要精心协调的辅助组件才能发挥其性能。每个组件都解决了从原始数据到人工智能预测流程中的特定瓶颈。

让我们来看看最重要的组成部分：

矩阵乘法单元

矩阵乘法单元（MXU）就是脉动阵列本身。

TPU v1 使用单个 256×256 数组，处理 8 位整数。后续版本采用 128×128 数组，并使用 Google 的 BFloat16 格式处理训练工作负载，然后在 v6 版本中又回归到 256×256 数组，吞吐量提升了四倍。这种权重固定设计最大限度地减少了数据移动，而数据移动是计算中主要的能耗来源。

统一缓冲区

统一缓冲区提供 24 兆字节的片上 SRAM，作为​​速度较慢的外部存储器和容量较大的 MXU 之间的高速过渡区域。

该缓冲区存储来自主机的输入激活值、神经网络层之间的中间结果以及传输前的最终输出。由于该存储器直接位于芯片上，因此其带宽高于外部存储器。这种带宽差异对于保持 MXU 持续接收数据至关重要，避免其闲置等待内存访问。

矢量处理单元

向量处理单元 (VPU) 可以处理 MXU 无法处理的操作，包括 ReLU、sigmoid 和 tanh 等激活函数。

神经网络需要非线性运算来学习复杂的模式。如果没有非线性运算，多层神经网络在数学上就会简化为单一的线性变换。TPU 没有采用软件实现这些功能，而是使用专用的硬件电路在一个周期内计算激活值。数据通常从 MXU 流向 VPU 进行激活值处理，然后再传递到下一层。

累加器

累加器收集来自 MXU 的 32 位结果。

当 8 位输入相乘时，乘积为 16 位，但累加和会因重复加法而不断增大。使用 32 位累加器可以防止矩阵乘法所需的多次加法运算过程中发生溢出。累加器内存总共占用 4 兆字节，分布在 4096 个向量中，每个向量包含 256 个元素。

重量先进先出缓冲

权重 FIFO 缓冲区使用称为双缓冲的技术在外部存储器和 MXU 之间暂存权重。

MXU 存储两组权重块：一组正在实时计算，另一组则从内存加载。这种重叠存储方式完全消除了内存延迟，确保计算单元无需等待数据。

高带宽内存

高带宽内存随着TPU的代际发展而不断演进。

最初的v1版本采用DDR3内存，传输速度为34GB/s。而现代的Ironwood TPU则达到了7.4TB/s，速度提升了217倍。HBM通过将多个DRAM芯片垂直堆叠，并在芯片间建立数千个连接来实现如此高的带宽，这是传统内存封装方式无法实现的。

精准优势

TPU 通过量化显著提升效率，使用比传统浮点运算精度更低的数字。这种选择对硬件有着深远的影响，并波及整个设计。

科学计算通常需要高精度。例如，将圆周率计算到小数点后十位需要精确表示极小的差异。然而，神经网络的运行方式不同。它们计算的是概率和模式。例如，模型预测图像是猫的概率是 85% 还是 85.3472%，对分类结果实际上没有区别。

乘法器电路的硅片面积与位宽的平方成正比。一个 8 位乘法器大约需要 64 个硅片面积单位，而一个 32 位乘法器则需要大约 576 个硅片面积单位。这种数学关系解释了为什么 TPU v1 可以在一块尺寸适中的芯片上集成 65,536 个乘加单元，而 GPU 包含的浮点运算单元数量则少得多。更多的乘法器意味着每个周期可以执行更多的并行运算。

第一代TPU使用8位整数进行推理，与32位浮点数相比，内存需求减少了四倍。一个91兆字节的模型量化后变为23兆字节。研究表明，推理很少需要32位精度。额外的小数位对预测结果几乎没有影响。

由于微小的梯度更新会在数百万次迭代中累积，训练需要更高的精度。谷歌通过发明 BFloat16（即 Brain Floating-Point 16）解决了这个问题。这种格式与 32 位浮点数一样保留了 8 位指数，但尾数仅使用 7 位。其关键在于，神经网络对由指数控制的动态范围比对由尾数控制的精度更为敏感。BFloat16 以一半的位数提供了多种浮点格式，从而实现了高效的训练，避免了其他 16 位格式常见的溢出问题。

请看下图：

现代TPU支持多种精度模式。

  • BFloat16 用于训练。

  • 在 TPU v5e 上，INT8 推理的运行速度提高了一倍。

  • 最新的FP8格式。

Ironwood 是首款原生支持 FP8 的 TPU，避免了早期几代产品中的模拟开销。

进化之旅

TPU 的发展遵循着清晰的轨迹。

每一代人工智能都在提升性能的同时提高了能效。这一演变过程揭示了随着模型规模的扩大，人工智能硬件需求是如何变化的。

  • TPU v1 于 2015 年秘密发布，专注于推理。它采用 28 纳米制程工艺，功耗仅为 40 瓦，每秒可执行 92 万亿次 8 位运算。该芯片通过 PCIe 连接到标准服务器，并在谷歌以外的任何人知晓其存在之前，就开始为谷歌搜索、照片、翻译和 YouTube 等应用提供技术支持。2016 年 3 月，TPU v1 助力 AlphaGo 战胜李世石，证明专用芯片在速度上可以比通用 GPU 快 15 到 30 倍，在能效上可以快 30 到 80 倍。

  • TPU v2 于 2017 年发布，其架构进行了根本性变革以支持训练。它用两个 128 x 128 的 BFloat16 阵列取代了原有的 256 x 256 8 位阵列，从而实现了训练所需的浮点精度。此外，它还增加了 16 GB、600 GB/s 的高带宽内存，彻底解决了 v1 版本中存在的内存瓶颈问题。更重要的是，v2 引入了芯片间互连技术，这是一种定制的高速链路，可将各个 TPU 直接连接起来。这使得 TPU Pod 成为可能，其中 256 个芯片作为一个独立的加速器运行，可提供 11.5 petaflops 的性能。

  • 2018 年推出的 TPU v3 将性能提升了一倍，单芯片性能达到 420 万亿次浮点运算，并引入了液冷技术以应对更高的功率密度。Pod 的芯片数量扩展到 1024 个，性能超过 100 千万亿次浮点运算，足以在合理的时间范围内训练当时规模最大的 AI 模型。

  • 2021 年发布的 TPU v4 带来了多项创新。SparseCores 技术仅占用芯片面积的 5%，就将推荐系统和语言模型的关键嵌入操作速度提升了 5 到 7 倍。光路交换机实现了动态网络拓扑重构。它不再使用固定电缆，而是利用机械反射镜在光纤间引导光束。这使得互连能够绕过故障，并可扩展至 4,096 个芯片的 Pod，性能接近 1 exaflop。3D 环面拓扑结构中，每个芯片连接到六个相邻芯片而非四个，从而降低了分布式训练的通信延迟。

  • Ironwood（或称 TPU v7）将于 2025 年发布，代表着迄今为止最重大的飞跃。它专为推理时代而设计，在这个时代，大规模部署 AI 比训练更为重要。每颗芯片可提供 4,614 万亿次浮点运算的性能，配备 192 GB 的 HBM 显存，带宽高达 7.4 TB/s。

结论

TPU部署在各种应用中都展现出了实际效果。

作为参考，单个 TPU 每天可处理超过 1 亿张 Google Photos 照片。AlphaFold 团队利用 TPU 解决了困扰人们 50 年的蛋白质折叠难题，并因此荣获 2024 年诺贝尔化学奖。PaLM 是一个拥有 5400 亿个参数的语言模型，在 6144 个 TPU v4 芯片上训练，50 天内硬件利用率达到 57.8%，如此大规模的分布式训练效率令人瞩目。除了 Google 之外，TPU 还为 Anthropic 的 Claude 助手、Midjourney 的图像生成模型以及众多研究突破提供支持。

然而，TPU并非在所有方面都优于GPU。它们在处理大规模语言模型训练和推理、需要大量矩阵运算的CNN和Transformer模型、高吞吐量批量处理以及优先考虑能效的工作负载方面表现出色。另一方面，对于PyTorch原生开发而言，GPU仍然是更好的选择，因为原生开发需要PyTorch/XLA桥接，而这会带来一些不便。小批量处理、混合AI和图形工作负载、多云部署以及快速原型开发通常更适合使用GPU。

TPU 代表着整个行业向特定领域加速器转变的趋势。

通用计算模型（即CPU能够相当流畅地运行任何程序）在工作负载扩展到每次查询数万亿次运算时，就会达到物理极限。而牺牲灵活性以换取效率的专用芯片，则能带来数量级的性能提升，这是任何通用处理器优化都无法比拟的。

参考：

赞助我们

让您的产品触达超过 100 万名科技专业人士。

我们的新闻简报会将您的产品和服务直接推送给重要的受众——数十万名工程领导者和高级工程师——他们对重大的技术决策和大宗采购具有影响力。

名额有限，立即预订！

广告位通常提前约 4 周售罄。为确保您的广告触达这一极具影响力的受众群体，请立即发送电子邮件至[email protected] 预订您的广告位。

原文： https://blog.bytebytego.com/p/how-googles-tensor-processing-unit