“本文轉(zhuǎn)載自 Henry Ko 的 Blog，深入地解析了 TPU 的相關(guān)知識(shí)。”

我最近一直在大量使用 TPU，很有趣地發(fā)現(xiàn)它們與 GPU 的設(shè)計(jì)理念有多么不同。

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

背景

簡(jiǎn)單來說，TPU 是谷歌的專用集成電路 (ASIC)，專注于兩個(gè)因素：極高的矩陣乘法吞吐量 + 高能效。

它們的起源可以追溯到 2006 年的谷歌，當(dāng)時(shí)他們首次評(píng)估應(yīng)該采用 GPU、FPGA 還是定制 ASIC。那時(shí)只有少數(shù)應(yīng)用需要專門的硬件，他們認(rèn)為這些需求可以通過利用其大型數(shù)據(jù)中心多余的 CPU 計(jì)算能力來滿足。但情況在 2013 年發(fā)生了變化，當(dāng)時(shí)谷歌的語(yǔ)音搜索功能開始在神經(jīng)網(wǎng)絡(luò)上運(yùn)行，內(nèi)部預(yù)測(cè)顯示，如果該功能大受歡迎，他們將需要更多的計(jì)算能力。

快進(jìn)到今天，TPU 為谷歌的大部分 AI 服務(wù)提供了動(dòng)力。當(dāng)然，這包括 Gemini 或 Veo 的訓(xùn)練和推理，也包括部署他們的推薦模型 (DLRM)。

讓我們從底層向上深入了解 TPU 的內(nèi)部結(jié)構(gòu)。

TPU 單芯片層面

我的圖表將主要關(guān)注 TPUv4，但這種布局或多或少也適用于最新一代的 TPU（例如，TPUv6p "Trillium"；截至 2025 年 6 月撰寫本文時(shí)，TPUv7 "Ironwood" 的細(xì)節(jié)尚未公布）。

以下是單個(gè) TPUv4 芯片的布局：

在每個(gè)芯片中，有兩個(gè) TPU TensorCore，它們負(fù)責(zé)計(jì)算。（注意：專用于推理的 TPU 只有一個(gè) TensorCore）。兩個(gè) TensorCore 共享內(nèi)存單元：CMEM (128MiB) 和 HBM (32GiB)。

在每個(gè) TensorCore 內(nèi)部，是我們的計(jì)算單元和更小的內(nèi)存緩沖區(qū)：

1. 矩陣乘法單元 (Matrix Multiply Unit, MXU)

• 這是 TensorCore 的關(guān)鍵組件，是一個(gè) 128x128 的脈動(dòng)陣列 (systolic array)。

• 我們將在下面介紹脈動(dòng)陣列。

• 向量單元 (Vector Unit, VPU)

• 用于通用的逐元素操作（例如 ReLU、逐點(diǎn)加/乘、規(guī)約操作）。

• 向量?jī)?nèi)存 (Vector Memory, VMEM; 32MiB)

• 內(nèi)存緩沖區(qū)。數(shù)據(jù)在 TensorCore 進(jìn)行任何計(jì)算之前，會(huì)從 HBM 復(fù)制到 VMEM 中。

• 標(biāo)量單元 + 標(biāo)量?jī)?nèi)存 (Scalar Unit + SMEM; 10MiB)

• 告訴 VPU 和 MXU 該做什么。

• 管理控制流、標(biāo)量操作和內(nèi)存地址生成。

如果你熟悉 NVIDIA GPU，可能會(huì)對(duì)一些初步觀察感到困惑：

1. TPU 上的片上內(nèi)存單元（CMEM、VMEM、SMEM）比 GPU 上的 L1、L2 緩存大得多。

2. TPU 上的 HBM 也比 GPU 上的 HBM 小得多。

3. 負(fù)責(zé)計(jì)算的“核心”似乎要少得多。

這與 GPU 的情況正好相反，GPU 擁有較小的 L1、L2 緩存（H100 分別為 256KB 和 50MB）、更大的 HBM（H100 為 80GB）以及數(shù)以萬計(jì)的核心。

