編者按： 在人工智能算力軍備競(jìng)賽愈演愈烈的今天，為什么 Google 會(huì)選擇與主流 GPU 截然不同的技術(shù)路線，開發(fā)出架構(gòu)獨(dú)特的 TPU？這種專用芯片究竟憑借什么優(yōu)勢(shì)，能夠支撐起 Gemini、Veo 等 AI 模型的訓(xùn)練與推理？

文章從單芯片架構(gòu)出發(fā)，深入剖析了 TPU 的核心設(shè)計(jì)理念：首先解釋了 TPU 如何通過脈動(dòng)陣列和流水線技術(shù)優(yōu)化矩陣運(yùn)算，然后闡述了 XLA 編譯器如何通過預(yù)先編譯減少緩存依賴，大幅降低能耗。在多芯片層面，作者詳細(xì)介紹了 TPU 從托盤、機(jī)架、Pod 到 Multi-Pod 的層級(jí)擴(kuò)展架構(gòu)，特別是 OCS 光交換技術(shù)如何實(shí)現(xiàn)靈活的拓?fù)渲貥?gòu)和故障容錯(cuò)。文章還通過具體案例展示了不同拓?fù)浣Y(jié)構(gòu)對(duì)并行訓(xùn)練策略的影響，以及 Multi-Pod 架構(gòu)如何支撐超大規(guī)模模型訓(xùn)練。

作者 | Henry Ko

編譯 | 岳揚(yáng)

最近我大量使用 TPU，發(fā)現(xiàn)它們與 GPU 的設(shè)計(jì)理念非常不同，感覺很有趣。

TPU 的主要優(yōu)勢(shì)在于其可擴(kuò)展性。這是通過硬件層面（例如能效方面和模塊化）與軟件層面（例如 XLA compiler）的協(xié)同設(shè)計(jì)實(shí)現(xiàn)的。

01 背景信息

簡(jiǎn)單介紹一下 TPU，它是谷歌的專用集成電路（ASIC），其設(shè)計(jì)聚焦于兩大要素：極高的矩陣運(yùn)算（matmul）吞吐量和能源效率。

它們的起源可追溯到 2006 年的谷歌。當(dāng)時(shí)，他們正在評(píng)估是采用 GPU、FPGA 還是定制的 ASIC。當(dāng)時(shí)，只有少數(shù)應(yīng)用需要使用專用硬件，他們判斷通過從大型數(shù)據(jù)中心調(diào)配多余的 CPU 算力即可滿足這些需求。但這一情況在 2013 年發(fā)生了變化，當(dāng)時(shí)谷歌的語音搜索功能運(yùn)行在神經(jīng)網(wǎng)絡(luò)上，而內(nèi)部預(yù)測(cè)認(rèn)為，如果該功能發(fā)展起來，將需要遠(yuǎn)超以往的算力。

時(shí)至今日，TPU 已為谷歌的大多數(shù)人工智能服務(wù)提供算力支撐。當(dāng)然，也包括 Gemini 或 Veo 的訓(xùn)練和推理，也包括他們的推薦模型。

讓我們從底層開始，深入了解一下 TPU 的內(nèi)部構(gòu)造。

02 單個(gè) TPU 芯片內(nèi)部的架構(gòu)層級(jí)

下文圖示均以 TPUv4 為例，但其整體布局基本也適用于最新一代 TPU（如 TPUv6p “Trillium”。TPUv7 “Ironwood” 的細(xì)節(jié)截至 2025 年 6 月尚未公布）。

單顆 TPUv4 芯片的結(jié)構(gòu)如下：

TPU Single Chip + TensorCore

每顆芯片內(nèi)含兩個(gè) TPU TensorCore，負(fù)責(zé)所有計(jì)算。（注：面向推理的專用 TPU 僅有一個(gè) TensorCore）。兩個(gè) TensorCore 共享同一份內(nèi)存：CMEM（128 MiB）和 HBM（32 GiB）。

而在每個(gè) TensorCore 內(nèi)部，都有計(jì)算單元和較小的內(nèi)存緩沖區(qū)：

1）矩陣乘法單元 (MXU)

• 這是 TensorCore 的核心部件，是一個(gè) 128x128 的脈動(dòng)陣列（systolic array）。

脈動(dòng)陣列的原理稍后說明。

2）向量單元（VPU）

• 負(fù)責(zé)執(zhí)行通用的逐元素操作（例如 ReLU、點(diǎn)加/點(diǎn)乘、歸約操作）

3）向量?jī)?nèi)存（VMEM；32 MiB）

• 內(nèi)存緩沖區(qū)。HBM 中的數(shù)據(jù)需先復(fù)制到 VMEM，TensorCore 才能開始計(jì)算。

