錯誤代碼：1000

類別：編譯時間：HBM OOM

這項錯誤表示程式需要的高頻寬記憶體 (HBM) 超過 TPU 裝置的實際可用量。

錯誤訊息示例：

RESOURCE_EXHAUSTED: XLA:TPU compile permanent error. Ran out of memory in memory space hbm. Used 49.34G of 32.00G hbm. Exceeded hbm capacity by 17.34G.

RESOURCE_EXHAUSTED: TPU TensorCore Hbm usage: 34.82G, SparseCore Hbm usage 174.10G, exceeding available bytes: 95.74G

XLA 後端：TPU

總覽

XLA 會執行檢查，確保所有必要靜態配置的總大小符合裝置的 HBM。

編譯器會管理 TPU 的固定 HBM 容量，用於幾種類型的配置：

  • 程式輸入和輸出內容：訓練批次、最佳化工具狀態等。
  • TensorCore + SparseCore Temporaries：中間計算 (例如啟用、梯度等) 需要動態記憶體。
  • 已編譯的二進位檔：TensorCore (TC) 和 SparseCore (SC) 的機器碼。
  • 系統負荷：為 XLA 執行階段保留的空間 (例如舊版 TPU 的輸入緩衝區)。
  • 常數：嵌入 HLO IR 的常數值會分配到 HBM。
  • 編譯器內部：程式層級和每個 HLO 的分配 (例如網格中節點的路由資訊)

如果 XLA 編譯器無法將上述所有分配項目納入裝置 HBM，就會發生這項錯誤。

偵錯

請仔細分析錯誤訊息和記錄，判斷下列哪一類 HBM OOM 最符合您的錯誤：

  • TensorCore (TC) + SparseCore (SC) HBM 使用量超出上限： 如果錯誤明確列出使用量，例如 「TC Hbm usage: X, SC Hbm usage Y」。→ 跳至 第 1 節。平衡 TC 和 SC HBM 用量。
  • 分配的記憶體超出預期： 如果錯誤訊息顯示「Ran out of memory in memory space HBM」(記憶體空間 HBM 記憶體不足)，請檢查記錄檔，瞭解 HBM 上分配的記憶體大小。如果出現一或多個出乎意料的大型張量 (例如超過 HBM 限制的 50%) → 請跳至第 2 節。Unexpectedly Large Allocations。
  • 匯總分配量超出 HBM 限制： 如果錯誤訊息顯示「記憶體空間 HBM 中記憶體不足」，但記錄中沒有任何異常大的張量 → 跳至第 3 節。匯總分配量超過 HBM 上限。

第 1 節：平衡 TC 和 SC HBM 用量

如果錯誤訊息明確列出用量明細，例如「TC Hbm usage: X, SC Hbm usage Y」比較這兩個值，找出瓶頸

  • SparseCore 用量過高：
    • 最佳化 HBM 堆疊用量：HBM 堆疊記憶體耗用量會隨著 feature_width、max_unique_nz_per_row 和 logical_replica_count 擴充。您可以調整 --xla_sc_num_serialized_tables_to_optimize_hbm 旗標，將資料表的處理程序序列化，藉此減少尖峰堆疊用量。但這會導致平行處理量減少。
    • 檢查填補負擔：SparseCore 會將嵌入資料表對齊 32B (8 個浮點數)。特徵寬度較小的表格 (例如< 8 個浮點數) 會產生大量填補空間，浪費 HBM。
    • 減少堆積使用量：maximum_parallel_iterations 值越高，預先擷取到 HBM 堆積的輸入資料量就越大。降低這個值可以釋出大量記憶體。
    • 驗證分片：確認嵌入資料表已在所有晶片上正確進行 mod 分片。請參閱「限制如何轉換為表格」。
    • 如需更多想法，請參閱 SC：效能和記憶體瓶頸。
  • TensorCore 使用率偏高：
  • 平衡
    • 如果兩者個別都不會過高，但總和過高，則表示晶片容量已達上限。請務必嘗試降低這兩項元件的用量。請按照這三個部分的建議操作。

第 2 節。意外的大量配置

如果記錄中出現一或多個異常大的配置 (> 50% 的 HBM 限制)，幾乎不會是硬體容量問題。這通常是設定錯誤。檢查大型配置的 XLA 標籤 (如有)，取得 JAX 原始碼的提示。

  • 移除偵錯構件：
    • 在大型執行作業中使用 jax.debug.print()，可能會強制編譯器在 HBM 中具體化完整張量，並將其傳輸至 CPU，導致融合中斷並增加尖峰記憶體用量。移除所有剩餘的 jax.debug.print()。
  • 修正網格形狀或分片效率不彰的問題：
    • 如果網格形狀有誤或缺少分片註解，編譯器可能會預設為「複製」，導致編譯器嘗試將非常大的張量放在單一晶片上
    • 檢查大型分配的形狀，並確認 XLA 正確指定及傳播分片。

第 3 節。總分配量超過 HBM 上限

