TPU 架構

Tensor Processing Unit (TPU) 是 Google 設計的特殊應用積體電路 (ASIC),用於加速機器學習工作負載。Cloud TPU 是 Google Cloud 服務,可讓 TPU 做為可調整的資源使用。

TPU 可快速執行矩陣運算,因此非常適合機器學習工作負載。您可以使用 PytorchJAX 等架構,在 TPU 上執行機器學習工作負載。

TPU 如何運作?

如要瞭解 TPU 的運作方式,建議您先瞭解其他加速器如何解決訓練機器學習模型的運算挑戰。

CPU 的運作方式

CPU 是基於 von Neumann 架構的通用處理器。這表示 CPU 可與軟體和記憶體搭配使用,如以下所示:

CPU 運作方式的插圖

CPU 的最大優點是其彈性。您可以為許多不同類型的應用程式,在 CPU 上載入任何類型的軟體。舉例來說,您可以使用 CPU 在電腦上進行文字處理、控制火箭引擎、執行銀行交易,或使用類神經網路將圖片分類。

CPU 會從記憶體載入值,對值執行計算,並將每項計算的結果儲存回記憶體。與計算速度相比,記憶體存取速度較慢,可能會限制 CPU 的總傳輸量。這通常稱為 von Neumann 瓶頸

GPU 運作方式

為了獲得更高的總處理量,GPU 在單一處理器中包含數千個算術邏輯單元 (ALU)。新型 GPU 通常含有 2,500 到 5,000 個 ALU。大量的處理器表示您可以同時執行數千個乘積和加法。

GPU 運作方式的插圖

此 GPU 架構可在具有大量平行運作 (例如類神經網路中的矩陣運算) 的應用程式上順利運作。實際上,在深度學習的典型訓練工作負載上,GPU 可以提供較 CPU 高幾十倍的總處理量。

但 GPU 仍是通用處理器,可支援許多不同的應用程式和軟體。因此,GPU 與 CPU 有相同的問題。對於數千個 ALU 中的每個計算,GPU 必須存取暫存器或共用記憶體,以讀取運算元並儲存中繼計算結果。

TPU 運作方式

Google 將 Cloud TPU 設計為類神經網路工作負載專用的矩陣處理器。TPU 無法執行文書處理程式、控制火箭引擎或執行銀行交易,但可以用非常快的速度來處理類神經網路中使用的大量矩陣運算。

TPU 的主要工作是矩陣處理,也就是乘法和累加運算的組合。TPU 包含數千個乘法累計器,彼此直接連接,形成大型物理矩陣。這稱為脈動陣列架構。Cloud TPU v3 在單一處理器上含有兩個 128 x 128 ALU 的脈動陣列。

TPU 主機會將資料串流至饋入佇列。TPU 會從 infeed 佇列載入資料,並將資料儲存在 HBM 記憶體中。運算完成後,TPU 會將結果載入輸出佇列。接著,TPU 主機會從外送佇列讀取結果,並將結果儲存在主機的記憶體中。

為了執行矩陣運算,TPU 會將 HBM 記憶體中的參數載入至矩陣乘法單元 (MXU)。

TPU 如何從記憶體載入參數的插圖

接著,TPU 會從 HBM 記憶體載入資料。執行每個乘法時,結果會傳遞到下一個乘法累加器。輸出結果是資料和參數間所有乘法結果的總和。矩陣相乘過程中不需要記憶體存取權。

TPU 如何從記憶體載入資料的插圖

因此,TPU 可以在類神經網路計算時達到較高的運算總處理量。

TPU 系統架構

以下各節將說明 TPU 系統的重要概念。如要進一步瞭解常見的機器學習術語,請參閱機器學習術語彙

如果您是 Cloud TPU 新手,請參閱 TPU 說明文件首頁

TPU 晶片

TPU 晶片包含一或多個 TensorCore。TensorCore 的數量取決於 TPU 晶片的版本。每個 TensorCore 都包含一或多個矩陣乘法單元 (MXU)、一個向量單元和一個純量單元。如要進一步瞭解 TensorCore,請參閱「專門用於訓練深度神經網路的超級電腦」。

MXU 由 256 x 256 (TPU v6e) 或 128 x 128 (TPU 6e 之前的版本) 乘法累計器組成,位於脈動陣列中。MXU 可在 TensorCore 中提供大部分的運算能力。每個 MXU 在每個週期中都能執行 16000 次乘積累加運算。所有乘法運算都會使用 bfloat16 輸入值,但所有累加運算都是以 FP32 數字格式執行。

