在谷歌，我們（指代本文作者谷歌，下同）的 TPU 設(shè)計理念始終圍繞三大支柱：可擴展性、可靠性和效率。隨著人工智能模型從密集型大型語言模型 (LLM) 演進到大規(guī)模混合專家模型 (MoE) 和推理密集型架構(gòu)，硬件不僅要提升每秒浮點運算次數(shù) (FLOPS)，還必須不斷演進以滿足最新工作負載的特定運算強度。
智能體人工智能的興起需要能夠處理長上下文窗口和復(fù)雜序列邏輯的基礎(chǔ)設(shè)施。與此同時，世界模型作為當前基于數(shù)據(jù)序列的架構(gòu)的必然演進而出現(xiàn)，這意味著新一代智能體能夠模擬未來場景、預(yù)測后果，并通過“想象”而非冒險的試錯法進行學習。第八代TPU（TPU 8t和TPU 8i）正是我們應(yīng)對這些挑戰(zhàn)的解決方案，確保從訓(xùn)練的第一個標記到多輪推理鏈的最后一步，每個工作負載都能以最高效的方式運行。它們旨在高效地訓(xùn)練和運行諸如谷歌DeepMind的Genie 3之類的世界模型，使數(shù)百萬智能體能夠在各種模擬環(huán)境中練習和完善其推理能力。
TPU 8：專為特定設(shè)計而生
鑒于預(yù)訓(xùn)練、后訓(xùn)練和實時服務(wù)的基礎(chǔ)設(shè)施需求已然分化，我們的第八代 TPU 引入了兩個截然不同的系統(tǒng)：TPU 8t 和 TPU 8i。這些新系統(tǒng)是 Google Cloud AI 超級計算機的關(guān)鍵組件，AI 超級計算機是一種集成的超級計算架構(gòu)，它結(jié)合了硬件、軟件和網(wǎng)絡(luò)，為完整的 AI 生命周期提供支持。雖然這兩個系統(tǒng)都繼承了 Google AI 技術(shù)棧的核心 DNA 并支持完整的 AI 生命周期，但它們各自針對不同的瓶頸進行設(shè)計，并針對關(guān)鍵開發(fā)階段優(yōu)化效率。此外，通過在第八代 TPU 系統(tǒng)中集成基于 Arm 的 Axion CPU 接口，我們消除了由數(shù)據(jù)準備延遲造成的主機瓶頸。Axion 提供了足夠的計算能力來處理復(fù)雜的數(shù)據(jù)預(yù)處理和編排，從而確保 TPU 始終有充足的資源，避免出現(xiàn)卡頓。
TPU 8t：pre-training王者
TPU 8t 針對大規(guī)模預(yù)訓(xùn)練和嵌入密集型工作負載進行了優(yōu)化，它采用我們成熟的 3D 環(huán)面網(wǎng)絡(luò)拓撲結(jié)構(gòu)，并在單個超級節(jié)點中集成了 9600 個芯片，規(guī)模更大。TPU 8t 旨在實現(xiàn)數(shù)百個超級節(jié)點的最大吞吐量，確保訓(xùn)練運行按計劃進行。
以下是TPU 8t相對于上一代TPU的一些主要改進：
SparseCore 的優(yōu)勢：TPU 8t 的核心是 SparseCore，這是一種專門設(shè)計的加速器，用于處理嵌入查找中不規(guī)則的內(nèi)存訪問模式。矩陣乘法單元 (MXU) 負責處理矩陣運算，而 SparseCore 則負責卸載數(shù)據(jù)相關(guān)的全集操作以及其他一些集中操作，從而避免了通用芯片中常見的零操作瓶頸。
VPU/MXU 重疊和均衡擴展：TPU 8t 旨在最大限度地利用已分配的浮點運算能力。通過實現(xiàn)更均衡的向量處理單元 (VPU) 擴展，該架構(gòu)最大限度地減少了暴露的向量運算時間。這使得量化、softmax 和層范數(shù)運算與 MXU 中的矩陣乘法運算能夠更好地重疊，從而幫助芯片保持繁忙狀態(tài)，而不是等待順序向量任務(wù)。
原生 FP4 ：TPU 8t 引入了原生 4 位浮點運算 (FP4)，以克服內(nèi)存帶寬瓶頸，在保持大型模型精度（即使在低精度量化下）的同時，將 MXU 吞吐量提升一倍。通過減少每個參數(shù)的位數(shù)，該平臺最大限度地減少了耗能的數(shù)據(jù)傳輸，并允許更大的模型層適應(yīng)本地硬件緩沖區(qū)，從而實現(xiàn)最佳計算利用率。
Virgo Network 拓撲結(jié)構(gòu)及高達 4 倍的數(shù)據(jù)中心網(wǎng)絡(luò)帶寬提升：為了支持 TPU 8t 的海量數(shù)據(jù)需求，我們推出了 Virgo Network。這種全新的網(wǎng)絡(luò)架構(gòu)使 TPU 8t 訓(xùn)練的數(shù)據(jù)中心網(wǎng)絡(luò) (DCN) 帶寬提升高達 4 倍。Virgo Network 是一種橫向擴展（scale-out）架構(gòu)，專為滿足現(xiàn)代 AI 工作負載的極端需求而設(shè)計。它基于高基數(shù)交換機構(gòu)建，通過增加每個交換機的端口數(shù)量來減少網(wǎng)絡(luò)層數(shù)，并采用扁平化的兩層無阻塞拓撲結(jié)構(gòu)。與傳統(tǒng)數(shù)據(jù)中心網(wǎng)絡(luò)相比，這種架構(gòu)通過最大限度地減少網(wǎng)絡(luò)層數(shù)，顯著降低了延遲。它采用多平面設(shè)計，并具有獨立的控制域來連接 TPU 8t 芯片。TPU 8t 機架還連接到 Jupiter 南北向網(wǎng)絡(luò)架構(gòu)，以訪問計算和存儲服務(wù)。這種精簡的架構(gòu)共同提供了海量的二分帶寬和確定性的低延遲，從而能夠支持全球最大規(guī)模的高可用性訓(xùn)練集群。
與上一代產(chǎn)品相比，TPU 8t 的芯片間互連 (ICI) 帶寬提升了 2 倍，原始 DCN 橫向擴展帶寬提升了高達 4 倍，從而顯著降低了數(shù)據(jù)瓶頸。為了進一步加速前沿模型的開發(fā)，我們將分布式訓(xùn)練擴展到單個集群之外。借助JAX和Pathways，我們現(xiàn)在可以在單個訓(xùn)練集群中擴展超過 100 萬個 TPU 芯片。Virgo Network 能夠在單個網(wǎng)絡(luò)架構(gòu)中連接超過 134,000 個 TPU 8t 芯片，提供高達 47 PB/s 的無阻塞雙向帶寬。該架構(gòu)可提供超過 160 萬 ExaFlops 的計算能力，并具有近乎線性的擴展性能。
更快的存儲訪問：我們在 TPU 8t 中引入了TPUDirect RDMA和TPU Direct Storage。TPU Direct RDMA 支持 TPU 內(nèi)存（HBM）與網(wǎng)絡(luò)接口卡（NIC）之間的直接數(shù)據(jù)傳輸，繞過主機 CPU 和 DRAM。這降低了延遲和主機系統(tǒng)瓶頸，提高了 TPU 間通信的有效帶寬。同樣，TPUDirect Storage 通過實現(xiàn) TPU 與高速托管存儲（例如 10T Lustre）之間的直接內(nèi)存訪問，繞過了 CPU 主機瓶頸，有效地將海量數(shù)據(jù)傳輸?shù)膸捥岣吡艘槐丁＿@種架構(gòu)使芯片能夠以線速攝取訓(xùn)練數(shù)據(jù)，確保即使在處理大型多模態(tài)數(shù)據(jù)集時，MXU 也能保持滿負荷運行。
通過結(jié)合Managed Lustre 10T和 TPUDirect Storage，將百 PB 級數(shù)據(jù)集直接路由到芯片，TPU 8t 避免了因數(shù)據(jù)攝取瓶頸導(dǎo)致的訓(xùn)練延遲。與在第七代 Ironwood TPU 上進行訓(xùn)練相比，這可實現(xiàn) 10 倍的存儲訪問速度提升。
TPU 8i：采樣（sampling ）和服務(wù)（serving）專家
TPU 8i 針對訓(xùn)練后處理和高并發(fā)推理進行了優(yōu)化，采用了我們最高的片上 SRAM、新的集體加速引擎 (CAE) 和名為 Boardfly 的新型服務(wù)優(yōu)化網(wǎng)絡(luò)拓撲結(jié)構(gòu)。
大型片上 SRAM：與上一代相比，TPU 8i 的片上 SRAM 容量增加了 3 倍，因此可以完全在硅片上容納更大的 KV 緩存，從而顯著減少長上下文解碼期間內(nèi)核的空閑時間。
集體加速引擎 (CAE：Collectives Acceleration Engine ) ：為了解決采樣瓶頸，TPU 8i 采用了 CAE，它能夠以近乎零延遲聚合跨核心的結(jié)果，從而顯著加速自回歸解碼和“思維鏈”處理過程中所需的歸約和同步步驟。每個 TPU 8i 芯片都包含兩個位于核心芯片上的張量核心 (TC) 和一個位于芯片組芯片上的 CAE，取代了上一代 Ironwood TPU 中位于核心芯片上的四個稀疏核心 (SC)。通過集成專用 CAE，TPU 8i 將集體操作的片上延遲進一步降低了 5 倍。每次集體操作的延遲越低，等待時間就越短，從而直接提高了運行數(shù)百萬個代理所需的吞吐量。
Boardfly ICI拓撲結(jié)構(gòu)：雖然3D環(huán)面結(jié)構(gòu)允許連接數(shù)千個芯片以協(xié)同工作，但大型網(wǎng)狀結(jié)構(gòu)會增加芯片間的跳數(shù)，從而導(dǎo)致更高的全連接延遲。在8i版本中，我們改變了芯片的連接方式，采用全連接板，然后將這些板聚合成組。利用高基數(shù)設(shè)計，我們最多可以連接1152個這樣的芯片，從而減小網(wǎng)絡(luò)直徑，并減少數(shù)據(jù)包穿越系統(tǒng)所需的跳數(shù)。通過大幅減少全連接通信（MoE和推理模型的核心）所需的跳數(shù)，Boardfly在通信密集型工作負載下實現(xiàn)了高達50%的延遲降低。
Boardfly 由以下元素組成，其拓撲結(jié)構(gòu)本質(zhì)上是分層的：
構(gòu)建模塊 (BB：Building Block )：每個托盤使用內(nèi)部 ICI 連接形成一個四芯片環(huán)，提供 16 個外部連接，以實現(xiàn)更廣泛的聯(lián)網(wǎng)。
組 (G：Group)：八塊電路板通過銅纜完全連接，形成一個本地組，利用 11 個可用的外部鏈路進行組內(nèi)通信。
Pod 結(jié)構(gòu)：最終架構(gòu)可擴展至 36 個組（最多 1,024 個活動芯片），通過光路交換機 (OCS) 連接，確保任何芯片間通信的最大延遲為 7 跳。
深入探討：Boardfly 與環(huán)面數(shù)學（torus math）
為什么 TPU 8i 要放棄環(huán)形結(jié)構(gòu)？這歸根結(jié)底是網(wǎng)絡(luò)直徑的問題。
在三維環(huán)面中，節(jié)點排列成網(wǎng)格狀，每個維度都像環(huán)一樣環(huán)繞。在 8 x 8 x 16（1024 個芯片）的配置中，要到達最遠的芯片，數(shù)據(jù)包必須遍歷每個環(huán)一半的距離：
3D torus = 8/2(X) + 8/2(Y) + 16/2(Z) = 16 hops
雖然環(huán)面對于密集訓(xùn)練中常見的鄰居間通信非常高效，但它會給所有芯片間的通信模式帶來額外的延遲。在推理模型和迭代優(yōu)化（MoE）時代，任何芯片都可能需要與其他任何芯片通信來路由令牌，因此跳數(shù)至關(guān)重要。
Boardfly 的高基數(shù)拓撲結(jié)構(gòu)借鑒了Dragonfly 的拓撲原理。通過增加板組之間的直接光纖長距離鏈路數(shù)量，我們實現(xiàn)了網(wǎng)絡(luò)扁平化。對于同樣的 1024 芯片組，Boardfly 將網(wǎng)絡(luò)直徑從 16 跳減少到僅 7 跳。
網(wǎng)絡(luò)直徑減少 56% 可直接轉(zhuǎn)化為更低的尾部延遲，因此 TPU 8i CAE 不會等待來自整個 pod 的數(shù)據(jù)到達。
TPU 8t和TPU 8i 一覽
軟件賦能：性能優(yōu)先的 AI 堆棧
硬件的性能取決于驅(qū)動它的軟件。第八代 TPU 基于我們首創(chuàng)于第七代 Ironwood TPU 的性能優(yōu)先技術(shù)棧構(gòu)建，旨在讓用戶能夠輕松進行自定義內(nèi)核開發(fā)，同時又不犧牲高級框架的抽象性。該技術(shù)棧包括：
Pallas 和 Mosaic ：我們?yōu)镻allas提供一流的支持，Pallas 是我們自定義的內(nèi)核語言，可讓您使用 Python 編寫硬件感知內(nèi)核。這使您能夠充分發(fā)揮 TPU 8i CAE 和 TPU 8t SparseCore 的性能。
原生 PyTorch 體驗：我們很高興地宣布，TPU 的原生 PyTorch 支持現(xiàn)已推出預(yù)覽版。如果您目前正在使用 PyTorch 構(gòu)建和運行模型，我們已讓您比以往任何時候都更容易上手 TPU。您可以將現(xiàn)有模型原封不動地遷移到我們的 TPU，并完全支持您所依賴的原生功能，例如 Eager Mode。
可移植性：在 Ironwood 上運行的 JAX、PyTorch 或 Keras 代碼同樣適用于這一代產(chǎn)品。加速線性代數(shù) (XLA) 在后臺處理 Broadly 拓撲的復(fù)雜轉(zhuǎn)換和 CAE 同步，讓您可以專注于模型本身，而無需關(guān)注互連。
代代相傳：性能飛躍
我們致力于軟硬件協(xié)同設(shè)計，這一理念持續(xù)帶來豐厚回報。與第七代 Ironwood TPU 相比，第八代 TPU 的性能提升巨大：
訓(xùn)練性價比：TPU 8t 的性價比比 Ironwood TPU 提高了 2.7 倍，適用于大規(guī)模訓(xùn)練。
推理性價比：TPU 8i 比 Ironwood TPU 的性價比提高了 80%，尤其是在大型 MoE 模型的低延遲目標方面。
能源效率：兩款芯片的每瓦性能提升高達 2 倍，這對于可持續(xù)地擴展下一代人工智能至關(guān)重要。
展望未來
為了助力 Google Cloud 客戶引領(lǐng)下一波創(chuàng)新浪潮，我們設(shè)計了 TPU 8t 和 TPU 8i 這兩款截然不同的專業(yè)系統(tǒng)，以滿足人工智能生命周期中多方面的未來需求。TPU 8t 和 8i 均專為最苛刻的服務(wù)和訓(xùn)練工作負載而打造，并與人工智能超級計算機軟件棧（包括 JAX、PyTorch、vLLM、XLA 和 Pathways）完全集成。這種專業(yè)化和徹底的重新設(shè)計，以及與 Google DeepMind 的深度合作，實現(xiàn)了卓越的性價比和能效。
我們第八代架構(gòu)的模塊化設(shè)計為未來提供了清晰而獨特的發(fā)展路線圖。正如計算領(lǐng)域的每一次重大變革都需要基礎(chǔ)設(shè)施的突破一樣，智能體時代也不例外。在持續(xù)反饋循環(huán)中進行規(guī)劃、執(zhí)行和學習的推理智能體，無法在最初針對傳統(tǒng)訓(xùn)練或事務(wù)推理優(yōu)化的硬件上高效運行；它們的運行強度從根本上就有所不同。我們的第八代 TPU 基礎(chǔ)設(shè)施正是為了滿足這些特定需求而發(fā)展而來的。
京公網(wǎng)安備11010802046783號
