Cloud TPU 效能指南

排解 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 大小在 v6e 中為 256 x 256，在 v6e 之前的版本中則為 128 x 128。詳情請參閱「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 效能的完整指南，請參閱「TPU7x (Ironwood) 效能最佳化」。