Cloud TPU 效能指南

排解 TPU 效能問題的第一步是剖析模型。如要進一步瞭解如何擷取效能設定檔，請參閱「在 Cloud TPU 上剖析模型」。

TPU 模型效能

本節說明可能導致模型效能降低的一般問題，以及解決方法。

輸入受限模型

TPU 的運算速度非常快，為確保 TPU 不會閒置，請務必將資料穩定載入 TPU。實際做法取決於您載入及預先處理資料集的方式。 舉例來說，您可以使用 tf.data.TFRecordset() 和 num_parallel_reads 參數，平行讀取資料檔案。

因分片而導致批次大小較小

TPU 執行階段會將批次作業分割到 TPU 裝置的所有 8 個核心 (例如 v2-8 或 v3-8)。如果您指定全域批次大小為 128，則每個核心會收到批次大小為 16 (128 / 8)。

如要達到最佳記憶體用量，請使用可容納 TPU 記憶體空間的最大批次大小。每個 TPU 核心都會使用二維 8 X 128 向量暫存器，處理矩陣乘法。一般而言，批量大小應可被 8 或 128 整除。

記憶體管理調整

您可以使用記憶體相關的環境變數，微調低階執行階段行為。

TPU_PREMAPPED_BUFFER_SIZE

TPU_PREMAPPED_BUFFER_SIZE 會設定主機記憶體緩衝區的大小 (以位元組為單位)，該緩衝區會預先對應並固定，供 TPU 執行階段用於資料傳輸 (例如 DMA)。預設值為 4294967296 個位元組。值必須是 2^12 的倍數 (4KB = 4 * 1024 Bytes = 4096 = 2^12)。

以下範例為有效的 TPU_PRE_MAPPED_BUFFER_SIZE 值。

17179869184 = 2^34 = 2^22 * 2^12 (2^22 4KB pages will be premapped).
40000000000 = 5^10 * 2^12 = (5^10 4KB pages will be premapped).

增加這個大小可能會提升主機與 TPU 裝置之間的資料傳輸效能，尤其是張量較大或主機與裝置之間經常通訊的工作負載。不過，這也會增加釘選的主機記憶體量，減少其他程序可用的記憶體。

排解記憶體問題

如果預先對應的緩衝區不夠大，無法在程式執行階段分配記憶體，工作負載就會失敗，並傳回類似以下的 RESOURCE_EXHAUSTED 錯誤：

「Allocating buffer from premmaped region failed with: RESOURCE_EXHAUSTED: Attempting to allocate allocation_size. 這是不可能的。有 available_size 免費。」

如果緩衝區過大，TPU 初始化時間可能會大幅延長 (可能超過 15 秒)，導致 TPU 似乎卡住。

如要診斷這個問題，請檢查 TPU 執行階段記錄。這些記錄檔會詳細說明執行的作業，包括緩衝區的預先對應。您可以在 /tmp/tpu_logs/tpu_driver.INFO 中找到記錄，也可以設定環境變數 TPU_STDERR_LOG_LEVEL=0，直接將記錄列印到控制台。這項設定會產生類似以下的輸出內容：

I0604 12:45:24.926233   62136 tpu_hal.cc:214] Starting premapped memory manager initialization...
I0604 12:45:29.411218   62136 system.cc:1059] tpu::System initialized, current host id: 0, logical device ids: 0
I0604 12:45:29.411244   61600 tfrt_tpu_system_state.cc:216] CreateTpuSystemState: TPU initialization is successful and it took 5.583190661s
I0604 12:45:29.411267   61600 tfrt_tpu_system_state.cc:220] CreateTpuSystemState: using TPU host premapped buffer of size: 4294967296

這項輸出內容會顯示初始化 TPU 所需的時間，以及預先對應的緩衝區大小。

設定緩衝區大小

如果預先對應的緩衝區太小或太大，您可以使用下列環境變數手動設定緩衝區大小。

  • TPU_PREMAPPED_BUFFER_SIZE：設定預先對應緩衝區區域的總大小 (以位元組為單位)。

  • TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES：設定可從預先對應區域分配的單一緩衝區大小上限。

例如，您可以：

export TPU_PREMAPPED_BUFFER_SIZE=4294967296

設定緩衝區大小，並：

export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES

即可啟用。

匯出時，系統會將大小設為預設值。