向量單元用於一般運算,例如啟用和軟性最大值。標量單位用於控制流程、計算記憶體位址和其他維護作業。

TPU Pod

TPU Pod 是透過專用網路組合在一起的連續 TPU 組合。TPU Pod 中的 TPU 晶片數量取決於 TPU 版本。

配量

配量是一組晶片,這些晶片都位於同一個 TPU Pod 中,並透過高速晶片間互連網路 (ICI) 連接。視 TPU 版本而定,配量會以晶片或 TensorCore 來說明。

方塊形狀方塊拓撲也指的是切片形狀。

多配量與單一配量

Multislice 是一組切片,可將 TPU 連線擴展至晶片間互連服務 (ICI) 連線之外,並運用資料中心網路 (DCN) 傳輸超出切片的資料。每個時間片段中的資料仍會透過 ICI 傳送。透過這種混合式連線方式,多配量可在各個切片中啟用平行處理,讓您在單一工作中使用比單一切片可容納的更多 TPU 核心。

您可以使用 TPU 在單一切片或多個切片上執行工作。詳情請參閱多切面簡介

TPU 立方體

互連 TPU 晶片的 4x4x4 拓撲。這項功能僅適用於 3D 拓樸 (從 TPU v4 開始)。

SparseCore

SparseCores 是資料流處理器,可加快仰賴建議模型中嵌入項目的模型處理速度。v5p 每個晶片包含四個 SparseCores,v6e 每個晶片包含兩個 SparseCores。

Cloud TPU ICI 復原能力

ICI 彈性有助於改善光纖連結和光纖電路切換器 (OCS) 的容錯能力,這些連結和切換器可在立方體之間連結 TPU。(立方體內的 ICI 連線會使用銅線連結,因此不會受到影響)。ICI 復原功能可讓 ICI 連線繞過 OCS 和光學 ICI 故障。因此,雖然 ICI 效能會暫時降低,但可改善 TPU 區塊的排程可用性。

對於 Cloud TPU v4 和 v5p,系統會預設啟用 ICI 復原功能,以便處理一個立方體或更大的切片,例如:

  • 指定加速器類型時的 v5p-128
  • 指定加速器設定時為 4x4x4

TPU 版本

TPU 晶片的確切架構取決於您使用的 TPU 版本。每個 TPU 版本也支援不同的切片大小和設定。如要進一步瞭解系統架構和支援的設定,請參閱以下頁面:

TPU 雲端架構

Google Cloud 會透過 TPU VM 提供 TPU 做為運算資源。您可以直接使用 TPU VM 執行工作負載,也可以透過 Google Kubernetes Engine 或 Vertex AI 使用 TPU VM。以下各節將說明 TPU 雲端架構的主要元件。

TPU VM 架構

TPU VM 架構可讓您使用 SSH 直接連線至與 TPU 裝置實際連線的 VM。TPU VM (也稱為 worker) 是執行 Linux 的虛擬機器,可存取基礎 TPU。您擁有 VM 的根目錄存取權,因此可以執行任意程式碼。您可以存取編譯器和執行階段偵錯記錄和錯誤訊息。

TPU VM 架構

單一主機、多主機和子主機

TPU 主機是指在連接 TPU 硬體的實體電腦上執行的 VM。TPU 工作負載可使用一或多部主機。

單主機工作負載僅限於一個 TPU VM。多主機工作負載會將訓練工作分散至多個 TPU VM。子主機工作負載不會使用 TPU VM 上的所有晶片。

TPU 節點架構 (已淘汰)

TPU 節點架構包含使用者 VM,可透過 gRPC 與 TPU 主機通訊。使用這種架構時,您無法直接存取 TPU 主機,因此很難對訓練和 TPU 錯誤進行偵錯。

TPU 節點架構

從 TPU 節點轉換至 TPU VM 架構

如果您有使用 TPU 節點架構的 TPU,請按照下列步驟將其識別、刪除,然後重新佈建為 TPU VM。

  1. 前往 TPU 頁面:

    前往 TPU

    在「Architecture」標題下方找出 TPU 及其架構。如果架構為「TPU VM」,您就不需要採取任何行動。如果架構為「TPU 節點」,您需要刪除並重新佈建 TPU。

  2. 刪除並重新佈建 TPU。

    如要瞭解如何刪除重新配置 TPU,請參閱「管理 TPU」。

後續步驟