在我們進(jìn)一步討論之前，請(qǐng)記住 TPU 能夠像 GPU 一樣實(shí)現(xiàn)極高的吞吐量。TPU v5p 每芯片可達(dá)到 500 TFLOPs/秒，一個(gè)包含 8960 個(gè)芯片的完整 Pod 大約可以達(dá)到 4.45 ExaFLOPs/秒。據(jù)說最新的 "Ironwood" TPUv7 每個(gè) Pod（9216 chips）最高可達(dá) 42.5 ExaFLOPS/秒。

要理解 TPU 是如何實(shí)現(xiàn)這一點(diǎn)的，我們需要了解它們的設(shè)計(jì)理念。

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

TPU 依靠?jī)纱笾е鸵粋€(gè)關(guān)鍵假設(shè)，實(shí)現(xiàn)了驚人的吞吐量和能效：脈動(dòng)陣列 + 流水線技術(shù)、預(yù)先編譯 (Ahead-of-Time, AoT)，以及假設(shè)大多數(shù)操作都可以用一種能很好地映射到脈動(dòng)陣列的方式來表達(dá)。幸運(yùn)的是，在我們現(xiàn)代的深度學(xué)習(xí)時(shí)代，矩陣乘法占據(jù)了計(jì)算的大部分，這非常適合脈動(dòng)陣列。

TPU 設(shè)計(jì)選擇 #1: 脈動(dòng)陣列 + 流水線技術(shù)

問：什么是脈動(dòng)陣列 (Systolic Array)？

脈動(dòng)陣列是一種硬件設(shè)計(jì)架構(gòu)，由一個(gè)網(wǎng)格狀的互連處理單元 (Processing Element, PE) 組成。每個(gè) PE 執(zhí)行一個(gè)小的計(jì)算（例如，乘法和累加），并將結(jié)果傳遞給相鄰的 PE。

這種設(shè)計(jì)的好處在于，一旦數(shù)據(jù)被送入脈動(dòng)陣列，就不再需要額外的控制邏輯來決定如何處理數(shù)據(jù)。此外，當(dāng)脈動(dòng)陣列足夠大時(shí)，除了輸入和輸出，沒有其他的內(nèi)存讀/寫操作。

由于其剛性的組織結(jié)構(gòu)，脈動(dòng)陣列只能處理具有固定數(shù)據(jù)流模式的操作，但幸運(yùn)的是，矩陣乘法和卷積完美地契合了這一范疇。

此外，通過流水線技術(shù)，可以將計(jì)算與數(shù)據(jù)移動(dòng)重疊起來。下面是一個(gè)在 TPU 上進(jìn)行流水線式逐點(diǎn)操作的圖示。

題外話：脈動(dòng)陣列的缺點(diǎn) - 稀疏性

你可以看到脈動(dòng)陣列非常喜歡密集矩陣（即每個(gè) PE 在幾乎每個(gè)周期都處于活動(dòng)狀態(tài)）。然而，其缺點(diǎn)在于，對(duì)于同樣大小的稀疏矩陣，性能沒有提升：即使是對(duì)于值為零的元素，PE 仍然需要執(zhí)行相同數(shù)量的周期來進(jìn)行計(jì)算。

如果深度學(xué)習(xí)社區(qū)更青睞不規(guī)則的稀疏性（例如 MoE），那么處理脈動(dòng)陣列的這種系統(tǒng)性限制將變得更加重要。

TPU 設(shè)計(jì)選擇 #2: 預(yù)先編譯 (AoT) + 減少對(duì)緩存的依賴

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

首先，回想一下，傳統(tǒng)緩存是為了處理不可預(yù)測(cè)的內(nèi)存訪問模式而設(shè)計(jì)的。一個(gè)應(yīng)用程序的程序與另一個(gè)應(yīng)用程序的程序可能有著截然不同的內(nèi)存訪問模式。本質(zhì)上，緩存使硬件具有靈活性，能夠適應(yīng)各種應(yīng)用。這是 GPU 成為非常靈活的硬件的一大原因（注：與 TPU 相比）。

然而，緩存訪問（以及一般的內(nèi)存訪問）會(huì)消耗大量能量。下面是一個(gè)芯片上操作的能耗粗略估算（45nm, 0.9V; ）。這里的關(guān)鍵信息是，內(nèi)存訪問和控制占據(jù)了我們大部分的能量，而算術(shù)運(yùn)算的能耗則要低得多。