如果懷疑主機裝置資料傳輸是瓶頸，請調整 TPU_PREMAPPED_BUFFER_SIZE 的值。監控主機記憶體用量和模型效能，找出最佳平衡點。一般來說，預設值已足以應付大多數用途。

tcmalloc 設定

在 Cloud TPU VM 上，系統預設會使用 tcmalloc 程式庫，改善模型效能，避免頻繁分配大量記憶體。這是透過 LD_PRELOAD 環境變數設定。

不過，對於某些工作負載 (例如具有非常大的嵌入資料表分配的 DLRM)，tcmalloc 可能會導致速度變慢。在這種情況下，您可以在執行訓練指令碼前，取消設定殼層工作階段中的 LD_PRELOAD 變數，藉此還原為標準 malloc 函式：

unset LD_PRELOAD

網路效能最佳化

以下各節說明如何設定最大傳輸單元 (MTU) 並在多重切片環境中使用多重 NIC，進而提升網路效能。

設定 MTU

如要獲得最佳網路效能，請使用 MTU (最大傳輸單位) 為 8,896 的網路。

根據預設，虛擬私有雲 (VPC) 僅提供 1,460 位元組的 MTU，這會導致網路效能不佳。您可以將虛擬私有雲網路的 MTU 設為 1,300 到 8,896 位元組 (含) 之間的任何值。常見的自訂 MTU 大小為 1,500 個位元組 (標準乙太網路) 或 8,896 個位元組 (最大可能值)。詳情請參閱「有效的虛擬私有雲網路 MTU 大小」。

如要進一步瞭解如何變更現有或預設網路的 MTU 設定，請參閱「變更虛擬私有雲網路的 MTU 設定」。

為多配量使用多 NIC 選項

在由數千個 TPU 晶片組成的 Multislice 環境中訓練大型模型時，透過資料中心網路 (DCN) 進行切片間通訊可能會造成瓶頸。如要提升網路繫結工作負載的網路頻寬，可以使用多個 NIC 增加 TPU VM 的網路介面數量。使用多個 NIC 時，每個 TPU VM 都會分配到額外的網路介面，且每個介面都會連線至專屬的 VPC 網路，進而提高整體網路輸送量。額外 NIC 必須位於互斥的 IP 範圍內。

如要進一步瞭解如何在使用 Google Kubernetes Engine (GKE) 時啟用多網路功能，請參閱「在 TPU Trillium 或 Ironwood (TPU7x) 上，不使用 hostNetwork 提升網路效能」。如要查看使用多重 NIC 和 XPK 的範例，請參閱使用 XPK 建立支援多重 NIC 的叢集。

XLA 編譯器最佳化

XLA 是機器學習專用的編譯器，可為 TPU、CPU、GPU 和其他平台產生二進位檔。XLA 是標準 TensorFlow 程式碼集的一部分，但也可以用於 PyTorch 和 JAX 模型。Cloud TPU 的模型會轉譯為 XLA 圖表，然後由 XLA 編譯為 TPU 可執行檔。如要進一步瞭解 XLA，請參閱「XLA：機器學習的最佳化編譯器」。

邊框間距

如要有效運用 TPU 記憶體，請將資料結構化，以便將資料分塊為 128 x 8 的大小。如果矩陣計算的資料未填滿整個 128 x 8 區塊，XLA 編譯器會填補張量。填充有兩項缺點：

  1. 填充張量會使 TPU 核心使用率過低。
  2. 填充會增加張量所需的晶載記憶體儲存量，可能會造成記憶體不足的錯誤。

雖然填充是由 XLA 編譯器於必要時自動執行，但您可利用記憶體檢視器工具決定執行的填充量。您可以選擇適合 TPU 的張量維度，避免填補。

張量維度

如要達到尖峰 FLOP，矩陣乘法運算的維度應大於所用 TPU 版本的 MXU 大小。MXU 大小為 256 x 256 (適用於 v6e)，以及 128 x 128 (適用於 v6e 之前的版本)。詳情請參閱「Cloud TPU 系統架構」。

批量

XLA 編譯器會將儲存在 TPU HBM 記憶體中的張量大小向上取整，以便更有效率地執行運算。這項填補作業會在硬體層級以透明方式進行，不會影響結果。但在某些情況下，填充將造成記憶體用量和執行時間大幅增加。

