TPU v4
本文說明 Cloud TPU v4 的架構和支援的設定。
系統架構
每顆 TPU v4 晶片包含 2 個 TensorCore。每個 TensorCore 都有四個矩陣乘法單元 (MXU)、一個向量單元和一個純量單元。下表列出 v4 TPU Pod 的主要規格。
| 規格 | 值 |
|---|---|
| 每個晶片的尖峰運算效能 | 每秒 275 兆次浮點運算 (bf16 或 int8) |
| HBM2 容量和頻寬 | 32 GiB，每秒 1200 GB |
| 測得的最小/平均/最大功率 | 90/170/192 瓦 |
| TPU Pod 大小 | 4096 個晶片 |
| 互連網路拓撲 | 3D 網格 |
| 每個 Pod 的尖峰運算效能 | 每秒 110 萬兆次浮點運算 (bf16 或 int8) |
| 每個 Pod 的 All-reduce 頻寬 | 1.1 PB/秒 |
| 每個 Pod 的對分頻寬 | 24 TB/秒 |
下圖說明 TPU v4 晶片。
如要進一步瞭解 TPU v4 的架構詳細資料和效能特徵，請參閱「TPU v4：適用於機器學習的光學可重設超級電腦，支援嵌入」。
3D 網格和 3D 環面
v4 TPU 會直接連線至 3 維空間中最鄰近的晶片，形成 3D 網格的網路連線。如果拓撲 AxBxC 為 2A=B=C 或 2A=2B=C，且每個維度都是 4 的倍數，則可將連線配置為 3D 環面。例如 4x4x8、4x8x8 或 12x12x24。一般來說，3D 環面設定的效能會優於 3D 網格設定。詳情請參閱「扭曲環面拓撲」。
TPU v4 相較於 v3 的效能優勢
本節將說明在 TPU v4 上執行範例訓練指令碼的記憶體效率方式，以及 TPU v4 相較於 TPU v3 的效能提升幅度。
記憶體系統
非統一記憶體存取 (NUMA) 是適用於多 CPU 電腦的電腦記憶體架構。每個 CPU 都能直接存取高速記憶體區塊。CPU 和記憶體稱為 NUMA 節點。NUMA 節點會連接至彼此直接相鄰的其他 NUMA 節點。一個 NUMA 節點的 CPU 可以存取另一個 NUMA 節點的記憶體，但這種存取方式比存取 NUMA 節點內的記憶體慢。
在多 CPU 機器上執行的軟體可將 CPU 需要的資料放在 NUMA 節點中，進而提高記憶體輸送量。如要進一步瞭解 NUMA，請參閱維基百科的「Non Uniform Memory Access」。
您可以將訓練指令碼繫結至 NUMA 節點 0，充分發揮 NUMA 本機性的優勢。
如要啟用 NUMA 節點繫結：
安裝 numactl 指令列工具。numactl 可讓您執行程序，並採用特定 NUMA 排程或記憶體放置政策。
$ sudo apt-get update $ sudo apt-get install numactl
將指令碼程式碼繫結至 NUMA 節點 0。將 your-training-script 替換為訓練指令碼的路徑。
$ numactl --cpunodebind=0 python3 your-training-script
在下列情況下啟用 NUMA 節點繫結：
- 如果工作負載高度依賴 CPU 工作負載 (例如圖片分類、建議工作負載)，無論架構為何。
- 如果使用的 TPU 執行階段版本沒有 -pod 後置字元 (例如
tpu-vm-tf-2.10.0-v4)。
其他記憶體系統差異：
- v4 TPU 晶片在整個晶片中共有 32 GiB 的統合式 HBM 記憶體空間，可增進晶片上兩個 TensorCore 的協調效率。
- 採用最新記憶體標準和速度，提升 HBM 效能。
- 改善 DMA 效能設定檔，內建支援以 512B 粒度進行高效能步進。
TensorCores
- MXU 數量加倍，時脈頻率更高，最高可達 275 TFLOPS。
- 轉置與排序頻寬是原先的 2 倍。
- Common Memory (Cmem) 的載入/儲存記憶體存取模型。
- MXU 權重載入頻寬更快，且支援 8 位元模式，可降低批次大小並縮短推論延遲時間。
晶片間互連
每個晶片有六個互連連結，可啟用網路直徑較小的網路拓撲。
其他
- x16 PCIE Gen3 介面，可連線至主機 (直接連線)。
- 改良安全性模型。
- 提高能源效率。
設定
TPU v4 Pod 由 4096 個晶片組成，並以可重新設定的高速連結互連。TPU v4 的彈性網路可讓您以多種方式連結相同大小配量中的晶片。建立 TPU 配量時，您會指定所需的 TPU 版本和 TPU 資源數量。
您可以使用 3 元組 AxBxC 指定 TPU 拓撲，其中 A<=B<=C，且 A、B、C 皆 <= 4，或皆為 4 的整數倍數。A、B 和 C 值分別是三個維度中的晶片數量。
對於 2A=B=C 或 2A=2B=C 的拓撲，也有針對全對全通訊最佳化的拓撲變體，例如 4×4×8、8×8×16 和 12×12×24。 這些稱為「扭曲環面」拓撲。
下圖顯示一些常見的 TPU v4 拓撲。
較大的晶片切片可由一或多個 4x4x4 的晶片「立方體」建構而成。
如要進一步瞭解如何管理 TPU，請參閱「管理 TPU」。如要進一步瞭解 Cloud TPU 的系統架構，請參閱「系統架構」。
扭曲環面拓撲
部分 v4 3D 環面配量形狀可選擇使用所謂的扭曲環面拓撲。舉例來說，兩個 v4 立方體可以排列成 4x4x8 的配量或 4x4x8_twisted。扭曲拓撲可提供高得多的二分頻寬。舉例來說，與非扭曲的 4x4x8 切片相比，採用 4x4x8_twisted 拓撲的切片在理論上可增加 70% 的平分頻寬。如果工作負載使用全域通訊模式，增加二分頻寬就很有幫助。扭曲拓撲可提升大多數模型的效能，其中以大型 TPU 嵌入工作負載的效益最顯著。
如果工作負載只使用資料平行處理策略，扭曲拓撲的效能可能會稍微好一點。對於 LLM，使用扭曲拓撲的效能可能因平行處理類型 (DP、MP 等) 而異。最佳做法是使用扭曲拓撲訓練 LLM，並在不使用扭曲拓撲的情況下訓練 LLM，判斷哪種方式能為模型帶來最佳效能。在 FSDP MaxText 模型上進行的實驗顯示，使用扭曲拓撲可提升 1 到 2 個 MFU。
扭曲拓撲的主要優點是，可將非對稱環面拓撲 (例如 4×4×8) 轉換為密切相關的對稱拓撲。對稱拓撲有許多優點：
- 改善負載平衡
- 更高的二分頻寬
- 封包路徑較短
這些優勢最終將轉化為許多全球通訊模式的效能提升。
TPU 軟體支援切片上的扭曲環面，其中每個維度的大小等於或兩倍於最小維度的大小。例如 4x4x8、4×8×8 或 12x12x24。
舉例來說，請參考這個 4×2 環面拓撲，TPU 標示為配量中的 (X,Y) 座標：
為清楚起見，這個拓撲圖中的邊緣會顯示為無向邊緣。在實務上，每個邊緣都是 TPU 之間的雙向連線。如圖所示，我們將這個格線一側與對側之間的邊緣稱為環繞邊緣。
扭曲這個拓撲後，我們最終會得到完全對稱的 4×2 扭曲環面拓撲：
這個圖表與上一個圖表之間唯一的差異，就是 Y 軸的環繞邊緣。這些晶片不會連線至 X 座標相同的其他 TPU，而是會連線至座標為 X+2 mod 4 的 TPU。
這個概念可推廣至不同大小的維度和不同數量的維度。只要每個維度等於或兩倍於最小維度的大小，產生的網路就會對稱。
下表列出支援的扭曲拓撲，以及與非扭曲拓撲相比，這些拓撲的理論二分頻寬增加量。
| 扭曲拓撲 | 理論上，對分頻寬會增加 (與非扭曲環面相比) |
|---|---|
| 4×4×8_twisted | 約 70% |
| 8x8x16_twisted | |
| 12×12×24_twisted | |
| 4×8×8_twisted | 約 40% |
| 8×16×16_twisted |
TPU v4 拓撲變體
有些拓撲包含相同數量的晶片，但排列方式可能不同。舉例來說，您可以使用下列拓撲設定具有 512 個晶片 (1024 個 TensorCore) 的 TPU 配量：4x4x32、4x8x16 或 8x8x8。含有 2048 個晶片 (4096 個 TensorCore) 的 TPU 配量提供更多拓撲選項：4x4x128、4x8x64、4x16x32 和 8x16x16。
與特定晶片數量相關聯的預設拓撲，最類似於立方體。這個形狀可能是資料平行機器學習訓練的最佳選擇。其他拓撲可用於具有多種平行處理作業的工作負載 (例如模型和平行資料，或模擬的空間分割)。如果拓撲與使用的平行處理相符，這些工作負載的效能最佳。舉例來說，在 X 維度上放置 4 向模型平行處理，在 Y 和 Z 維度上放置 256 向資料平行處理，即符合 4x16x16 拓撲。
如果模型有多個平行處理維度，將平行處理維度對應至 TPU 拓撲維度，可獲得最佳效能。這些通常是資料和模型並行的大型語言模型 (LLM)。舉例來說，如果 TPU v4 配量的拓撲為 8x16x16，則 TPU 拓撲維度為 8、16 和 16。使用 8 向或 16 向模型平行處理 (對應至其中一個實體 TPU 拓撲維度) 的效能較佳。如果使用 4 向模型平行化，由於與任何 TPU 拓撲維度都不一致，因此並非最佳做法，但如果使用相同晶片數量的 4x16x32 拓撲，則會是最佳做法。
TPU v4 配置包含兩組，分別是拓撲小於 64 個晶片的配置 (小型拓撲)，以及拓撲大於 64 個晶片的配置 (大型拓撲)。
小型 v4 拓撲
Cloud TPU 支援下列小於 64 個晶片的 TPU v4 配量，也就是 4x4x4 立方體。您可以使用 TensorCore 型名稱 (例如 v4-32) 或拓撲 (例如 2x2x4) 建立這些小型 v4 拓撲：
| 名稱 (根據 TensorCore 數量) | 晶片數量 | 拓撲 |
| v4-8 | 4 | 2x2x1 |
| v4-16 | 8 | 2x2x2 |
| v4-32 | 16 | 2x2x4 |
| v4-64 | 32 | 2x4x4 |
大型 v4 拓撲
TPU v4 配量以 64 個晶片為增量單位，且所有三個維度的形狀都是 4 的倍數。維度必須遞增排序。下表列出幾個範例。其中幾種拓撲是「自訂」拓撲，只能使用 --type 和 --topology 標記建立，因為晶片排列方式不只一種。
| 名稱 (根據 TensorCore 數量) | 晶片數量 | 拓撲 |
| v4-128 | 64 | 4x4x4 |
| v4-256 | 128 | 4x4x8 |
| v4-512 | 256 | 4x8x8 |
自訂拓撲：必須使用 --type 和 --topology 標記 |
256 | 4x4x16 |
| v4-1024 | 512 | 8x8x8 |
| v4-1536 | 768 | 8x8x12 |
| v4-2048 | 1024 | 8x8x16 |
自訂拓樸：必須使用 --type 和 --topology 標記 |
1024 | 4x16x16 |
| v4-4096 | 2048 | 8x16x16 |
| … | … | … |