谷歌撰文：深度剖析TPU技術

DeepSeek核心技術：模型訓練、優化及數據處理的技術精髓（合集）

《42篇半導體行業深度報告&圖譜（合集）

TPU 基本上是一個專門用于矩陣乘法的計算核心（稱為 TensorCore），連接到一堆快速內存（稱為高帶寬內存或 HBM），以下是圖表：

圖： TPU 芯片的基本組件。TensorCore 是左側灰色框，包含矩陣乘法單元 (MXU)、矢量單元 (VPU) 和矢量內存 (VMEM)。

您可以將 TensorCore 視為一款出色的矩陣乘法機，但它還有一些其他值得注意的功能。TensorCore 有三個關鍵單元：

MXU （矩陣乘法單元）是 TensorCore 的核心。對于大多數 TPU 代，它執行一次矩陣bfloat16[8,128] @ bf16[128,128] -> f32[8,128]乘法1每 8 個周期使用一個脈動陣列。

在 TPU v5e 上，1.5GHz 時每個 MXU大約為5e13bf16 FLOPs/s。大多數 TensorCores 有 2 個或 4 個 MXU，因此例如 TPU v5e 的總 bf16 FLOPs/s 為2e14。

TPU 還支持具有更高吞吐量的較低精度矩陣乘法（例如，每個 TPU v5e MXU 可以執行4e14int8 OPs/s）。

VPU （矢量處理單元）執行一般數學運算，如 ReLU 激活或矢量之間的逐點加法或乘法。這里還執行縮減（求和）。

VMEM（矢量內存）是位于 TensorCore 中靠近計算單元的片上暫存器。它比 HBM 小得多（例如，TPU v5e 上為 128 MiB），但到 MXU 的帶寬要高得多。VMEM 的運行方式有點像 CPU 上的 L1/L2 緩存，但要大得多，而且由程序員控制。HBM 中的數據需要復制到 VMEM 中，然后 TensorCore 才能使用它進行任何計算。

TPU 在矩陣乘法方面非常非常快。這主要是它們所做的事情，而且它們做得很好。TPU v5p是迄今為止最強大的 TPU 之一，可以達到2.5e14bf16 FLOPs/秒/核心或5e14bf16 FLOPs/秒/芯片。一個由 8960 個芯片組成的 pod 可以達到 4 exaflops/秒。這可真不少。這是世界上最強大的超級計算機之一。而且谷歌有很多這樣的超級計算機。

上圖還包括一些其他組件，如 SMEM 和標量單元，它們用于控制流處理，并在附錄 C中進行了簡要討論，但理解起來并不重要。另一方面，HBM 很重要，而且相當簡單：

HBM（高帶寬內存）是一大塊快速內存，用于存儲供 TensorCore 使用的張量。HBM 的容量通常為數十 GB 的數量級（例如，TPU v5e 有 16GiB 的 HBM）。

當需要計算時，張量會通過 VMEM（見下文）從 HBM 流出到 MXU，然后將結果從 VMEM 寫回 HBM。

HBM 和 TensorCore（通過 VMEM）之間的帶寬稱為“HBM 帶寬”（通常約為 1-2TB/秒），限制了在內存受限工作負載中計算的速度。

通常，所有 TPU 操作都是流水線式和重疊式的。要執行 matmul X?A → Y，TPU 首先需要復制矩陣塊和從 HBM 傳輸到 VMEM，然后將它們加載到 MXU 中，MXU 將 8x128 的塊相乘（例如）和 128x128（用于)，然后將結果逐塊復制回 HBM。為了高效地完成此操作，matmul 是流水線式的，因此往返于 VMEM 的復制與 MXU 工作重疊。這允許 MXU 繼續工作，而不必等待內存傳輸，從而使 matmul 受計算限制，而不是受內存限制。

以下是如何從 HBM 執行元素積的示例：

圖：上圖展示了在 TPU 上執行的逐點乘積，其中字節從 HBM 加載。請注意字節如何以塊的形式從內存中流出，并且部分結果如何通過管道傳輸回來，而無需等待整個數組實現。

