Performance Profiling (XProf)

Performance Profiling (XProf)#

A job that runs on a TPU or GPU isn’t necessarily using it well. Profiling shows where the accelerator time actually goes. This guide covers XProf — the profiler for XLA workloads — and how to capture a trace from a Kinetic job and view it on your own machine.

Because a Kinetic pod scales to zero the moment your job finishes, the workflow is: capture inside the function, write the trace to KINETIC_OUTPUT_DIR (durable GCS), then view it locally.

What XProf is#

XProf is the open-source accelerator profiler from OpenXLA (formerly the TensorBoard “profile” plugin). It runs standalone or as a TensorBoard tab, reads hardware counters rather than wall-clock timers, and adds little overhead during the capture window.

It is an XLA profiler, and that decides which tool you capture with:

Backend / workload

Capture with

View in

Keras-on-JAX (Kinetic default), native JAX

jax.profiler

XProf

Keras-on-TensorFlow

TensorBoard(profile_batch=…)

XProf

PyTorch/XLA (torch_xla)

torch_xla.debug.profiler

XProf

Native PyTorch (eager CUDA)

torch.profiler

Perfetto

Native eager PyTorch doesn’t compile through XLA, so it uses torch.profiler and views in Perfetto — not XProf.

Capture a profile#

Capture a handful of steps after a warm-up, inside the decorated function, and write to $KINETIC_OUTPUT_DIR/profile so the trace survives the pod (see Checkpointing). The job needs no extra packages — jax.profiler ships with JAX, and xprof is a local viewer, not a job dependency. If a profiling job does need an extra package, add it to a requirements.txt in that script’s own directory (see Dependencies).

This is a deliberately minimal demo — a tiny JAX training loop — but the same capture pattern drops into any job (real training, KerasHub fine-tuning, vLLM serving, multi-host runs):

examples/jax_profiling_demo.py#
import kinetic


@kinetic.run(accelerator="tpu-v5litepod-1x1")
def train_and_profile():
  """Capture an XProf trace of a few JAX training steps, saved to GCS."""
  import os

  import jax
  import jax.numpy as jnp

  trace_dir = os.path.join(
    os.environ.get("KINETIC_OUTPUT_DIR", "/tmp/kinetic-out"), "profile"
  )

  # A small MLP, params held as a pytree.
  key = jax.random.PRNGKey(0)
  k1, k2, k3 = jax.random.split(key, 3)
  params = {
    "w1": jax.random.normal(k1, (1024, 2048)) * 0.02,
    "w2": jax.random.normal(k2, (2048, 2048)) * 0.02,
    "w3": jax.random.normal(k3, (2048, 10)) * 0.02,
  }
  x = jax.random.normal(key, (4096, 1024))
  y = jax.random.normal(key, (4096, 10))

  def loss_fn(p, x, y):
    h = jnp.tanh(x @ p["w1"])
    h = jnp.tanh(h @ p["w2"])
    pred = h @ p["w3"]
    return jnp.mean((pred - y) ** 2)

  @jax.jit
  def update(p, x, y, lr=1e-3):
    grads = jax.grad(loss_fn)(p, x, y)
    return {k: p[k] - lr * grads[k] for k in p}

  # Warm up once so XLA compilation doesn't pollute the trace.
  params = update(params, x, y)
  jax.block_until_ready(params)

  # The context manager flushes the trace even if a step raises.
  with jax.profiler.trace(trace_dir):
    for _ in range(10):
      params = update(params, x, y)
    # force the async work to land before the trace closes
    jax.block_until_ready(params)

  print(f"final loss: {float(loss_fn(params, x, y)):.4f}")
  print(f"trace written to: {trace_dir}")
  return trace_dir


if __name__ == "__main__":
  train_and_profile()

For Keras-on-JAX, it’s the same idea: wrap a short model.fit(...) in with jax.profiler.trace(trace_dir): after a warm-up epoch.

Note

JAX dispatches asynchronously, so always block_until_ready() inside the trace region — otherwise the profiler can close before the device work lands. Keep the window to a few steps; traces grow fast.

Other backends: native PyTorch uses torch.profiler with tensorboard_trace_handler(trace_dir) (view in Perfetto); PyTorch/XLA uses torch_xla.debug.profiler, which produces XProf-readable traces.

View the trace#

Install the viewer and point it at the trace path your job printed:

pip install xprof gcsfs        # gcsfs lets XProf read gs:// directly
xprof --logdir gs://<project>-kn-<cluster>-jobs/outputs/<job_id>/profile --port 6006
# --logdir is the directory that contains plugins/ (the path your job printed)
# then open http://localhost:6006

Or copy it down first — gcloud storage cp -r <gs-path> ./trace and --logdir ./trace. In the UI, use the tool dropdown: Trace Viewer for the step timeline, Overview Page for the summary. (The Capture Profile button does live, on-demand capture against a running profiler server, so it isn’t used here — the trace was already captured inside the job.)

What the tools show#

  • Overview Page — top-level summary; whether you’re host- or device-bound. Start here.

  • Trace Viewer — per-event timeline across host / TPU / GPU; where you spot gaps and stalls.

  • Roofline — memory-bound vs. compute-bound, which decides your optimization strategy.

  • Framework / HLO Op Stats — cost by framework op and by compiled HLO op.

  • Memory Viewer / Profile — usage over time and at peak; first stop after an OOM.

  • Megascale Stats — cross-slice (DCN) communication on multi-host Pathways runs.