4）標(biāo)量單元 + 標(biāo)量?jī)?nèi)存（SMEM；10 MiB）

• 用于調(diào)度 VPU 和 MXU 的執(zhí)行指令。
• 負(fù)責(zé)管理控制流、標(biāo)量運(yùn)算和內(nèi)存地址生成。

如果你使用的是英偉達(dá)（NVIDIA）GPU，那么一些初步觀察結(jié)果可能會(huì)讓你大吃一驚：

1）TPU 的片上內(nèi)存單元（CMEM、VMEM、SMEM）遠(yuǎn)大于 GPU 的 L1/L2 緩存。

2）TPU 的 HBM 容量卻遠(yuǎn)小于 GPU 的 HBM。

3）負(fù)責(zé)計(jì)算的"核心"（cores）數(shù)量明顯更少。

這與 GPU 架構(gòu)完全相反 —— GPU 擁有較小的 L1/L2 緩存（以 H100 為例，分別為 256KB 和 50MB）、更大的 HBM（H100 為 80GB）以及數(shù)以萬計(jì)的計(jì)算核心（cores）。

在我們進(jìn)一步討論之前，需明確的是，TPU 與 GPU 同樣具備極高的吞吐量。單顆 TPU v5p 芯片可達(dá) 500 TFLOPs/sec，由 8960 顆芯片組成的完整 pod 集群可實(shí)現(xiàn)約 4.45 ExaFLOPs/sec。而最新的 "Ironwood" TPUv7 每個(gè) pod（9216 顆芯片）據(jù)稱可達(dá) 42.5 ExaFLOPS/sec。

要理解 TPU 如何實(shí)現(xiàn)這種性能，我們需要深入探究其設(shè)計(jì)理念。

03 TPU 的設(shè)計(jì)理念

TPU 通過兩大技術(shù)支柱和一個(gè)核心前提實(shí)現(xiàn)了驚人的吞吐量與能源效率：systolic array（脈動(dòng)陣列） + pipelining（流水線）、Ahead-of-Time (AoT) compilation（預(yù)先編譯），以及假設(shè)絕大多數(shù)運(yùn)算都可通過適配 systolic array（脈動(dòng)陣列）的方式表達(dá)。幸運(yùn)的是，在現(xiàn)代深度學(xué)習(xí)（DL）領(lǐng)域，計(jì)算的大部分都是矩陣運(yùn)算，而這些運(yùn)算都適合使用 systolic array（脈動(dòng)陣列）。

3.1 TPU 設(shè)計(jì)選擇之一：Systolic Array + Pipelining

問：什么是 Systolic Array？

答：Systolic Array 是一種硬件設(shè)計(jì)架構(gòu)，由相互連接的處理單元（PE）網(wǎng)格組成。每個(gè) PE 執(zhí)行少量運(yùn)算（例如乘法和累加運(yùn)算），并將結(jié)果傳遞給相鄰 PE。

這種設(shè)計(jì)的好處是，數(shù)據(jù)一旦輸入 systolic array（脈動(dòng)陣列），便無需額外的控制邏輯來處理數(shù)據(jù)。此外，當(dāng)脈動(dòng)陣列的規(guī)模足夠大時(shí)，除輸入輸出外再無內(nèi)存讀寫操作。

由于脈動(dòng)陣列的剛性結(jié)構(gòu)設(shè)計(jì)（rigid organization），其僅能處理具有固定數(shù)據(jù)流模式的操作，但幸運(yùn)的是，矩陣乘法和卷積運(yùn)算（convolutions）恰好完美適配這種架構(gòu)范式。

不僅如此，pipelining（流水線技術(shù)）顯然有機(jī)會(huì)將計(jì)算與數(shù)據(jù)移動(dòng)重疊執(zhí)行。下圖展示了 TPU 架構(gòu)上 pipelined pointwise operation （通過流水線技術(shù)，加速 pointwise operation（逐點(diǎn)操作） 的執(zhí)行過程。）的示意圖。

Pipelined Pointwise Operation (from "How to Scale Your Model" [4])

旁注：Systolic Arrays（脈動(dòng)陣列）的局限性 —— 稀疏性

我們可以看到，脈動(dòng)陣列（systolic arrays）非常喜歡稠密矩陣（dense matrices）（即每個(gè) PE 幾乎每個(gè)時(shí)鐘周期都處于活躍狀態(tài)）。然而，其劣勢(shì)是，相同規(guī)模的稀疏矩陣（sparse matrices）無法獲得性能提升 —— 即使對(duì)于零值元素（zero-valued elements），PE 仍需執(zhí)行相同數(shù)量的計(jì)算周期（cycles），導(dǎo)致資源浪費(fèi)。

