錯誤代碼：0100

類別：執行階段：緩衝區分配失敗

這項錯誤表示 XLA：TPU 執行階段的記憶體分配器無法在加速器的 HBM 上，為要求分配作業找到合適的記憶體區塊。

錯誤訊息示例：

ValueError: RESOURCE_EXHAUSTED: Error allocating device buffer: Attempting to allocate 8.00M. That was not possible. There are 6.43M free.; (0x0x1_HBM0)

XLA 後端：TPU

總覽

這個錯誤會在下列情況擲回：

  • 使用者透過 jax.device_put 或
  • 使用者排定程式的輸出分配失敗。

這類失敗通常有以下幾個原因：

  • 記憶體不足 (OOM)：使用者嘗試分配的記憶體區塊大於 TPU HBM 可用的記憶體總量。
  • 記憶體片段化：由於記憶體空間中沒有足夠大的單一連續可用區塊，因此分配作業失敗。可用記憶體總量足以分配，但分散在記憶體空間中，以不連續的小區塊形式存在。

TPU 執行階段有多種機制可重試分配失敗作業，包括：

  • 如果佇列中有待取消分配的項目，執行階段會重試失敗的分配作業，
  • 如果 OOM 是由片段化所導致，執行階段可以自動觸發重組作業並重試。
  • TPU 執行階段會優先分配緩衝區，而非讓程式保持載入狀態。如果 HBM 不足而導致緩衝區分配失敗，系統會逐出已載入的 TPU 程式，直到緩衝區有足夠的記憶體為止。

因此，在採取上述緩解措施後，如果仍發生錯誤，通常需要使用者採取行動。

偵錯

  • 減少模型的記憶體用量：
    • 縮減批次大小：直接縮減批次大小可降低記憶體用量。
    • 參數分片：對於非常大型的模型，請使用模型平行處理或分片等技術，將參數分配到多個 TPU 核心或主機的 HBM。
    • 縮短序列/內容長度：對於處理序列的模型 (例如語言模型)，減少輸入序列長度可大幅減少記憶體用量。
    • 緩衝區捐贈：利用架構功能 (例如：jax.jit(..., donate_argnums=...)) 向 XLA 發出信號，表示特定輸入緩衝區可以覆寫並重複用於輸出。
    • 最佳化檢查點策略：不要一次儲存整個模型狀態，考慮只儲存模型權重，或使用分片檢查點策略。
  • 位址記憶體配置和填補：
    • TPU 記憶體是以區塊形式分配，而填充可能會增加張量的實際大小。
  • 確保沒有記憶體流失：
    • 請確保對 jax.Array 物件的參照不會保留超過預期時間。即使程式編譯完成，保留 jax.Array 物件也可能會導致系統無法自動取消分配。

工具

啟用 tpu_log_allocations_on_oom 標記，當發生 OOM 時，分配器會傾印所有目前分配的詳細報表，這對偵錯來說非常寶貴。