openxla/xprof

XProf – A Scalable Profiler for the Modern ML Stack

What it is – XProf is an open‑source profiling system (with a TensorBoard plugin) that lets you collect, analyse and visualise performance data from machine‑learning workloads running on CPUs, GPUs, TPUs or other accelerators. It works with JAX, TensorFlow, PyTorch/XLA and other frameworks that emit XLA‑style profiling traces.

Key capabilities

• Overview page – high‑level summary of step times, environment details and a graph of per‑step latency.
• Trace viewer – timeline showing the duration of each operation, whether it ran on host or device, and inter‑device communication.
• Memory profile – live view of memory consumption per device.
• Graph viewer – visualisation of the HLO (High‑Level Optimiser) graph of the model.
• Distributed profiling – an aggregator can split heavy‑weight processing across multiple worker nodes via gRPC.
• TensorBoard integration – the UI appears under the #profile tab when you run TensorBoard with a compatible log directory.

Installation

# stable release
pip install xprof
# optional TensorBoard integration
pip install xprof tensorboard

For Python 3.12+ you may need an older setuptools (pip install "setuptools<70").

Running the profiler

• Standalone server

  xprof --logdir=path/to/profile_data --port=6006
  # or explicitly
  xprof server -l path/to/profile_data -p 6006
  

• Via TensorBoard

  tensorboard --logdir=path/to/profile_data
  

  Then open http://localhost:6006/#profile.

How to feed data XProf expects .xplane.pb files (XLA profiling traces) under:

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

Multiple sessions appear as separate entries in the UI. You can also load data on‑the‑fly using URL parameters:

• session_path – points to a single session directory.
• run_path – points to a directory containing many session sub‑folders.

Distributed mode

1. Start one or more worker nodes:

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

2. Start an aggregator node that knows the workers:

   xprof server \
     --worker_service_address=worker1:50051,worker2:50051 \
     --port=6006 --logdir=profiler/demo
   

The aggregator UI distributes heavy processing (overview, op stats, input‑pipeline, pod viewer) across the workers.

Resources & docs

• Full documentation: https://openxla.org/xprof
• JAX profiling guide, PyTorch/XLA guide, TensorFlow profiling guide, Cloud TPU profiling guide (all linked in the README).
• A Colab demo for newcomers: https://docs.jaxstack.ai/en/latest/JAX_for_LLM_pretraining.html

Citation If you use XProf in a research project, cite the MLSys 2026 paper:

@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 is a genuine, production‑grade tool for profiling ML workloads, not a tutorial collection or a trivial demo.