The training run

Every piece of this capstone already exists somewhere earlier in this path. Chapter 4 gave you the loop: forward, loss, zero the gradients, backward, step. Chapter 5 gave you the pipeline that keeps a device fed instead of starved. Chapter 6 gave you compilation, and chapter 8 gave you distribution across ranks. Chapter 9 gave you the discipline of measuring instead of feeling. Chapter 10 gave you the bridge that carries all of it onto a TPU. The capstone asks you to run all six at once, over real hardware, and that is the assignment in full.

The loop you own, on the machine this site teaches.

This is deliberate, not a shortcut taken because nothing new was left to teach. The jax path and the xla path both close the same way, with a capstone that composes rather than introduces, because composing correctly under real constraints, real shapes, a real chip, a real interruption, is a harder test than any single new concept would be.

The bar is specific enough to fail cleanly. A small model, trained end to end on a TPU, through either bridge, torch_xla or torchax, your choice. Checkpointed with its full state, not just the parameters: the model's state_dict, the optimizer's state_dict, the step count, and RNG state if a bitwise-identical resume matters to you, the same rule chapter 4 taught, aimed now at a different serialization surface. Killed on purpose partway through. Resumed from that checkpoint.

And the resume proven, not asserted: overlay the loss curve from before the kill against the curve after resume and show they are the same curve, not two curves that happen to end up close.

Until the Colab labs for this path land, the capstone ships the way the kernel path's own capstone ships: as criteria and starting points rather than a gate this site measures for you. Three things you already have are everything the assignment requires: the loop from chapter 4, the bridge chapter you just read, and a Colab notebook with a TPU attached to it. Nothing else on this site needs to exist first.

That is a genuine constraint, not an excuse. The criteria are exact: the run has to actually train on a TPU, the checkpoint has to actually hold the full state, and the resume has to actually prove itself in an overlaid curve, not a claim in a comment. A capstone that skips any one of those three is not a smaller version of this assignment; it is a different, easier assignment that happens to look similar.

Write a half-page provenance report when the run is done, the same habit this whole path has been building one measured number at a time: which chip, which dtype, the shapes you trained at, which bridge and which method you used, and the overlaid curve itself as the proof. A number without that context is not a result. It is a claim nobody else can check, and this path has spent eleven chapters teaching you not to trust those either.

The site publishes numbers with their provenance, so here is this one. The loop above ran on a Colab TPU v6e-1 through torch_xla, with jax 0.7.2 in the runtime and torch and torch_xla pinned as a matched 2.8.0 pair, two hundred steps of the chapter 4 MLP against y = x.sum(dim=1). It finished at a loss of 0.00370. The same loop on plain CPU torch 2.2.2, the machine this course was written on, finished at 0.00372. Same seed, same shapes, two different runtimes, and the answers agree to the fourth decimal.

The capstone closed on that same runtime. The run checkpointed the whole state at step 99 with a loss of 0.01456, dropped every reference it had built, rebuilt the model and the optimizer from scratch, restored from the checkpoint, and reported 0.01437 at step 100. That is the proof the bar asks for. A fresh model on this problem starts near 1.0, so a resumed first loss sitting one ordinary step below the checkpointed one means the optimizer's moments came back with the weights rather than the weights arriving alone.

A resume you have not plotted is a resume you have not proven.

One caveat the honest version keeps: this run's data is fixed, so nothing in it depended on generator state and the checkpoint carried none. A loop with augmentation or dropout has to put its generator state into the same dictionary before its resume means what this one means.