TPU v4

本文說明 Cloud TPU v4 的架構和支援的設定。

系統架構

每個 TPU v4 晶片都包含兩個 TensorCore。每個 TensorCore 都有四個矩陣乘法單元 (MXU)、一個向量單元和一個純量單元。下表列出 v4 TPU Pod 的主要規格。

下圖說明 TPU 4 晶片。

如要進一步瞭解 TPU v4 的架構細節和效能特徵，請參閱「TPU v4：可光學重新設定的超級電腦，可透過硬體支援嵌入機器學習」。

3D 網格和 3D 環面

v4 TPU 可直接連線至 3D 空間中最近的鄰近晶片，進而形成網路連線的 3D 網格。在切片中，連線可設為 3D 環面，其中拓樸 AxBxC 為 2A=B=C 或 2A=2B=C，且每個維度都是 4 的倍數。例如 4x4x8、4x8x8 或 12x12x24。一般來說，3D 圓環設定的效能會優於 3D 網格設定。詳情請參閱「扭曲的圓環拓撲」。

TPU v4 相較於 v3 的效能優勢

本節將說明在 TPU v4 上執行訓練範例指令碼的記憶體節省方式，以及 TPU v4 相較於 TPU v3 的效能提升情形。

記憶體系統

非均勻記憶體存取 (NUMA) 是一種電腦記憶體架構，適用於有多個 CPU 的機器。每個 CPU 都能直接存取一組高速記憶體。CPU 及其記憶體稱為 NUMA 節點。NUMA 節點會連結至彼此相鄰的其他節點。來自一個 NUMA 節點的 CPU 可以存取另一個 NUMA 節點中的記憶體，但這種存取速度比在 NUMA 節點內存取記憶體的速度慢。

在多 CPU 機器上執行的軟體可將 CPU 所需的資料放入其 NUMA 節點，藉此提高記憶體吞吐量。如要進一步瞭解 NUMA，請參閱維基百科的「Non Uniform Memory Access」一文。

您可以將訓練指令碼繫結至 NUMA 節點 0，充分利用 NUMA 區域優勢。

如要啟用 NUMA 節點繫結，請按照下列步驟操作：

1. 建立 TPU v4。

2. 使用 SSH 連線至 TPU。

3. 安裝 numactl 指令列工具。numactl 可讓您使用特定的 NUMA 排程或記憶體放置政策執行程序。

    $ sudo apt-get update
    $ sudo apt-get install numactl

4. 將指令碼繫結至 NUMA 節點 0。將 your-training-script 替換為訓練指令碼的路徑。

    $ numactl --cpunodebind=0 python3 your-training-script

啟用 NUMA 節點繫結的條件如下：

• 無論架構為何，如果工作負載高度依賴 CPU 工作負載 (例如圖像分類、推薦工作負載)。
• 如果您使用的是沒有 -pod 後置字元的 TPU 執行階段版本 (例如 tpu-vm-tf-2.10.0-v4)。

其他記憶體系統差異：

• v4 TPU 晶片在整個晶片中提供統一的 32 GiB HBM 記憶體空間，可讓兩個晶片上的 TensorCore 之間的協調更加順暢。
• 採用最新的記憶體標準和速度，改善 HBM 效能。
• 改善 DMA 效能設定檔，內建支援以 512B 精細度進行高效能跨步。

TensorCores

• 提供兩倍的 MXU 數量和更高的時脈率，可達到 275 個 TFLOPS 的最高效能。
• 2 倍的轉置和置換頻寬。
• 通用記憶體 (Cmem) 的載入-儲存記憶體存取模式。
• 支援更快的 MXU 權重載入頻寬和 8 位元模式，以便縮小批次大小並縮短推論延遲時間。

晶片間互連

每個晶片有六個互連連結，可支援網路直徑較小的網路拓撲。

其他

• x16 PCIE Gen 3 介面至主機 (直接連線)。
• 改善安全性模型。
• 提高能源效率。

設定

TPU v4 Pod 由 4096 個晶片組成，這些晶片透過可重新設定的高速連結互相連接。TPU v4 的彈性網路可讓您以多種方式連結相同大小的配量中的晶片。建立 TPU 配量時，您會指定所需的 TPU 版本和 TPU 資源數量。建立 TPU v4 配量時，您可以使用 AcceleratorType 和 AccleratorConfig 這兩種方式之一指定其類型和大小。

正在使用 AcceleratorType

如未指定拓撲，請使用 AcceleratorType。如要使用 AcceleratorType 設定 v4 TPU，請在建立 TPU 區塊時使用 --accelerator-type 標記。將 --accelerator-type 設為包含 TPU 版本和要使用的 TensorCore 數量的字串。舉例來說，如要建立含有 32 個 TensorCore 的 v4 切片，請使用 --accelerator-type=v4-32。

使用 gcloud compute tpus tpu-vm create 指令，搭配 --accelerator-type 標記建立具有 512 個 TensorCore 的 v4 TPU 切片：

  $ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-central2-b \
    --accelerator-type=v4-512 \
    --version=tpu-ubuntu2204-base

TPU 版本 (v4) 後面的數字會指定 TensorCore 數量。v4 TPU 中有兩個 TensorCore，因此 TPU 晶片數量為 512/2 = 256。

