发表时间: 2026-06 · arXiv:2606.15870 (Google)
作者/机构: Norman P. Jouppi, Sridhar Lakshmanamurthy, Cliff Young, and David Patterson, Google, LLC
本文回顾了从TPU v2到Ironwood的五代谷歌TPU训练超级计算机的发展历程,旨在展示其作为可扩展、高弹性、高能效和可持续的AI训练超级计算机的演进。
核心问题与研究目标:十年前,业界普遍怀疑像谷歌这样的软件公司是否会自研芯片。TPU v1的发布及其卓越的能效比引发了行业震动,但随之而来的是对专用集成电路(ASIC)能否适应快速迭代的AI模型的质疑。本文旨在回应这一质疑,通过详细记录TPU v2至Ironwood五代产品的发展,证明其架构的持久生命力。研究目标是阐述TPU如何在架构保持稳定的前提下,成功实现了规模、弹性、能效和可持续性的巨大提升,以适应从MLP到Transformer等不断变化的DNN工作负载。
主要创新与贡献:
1. 证明了架构稳定性:论文的核心论点是,TPU v2的初始微架构设计原则具有非凡的持久性。尽管DNN模型在过去八年中发生了巨大变化(例如Transformer模型的兴起),但TPU的基本架构并未进行根本性改变,后代产品主要通过提升组件的速度和规模来“驾驭技术突破”。这种稳定性极大地降低了为新硬件优化软件和模型的难度。
2. 展示了惊人的规模扩展:在摩尔定律放缓和登纳德缩放定律失效的时代,TPU超级计算机的系统级性能实现了3600倍的增长。论文详细记录了各项指标的扩展:HBM容量和带宽提升约10倍,单节点峰值性能提升约100倍,超级计算机规模扩大36倍。这成功跨越了所谓的“加速器墙”(Accelerator Wall)。
3. 阐述了弹性增强机制:随着系统规模扩大36倍,保持系统可靠性成为巨大挑战。论文介绍了TPU如何通过引入光学电路交换机(OCS)、功能性内置自测试(FBIST)和硬件重放(hardware replay)等技术,增强模块化隔离能力、故障绕行能力和对静默数据损坏(SDC)的抵御能力,从而在超大规模下实现超过90%的有效吞吐率(Goodput)。
4. 量化了能效和可持续性改进:面对数据中心电力供应日益紧张的挑战,论文强调了“每瓦性能”的重要性,并展示了TPU五代产品在该指标上约30倍的提升。更进一步,论文引入并推广了“计算碳强度”(CCI)这一新指标,它综合了硬件生命周期内的隐含碳排放和运营碳排放,更全面地评估了AI硬件的环境影响。
5. 总结了成功加速器的关键特征:论文最后提炼了六项可能定义本十年成功训练加速器的关键设计决策,包括收缩阵列、面向范围的窄浮点格式、HBM内存等,为未来AI硬件的设计提供了宝贵经验。
表1展示了TPU节点和超级计算机在8年间的规模扩展情况:
DNN工作负载的剧变。图1显示,DNN模型发生了迅速而剧烈的变化。如今,Transformer的变体主导了谷歌的工作负载。尽管领域特定架构存在与最新DNN趋势不匹配的潜在风险(设计、制造和部署加速器需要2-3年),但最初的TPU v2微架构在这个快速发展的领域中展示了其长期可行性。架构的稳定性降低了为新一代TPU优化软件和模型的难度。
TPU v2 核心架构。图2展示了TPU v2的框图【7,The Design Process For Google's Training Chips: TPU v2 and TPU v3,2021,Micro】,这个框图惊人地适用于从TPU v2到Ironwood的每一代训练TPU。我们首先回顾TPU v2,然后描述其演进。(接下来的几段内容源自【7】)。
双核设计。TPU v2拥有两个TensorCore。选择两个核心是在单个大核心较长的布线延迟与需要软件将众多微小核心连接起来的复杂性之间取得了恰当的平衡。通过两个大核心操作大数据块是一种更简单的编程模型。
TensorCore内部结构。TensorCore的标量单元从本地指令存储器中获取完整的322位VLIW(超长指令字)束,本地执行标量操作槽,然后将解码后的指令转发给向量和矩阵单元,以实现与标量执行解耦的后续执行。标量执行后,指令束和标量寄存器值被转发到向量单元。向量单元具有128个向量通道(lane),每个通道包含一个额外的8路数据并行维度,称为子通道(sublane),使得每个时钟周期可以对8组128宽的向量进行操作。每个通道的寄存器文件对其本地的向量内存(VMEM)切片执行加载和存储。一个异步DMA(直接内存访问)单元负责在HBM和本地向量内存之间传输数据。与一些依赖缓存的其他加速器不同,TPU的内存层次结构是由编译器控制的。由于HBM存储向量和矩阵,DMA可以跨步访问内存。DMA完成后,一个通知会送达核心的同步标志,允许程序暂停直到数据到达。
矩阵乘法单元(MXU)。MXU是TPU的计算核心。在TPU v2中,它是一个128x128的乘法器和加法器收缩阵列,每个周期可提供32,768次操作。TPU v2是首个偏离IEEE浮点标准的DNN加速器,因为谷歌认为对于DNN而言,数值范围比精度更重要。在16位Brain Float格式(BF16)中,指数(8位)首次大于其尾数(7位)。相比之下,IEEE标准的指数较小,例如FP16使用5位指数和10位尾数,FP32使用8位指数和23位尾数。许多后续的窄浮点格式都遵循了BF16扩大指数的先例。乘法以BF16格式进行,累加以完整的IEEE FP32格式进行。除了矩阵乘法,其他单元还能高效执行各种矩阵原语,如转置、行规约或列置换。
芯片间互连(ICI)。TPU v2具有四个片外链路(芯片间互连或ICI)和两个到片上路由器的片内链路。这四个链路构成了用于256个TPU的2D环面(torus)系统互连,支持如AllReduce等常见的机器学习通信模式。向其他TPU的DMA操作与向本地HBM的DMA操作类似,但为了简化设计,存在“仅推送”(push-only)的限制。这种专用的TPU互连实现了在所有TPU上进行可扩展的同步训练。
SparseCore简介。SparseCore是一种最初为嵌入(embedding)训练设计的领域特定架构【8,TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings,2023,50th ISCA】。SparseCore相对廉价,通常只占用约5%的芯片面积和约5%的功耗。它们以“核之海”(sea-of-cores)的配置方式运行,集成了超级计算机规模的HBM和ICI,创建了一个扁平的、全局可寻址的内存空间。与密集训练中大型参数张量的AllReduce不同,稀疏训练中较小嵌入向量的all-to-all传输利用HBM和ICI,通过更细粒度的访问模式进行散布/收集(scatter/gather)。作为独立的核,SparseCore可以在密集计算、SparseCore自身操作和ICI通信之间实现并行化。我们认为SparseCore是一种“数据流”架构,因为数据从内存流向各种专门的计算单元。
SparseCore组件。SparseCore单元包含16个计算瓦片(compute tile)。每个瓦片都有一个关联的HBM通道,并支持多个未完成的内存访问。每个瓦片还包括一个提取单元(Fetch Unit)、一个可编程的8路SIMD向量处理单元(Vector Processing Unit)和一个写回单元(Flush Unit)。提取单元从HBM读取激活值和参数到该瓦片所属的2.5 MiB稀疏向量内存(Sparse Vector Memory)切片中。在反向传播过程中,写回单元将更新后的参数写回HBM。与TPU v1类似,这些单元执行类CISC指令,并对可变长度的输入进行操作,其中指令执行时间依赖于数据。
SparseCore的应用演进。SparseCore最初是为深度学习推荐模型中的稀疏嵌入表而开发的,这些模型用于广告、搜索排名、YouTube和Google Play等应用。在2016年,这些模型占TPU v1工作负载的61%,到2022年,在TPU v4上仍占约25%【8】。随着Transformer模型的兴起(到2022年占工作负载的60%),SparseCore也开始承担新的角色,作为卸载引擎,用于处理集合操作(如AllReduce、AllGather、ReduceScatter和Broadcast)、数据摘要操作(如Top-K)以及小型稀疏张量操作(如Transformer解码)。此外,当TensorCore处理Transformer的前向传播路径中的密集注意力计算时,SparseCore可以通过并行操作来提升性能。
主机与存储连接。我们通过数据中心网络将TPU超级计算机连接到存储,以便通过一个PCIe连接的CPU主机为模型提供输入数据。在CPU、网络和存储之间保持系统平衡对于实现大规模端到端性能至关重要。
TPU软件栈从TPU v2到Ironwood一直在演进,其中XLA(加速线性代数)框架是唯一不变的常量。早期的TPU由TensorFlow驱动,通过一个“桥接器”将TensorFlow图转换为XLA的高级优化器(HLO)格式。在XLA内部,“融合”(fusions,类似于指令级并行编译器中的区域形成)技术允许跨多个操作进行优化,以节省内存访问次数。如今,JAX(即时自动微分XLA)已成为编程TPU的首选语言和系统,而Pallas内核语言为模型开发者增加了细粒度控制。
图3展示了五代TPU板卡以及最近三代产品的封装。它揭示了一个显著的稳定性:尽管发生了巨大变化,每块板上始终集成四颗TPU。前三张图捕捉了从风冷到液冷的转变,以及从液冷循环回路到分配和收集歧管的切换。不太明显的是在供电和稳压(包括垂直供电)、芯片尺寸、小芯片(chiplet)数量以及封装复杂性方面的进步。
TPU微架构跨代演进主要体现在组件的规模和数量上,而不是像其他加速器那样引入新的微架构特性:
TensorCores:每一代训练TPU都使用两个物理TensorCore,它们仅共享HBM。自TPU v4起,我们的XLA(加速线性代数)编译器支持张量并行化指令,这给出了一个单一大型核心的错觉——称为Megacore——在单个有效线程中统一了HBM容量和ICI带宽【9,JAX: TPU Pipelining,2024,The JAX Authors】。
矩阵乘法单元(MXU):收缩阵列被证明是基础构建模块,大多数加速器要么从一开始就采用,要么最终集成了它。利用从16纳米到更精细工艺带来的逻辑密度提升,MXU从TPU v2的两个128x128收缩阵列扩展到Ironwood的四个256x256(用于bf16)阵列。Ironwood还增加了对FP8算术的支持,这意味着它也可以计算四个512x512的FP8乘法。类似于为提高内存良率和降低成本而增加冗余行,Ironwood在MXU中增加了一个冗余行。
VPU(向量处理单元):逻辑密度的提升使得VPU从每个向量通道两个受限的ALU(每个只能执行部分ALU操作)演进为每个通道四个通用ALU。向量寄存器在两个维度上也翻了一番,从TPU v2–v5p的8x128演进到Ironwood的16x256。该向量架构非常适合非矩阵的DNN操作,如激活函数、softmax、批归一化以及用于较低精度数值的量化函数。
VLIW指令:通过VLIW指定指令级并行性效果很好。随着TPU获得更多并行硬件(如更多MXU)需要控制,我们只需在每一代中加宽VLIW指令;Ironwood的指令比TPU v2宽50%以上。DNN程序的大小并非大到无法接受,使得额外的指令内存成本低廉,而CPU所需的向后二进制兼容性对于加速器来说是多余的。
SparseCores:其架构变得更加通用,并且每代性能都有所提升——例如,从TPU v5p到Ironwood提升了2.4倍——此外,其数量从TPU v2的两个翻倍到TPU v4的四个,Ironwood也保持四个。
VMEM(向量内存):SRAM的密度增长慢于逻辑密度。因此,尽管芯片面积大幅增加,其大小仅从TPU v2的每节点32 MB翻了两番,达到TPU v5p和Ironwood的128 MB。
HBM(高带宽内存):与MXU的收缩阵列类似,HBM被证明是主存的明智选择。标准DRAM是TPU v1的瓶颈【1,In-data center Performance Analysis of a Tensor Processing Unit,2017,44th International Symposium on Computer Architecture (ISCA)】,因此TPU v2通过HBM将内存带宽提高了30倍。八年来,TPU的容量和带宽又扩展了十倍,从TPU v2使用4个HBM2堆栈的16 GiB @ 700 GB/sec,到Ironwood使用8个HBM3E堆栈的192 GiB @ 7300 GB/sec。再次证明,大多数加速器要么最初就使用HBM,要么最终采用了它。
ICI(芯片间互连):从TPU v2的每节点4个外部链路(速率62 GB/sec),形成2D环面互连,发展到TPU v4、TPU v5p和Ironwood的每节点6个外部链路(速率100 GB/sec),支持3D环面。与TPU v2相比,Ironwood的节点数增加了36倍,对剖带宽增加了39倍。
超级计算机规模:从TPU v2的256个节点增长到Ironwood的9216个节点(从0.25K到9K)。
弹性策略。谷歌长期以来使用分布式系统的商用组件来构建可靠的计算服务。我们对可靠训练超级计算机的方法遵循了高性能计算(HPC)领域的做法,因为两者都涉及运行大型、长时间的批处理作业,而非交互式服务。关键特性包括:
1. 增强的节点质量:减少执行期间发生故障的几率。
2. 错误检查:在执行期间发现错误。
3. 检查点/恢复:在发生故障后恢复长时间运行的批处理作业的继续执行。
4. 严格的确定性可重复性要求:以帮助系统测试和故障检测。
5. 模块化隔离:允许超级计算机在部分节点发生故障时仍能继续服务。
有效吞吐率(Goodput)。谷歌在多个数据中心的多个8960芯片TPU v5p pod上使用同步数据并行训练来并行化Gemini 2.5的训练,有效吞-吐率(goodput)达到了93%【15,Gemini 2.5: Pushing the Frontier,2025,Gemini Team】。在较小规模上,Gemini 1.0在TPU v4上的一个类似指标为97%【14,Gemini: A Family of Highly Capable Multimodal Models,2023,Gemini Team】。有效吞吐率是“良好吞吐率”的简称,在训练系统中指有效或有效果的训练进度速率。例如,一个系统在正常运行时训练吞吐量为X,但如果系统花费10%的总时间用于从错误或故障中恢复,那么其有效吞吐率将是0.9X。
光学电路交换机(OCS)。从TPU v4开始,通过使用光学电路交换机(OCS),第5点(模块化隔离)得到了显著改进【8】。(以下几段源自该论文)。为了增强数据中心网络,谷歌推动了基于3D微机电系统(MEMS)反射镜的光收发器在可靠性和成本方面的技术进步,这种反射镜能在毫秒级内切换。在机架内,电气连接提供了最佳的性价比。问题在于应使用多大规模的电气布线构建块?考虑到3D环面拓扑,3D立方体具有最佳的对剖带宽,这表明4×4×4(64芯片)或8×8×8(512芯片)是合适的选择。由于每个CPU主机带4个TPU,64个TPU芯片及其16个CPU主机可以舒适地容纳在一个机架中。而512个芯片需要多个机架,因此选择了4x4x4(即$4^3$)的构建块;我们称之为“立方体”(cubes)。
OCS的连接方式。图4展示了一个$4^3$立方体6个“面”的链路。每个面有16个链路,每个立方体总共有96个光纤链路连接到OCS。为了创建3D环面的环绕链路,相对面上的链路必须连接到同一个OCS。因此,每个立方体连接到 6 × 16 ÷ 2 = 48个OCS。TPU v4的OCS具有136×136端口(128个端口加上8个用于链路测试和修复的备用端口),因此48个OCS连接到来自64个立方体(每个立方体64个芯片)的48对电缆,从而构成了所需的4096个TPU v4芯片。与HPC超级计算机类似,工作负载由不同规模的“切片”(slices)组成,例如64、128、...、2048个芯片。
OCS的角色与优势。ICI,即TPU超级计算机互连,依赖于3D环面拓扑,每个TPU芯片都包含一个分布式路由器;它不需要额外的芯片进行TPU间的通信。OCS提供了一个位于ICI之下的物理层,通过绕过故障来增强可用性。TPU超级计算机主要的可用性挑战来自CPU主机;每个主机带4个TPU,意味着一个Ironwood超级计算机有2304个CPU主机。如果没有OCS,主机的可用性必须大于99.9%才能实现高切片有效吞吐率。
OCS简化调度与部署。OCS还简化了调度,从而提高了利用率。对于TPU v2和TPU v3,调度一个128芯片的切片需要调度器找到128个连续的空闲芯片。有了OCS,调度器可以从超级计算机的任何地方选择两个$4^3$立方体。调度难度随切片大小急剧增加。这种可用性的提高和调度的简化是Ironwood拥有9K节点而不是2的幂次方的原因。Ironwood可以运行四个流行的2K切片作业(每个需要32个立方体),即使某些节点宕机,因为还有16个备用立方体可用作替代。OCS还缩短了部署时间【8】。TPU v3系统直到所有1024个芯片和所有电缆都安装并测试完毕后才能使用。任何组件的交付延迟都会拖延整个超级计算机的进度。从TPU v4开始,OCS使每个机架独立,因此每个立方体在64个芯片和必要的电缆安装测试完毕后即可投入生产。增量部署极大地缩短了投入生产使用的时间,从而提高了TPU超级计算机的成本效益。
对抗静默数据损坏(SDC)。计算逻辑中的静默数据损坏(SDC)对大规模AI的可靠性构成了关键挑战【13,Silent Data Corruption in AI,2026,IEEE Micro】。由制造逃逸或硅老化引起的边缘缺陷会悄无声息地降低模型质量和收敛性。虽然前几代TPU依赖于基于软件的健康检查和在位工作负载监控器【15】,但Ironwood将这些缓解措施扩展到了硬件中:
1. 功能性内置自测试(FBIST):集成在MXU内的FBIST引擎,在制造和数据中心老化测试期间执行高覆盖率的功能测试模式,以拦截通过结构测试的芯片,防止其进入机群,并在运行期间发现新出现的有缺陷芯片。
2. 硬件重放单元:为了对抗由环境压力(如电压波动、温度变化或特定数据模式)引起的计算数据通路中的间歇性错误,Ironwood为VPU引入了硬件重放单元。这个对编译器透明的单元在VLIW指令现有的空闲槽内,随机抽样向量束进行机会性的冗余执行。通过在偶数通道上重放奇数通道的操作而不改变架构状态,该机制以零性能开销和可忽略的功耗影响提供了有效的错误检测。在整个生产机群中部署后,这种在位监控持续识别出所有其他筛选方法都未能发现的缺陷单元。一旦识别,这些单元会立即通过OCS从生产中移除,并随后进行修复。
能效指标的转变。Vahdat等人【10,New Computer Evaluation Metrics for a Changing World,2024,CACM】建议加速器设计者考虑每工作负载有效吞吐率的平均功耗,而不仅仅是每总拥有成本(TCO)的基准性能。这一转变是由为新数据中心获取足够电力的难度日益增加所驱动的,鼓励最大限度地利用可用电力。如今,每瓦性能比每TCO性能更受重视。
TPU的能效提升。图5展示了TPU各代产品在每瓦性能上的持续改进,尽管它使用的是峰值性能除以TDP(热设计功耗)瓦数,而不是Vahdat等人建议的运行生产工作负载时测量的性能和功耗。正如我们接下来将看到的,运营排放在TPU总体碳足迹中的重要性凸显了这些收益的重要性。
关注碳排放。Vahdat等人还建议关注每工作负载有效吞吐率的碳排放【10】。为此,谷歌最近完成了对几款TPU的生命周期评估(LCA)【12,An Introduction to Life-Cycle Emissions of AI hardware,2024,Micro】。LCA是对与TPU硬件相关的温室气体(GHG)排放的全面分析,包括从原材料提取到制造和能源使用的整个硬件生命周期。(本节内容源自该论文。)
计算碳强度(CCI)。一个关键挑战是如何平衡新一代TPU增加的制造成本和功耗与其不断增长的性能。新一代TPU可能每秒消耗更多电力,但训练所需的时间也更少。答案是一个新指标:计算碳强度(CCI)。它量化了每单位已用浮点运算的二氧化碳当量排放量(CO2e),即CO2e/FLOP。CCI的显著优势在于它很容易整合隐含碳排放和运营碳排放,即总CCI = 运营CCI + 隐含CCI。与CCI不同,每瓦性能不包括隐含排放。
CCI结果分析。图6显示了三代训练TPU的CCI【12】【6,AI infrastructure efficiency: Ironwood TPUs deliver 3.7x carbon efficiency gains,2026,Google Cloud Blog】。TPU v5p在运营和总体CCI上比TPU v4好1.1倍,在隐含CCI上好1.3倍。Ironwood在运营CCI上有更大的飞跃,约为3.7倍,正如-图5所暗示的,在隐含CCI上约为3.8倍。
运营排放与隐含排放。对于这三代TPU,运营CCI约占总体CCI的75%,这反映了数据中心AI加速器相对较高的功耗和较长的使用寿命。不出所料,鉴于移动设备的低功耗和短寿命,情况正好相反:智能手机87%的排放是隐含排放【11,Energy and Emissions of Machine Learning on Smartphones vs. the Cloud,2024,CACM】。
CCI与每瓦性能的关系。顺便说一下,运营CCI包含了每瓦性能,因为:
其中,电力排放因子代表每消耗单位电力所排放的温室气体量(例如,gCO2e/kWh)。
CCI的应用。给定任务的FLOPs,CCI还可以提供排放量的大致估算。例如,训练GPT-3大约需要$3.14 \times 10^{23}$ FLOPs,而TPU v5p的CCI是$265 \times 10^{-18}$ gCO2e/FLOP,所以排放量就是它们的乘积:约$83 \times 10^6$ gCO2e,即约8300万公吨二氧化碳当量排放(mtCO2e)。每瓦性能仅有助于计算运营排放,而CCI则明确了每瓦性能和隐含排放,并使估算碳排放变得容易。
本论文并未设置传统的独立实验章节,其关键的性能与评估数据分布于全文各部分。
系统规模与性能扩展(表1):
系统弹性与效率:
能效与可持续性:
本文回顾了谷歌从TPU v2到Ironwood(TPU v7)五代训练TPU的演进历程和持久效能。在八年时间里,尽管登纳德缩放定律终结且摩尔定律放缓,TPU系统仍在规模、弹性、能效和可持续性方面取得了巨大进步:HBM容量和带宽增加了10倍,超级计算机规模扩大了36倍,TPU节点峰值性能提升了100倍,共同促成了超级计算机峰值性能3600倍的增长。
关键技术与设计选择:
架构稳定性与未来展望:
尽管最初存在对ASIC寿命的疑虑,TPU v2的微架构表现出卓越的稳定性和适应性,成功应对了十年间从Transformer到Diffusion等多样化工作负载的演变。这种稳定性使得模型、编译器和软件栈的优化成果可以轻松地在新一代TPU上复用。
论文总结了六项被证明具有持久价值并被其他AI加速器广泛采用的关键初始决策:
1. 收缩阵列用于矩阵乘法。
2. 采用面向范围的窄浮点算术(如BF16, FP8)而非面向精度的宽IEEE浮点算术(如FP16, FP32)。
3. 使用HBM作为主内存。
4. 通过定制高速链路(ICI)将AI加速器组装成AI超级计算机。
5. 通过DMA和暂存SRAM实现软件控制的内存层次结构,而非传统的类CPU缓存层次结构。
6. 使用向量单元执行非矩阵计算。
此外,OCS和SparseCore仍是TPU的独特创新。论文推测,上述六大特征可能将成为2020年代训练加速器的定义性特征。
TPU的五大优势:
最终,TPU的成功可归结为五大优势的强力组合:
1. 简化的编程模型:每TPU仅有两个大型处理器(可统一为更大的Megacore),简化了新模型的开发。
2. 架构稳定性与软硬件协同设计:确保新TPU发布时即有优化好的模型可用。
3. 万级节点规模与高弹性:支持在单层网络上以超过90%的有效吞吐率进行大规模、长时间的模型训练。
4. 持续迭代的记录:五代产品在计算、内存、带宽和系统规模上持续增长。
5. 减少环境影响:通过降低每浮点运算的碳排放实现可持续发展。
这些优势使TPU能够持续满足AI发展的计算需求,并以可持续的方式释放其巨大潜力。