Zero-infrastructure managed XProf: Profiling ML workloads on Cloud TPU with ML Diagnostics

Google Cloud Community Articles
, , ,
1

TL;DR: Capture XProf performance profiles and monitor real-time ML training metrics on Cloud TPU VMs with zero infrastructure setup. Skip self-hosted profiler servers, stream hardware telemetry directly to Google Cloud, and analyze multi-host traces up to 10x faster in the managed XProf viewer at zero compute charge.

Waiting for a JAX training loop to finish on a Cloud TPU slice while step times spike or accelerator utilization drops is frustrating. Debugging these bottlenecks often pulls ML engineers away from modeling and into operations: setting up self-hosted TensorBoard and XProf servers, configuring SSH port-forwarding, and manually syncing profile traces across storage buckets.

The Google Cloud ML Diagnostics SDK (part of the ML Diagnostics platform) eliminates this overhead. It delivers a fully managed XProf profiling and observability experience directly in the Google Cloud console. Best of all, there is no compute charge for the managed XProf backend deployed by ML Diagnostics—you only pay for Cloud Logging metric ingestion and Cloud Storage trace storage.

In this 5-minute walkthrough, you will provision a Cloud TPU VM, instrument a JAX training script to record metrics and capture XProf profiles, and inspect your hardware telemetry in the Google Cloud console.

Prerequisites

Configure the required APIs, IAM permissions, and Cloud Storage bucket in your Google Cloud project so your workload can stream telemetry without permission errors.

Required APIs and IAM roles

Type Resource Purpose
API hypercomputecluster.googleapis.com Cluster Director API for MLRun registration, telemetry ingestion, and UI visualization
API tpu.googleapis.com Cloud TPU API for provisioning and managing Cloud TPU VMs
API compute.googleapis.com Compute Engine API for VM networking and SSH metadata
IAM role roles/hypercomputecluster.editor Creates and manages MLRun resources and grants UI access (displayed as Cluster Director Editor in the Google Cloud console)
IAM role roles/logging.logWriter Writes workload configurations and metrics to Cloud Logging
IAM role roles/storage.objectUser Uploads and reads XProf profile traces in Cloud Storage

1. Enable required APIs

Enable the required services in your Google Cloud project:

gcloud services enable hypercomputecluster.googleapis.com tpu.googleapis.com compute.googleapis.com

:memo: Note: The hypercomputecluster.googleapis.com API powers MLRun registration, telemetry ingestion, and the console UI. You do not need the Cluster Director product to manage your clusters—ML Diagnostics works independently with standalone Cloud TPU VMs and Google Kubernetes Engine (GKE). See the official overview for details.

API enablement in the console

API enablement in the Google Cloud console
API enablement in the Google Cloud console2022×1406 285 KB

2. Enable Log Analytics on Cloud Logging

The Run Diagnostics dashboard uses Log Analytics to render interactive time-series charts for model and system metrics. Upgrade your _Default log bucket before running your workload (Log Analytics does not backfill historical logs):

gcloud logging buckets update _Default --location=global --enable-analytics --project=YOUR_PROJECT_ID

3. Set up a service account and Cloud Storage bucket

Create a dedicated service account for your Cloud TPU VM and a Cloud Storage bucket to store XProf trace files:

  1. Create a service account:

    gcloud iam service-accounts create diagon-walkthrough-sa --display-name "ML Diagnostics Walkthrough Service Account"

  2. Grant required IAM roles:

    export PROJECT_ID=$(gcloud config get-value project)
export SA_EMAIL="diagon-walkthrough-sa@${PROJECT_ID}.iam.gserviceaccount.com"

for role in roles/hypercomputecluster.editor roles/logging.logWriter roles/storage.objectUser; do
  gcloud projects add-iam-policy-binding ${PROJECT_ID} --member="serviceAccount:${SA_EMAIL}" --role="$role"
done

  3. Create a Cloud Storage bucket:

    export BUCKET_NAME="${PROJECT_ID}-diagon-profiles"
gcloud storage buckets create gs://${BUCKET_NAME} --location=us-central1

:light_bulb: Tip: If your organization restricts roles/hypercomputecluster.editor, create a custom IAM role with only the required hypercomputecluster.machineLearningRuns.* permissions. See the ML Diagnostics IAM guide.

Step 1: Provision and connect to your Cloud TPU VM

Provision a single-host Cloud TPU v5e VM (v5litepod-4) in us-central1-a, attaching your service account and enabling the cloud-platform OAuth scope:

