今年早些時候，谷歌發佈了第六代也是迄今爲止性能最強大的 TPU——Trillium。今天，Trillium 正式面向 Google Cloud 客戶開放使用。

谷歌使用 Trillium TPU 訓練了最新的 Gemini2.0，這是谷歌迄今爲止最強大的 AI 模型。現在，企業和初創公司都可以利用同樣強大、高效和可持續的基礎設施。

image.png

AI 超級計算機的核心:Trillium TPU

Trillium TPU 是 Google Cloud AI Hypercomputer 的關鍵組成部分。AI Hypercomputer 是一種突破性的超級計算機架構，它採用性能優化的硬件、開放軟件、領先的 ML 框架和靈活的消費模型集成系統。隨着 Trillium TPU 的正式推出，谷歌還對 AI Hypercomputer 的開放軟件層進行了關鍵增強，包括優化 XLA 編譯器和 JAX、PyTorch 和 TensorFlow 等流行框架，以在 AI 訓練、調整和服務方面實現領先的性價比。

此外，使用大規模主機 DRAM（補充高帶寬內存或 HBM）的主機卸載等功能可提供更高水平的效率。AI Hypercomputer 使您能夠從每個 Jupiter 網絡架構中前所未有的超過10萬個 Trillium 芯片部署中提取最大價值，該架構具有13Petabits/秒 的雙向帶寬，能夠將單個分佈式訓練作業擴展到數十萬個加速器。

image.png

AI21Labs 等客戶已經在使用 Trillium，以更快地向其客戶交付有意義的 AI 解決方案:

AI21Labs 首席技術官 Barak Lenz表示:“在 AI21，我們不斷努力提高 Mamba 和 Jamba 語言模型的性能和效率。作爲 TPU v4的長期用戶，我們對 Google Cloud 的 Trillium 的功能印象深刻。在規模、速度和成本效率方面的進步非常顯著。我們相信 Trillium 將在加速我們下一代複雜語言模型的開發方面發揮至關重要的作用，使我們能夠爲客戶提供更強大和更易於訪問的 AI 解決方案。” 

image.png

Trillium 性能大幅提升，多項指標刷新紀錄

與上一代相比，Trillium 在以下方面進行了顯著改進:

訓練性能提高4倍以上

推理吞吐量提高3倍

能源效率提高67%

每個芯片的峯值計算性能提高4.7倍

高帶寬內存 （HBM） 容量翻倍

芯片間互連 （ICI） 帶寬翻倍

單個 Jupiter 網絡架構中包含10萬個 Trillium 芯片

每美元訓練性能提高2.5倍，每美元推理性能提高1.4倍

這些增強功能使 Trillium 能夠在各種 AI 工作負載中表現出色，包括:

擴展 AI 訓練工作負載

訓練 LLM，包括密集型模型和混合專家 （MoE） 模型

推理性能和集合調度

嵌入密集型模型

提供訓練和推理性價比

image.png

Trillium 如何在不同工作負載中表現出色?

擴展 AI 訓練工作負載

訓練像 Gemini2.0這樣的大型模型需要大量的數據和計算。Trillium 的近乎線性的擴展能力使這些模型可以通過在多個 Trillium 主機之間有效且高效地分配工作負載來顯著加快訓練速度，這些主機通過高速芯片間互連連接在256芯片 pod 和我們最先進的 Jupiter 數據中心網絡中。這通過 TPU 多片、用於大規模訓練的全棧技術實現，並通過 Titanium 進一步優化，Titanium 是一種動態數據中心級卸載系統，範圍從主機適配器到網絡架構。

Trillium 在由3072個芯片組成的12個 pod 的部署中實現了99% 的擴展效率，並在具有6144個芯片的24個 pod 中展示了94% 的擴展效率，以預訓練 gpt3-175b，即使在跨數據中心網絡運行時也是如此。

訓練 LLM，包括密集型模型和混合專家 （MoE） 模型

像 Gemini 這樣的 LLM 本身就非常強大和複雜，具有數十億個參數。訓練這種密集型 LLM 需要巨大的計算能力以及協同設計的軟件優化。Trillium 在訓練 Llama-2-70b 和 gpt3-175b 等密集型 LLM 時，速度比上一代 Cloud TPU v5e 快4倍。