如若深度學(xué)習(xí)（DL）領(lǐng)域更傾向于采用更不規(guī)則的稀疏性（例如 MoE 架構(gòu)），應(yīng)對(duì)脈動(dòng)陣列的這一系統(tǒng)性局限將變得愈發(fā)重要。

3.2 TPU 設(shè)計(jì)選擇之二：預(yù)先（AoT）編譯 + 減少對(duì)緩存的依賴

本節(jié)將回答 TPU 如何通過軟硬件協(xié)同設(shè)計(jì)（TPU + XLA 編譯器）來避免使用緩存，從而實(shí)現(xiàn)高能效。

首先，請(qǐng)記住傳統(tǒng)緩存是為了處理不可預(yù)測(cè)的內(nèi)存訪問模式而設(shè)計(jì)的。一個(gè)應(yīng)用程序的內(nèi)存訪問模式（memory access patterns），可能與另一個(gè)應(yīng)用程序大相徑庭。從本質(zhì)上講，緩存允許硬件靈活地適應(yīng)各種應(yīng)用場(chǎng)景。這也是 GPU（相較于 TPU）靈活性極高的一個(gè)重要原因。

然而，緩存訪問（以及一般意義上的內(nèi)存訪問）會(huì)消耗大量能源。下面是對(duì)芯片（45納米，0.9V；[18]）上各類操作的能耗粗略估計(jì)。這里的主要啟示是，內(nèi)存的訪問和控制占用了大部分的能耗，而算術(shù)操作本身的能耗占比則小得多。

但是，如果你的應(yīng)用非常特殊，而且其計(jì)算和內(nèi)存訪問模式具有很高的可預(yù)測(cè)性呢？

舉個(gè)極端的例子，如果我們的編譯器能提前確定所有需要的內(nèi)存訪問，那么硬件僅需一個(gè)暫存器作為緩沖區(qū)就足以滿足需求，根本不需要緩存。

這正是 TPU 的設(shè)計(jì)理念所追求的，也是 TPU 使用 XLA 編譯器設(shè)計(jì)以實(shí)現(xiàn)這一目標(biāo)的根本原因。XLA 編譯器通過提前分析計(jì)算圖來生成優(yōu)化過的程序。

問：但 JAX 在 TPU 上也運(yùn)行良好，它們使用 @jit 嗎？

TPU 上的 JAX+XLA 實(shí)際處于 JIT 與 AOT 的混合模式，因此容易產(chǎn)生混淆。當(dāng)首次調(diào)用 JAX 中被 @jit 修飾的函數(shù)時(shí)，JAX 會(huì)進(jìn)行代碼追蹤并生成靜態(tài)計(jì)算圖。然后將其傳遞給 XLA 編譯器，在那里被轉(zhuǎn)化為適用于 TPU 的完全靜態(tài)二進(jìn)制文件。在最后的轉(zhuǎn)化階段，編譯器會(huì)實(shí)施針對(duì) TPU 的優(yōu)化（例如，最大限度地減少內(nèi)存訪問），使整個(gè)過程適合 TPU。

但有一點(diǎn)需要注意：當(dāng)輸入張量的形狀（shape）發(fā)生變化時(shí)，已編譯的 JIT 函數(shù)需重新編譯并緩存。這就是為什么 JAX 在處理動(dòng)態(tài)填充（dynamic padding）或長(zhǎng)度隨輸入變化的 for 循環(huán)層時(shí)表現(xiàn)不佳。

當(dāng)然，這種方案雖有優(yōu)勢(shì)，卻也存在明顯的局限。它缺乏靈活性，而對(duì)編譯器的重度依賴猶如一把雙刃劍。

那么，Google 為何仍要堅(jiān)持這種設(shè)計(jì)理念？

TPU 及其能源效率（TPUv4）

前文的能耗示意圖并不能精確反映 TPU 的實(shí)際情況，此處是 TPUv4 的能耗細(xì)目。注意，TPUv4 采用 7nm 工藝，表中 45nm 的數(shù)據(jù)僅用于對(duì)比（[3], [16]）。

單次操作能耗對(duì)比（TPUv4, 7 nm）

上方的柱狀圖展示了具體數(shù)值，但需注意，現(xiàn)代芯片采用的是 HBM3 內(nèi)存，其能耗遠(yuǎn)低于本圖表中顯示的 DDR3/4 DRAM。盡管如此，該圖仍表明內(nèi)存操作的能耗仍高出計(jì)算操作數(shù)個(gè)數(shù)量級(jí)。

這恰與 scaling laws 形成呼應(yīng)：我們非常樂意通過增加浮點(diǎn)運(yùn)算量（FLOPS）來換取更少的內(nèi)存操作。因此減少內(nèi)存操作能帶來雙重優(yōu)化收益——不僅提升程序運(yùn)行速度，還可顯著降低能耗。

