openxla/xprof

XProf – 現代MLスタック向けスケーラブルプロファイラ

何であるか – XProfは、オープンソースのプロファイリングシステム（TensorBoardプラグイン付き）であり、CPU、GPU、TPU、またはその他のアクセラレータ上で実行される機械学習ワークロードから、パフォーマンスデータを収集・分析・可視化できる。JAX、TensorFlow、PyTorch/XLA、およびXLAスタイルのプロファイリングトレースを出力するその他のフレームワークと互換性がある。

主な機能

概要ページ – ステップ時間、環境情報の高レベルな要約、およびステップごとのレイテンシのグラフ。
トレースビューア – 各操作の実行時間、ホストまたはデバイスで実行されたか、およびデバイス間通信を示すタイムライン。
メモリプロファイル – 各デバイスのメモリ使用量のライブ表示。
グラフビューア – モデルのHLO（High-Level Optimiser）グラフの可視化。
分散プロファイリング – 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 – 複数のセッションサブフォルダを含むディレクトリを指す。

分散モード

1つ以上のワーカーノードを起動する：

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スタック向けスケーラブルプロファイラ

関連