1. Introduction
In the "Control JAX compilation with jax.jit" codelab, you learnt that compilation is one reason a JAX program seems slow, and how to use jax.jit from stopping recompile. But compilation is only one candidate.
When a GPU workload is slow, the Python code almost never tells you which cause you actually have. It could be compilation, a timing measurement that never waited for the GPU, host-device transfers hidden inside your logging, a batch too small to keep the GPU busy, or memory pressure.
In this codelab you start measuring these aspects. You profile a real training step, read the trace in XProf, and drop down to the CUDA timeline with Nsight Systems.
What you'll do
- Separate first-call compilation time from cached execution time, and time JAX honestly with
block_until_ready() - Capture a JAX profiler trace with
jax.profiler.traceand named step annotations - Open that trace in XProf through a Cloud Shell port-forward, with TensorBoard as an alternative front end
- Diagnose four common slowness patterns: too many small operations, host-device transfers, small batches, and memory pressure
- Capture a CUDA-level timeline with Nsight Systems, using NVTX ranges to mark the region you care about
- Summarize the Nsight report inside the notebook with
nsys stats
What you'll need
- A Google Cloud project with billing enabled, and workshop credits or a reservation covering GPU usage
- Quota for at least 2 NVIDIA L4 GPUs in your chosen region (how to check GPU quota)
- Completion of codelabs 1 and 2, or an equivalent JAX GPU environment
- Optionally, Nsight Systems installed on your own machine so you can open the CUDA report in the GUI. It is a free download, and the codelab also prints a text summary of the same report.
Estimated time to complete: 70 minutes.
The profiling workflow
A useful way to approach JAX performance debugging is to move from simple checks to deeper tools.
First make your timing by blocking until GPU work finishes. Then capture a JAX profiler trace so you can see compilation, host activity, and device execution together. Next, open that trace in XProf or TensorBoard to inspect timelines, memory, graphs, and operation statistics. When you need the CUDA-level view, use Nsight Systems to see streams, kernels, memory copies, library calls, and communication.
What to look for
When you open a trace, do not try to understand every event at once. Start by scanning for a few common visual patterns.
Compile spans tell you whether time is going into XLA compilation instead of execution. Empty gaps on GPU rows often mean the host is not feeding the device quickly enough. Transfer activity can point to accidental host-device synchronization, such as logging with float(loss). Memory peaks help you spot batches, activations, or temporary buffers pushing the GPU close to its limit.
2. Before you begin
Select your project
In the Google Cloud console, select or create a project with billing enabled.
Open Cloud Shell
Click Activate Cloud Shell (the terminal icon at the top right of the console) to start a Cloud Shell session, then point it at your project:
gcloud config set project <YOUR_PROJECT_ID>
This codelab runs in the same environment as Codelab 1: Run your first JAX program on NVIDIA GPUs with GKE. If your GKE cluster and JupyterLab Pod are still running, skip ahead to Install what this codelab needs. Otherwise, provision the environment now.
Provision the GPU environment
Run the following in Cloud Shell.
gcloud services enable \
container.googleapis.com \
compute.googleapis.com \
iam.googleapis.com \
cloudresourcemanager.googleapis.com \
logging.googleapis.com \
monitoring.googleapis.com
git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu/terraform
cp terraform.tfvars.example terraform.tfvars
Edit terraform.tfvars and set project_id = ", then provision the cluster and deploy JupyterLab:
terraform init
terraform apply
$(terraform output -raw get_credentials_command)
cd ..
kubectl apply -f deploy/jupyter.yaml
terraform apply takes about 10 minutes. When it finishes, wait for the Pod and the LoadBalancer, then read the one-time JupyterLab token from the Pod log:
kubectl get pod jax-jupyter -w # wait for Running, then Ctrl+C
kubectl get svc jax-jupyter-svc -w # wait for EXTERNAL-IP, then Ctrl+C
kubectl logs jax-jupyter | grep -o 'token=[a-z0-9]*' | head -1
Open http://, paste the token, and create a new Python 3 notebook in /workspace. Every code block in this codelab goes into a cell of that notebook.
Install what this codelab needs
!pip install --quiet xprof nvtx
nsys, the Nsight Systems command-line profiler, already ships in the NVIDIA JAX container, so there is nothing to install for the CUDA timeline later on. Add tensorboard to that install command if you would rather use the TensorBoard profile tab than standalone XProf as your trace viewer.
Set up and verify the GPU
Import the tools, check the GPU, and define the small display helpers the rest of this codelab uses. If this cell fails, fix the environment before continuing; profiling a CPU fallback is misleading.
import csv
import importlib.util
import io
import os
import pathlib
import shutil
import socket
import subprocess
import sys
import tempfile
import textwrap
import time
from IPython.display import HTML, Javascript, display
import jax
import jax.numpy as jnp
import numpy as np
def require_executable(name):
"""Look up `name` on PATH and assert it's found; returns the absolute path or fails fast."""
path = shutil.which(name)
assert path, f"Required executable '{name}' not found on PATH."
return path
XPROF_BIN = require_executable("xprof")
NSYS_BIN = require_executable("nsys")
TENSORBOARD_BIN = shutil.which("tensorboard")
assert importlib.util.find_spec("nvtx"), "Required Python package 'nvtx' is missing. Install with: pip install nvtx"
def show_bars(rows, title, unit="", lower_is_better=False):
"""Render (label, value) pairs as a horizontal bar chart in HTML, scaled to the largest value."""
max_value = max(float(value) for _, value in rows) or 1.0
html = ["<div style='font-family: Arial, sans-serif; max-width: 760px;'>"]
html.append(f"<h4 style='margin: 0 0 8px 0;'>{title}</h4>")
for label, value in rows:
width = max(3, 100 * float(value) / max_value)
html.append(
"<div style='display:grid; grid-template-columns: 190px 1fr 115px; gap: 8px; "
"align-items:center; margin: 6px 0;'>"
f"<div style='font-size:13px;'>{label}</div>"
"<div style='background:#f6f8fa; border-radius:6px; overflow:hidden; height:22px;'>"
f"<div style='height:22px; width:{width:.1f}%; background:#0969da;'></div></div>"
f"<div style='font-size:13px; font-variant-numeric: tabular-nums;'>{value:.3f} {unit}</div>"
"</div>"
)
html.append(f"<div style='font-size:12px; color:#57606a;'>{'Lower' if lower_is_better else 'Higher'} is better.</div></div>")
display(HTML("".join(html)))
def show_table(headers, rows, title=None, aligns=None):
"""Render rows as an HTML table; `aligns` is an optional per-column list of "left"/"right"/"center"."""
aligns = aligns or ["left"] * len(headers)
html = ["<div style='font-family: system-ui; max-width: 980px;'>"]
if title:
html.append(f"<h4 style='margin-bottom: 8px;'>{title}</h4>")
html.append("<table style='border-collapse: collapse; width: 100%; font-size: 13px;'>")
html.append("<thead><tr>")
for h, a in zip(headers, aligns):
html.append(f"<th style='text-align:{a}; border-bottom:1px solid #d0d7de; padding:6px;'>{h}</th>")
html.append("</tr></thead><tbody>")
for row in rows:
html.append("<tr>")
for cell, a in zip(row, aligns):
html.append(f"<td style='text-align:{a}; border-bottom:1px solid #edf0f2; padding:6px; vertical-align:top; white-space:nowrap;'>{cell}</td>")
html.append("</tr>")
html.append("</tbody></table></div>")
display(HTML("".join(html)))
def show_file_list(root, limit=12):
"""Print files under `root` with their sizes (KB); truncates after `limit` entries."""
root = pathlib.Path(root)
files = [p for p in sorted(root.rglob("*")) if p.is_file()]
for p in files[:limit]:
print(f"{p.stat().st_size / 1024:8.1f} KB {p.relative_to(root)}")
if len(files) > limit:
print(f"... {len(files) - limit} more files")
devices = jax.devices()
gpu_devices = [d for d in devices if d.platform == "gpu"]
print(f"JAX version: {jax.__version__}")
print(f"Default backend: {jax.default_backend()}")
print(f"Devices: {devices}")
print(f"xprof: {XPROF_BIN}")
print(f"tensorboard: {TENSORBOARD_BIN or 'not found, XProf standalone is enough'}")
print(f"nsys: {NSYS_BIN}")
assert gpu_devices, f"This lesson assumes a GPU backend. Available devices: {devices}"
print(f"GPU devices: {gpu_devices}")
The cell prints the resolved paths for xprof and nsys and the list of JAX devices, then asserts that at least one of them is a GPU. If the xprof lookup fails, re-run the pip install above and restart the kernel.
3. Build a workload to profile
You need something to be worth profiling, but small enough to rerun in seconds. The cell below defines a two-layer MLP, a mean-squared-error loss, gradients from jax.value_and_grad, and a plain SGD update, all wrapped in one jitted train_step.
The last few lines matter as much as the model. The warmup call compiles train_step once, so every timing cell that follows measures execution rather than accidentally measuring setup.
BATCH = 256
IN_DIM = 1024
HIDDEN = 1024
OUT_DIM = 256
LR = 1e-3
def init_params(key):
"""Initialize the two-layer MLP weights this lesson profiles."""
k1, k2 = jax.random.split(key)
return {
"w1": jax.random.normal(k1, (IN_DIM, HIDDEN), dtype=jnp.float32) * 0.02,
"w2": jax.random.normal(k2, (HIDDEN, OUT_DIM), dtype=jnp.float32) * 0.02,
}
def make_batch(key, batch_size=BATCH):
"""Generate a random (x, target) batch with the standard input/output dims."""
kx, ky = jax.random.split(key)
x = jax.random.normal(kx, (batch_size, IN_DIM), dtype=jnp.float32)
y = jax.random.normal(ky, (batch_size, OUT_DIM), dtype=jnp.float32)
return x, y
def loss_fn(params, batch):
"""Forward pass plus MSE loss"""
x, target = batch
hidden = jax.nn.gelu(x @ params["w1"])
pred = hidden @ params["w2"]
return jnp.mean((pred - target) ** 2)
# This is the function we will profile.
@jax.jit
def train_step(params, batch):
"""One compiled SGD step"""
loss, grads = jax.value_and_grad(loss_fn)(params, batch)
params = jax.tree.map(lambda p, g: p - LR * g, params, grads)
return params, loss
key = jax.random.key(0)
params = init_params(key)
batch = make_batch(jax.random.fold_in(key, 1))
# Warm up once so later timing is on execution
params, loss = train_step(params, batch)
jax.block_until_ready((params, loss))
print(f"x shape/device: {batch[0].shape} on {batch[0].device}")
print(f"target shape/device: {batch[1].shape} on {batch[1].device}")
print(f"w1 shape/device: {params['w1'].shape} on {params['w1'].device}")
print(f"warmup loss: {float(loss):.4f}")
Every line of the printout should name a CUDA device: the inputs at (256, 1024), the targets at (256, 256), and w1 at (1024, 1024), followed by a single warmup loss value.
4. Separate compilation from execution and time honestly
A jitted function has two very different modes:
- The first call for a new input signature traces and compiles, then executes.
- Later calls with the same shapes and dtypes reuse the compiled executable.
Changing the batch shape creates a new signature, so JAX has to compile again. This is the root cause of many training loop keeps pausing.
Measure compile cost against cached execution
Here you have the same train_step three times: once with a cleared cache, once with the cache warm, and once with a batch of a different shape.
jax.clear_caches()
compile_params = init_params(jax.random.key(101))
compile_batch = make_batch(jax.random.key(102), batch_size=BATCH)
# Trace + compile + execute.
t0 = time.perf_counter()
compile_params, compile_loss = train_step(compile_params, compile_batch)
jax.block_until_ready((compile_params, compile_loss))
first_ms = (time.perf_counter() - t0) * 1000
# Execute using the cached compiled executable.
t0 = time.perf_counter()
compile_params, compile_loss = train_step(compile_params, compile_batch)
jax.block_until_ready((compile_params, compile_loss))
cached_ms = (time.perf_counter() - t0) * 1000
# Different batch shape which triggers another compile.
smaller_batch = make_batch(jax.random.key(103), batch_size=BATCH // 2)
t0 = time.perf_counter()
_, shape_loss = train_step(compile_params, smaller_batch)
shape_change_ms = (time.perf_counter() - t0) * 1000
print(f"First call, same shape (compile + execute): {first_ms:8.2f} ms")
print(f"Cached call, same shape (execute only): {cached_ms:8.2f} ms")
print(f"New batch shape (compile + execute): {shape_change_ms:8.2f} ms")
show_bars(
[
("first call", first_ms),
("cached call", cached_ms),
("new shape", shape_change_ms),
],
title="Compilation Cost vs. Cached Execution",
unit="ms",
lower_is_better=True,
)
# Re-warm the original shape.
params, loss = train_step(params, batch)
jax.block_until_ready((params, loss))
You should get three numbers and a bar chart. Two of the three bars — the first call and the new shape — include compilation, so both should be much larger than the cached call. That middle bar is the only one that reflects what a steady-state training step actually costs.
Block before you stop the clock
JAX dispatches GPU work asynchronously, so Python can return before the GPU has finished. Without block_until_ready(), you usually measure how long Python took to enqueue the work, not how long the GPU took to execute it.
Here you have the same training step twice: once incorrectly, and once with a block on the result.
t0 = time.perf_counter()
params_dispatch, loss_dispatch = train_step(params, batch)
dispatch_ms = (time.perf_counter() - t0) * 1000
t0 = time.perf_counter()
params_ready, loss_ready = train_step(params, batch)
jax.block_until_ready((params_ready, loss_ready))
ready_ms = (time.perf_counter() - t0) * 1000
print(f"Dispatch-only timing: {dispatch_ms:.3f} ms")
print(f"Blocked timing: {ready_ms:.3f} ms")
show_bars(
[("dispatch only", dispatch_ms), ("block_until_ready", ready_ms)],
title="Timing the Same JAX Step",
unit="ms",
lower_is_better=True,
)
The dispatch-only bar should be the shorter of the two. That shorter number is not a faster program, it is an unmeasured one.
5. Fix the too-many-small-operations pattern
Another common GPU performance challenge is launching many tiny operations from Python. Each small JAX operation carries Python dispatch overhead and may produce a small GPU kernel that finishes before the device is properly busy.
jax.jit helps because it lets XLA see the whole chain and fuse or schedule it as one compiled unit. The two functions below do the same calculation. The first dispatches the operations from Python and the second compiles the chain.
SMALL_OP_STEPS = 50
small_x = jnp.ones((4096,), dtype=jnp.float32)
def many_small_ops(x):
"""Un-jitted chain of small operations."""
# Each loop iteration dispatches JAX operations from Python.
y = x
for _ in range(SMALL_OP_STEPS):
y = jnp.sin(y) + 0.01 * y
return y
@jax.jit
def compiled_chain(x):
"""Same operations under `jax.jit` so XLA can fuse."""
# JAX traces the whole chain, and XLA can optimize it as one compiled computation.
y = x
for _ in range(SMALL_OP_STEPS):
y = jnp.sin(y) + 0.01 * y
return y
_ = compiled_chain(small_x).block_until_ready()
# Time the non-jitted chain.
t0 = time.perf_counter()
_ = many_small_ops(small_x).block_until_ready()
small_ops_ms = (time.perf_counter() - t0) * 1000
# Time the compiled chain.
t0 = time.perf_counter()
_ = compiled_chain(small_x).block_until_ready()
compiled_chain_ms = (time.perf_counter() - t0) * 1000
print(f"Many small Python-dispatched ops: {small_ops_ms:8.3f} ms")
print(f"Compiled chain with jax.jit: {compiled_chain_ms:8.3f} ms")
show_bars(
[("many small ops", small_ops_ms), ("compiled chain", compiled_chain_ms)],
title="Too Many Small Operations",
unit="ms",
lower_is_better=True,
)
The compiled chain should be the shorter bar. The array here has only 4096 elements and the loop runs 50 times, so the un-jitted version pays Python dispatch overhead 50 times over for very little arithmetic each time. This is the pattern to recognize in your own code: a Python loop over small JAX operations that could be one compiled function instead.
6. Capture a JAX profiler trace
Timings tell you that something is slow. A trace tells you what happened over time: when Python was active, when XLA compiled, when the GPU ran kernels, and where the gaps were.
Raw traces are full of low-level operation names, so annotate the workload before you capture it. Marking each training step gives you human-readable landmarks to navigate by. JAX offers three annotations:
StepTraceAnnotationnames repeated steps, such as training iterations.TraceAnnotationnames a region inside a step, such as batch prep or the optimizer update.annotate_functionnames a Python function in the trace.
This cell traces six annotated training steps into a temporary directory:
trace_dir = pathlib.Path(tempfile.mkdtemp(prefix="jax-trace-"))
trace_params = params
trace_key = key
with jax.profiler.trace(str(trace_dir), create_perfetto_link=False):
for step_num in range(6):
with jax.profiler.StepTraceAnnotation("train_step", step_num=step_num):
with jax.profiler.TraceAnnotation("make_batch"):
trace_key = jax.random.fold_in(trace_key, step_num + 10)
trace_batch = make_batch(trace_key)
with jax.profiler.TraceAnnotation("sgd_update"):
trace_params, trace_loss = train_step(trace_params, trace_batch)
jax.block_until_ready((trace_params, trace_loss))
print(f"Trace directory: {trace_dir}")
show_file_list(trace_dir)
You should see the path of the trace directory followed by a short listing of the trace files written inside it, each with its size. That directory path is what you point the viewer at in the next step, so keep the cell output visible.
7. Open the trace in XProf
XProf is the profiler UI behind the JAX trace format. It reads the trace directory directly and serves a web interface with a trace viewer, a memory viewer, a graph viewer, and operation statistics.
Start it inside the Pod:
xprof_port = 6007
xprof_proc = subprocess.Popen(
[XPROF_BIN, "--port", str(xprof_port), str(trace_dir)],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
time.sleep(3)
xprof_url = f"http://localhost:{xprof_port}"
print(f"XProf is running at {xprof_url} (PID {xprof_proc.pid})")
print(f"Stop it later with: kill {xprof_proc.pid}")
display(Javascript(f'window.open("{xprof_url}", "_blank");'))
display(HTML(f'<a href="{xprof_url}" target="_blank" rel="noopener">Open XProf</a>'))
The cell prints the URL XProf is serving on, the process ID, and the kill command that stops it. Write that PID down — you need it in the Clean up step.
Reach the viewer from your laptop
Nothing listening inside the Pod is reachable from your machine by default. kubectl port-forward opens a tunnel from Cloud Shell to a port inside the Pod, and Cloud Shell's Web Preview publishes that tunnel to your browser.
Open a new Cloud Shell tab — the command runs in the foreground until you stop it, and your notebook keeps running in the Pod meanwhile:
kubectl port-forward pod/jax-jupyter 8080:6007
Then click Web Preview - Preview on port 8080 in the Cloud Shell toolbar and you should be able to access the UI.
Read the trace
Once the Web Preview tab loads XProf:
- Choose the run from the Runs dropdown.
- Open Tools - trace_viewer.
- Find the
train_stepspans — these are theStepTraceAnnotationnames from the previous step. - Look for compile spans, GPU gaps, copies, and kernel activity.
Trace reading checklist
The table below shows how visual pattern in the chart trigger the next action:
Visual clue | What it probably means | What to try next |
Long compile span before first step | Normal first-call JIT compilation | Warm up before measuring. |
Compile spans between steps | Recompilation | Check changing shapes, dtypes, or static args from codelab 2. |
GPU rows have white gaps | Host is not feeding the GPU | Look for data loading, printing, |
Many tiny kernels | Launch overhead or too-small compiled regions | JIT a larger function; batch work. |
Memcpy activity between steps | Host-device transfers | Keep metrics on device; log less often; use transfer guard. |
High peak memory | Activations or temporary buffers may dominate | Open |
Three other entry points are worth mentioning even though this codelab does not use them: start_trace() and stop_trace() for programmatic trace regions that do not fit into a with jax.profiler.trace(...) block, start_server() plus python -m jax.collect_profile for profiling long-running jobs, and Perfetto export for opening traces in the Perfetto UI. The JAX profiling guide covers all three.
8. Diagnose host transfers and batch size
Two of the checklist items above are common enough to be worth reproducing : accidental host transfers, and a batch too small to fill the GPU.
Host-device transfers
Pulling a JAX value back to Python inside a loop forces synchronization. Common examples are float(loss), .item(), np.asarray(...), and printing arrays.
This code compares two logging styles. The bad version turns the loss into a Python float every step. The better one keeps losses as JAX arrays and synchronizes once at the end.
N = 30
# Convert the loss to a Python float every step.
t0 = time.perf_counter()
bad_params = params
bad_losses = []
for _ in range(N):
bad_params, bad_loss = train_step(bad_params, batch)
bad_losses.append(float(bad_loss))
bad_ms = (time.perf_counter() - t0) * 1000 / N
# Keep metrics as JAX arrays and synchronize once.
t0 = time.perf_counter()
good_params = params
good_losses = []
for _ in range(N):
good_params, good_loss = train_step(good_params, batch)
good_losses.append(good_loss)
jax.block_until_ready((good_params, good_losses))
good_ms = (time.perf_counter() - t0) * 1000 / N
print(f"float(loss) every step: {bad_ms:.3f} ms / step")
print(f"deferred sync: {good_ms:.3f} ms / step")
show_bars(
[("float(loss) every step", bad_ms), ("deferred sync", good_ms)],
title="Cost of Pulling Metrics to Python",
unit="ms/step",
lower_is_better=True,
)
# Transfer guard can catch accidental transfers.
try:
_, guard_loss = train_step(params, batch)
guard_loss.block_until_ready()
with jax.transfer_guard("disallow"):
_ = float(guard_loss)
except RuntimeError as e:
print("\nTransfer guard caught an implicit transfer:")
print(str(e).splitlines()[0])
The per-step cost of the float(loss) version should be the higher of the two bars, and the cell should finish by printing the first line of the RuntimeError that the transfer guard raised.
Batch size and throughput
Small batches often do not give the GPU enough parallel work. Throughput usually improves as the batch grows, then flattens once the GPU is saturated or memory becomes the limit.
Batch size also affects memory. The sweep below includes a simple estimate for the batch-shaped buffers in this forward pass: input, hidden activation, and output. Real training uses more than that, because gradients and temporary buffers count too.
@jax.jit
def forward_only(params, x):
"""Forward pass only."""
hidden = jax.nn.gelu(x @ params["w1"])
return hidden @ params["w2"]
batch_results = []
bytes_per_float32 = np.dtype(np.float32).itemsize
for batch_size in (1, 8, 32, 128, 256, 512, 1024):
xb = jax.random.normal(jax.random.key(batch_size), (batch_size, IN_DIM), dtype=jnp.float32)
_ = forward_only(params, xb).block_until_ready()
reps = 80 if batch_size <= 128 else 30
# Async dispatch lets JAX queue all `reps` calls without blocking. We block once at the end and divide by reps, so this measures *amortized* time per call when the executor stays busy.
t0 = time.perf_counter()
for _ in range(reps):
yb = forward_only(params, xb)
yb.block_until_ready()
ms = (time.perf_counter() - t0) * 1000 / reps
examples_per_sec = batch_size / (ms / 1000)
# Simple memory estimate for batch-shaped forward buffers.
estimated_forward_mib = batch_size * (IN_DIM + HIDDEN + OUT_DIM) * bytes_per_float32 / 2**20
batch_results.append((batch_size, ms, examples_per_sec, estimated_forward_mib))
show_table(
["Batch", "ms/call (avg)", "examples/sec", "estimated batch buffers"],
[
(batch_size, f"{ms:.3f}", f"{examples_per_sec:,.0f}", f"{estimated_mib:.1f} MiB")
for batch_size, ms, examples_per_sec, estimated_mib in batch_results
],
title="Batch Size, Throughput, and Memory",
aligns=["right", "right", "right", "right"],
)
show_bars(
[(f"batch {batch_size}", examples_per_sec) for batch_size, _, examples_per_sec, _ in batch_results],
title="Throughput by Batch Size",
unit="ex/s",
lower_is_better=False,
)
You get a table of seven rows and a bar chart. Look the examples/sec column, not ms/call: throughput should climb sharply over the first few batch sizes and then flatten as the GPU saturates. The ms/call column grows the whole way, which is expected.
9. Check GPU memory pressure
JAX usually preallocates a certain % of GPU memory on first use. That is deliberate because it reduces allocation overhead and fragmentation. It also means nvidia-smi can look almost full even when your model is tiny, which sends a lot of people chasing a memory leak that does not exist.
Use memory_stats for a quick process-level view, then XProf's memory tools for deeper analysis.
def gib(value):
"""Convert a byte count to gibibytes."""
return value / 2**30
print(f"{'device':<26} {'limit':>12} {'in use':>12} {'peak':>12}")
print("-" * 66)
for device in gpu_devices:
stats = device.memory_stats()
if not stats:
print(f"{str(device):<26} memory_stats unavailable")
continue
limit = stats.get("bytes_limit")
in_use = stats.get("bytes_in_use")
peak = stats.get("peak_bytes_in_use")
limit_s = f"{gib(limit):.2f} GiB" if limit is not None else "n/a"
in_use_s = f"{gib(in_use):.2f} GiB" if in_use is not None else "n/a"
peak_s = f"{gib(peak):.2f} GiB" if peak is not None else "n/a"
print(f"{str(device):<26} {limit_s:>12} {in_use_s:>12} {peak_s:>12}")
You should get one row per visible GPU. The limit column is the allocator's reservation rather than the card's total memory, and in use for this small MLP should be a small fraction of it.
Memory settings must be set before importing JAX. In the notebook, that means setting them before kernel startup and then restarting the kernel.
Variable | Example | Use when |
|
| You share a GPU and want JAX to reserve less memory. |
|
| You want on-demand allocation, accepting more fragmentation risk. |
|
| You are debugging memory and want deallocation; too slow for normal training. |
This cell shows which of those are set in the current kernel:
for name in (
"XLA_PYTHON_CLIENT_MEM_FRACTION",
"XLA_PYTHON_CLIENT_PREALLOCATE",
"XLA_PYTHON_CLIENT_ALLOCATOR",
):
print(f"{name}={os.environ.get(name, '<unset>')}")
print("\nExample for a shared GPU, set before launching Python/Jupyter:")
print("export XLA_PYTHON_CLIENT_MEM_FRACTION=0.50")
All three will most likely print
10. Capture a CUDA timeline with Nsight Systems
XProf is the right first profiler for JAX. Nsight Systems is the second view, for when you need the CUDA timeline: streams, CUDA API calls, kernels, memcopies, cuBLAS and cuDNN calls, and eventually NCCL.
The workflow has four parts:
- Put the workload in a short script.
- Add NVTX ranges around the steps you care about.
- Run the script with
nsys profile. - Open the
.nsys-repfile with the Nsight Systems GUI.
Write the capture script
Profiling a notebook directly is not easy, so this code below writes a small standalone script instead. It warms up once, then uses NVTX ranges to mark the region worth capturing.
You do not need to read every line. Most of it repeats the same tiny model from earlier so nsys can profile a fresh Python process.
nsight_dir = pathlib.Path(tempfile.mkdtemp(prefix="jax-nsight-"))
script_path = nsight_dir / "nsight_train_step.py"
report_base = nsight_dir / "jax_train_step"
report_path = pathlib.Path(f"{report_base}.nsys-rep")
script_source = f"""
import nvtx
import jax
import jax.numpy as jnp
BATCH = {BATCH}
IN_DIM = {IN_DIM}
HIDDEN = {HIDDEN}
OUT_DIM = {OUT_DIM}
LR = {LR}
def init_params(key):
k1, k2 = jax.random.split(key)
return {{ "{{" }}
"w1": jax.random.normal(k1, (IN_DIM, HIDDEN), dtype=jnp.float32) * 0.02,
"w2": jax.random.normal(k2, (HIDDEN, OUT_DIM), dtype=jnp.float32) * 0.02,
{{ "}}" }}
def make_batch(key):
kx, ky = jax.random.split(key)
return (
jax.random.normal(kx, (BATCH, IN_DIM), dtype=jnp.float32),
jax.random.normal(ky, (BATCH, OUT_DIM), dtype=jnp.float32),
)
def loss_fn(params, batch):
x, target = batch
hidden = jax.nn.gelu(x @ params["w1"])
pred = hidden @ params["w2"]
return jnp.mean((pred - target) ** 2)
@jax.jit
def train_step(params, batch):
loss, grads = jax.value_and_grad(loss_fn)(params, batch)
params = jax.tree.map(lambda p, g: p - LR * g, params, grads)
return params, loss
key = jax.random.key(0)
params = init_params(key)
batch = make_batch(key)
# Warm up before the region we care about.
params, loss = train_step(params, batch)
jax.block_until_ready((params, loss))
with nvtx.annotate("profile_region", domain="jax_course"):
for step in range(8):
with nvtx.annotate(f"train_step_{{ "{{" }}step{{ "}}" }}", domain="jax_course"):
key = jax.random.fold_in(key, step)
batch = make_batch(key)
params, loss = train_step(params, batch)
jax.block_until_ready((params, loss))
"""
script_path.write_text(textwrap.dedent(script_source))
print(f"Wrote script: {script_path}")
print(f"Report path: {report_path}")
The cell prints the two paths it will use. Nothing has run yet.
Run Nsight Systems
The invocation below starts collection only when the NVTX range profile_region@jax_course begins, and stops when that range ends. That keeps the report focused on post-warmup execution instead of burying it under process startup and JIT compilation.
import base64
nsys_cmd = [
NSYS_BIN,
"profile",
"--trace=cuda,nvtx,osrt,cudnn,cublas",
"--capture-range=nvtx",
"--capture-range-end=stop",
"--nvtx-capture=profile_region@jax_course",
"--force-overwrite=true",
f"--output={report_base}",
"-e",
"NSYS_NVTX_PROFILER_REGISTER_ONLY=0",
sys.executable,
str(script_path),
]
print("Running Nsight Systems:")
print(" ".join(nsys_cmd))
nsys_result = subprocess.run(nsys_cmd, capture_output=True, text=True, check=False)
if nsys_result.returncode != 0:
print("nsys profile failed")
print(nsys_result.stdout[-2000:])
print(nsys_result.stderr[-4000:])
raise RuntimeError("Nsight Systems capture failed. This lesson requires nsys to run successfully.")
print(f"Nsight report: {report_path} ({report_path.stat().st_size / 1024:.1f} KB)")
b64 = base64.b64encode(report_path.read_bytes()).decode()
download_html = (
f'<a href="data:application/octet-stream;base64,{b64}" '
f'download="{report_path.name}" '
f'style="font-weight:600;">Download {report_path.name}</a>'
)
display(HTML(download_html))
The capture takes a while, because it launches a fresh Python process that imports JAX and compiles train_step again. When it finishes you get the report path and size, and a Download link. That link is a data URL, so it works straight from the notebook in your browser — no port-forward needed.
Read the Nsight timeline
To open the report in the GUI:
- Install Nsight Systems on your local machine from developer.nvidia.com/nsight-systems. It is free and runs on Linux, macOS, and Windows.
- Download the generated
.nsys-repreport using the link above. - Launch
nsys-uiand open the file with File - Open, or drag and drop it into the window.
Start with these rows:
- NVTX: find
profile_regionandtrain_step_*. - CUDA GPU rows: check kernel coverage and white gaps, where a white gap means idle GPU time.
- CUDA API: look for
cudaLaunchKernel,cudaMemcpy*, and synchronization calls. - cuBLAS / cuDNN: matmul, convolution, and attention library calls appear here when traced.
- OS Runtime (
osrt): host waits, locks, sleeps, and other blocking behavior.
For GPU utilization metrics, rerun the capture with --gpu-metrics-devices=cuda-visible if your environment allows GPU metrics collection.
Summarize the report
The GUI is the main Nsight experience, but nsys stats gives a useful text summary you can read. This code prints the top kernel, memory-operation, and CUDA API tables.
stats_cmd = [
NSYS_BIN,
"stats",
"--quiet",
"--report",
"cuda_gpu_kern_sum,cuda_gpu_mem_time_sum,cuda_api_sum",
"--format",
"csv",
"--timeunit",
"msec",
str(report_path),
]
stats_result = subprocess.run(stats_cmd, capture_output=True, text=True, check=False)
if stats_result.returncode != 0:
print("nsys stats failed")
print(stats_result.stderr[-3000:])
else:
sections = [s for s in stats_result.stdout.strip().split("\n\n") if s.strip()]
for idx, section in enumerate(sections[:3], start=1):
reader = csv.DictReader(io.StringIO(section))
rows = list(reader)
if not rows:
continue
headers = reader.fieldnames or []
print(f"\nReport section {idx}: {headers}")
compact_rows = []
for row in rows[:8]:
name = row.get("Name") or row.get("Operation") or row.get("Name:Demangled") or ""
time_pct = row.get("Time (%)", "")
total = row.get("Total Time (ms)") or row.get("Total Time (msec)") or row.get("Total Time") or row.get("Total Time (ns)") or ""
count = row.get("Instances") or row.get("Count") or row.get("Num Calls") or row.get("Calls") or ""
compact_rows.append((time_pct, total, count, name[:90]))
show_table(["Time (%)", "Total time (ms)", "Count", "Name / operation"], compact_rows, title=f"Nsight stats section {idx}", aligns=["right", "right", "right", "left"])
You should get up to three tables, each showing the top eight rows by time. Read the kernel summary first: it ranks GPU kernels by total time, which tells you where the device actually spent it. Compare that against the memory-operation table — if copies are competing with kernels for time, you are back to the host-transfer problem from earlier.
11. Clean up
The GPU node bills whether or not you are running anything on it, so do not skip this step.
Stop the viewers you started inside the Pod first: run the kill commands that the XProf and TensorBoard cells printed, and press Ctrl+C in any Cloud Shell tab still running kubectl port-forward.
Delete the Jupyter workload, including the LoadBalancer and the persistent volume:
kubectl delete -f deploy/jupyter.yaml
Destroy the cluster, node pool, VPC, and service account:
cd terraform
terraform destroy
Type yes when prompted, then confirm nothing is left behind:
gcloud container clusters list
gcloud compute instances list
Both should be empty for this project. If you created a project just for this series, you can instead delete the whole project from the Cloud console.
12. Congratulations
You profiled a JAX training step end to end, from a time.perf_counter measurement all the way down to the CUDA timeline.
What you've learned
- How to time JAX honestly with
block_until_ready(), and why a dispatch-only measurement is meaningless - How to separate first-call compilation time from cached execution time, and how a new input shape forces a recompile
- How to capture a trace with
jax.profiler.traceand label it withStepTraceAnnotationandTraceAnnotation - How to open that trace in XProf, and in TensorBoard as an alternative front end, through a Cloud Shell port-forward
- How to recognize the common slowness patterns: recompilation, too many small operations, host-device transfers, inefficient batch sizes, and memory pressure
- How to inspect GPU memory with
memory_stats(), and howXLA_PYTHON_CLIENT_MEM_FRACTIONchanges JAX's reservation before startup - How to capture and read a CUDA timeline with Nsight Systems using NVTX ranges, and how to summarize it with
nsys stats
Next steps
- Codelab 4: Train a model on GPU with JAX, Optax, and Fashion-MNIST where you will experience a real training loop on real data, with the profiling habits from this lab already in place
- Re-run the Nsight capture with
--gpu-metrics-devices=cuda-visibleadded tonsys_cmd, and compare the GPU utilization counters against the kernel coverage you saw in the timeline - Set
XLA_PYTHON_CLIENT_MEM_FRACTION=0.50before restarting the kernel, then re-run thememory_stats()cell and check how thelimitcolumn changed