04 TPU 的多芯片互聯(lián)層級(jí)結(jié)構(gòu)

現(xiàn)在升級(jí)到更高層級(jí)，觀察 TPU 在多芯片環(huán)境中的運(yùn)作方式。

4.1 托盤層級(jí)（即"板卡"；含4個(gè)芯片）

單塊 TPU 托盤包含 4 個(gè) TPU 芯片或 8 個(gè) TensorCore（簡(jiǎn)稱"核心"）。每塊托盤配備獨(dú)立 CPU 主機(jī)（注：推理型 TPU 的每個(gè)主機(jī)可訪問 2 塊托盤，因其每芯片僅含 1 個(gè)核心）。

主機(jī)（Host） ? 芯片（Chip）的連接采用 PCIe 接口，但芯片（Chip）?芯片（Chip）之間通過 Inter-Core Interconnect（ICI）連接，該接口具備更高帶寬。

不過 ICI 連接還可進(jìn)一步擴(kuò)展至多塊托盤。為此，我們需要繼續(xù)提升到機(jī)架層級(jí)（Rack level）。

4.2 機(jī)架層級(jí)（4x4x4 芯片）

TPU 最令人興奮的特性在于其可擴(kuò)展性，這一點(diǎn)從機(jī)架層級(jí)開始顯現(xiàn)。

一個(gè) TPU 機(jī)架包含 64 個(gè) TPU 芯片，通過 4x4x4 三維環(huán)面網(wǎng)絡(luò)互聯(lián)。如果您看過谷歌的 TPU 宣傳資料（如下圖），這張圖展示的是 8 個(gè) TPU 機(jī)架的集群。

8 個(gè) TPU 機(jī)架（TPUv4）

但在深入討論機(jī)架之前，我們需要澄清幾個(gè)容易混淆的術(shù)語：機(jī)架（Rack）、Pod 和切片（Slice）的區(qū)別。

問：TPU 機(jī)架、TPU Pod 和 TPU 切片有何不同？

不同谷歌資料對(duì)這些術(shù)語的使用存在差異，有時(shí)甚至混用"TPU Pod"和"TPU Slice"。本文采用谷歌 TPU 論文和 GCP 官方文檔的定義（[3][7][9]）：

1）TPU 機(jī)架（Rack）

• 包含 64 塊芯片的物理單元，也稱為“立方體（cube）”。

2）TPU Pod

• 通過 ICI 和光纖連接的 TPU 最大單元。
• 又稱"Superpod"或"Full Pod"。例如 TPUv4 的 TPU Pod 包含 4096 塊芯片（或 64 個(gè)機(jī)架）。

3）TPU 切片（Slice）

• 介于 4 塊芯片到 Superpod 規(guī)模之間的任何 TPU 配置組合。

主要區(qū)別在于，TPU 機(jī)架和 TPU Pod 是物理計(jì)量單位，而 TPU 切片是抽象計(jì)量單位。當(dāng)然，TPU 切片的設(shè)置涉及重要的物理拓?fù)浼s束，但現(xiàn)階段我們暫不展開討論。

現(xiàn)在，我們將聚焦物理計(jì)量單位：TPU 機(jī)架和 TPU Pod。這是因?yàn)椋斫?TPU 系統(tǒng)的物理連接方式，能更深入地掌握其設(shè)計(jì)哲學(xué)。

現(xiàn)在回到 TPUv4 機(jī)架的具體結(jié)構(gòu)：

單個(gè) TPU 機(jī)架通過 ICI 和 OCS（Optical Circuit Switching）技術(shù)連接 64 個(gè)芯片。實(shí)質(zhì)上，我們通過組合多個(gè)托盤（trays）來構(gòu)建一個(gè) 64 芯片的完整系統(tǒng)。這種"將小型單元組裝成超級(jí)計(jì)算機(jī)"的設(shè)計(jì)理念將持續(xù)貫穿后續(xù)層級(jí)。

下圖展示了 TPUv4 單個(gè)機(jī)架的拓?fù)浣Y(jié)構(gòu)。它采用 4x4x4 三維環(huán)面網(wǎng)絡(luò)，其中每個(gè)節(jié)點(diǎn)都代表一塊芯片，藍(lán)色箭頭表示 ICI 鏈路，而各個(gè)面上的連接線則代表 OCS（根據(jù)文獻(xiàn) [7] 重繪）。

使用 OCS 的 TPU 單機(jī)架架構(gòu)

然而，這張圖表引出了兩個(gè)關(guān)鍵問題：為何 OCS 僅應(yīng)用于環(huán)面結(jié)構(gòu)的表面？換句話說 —— 使用 OCS 的核心優(yōu)勢(shì)是什么？共有三大核心優(yōu)勢(shì)，我們將在后文再詳述另外兩點(diǎn)。

