硬件革新:谷歌的智能体第八代 TPU

今天在 Google Cloud Next 大会上,谷歌正式推出了谷歌第八代定制张量处理单元(TPU),即将推出两种专为训练和推理设计的不同架构:TPU 8t 和 TPU 8i。这两款芯片旨在驱动谷歌的定制超级计算机,支持从前沿模型训练、智能体开发到大规模推理任务的一切需求。多年来,TPU 一直是包括 Gemini 在内的领先基础模型的动力源泉。第八代 TPU 将在训练、服务和智能体工作负载方面提供更强的规模、效率和能力。

在这个智能体时代,模型必须能够推理问题、执行多步工作流,并在持续的循环中从自身行动中学习。这对基础设施提出了全新的要求。TPU 8t 和 TPU 8i 是与 Google DeepMind 合作设计的,旨在应对最苛刻的 AI 工作负载,并大规模适应不断演进的模型架构。

TPU 为许多机器学习超级计算组件设定了标准,包括定制数值、液冷技术、定制互连等,而谷歌的第八代 TPU 是十多年研发的结晶。原始 TPU 设计背后的关键见解至今仍然适用:通过将硅片与硬件、网络和软件(包括模型架构和应用需求)进行定制化协同设计,谷歌可以实现显著更高的能效和绝对性能。

谷歌很高兴看到十年的创新转化为现实世界的突破。今天,像城堡证券这样具有开拓精神的组织正在挑战可能性的极限,选择 TPU 来驱动其前沿的 AI 工作负载。

两款芯片,应对当下

硬件开发周期远长于软件。在设计每一代 TPU 时,谷歌需要考虑产品上市时的技术和需求。几年前,谷歌预见到随着前沿 AI 模型在生产环境中的大规模部署,客户对推理的需求将大幅增长。随着 AI 智能体的兴起,谷歌确定社区将从分别针对训练和服务需求进行专门优化的芯片中受益。

TPU 8t 专为大规模、计算密集型的训练任务而生,具备更大的计算吞吐量和更强的扩展带宽。TPU 8i 专为对延迟最敏感的推理工作负载设计,拥有更高的内存带宽。这至关重要,因为大规模智能体之间的交互会放大哪怕是微小的低效。

重要的是,两款芯片都可以运行各种工作负载,但专业化分工释放了显著的效率和提升。

TPU 8t:训练领域的性能怪兽

TPU 8t 的目标是将前沿模型的开发周期从数月缩短至数周。通过平衡最高的计算吞吐量、共享内存和片间带宽,以及最佳的能效和有效计算时间,谷歌打造了一个系统,其单个机架(Pod)的计算性能比上一代提升了近 3 倍,从而实现更快的创新,确保谷歌的客户继续保持行业领先。

  • 巨大规模:单个 TPU 8t 超级机架(Superpod)现可扩展至 9,600 个芯片和 2PB 的共享高带宽内存,片间带宽是上一代的两倍。该架构可提供 121 ExaFlops 的算力,并允许最复杂的模型利用单个巨大的内存池。
  • 最大化利用率:通过集成 10 倍速的存储访问,并结合 TPUDirect 技术直接将数据拉入 TPU,TPU 8t 确保了端到端系统的最高利用率。
  • 近线性扩展:谷歌全新的 Virgo 网络结合 JAX 和谷歌的 Pathways 软件,意味着 TPU 8t 可以在单个逻辑集群中为多达一百万个芯片提供近乎线性的扩展。

除了原始性能外,TPU 8t 还通过一套全面的可靠性、可用性和可服务性(RAS)功能,将目标定在超过 97% 的“有效吞吐量”(Goodput),即衡量有用、生产性计算时间的指标。这些功能包括横跨数万个芯片的实时遥测、自动检测并绕过故障 ICI 链路而不中断任务,以及在无需人工干预的情况下围绕故障重新配置硬件的光路切换(OCS)。

每一次硬件故障、网络停顿或检查点重启,都是集群未进行训练的时间。在前沿训练规模下,每一个百分点的提升都可以转化为数天的活跃训练时间。

TPU 8t 超级机架架构示意图

TPU 8i:推理引擎