matmul 看起來幾乎完全相同，只是它會加載到 MXU 而不是 VPU/矢量單元，并且加載和存儲會以不同的順序發生，因為相同的權重塊用于多個激活塊。您可以看到數據塊流入 VMEM，然后流入 VREG（矢量寄存器），然后流入矢量單元，然后返回 VMEM 和 HBM。正如我們將要看到的，如果從 HBM 到 VMEM 的加載速度比矢量單元（或 MXU）中的?FLOP?慢，我們就會變得“帶寬受限”，因為我們正在耗盡 VPU 或 MXU 的工作量。

關鍵要點： TPU 非常簡單。它們將權重從 HBM 加載到 VMEM，然后從 VMEM 加載到脈動陣列中，該陣列每秒可執行約 200 萬億次乘加運算。HBMVMEM 和 VMEM脈動陣列帶寬對 TPU 能夠有效執行的計算設定了基本限制。

VMEM 和算術強度： VMEM 比 HBM 小得多，但其到 MXU 的帶寬高得多。正如我們在第 1 節中看到的，這意味著如果算法可以將其所有輸入/輸出放入 VMEM 中，則不太可能遇到通信瓶頸。當計算的算術強度較差時，這尤其有用：VMEM 帶寬大約比 HBM 帶寬高 22 倍，這意味著從 VMEM 讀取/寫入 MXU 操作只需要 10-20 的算術強度即可實現峰值 FLOPs 利用率。這意味著如果我們可以將權重放入 VMEM 而不是 HBM，我們的矩陣乘法可以在更小的批量大小下受 FLOPs 約束。這意味著從根本上具有較低算術強度的算法仍然可以高效。VMEM 太小了，這通常是一個挑戰。

TPU 芯片通常（但并非總是）由兩個共享內存的 TPU 核心組成，可以視為一個具有兩倍 FLOP 的大型加速器。自 TPU v4（稱為“兆核”）以來一直如此。在較舊的 TPU 芯片上，它們具有單獨的內存，被視為兩個單獨的加速器（TPU v3 及更早版本）。像 TPU v5e 這樣的推理優化芯片每個芯片只有一個 TPU 核心。

芯片以 4 個為一組排列在通過 PCIe 網絡連接到CPU 主機的“托盤”（tray）上。這是大多數讀者熟悉的格式，即通過?Colab?或單個 TPU-VM 公開的 4 個芯片（8 個核心，但通常被視為 4 個邏輯兆核）。對于像 TPU v5e 這樣的推理芯片，我們每個主機有 2 個托盤，而不是 1 個，但每個芯片只有 1 個核心，這樣我們就有 8 個芯片 = 8 個核心。

PCIe 帶寬有限：與 HBM 一樣VMEM 鏈接，CPUHBM PCIe 連接具有特定帶寬，限制了從主機內存加載到 HBM 或反之亦然的速度。例如，TPU v4 的 PCIe 帶寬為單向 16GB/秒，因此比 HBM 慢近 100 倍。我們可以將數據加載/卸載到主機 (CPU) RAM 中，但速度不是很快。

TPU 網絡

芯片通過 Pod 中的 ICI 網絡相互連接。在老一代（TPU v2 和 TPU v3）、推理芯片（例如 TPU v5e）和 Trilium（TPU v6e）中，ICI（“芯片間互連”）連接 4 個最近的鄰居（通過邊緣鏈接形成 2D 環面）。TPU v4 和 TPU v5p 連接到最近的 6 個鄰居（形成 3D 環面）。請注意，這些連接不通過其主機，它們是芯片之間的直接鏈接。

環形結構減少了任意兩個節點之間的最大距離到，使通信速度更快。TPU 還具有“扭曲圓環”配置，該配置將圓環包裹在類似莫比烏斯帶的拓撲中，以進一步縮短節點之間的平均距離。

TPU pod（由 ICI 連接）可以變得非常大：最大 pod 尺寸（稱為超級 pod）適用16x16x16于 TPU v4 和16x20x28TPU v5p。這些大型 pod 由可重構的芯片立方體組成，4x4x4通過光學環繞鏈路連接5我們可以重新配置它來連接非常大的拓撲。

