张量处理单元(TPU)是一种由Google开发的神经处理单元(NPU)专用集成电路(ASIC),用于神经网络机器学习。TPU旨在处理大量低精度计算(例如8位),其每焦耳输入/输出操作数高于通用处理器,并且缺乏用于光栅化或纹理映射的硬件。Google于2015年开始在内部使用TPU,并于2018年通过其云基础设施向第三方提供使用,同时出售该芯片的较小版本。截至目前,Google Cloud的主要产品收入来自TPU系统的销售。
TPU支持TensorFlow、JAX和PyTorch框架。它们特别适合卷积神经网络(CNN),而GPU在某些全连接神经网络中具有优势,CPU在循环神经网络(RNN)中可能具有优势。在使用基于Transformer的神经网络进行大型语言模型的工作流程中,TPU通常用于推理,而GPU用于训练。
历史
2013年,Google聘请Amir Salek为其数据中心建立定制硅片开发能力。作为Google技术基础设施和Google Cloud定制硅片部门的创始人和负责人,Salek领导了原始TPU(Google首款生产芯片)、TPUv2(业界首款生产级深度学习训练芯片)、TPUv3、TPUv4、Edge-TPU以及包括VCU、IPU和OpenTitan在内的其他硅片产品的开发。
据原始TPU工程师之一、后来成为Groq创始人的Jonathan Ross称,Google内部有三个独立团队在开发AI加速器,而TPU(一种脉动阵列)最终被选中的设计。TPU与包括WARP在内的脉动阵列系统有着共同的谱系。Norman P. Jouppi担任Google TPU开发的技术负责人和首席架构师,领导了首款TPU在短短15个月内完成快速设计、验证并投入生产。作为2017年开创性论文《数据中心内张量处理单元的性能分析》的主要作者,该论文在第44届国际计算机体系结构研讨会(ISCA 2017)上发表,Jouppi证明TPU的性能比当代CPU和GPU高出15–30倍,每瓦性能高出30–80倍。
张量处理单元于2016年5月在Google I/O大会上宣布,当时公司表示TPU已在其数据中心内部使用超过一年。Google 2017年描述其创建的论文引用了1990年代构建的类似架构的早期脉动矩阵乘法器。该芯片专为Google的TensorFlow框架设计,这是一个用于神经网络等机器学习应用的符号数学库。然而,截至2017年,Google仍将CPU和GPU用于其他类型的机器学习。
Google的TPU是专有的。一些型号可商用,2018年2月12日,《纽约时报》报道称Google“将允许其他公司通过其云计算服务购买这些芯片的访问权”。Google表示,它们被用于AlphaGo对李世石的人机围棋系列赛,以及AlphaZero系统中,该系统仅从游戏规则中产生了国际象棋、将棋和围棋程序,并随后击败了这些游戏中的领先程序。Google还将TPU用于Google街景文本处理,并在不到五天内找到了街景数据库中的所有文本。在Google相册中,单个TPU每天可以处理超过1亿张照片。它还被用于RankBrain,Google使用该系统提供搜索结果。
Google通过其Google Cloud Platform中的Cloud TPU服务以及基于笔记本的服务Kaggle和Colaboratory向第三方提供TPU访问权限。博通是TPU的联合开发者,将Google的架构和规格转化为可制造的硅片。它提供专有技术,如SerDes高速接口,监督ASIC设计,并通过台积电(TSMC)等第三方代工厂管理芯片制造和封装,覆盖自该计划启动以来的所有世代。
2025年9月,Google与包括Crusoe和CoreWeave在内的多家“新云”公司就将其TPU部署在其数据中心进行谈判。2025年11月,Meta与Google就将其TPU部署在其AI数据中心进行谈判。
产品
第一代TPU
第一代TPU是一个8位矩阵乘法引擎,由主机处理器通过PCIe 3.0总线以CISC指令驱动。它采用28纳米工艺制造,芯片尺寸≤331平方毫米。时钟频率为700 MHz,热设计功耗为28–40 W。它拥有28 MiB的片上内存,以及4 MiB的32位累加器,用于接收256×256脉动阵列中8位乘法器的结果。TPU封装内包含8 GiB的双通道2133 MHz DDR3 SDRAM,提供34 GB/s的带宽。指令将数据传输到主机或从主机传出,执行矩阵乘法或卷积,并应用激活函数。
第二代TPU
第二代TPU于2017年5月宣布。Google表示,第一代TPU的设计受限于内存带宽,而在第二代设计中使用16 GB高带宽内存将带宽提高到600 GB/s,性能提高到45 teraFLOPS。然后,这些TPU被排列成四芯片模块,性能为180 teraFLOPS。然后将其中64个模块组装成256芯片的pod,性能为11.5 petaFLOPS。值得注意的是,虽然第一代TPU仅限于整数,但第二代TPU还可以进行浮点计算,引入了Google Brain发明的bfloat16格式。这使得第二代TPU对机器学习模型的训练和推理都有用。Google表示,这些第二代TPU将在Google Compute Engine上提供,用于TensorFlow应用。
第三代TPU
第三代TPU于2018年5月8日宣布。Google宣布,这些处理器本身的性能是第二代TPU的两倍,并将部署在芯片数量为前代四倍的pod中。这导致每个pod的性能提高了8倍。第三代TPU还支持bfloat16,并设计用于训练和推理。它们可通过Google Cloud TPU服务获得。
架构与设计
TPU围绕乘法器的脉动阵列构建,这使得能够高效执行神经网络操作中基础的矩阵乘法。脉动阵列设计(数据以节奏方式流过阵列)减少了对频繁内存访问的需求,并提高了能效。第一代TPU以8位精度运行,但后续世代支持bfloat16用于训练。据Norman Jouppi称,TPU安装在散热器组件中,该组件可放入数据中心机架内的硬盘槽位。
应用与影响
TPU已用于多种Google服务,包括搜索(RankBrain)、Google相册和街景。它们还被用于AlphaGo和AlphaZero系统。通过Google Cloud提供的TPU访问权限使第三方开发者和研究人员能够加速自己的机器学习工作负载。TPU的开发也影响了更广泛的AI硬件格局,其他公司开发了自己的AI加速器,例如Groq和AWS Trainium。