吳明宜在週三舉行的 Google Cloud Next 大會上，Google 宣布第 8 代自製 TPU，將用於自家超級電腦。
TPU 向來用於 Google 基礎模型如執行 Gemini 。在代理式 AI 應用中，模型必須在不同工作間切換，像是理解問題、執行多步驟工作及從行為錯誤中學習，因此對基礎架構硬體的要求也更高。最新的 TPU 8t 及 TPU 8i 晶片是由 DeepMind 和 Google Cloud 合作開發以滿足多變的任務，未來將用於執行超級電腦任務，從模型訓練、代理人開發、到大規模推理任務。
其中 TPU 8t 適合大吞吐量及垂直擴充 (scale-up) 頻寬的模型訓練這類運算密集任務。 TPU 8i 則適合記憶體頻寬、要求低延遲性的推論任務，這類能力對大規模代理人互動極為重要，失之毫釐，差之千里。兩款晶片既能執行多樣化任務，專一執行則可提升效率和效能。
TPU 8t：百萬晶片叢集驅動模型訓練
TPU 8t 現在可擴充至 9,600 顆晶片、 2 petabyte 共享頻寬記憶體，晶片間頻寬是前代的 2 倍。 TPU 8t 架構可提供 121 ExaFlops 算力，能以單一巨大記憶池支援最複雜的模型。它還整合 TPUDirect 技術，使 TPU 資料存取速度擴大 10 倍，而其新的 Virgo Network 整合 JAX 及 Google Pathways 軟體，意味能幾近直線垂直擴充，在單一邏輯叢集中擴展到百萬顆晶片。此外，利用新的穩定性、可用性及服務性 (RAS) 功能，TPU 8t 的「goodput」（具生產力的運算時間）可達 97%，大幅推升原始效能。
TPU 8i：四大創新加速代理式 AI 推論
TPU 8i 則有利於代理式 AI 任務執行。它能處理複合式流程中，多個專用代理人群集式地協同進行複雜而反覆的動作。 Google 透過 4 種創新來避免記憶體「waiting room」效應。一是打破記憶體牆 (memory wall)，TPU 8i 可平行執行 288GB 高頻寬記憶體及 384MB 晶片上 SRAM，速度為前代的 3 倍，使模型「Active Working Set」（活躍工作集）完全在晶片上運行。
第二是使用自製 Axion Arm-based CPU，並將每台伺服器的實體 CPU 主機加倍。由於隔離使用非統一主記憶體存取架構 (non-uniform memory access, NUMA)，讓整台系統效能得以攀升。
其次是擴大 MoE (Mixture of Experts) 模型，Google 將 ICI (interconnect) 頻寬擴大到 19.2 Tb/s 。 Google 的 Boardfly 架構使最大網路直徑減少超過 50%，使整台系統有如高度聚合的單一低延遲性運算單元。最後其新的集合式加速引擎 (Collectives Acceleration Engine, CAE) 得以卸載整體運作負擔，使晶片上延遲性減少到 1/5 。這些創新使 TPU 8i 的單位效能 (performance-per-dollar) 提升 80%，企業同樣成本的客戶服務量翻為 2 倍。
Google 強調，兩顆新晶片將可提升 Gemini 應用效能，例如 Boardfly 拓撲有利於推理模型的溝通性能要求，TPU 8i 的 SRAM 容量提供推論模型的 KV cache 。 Virgo Networks 頻寬可支援訓練兆級參數模型需要的平行運算。
這也是 Google TPU 首次在自家 Axion Arm-based CPU 主機上執行，展現 Google 全堆疊的技術自主里程碑。兩款晶片都原生支援 JAX 、 MaxText 、 PyTorch 、 SGLang 和 vLLM，也能裸機部署，使硬體可直接執行，減少虛擬化軟體的額外負擔，也能使用開原軟體如 MaxText 執行推論或 Tunix 進行強化式學習。
來源: Google