OCS 的優(yōu)勢(shì) #1：環(huán)繞連接 (Wraparound)

通過環(huán)形拓?fù)鋬?yōu)化節(jié)點(diǎn)間的通信效率。

OCS 還承擔(dān)特定 TPU 配置的環(huán)繞連接功能。該設(shè)計(jì)將兩節(jié)點(diǎn)間的跳數(shù)從最壞情況下 N-1 跳降至每軸 (N-1)/2 跳，因?yàn)槊織l軸均形成一個(gè)環(huán)形（一維環(huán)面拓?fù)洌?/p>

隨著規(guī)模的進(jìn)一步擴(kuò)大，這種影響變得更加重要，因?yàn)榻档托酒g的通信延遲對(duì)于高度并行化的實(shí)現(xiàn)至關(guān)重要。

附注：并非所有 TPU 都采用 3D 環(huán)面拓?fù)?/strong>

注意，早期 TPU（如 TPUv2/v3）及推理專用 TPU（如 TPUv5e/v6e）使用 2D 環(huán)面拓?fù)涠窍挛乃龅?3D 環(huán)面。不過 TPUv7"Ironwood" 雖定位為推理芯片，但其拓?fù)湟伤?3D 環(huán)面（注：僅根據(jù)官方宣傳材料推測(cè)）。

2D環(huán)面拓?fù)涫疽鈭D

4.3 Full Pod 層級(jí)（又稱 "Superpod"；TPUv4 為 4096 塊芯片）

正如我們通過互聯(lián)多個(gè)芯片構(gòu)建 TPU 機(jī)架，我們也可連接多個(gè)機(jī)架組成大型 Superpod。

Superpod 特指僅通過 ICI 和 OCS 互聯(lián)的最大 TPU 集群規(guī)模。雖然存在 multi-pod 層級(jí)，但這種層級(jí)需依賴更慢速的連接方式，后續(xù)將展開說明。

芯片數(shù)量會(huì)因版本不同而變化，但 TPUv4 的芯片數(shù)量為 4096（即 64 個(gè) 4x4x4 芯片的機(jī)架）。最新的 TPUv7 "Ironwood" 則高達(dá) 9216 塊芯片。

下圖展示了 TPUv4 的一個(gè) Superpod：

TPUv4 Superpod 架構(gòu)（64 個(gè)機(jī)架）

請(qǐng)注意，每個(gè)立方體（即 TPU 機(jī)架）是如何通過 OCS 相互連接的，這種設(shè)計(jì)也支持在 Pod 內(nèi)靈活劃分 TPU 切片。

采用 OCS 的 TPU 切片

我們可在 Pod 內(nèi)申請(qǐng) TPU 子集，即 TPU 切片。但即使所需芯片數(shù)(N)相同，也存在多種拓?fù)浣Y(jié)構(gòu)可供選擇。

例如，若總共需要 512 塊芯片，可選擇立方體(8x8x8)、條狀拓?fù)?4x4x32)或矩形拓?fù)?4x8x16)。選擇切片的拓?fù)浣Y(jié)構(gòu)本身就是一個(gè)超參數(shù)。

所選拓?fù)浣Y(jié)構(gòu)直接影響節(jié)點(diǎn)間通信帶寬，進(jìn)而影響各類并行策略的性能表現(xiàn)。

以立方體結(jié)構(gòu)（如8x8x8）為例，它特別適合需要全連接通信的并行計(jì)算模式，比如數(shù)據(jù)并行或張量并行，因?yàn)檫@種拓?fù)浣Y(jié)構(gòu)能提供最高的二分帶寬（bisection bandwidth）。而條狀結(jié)構(gòu)（如4x4x32）則更適用于流水線計(jì)算，這種布局可以讓順序排列的計(jì)算層之間實(shí)現(xiàn)更快速的數(shù)據(jù)傳輸（前提是單個(gè)計(jì)算層能夠適配 4x4 芯片的子切片配置）。

典型 TPU 拓?fù)涫纠?/p>

當(dāng)然，最優(yōu)拓?fù)淙Q于具體模型結(jié)構(gòu)，其尋優(yōu)過程本身即是一門學(xué)問。TPUv4 論文[9]實(shí)測(cè)表明，拓?fù)鋬?yōu)化可大大提升吞吐量（注：我不確定第一行指的是哪種 LLM 架構(gòu)，因?yàn)闆]有具體說明）。

不同拓?fù)浣Y(jié)構(gòu)的吞吐量?jī)?yōu)化對(duì)比