如要進一步瞭解如何管理 TPU，請參閱「管理 TPU」。如要進一步瞭解 Cloud TPU 的系統架構，請參閱「系統架構」。

正在使用 AcceleratorConfig

如要自訂 TPU 配量的實體拓撲，請使用 AcceleratorConfig。一般來說，如果切片數量超過 256 個晶片，就必須進行效能調整。

如要使用 AcceleratorConfig 設定 v4 TPU，請使用 --type 和 --topology 標記。將 --type 設為要使用的 TPU 版本，並將 --topology 設為配量中 TPU 晶片的實體排列方式。

您可以使用 3 元組 AxBxC 指定 TPU 拓撲，其中 A<=B<=C，且 A、B、C 全都小於或等於 4，或全都是 4 的整數倍數。值 A、B 和 C 是三個維度中每個維度的積木數量。舉例來說，如要建立含有 16 個晶片的 v4 區塊，您必須設定 --type=v4 和 --topology=2x2x4。

使用 gcloud compute tpus tpu-vm create 指令建立 v4 TPU 切片，其中 128 個 TPU 晶片以 4x4x8 陣列排列：

  $ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-central2-b \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-ubuntu2204-base

拓樸圖中 2A=B=C 或 2A=2B=C 也具有拓樸圖變化版本，可針對全對全通訊進行最佳化，例如 4×4×8、8×8×16 和 12×12×24。這就是所謂的「扭曲圓環」拓樸。

下圖顯示一些常見的 TPU 4 版拓撲。

較大的切片可由一或多個 4x4x4 的「立方體」組成。

如要進一步瞭解如何管理 TPU，請參閱「管理 TPU」。如要進一步瞭解 Cloud TPU 的系統架構，請參閱「系統架構」。

扭曲的圓環拓撲

部分 v4 3D 環面形切片形狀可選擇使用所謂的扭曲環面形拓撲。舉例來說，兩個 v4 立方體可以排列為 4x4x8 切片或 4x4x8_twisted。扭曲拓撲結構可提供更高的分割頻寬。舉例來說，相較於非扭曲的 4x4x8 切片，採用 4x4x8_twisted 拓撲結構的切片可理論上提高 70% 的雙分頻寬。增加二分頻寬對於使用全域通訊模式的工作負載相當實用。扭曲拓撲結構可改善大部分模型的效能，其中以TPU 嵌入工作負載最為顯著。

對於使用資料平行處理做為唯一平行處理策略的工作負載，扭曲的拓樸結構可能會略為提升效能。對於 LLM，使用扭曲拓樸結構的效能可能因並行處理類型 (DP、MP 等) 而異。最佳做法是訓練有/無扭曲拓樸結構的大型語言模型，以便判斷哪一種做法可為模型提供最佳效能。在 FSDP MaxText 模型上進行的部分實驗中，使用扭曲拓撲結構可改善 1 到 2 個 MFU。

扭曲拓樸的主要優點是，它可將非對稱的圓柱拓樸 (例如 4×4×8) 轉換為密切相關的對稱拓樸。對稱拓撲結構有許多優點：

• 改善負載平衡
• 更高的二分頻寬
• 縮短封包路徑

這些優點最終將轉化為許多全球通訊模式的效能提升。

TPU 軟體支援在切片上扭曲的圓環，其中每個維度的大小等於或為最小維度的兩倍。例如 4x4x8、4×8×8 或 12x12x24。

舉例來說，請考慮這個 4×2 環面拓撲，其中 TPU 在切片中標示了 (X,Y) 座標：

為了方便說明，這張拓樸圖中的邊都以無向邊顯示。實際上，每個邊緣都是 TPU 之間的雙向連線。如圖所示，我們將這個格線一側與另一側之間的邊緣稱為「迴轉邊緣」。

透過扭曲這個拓樸結構，我們最終得到完全對稱的 4×2 扭曲圓柱拓樸結構：

這張圖與上一個圖之間唯一的差異，就是 Y 包覆邊緣。這些 TPU 已改為連線至 X+2 mod 4 座標的 TPU，而非連線至具有相同 X 座標的 TPU。

同樣的概念也適用於不同的維度大小和維度數量。只要每個維度的大小等於或為最小維度的兩倍，產生的網路就會對稱。

如要進一步瞭解如何在建立 Cloud TPU 時指定扭曲多邊形設定，請參閱「使用 AcceleratorConfig」。

下表列出支援的扭曲拓撲，以及與未扭曲拓撲相比，使用扭曲拓撲時理論上可增加的二分頻寬。

TPU v4 拓撲變化版本

含有相同數量晶片的某些拓樸結構可以不同方式排列。舉例來說，如果 TPU 配量有 512 個晶片 (1024 個 TensorCore)，則可使用下列拓撲進行設定：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 支援以下 TPU v4 子集，其大小小於 64 個晶片，也就是 4x4x4 立方體。您可以使用 TensorCore 名稱 (例如 v4-32) 或拓樸 (例如 2x2x4) 建立這些小型 v4 拓樸：

大型 v4 拓撲

TPU v4 切片可以 64 個晶片為單位增加，且形狀在所有三個維度上都是 4 的倍數。維度必須以遞增方式排序。下表列出幾個範例。其中幾個拓撲是「自訂」拓撲，只能使用 --type 和 --topology 標記建立，因為安排方塊的方式不只一種。