JAX Pathways

Course · TPU Setup: Provisioning & Runtime VerificationPhase 19 · Study guide

Can you provision a Cloud TPU VM, verify that JAX actually executes on its chips, and prove clean teardown?

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

Revision v0.7.0 ·

Every accelerator experiment starts with knowing how many hosts and GiB of HBM your slice provides, verifying that libtpu is active inside an isolated .venv, and deleting billable VMs when finished.

Preparation: Setup & first steps

Runtime: CPU topology/contract verifier; Cloud TPU VM (v5litepod-4 / v6e-4) or Colab TPU for target runs

Start the phaseSee the project

Check your starting point

You created a Cloud TPU VM and gcloud says READY. What still has to happen before you trust a JAX run on it?

Compare your reasoning

Bootstrap ~/jax-tpu-lab/.venv with jax[tpu], verify inside Python that jax.default_backend() == 'tpu' and jax.device_count() matches your slice, and run a synchronized pmap reduction receipt.

If this is unfamiliar, work through Select a TPU slice and provision your Cloud TPU VM before attempting the project.

Build in stages

1. Select a single-host TPU topology and provision the VM

Compare v5litepod-1, v5litepod-4, v5litepod-8, and v6e-4 HBM totals, distinguish single-host VMs from multi-host Pod slices, and run gcloud create/describe/delete.

2. Bootstrap .venv and verify the runtime contract receipt

Install jax[tpu] in ~/jax-tpu-lab/.venv, run verify_runtime_contract with a synchronized pmap sum of squares, and copy back the SHA-256 receipt.

Try a changed condition

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

A script runs on a Cloud TPU VM without raising an error, but throughput is identical to a laptop CPU. What should your receipt check first?

Compare an approach

Inspect observed_backend and device_count in the JSON receipt; if libtpu was missing or JAX_PLATFORMS=tpu was not set, JAX silently fell back to CPU.

Diagnose a likely failure

Symptom: You closed your SSH terminal after finishing a run, assuming the Cloud TPU VM stopped billing.

Check next: Run gcloud compute tpus tpu-vm delete and confirm gcloud compute tpus tpu-vm list returns 0 items.

Decide what is ready

Retain the slice HBM table from tpu-01, the verified JSON runtime receipt from tpu-04, and the empty gcloud compute tpus tpu-vm list confirmation.

Where to put your project code

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

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

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

Open Inspect CPU device inventory, BF16 storage, and collective participant arithmetic

Open Operate and recover a local JAX workload

What still needs further work

Continue immediately to Phase 20 (TPU Workflows: Launching Jobs & Checkpointed Experiments) to stage code and run resumable training loops.