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

2026-07-06 · Source: Google Developers Blog - AI · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Cloud Computing & IT Infrastructure · Depth: Advanced, extended

Summary

The article introduces elastic training for large language models using MaxText, Pathways, and Orbax on Cloud TPUs, demonstrating rapid recovery from hardware failures. A live experiment showed a qwen3-0.6b model training across three TPU v5e-16 slices (48 chips) on Google Kubernetes Engine recovered from a force-killed worker pod in under two minutes. This in-process recovery, facilitated by a single controller architecture, caught the failure as a "JaxRuntimeError" and restored from the last valid Orbax checkpoint (step 3300) without a full job restart, losing only 88 steps. This approach avoids the significant overhead of rescheduling the entire workload, saving hundreds of seconds compared to traditional methods.

Key takeaway

For MLOps Engineers managing large-scale distributed training on Cloud TPUs, implementing elastic training with MaxText and Pathways significantly enhances resilience. You can avoid costly full job restarts by configuring "elastic_enabled=true" and "enable_single_controller=true", allowing in-process recovery from unplanned hardware failures. For models larger than 6.7 GiB, ensure "ENABLE_PATHWAYS_PERSISTENCE=1" to prevent proxy memory issues and leverage parallel checkpoint I/O. This approach minimizes downtime and preserves training progress, making your LLM development more robust.

Key insights

Elastic training enables rapid, in-process recovery from unplanned hardware failures in distributed LLM training.

Principles

• A single controller architecture allows Python process persistence during worker failures.
• Catching hardware failures as exceptions enables in-process recovery.
• Validated checkpoints prevent loading corrupted model states.

Method

Pathways detects hardware failure as a "jax.errors.JaxRuntimeError". The "elastic_retry" decorator catches this, cleans partial state, and Orbax restores the last valid checkpoint, re-entering the training function.

In practice

• Use "enable_single_controller=True" and "elastic_enabled=true" in MaxText.
• Set "ENABLE_PATHWAYS_PERSISTENCE=1" for large models to avoid proxy OOM.
• Configure "checkpoint_period" for frequent, reliable recovery points.

Topics

• Elastic Training
• MaxText
• Pathways
• Cloud TPUs
• Distributed Training
• Kubernetes Engine
• LLM Training

Code references

• AI-Hypercomputer/pathways-utils
• AI-Hypercomputer/maxtext
• AI-Hypercomputer/xpk

Best for: Machine Learning Engineer, MLOps Engineer, AI Architect

Related on AIssential

• Boost Training Goodput: How Continuous Checkpointing Optimizes Reliability in Orbax and MaxText — Continuous checkpointing in Orbax and MaxText dynamically optimizes training reliability and performance by adapting save frequency.
• Open model, open metrics: How Lambda and the Olmo team trained Olmo Hybrid — Hybrid RNN-transformer architectures can enhance model expressivity and performance, particularly in structured reasoning tasks.
• MaxText-Slurm: Production-Grade LLM Training with Built-In Observability — MaxText-Slurm provides a unified, observable, and extensible system for production-grade LLM training on AMD GPUs.
• Decoupled DiLoCo: A new frontier for resilient, distributed AI training — Decoupled DiLoCo enables resilient, low-bandwidth, and asynchronous LLM training across distributed, heterogeneous hardware.
• Part 12 -The 80GB Wall: GPU Infrastructure and Scheduling, Worked End to End — Infrastructure decisions regarding GPU memory, scheduling, and priority tiers are critical and must be made proactively to prevent costly LLM training failures.
• Resilient Large-Scale Training: Integrating TorchFT with TorchTitan on AMD GPUs — Decoupling parallelism from fault tolerance enables resilient, checkpoint-less distributed training for large AI models.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by Google Developers Blog - AI.