We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText

Summary

Distributed training on multiple TPU machines normally fails as a whole when one worker disappears: collective communication times out, processes exit, and the job restarts from a checkpoint. The article demonstrates elastic training with JAX, MaxText, PathwayS, Orbax, and Cloud TPUs on GKE, using a single CPU-side controller that remains alive while TPU workers execute compiled programs. A Pathways elastic_retry wrapper catches the failure, cleans up partial state, restores the last viable checkpoint, and re-enters the training function without relaunching the workload; the demo returned to training in under two minutes, with no JobSet restarts, after losing 88 steps. It also contrasts pause-and-resume with replica resize, explains direct checkpoint persistence through ENABLE_PATHWAYS_PERSISTENCE=1, and describes planned snapshot-based recovery from host memory.

Takeaways

1. 1

   PathwayS uses one persistent Python controller while TPU hosts run a thin worker binary, allowing a hardware failure to surface as a Python exception rather than terminating every process.

2. 2

   Pause-and-resume recovered the demo without a JobSet restart, reaching the next training step in under two minutes; 88 steps were lost because the last committed checkpoint was at step 3300.

3. 3

   Setting ENABLE_PATHWAYS_PERSISTENCE=1 makes TPU hosts read and write their own checkpoint shards directly to Cloud Storage, avoiding controller memory pressure and enabling parallel checkpoint I/O.