TPU 執行階段會在記憶體中配置張量，盡量提高運算效率並減少填充。如要將記憶體負擔降至最低，並同時獲得最高的運算效率，就必須滿足下列其中一項條件：

  1. 總批量應為 64 的倍數 (每個 TPU 核心為 8)，且特徵維度大小應為 128 的倍數。

  2. 總批量應為 1024 的倍數 (每個 TPU 核心為 128)，且特徵維度大小應為 8 的倍數。

批量為 1024 且特徵維度為 128 的倍數時，將可獲得最佳效率。不過這樣的配置可能不適用於所有模型。

Fusion

融合是 XLA 編譯器用來最佳化程式的一般技術。融合作業是多個組成作業的組合，這些作業會一併執行。

例如請考量以下的運算系列：

    tmp = tf.add(x, y)
    result = tf.multiply(tmp, z)

此程式碼大致等同於下列虛擬程式碼：

    for (i = 0; i < element_count; i++) {
      tmp[i] = x[i] + y[i];
    }

    for (i = 0; i < element_count; i++) {
      result[i] = tmp[i] * z[i];
    }

若進行融合，陣列存取就會同時發生：

    for (i = 0; i < element_count; i++) {
      result[i] = (x[i] + y[i]) * z[i];
    }

在這個範例中，記憶體往返次數減少，XLA 也不需要為「tmp」分配任何空間。

融合是關鍵最佳化，以多種方式讓 Cloud TPU 受益：

  • 不需要在緩慢的主記憶體儲存中繼結果，減少記憶體傳輸。
  • 可提升未利用硬體單元的使用率。
  • 可減少模型記憶體使用率，原因是同時間需要上線的緩衝區數量減少。

廣播

如果兩個張量結合，其中各具有不同但相容的形狀，就會隱含產生廣播。

舉例來說，tf.add(vector, matrix) 需要將向量廣播至矩陣的形狀。運算結果的形狀與矩陣相同。詳情請參閱廣播陣列指南。

雖然廣播通常可以與消費者融合，但強制廣播可能會導致效能不佳和記憶體用量增加。

在下列範例中，廣播在新增向量情況下隱含，而矩陣無法與 argmax 融合，造成具體化的廣播：

`tf.argmax(tf.add(vector, zero_matrix), axis=0)`

Ironwood 雙晶片架構的效能建議

Ironwood 程式設計模型可讓您存取兩個 TPU 裝置，而非前幾代 (TPU v4 和 v5p) 使用的單一邏輯核心 (也稱為 MegaCore) 架構。這項變更可提高晶片的製造成本效益和效率。雖然這代表架構的轉移，但新設計可確保您能以最少的變更重複使用現有軟體模型。

如要讓雙晶片架構發揮最佳效能，建議採取下列做法：

  • 在晶片間使用張量平行處理：高頻寬 D2D 介面專為高效張量平行處理而設計。建議您將張量分割到兩個晶片上裝置。

  • 運用階層式集合：為盡可能提高通訊效率，請善用雙層網路階層：晶片內晶片間的超快速 D2D 連結，以及切片內的快速 ICI 連結。使用 SPMD (單一程式、多個資料) 的自動平行處理時，XLA 編譯器會自動產生階層式集合作業，為您處理這項作業。手動分割模型時，也應根據這個階層設計通訊模式。優先處理同一晶片上兩部裝置之間的通訊，再處理其他晶片上裝置之間的通訊。

  • 將通訊與運算重疊：為盡量提高硬體使用率，請將集體通訊作業 (例如 all-reduce) 卸載至 SparseCore。這些作業並未繫結至矩陣乘法單元 (MXU)，因此可在 TensorCore 繼續運算時，同時在 SparseCore 上執行。這項技術可恢復先前 MegaCore 架構中融合作業的固有部分效能優勢。

  • 將嵌入卸載至 SparseCore：在雙晶片設計中，嵌入資料表可跨越兩個晶片的 HBM 分區。為避免共用記憶體不足導致效能下降，請將嵌入式聚集作業卸載至 SparseCore。這項策略會運用高速 D2D 互連技術，在晶片間有效率地轉移嵌入向量。如要進一步瞭解 SparseCore 和嵌入模型，請參閱「深入瞭解大型嵌入模型 (LEM) 的 SparseCore」。

如要進一步瞭解 TPU7x 中的 Ironwood 架構，請參閱「TPU7x (Ironwood)」。