Run a batch workload with Pathways

For the purpose of this document, batch workloads are defined as JAX workloads that execute to completion and are deployed within the same GKE cluster as the Pathways cluster, specifically alongside the Pathways controller components (IFRT proxy server and Pathways resource manager). Completion of the JAX workload terminates the Pathways cluster components. This guide uses a JAX training workload to demonstrate this.

Before you begin

Make sure you have:

Build a training image using Maxtext

MaxText is an open-source, large language model (LLM) project developed by Google. It's written in JAX and designed to be highly performant and scalable, running efficiently on Google Cloud TPUs and GPUs.

To build a MaxText Docker image by using the latest version of stable JAX from the OSS GitHub repository, run the following command:

git clone https://github.com/AI-Hypercomputer/maxtext
cd maxtext
gcloud config set project PROJECT_ID
bash src/dependencies/scripts/docker_build_dependency_image.sh MODE=stable
gcloud auth configure-docker
# This script needs bash version >= 4.2 to execute.
bash src/dependencies/scripts/docker_upload_runner.sh \
    CLOUD_IMAGE_NAME=USER_runner

This command pushes the MaxText Kubernetes image to gcr.io/$PROJECT_ID/${USER}_runner. You can use this Docker image to run training on TPUs by using the Pathways backend.

Run a batch workload with Cluster Toolkit

If the gcluster binary isn't in your PATH environment variable, add the directory containing gcluster to your PATH:

export PATH="$PATH:CLUSTER_TOOLKIT_PATH"

Replace CLUSTER_TOOLKIT_PATH with the absolute path to the directory containing the gcluster binary, such as $HOME/cluster-toolkit.

Submit the prebuilt MaxText Docker image by using the gcluster job submit command:

gcluster job submit \
    --pathways \
    --pathways-gcs-location="gs://BUCKET_NAME/pathways-artifacts" \
    --name=WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION \
    --num-slices=WORKLOAD_NODEPOOL_COUNT \
    --compute-type=COMPUTE_TYPE \
    --topology=TOPOLOGY \
    --image="gcr.io/PROJECT_ID/USER_runner" \
    --command="python3 -m maxtext.trainers.pre_train.train /deps/src/maxtext/configs/base.yml base_output_directory=gs://BUCKET_NAME per_device_batch_size=1 enable_checkpointing=false remat_policy=full global_parameter_scale=1 steps=20 max_target_length=2048 use_iota_embed=true reuse_example_batch=1 dataset_type=synthetic attention=flash gcs_metrics=True enable_single_controller=True run_name=RUN_NAME-pathways-job"

For more information about job submission options, see the Cluster Toolkit Job Submission Guide.

Replace the following:

  • WORKLOAD: a unique name to identify your workload; because of DNS label limits, this must be 22 characters or fewer
  • CLUSTER: the name of your GKE cluster
  • WORKLOAD_NODEPOOL_COUNT: the number of TPU slice node pools
  • COMPUTE_TYPE: the TPU machine type (for example, ct6e-standard-4t). For more information about supported TPU types for each TPU version, see TPU versions.
  • TOPOLOGY: the TPU placement topology (for example, 2x4)
  • PROJECT_ID: your Google Cloud project ID
  • LOCATION: the region or zone of your GKE cluster
  • USER: your Google Cloud user ID
  • BUCKET_NAME: the Cloud Storage bucket for outputs
  • RUN_NAME: a user-assigned name to identify the workflow run

Follow the progress of your workload by using the gcluster job logs command:

gcluster job logs WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION \
    --main-only \
    --follow

completed step: 0, seconds: 1.259, TFLOP/s/device: 10.792, Tokens/s/device: 1626.636, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 1, seconds: 3.366, TFLOP/s/device: 4.036, Tokens/s/device: 608.366, total_weights: 65536, loss: 10.876, lm_loss: 10.876, perplexity: 52906.090
completed step: 2, seconds: 0.264, TFLOP/s/device: 51.415, Tokens/s/device: 7749.885, total_weights: 65536, loss: 10.563, lm_loss: 10.563, perplexity: 38665.316
completed step: 3, seconds: 1.025, TFLOP/s/device: 13.257, Tokens/s/device: 1998.187, total_weights: 65536, loss: 9.983, lm_loss: 9.983, perplexity: 21646.125
completed step: 4, seconds: 1.194, TFLOP/s/device: 11.376, Tokens/s/device: 1714.750, total_weights: 65536, loss: 9.462, lm_loss: 9.462, perplexity: 12864.218
completed step: 5, seconds: 1.185, TFLOP/s/device: 11.465, Tokens/s/device: 1728.203, total_weights: 65536, loss: 8.961, lm_loss: 8.961, perplexity: 7791.373

To cancel the workload before it completes, use the gcluster job cancel command:

gcluster job cancel WORKLOAD --cluster=CLUSTER --project=PROJECT_ID --location=LOCATION

What's next