How to Profile TPU Code

Chapter 1: The TPU Software Stack

Google exposes a stack of APIs for programming TPUs, from high-level JAX code to low-level Pallas or HLO. Most programmers write JAX code exclusively, which lets you write abstract NumPy-style linear algebra that is compiled automatically to run efficiently on TPUs.

Here is a simple example — a JAX program that multiplies two matrices:

import jax
import jax.numpy as jnp

def multiply(x, y):
  return jnp.einsum('bf,fd->db', x, y)

y = jax.jit(multiply)(jnp.ones((128, 256)),
                       jnp.ones((256, 16), dtype=jnp.bfloat16))

By calling jax.jit, we tell JAX to trace this function and emit a lower-level IR called StableHLO, a platform-agnostic IR for ML computation. This is then lowered to HLO by the XLA compiler.

The compiler runs many optimization passes to determine fusions (grouping related ops together), layouts (how arrays are tiled in memory), and other factors. The resulting HLO is what you observe in a JAX profile.

Here is the HLO for the matmul above (abridged):

ENTRY %main.5 (Arg_0.1: f32[128,256], Arg_1.2: bf16[256,16]) -> f32[16,128] {
  %Arg_1.2 = bf16[256,16]{1,0} parameter(1)
  %convert.3 = f32[256,16]{1,0} convert(bf16[256,16] %Arg_1.2)
  %Arg_0.1 = f32[128,256]{1,0} parameter(0)
  ROOT %dot.4 = f32[16,128]{1,0} dot(
    f32[256,16] %convert.3,
    f32[128,256] %Arg_0.1),
    lhs_contracting_dims={0}, rhs_contracting_dims={1}
}

Notice how closely the HLO matches the original JAX code. The ROOT %dot.4 instruction is the actual matmul, contracting along dimensions 0 and 1 of the two f32 matrices.

When something goes wrong at a lower level, you have one more escape hatch: writing custom kernels in Pallas. But most of the time, you work at the JAX level and use the profiler to understand the HLO below.

Chapter 2: The JAX Profiler

JAX provides a multi-purpose TPU profiler that records everything from the duration of each subcomponent, the HLO of each program, memory usage, and more. You use the jax.profiler module to capture a trace while your program runs.

Here is the minimal recipe for capturing a profile:

import jax

with jax.profiler.trace("/tmp/tensorboard"):
  key = jax.random.key(0)
  x = jax.random.normal(key, (1024, 1024))
  y = x @ x
  y.block_until_ready()

# View in TensorBoard:
# pip install tensorboard tensorboard-plugin-profile
# tensorboard --logdir=/tmp/tensorboard

Once in TensorBoard, the profiler exposes several key views:

For annotating your traces, use jax.named_scope to label regions of code. These names show up in the Trace Viewer as nested scopes, making it much easier to identify which part of a Transformer corresponds to which operations:

with jax.named_scope("attention"):
  q = x @ w_q
  k = x @ w_k
  v = x @ w_v
  attn = softmax(q @ k.T / sqrt_dk) @ v

with jax.named_scope("ffn"):
  h = gelu(attn @ w_up)
  out = h @ w_down

Chapter 3: The Trace Viewer

The Trace Viewer is probably the most useful part of the profiler. It shows a chronological timeline of all actions on each TPU core, letting you see exactly what the hardware is doing and for how long.

A few key things to understand about the Trace Viewer:

In a typical Transformer profile, you will see repeated blocks — one per layer. Within each block, you can identify the attention portion and the MLP/FFN portion. The XLA ops row tells you what is actually executing, while the scope rows help you map back to your code.

When you click on an XLA op in the trace, you get detailed information:

• Op name and HLO text — the full fusion or operation
• Duration — how long this op took in microseconds
• Source code location — which line of Python generated this op
• Link to Graph Viewer — jump to see the dataflow graph around this op

Since all TPUs typically execute the same instructions in SPMD programs, you usually only need to look at TPU:0. The timeline for other cores will be nearly identical.

Common Trace Patterns

Annotating Your Code

Good annotations are the difference between a readable profile and a wall of inscrutable HLO names. The key tools:

