谷歌在拉斯维加斯举办的年度Cloud Next大会上发布了两款全新自研AI加速器:一款专为提升训练速度而设计,另一款则致力于降低模型推理成本。
谷歌表示,第八代张量处理单元(TPU)在训练性能上最高提升2.8倍,与上一代Ironwood TPU相比,大语言模型推理的性价比提高了80%。
为实现这一目标,谷歌将加速器研发分为两条并行路线,分别打造了面向训练的TPU 8t与面向推理的TPU 8i。两款芯片虽基于相似架构,但各自针对不同工作负载的瓶颈进行了专项优化。
谷歌并非首家走这条路的企业。亚马逊云服务(AWS)早在AI芯片研发初期,就意识到推理与训练需要分别优化的加速器。英伟达也尝试过类似的差异化设计,尽管程度不及谷歌彻底。其Blackwell Ultra系列专为AI推理优化,相较于上一代Blackwell产品,内存容量提升50%,FP4计算能力也有所增强,但牺牲了部分高精度算力。
除了双轨并行策略,谷歌还将TPU主机处理器从x86架构切换为自研的基于Arm架构的Axion CPU。亚马逊今年早些时候也在Graviton与Trainium 3的组合中采用了类似做法。
谷歌的专业化思路远不止于芯片本身。这家搜索与广告巨头为不同工作场景开发了具有独特网络拓扑结构的计算集群,以最大程度降低推理与训练任务在扩展过程中的性能损耗。
如今,现代AI工作负载几乎不会在单个加速器上运行,因此跨多芯片高效扩展的能力,往往比单颗芯片的速度更为关键。
TPU 8t:专为大规模训练而生
与此前历代加速器相同,TPU 8t专为超大规模训练而设计。谷歌或许不像OpenAI或Anthropic那样备受关注,但其依然是业内高产的模型构建者之一。
在硬件层面,谷歌对Ironwood中引入的向量运算、矩阵乘法与SparseCore嵌入加速器的组合进行了调整,以最大化有效浮点吞吐量。
每颗TPU 8t配备216 GB高带宽内存(HBM),带宽达6.5 TB/s,128 MB片上SRAM,4位浮点算力最高达12.6 PetaFLOPS,芯片间互联带宽最高19.2 Tbps。
与谷歌同期宣布将在云基础设施中部署的英伟达Rubin GPU相比,TPU 8t的参数略显保守。Rubin的FP4训练性能最高可达35 PetaFLOPS,搭载288 GB HBM4,带宽高达22 TB/s。从单颗芯片来看,英伟达的性能更胜一筹,但这并不重要——训练一个前沿模型从来不只用一颗GPU,而是同时调用数以千计的算力资源。在规模扩展方面,谷歌才是真正的优势方。
英伟达最新GPU在单个NVLink域内最多支持576个加速器,超出后需通过以太网或InfiniBand进行扩展。TPU 8t则采用光路交换技术,这是一种类似于电话交换机的光机械切换方案,可在单个统一计算集群中连接多达9,600个加速器。
多个集群再通过全新的Virgo网络互联,支持更大规模的计算域。谷歌采用的并非单纯扩大芯片间互联网格的思路,而是使用具备极高端口密度的包交换技术,构建了近似扁平化的两层全互联拓扑,据称可在单个数据中心内连接最多134,000颗TPU,跨站点时更可达百万颗TPU规模。
与新网络架构协同工作的,还有谷歌自研的托管式Lustre存储系统,可提供高达10 TB/s的聚合数据带宽,直接向加速器内存供给数据。
谷歌表示,上述技术结合改进后的可靠性、可用性与可维护性(RAS)能力,能够使训练优化型TPU实现高达97%的"有效利用率"。在此语境下,有效利用率指TPU真正用于训练任务的时间占比。
谷歌在一篇博客文章中解释道:"每一次硬件故障、网络阻塞或检查点重启,都是集群停止训练的时间损耗。在前沿训练规模下,每一个百分点的差异都可能意味着数天的有效训练时间。"
TPU 8i:推理场景的全面优化
推理是一种自回归工作负载,这意味着每生成一个Token,都需要将整个模型的激活权重从内存中流式读取一遍。虽然算力依然重要,但真正的瓶颈往往在于内存带宽。
在推理专用的TPU 8i上,谷歌以牺牲部分FLOPS为代价,换来了更大容量的SRAM缓存与更快、更大的内存池。
从纸面数据来看,TPU 8i与英伟达的Blackwell系列基本处于同一水平:配备FP4算力10.1 PetaFLOPS,片上SRAM达384 MB,288 GB HBM带宽8.6 TB/s,芯片间互联带宽与训练版TPU 8t保持一致。
谷歌表示,大容量SRAM使TPU 8i能够将更多的键值缓存(即模型的短期记忆)保留在片上,从而减少计算核心等待数据的时间。
TPU 8i还以集合加速引擎(CAE)取代了谷歌的SparseCore设计。顾名思义,CAE通过卸载"全归约"或"全聚集"等集合通信操作来加速推理,缩短同步停顿时间,避免芯片的张量核心长时间处于空闲状态。
事实上,谷歌在TPU 8i上的核心设计目标,就是尽可能消除推理链路中的各类延迟。
随着混合专家(MoE)架构的普及,集合通信的开销变得愈发突出。MoE模型由多个子模型(即"专家")组成,每次生成Token时只激活其中的一个子集。这种设计的优势在于,模型可以持续扩大规模而无需相应增加内存带宽;劣势则在于芯片间通信的可预测性较差——生成不同Token时可能调用不同的专家,而这些专家可能分布在不同的加速器上。
谷歌表示,CAE将集合通信延迟降低了5倍,使同一硬件能够承载更多用户,从而显著提升了推理经济性。
在推理优化芯片之外,谷歌还开发了一种名为Boardfly的新型网络拓扑结构。该架构与高性能计算集群中常见的Dragonfly拓扑有相似之处,借助光路交换技术,可将1,152颗芯片(同一时刻有效激活1,024颗)互联成一个计算域。Boardfly的核心优势在于,将最大芯片间跳数从3D环面拓扑的16跳降至仅7跳,进一步降低了运行MoE或推理模型时的通信延迟。
谷歌并非唯一一家为适应新兴AI架构而调整推理集群设计的厂商。去年12月,AWS出于类似原因放弃了推理场景下的3D环面拓扑,转而采用更传统的包交换网络架构。
TPU 8两款加速器预计将于今年晚些时候在谷歌云平台上作为实例正式上线,也将作为该云厂商全栈式AI超算(AI Hypercomputer)平台的组成部分推出。该平台集网络、存储、计算与软件于一体,支持大语言模型的大规模部署与训练。
Q&A
Q1:谷歌TPU 8t和TPU 8i有什么区别?
A:TPU 8t专为大规模模型训练设计,配备216 GB HBM和高达12.6 PetaFLOPS的FP4算力,并通过光路交换技术支持最多9,600个加速器互联;TPU 8i则专为推理优化,以较低的FLOPS换取384 MB超大片上SRAM和288 GB HBM,同时引入集合加速引擎(CAE)将集合通信延迟降低5倍,两者各司其职,分别针对不同工作负载的核心瓶颈进行了专项优化。
Q2:谷歌TPU 8的训练集群规模能达到多大?
A:谷歌TPU 8t通过光路交换技术,可在单个统一计算集群中连接最多9,600个加速器,多个集群再经由新型Virgo网络互联。在拓扑层面,谷歌采用近似扁平化的两层全互联架构,据称单个数据中心内最多可连接134,000颗TPU,跨多个数据中心站点时更可扩展至百万颗TPU规模,大幅超越英伟达NVLink域最多576个加速器的单域上限。
Q3:Boardfly网络拓扑在推理场景下有什么优势?
A:Boardfly是谷歌为TPU 8i专门开发的网络拓扑,灵感来源于高性能计算中的Dragonfly架构。它利用光路交换将1,152颗芯片互联,并将最大芯片间跳数从传统3D环面拓扑的16跳压缩至7跳,显著降低了通信延迟。这对混合专家(MoE)模型尤为关键,因为MoE在生成不同Token时会激活分布在不同加速器上的专家,跳数减少意味着延迟更低、推理效率更高。