JAX Pathways

Course · TPU Systems: Generations, Memory, Precision & XProfPhase 21 · Study guide

How do you size TPU HBM capacity, bound BF16/INT8 precision error, and capture a clean XProf trace?

Cloud TPU Mini-Course · Step 3 of 3Cloud TPU VM + CPU verifierBite-sized · 2 lessons (~50m)

Revision v0.7.0 ·

Provisioning a TPU slice without budgeting HBM wastes time on out-of-memory crashes, and profiling without warming up measures compiler overhead instead of MXU execution.

Preparation: TPU Workflows: Launching Jobs & Checkpointed Experiments

Runtime: CPU numerical & profiler harness; Cloud TPU VM for target XProf traces

Start the phaseSee the project

Check your starting point

A matrix multiplication has arithmetic intensity 170 FLOP/byte on TPU v5e (ridge point 240.5 FLOP/byte). Is it compute-bound or HBM bandwidth-bound?

Compare your reasoning

HBM bandwidth-bound, because its arithmetic intensity (170 FLOP/byte) is below the TPU v5e ridge point (240.5 FLOP/byte). Increasing batch size or sequence length raises operational intensity toward the MXU ceiling.

If this is unfamiliar, work through TPU generations and HBM memory budgeting before attempting the project.

Build in stages

1. Compare TPU generations and budget HBM capacity

Compute TPU v5e, v5p, and v6e ridge points, calculate FP32 vs mixed-BF16 training state and KV-cache GiB, and determine the minimum chip count at an 80% HBM budget.

2. Audit BF16/INT8 precision and capture a warmed XProf trace

Compare BF16 and symmetric INT8 matmul errors against FP32, verify FP32 accumulation inside jax.lax.dot_general, and capture a warmed *.xplane.pb trace.

Try a changed condition

Make a prediction before opening the reasoning. Keep your answer and the evidence you would collect.

An XProf trace on your TPU VM shows a 12-second host compile block followed by five 9-millisecond step blocks. How do you get a clean steady-state profile?

Compare an approach

Run at least one warmup call and call jax.block_until_ready() before entering with jax.profiler.trace(trace_dir): so only steady-state execution is recorded.

Diagnose a likely failure

Symptom: Accumulating a long bfloat16 dot product directly in bfloat16 produces much larger drift than expected.

Check next: Pass preferred_element_type=jnp.float32 to jax.lax.dot_general so the TPU MXU accumulates products in FP32 before any output cast.

Decide what is ready

Retain the TPU ridge-point and HBM budget tables from tpu-03, plus the BF16/INT8 error summary and *.xplane.pb trace receipt from tpu-06.

Where to put your project code

Write your code in: projects/tpu-systems-audit/starter/model.py (copy projects/tpu-systems-audit/starter/model.py first: cp projects/tpu-systems-audit/starter/model.py projects/tpu-systems-audit/starter/model.py).

Verify your implementation: python3 projects/tpu-systems-audit/tests/check.py --stage all --impl projects/tpu-systems-audit/starter/model.py

Run the staged project checks on your implementation in projects/tpu-systems-audit/starter/model.py and record your transfer experiment results in your portfolio.

Open Estimate optimizer-state bytes and calculate supplied-rate scaling efficiency

Open Operate and recover a local JAX workload

What still needs further work

Apply the 3-part Cloud TPU mini-course across Phases 07, 08, 09, 13, 15, 16, 17, and 18 by staging each lesson script in exercises/ to your TPU VM.