除了密集型 LLM 之外，使用混合專家 （MoE） 架構訓練 LLM 是一種越來越流行的方法，它結合了多個“專家”神經網絡，每個神經網絡都專門負責 AI 任務的不同方面。與訓練單個整體模型相比，在訓練期間管理和協調這些專家增加了複雜性。Trillium 在訓練 MoE 模型時的速度比上一代 Cloud TPU v5e 快3.8倍。

此外，與 Cloud TPU v5e 相比，Trillium TPU 提供了3倍的主機動態隨機存取存儲器 （DRAM）。這會將一些計算卸載到主機，有助於最大限度地提高大規模性能和良好吞吐量。Trillium 的主機卸載功能在訓練 Llama-3.1-405B 模型時，在模型 FLOP 利用率 (MFU) 方面提供了超過50% 的性能提升。

推理性能和集合調度

在推理時，多步驟推理的重要性日益提高，這需要加速器能夠有效地處理增加的計算需求。Trillium 爲推理工作負載提供了重大進步，可以更快、更高效地部署 AI 模型。事實上，Trillium 爲圖像擴散和密集型 LLM 提供了我們最佳的 TPU 推理性能。我們的測試表明，與 Cloud TPU v5e 相比，Stable Diffusion XL 的相對推理吞吐量（每秒圖像數）高出3倍以上，而 Llama2-70B 的相對推理吞吐量(每秒標記數)高出近2倍。

Trillium 是我們用於離線和服務器推理用例的性能最高的 TPU。下圖顯示，與 Cloud TPU v5e 相比，Stable Diffusion XL 的離線推理相對吞吐量（每秒圖像數）高出3.1倍，服務器推理相對吞吐量高出2.9倍。

除了更好的性能之外，Trillium 還引入了新的集合調度功能。此功能允許 Google 的調度系統做出智能的作業調度決策，以提高集合中存在多個副本時推理工作負載的整體可用性和效率。它提供了一種管理運行單主機或多主機推理工作負載的多個 TPU 片段的方法，包括通過 Google Kubernetes Engine （GKE）。將這些片段分組到一個集合中可以輕鬆地調整副本的數量以匹配需求。

嵌入密集型模型

通過添加第三代 SparseCore，Trillium 使嵌入密集型模型的性能提高了2倍，DLRM DCNv2的性能提高了5倍。

SparseCore 是數據流處理器，它爲嵌入密集型工作負載提供了更具適應性的架構基礎。Trillium 的第三代 SparseCore 擅長加速動態和數據相關的操作，例如分散收集、稀疏段求和和分區。

提供訓練和推理性價比

除了訓練世界上一些最大規模 AI 工作負載所需的絕對性能和規模外，Trillium 還旨在優化每美元的性能。迄今爲止，Trillium 在訓練 Llama2-70b 和 Llama3.1-405b 等密集型 LLM 時，每美元的性能比 Cloud TPU v5e 高出2.1倍，比 Cloud TPU v5p 高出2.5倍。

Trillium 擅長以經濟高效的方式並行處理大型模型。它的設計目的是使研究人員和開發人員能夠以比以前低得多的成本提供強大而高效的圖像模型。在 Trillium 上生成一千張圖像的成本，對於離線推理，比 Cloud TPU v5e 低27%，對於 SDXL 上的服務器推理，則比 Cloud TPU v5e 低22%。

將 AI 創新提升到新的水平

Trillium 代表了 Google Cloud AI 基礎設施的重大飛躍，爲各種 AI 工作負載提供了令人難以置信的性能、可擴展性和效率。藉助其使用世界一流的協同設計軟件擴展到數十萬個芯片的能力，Trillium 使您能夠實現更快的突破並交付卓越的 AI 解決方案。此外，Trillium 卓越的性價比使其成爲希望最大限度地提高 AI 投資價值的組織的經濟高效的選擇。隨着 AI 格局的不斷髮展，Trillium 證明了 Google Cloud 致力於提供尖端基礎設施，以幫助企業釋放 AI 的全部潛力。

官方介紹：https://cloud.google.com/blog/products/compute/trillium-tpu-is-ga