在智能体时代,用户期望能够提问、委托任务并获得结果。TPU 8i 旨在处理许多专业智能体之间的复杂、协作和迭代工作,这些智能体通常在复杂的流程中“成群”协作,为最具挑战性的任务提供解决方案和见解。谷歌重新设计了堆栈,通过四项关键创新消除了“等待室”效应:

  • 打破“存储墙”:为了防止处理器处于闲置状态,TPU 8i 配备了 288 GB 的高带宽内存和 384 MB 的片上 SRAM,是上一代的 3 倍,从而将模型的活动工作集完全保留在芯片上。
  • Axion 驱动的效率:谷歌将每台服务器的物理 CPU 主机数量翻倍,迁移到了谷歌定制的基于 Arm 的 Axion CPU。通过使用非统一内存架构(NUMA)进行隔离,谷歌优化了整个系统的卓越性能。
  • 扩展 MoE 模型:针对现代混合专家(MoE)模型,谷歌将互连(ICI)带宽翻倍至 19.2 Tb/s。全新的 Boardfly 架构将最大网络直径减少了 50% 以上,确保系统作为一个内聚、低延迟的单元运行。
  • 消除延迟:谷歌全新的片上集合加速引擎(CAE)可卸载全局操作,将片上延迟降低多达 5 倍,从而最大限度地减少滞后。

这些创新使性价比比上一代提高了 80%,让企业能够以同样的成本服务近两倍的客户量。

TPU 8i Boardfly 拓扑结构示意图

为 Gemini 协同设计,对所有人开放

第八代 TPU 也是谷歌协同设计理念的最新体现,每一个规格都是为了解决 AI 最大的障碍而打造。

  • Boardfly 拓扑结构专为当今最强大的推理模型的通信需求设计。
  • TPU 8i 中的 SRAM 容量是根据生产规模下推理模型的 KV 缓存占用空间量身定制的。
  • Virgo 网络矩阵的带宽目标源自万亿参数规模训练的并行性需求。

此外,这两款芯片首次均在谷歌自研的 Axion ARM 架构 CPU 主机上运行,使谷歌能够针对性能和效率优化整个系统,而不只是芯片本身。

两个平台均支持原生 JAX、MaxText、PyTorch、SGLang 和 vLLM(这些都是开发者已经在使用的框架),并提供裸金属访问权限,让客户能够直接访问硬件,而无需承担虚拟化开销。包括 MaxText 参考实现和用于强化学习支持的 Tunix 在内的开源贡献,为从能力开发到生产部署提供了关键路径。

规模化的能效设计

在当今的数据中心,电力而非芯片供应已成为主要的约束条件。为了解决这个问题,谷歌在整个堆栈优化了效率,集成了功率管理系统,可根据实时需求动态调整功耗。TPU 8t 和 TPU 8i 的每瓦性能比上一代(Ironwood)提升了多达两倍。

但在谷歌,效率不仅仅是一个芯片级的指标,它还是一个从硅片贯穿到数据中心的系统级承诺。例如,谷歌将网络连接与计算集成在同一块芯片上,显著降低了在 TPU 机架间移动数据的电力成本。甚至谷歌的数据中心也是与 TPU 协同设计的。通过硬件和软件的创新,谷歌使数据中心每单位电力提供的计算能力比五年前提高了六倍。

TPU 8t 和 TPU 8i 延续了这一轨迹。两者均由谷歌的第四代液冷技术支持,该技术能维持风冷无法企及的性能密度。通过拥有从 Axion 主机到加速器的全堆栈,谷歌可以以主机和芯片独立设计时根本无法实现的方式,优化系统级的能源效率。

Google Cloud 第四代液冷分配单元

智能体时代的基础设施

每一次重大的计算转型都需要基础设施的突破,智能体时代也不例外。基础设施必须进化,以满足自主智能体在推理、规划、执行和学习的持续循环中运作的需求。

TPU 8t 和 TPU 8i 是谷歌对这一挑战的回答:两种专门的架构旨在重新定义 AI 的可能性,从构建最强大的 AI 模型,到完美编排的智能体集群,再到管理最复杂的推理任务。两款芯片将于今年晚些时候正式上市,并可作为谷歌 AI 超级计算机的一部分使用。AI 超级计算机将专门构建的硬件(计算、存储、网络)、开放软件(框架、推理引擎)和灵活消耗模式(编排、集群管理和交付模式)整合进一个统一的堆栈中。

智能体计算将重新定义可能性。谷歌很高兴宣布谷歌为推动这一转型而进行不懈创新的最新成果:TPU 8i 和 8t。