但是，如果你的應(yīng)用非常特定，其計(jì)算/內(nèi)存訪問模式高度可預(yù)測(cè)呢？

舉個(gè)極端的例子，如果我們的編譯器能夠提前計(jì)算出所有需要的內(nèi)存訪問，那么我們的硬件只需要一個(gè)便簽式內(nèi)存 (scratchpad memory) 作為緩沖區(qū)就足夠了，完全不需要緩存。

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

問：但是 JAX 也能很好地與 TPU 配合，可它們用的是 @jit？

JAX+XLA 在 TPU 上處于即時(shí)編譯 (JIT) 和預(yù)先編譯 (AOT) 的混合狀態(tài)，因此會(huì)產(chǎn)生困惑。當(dāng)我們第一次在 JAX 中調(diào)用一個(gè) jit 函數(shù)時(shí)，JAX 會(huì)追蹤它以創(chuàng)建一個(gè)靜態(tài)計(jì)算圖。這個(gè)圖被傳遞給 XLA 編譯器，在那里它被轉(zhuǎn)換成一個(gè)完全靜態(tài)的 TPU 二進(jìn)制文件。正是在這最后的轉(zhuǎn)換階段，會(huì)進(jìn)行 TPU 特定的優(yōu)化（例如，最小化內(nèi)存訪問）來為 TPU 定制處理過程。

但有一個(gè)注意事項(xiàng)：如果 jit 函數(shù)以不同的輸入形狀運(yùn)行，就必須重新編譯和緩存。這就是為什么 JAX 在處理任何動(dòng)態(tài)填充或依賴于輸入的不同長(zhǎng)度的 for 循環(huán)層時(shí)表現(xiàn)不佳的原因。

當(dāng)然，這種方法聽起來很好，但也有不便的缺點(diǎn)。它缺乏靈活性，并且對(duì)編譯器的重度依賴是一把雙刃劍。

但為什么谷歌仍然堅(jiān)持這種設(shè)計(jì)理念呢？

TPU 與能效 (TPUv4)

前面那張能耗圖并不能準(zhǔn)確代表 TPU，所以這里是 TPUv4 的能耗分解。請(qǐng)注意，TPUv4 是 7nm 工藝，而 45nm 的數(shù)據(jù)僅作對(duì)比 。

左邊的條形圖直觀地顯示了數(shù)值，但需要注意的一點(diǎn)是，現(xiàn)代芯片使用 HBM3，其能耗遠(yuǎn)低于這里顯示的 DDR3/4 DRAM。盡管如此，這表明內(nèi)存操作的能耗要高出幾個(gè)數(shù)量級(jí)。

這與現(xiàn)代規(guī)模法則 (scaling laws) 有很好的聯(lián)系：我們非常樂意增加浮點(diǎn)運(yùn)算次數(shù) (FLOPS) 來?yè)Q取內(nèi)存操作的減少。因此，減少內(nèi)存操作具有雙重的優(yōu)化效益，因?yàn)樗鼈儾粌H使程序運(yùn)行更快，而且能耗也更低。

TPU 多芯片層面

讓我們?cè)僖粚?，看?TPU 在多芯片環(huán)境中的工作方式。

Tray 層面 (又稱 "板卡"; 4 芯片)

一個(gè) TPU tray 包含 4 個(gè) TPU 芯片或 8 個(gè) TensorCore（簡(jiǎn)稱為“核心”）。每個(gè) tray 都有自己的 CPU 主機(jī)（注意：對(duì)于推理型 TPU，一個(gè)主機(jī)訪問 2 個(gè)tray，因?yàn)樗鼈兠總€(gè)芯片只有 1 個(gè)核心）。

主機(jī)與芯片的連接是 PCIe，但芯片與芯片之間的連接是核間互連 (Inter-Core Interconnect, ICI)，它具有更高的帶寬。

但 ICI 連接可以延伸到更遠(yuǎn)的多個(gè)托盤。為此，我們需要上升到機(jī)架 (Rack) 層面。

