openxla/xprof

XProf – 現代 ML 堆疊用的可擴展分析工具

是什麼 – XProf 是一個開源的分析系統（附帶 TensorBoard 插件），可讓您收集、分析並可視化在 CPU、GPU、TPU 或其他加速器上執行的機器學習工作負載的效能資料。它支援 JAX、TensorFlow、PyTorch/XLA 及其他產生 XLA 風格分析追蹤的框架。

主要功能

概觀頁面 – 步驟時間、環境資訊的高階摘要，以及每步延遲的圖表。
追蹤檢視器 – 顯示每個運算的持續時間、是在主機還是裝置上執行，以及裝置間通訊的時間軸。
記憶體分析 – 各裝置記憶體使用量的即時視圖。
圖形檢視器 – 模型 HLO（高階優化器）圖的可視化。
分散式分析 – 透過 gRPC，聚合器可將重負載處理分散至多個工作節點。
TensorBoard 整合 – 當您使用相容的日誌目錄執行 TensorBoard 時，UI 會出現在 #profile 選項卡下。

安裝

# 穩定版本
pip install xprof
# 可選：TensorBoard 整合
pip install xprof tensorboard

對於 Python 3.12+，您可能需要較舊版本的 setuptools（pip install "setuptools<70"）

執行分析器

獨立伺服器

xprof --logdir=path/to/profile_data --port=6006
# 或明確指定
xprof server -l path/to/profile_data -p 6006

透過 TensorBoard
```
tensorboard --logdir=path/to/profile_data
```
然後開啟 http://localhost:6006/#profile。

如何提供資料 XProf 期望在以下路徑下存在 .xplane.pb 檔案（XLA 分析追蹤）：

<log_dir>/plugins/profile/<session_name>/host*.xplane.pb

多個會話會在 UI 中顯示為獨立項目。您也可以透過 URL 參數即時載入資料：

session_path – 指向單一會話目錄。
run_path – 指向包含多個會話子資料夾的目錄。

分散式模式

啟動一個或多個 工作節點：

xprof server --grpc_port=50051 --port=9999 --hide_capture_profile_button

啟動一個知道工作節點的 聚合節點：
```
xprof server \
  --worker_service_address=worker1:50051,worker2:50051 \
  --port=6006 --logdir=profiler/demo
```
聚合節點 UI 會將重負載處理（概觀、運算統計、輸入管道、Pod 檢視器）分散至各工作節點。

資源與文件

完整文件： https://openxla.org/xprof
JAX 分析指南、PyTorch/XLA 指南、TensorFlow 分析指南、Cloud TPU 分析指南（均連結於 README）。
新手的 Colab 示範： https://docs.jaxstack.ai/en/latest/JAX_for_LLM_pretraining.html

引用如果您在研究專案中使用 XProf，請引用 MLSys 2026 論文：

@inproceedings{1076558,
  title     = {XProf: An Open, Scalable and Extensible Profiling System for the Modern ML Stack},
  author    = {Robert Hundt and Naveen Kumar and Jose Baiocchi Paredes and Scott Goodson and Clive Verghese and Prasanna Rengasamy and Kelvin Le and Jiya Zhang and Charles Alaras and Yin Zhang and Kan Cai and Jiten Thakkar and Sai Ganesh Bandiatmakuri and Yogesh SY and Ani Udipi and Vikas Aggarwal},
  year      = {2026},
  booktitle = {Ninth Conference on Machine Learning and Systems}
}

XProf 是一個真正的、生產級的 ML 工作負載分析工具，而非教學範例集合或簡單示範。

openxla/xprof

XProf – 現代 ML 堆疊用的可擴展分析工具

相關