前文闡述了 TPU 切片，但另有一項(xiàng)重要的特性有助于提高 TPU 的運(yùn)行穩(wěn)定性。

借助 OCS 技術(shù)，這些切片無需占據(jù)物理連續(xù)的機(jī)架空間。這正是 OCS 的第二大優(yōu)勢(shì) —— 可能也是其最大優(yōu)勢(shì)，但我們此前尚未展開討論。

OCS 的優(yōu)勢(shì) #2：可重新配置的非連續(xù)多節(jié)點(diǎn)切片

需注意，這不同于將多個(gè)節(jié)點(diǎn)硬連在一起來模擬非連續(xù)切片。由于 OCS 采用光交換技術(shù)而非硬連線架構(gòu)，跨節(jié)點(diǎn)間的物理線纜數(shù)量大幅減少，從而支持更大規(guī)模的集群擴(kuò)展（即可構(gòu)建超大規(guī)模 TPU Pod）。

這樣就可以進(jìn)行靈活的節(jié)點(diǎn)規(guī)模配置。例如，假設(shè)我們想在單個(gè) Pod 上運(yùn)行三個(gè)任務(wù)。雖然傳統(tǒng)的調(diào)度方式不允許這樣做，但 OCS 連接允許我們抽象出節(jié)點(diǎn)的物理位置，使整個(gè) Pod 可視為一個(gè)"節(jié)點(diǎn)資源池"（根據(jù)參考文獻(xiàn)[6]重繪）。

單任務(wù)可將 Pod 內(nèi)機(jī)架視為"節(jié)點(diǎn)資源池"

此舉不僅提高了 Pod 的利用率，而且能在節(jié)點(diǎn)出現(xiàn)故障的情況下簡(jiǎn)化維護(hù)流程。谷歌將其描述為"故障節(jié)點(diǎn)的影響范圍很小"。但尚不確定其液冷系統(tǒng)在部分節(jié)點(diǎn)停機(jī)時(shí)如何運(yùn)作。

最后，這種靈活的 OCS 還有項(xiàng)延伸應(yīng)用：我們還可以改變 TPU 切片的拓?fù)浣Y(jié)構(gòu)（例如將規(guī)則環(huán)面調(diào)整為扭曲環(huán)面）。

OCS 的優(yōu)勢(shì) #3：扭曲環(huán)面拓?fù)?/strong>

此前我們通過改變固定芯片數(shù)量下的 (x,y,z) 維度來實(shí)現(xiàn)不同的 TPU 切片拓?fù)浣Y(jié)構(gòu)。本節(jié)則聚焦固定維度配置，通過改變布線方式構(gòu)造新型拓?fù)洹?/p>

典型案例如下：將常規(guī)條狀環(huán)面改造為扭曲條狀環(huán)面。

常規(guī)環(huán)面 vs 扭曲環(huán)面（來源：TPUv4論文[9]）

扭曲環(huán)面拓?fù)浣Y(jié)構(gòu)能加速扭曲二維平面上的芯片之間的通信，該特性對(duì)提升全局通信效率尤其有用。

下文將深入分析其具體應(yīng)用場(chǎng)景。

使用扭曲環(huán)面加速訓(xùn)練

理論上，扭曲環(huán)面對(duì)張量并行（TP）的加速效益最大，因?yàn)槊繉由婕岸啻?all-gather 和 reduce-scatter 操作。對(duì)數(shù)據(jù)并行（DP）也有適度提升，因?yàn)槊總€(gè)訓(xùn)練步需執(zhí)行 all-reduce 操作，但發(fā)生頻率較低。

想象一下，假設(shè)我們訓(xùn)練一個(gè)標(biāo)準(zhǔn)的僅解碼器架構(gòu)的 Transformer 模型，并采用多種并行策略來加速訓(xùn)練。下面我們將看到兩種場(chǎng)景：

場(chǎng)景 #1：4x4x16 拓?fù)浣Y(jié)構(gòu)（TP+PP；共 256 塊芯片）

設(shè)定 z 軸為流水線(PP)維度，二維 TP 維度為 4x4。本質(zhì)上，假設(shè)第 k 層位于 z=k 平面，且每層分片至 16 塊芯片。若未明確繪制，默認(rèn)采用 OCS 最近鄰連接。

TP+PP 的 4x4x16 拓?fù)浼軜?gòu)

通過在每個(gè) z=k 平面實(shí)施 2D 環(huán)面扭曲，可加速 TP 層內(nèi)芯片通信。由于 PP 層主要依靠點(diǎn)對(duì)點(diǎn)通信，因此沒有必要沿 PP 層扭曲。

注：實(shí)際應(yīng)用中，扭曲環(huán)面在芯片數(shù)＞4x4 時(shí)效益顯著。本示例使用 4x4 僅出于可視化的目的。