機(jī)架 (Rack) 層面 (4x4x4 芯片)

TPU 特別令人興奮的部分在于其可擴(kuò)展性，我們從機(jī)架層面開始看到這一點(diǎn)。

一個(gè) TPU 機(jī)架由 64 個(gè) TPU 組成，它們以 4x4x4 的 3D 環(huán)面 (torus) 結(jié)構(gòu)連接。如果你看過谷歌下面這樣的 TPU 宣傳材料，那就是 8 個(gè) TPU 機(jī)架的圖像。

但在我們深入了解機(jī)架之前，需要澄清一些容易混淆的術(shù)語(yǔ)：機(jī)架 (rack) vs. Pod vs. Slice。

問：“TPU Rack”、“TPU Pod” 和 “TPU Slice” 之間有什么區(qū)別？

不同的谷歌資料對(duì)它們的使用略有不同，有時(shí)會(huì)將 "TPU Pods" 和 "TPU Slices" 互換使用。但在本文中，我們將遵循谷歌 TPU 論文和 GCP TPU 文檔中的定義。

• TPU Rack (機(jī)架):

  • 包含 64 個(gè)芯片的物理單元。也稱為“立方體 (cube)”。

• TPU Pod:

  • 可通過 ICI 和光纖連接的 TPU 的最大單元。

  • 也稱為 "Superpod" 或 "Full pod"。例如，TPUv4 的一個(gè) TPU Pod 將由 4096 個(gè)芯片或 64 個(gè) TPU 機(jī)架組成。

• TPU Slice (切片):

  • 介于 4 個(gè)芯片和 Superpod 大小之間的任何 TPU 配置。

關(guān)鍵區(qū)別在于，TPU Rack 和 TPU Pod 是物理度量單位，而 TPU Slice 是一個(gè)抽象單位。當(dāng)然，設(shè)置 TPU Slice 有重要的物理屬性，但我們暫時(shí)將其抽象化。

現(xiàn)在，我們將使用物理度量單位：TPU Racks 和 TPU Pods。這是因?yàn)榱私?TPU 系統(tǒng)是如何物理連接的，可以幫助我們更好地理解 TPU 的設(shè)計(jì)理念。

現(xiàn)在回到 TPU 機(jī)架 (針對(duì) TPUv4)：

一個(gè) TPU 機(jī)架由 64 個(gè)芯片組成，通過 ICI 和光路交換 (Optical Circuit Switching, OCS) 連接在一起。本質(zhì)上，我們連接多個(gè) tray 來模擬一個(gè) 64 芯片的系統(tǒng)。這種將小部件組合成超級(jí)計(jì)算機(jī)的主題在后面會(huì)繼續(xù)出現(xiàn)。

下面是單個(gè) TPUv4 機(jī)架的圖示。它是一個(gè) 4x4x4 的 3D 環(huán)面，每個(gè)節(jié)點(diǎn)是一個(gè)芯片，藍(lán)色的箭頭是 ICI，而各個(gè)面上的線是 OCS。

然而，這張圖引出了幾個(gè)問題。為什么 OCS 只用于各個(gè)面？換句話說，使用 OCS 有什么好處？有 3 大好處，我們稍后會(huì)介紹另外兩個(gè)。

OCS 的好處 #1: 環(huán)繞連接 (Wraparound)

通過環(huán)繞連接實(shí)現(xiàn)節(jié)點(diǎn)間更快的通信。

OCS 還充當(dāng)給定 TPU 配置的環(huán)繞連接。這將兩個(gè)節(jié)點(diǎn)之間的最壞情況跳數(shù)從 N-1 跳減少到每個(gè)軸 (N-1)/2 跳，因?yàn)槊總€(gè)軸都變成了一個(gè)環(huán)（1D 環(huán)面）。

隨著我們進(jìn)一步擴(kuò)展，這種效應(yīng)變得更加重要，因?yàn)闇p少芯片間通信延遲對(duì)于高并行化至關(guān)重要。

題外話：并非所有 TPU 都具有 3D 環(huán)面拓?fù)?/span>

