調查 JAX 追蹤
概述
對單一擷取進行七步驟調查：結論、資料核對、泳道圖、核心 → HLO → Python 程式行、屋頂線圖、應改變之處，以及信任清單。
一切皆來自使用者提供的擷取：<host>.xplane.pb 和/或 <host>.trace.json.gz。請勿重新執行工作負載、傾印 XLA 或執行其他分析器。若追蹤無法回答問題，則將其列入「不可信任」清單。
目標
報告完成後，讀者將了解：
- (a) 應改變什麼以加速該步驟，以及其假設下的估計值；
- (b) 此追蹤中無法支援歸因的部分，並附上計數。
工具
腳本路徑相對於此技能的基本目錄，即技能載入時顯示的目錄。請在暫存目錄中執行：兩個腳本都會將其輸出寫入當前目錄。
B=<此技能的基本目錄>
python3 -P $B/../analyzing-jax-launch-overhead/scripts/skip_stats.py <擷取>/<host>.trace.json.gz
python3 -P $B/scripts/attribute_kernels.py <擷取目錄> [--program ID] [--user-re REGEX]
skip_stats.py（僅需標準函式庫）：結論、每個串流的時間預算、閒置期間的主機事件，以及 SKIP 論文指標。閱讀方式請參閱 analyzing-jax-launch-overhead 技能。attribute_kernels.py（需要jax和xprof套件）：每個 GPU 核心 → HLO 指令 → op_name → Python 程式行。它只讀取.xplane.pb，並先將其複製到./attribution/xplane/，因為 XProf 會在其旁邊寫入 HLO 原型。它會將kernels.csv、ops.csv、lines.csv、how.csv、summary.txt以及該程式的 HLO 文字寫入./attribution/。- 精確（exact）： 追蹤器本身的
hlo_op，適用於每次呼叫啟動一個核心的情況。 - 名稱（name）： 一個圖形核心，其名稱是其時間視窗內唯一的候選者。
- 順序（order）： 共享核心、程式庫 GEMM、複製與集體操作，按排程順序匹配。
- 無（none）： 沒有剩餘的候選者。
- 時間視窗是程式的排程中（while 迴圈已依據已知迭代次數展開）相鄰精確核心之間的指令。
- 驗證： 在一個 MiniGPT 擷取中，它為所有 147,840 個核心選擇了與某個管線相同的指令，而該管線在關閉指令緩衝區的相同程式執行中具有 100% 的一致性。
- 精確（exact）： 追蹤器本身的
五項技術，各附工具
| 領域 | 技術 | 此處工具 |
|---|---|---|
| 效能工程 | 關鍵路徑；屋頂線圖 | skip_stats.py 提供的串流預算；僅含正確峰值的屋頂線圖（參見陷阱） |
| 分散式追蹤 | 透過關聯 ID 連接主機與 GPU 事件；找出落後者 | skip_stats.py → rows.csv；各 GPU 步驟結束時間的差異 |
| 編譯器 | 透過編譯器傳遞的原始碼追溯 | attribute_kernels.py → lines.csv |
| 法務會計 | 總和必須與獨立記錄相符 | 步驟 2 中的檢查 |
| 資訊設計 | 對齊的泳道圖，每個 GPU 串流一條 | 一個步驟，所有 GPU × 串流，共用一個時間軸 |
步驟
計劃包含以下七個步驟，全部針對給定的檔案執行。只有新的擷取才能回答的問題（例如主機呼叫內部的行為、修復的效果、批次掃描）則歸入步驟 7 的「需要新擷取」清單，是否要擷取由使用者決定。
- 結論與主機端原因： 執行
skip_stats.py。 - 資料核對： 每個檢查若成立則保留，否則列入信任清單：
- 每個裝置的 GPU 事件：json = xplane；
- 每個 GPU 事件都與一個主機啟動呼叫相關聯；
- 每個步驟的核心數是固定的；
- 所有 GPU 執行相同的工作（每個步驟的核心計數相等，或可解釋其差異）；
- 對於每個串流，計算 + 集體操作 + memcpy + 閒置時間佔比總和等於該時間跨度。
- 泳道圖： 對於中位數步驟，將所有 GPU × 串流放在一個時間軸上，並加上啟動執行緒的主機事件。這顯示了關鍵路徑與落後者：哪個 GPU 最後完成，以及該 GPU 在等待什麼。
- 核心 → HLO → Python 程式行： 執行
attribute_kernels.py。- 報告每個匹配類別佔用的 GPU 時間比例，以及沒有候選者的核心。
- 使用
lines.csv和ops.csv來製作每行程式行和每個操作的表。 --user-re選擇哪些框架計入使用者程式碼。預設為 site-packages 和 jax 之外的第一個框架。
- 屋頂線圖， 僅針對可從其 HLO 形狀獲得 FLOPs 和位元組數的操作：使用
attribute_kernels.py寫入./attribution/的<程式>(<id>).hlo.txt。使用該操作資料類型的峰值，引用自供應商規格，並標記為讀取；追蹤無法量測峰值。 - 應改變之處： 根據各原因在關鍵路徑中的佔比進行排序。針對每個原因，估計如果它消失後的步驟時間，例如步驟時間 − 關鍵串流上的主機閒置時間。將其標記為推斷，未執行，並命名一個調整參數作為假設供使用者測試（例如
donate_argnums、關閉指令緩衝區、使用更少或更大的集體操作）。 - 此追蹤不可信任之處： 列出每個項目並附上計數或佔比。然後加入「需要新擷取」清單：每個待解決問題，以及如何擷取才能回答該問題。
陷阱（在 ROCm MI355X 擷取上量測）
- XProf 的操作檢視會遺漏圖形核心。
hlo_stats在一個 4-GPU 擷取中僅涵蓋了 37.6% 的 GPU 核心時間。另外 62.4%（指令緩衝區核心）在kernel_stats中顯示為操作 "XlaModule"。在引用hlo_stats、op_profile或roofline_model之前，請先檢查覆蓋率。 - XProf 對於 gfx950 的屋頂線圖峰值 為 2,517 TFLOP/s 和 3,624 GiB/s，與資料類型無關。FP32 GEMM 需要 FP32 峰值。
- RCCL/NCCL 核心時間包含等待同儕的時間。 這不是傳輸時間。
- 非同步集體操作： 它們的
hlo_op（例如all-gather.34）不在擷取中儲存的 HLO 裡。attribute_kernels.py會保留該標籤，並根據順序從匹配的 collective-start 指令中取得 Python 程式行。 - ROCm memcpy 行 可能顯示
kind:Unknown size:0。 - Python 追蹤器 會誇大主機時間。
- 一個程式可能出現兩次（兩個不同的
.hlo_proto.pb檔案，具有不同的 ID）。attribute_kernels.py使用在 xplane 中佔用最多 GPU 時間的程式 ID；json 沒有program_id。
報告
- Markdown： 結果優先。
- 數字： 每個數字都附帶其在所有步驟和 GPU 中的範圍，並標記為量測、讀取或推斷。
- 章節： 結論與原因、資料核對、每個 GPU/串流表格、熱門操作與 Python 程式行、應改變之處（估計值），以及信任清單。
- 事實查核： 在發送前，請讓一個新的子代理根據表格檢查每一項陳述。
- HTML 頁面： 如果使用者需要，writing-trace-report-pages 技能可以根據此報告及其表格建構頁面。