場(chǎng)景 #2：16x4x16 拓?fù)洌―P+TP+PP；共 1024 塊芯片）

作為延伸方案，我們?cè)谇耙粓?chǎng)景基礎(chǔ)上增加 DP 維度（x 軸 4 個(gè)實(shí)例），即沿 x 軸部署 4 組場(chǎng)景 #1 的模型。

DP+TP+PP 的 16x4x16 拓?fù)浼軜?gòu)

請(qǐng)注意，扭曲環(huán)面僅應(yīng)用于每個(gè) DP 模型內(nèi)的每個(gè) TP 維度（即對(duì)每個(gè) z=k 平面實(shí)施 4x4 二維扭曲，k 取值 1…16）。DP 維度僅維持基礎(chǔ)的環(huán)繞連接，使每行構(gòu)成長(zhǎng)度為 16 的水平環(huán)。

你可能已經(jīng)發(fā)現(xiàn)還有一種拓?fù)浣Y(jié)構(gòu)方案（如 8x8x16，即 2x2 DP 維度），但這會(huì)混合 DP 與 TP 維度 —— 這就變得更加復(fù)雜了。具體來說，我們還不清楚如何在 y 軸構(gòu)建 OCS 環(huán)繞連接的同時(shí)兼容各 TP 維度的扭曲環(huán)面？

4.4 Multi-Pod 層級(jí)（即"Multislice"；TPUv4 支持 4096+ 塊芯片）

TPU 層次結(jié)構(gòu)的最終層級(jí)是 Multi-pod 架構(gòu)。此時(shí)可將多個(gè) Pod 視為一臺(tái)大型機(jī)器，但 Pod 之間的通信需通過數(shù)據(jù)中心網(wǎng)絡(luò)（DCN） 進(jìn)行 —— 其帶寬低于 ICI。

通過 DCN 互聯(lián)的雙 Pod 架構(gòu) [1]

PaLM 模型即采用此方案進(jìn)行訓(xùn)練。在 6144 個(gè) TPUv4 芯片（2 個(gè) Pod）上耗時(shí) 56 天完成。下圖是 6 個(gè) Pod 中的 TPU 任務(wù)分配情況：綠色為 PaLM 任務(wù)，紅色為空閑狀態(tài)，其余為其他任務(wù)。注意每個(gè)方格代表一個(gè) 4x4x4 的 TPU 芯片立方體。

PaLM 訓(xùn)練過程中的 TPU Pod 利用率 [6]

實(shí)現(xiàn)這一架構(gòu)已屬不易，但更關(guān)鍵的是開發(fā)者體驗(yàn)設(shè)計(jì)，具體來說，就是要關(guān)注：如何實(shí)現(xiàn)模型擴(kuò)展過程中系統(tǒng)/硬件層面的最大程度抽象化？

谷歌的解決方案是：由 XLA 編譯器在大規(guī)模計(jì)算場(chǎng)景下協(xié)調(diào)芯片間的通信。研究人員只需配置相關(guān)參數(shù)（如 DP、FSDP、TP 等并行維度及切片數(shù)量），XLA 編譯器即會(huì)根據(jù)當(dāng)前 TPU 拓?fù)浣Y(jié)構(gòu)自動(dòng)插入分層集合通信操作（Xu et al, 2021: GSPMD [2]）。我們的目標(biāo)是在盡可能少修改代碼的情況下實(shí)現(xiàn)大規(guī)模訓(xùn)練。

例如，谷歌博客[1]展示了跨多 TPU 切片的 all-reduce 操作分解流程：

XLA 實(shí)現(xiàn)的跨 Pod All-Reduce 規(guī)約操作

這表明 XLA 編譯器可以同時(shí)處理切片內(nèi)與切片間的集合通信操作。

舉個(gè)具體例子，在訓(xùn)練模型時(shí)，TPU 的拓?fù)浣Y(jié)構(gòu)可能如下所示。激活值的通信在切片內(nèi)通過 ICI 進(jìn)行，而梯度的通信則需跨切片通過 DCN 完成（即在 DCN 的 DP 維度上）[1]。

05 實(shí)物圖示對(duì)照解析

結(jié)合硬件實(shí)拍圖理解架構(gòu)圖會(huì)更直觀，以下為綜合解析。

若看過谷歌 TPU 宣傳資料，可能見過下圖：

8 個(gè) TPU 機(jī)架（TPUv4）

此圖為 8 個(gè) TPU Pods 的集群，每個(gè)單元即前述的 4x4x4 三維環(huán)面架構(gòu)。一個(gè) Pod 中的每一行有 2 個(gè)托盤，這意味著每一行有 8 個(gè) TPU 芯片。

單塊 TPUv4 托盤實(shí)拍圖：