注意：較早的 TPU 代（例如 TPUv2, v3）和推理 TPU（例如 TPUv5e, TPUv6e）具有 2D 環(huán)面拓?fù)?，而不是像下面這樣的 3D 環(huán)面。然而，TPUv7 "Ironwood" 似乎是 3D 環(huán)面，盡管它被宣傳為推理芯片（注意：我只是根據(jù)他們的宣傳材料進(jìn)行假設(shè)）。

完整 Pod 層面 (又稱 "Superpod"; TPUv4 為 4096 芯片)

就像我們將多個(gè)芯片連接起來組成一個(gè) TPU 機(jī)架一樣，我們可以連接多個(gè)機(jī)架來組成一個(gè)大型的 Superpod。

Superpod 也指 TPU 可以達(dá)到的（僅使用 ICI 和 OCS）互連芯片的最大配置。接下來還有一個(gè)多 Pod 層面，但這必須通過較慢的互連，我們稍后會(huì)討論。

這個(gè)大小因代而異，但對(duì)于 TPUv4 是 4096 個(gè)芯片（即 64 個(gè) 4x4x4 芯片的機(jī)架）。對(duì)于最新的 TPUv7 "Ironwood"，是 9216 個(gè)芯片。

下圖顯示了一個(gè) TPUv4 的 Superpod。

請(qǐng)注意每個(gè)立方體（即一個(gè) TPU 機(jī)架）是如何通過 OCS 相互連接的。這也允許我們?cè)谝粋€(gè) Pod 中獲取 TPU 的“切片”(slices)。

帶 OCS 的 TPU 切片

我們可以在 Pod 內(nèi)請(qǐng)求 TPU 的子集，這些就是 TPU 切片。但即使你想要 N 個(gè)芯片，也有多種拓?fù)淇晒┻x擇。

例如，假設(shè)你總共需要 512 個(gè)芯片。你可以要求一個(gè)立方體 (cube) (8x8x8)、一個(gè)雪茄形 (cigar shape) (4x4x32) 或一個(gè)矩形 (rectangle) (4x8x16)。選擇切片的拓?fù)浔旧砭褪且粋€(gè)超參數(shù)。

你選擇的拓?fù)鋾?huì)影響節(jié)點(diǎn)之間的通信帶寬。這直接影響不同并行化方法的性能。

例如，對(duì)于全局通信 (all-to-all)，如數(shù)據(jù)并行或張量并行，立方體 (例如 8x8x8) 會(huì)是首選，因?yàn)樗哂凶罡叩膶?duì)分帶寬 (bisection bandwidth)。然而，對(duì)于流水線并行，雪茄形 (例如 4x4x32) 會(huì)更好，因?yàn)樗梢愿斓嘏c順序?qū)油ㄐ牛僭O(shè)一個(gè)層適合一個(gè) 4x4 芯片的子切片）。

當(dāng)然，最佳拓?fù)淙Q于模型，找到它本身就是一項(xiàng)工作。TPUv4 的論文 [9] 也對(duì)此進(jìn)行了測(cè)量，以顯示拓?fù)渥兓绾渭铀偻掏铝浚ㄗ⒁猓何也淮_定第一行指的是哪種 LLM 架構(gòu)，因?yàn)樗鼪]有具體說明）。

我們介紹了 TPU 切片，但還有一個(gè)重要特性有助于 TPU 的高運(yùn)行穩(wěn)定性。

那就是由于 OCS，這些切片不必是連續(xù)的機(jī)架。這是我們前面沒有提到的使用 OCS 的第二個(gè)好處——可能也是最大的好處。

OCS 的好處 #2: (可重構(gòu)的) 非連續(xù)多節(jié)點(diǎn)切片

請(qǐng)注意，這與硬連線多個(gè)節(jié)點(diǎn)來模擬非連續(xù)切片是不同的。由于 OCS 是一個(gè)交換機(jī)而不是硬連線，節(jié)點(diǎn)之間的物理線路要少得多，因此允許更高的可擴(kuò)展性（即更大的 TPU Pod 尺寸）。

這允許大規(guī)模靈活的節(jié)點(diǎn)配置。例如，假設(shè)我們想在一個(gè) Pod 上運(yùn)行三個(gè)作業(yè)。雖然樸素的調(diào)度不允許這樣做，但 OCS 連接允許我們抽象出節(jié)點(diǎn)的位置，并將整個(gè) Pod 僅僅看作一個(gè) “節(jié)點(diǎn)袋” (bag of nodes)。