如果由於分配的總和超過 HBM 限制，導致程式容量不足，通常可透過記憶體設定檔的視覺化呈現，找出造成用量高峰的特定緩衝區。如需逐步指南，瞭解如何找出尖峰記憶體貢獻者，請參閱「使用 XProf 偵錯 OOM 錯誤」。

找出一些主要貢獻者後，請按照下列步驟最佳化記憶體用量。

A. 檢查張量填補和對齊方式

張量形狀效率不彰是 TPU 發生 OOM 的常見原因，但通常不會顯示錯誤訊息。為在 TPU 上達到最高效能，XLA 會填充張量維度，通常是將最次要的維度填充為 128 的倍數，次要維度則填充為 8 的倍數。這項填補作業會影響輸入陣列和中間張量 (HLO 暫時)，可能會大幅增加記憶體用量，尤其是維度大小較小時。請參閱「陣列版面配置」。

  • 稽核大型緩衝區的形狀： (在 TPU v5 上使用預設版面配置)
    • 將滑鼠游標懸停在 Xprof Memory Viewer 的緩衝區上，會顯示緩衝區詳細資料資訊卡，其中包含緩衝區詳細資料，包括填補資訊。
    • 範例：(129, 1024) 的形狀可能會填補至 (256, 1024)，導致近 50% 的記憶體浪費。
    • 更正：(128, 1024) 形狀不需要填補，也不會浪費記憶體。
  • 對齊維度：確保所有大型張量維度 (批次大小、嵌入維度、隱藏大小) 都是 128 的倍數。

B. 調整設定

通常只要調整以下設定，就能解決 OOM 問題：

  • 縮減批次大小：中間啟用和梯度所需的記憶體與批次大小成正比。減少批次大小通常有助於減少記憶體用量。
  • 捐贈輸入緩衝區：使用 jax.jit 時，請為模型參數指定 donate_argnums。這樣 XLA 就能以輸出內容覆寫輸入記憶體。
  • 啟用混合精度 (bfloat16)：如果模型架構和品質要求允許，請對程式中最大的張量使用 bfloat16 或量化 (int8 等)。

C. 最佳化架構和分片

如果設定變更不足，模型拓撲可能過大，不適合目前的硬體設定。

  • 使用新一代 TPU：新一代 TPU 通常每個晶片可提供更多 HBM；請改用新一代 TPU (如有)。
  • 在較大的晶片拓撲上執行：如果模型權重對於現有拓撲而言過大，您可以嘗試將權重分散到更多晶片。
  • 導入進階分片技術：
    • 探索更進階的資料、張量或管道平行處理方法。
    • 為中間值和輸出內容指定分片提示。
  • 使用 JAX 主機卸載功能：將大型張量卸載至主機 CPU 記憶體。例如啟用卸載和最佳化工具狀態卸載。

D. 調整影響記憶體的重要 XLA 旗標：

您可以調整重要記憶體旗標，在效能與記憶體用量之間取得平衡。但這些方法可能會對效能造成負面影響，因此請做為最後手段。

E. 調整 XLA 重新實體化傳遞 / 手動檢查點

如果模型即將超出記憶體容量，您可以強制 XLA::Rematerialization 傳遞優先節省記憶體，但可能會導致編譯速度變慢：

檢舉 說明 影響 / 取捨
--xla_tpu_max_hbm_size_mib 手動設定 Rematerialization 傳遞作業使用的 HBM 大小上限。 強制編譯器更努力地將程式納入小於實際實體 HBM 的限制。
--xla_tpu_rematerialization_algo=PEAK_PRIORITY 將工作重點放在記憶體用量尖峰時段。 與預設演算法相比，可更有效率地大幅減少記憶體用量。
--xla_tpu_rematerialization_max_block_size_limit=32 控制區塊中可一次重新實現的指令數量上限。 增加這個值可節省記憶體，但編譯時間會大幅增加。
--xla_tpu_rematerialization_block_effort_factor=10.0 定義搜尋要重新實現的區塊所花費的精力 (編譯時間)。 值越高，搜尋記憶體節省空間的範圍就越廣，但編譯時間會增加。
--xla_tpu_pre_fusion_remat=true 在融合階段之前啟用額外的重新具體化階段。 可節省更多記憶體，但會增加編譯時間，且可能影響數值穩定性。

或者，您也可以搭配 jax.grad 使用 jax.checkpoint 修飾符，手動控管在正向傳遞時儲存哪些中介值，以及在反向傳遞時重新計算哪些中介值，藉此以運算週期換取 HBM。

F. 使用進階剖析工具

使用 XProf 偵錯記憶體不足錯誤教學課程，說明如何使用 XProf 記憶體檢視器，以視覺化方式呈現編譯器對 HBM 使用情形的看法。

這項工具可讓您查看尖峰記憶體配置和緩衝區生命週期，有助於瞭解尖峰使用率時，究竟是哪些項目耗用 HBM。如需一般剖析設定，請參閱「開始使用 Xprof」和「TensorBoard 剖析」。