谷歌撰文，深度剖析TPU
關於TPU 的工作原理、它們如何聯網以實現多晶片訓練和推理，以及它們如何限制我們最喜歡的演算法的效能。雖然這看起來有點枯燥，但對於真正提高模型效率來說，它非常重要。
什麼是TPU？
TPU 基本上是一個專門用於矩陣乘法的計算核心（稱為TensorCore），連接到一堆快速記憶體（稱為高頻寬記憶體或HBM），以下是圖表：
您可以將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 執行元素積的範例：
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每主機每秒字節數。為了完整起見，我們同時包括單向和雙向頻寬。通常，當我們能夠存取完整的環繞環時，雙向頻寬是更有用的數字，而單向頻寬更符合硬體。(半導體產業觀察)