這提高了 Pod 的利用率，并且在節(jié)點(diǎn)發(fā)生故障時(shí)可能使維護(hù)更容易。谷歌將其描述為**“故障節(jié)點(diǎn)的爆炸半徑很小”。然而，我不確定當(dāng)只有某些節(jié)點(diǎn)必須關(guān)閉時(shí)，其液體冷卻會(huì)受到怎樣的影響。

最后，這種靈活的 OCS 還有一個(gè)有趣的擴(kuò)展：我們還可以改變 TPU 切片的拓?fù)?，例如從常?guī)環(huán)面變?yōu)榕でh(huán)面 (twisted torus)。

OCS 的好處 #3: 扭曲的 TPU 拓?fù)?/span>

我們之前看到了如何通過改變固定芯片數(shù)量的 (x,y,z) 維度來獲得不同的 TPU 切片拓?fù)?。然而，這次我們將在固定的 (x,y,z) 維度下工作，但改變它們的連接方式以實(shí)現(xiàn)不同的拓?fù)洹?/span>

一個(gè)顯著的例子是從雪茄形的常規(guī)環(huán)面變?yōu)槿缦滤镜呐で┣循h(huán)面。

扭曲環(huán)面允許在扭曲的 2D 平面上的芯片之間進(jìn)行更快的通信。這對(duì)于加速全局通信 (all-to-all) 特別有用。

讓我們更深入地探討一下，想象一個(gè)具體的場(chǎng)景，這會(huì)有所幫助。

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

理論上，扭曲環(huán)面對(duì)張量并行 (Tensor Parallel, TP) 的好處最大，因?yàn)槊繉佣加卸鄠€(gè) all-gather 和 reduce-scatter 操作。它可能對(duì)數(shù)據(jù)并行 (Data Parallel, DP) 帶來中等的好處，因?yàn)槊總€(gè)訓(xùn)練步驟也有一個(gè) all-reduce，但這會(huì)不那么頻繁。

想象一下，我們正在訓(xùn)練一個(gè)標(biāo)準(zhǔn)的 decoder-only transformer，并且我們想采用大量的并行化來加速訓(xùn)練。我們將在下面看到兩種情況。

場(chǎng)景 #1: 4x4x16 拓?fù)?(TP + PP; 總共 256 芯片)

我們的 z 軸將是我們的流水線并行 (Pipeline Parallel, PP) 維度，我們的 2D TP 維度將是 4x4。本質(zhì)上，假設(shè)每個(gè)層 k 位于 z=k，并且每個(gè)層在 16 個(gè)芯片上分片。如果沒有明確繪制，則假定標(biāo)準(zhǔn)的 OCS 連接（即最近鄰）。

我們將在每個(gè) z=k 處扭曲 2D 環(huán)面，這使得每個(gè) TP 層中的芯片之間的通信更快。沿著我們的 PP 維度扭曲是不必要的，因?yàn)樗鼈冎饕蕾囉邳c(diǎn)對(duì)點(diǎn)通信。

注意: 實(shí)際上，當(dāng)芯片數(shù)量大于 4x4 時(shí)，扭曲環(huán)面才會(huì)帶來好處。我們?cè)谶@里使用 4x4 僅為可視化目的。

場(chǎng)景 #2: 16x4x16 拓?fù)?(DP + TP + PP; 總共 1024 芯片)

作為擴(kuò)展，我們將在之前的場(chǎng)景中添加一個(gè)大小為 4 的 DP 維度。這意味著沿著 x 軸有 4 個(gè)場(chǎng)景 #1 的模型。

請(qǐng)注意扭曲環(huán)面如何僅限于每個(gè) DP 模型的每個(gè) TP 維度（即，對(duì)于給定的 k=1…16，在每個(gè) z=k 處的一個(gè) 4x4 的 2D 平面）。DP 維度只有一個(gè)環(huán)繞連接，以便每行成為一個(gè)大小為 16 的水平環(huán)。

