Skip to content

Commit 8189d41

Browse files
inardiniinardini
committed
Add elastic training (Pathways + xpk) guide under docs/
Add docs/run_maxtext/run_maxtext_elastic_training.md and register it in docs/run_maxtext.md. The guide launches an elastic Qwen3 0.6B run via xpk workload create-pathways (--elastic-slices / --max-slice-restarts), assumes an existing Pathways cluster. Commands validated end to end on 48 v5e chips. Supersedes AI-Hypercomputer#4043.
1 parent df1b359 commit 8189d41

2 files changed

Lines changed: 199 additions & 0 deletions

File tree

‎docs/run_maxtext.md‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,13 @@ Run large-scale JAX jobs on TPUs using Pathways. Supports batch and headless (in
3939

4040
Run tests and local development without Google Cloud dependencies (no `gcloud`, GCS, or Vertex AI required).
4141
:::
42+
43+
:::{grid-item-card} ♻️ Elastic training (demo)
44+
:link: run_maxtext/run_maxtext_elastic_training
45+
:link-type: doc
46+
47+
Demonstrate fault-tolerant training with Pathways on GKE: lose a TPU slice mid-run and recover in-process from the last checkpoint, no job restart.
48+
:::
4249
::::
4350

4451
```{toctree}
@@ -51,4 +58,5 @@ run_maxtext/run_maxtext_single_host_gpu.md
5158
run_maxtext/run_maxtext_via_xpk.md
5259
run_maxtext/run_maxtext_via_pathways.md
5360
run_maxtext/decoupled_mode.md
61+
run_maxtext/run_maxtext_elastic_training.md
5462
```

‎docs/run_maxtext/run_maxtext_elastic_training.md‎

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
1+
<!--
2+
Copyright 2025 Google LLC
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
https://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
-->
16+
17+
(run-elastic-training)=
18+
19+
# Elastic training with Pathways
20+
21+
This guide shows how to run **elastic training** on a multi-slice TPU cluster: training that survives a slice failure *in-process*, without restarting the job. You launch a Qwen3 0.6B run across several TPU slices with Pathways, lose a slice mid-run, and watch training recover from the last checkpoint on the same controller.
22+
23+
```{important}
24+
This guide is a **demonstration of the elastic training mechanism**, not a production recipe. It uses a small model (Qwen3 0.6B) and synthetic data so you can see recovery happen on a short run, then tear everything down. The exact slice counts, timeouts, and checkpoint cadence here are illustrative; tune them for your own model and hardware. Treat it as a starting point to understand the feature, not a configuration to copy verbatim into a long-running job.
25+
```
26+
27+
## What is elastic training?
28+
29+
Large model training runs across many TPU slices. When one slice fails (a hardware fault, a preemption, a network blip), the default outcome is that the whole job crashes and restarts from scratch, losing the XLA compilation time plus everything since the last checkpoint.
30+
31+
Elastic training keeps the training process alive instead. Three components make that possible:
32+
33+
- **Pathways** orchestrates training across the slices. Its Resource Manager detects when a slice goes down and reports it to the training process.
34+
- **MaxText** wraps the training loop with `elastic_retry`. When Pathways reports a failure, it catches the exception *inside the same Python process*, cleans up, and restarts training without exiting.
35+
- **Orbax** handles checkpointing. Each checkpoint writes to GCS and creates a `commit_success` marker only after all data is flushed, so a checkpoint interrupted mid-write has no marker and is safely discarded on recovery.
36+
37+
Because the controller process never exits, the expensive XLA recompile is skipped and recovery is fast.
38+
39+
```{note}
40+
This demo shows recovery via *checkpoint restore* on a fixed mesh: when a slice is lost, Pathways waits for a replacement, then all slices restore from the last committed checkpoint. It does **not** show elastic *degradation* (continuing on fewer slices at reduced throughput), which requires dynamic mesh resize and is not covered here.
41+
```
42+
43+
## 1. Prerequisites
44+
45+
This guide assumes you already have a **Pathways-enabled GKE cluster** created with `xpk`, and a MaxText Docker image in your Artifact Registry. If you don't:
46+
47+
1. **Install XPK and create a Pathways GKE cluster.** Follow [Running MaxText with XPK](run_maxtext_via_xpk.md) and the [Pathways & XPK cluster guide](https://cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster#xpk). Cluster creation and management is out of scope for this page.
48+
2. **Build and upload the MaxText Docker image.** See [Build MaxText](../build_maxtext.md).
49+
50+
```{note}
51+
If you installed `xpk` inside a Python virtual environment (`venv`), reactivate it (e.g., `source <VENV_NAME>/bin/activate`) in any new terminal before running `xpk` commands, or you will hit a `Command xpk not found` error.
52+
```
53+
54+
## 2. Environment configuration
55+
56+
Set these environment variables in your shell. Replace the placeholders with your own values.
57+
58+
```bash
59+
# -- Google Cloud Configuration --
60+
export PROJECT_ID=<GCP project ID>
61+
export ZONE=<GCP location> # e.g., 'us-central1'
62+
export GKE_CLUSTER=<cluster name> # your Pathways-enabled cluster
63+
64+
# -- Workload Configuration --
65+
# Kubernetes requires workload names to be valid DNS labels (lowercase, no underscores/periods).
66+
export RUN_NAME="elastic-qwen3-$(date +%Y%m%d-%H%M%S)"
67+
68+
# TPU type and slice count. For supported types see src/maxtext/utils/accelerator_to_spec_map.py.
69+
export TPU_TYPE="v5litepod-16" # one slice = 16 v5e chips
70+
export NUM_SLICES=3 # total slices in the run
71+
72+
# -- MaxText & Storage Configuration --
73+
export BASE_OUTPUT_DIRECTORY=<gcs bucket path> # e.g., gs://my-bucket/maxtext-runs
74+
export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/<your maxtext image>"
75+
```
76+
77+
## 3. Launch the elastic workload
78+
79+
Submit the run with `xpk workload create-pathways`. Two sets of flags make it elastic:
80+
81+
- **On the `xpk` side**, `--elastic-slices` tells Pathways how many slices the workload is allowed to lose and keep going, and `--max-slice-restarts` caps how many times a slice's workers may be restarted.
82+
- **On the MaxText side** (inside `--command`), `elastic_enabled=true` turns on the `elastic_retry` wrapper, and `enable_single_controller=True` runs training through Pathways. `checkpoint_period` is kept small so a recovery rewinds only a little.
83+
84+
```bash
85+
xpk workload create-pathways \
86+
--workload=${RUN_NAME?} \
87+
--cluster=${GKE_CLUSTER?} \
88+
--project=${PROJECT_ID?} \
89+
--zone=${ZONE?} \
90+
--tpu-type=${TPU_TYPE?} \
91+
--num-slices=${NUM_SLICES?} \
92+
--docker-image=${DOCKER_IMAGE?} \
93+
--elastic-slices=1 \
94+
--max-slice-restarts=10 \
95+
--command="python3 -m maxtext.trainers.pre_train.train \
96+
src/maxtext/configs/base.yml \
97+
base_output_directory=${BASE_OUTPUT_DIRECTORY?} \
98+
run_name=${RUN_NAME?} \
99+
model_name=qwen3-0.6b \
100+
dataset_type=synthetic \
101+
per_device_batch_size=1 \
102+
max_target_length=2048 \
103+
attention=flash \
104+
remat_policy=full \
105+
steps=5000 \
106+
enable_checkpointing=true \
107+
checkpoint_period=100 \
108+
enable_single_controller=True \
109+
elastic_enabled=true \
110+
elastic_timeout_seconds=300 \
111+
elastic_max_retries=10"
112+
```
113+
114+
```{note}
115+
`--elastic-slices=1` means the run tolerates losing **one** slice at a time out of `${NUM_SLICES}`. Keep `--max-slice-restarts` and the MaxText `elastic_max_retries` consistent with how many failures you want to ride out.
116+
```
117+
118+
```{warning}
119+
**Do not enable profiling in an elastic run.** An elastic event (a slice going down and recovering) in the middle of a profile is not supported, so this example leaves the profiler off (`profiler` is unset). Profile a separate, non-elastic run if you need performance traces.
120+
```
121+
122+
### Watch training start
123+
124+
List the workload and follow its logs through the Cloud Console (**Kubernetes Engine → Workloads →** your run **→ Logs**), or:
125+
126+
```bash
127+
xpk workload list --cluster=${GKE_CLUSTER?} --project=${PROJECT_ID?} --zone=${ZONE?}
128+
```
129+
130+
After XLA compilation (a couple of minutes) you should see elastic training enabled and a steady stream of steps:
131+
132+
```
133+
Elastic utils: Elastic training enabled.
134+
Elastic Retry Enabled
135+
completed step: 8, seconds: 0.159, TFLOP/s/device: 43.430, loss: 220.774
136+
completed step: 9, seconds: 0.166, TFLOP/s/device: 41.524, loss: 217.296
137+
```
138+
139+
Let it run until the step counter passes the first checkpoint (here, step ~130, so `checkpoint_period=100` has committed once) before you inject a failure, so there is a complete checkpoint to recover from.
140+
141+
## 4. Simulate a slice failure
142+
143+
To see recovery, remove a worker on one slice. Connect to the cluster and delete a worker pod immediately (`--grace-period=0 --force`), so it does not drain gracefully. This mimics an abrupt hardware failure rather than a clean shutdown:
144+
145+
```bash
146+
gcloud container clusters get-credentials ${GKE_CLUSTER?} --location ${ZONE?} --project ${PROJECT_ID?}
147+
148+
# Pick a worker pod on one slice and remove it immediately.
149+
WORKER=$(kubectl get pods -o name | grep "${RUN_NAME?}" | grep worker | head -1)
150+
kubectl delete ${WORKER?} --grace-period=0 --force
151+
```
152+
153+
```{warning}
154+
This deliberate pod deletion is only for observing recovery in this demo. Do not remove pods this way against a real training job.
155+
```
156+
157+
## 5. Verify in-process recovery
158+
159+
Recovery shows up in the **same controller log** you were already watching, which is the point: the controller process never exited. Within seconds of the termination you should see Pathways report the slice down and `elastic_retry` restore the last committed checkpoint:
160+
161+
```
162+
Slice down event detected. Retrying.
163+
Found commit_success file. Keeping gs://.../checkpoints/100/.
164+
Elastic attempt 2 out of 10
165+
Restoring checkpoint from gs://.../checkpoints/100.
166+
completed step: 101, ...
167+
```
168+
169+
The step counter dropping (for example `150 -> 101`) is the rewind to the last committed checkpoint. Training then continues from there on the same controller, with no JobSet restart. That is the whole point of elastic training: a slice failure became a short rewind instead of a full job restart.
170+
171+
## 6. Clean up
172+
173+
Delete the workload to stop the meter. TPU slices are expensive, so don't skip this.
174+
175+
```bash
176+
xpk workload delete --workload=${RUN_NAME?} --cluster=${GKE_CLUSTER?} --project=${PROJECT_ID?} --zone=${ZONE?}
177+
```
178+
179+
If you created the cluster only for this demo, delete it too (see the [XPK documentation](https://github.com/AI-Hypercomputer/xpk) for `xpk cluster delete`).
180+
181+
## Going further
182+
183+
- **The elastic flags** are documented in `src/maxtext/configs/base.yml`: `elastic_enabled`, `elastic_timeout_seconds`, `elastic_max_retries`, plus `enable_single_controller` (runs training through Pathways) and `checkpoint_period`.
184+
- **A larger model** changes the checkpoint size that streams through Pathways during recovery; size the controller and adjust `checkpoint_period` accordingly.
185+
- **Custom Pathways server args** can be passed through `xpk` with `--custom-pathways-proxy-server-args` if you need finer control than `--elastic-slices` exposes.
186+
187+
## More information
188+
189+
- [Running MaxText with XPK](run_maxtext_via_xpk.md)
190+
- [Running MaxText via Pathways](run_maxtext_via_pathways.md)
191+
- [Pathways on Cloud documentation](https://cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro)

0 commit comments

Comments
 (0)