從擷取檔案除錯 JAX OOM
概述
JIT 編譯的 JAX 步驟會將其所有暫存記憶體規劃為一個暫存區塊。因此，該步驟的記憶體使用就是 XLA 的計劃，而 .xplane.pb 檔案同時包含了該計劃與分配器的執行時期事件。scripts/oom_report.py 僅使用標準函式庫，便能於 2–5 秒內讀取這兩部分資訊：
- 計劃： 每個模組的緩衝區分配，來自
/host:metadata。它會給出參數（以及有多少被捐贈）、非別名輸出、暫存區塊、以及堆疊峰值。對於峰值時存活的緩衝區，它會提供大小、形狀、角色以及 Python 程式碼行號。角色來自 JAX 的名稱堆疊：jvp(代表前向（為反向傳播而保留），transpose(代表反向，s8自訂呼叫的輸出是函式庫的工作空間。 - 執行時期： BFC 分配器的事件。它們會給出容量、每次執行開始時的已用位元組數、峰值、
AllocatorRetry等待時間、從未啟動 kernel 的執行、以及最後兩次執行之間的主機事件。它們還會給出最後一次執行開始時仍存活的內容，並按分配該內容的主機事件分組：outputs of jit_step execution 3、BufferFromHostBuffer(s32[…])（一個device_put），或「在擷取開始前就已分配」。
報告會以測量值或推估值標記的診斷行結尾。請先執行以下指令：
python3 -P <此技能基底目錄>/scripts/oom_report.py <擷取目錄或 .xplane.pb> [--module 'jit_train_step(123)'] [--top 25] [--json out.json]
python3 -P <此技能基底目錄>/scripts/oom_report.py --self-test
預設情況下，它會分解最後一次執行從未啟動 kernel 的模組。如果沒有此類模組，則會選取執行了的最大模組。使用 --module 可以選取其他模組，例如評估步驟。
從診斷行到原因與修復
| 診斷行 | 原因 | 要建議的修復；節省量來自報告中的位元組數 |
|---|---|---|
the step itself doesn't fit | 僅計劃就超過了記憶體池容量 | 依照峰值處的角色分解（如下） |
↳ 大部分為 forward, kept for the backward pass | activations：它們隨 batch × seq 增長；樸素注意力每層增加 B·H·S·S | jax.checkpoint / remat（節省大部分 forward-kept 的位元組，代價約為一個額外的前向傳遞）；使用更小的 micro-batch 搭配梯度累積；更短的序列長度（S 減半時 S² 緩衝區縮小 4 倍）。僅在「確認 B·H·S·S 緩衝區離開計劃」時才提供 fused attention |
↳ 大部分為 backward | 梯度及其暫存變數 | 使用 FSDP 分片或累積；檢查 bf16 梯度是否有 f32 副本 |
↳ workspace | GEMM / 卷積函式庫的暫存空間 | 通常很小；在 ROCm 上，gfx950 每個 hipBLASLt GEMM 為 64 MiB |
grow by X MiB per execution, = 1.00x the non-aliased outputs | 程式保留了每個步驟的輸出（跨步驟的記憶體洩漏）：存活清單顯示 outputs of … execution 1, 2, … | 在執行之間的主機事件中找到持有者（例如 $<unknown> append = list.append）。然後刪除引用、在步驟內部減少（一個 norm，而不是整個樹），或使用 jax.device_get。更大的記憶體池只會延遲 OOM |
no parameter is donated … can save up to X | 舊的和新的參數 / 優化器狀態同時存活 | 對狀態引數使用 donate_argnums |
other live arrays take the room; the largest … came from X | 步驟外部的陣列：放在 GPU 上的資料集或評估副本、預先提取、其他 jit 的輸出 | 將其保留在主機上並以串流方式傳輸、釋放它，或縮小其 dtype。「在擷取開始前就已分配」表示需要更早進行 profiling 以識別它 |
AllocatorRetry … longest N s, launched none | OOM 發生在擷取期間 | 僅為證據。失敗的請求通常是該模組的暫存區塊或某個輸出 |
XLA already rematerialized N instructions | XLA 自己的 rematerialization 已經在極限下執行過 | 預期修復來自程式本身，而非來自 XLA |
對於每項修復，請給出其估計的節省量與成本，並標記為推估值。檢查估計值是否超出容量。只有當 GPU 確實有空間時，「調高 XLA_PYTHON_CLIENT_MEM_FRACTION」才能修復計劃無法容納的問題。容量遠低於 HBM 暗示該比例已被設定；擷取檔案不會記錄環境變數。
以 JAX 術語解釋 Playbook 的記憶體結構
Ultra-Scale Playbook（「對記憶體使用進行 Profiling」）將訓練記憶體分為權重、梯度、優化器狀態和 activation。在報告中：
- 權重 + 優化器狀態 = 步驟的參數。Adam 在 fp32 下每個參數佔 12 B（p, m, v）。被捐贈意味著沒有第二份副本。
- 梯度 = 峰值處的
backward角色，如果步驟回傳梯度則還包括輸出。 - Activations =
forward, kept for the backward pass。Scan /nn.scan模型會將它們作為堆疊的[layers, …]緩衝區保留在while迴圈中。 - 「步驟 1 通過，步驟 N 失敗」： 在 PyTorch 中，這是因為第一步創建了優化器狀態。在 JAX 中，狀態在步驟 1 之前就已存在，因此請改為查看增長行。第一次呼叫還會編譯和自動調優，因此第一次執行之前的峰值是編譯時間。
規則
- 僅使用擷取檔案。 不要為了測試修復而重建模型或重新執行工作：根據計劃的位元組數進行估計，並說明如何在使用者端確認（
jit(...).lower(...).compile().memory_analysis()）。 - 不要寫入擷取目錄。 XProf 的轉換器（
xspace_to_tool_names、memory_viewer、xprofCLI）會在.xplane.pb旁邊寫入*.hlo_proto.pb快取檔案。如果您需要 XProf 的檢視，請先將.xplane.pb複製到暫存目錄。 - 為每個數字標記 為測量值、讀取值或推估值。
測量參考（單一 MI355X，小記憶體池，2026-11-06）
| 擷取 | 報告顯示的內容 |
|---|---|
| attention 計劃過大 | 需要 3.72 GiB > 2.79 GiB；堆疊峰值 81.6% forward-kept；f32[8,8,1024,1024] ×5 位於 softmax 行；暫存區塊 3,495 MiB = 錯誤訊息中的「3.41GiB」；AllocatorRetry 10 秒；89 個 remat 指令 |
| 每個步驟都保留梯度 | 步驟開始時使用量 316 → 1,252 MiB，每個步驟 +104.0 MiB = 1.00× 非別名輸出；計劃 788 MiB 可容納；步驟之間有 $<unknown> append |
| 沒有捐贈 | 需要 3.88 GiB > 2.79 GiB；1,902 MiB 參數 + 1,902 MiB 輸出，0 被捐贈；暫存區塊僅 172 MiB |
| MaxText llama2-7b，8 GPU，FSDP（可容納） | 暫存區塊 61.9 GiB；堆疊峰值 57% forward-kept，scan 堆疊 bf16[32,4,4096,11008] ×2 透過 while.50 |
限制
- 每個檔案僅限一個主機： 多主機執行會為每個主機產生一個
.xplane.pb，因此請為每個主機都執行一次。 - 僅限 HBM： 只有記憶體空間 0 會被分解。主機卸載顯示為
memory space 5。 - 連續性：「可容納」是比較位元組數。BFC 還需要連續的區塊，因此接近容量邊界時仍可能因碎片化而失敗。
- 「其他」角色： 由 XLA 創建的運算（無 JAX 名稱堆疊）且無別名的緩衝區。
- CUDA： 未經測試。它使用相同的 BFC 和 XLA 事件，但名稱尚未驗證。
- TPU / eager： TPU 擷取僅提供計劃總計，沒有堆疊峰值和執行資訊。Eager（逐運算）程式顯示為許多小模組，請在該處閱讀執行時期部分。