gcloud compute tpus tpu-vm create my-diagon-tpu \
  --zone=us-central1-a \
  --accelerator-type=v5litepod-4 \
  --version=v2-alpha-tpuv5-lite \
  --service-account=${SA_EMAIL} \
  --scopes=https://www.googleapis.com/auth/cloud-platform

Connect to your Cloud TPU VM over SSH:

gcloud compute tpus tpu-vm ssh my-diagon-tpu --zone=us-central1-a

:light_bulb: Tip: If you encounter SSH timeouts or firewall blocks, see Troubleshooting SSH on Compute Engine.

Step 2: Isolate dependencies with venv

Create and activate an isolated Python virtual environment (Python 3.10+ recommended):

python3 -m venv diagon-sdk-env
source diagon-sdk-env/bin/activate

:light_bulb: Tip: If python3 -m venv fails with a missing ensurepip error on Debian or Ubuntu images, run sudo apt-get update && sudo apt-get install -y python3-venv, or install packages in user space with pip install --user.

Step 3: Install the ML Diagnostics SDK and JAX

Install google-cloud-mldiagnostics (version 1.0.6 or higher), google-cloud-logging, and jax[tpu]:

pip install --upgrade pip
pip install "google-cloud-mldiagnostics>=1.0.6" "jax[tpu]" google-cloud-logging \
  -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

:light_bulb: Tip: In production containers, pin exact versions of JAX (for example, jax==0.4.35) and the SDK (google-cloud-mldiagnostics==1.0.6) in your requirements.txt or Dockerfile for reproducible builds.

Step 4: Instrument your workload

Initializing machinelearning_run registers your experiment with the ML Diagnostics control plane and automatically records software metadata (JAX version, XLA compiler flags) and hardware topology (TPU device type, slice count).

  • Setting log_system_metrics=True streams fine-grained TPU hardware counters (TPU TensorCore utilization, TPU duty cycle, HBM utilization) every 10 seconds to Cloud Logging.
  • Wrapping critical sections in a programmatic with xprof(): context manager captures cycle-accurate XProf hardware traces directly to your Cloud Storage bucket.

Create a file named test_workload.py on your Cloud TPU VM:

import time
import logging
import random
import jax
import jax.numpy as jnp
import google.cloud.logging
from google_cloud_mldiagnostics import machinelearning_run, metrics, xprof, metric_types

# 0. Route Python logs to Cloud Logging
logging_client = google.cloud.logging.Client()
logging_client.setup_logging()

# Initialize JAX distributed runtime (auto-detects multi-host TPU slices; safe no-op if single-host)
try:
    jax.distributed.initialize()
except Exception:
    pass
time.sleep(jax.process_index() * 5)  # Stagger multi-host startup to avoid contention

# 1. Initialize the machine learning run
run_name = f"tpu-v5e-walkthrough-run-{int(time.time())}"
my_run = machinelearning_run(
    name=run_name,
    run_group="tpu-getting-started",
    configs={"epochs": 30, "batch_size": 32, "precision": "bfloat16"},
    gcs_path="gs://YOUR_BUCKET_NAME/profiles",  # Replace with your Cloud Storage bucket name
    log_system_metrics=True,                    # Streams libTPU and host telemetry every 10 seconds
    project="YOUR_PROJECT_ID",                  # Replace with your Google Cloud project ID
    region="us-central1",                       # Region where MLRun metadata is stored
)

logging.info(f"MLRun created: {my_run.name} on Worker {jax.process_index()} (Devices: {jax.devices()})")

# Simulate a memory-bound attention block by allocating bfloat16 matrices on TPU HBM
key = jax.random.PRNGKey(42)
x = jax.random.normal(key, (4096, 4096), dtype=jnp.bfloat16)

