Elastic training with MaxText recovers a TPU failure in seconds

Google Cloud describes a demonstration of elastic training for largemodel workloads on Cloud TPUs using MaxText, Pathways, and Orbax. In the test, a TPU worker was forcibly terminated during training on Google Kubernetes Engine, and the job recovered without restarting the full workload. The setup used one CPU controller process and multiple TPU worker pods, with Pathways detecting the failure and raising a catchable Python exception. A MaxText retry wrapper then restored the latest valid checkpoint from Cloud Storage and resumed training after Kubernetes scheduled a replacement pod. The post also explains how this approach differs from a full restart, how checkpointing and recovery work, and what to change when scaling to larger models. It notes that direct storage access through Pathways Persistence can avoid controller memory bottlenecks during restore.