• jax.named_scope("name") — creates a visual scope in the Trace Viewer that groups all ops inside it
• jax.named_call(fn, name="name") — similar but wraps a function call
• Nesting scopes creates a hierarchy: with named_scope("layer_0"): with named_scope("attention"):

For production Transformer code, annotate at minimum: each layer, attention vs MLP, each projection, and each collective operation. This makes the profile immediately navigable.

Chapter 4: Reading HLO Ops

HLO is not actually very hard to read, and it is immensely helpful for understanding what a given part of the trace corresponds to. Let us break down the anatomy of an HLO op with a real example.

Consider this op called fusion.3:

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)}
  fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32),
  kind=kCustom, calls=%all-reduce-scatter.3

Let us unpack each component:

Understanding Tilings

The layout notation tells us how an N-dimensional array is laid out sequentially in memory. Consider a simpler example:

f32[3,5]{1,0:T(2,2)}

This means:

• {1,0} — read right-to-left: dimension 0 is most major, dimension 1 is most minor. The physical layout is [3, 5], identical to the logical shape in this case.
• T(2,2) — the array is tiled in 2x2 chunks. Within each chunk, elements are stored row-major.

For the more complex tiling bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)}, we have two levels of tiling: an outer (8, 128) tiling and an inner (2, 1) tiling within that. The inner (2, 1) tiling is used for bf16 so that loads are always multiples of 4 bytes.

Worked Tiling Example

Let us walk through a concrete padding calculation. Consider bf16[3, 5]{1,0:T(8,128)(2,1)}:

1. The outer tiling T(8, 128) means each tile is 8 rows × 128 columns
2. To fit our [3, 5] array, we need ceil(3/8)=1 row tile and ceil(5/128)=1 col tile
3. Padded shape: [8, 128] — much larger than [3, 5]
4. Memory overhead: (8×128) / (3×5) = 68x wasted memory for this tiny array

Common HLO Op Types

Chapter 5: The Graph Viewer

While some fusions in the Trace Viewer can look complicated, the XLA Graph Viewer makes them much easier to parse. It shows the HLO as a visual dataflow graph — boxes for operations, arrows for data dependencies.

The Graph Viewer is especially helpful for:

• Understanding fusions — seeing which ops are grouped together and why
• Tracing data flow — following an array from creation through transformations to consumption
• Identifying collectives — seeing where AllGather, AllReduce, and ReduceScatter ops appear in the graph
• Finding the source — hovering over a box often shows the line of code where the function was defined

A common pattern you will see in the Graph Viewer is a fusion containing a dot (matmul) followed by some pointwise operations. For example, an MLP up-projection might show:

XLA's fusion pass is quite aggressive. It will fuse the matmul with subsequent pointwise operations (ReLU, GeLU, bias adds, etc.) into a single operation, eliminating intermediate memory traffic. This is one of XLA's main advantages: you write clean, modular code and the compiler fuses it for free.

However, fusions are not always beneficial. Sometimes XLA fuses operations that would be better left separate, or it fails to fuse operations you expect. The Graph Viewer helps you verify that the compiler is doing what you want.

Chapter 6: A Worked Profile

Let us walk through a real profile of a Transformer and verify our theoretical predictions. We will look at a model running on 8 TPU v2 cores with 4-way data parallelism and 2-way model parallelism.

The FFN Block

Zooming into the FFN (feed-forward network) block, we see the up-projection matmul. The HLO tells us the per-shard shapes:

// Per-shard computation:
bf16[8, 1024, 8192] * bf16[8192, 16384] -> bf16[8, 1024, 16384]

// True (unsharded) shapes (4-way DP, 2-way MP):
X: bf16[32, 1024, 8192] * Win: bf16[8192, 32768]
  -> Tmp: bf16[32, 1024, 32768]

How long should this take? Our batch size per DP shard is 8 * 1024 = 8192 tokens, so we are solidly compute-bound.

The profile reports 96 ms. That is essentially perfect FLOPs utilization — we are hitting the roofline.

Communication: The ReduceScatter

At the end of the second matmul, there is a small fusion that is actually a ReduceScatter:

%fusion.1 = bf16[8,1024,4096]{2,1,0:T(8,128)(2,1)S(1)}
  fusion(bf16[8,1024,8192]{...} %fusion.31),
  kind=kCustom, calls=%all-reduce-scatter.1