# 2. Record training loop metrics and capture profiles
learning_rate = 0.001
total_steps = 30
for step in range(1, total_steps + 1):
    step_start = time.perf_counter()
    x = jnp.dot(x, x)
    x.block_until_ready()

    step_time = time.perf_counter() - step_start + random.uniform(0.05, 0.15)
    loss = 2.5 / step + random.uniform(-0.02, 0.02)
    accuracy = min(0.99, 0.50 + (step * 0.015))
    mfu_percent = min(95.0, 52.0 + (step * 0.5))  # MetricType.MFU expects percentage (0-100)

    # Batch record predefined MetricType keys and custom metrics
    metrics.record_metrics([
        {"metric_name": metric_types.MetricType.LOSS, "value": loss},
        {"metric_name": metric_types.MetricType.LEARNING_RATE, "value": learning_rate},
        {"metric_name": metric_types.MetricType.STEP_TIME, "value": step_time},
        {"metric_name": metric_types.MetricType.THROUGHPUT, "value": 32.0 / step_time},
        {"metric_name": metric_types.MetricType.MFU, "value": mfu_percent},
        {"metric_name": "ACCURACY", "value": accuracy},
    ], step=step)

    print(f"[Worker {jax.process_index()}] Step {step}/{total_steps} | Loss: {loss:.4f} | MFU: {mfu_percent:.1f}% | Step Time: {step_time:.3f}s")

    # 3. Capture a hardware performance profile programmatically at step 5
    if step == 5:
        print(f"[Worker {jax.process_index()}] Capturing programmatic XProf trace at step {step}...")
        with xprof():
            for _ in range(5):
                y = jnp.dot(x, x)
                y.block_until_ready()
            time.sleep(2)
        print(f"[Worker {jax.process_index()}] Programmatic profile uploaded to Cloud Storage! Training loop continuing...")

    learning_rate *= 0.95
    time.sleep(2)

Run the instrumented script:

python3 test_workload.py

Practical pro tips

  • Multi-host TPU slices: When running across multi-host TPU slices, stagger worker startup (time.sleep(jax.process_index() * 5)) and pass run_workload_id="my-shared-run-id" to machinelearning_run() so all workers group cleanly under a single unified run.
  • Unique run names: Standalone Cloud TPU VMs require a unique name for each new run (such as appending int(time.time())), whereas GKE workloads automatically append timestamps via the injection webhook.

Step 5: Verify results in the console

Once your script runs, inspect your telemetry across Cloud Logging, Cloud Storage, and the Run Diagnostics dashboard. For full details on console features and CLI queries, see View machine learning runs with ML Diagnostics.

:memo: Note: While Steps 1–4 provision a single-host v5litepod-4 VM for quick testing, the console screenshots below showcase a completed 4-host Cloud TPU v5e slice (tpu-v5e-walkthrough-run-completed, workers t1v-n-fef2b8f2-w-0 through w-3) to illustrate how multi-host system metrics, per-worker Cloud Storage artifacts, and synchronized 4-host XProf timelines render across a full TPU slice.

