JAX Pathways

Course · TPU Workflows: Launching Jobs & Checkpointed ExperimentsPhase 20 · Study guide

Can you launch a JAX job over SSH, separate XLA warmup from steady-state timing, and resume cleanly after interruption?

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

Revision v0.7.0 ·

Cloud TPU VMs can be preempted and multi-host Pod slices require lockstep execution across all workers. A reliable workflow saves atomic checkpoints, logs JSONL events, and copies artifacts back before teardown.

Preparation: TPU Setup: Provisioning & Runtime Verification

Runtime: CPU launcher rehearsal; Cloud TPU VM (single-host or --worker=all Pod slice) for target execution

Start the phaseSee the project

Check your starting point

Step 1 of your TPU job takes 14 seconds while steps 2–50 take 11 milliseconds each. How should your launcher record timing?

Compare your reasoning

Run and synchronize one warmup call before the timed loop and log it as a separate warmup event so XLA compilation does not pollute steady-state step latency.

If this is unfamiliar, work through Stage and launch checkpointed TPU jobs before attempting the project.

Build in stages

1. Launch resumable TPU jobs with warmup separation

Run run_tpu_ready_job, log warmup separately from timed steps, and verify that a checkpoint-resumed run matches an uninterrupted control run bit-for-bit.

2. Stage course scripts and launch multi-host Pod slices

Build scp-in, ssh-run, and scp-out commands for any course script, use --worker=all on multi-host slices, and gate host artifact writes on jax.process_index() == 0.

Try a changed condition

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

You move a script from a single-host v5litepod-4 VM to a 4-host v5litepod-16 Pod slice, and it hangs at startup. What is the first thing to check?

Compare an approach

Check that you launched the script across all hosts simultaneously using gcloud compute tpus tpu-vm ssh --worker=all and called jax.distributed.initialize() before querying jax.devices().

Diagnose a likely failure

Symptom: A resumed run loads the saved model weights w, but its next loss diverges from the uninterrupted control run.

Check next: Verify that the checkpoint also restored optimizer momentum v and step counter step so the deterministic batch schedule continues at the exact next step.

Decide what is ready

Retain the resumed vs control events.jsonl and state hashes from tpu-02, plus the single-host and --worker=all command workflows from tpu-05.

Where to put your project code

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

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

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

Open Gate simulated multi-worker checkpoint publication on complete receipts

Open Operate and recover a local JAX workload

What still needs further work

Continue to Phase 21 (TPU Systems: Generations, Memory, Precision & XProf) to budget HBM, test BF16/INT8 precision, and capture XProf traces.