Expected time: the array has size 2 * 32 * 1024 * 8192 = 128 MB with the batch axis sharded 4 ways on a TPU v2 4x2 (1.2 × 10¹¹ bidirectional bandwidth). One hop required:

Profile reports: 1.13 ms. Again, essentially at the roofline.

Attention

The attention component shows the Q, K, V projections as separate matmuls. The Q projection uses weight matrix W_Q of shape [d_model=8192, n_heads=32, d_qkv=256], with Megatron sharding along the head dimension.

Let us do the exercise. The Q projection has true shape:

With Megatron sharding along the 32 heads (2-way MP), each shard computes 16 heads. The FLOPs per shard:

At 23 TFLOPS per core: T ≈ 68.7 × 10¹² / (23 × 10¹² × 8) ≈ 373 ms. If the profile shows a similar number, we are at the roofline. If much longer, check for retiling or communication overhead.

Verifying Communication Predictions

The workflow for any operation in a profile:

Chapter 7: Memory Profile

The Memory Profile view shows program memory as a function of time. This is invaluable for debugging out-of-memory (OOM) errors — the most common failure mode when scaling models.

In a typical profile, you will see:

The Memory Profile also helps answer practical questions:

• How much headroom do I have? If parameters consume 7.5 GB and you have 8.5 GB free, you know you can increase batch size or model size.
• Is activation recomputation working? Without it, activation memory grows linearly with the number of layers. With it, you should see a sawtooth pattern as activations are freed and recomputed.
• Are there memory leaks? Memory should return to baseline between training steps. If it grows monotonically, something is being retained.

The Memory Viewer (a separate tab) provides a more granular view, showing individual buffer allocations. This is useful for identifying which specific tensors are consuming the most memory.

Practical Memory Debugging Workflow

When you hit an OOM, follow this systematic approach:

1. Estimate theoretical memory: Calculate expected memory for weights (2 bytes/param for bf16), optimizer states (8 bytes/param for Adam), and activations (depends on batch size and sequence length)
2. Capture a profile just before the OOM: Reduce batch size or sequence length until training succeeds, then profile
3. Compare Memory Profile peak to your estimate: If the peak is much higher, the difference is overhead from tiling, temporaries, or XLA's memory allocator
4. Use Memory Viewer for the culprits: Sort by buffer size to find the largest allocations
5. Check for activation recomputation: If enabled, verify the sawtooth pattern in the Memory Profile

Memory Estimation Formulas

This 112 GB already exceeds a single H100's 80 GB HBM — before any activations. This is why sharding (ZeRO, TP, PP) is essential for models ≥ 7B.

Chapter 8: Worked Problems

Let us put everything together with some worked problems that test your ability to read profiles and reason about what is happening.

Problem 1: Mystery Profile

You see a profile with two big fusions, a reduce, and an AllReduce. One fusion has this HLO:

%fusion.1 = bf16[4096]{0:T(1024)(128)(2,1)}
  fusion(
    bf16[4096,8192]{1,0:T(8,128)(2,1)} %param.1,
    bf16[8192]{0:T(1024)(128)(2,1)} %reduce.6
  ), kind=kLoop, calls=%fused_computation.1

The final AllReduce has replica_groups={{0,16,32,48,64,80,96,112}, ...}.

This is two chained matrix multiplications with 8-way MP sharding:

def matmul(w1, w2, x):
  return jnp.einsum('wf,bf->bw', w2,
                     jnp.einsum('fw,bw->bf', w1, x))

Problem 2: Fixing Sharding

You profile a Transformer and find an unexpectedly large AllGather taking 80% of the step time. This is a sign that XLA's automatic partitioner (GSPMD/Shardy) chose a poor sharding for some intermediate tensor.

The fix: use jax.lax.with_sharding_constraint to guide the compiler:

def ffn(x, w_up, w_down):
  h = jnp.einsum('bd,df->bf', x, w_up)
  # Force h to be sharded along F, not gathered
  h = jax.lax.with_sharding_constraint(h, P('X', 'Y'))
  h = jax.nn.gelu(h)
  return jnp.einsum('bf,df->bd', h, w_down)

Problem 3: Common Performance Anti-Patterns

Here are real-world patterns you will encounter and how to diagnose them from a profile:

Key Profiling Checklist