1. Check Cloud Logging and Cloud Storage

  • Cloud Logging: Open Logs Explorer and confirm the MLRun created entry alongside structured metric payloads under the ml_diagnostics_metric log stream.
  • Cloud Storage: Browse to your profile bucket path (gs://YOUR_BUCKET_NAME/.../plugins/profile/) and verify that compressed multi-host XPlane trace files (.xplane.pb, .trace.json.gz, and .op_stats_v2.pb) were uploaded automatically for each TPU worker host alongside pre-indexed .SSTABLE cache files generated by Managed XProf:

Cloud Storage profile artifacts

Cloud Storage profile artifacts
Cloud Storage profile artifacts1920×1340 530 KB

2. Explore the Run Diagnostics dashboard

Open the Run Diagnostics dashboard in the Google Cloud console via either navigation path:

Runs list

Compare all your machine learning runs side-by-side in the Runs table. Track run status (active in-progress spinner vs. completed checkmark), creation and update timestamps, hardware labels (accelerator_type : tpu), and direct View profiles shortcuts:

Runs list in the Run Diagnostics dashboard
Runs list in the Run Diagnostics dashboard3396×720 232 KB

Run details

Click your run (tpu-v5e-walkthrough-run-completed) to inspect the Details tab. Review run metadata (Run group: tpu-getting-started, Orchestrator: GCE), auto-attached SDK labels (accelerator_type : tpu, diagon_sdk_version : 1-0-7, framework : jax), and auto-collected software configurations (libtpu_version: 0.0.17, framework: JAX, framework_version: 0.6.2):

Run details metadata and configurations
Run details metadata and configurations2902×1875 401 KB

Model metrics

Switch to the Model metrics tab to analyze training convergence curves side-by-side. Every metric recorded via metrics.record_metrics() (such as Learning Rate and Loss) is plotted both over Time (left column) and over Steps (right column). In this completed 42-step run, the charts capture schedule transitions including the initial linear warmup ramp (0.0002 to 0.001 over steps 1–4), smooth exponential decay (0.95x per step over steps 4–42), and monotonic loss reduction (2.5 down to 0.38):

Model metrics convergence charts
Model metrics convergence charts2902×1588 236 KB

Performance metrics

Open the Performance metrics tab to track training speed and stability. Interactive charts for Step Time (in seconds) and Throughput (per second) are plotted over Time and over Steps, with automatic Mean and StdDev dashed horizontal reference lines that highlight step-time jitter at a glance:

Performance metrics step time and throughput charts
Performance metrics step time and throughput charts2902×1588 319 KB

System metrics

Select the System metrics tab to inspect fine-grained hardware counters streamed every 10 seconds across all TPU hosts in your slice (t1v-n-fef2b8f2-w-0 through w-3). Correlate per-host Top 5 TPU Duty Cycle (%) (peaking at 100% during active matrix multiplications), Top 5 TPU TensorCore Utilization (%) (peaking around ~2.8% for this lightweight synthetic benchmark), Top 5 HBM Utilization (%) (~41% peak), and Top 5 Host CPU Utilization (%) against Mean and StdDev baselines:

System metrics multi-host TPU utilization charts
System metrics multi-host TPU utilization charts2902×1600 438 KB

Captured profiles (programmatic XProf capture)

The Profiles tab lists every profile session captured programmatically during the run via with xprof(): (such as 2026_09_12_10_36_04), displaying capture timestamps and tracer levels (Host tracer level: Info, Device tracer level: Enabled). Click any session link to launch the interactive Managed XProf Trace Viewer directly in your browser:

Captured profiles list in Run Diagnostics
Captured profiles list in Run Diagnostics1850×754 76.4 KB

:memo: Note: On standalone Cloud TPU VMs, profile sessions are captured programmatically in your script using the with xprof(): context manager. Live UI-triggered profiling (+ Capture new profile session) is supported on GKE clusters (covered in Part 2).

Managed XProf trace viewer

Clicking any profile session link opens the Managed XProf Trace Viewer (XProf v2.23.1) directly in your browser. Powered by a distributed worker-aggregator backend, Managed XProf loads multi-gigabyte multi-host TPU traces up to 10x faster than self-hosted TensorBoard servers (~1.7 minutes vs. 17 minutes for an 8-host 1.4GB profile) and enables instant link sharing across your team:

Managed XProf trace viewer timeline
Managed XProf trace viewer timeline3450×1860 289 KB

Inside Managed XProf, switch between all 4 TPU hosts (Hosts (4)) and 10 profiling tools (Tools (10)) to pinpoint hardware bottlenecks:

  • Trace Viewer: Inspect a cycle-accurate timeline of host CPU PJRT dispatch threads (host:CPU), TPU TensorCore execution streams across chips (device:TPU:0 through device:TPU:3), XLA High-Level Optimizer (HLO) modules (XLA Modules, XLA Ops, XLA TraceMe), and framework operations.
  • Roofline Analysis: Identify whether operations are memory-bound or compute-bound based on operational intensity (FLOPs/Byte) vs. achieved TFLOPs/s.
  • Memory Viewer: Track High Bandwidth Memory (HBM) usage over time and inspect exact tensor buffer shapes at peak memory allocation to resolve Out-Of-Memory (OOM) errors.

Conclusion

In just a few minutes, you provisioned a Cloud TPU VM, instrumented a JAX workload with the Google Cloud ML Diagnostics SDK, streamed real-time hardware telemetry to Cloud Logging, and analyzed TPU execution traces in the zero-cost Managed XProf viewer.

Ready to try it on your own models? Replace the synthetic loop in test_workload.py with your actual JAX training step to start profiling your Cloud TPU workloads today. If you have questions or run into issues setting this up, drop a comment below!

What’s next: Scaling diagnostics to GKE clusters

While standalone Cloud TPU VMs are ideal for interactive prototyping and single-slice workloads, large-scale foundation model training often spans multi-slice TPU clusters orchestrated by Google Kubernetes Engine (GKE).

In Part 2 of this series, we will cover using ML Diagnostics with GKE clusters (including one-command cluster setup via xpk cluster create --managed-mldiagnostics), featuring:

  • Live UI-triggered on-demand profiling: Trigger interactive multi-host XProf trace sessions on the fly directly from the Google Cloud console (+ Capture new profile session) without modifying training code or restarting jobs.
  • Automated GKE workload correlation: Automatically map TPU slices and GKE pods to Kubernetes JobSet and LeaderWorkerSet hierarchies via the ML Diagnostics GKE webhook.
  • Multi-slice observability: Stream cluster-wide TPU duty cycle, TensorCore utilization, and HBM telemetry across multi-host GKE workloads in the Run Diagnostics dashboard.

Resources and further reading

1 Like