你可能已經(jīng)注意到，還有一種 8x8x16 的替代拓?fù)洌?2x2 DP 維度），但這變得更加復(fù)雜，因?yàn)槲覀兓旌狭?DP 和 TP 維度。具體來說，不清楚我們應(yīng)該如何為 y 軸構(gòu)建 OCS 環(huán)繞連接，同時(shí)為每個(gè) TP 維度容納扭曲環(huán)面。

多 Pod 層面 (又稱 "Multislice"; TPUv4 為 4096+ 芯片)

TPU 層次結(jié)構(gòu)的最后一層是多 Pod 層面。在這里，你可以將多個(gè) Pod 視為一臺(tái)大型機(jī)器。然而，Pod 之間的通信是通過數(shù)據(jù)中心網(wǎng)絡(luò) (Data-Center Network, DCN)完成的，其帶寬低于 ICI。

圖示多 Pod 訓(xùn)練如何配置

PaLM 就是這樣訓(xùn)練的。它花了 56 天在 6144 個(gè) TPUv4（2個(gè) Pod）上進(jìn)行訓(xùn)練。下面你可以看到 6 個(gè) Pod 上的 TPU 作業(yè)分配：綠色是 PaLM，紅色是未分配，其余是其他作業(yè)。注意，每個(gè)方塊是一個(gè) 4x4x4 的 TPU 立方體。

實(shí)現(xiàn)這一點(diǎn)本身就很困難，但更令人印象深刻的是對(duì)開發(fā)者體驗(yàn)的關(guān)注。具體來說，是關(guān)注 “我們?nèi)绾尾拍鼙M可能地抽象化模型擴(kuò)展的系統(tǒng)/硬件部分？” 這個(gè)問題。

谷歌的答案是讓 XLA 編譯器負(fù)責(zé)協(xié)調(diào)大規(guī)模芯片間的通信。通過研究人員提供的正確標(biāo)志（即 DP、FSDP、TP 的并行維度、切片數(shù)量等），XLA 編譯器會(huì)為手頭的 TPU 拓?fù)洳迦胝_的分層集合通信操作。目標(biāo)是以盡可能少的代碼更改來實(shí)現(xiàn)大規(guī)模訓(xùn)練。

例如，這里是谷歌博客 [1] 中跨多個(gè)切片的 all-reduce 操作的分解。

這表明 XLA 編譯器負(fù)責(zé)處理切片之間和切片內(nèi)部的通信集合操作。

舉一個(gè)具體的例子，訓(xùn)練模型可能存在如下的 TPU 拓?fù)?。激活值通信通過 ICI 在切片內(nèi)發(fā)生，而梯度通信將通過 DCN 跨切片發(fā)生（即跨 DCN DP 維度）。

將圖表與現(xiàn)實(shí)聯(lián)系起來

當(dāng)你看到硬件的實(shí)際照片時(shí)，將圖表與現(xiàn)實(shí)聯(lián)系起來會(huì)很有幫助。下面是一個(gè)總結(jié)。

如果你看過谷歌 TPU 宣傳材料的圖片，你可能見過下面這張圖。

這是 8 個(gè) TPU 機(jī)架，每個(gè)單元是我們上面看到的 4x4x4 的 3D 環(huán)面。一個(gè)機(jī)架中的每一行有 2 個(gè)托盤，這意味著每行有 8 個(gè) TPU 芯片。

這是一個(gè) TPUv4 的單托盤：

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

下面是一個(gè)單芯片：

中間部分是 ASIC，周圍的 4 個(gè)塊是 HBM 堆棧。我們看到的是一個(gè) TPU v4，所以它內(nèi)部有 2 個(gè) TensorCore，因此總共有 4 個(gè) HBM 堆棧。

我沒找到 TPUv4 的芯片平面圖，所以這里有一個(gè) TPUv4i 的，它很相似，只是因?yàn)樗且粋€(gè)推理芯片，所以只有一個(gè) TensorCore。

請(qǐng)注意 CMEM 在 TPUv4i 的布局上占了相當(dāng)大的空間。

---

原文轉(zhuǎn)載自：https://henryhmko.github.io/posts/tpu/tpu.html，經(jīng)過翻譯、校對(duì)。