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 栈的可扩展分析器

相关