也可以請求較小的拓撲（例如2x2x1，2x2x2），盡管沒有回繞。這是一個重要的警告，因為它通常會使大多數通信的時間加倍。任何完整立方體的倍數（例如4x4x4或4x4x8）都將具有由光開關提供的回繞。

對于 TPU v5e 和?Trillium，我們有由 2D 圓環組成的 pod。TPU 16x16v5e 和 v6e (Trillium) 無法擴展到 16x16 圓環之外，但 pod 仍可通過標準數據中心網絡 (DCN) 相互通信。同樣，可以請求較小的拓撲，而無需繞回 dims< 16。

這種最近鄰連接是 TPU 和 GPU 之間的一個關鍵區別。GPU 以全對全配置（稱為節點）連接多達 256 個 H100，而不是使用本地連接。一方面，這意味著 GPU 可以在單個低延遲跳躍中在節點內發送任意數據。另一方面，TPU 連接起來的成本要低得多，也更簡單，并且可以擴展到更大的拓撲，因為每個設備的鏈接數量是恒定的。

ICI相對于 DCN 非常快，但仍比 HBM 帶寬慢。例如，TPU v5p2.5e12每軸的HBM 帶寬為字節/秒（2.5 TB/秒），1e11ICI 帶寬為字節/秒（100 GB/秒），大約低 25 倍。這意味著當我們將模型拆分到多個芯片上時，我們需要小心，避免因跨設備通信速度較慢而導致 MXU 出現瓶頸。

多切片訓練：一組 ICI 連接的 TPU 稱為一個切片。不同的切片可以使用 DCN 相互連接，例如鏈接不同 pod 上的切片。由于 DCN 的連接速度比 ICI 慢得多，因此應該嘗試限制我們的計算需要等待 DCN 數據的時間。

TPU 很簡單，在大多數情況下可以被認為是連接到內存（超快）、通過 ICI 連接到其他芯片（相當快）以及通過 DCN 連接到數據中心其余部分的矩陣乘法單元（有點快）。

通信受到我們各種網絡帶寬的限制，速度依次為：

HBM 帶寬：TensorCore 與其相關的 HBM 之間。

ICI 帶寬：TPU 芯片與其最近的 4 個或 6 個鄰居之間。

PCIe 帶寬：CPU 主機與其關聯的芯片托盤之間。

DCN 帶寬：多個 CPU 主機之間，通常是未通過 ICI 連接的主機。

在一個切片內，TPU 僅通過 ICI 與其最近的鄰居相連。這意味著切片內遠距離芯片之間的 ICI 通信需要先跳過中間芯片。

權重矩陣需要在兩個維度上至少填充到 128 大小（TPU v6 上為 256）才能填滿 MXU（事實上，較小的軸會填充到 128）。

精度較低的矩陣乘法往往更快。對于支持此功能的幾代產品，TPU 執行 int8 或 int4 FLOP 的速度大約比執行 bfloat16 FLOP 快 2 倍/4 倍。VPU 操作仍在 fp32 中執行。

為了避免 TPU 計算單元出現瓶頸，我們需要確保每個通道上的通信量與其速度成正比。以下是TPU的一些具體數字：

主機大小是指連接到單個主機的 TPU 的拓撲結構（例如，TPU v5e 有一個 CPU 主機，以 4x2 拓撲結構連接到 8 個 TPU）。以下是互連圖：

我們同時包括單向（單向）帶寬和雙向（雙向）帶寬，因為單向帶寬更符合硬件，但雙向帶寬在涉及全環的方程中更常出現。

PCIe 帶寬通常約為1.5e10每芯片每秒字節數，而 DCN 帶寬通常約為2.5e10每主機每秒字節數。為了完整起見，我們同時包括單向和雙向帶寬。通常，當我們能夠訪問完整的環繞環時，雙向帶寬是更有用的數字，而單向帶寬更符合硬件。

內容轉載：半導體行業觀察

參考鏈接：https://jax-ml.github.io/scaling-book/tpus/