Cloud TPU 效能指南

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

TPU 模型效能

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

1. 模型受輸入內容限制

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

2. 由於分片 (在核心之間分割批次)，批次大小過小

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

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

3. 記憶體管理調整

   您可以使用 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
   ```
  
  This output will tell you how long it took to initialize the TPU and
  the size of the premapped buffer.
  

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

  TPU_PREMAPPED_BUFFER_SIZE: Sets the total size (in bytes) of the
  pre-mapped buffer region.
  TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: Sets the maximum size of
  a single buffer that can be allocated from the pre-mapped region.
  

  例如，您可以：

   export TPU_PREMAPPED_BUFFER_SIZE=4294967296
  

  設定緩衝區大小，並：

   export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES
   ```
   to enable it.
  
   This export sets the size to the default.
  

• 指引：如果懷疑主機與裝置間的資料傳輸是瓶頸，請調整 TPU_PREMAPPED_BUFFER_SIZE 的值。監控主機記憶體用量和模型效能，找出最佳平衡點。在大多數情況下，預設值已足夠。

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 的張量維度，避免填補。

張量維度

如要達到尖峰 FLOPS，矩陣乘積的維度應大於所用 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 = tmp[i] * z[i];
    }

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

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

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

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

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

廣播

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

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

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

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

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