請(qǐng)注意，圖中簡(jiǎn)化為只有一個(gè) PCIe 端口，但實(shí)際托盤上有 4 個(gè) PCIe 端口（在左側(cè)） —— 每個(gè) TPU 一個(gè)。

單芯片結(jié)構(gòu)圖：

TPUv4 芯片：中央是 ASIC + 4 組 HBM 內(nèi)存堆棧

中央?yún)^(qū)域?yàn)?ASIC 芯片，周圍 4 個(gè)區(qū)塊為 HBM 內(nèi)存堆棧。因 TPUv4 內(nèi)含 2 個(gè) TensorCore，故配置 4 組 HBM 內(nèi)存堆棧。

未找到 TPUv4 芯片平面圖，此處展示結(jié)構(gòu)近似的 TPUv4i（推理芯片），其僅含 1 個(gè)TensorCore[3]：

可見 CMEM（芯片內(nèi)存）在 TPUv4i 的布局中占據(jù)了相當(dāng)大的空間。

06 致謝

感謝 Google TPU Research Cloud（TRC）提供的 TPU 資源支持！

References

[1] Google Blog: TPU Multi-Slice Training（??https://cloud.google.com/blog/products/compute/using-cloud-tpu-multislice-to-scale-ai-workloads）??

[2] Xu, et al. "GSPMD: General and Scalable Parallelizaton for ML Computation Graphs"（??https://arxiv.org/pdf/2105.04663）??

[3] Jouppi et al. "Ten Lessons From Three Generations Shaped Google's TPUv4i"（??https://gwern.net/doc/ai/scaling/hardware/2021-jouppi.pdf）??

[4] How to Scale Your Model - TPUs（??https://jax-ml.github.io/scaling-book/tpus/）??

[5] Domain Specific Architectures for AI Inference - TPUs（??https://fleetwood.dev/posts/domain-specific-architectures#google-tpu）??

[6] HotChips 2023: TPUv4（??https://hc2023.hotchips.org/assets/program/conference/day2/ML+training/HC2023.Session5.ML_Training.Google.Norm_Jouppi.Andy_Swing.Final_2023-08-25.pdf）??

[7] Google Cloud Docs: TPUv4（??https://cloud.google.com/tpu/docs/v4）??

[8] Jouppi et al. "In-Datacenter Performance Analysis of a Tensor Processing Unit" -- TPU origins paper（??https://arxiv.org/abs/1704.04760）??

[9] Jouppi et al. "TPU v4"-- TPUv4 paper（??https://arxiv.org/abs/2304.01433）??

[10] PaLM training video（??https://www.youtube.com/watch?v=0yPFBxkOKRY）??

[11] HotChips 2021: "Challenges in large scale training of Giant Transformers on Google TPU machines"（??https://hc33.hotchips.org/assets/program/tutorials/HC2021.Google.Sameer+Kumar.pdf）??

[12] HotChips 2020: "Exploring Limits of ML Training on Google TPUs"（??https://hc32.hotchips.org/assets/program/tutorials/HC2020.Google.SameerKumarDehaoChen.v02.pdf）??

[13] Google Blog: Ironwood（??https://blog.google/products/google-cloud/ironwood-tpu-age-of-inference/）??

[14] HotChips 2019: "Cloud TPU: Codesigning Architecture and Infrastructure"（??https://old.hotchips.org/hc31/HC31_T3_Cloud_TPU_Codesign.pdf）??

[15] ETH Zurich's Comp Arch Lecture 28: Systolic Array Architectures（??https://www.youtube.com/watch?v=XkgtANeDrm8）??

[16] Patterson presentation: "A Decade of Machine Learning Accelerators: Lessons Learned and Carbon Footprint"（??https://www.cs.ucla.edu/wp-content/uploads/cs/PATTERSON-10-Lessons-4-TPU-gens-CO2e-45-minutes.pdf）??

[17] Camara et al. "Twisted Torus Topologies for Enhanced Interconnection Networks."（??

[18] Horowitz article: "Computing's Energy Problem(and what we can do about it)"（??https://gwern.net/doc/cs/hardware/2014-horowitz-2.pdf）??

END

本期互動(dòng)內(nèi)容 ??

?您更傾向 TPU 的專用化路線（犧牲靈活性換取能效），還是 GPU 的通用化路線（保留靈活性但能耗較高）？請(qǐng)結(jié)合您的應(yīng)用場(chǎng)景說明理由。

本文經(jīng)原作者授權(quán)，由 Baihai IDP 編譯。如需轉(zhuǎn)載譯文，請(qǐng)聯(lián)系獲取授權(quán)。

原文鏈接：

??https://henryhmko.github.io/posts/tpu/tpu.html??