--- title: We terminated a TPU mid-training and it recovered in... description: Elastic training on Cloud TPUs allows a multi-node LLM training job to recover from hardware failures without restarting the entire workload. Using Google's... canonical: https://daily.dev/posts/we-terminated-a-tpu-mid-training-and-it-recovered-in-seconds-introduction-to-elastic-training-with--0jbckhdn6 twitter:card: summary_large_image twitter:site: @dailydotdev og:type: website og:site_name: daily.dev og:title: We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText | daily.dev og:description: Elastic training on Cloud TPUs allows a multi-node LLM training job to recover from hardware failures without restarting the entire workload. Using Google's... og:url: https://daily.dev/posts/we-terminated-a-tpu-mid-training-and-it-recovered-in-seconds-introduction-to-elastic-training-with--0jbckhdn6 og:image: https://api.daily.dev/og/posts/0jBCKHDN6.png og:image:alt: We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText og:image:width: 1200 og:image:height: 630 og:locale: en --- > ## Documentation Index > Fetch the complete documentation index at: https://daily.dev/llms.txt > Use this file to discover all available pages before exploring further. # We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText **[Google Developers](https://daily.dev/sources/googledevs)** · 23 min read · 1 upvotes · 0 comments ## Summary Elastic training on Cloud TPUs allows a multi-node LLM training job to recover from hardware failures without restarting the entire workload. Using Google's JAX AI stack (MaxText + Pathways) on GKE, a single Python controller process survives worker failures and catches them as exceptions rather than process terminations. Two recovery modes are available: pause-and-resume (waits for the failed slice to be replaced, then reloads the last checkpoint) and replica-resize (continues training on surviving slices at reduced throughput). A live demo shows a force-killed TPU worker pod recovered in under 2 minutes with zero JobSet restarts, rewinding only 88 steps to the last committed Orbax checkpoint. Key components are Pathways (failure detection via heartbeat/in-flight errors), the elastic_retry decorator from pathwaysutils, and Orbax checkpointing with commit_success markers. For larger models, enabling ENABLE_PATHWAYS_PERSISTENCE=1 is critical to avoid routing checkpoint data through the proxy container. The post also previews in-memory snapshot-based recovery as a future improvement to reduce rewind distance and eliminate GCS round-trips. ## Full article daily.dev links to this article rather than hosting it. Read it at the original source: ## Questions this post answers ### How does elastic training recover from a TPU worker failure without restarting the whole job in MaxText? Pathways detects the failure via a JaxRuntimeError (either a DATA_LOSS error on an in-flight operation or a DEADLINE_EXCEEDED after roughly 10 seconds of missed heartbeats), the pathwaysutils elastic_retry decorator catches that exception around the MaxText training function, and Orbax restores the last checkpoint with a verified commit_success marker, all inside the same Python process without a full JobSet restart. _Anyone hardening a TPU training pipeline against hardware failures can track techniques like this through daily.dev._ ### What MaxText flags enable elastic training with pause and resume? Set enable_single_controller=true to route JAX through the Pathways proxy, elastic_enabled=true to wrap the training function in the elastic_retry decorator, elastic_timeout_seconds to bound how long the retry waits for a replacement slice, and elastic_max_retries to cap total tolerated failures; leaving elastic_min_slice_count at its default of -1 requires all slices before resuming (pause and resume mode). _Developers configuring resilient TPU jobs can follow config changes like these via daily.dev._ ### Why does checkpoint recovery fail with an out-of-memory error on larger models in MaxText elastic training? By default Orbax restores the checkpoint into host RAM on the CPU controller and JAX pushes it through the pathways-proxy container to the TPUs, so for a model like Qwen3-4B whose checkpoint (params plus Adam optimizer moments) totals about 135 GiB, a proxy memory limit of 100 GB causes an OOM kill. The fix is setting ENABLE_PATHWAYS_PERSISTENCE=1 so each TPU host reads and writes its own checkpoint shards directly to Cloud Storage, bypassing the controller entirely. _Teams scaling checkpoint recovery for large models can keep tabs on fixes like this on daily.dev._ ## Similar posts on daily.dev - [Amazon SageMaker HyperPod now supports checkpointless training](https://daily.dev/posts/amazon-sagemaker-hyperpod-now-supports-checkpointless-training-s7b6vf6do) · AWS · 0 upvotes · 0 comments - [Reproducing OLMo 3 7B Pre-training in MaxText: case study of large scale training on TPUs](https://daily.dev/posts/reproducing-olmo-3-7b-pre-training-in-maxtext-case-study-of-large-scale-training-on-tpus-gqxayxq2i) · Google Developers · 2 upvotes · 0 comments - [Cluster reliability for trillion parameter models on TPUs](https://daily.dev/posts/cluster-reliability-for-trillion-parameter-models-on-tpus-gtbp3jngd) · Google Cloud · 1 upvotes · 0 comments --- Tags: [#machine-learning](https://daily.dev/tags/machine-learning), [#kubernetes](https://daily.dev/tags/kubernetes) [View this post on daily.dev](https://daily.dev/posts/we-terminated-a-tpu-mid-training-and-it-recovered-in-seconds-introduction-to-elastic-training-with--0jbckhdn6) ```json {"@context":"https://schema.org","@graph":[{"@type":"Organization","@id":"https://daily.dev/#organization","name":"daily.dev","url":"https://daily.dev","logo":{"@type":"ImageObject","url":"https://daily.dev/apple-touch-icon.png","width":180,"height":180},"sameAs":["https://twitter.com/dailydotdev","https://github.com/dailydotdev","https://www.linkedin.com/company/daily-dev-ltd"]},{"@type":"WebSite","@id":"https://daily.dev/#website","url":"https://daily.dev","name":"daily.dev","publisher":{"@id":"https://daily.dev/#organization"},"potentialAction":{"@type":"SearchAction","target":{"@type":"EntryPoint","urlTemplate":"https://daily.dev/search?q={search_term_string}"},"query-input":"required name=search_term_string"}}]} {"@context":"https://schema.org","@type":"TechArticle","headline":"We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText","url":"https://daily.dev/posts/we-terminated-a-tpu-mid-training-and-it-recovered-in-seconds-introduction-to-elastic-training-with--0jbckhdn6","mainEntityOfPage":{"@type":"WebPage","@id":"https://daily.dev/posts/we-terminated-a-tpu-mid-training-and-it-recovered-in-seconds-introduction-to-elastic-training-with--0jbckhdn6"},"datePublished":"2026-07-06T16:09:54.933Z","dateModified":"2026-09-13T20:02:52.048Z","description":"Elastic training on Cloud TPUs allows a multi-node LLM training job to recover from hardware failures without restarting the entire workload. Using Google's...","image":"https://media.daily.dev/image/upload/f_auto,q_auto/v1/posts/484274848b66762531a85c95e0747f94?_a=AQAEuop","thumbnailUrl":"https://media.daily.dev/image/upload/f_auto,q_auto/v1/posts/484274848b66762531a85c95e0747f94?_a=AQAEuop","isAccessibleForFree":true,"articleSection":"Google Developers","inLanguage":"en","publisher":{"@type":"Organization","name":"daily.dev","url":"https://daily.dev","logo":{"@type":"ImageObject","url":"https://daily.dev/apple-touch-icon.png","width":180,"height":180}},"author":{"@type":"Organization","name":"Google Developers","logo":"https://media.daily.dev/image/upload/t_logo,f_auto/v1/logos/ea4bbded4e7f45ccb82e121a3156535f","url":"https://daily.dev/sources/googledevs"},"commentCount":0,"discussionUrl":"https://daily.dev/posts/we-terminated-a-tpu-mid-training-and-it-recovered-in-seconds-introduction-to-elastic-training-with--0jbckhdn6","interactionStatistic":[{"@type":"InteractionCounter","interactionType":{"@type":"LikeAction"},"userInteractionCount":1},{"@type":"InteractionCounter","interactionType":{"@type":"CommentAction"},"userInteractionCount":0}],"keywords":"machine-learning,kubernetes","timeRequired":"PT23M"} {"@context":"https://schema.org","@type":"BreadcrumbList","itemListElement":[{"@type":"ListItem","position":1,"name":"Home","item":"https://daily.dev"},{"@type":"ListItem","position":2,"name":"Google Developers","item":"https://daily.dev/sources/googledevs"},{"@type":"ListItem","position":3,"name":"We terminated a TPU mid-training and it recovered in seconds: Introduction to elastic training with MaxText"}]} {"@context":"https://schema.org","@type":"FAQPage","@id":"https://daily.dev/posts/we-terminated-a-tpu-mid-training-and-it-recovered-in-seconds-introduction-to-elastic-training-with--0jbckhdn6#faq","mainEntity":[{"@type":"Question","name":"How does elastic training recover from a TPU worker failure without restarting the whole job in MaxText?","acceptedAnswer":{"@type":"Answer","text":"Pathways detects the failure via a JaxRuntimeError (either a DATA_LOSS error on an in-flight operation or a DEADLINE_EXCEEDED after roughly 10 seconds of missed heartbeats), the pathwaysutils elastic_retry decorator catches that exception around the MaxText training function, and Orbax restores the last checkpoint with a verified commit_success marker, all inside the same Python process without a full JobSet restart. Anyone hardening a TPU training pipeline against hardware failures can track techniques like this through daily.dev."}},{"@type":"Question","name":"What MaxText flags enable elastic training with pause and resume?","acceptedAnswer":{"@type":"Answer","text":"Set enable_single_controller=true to route JAX through the Pathways proxy, elastic_enabled=true to wrap the training function in the elastic_retry decorator, elastic_timeout_seconds to bound how long the retry waits for a replacement slice, and elastic_max_retries to cap total tolerated failures; leaving elastic_min_slice_count at its default of -1 requires all slices before resuming (pause and resume mode). Developers configuring resilient TPU jobs can follow config changes like these via daily.dev."}},{"@type":"Question","name":"Why does checkpoint recovery fail with an out-of-memory error on larger models in MaxText elastic training?","acceptedAnswer":{"@type":"Answer","text":"By default Orbax restores the checkpoint into host RAM on the CPU controller and JAX pushes it through the pathways-proxy container to the TPUs, so for a model like Qwen3-4B whose checkpoint (params plus Adam optimizer moments) totals about 135 GiB, a proxy memory limit of 100 GB causes an OOM kill. The fix is setting ENABLE_PATHWAYS_PERSISTENCE=1 so each TPU host reads and writes its own checkpoint shards directly to Cloud Storage, bypassing the controller entirely. Teams scaling checkpoint recovery for large models can keep tabs on fixes like this on daily.dev."}}]} ```