類別：編譯時間：範圍內 Vmem OOM
這項錯誤表示程式需要的範圍向量記憶體 (Vmem) 超出分配到的數量。
錯誤訊息示例：
RESOURCE_EXHAUSTED: Ran out of memory in memory space vmem while allocating on stack for %my-custom-kernel = bf16[2048,4096]{1,0:T(8,128)(2,1)} custom-call(...) ...
XLA 後端：TPU
總覽
TPU 具有向量記憶體 (VMEM)，這是 TensorCore (TC) 專用的本機暫存記憶體。編譯器會管理不同類型的 Vmem 分配：
- 指令範圍的分配：執行單一 HLO 指令時，Vmem 中的暫時儲存空間。這包括運算元跨度緩衝區 (例如用於雙重緩衝) 和暫存器溢出。
- 以程式為範圍的分配：超出單一 HLO 指令範圍的分配。這些通常是 HLO 暫時項目和中間結果，是 HLO 指令的輸入和/或輸出內容。
如果指令範圍的分配量超過該指令的分配量上限，就會發生編譯時間範圍 Vmem OOM。這項限制受到控管
- 透過 --xla_tpu_scoped_vmem_limit_kib 標記，為整個程式全域設定
和
- 每個自訂核心透過 vmem_limit_bytes 參數。
這類錯誤通常是由內部編譯器錯誤或超出分配限制的自訂核心所造成。
偵錯
請仔細分析錯誤訊息，判斷錯誤是源自於自訂核心還是標準 HLO。自訂核心造成的錯誤應具有下列簽章：
Ran out of memory in memory space vmem while allocating on stack for %my-custom-call = <output-shape> custom-call(<params>), custom_call_target="tpu_custom_call" ...
- 自訂核心範圍的 Vmem OOM：如果錯誤指向自訂核心 → 請跳至「重新調整核心」。
- 非核心 Vmem 問題：如果 Vmem OOM 是由非自訂核心運算所致，很可能是內部編譯器錯誤。請在 XLA 上提交錯誤報告，並附上 HLO 傾印。
重新調整核心
如果錯誤源自於自訂核心，請使用下列技術來減少核心的記憶體需求：
- 調整區塊大小：在核心設定中縮減區塊大小 (圖塊大小)，以降低範圍 Vmem 用量。
- 設定每個核心範圍的 Vmem 限制：使用 vmem_limit_bytes 參數，明確要求該特定核心所需的記憶體量
- 修改記憶體著色：使用 pallas.tpu.with_memory_space_constraint，明確將核心的輸入/輸出內容著色/限制為 Vmem。 請注意，不要為 Vmem 著色太多輸入/輸出內容，否則可能會導致整體 Vmem OOM。
- 調整 Vmem 限制：如果難以進行核心專屬的重新調整，或問題影響多個核心，您可以使用 --xla_tpu_scoped_vmem_limit_kib 標記調整全域 Vmem 限制。