Skip to content

README.md

Latest commit

 

History

History
171 lines (146 loc) · 6.99 KB

README.md

File metadata and controls

171 lines (146 loc) · 6.99 KB
Raw

Elastic Training with Pathways

This document demonstrates how to leverage the elasticity primitives within pathwaysutils.elastic to create a resilient JAX training loop that can handle hardware failures gracefully. We illustrate this using an example based on the MaxText training loop running on TPUs provisioned by GKE via PathwaysJob API.

Overview

Distributed training jobs, especially long-running ones, are susceptible to various failures, such as machine preemptions and hardware issues. Elasticity allows a training job to adapt to changes in the number of available accelerators without crashing. It typically involves:

  1. Training State Management: Regularly snapshotting the training state (model params, optimizer state, data iterator state).
  2. Failure Detection: Pathways Resource Manager detects when workers join or leave.
  3. Failure Propogation: Pathways runtime propagates the error to JAX client.
  4. Training Reconfiguration: Adapting the training computation distribution to the current set of healthy workers.
  5. Resumption: Continuing training from the last valid snapshot with the new configuration.

The pathwaysutils.elastic primitives provide elastcity building blocks to use within your JAX training loop when using the Pathways' Proxy JAX backend.

Prerequisites

  • A Pathways compatible GKE cluster with TPU and CPU nodepools.
  • kubectl configured to interact with your cluster.
  • Access to a container image containing JAX, your model code (e.g., MaxText), and the pathwaysutils package with elasticity features integrated.

Elastic MaxText Training with Pathways on GKE

This example demonstrates running an elastic MaxText job on 3 x v5e-32 slices using Pathways. See the PathwaysJob docs for more details about the various attributes set in the YAML below.

1. Elastic PathwaysJob Definition (pathwaysjob-elastic.py)

Please set the variables marked with <> below before executing the script.

apiVersion: pathways-job.pathways.domain/v1
kind: PathwaysJob
metadata:
  name: pathways-<USER>
spec:
  maxRestarts: 0
  workers:
  - type: ct5lp-hightpu-4t
    topology: 4x8
    numSlices: 3
    maxSliceRestarts: 2
  pathwaysDir: "gs://<BUCKET>" # Pre-create this bucket.
  controller:
    deploymentMode: default
    elasticSlices: 1
    template:
      spec:
        containers:
        - name: main
          image: <MAXTEXT_IMAGE>
          imagePullPolicy: Always
          command:
          - bash
          - -c
          - >
            python3 -m MaxText.elastic_train MaxText/configs/base.yml
            base_output_directory=gs://<BUCKET>
            per_device_batch_size=4
            enable_checkpointing=false
            remat_policy=full
            global_parameter_scale=8
            steps=50
            max_target_length=2048
            use_iota_embed=true
            reuse_example_batch=1
            dataset_type=synthetic
            attention=flash
            gcs_metrics=True
            enable_pathways_goodput=True
            run_name=pathways-<USER>

The MaxText elastic training script invoked by the main container above is integrated with pathwaysutils.elastic primitives.

2. Running the Elastic Training Loop and Simulating hardware failures

The following bash script demonstrates launching the above elastic maxtext job with Pathways, monitoring its progress, simulating a hardware failure by issuing a kubectl drain to a randomly selected TPU node, and observing the recovery. Please set the variables marked as <> below before executing the script. At the end of the script, we verify elasticity worked as expected.

#!/bin/bash
WORKING_DIR=</LOCAL/DIRECTORY/PATH>
USER_LABEL_SELECTOR="<USER>"
LOG_DIR="${WORKING_DIR}/logs"
RUN_ID=pathways-${USER_LABEL_SELECTOR}
LOG_FILE="${LOG_DIR}/logs_${RUN_ID}.log"
JOB_DEFINITION_FILE="${WORKING_DIR}/pathwaysjob-elastic.yaml" # Copy the above yaml into this file

mkdir -p ${LOG_DIR}

echo "Running Elastic MaxText with Run ID: ${RUN_ID}"

# 1. Launch the PathwaysJob
kubectl apply -f "$JOB_DEFINITION_FILE"

# 2. Monitor the PathwaysJob
echo "Waiting for pods to start..."
head_pod=""
for i in $(seq 1 10)
do
  head_pod=$(kubectl get pods -o=name --field-selector='status.phase==Running' | grep "$USER_LABEL_SELECTOR" | grep 'head' | head -n 1)
  if [ -n "$head_pod" ]; then
    echo "Found head pod: $head_pod"
    break
  fi
  echo "Head pod not found yet, retrying..."
  sleep 10s
done

if [ -z "$head_pod" ]; then
  echo "Error: Could not find running head pod after multiple attempts. Cleaning up..." 1>&2
  kubectl delete -f "$JOB_DEFINITION_FILE"
  exit 1
fi

echo "Streaming logs from $head_pod to ${LOG_FILE}"
kubectl logs -f "$head_pod" >> "${LOG_FILE}" &
logs_pid=$!
echo "Waiting for job to start making progress..."
sleep 90s

# 3. Simulate Failure: Evict a Worker Pod
echo "Randomly select a worker pod to disrupt..."
read -r node_name pod_name <<<$(kubectl get pods -o wide --field-selector='status.phase==Running' | grep "$USER_LABEL_SELECTOR" | grep worker | shuf | head -n 1 | awk '{print $7, $1}')

if [ -z "$pod_name" ] || [ -z "$node_name" ]; then
  echo "Warning: Could not find a running worker pod to disrupt. Skipping disruption."
else
  echo "Attempting to cordon '$node_name' and kill pod '$pod_name'..."
  kubectl cordon "$node_name"
  kubectl exec -it "$pod_name" -c pathways-worker -- /bin/sh -c "kill -s SIGILL 1"
  echo "Node cordoned. Waiting briefly for training to reconfigure to N-1 slices..."
  sleep 90s

  # 4. Allow Recovery: Uncordon the Node
  echo "Uncordoning node '$node_name' to allow scheduling again."
  kubectl uncordon "$node_name"
fi

# 5. Wait for Training to resume on all slices
sleep 90s

# 6. Terminate the Job and Cleanup
echo "Terminating Run ID ${RUN_ID}"
kubectl delete -f "$JOB_DEFINITION_FILE"
# Ensure log streaming process is killed
kill "$logs_pid" 2>/dev/null 
echo "Completed Run ID ${RUN_ID}."

# 6. Verify by printing steps where training reconfigured from N to N-1 slices and later back to N slices
# Expect output like:
# Step: 5, Old Slice Count: 3, New Slice Count: 2 (3 -> 2 slices)
# Step: 17, Old Slice Count: 2, New Slice Count: 3 (2 -> 3 slices)
awk '
  /step=/ && /elastic_manager\.elastic_down_event_count=/ {
    split($0, fields, " ")
    step = ""
    good_slice_count = ""
    for (i in fields) {
      split(fields[i], kv, "=")
      if (kv[1] == "step") {
        step = kv[2]
      } else if (kv[1] == "elastic_manager.good_slice_count") {
        good_slice_count = kv[2]
      }
    }
    if (prev_good_slice_count != "" && prev_good_slice_count != good_slice_count) {
      print "Step: " step ", Old Slice Count: " prev_good_slice_count ", New Slice Count: " good_slice_count
    }
    prev_step = step
    prev_good_slice_count = good_slice_count
  }
' "${LOG_FILE}"