disaggregated llama3.1-70b w pathways

Author: A9isha

docker_build_dependency_image.sh

@@ -27,9 +27,21 @@
 # works with any custom wheels.
 # bash docker_build_dependency_image.sh MODE=custom_wheels
 
+# bash docker_build_dependency_image.sh MODE=grpo
+
 # Enable "exit immediately if any command fails" option
 set -e
 
+# Check for docker permissions
+if ! docker info > /dev/null 2>&1; then
+  echo "ERROR: Permission denied while trying to connect to the Docker daemon." >&2
+  echo "You can fix this by:" >&2
+  echo "1. Running this script with sudo: 'sudo bash $0 $@'" >&2
+  echo "2. Adding your user to the 'docker' group: 'sudo usermod -aG docker \${USER}' (requires a new login session)." >&2
+  echo "3. Running `newgrp docker` in your current terminal." >&2
+  exit 1
+fi
+
 export LOCAL_IMAGE_NAME=maxtext_base_image
 echo "Building to $LOCAL_IMAGE_NAME"
 
@@ -40,9 +52,10 @@ echo "Starting to build your docker image. This will take a few minutes but the
 
 # Set environment variables
 for ARGUMENT in "$@"; do
-    IFS='=' read -r KEY VALUE <<< "$ARGUMENT"
+    IFS='=' read -r RAW_KEY VALUE <<< "$ARGUMENT"
+    KEY=$(echo "$RAW_KEY" | tr '[:lower:]' '[:upper:]')
     export "$KEY"="$VALUE"
-    echo "$KEY"="$VALUE"
+    echo "$KEY=$VALUE"
 done
 
 
@@ -54,11 +67,18 @@ fi
 if [[ -z ${MODE} ]]; then
   export MODE=stable
   echo "Default MODE=${MODE}"
+  export CUSTOM_JAX=0
+  export INSTALL_GRPO=0
 elif [[ ${MODE} == "custom_wheels" ]] ; then
   export MODE=nightly
   export CUSTOM_JAX=1
+  export INSTALL_GRPO=0
+elif [[ ${MODE} == "grpo" || ${MODE} == "grpo-experimental" ]] ; then
+  export INSTALL_GRPO=1
+  export CUSTOM_JAX=0
 else
   export CUSTOM_JAX=0
+  export INSTALL_GRPO=0
 fi
 
 if [[ -z ${DEVICE} ]]; then
@@ -104,6 +124,9 @@ if [[ -z ${LIBTPU_GCS_PATH+x} ]] ; then
     elif [[ ${MANTARAY} == "true" ]]; then
       echo "Building with benchmark-db"
       docker build --network host --build-arg MODE=${MODE} --build-arg JAX_VERSION=$JAX_VERSION --build-arg LIBTPU_GCS_PATH=$LIBTPU_GCS_PATH --build-arg DEVICE=$DEVICE -f ./maxtext_db_dependencies.Dockerfile -t ${LOCAL_IMAGE_NAME} .
+    elif [[ ${INSTALL_GRPO} -eq 1 && ${DEVICE} == "tpu" ]]; then
+      echo "Installing MaxText stable mode dependencies for GRPO"
+      docker build --network host --build-arg MODE=stable --build-arg JAX_VERSION=$JAX_VERSION --build-arg LIBTPU_GCS_PATH=$LIBTPU_GCS_PATH --build-arg DEVICE=$DEVICE -f ./maxtext_dependencies.Dockerfile -t ${LOCAL_IMAGE_NAME} .
     else
       docker build --network host --build-arg MODE=${MODE} --build-arg JAX_VERSION=$JAX_VERSION --build-arg LIBTPU_GCS_PATH=$LIBTPU_GCS_PATH --build-arg DEVICE=$DEVICE -f ./maxtext_dependencies.Dockerfile -t ${LOCAL_IMAGE_NAME} .
     fi
@@ -113,6 +136,32 @@ else
   docker build --network host --build-arg CUSTOM_LIBTPU=true -f ./maxtext_libtpu_path.Dockerfile -t ${LOCAL_IMAGE_NAME} .
 fi
 
+if [[ ${INSTALL_GRPO} -eq 1 ]] ; then
+  if [[ ${DEVICE} != "tpu" ]] ; then
+    echo "Error: MODE=grpo is only supported for DEVICE=tpu"
+    exit 1
+  fi
+
+  # # To install tpu_commons from a local path, we copy it into the build context, excluding __pycache__.
+  # # This assumes vllm, tunix, tpu_commons is a sibling directory to the current one (maxtext).
+  # rsync -a --exclude='__pycache__' ../tpu_commons .
+  # # To install vllm from a local path, we copy it into the build context, excluding __pycache__.
+  # # This assumes vllm is a sibling directory to the current one (maxtext).
+  # rsync -a --exclude='__pycache__' ../vllm .
+
+  # rsync -a --exclude='__pycache__' ../tunix .
+
+  # # The cleanup is set to run even if the build fails to remove the copied directory.
+  # trap "rm -rf ./tpu_commons ./vllm ./tunix" EXIT INT TERM
+
+  docker build \
+    --network host \
+    --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} \
+    --build-arg MODE=${MODE} \
+    -f ./maxtext_grpo_dependencies.Dockerfile \
+    -t ${LOCAL_IMAGE_NAME} .
+fi
+
 if [[ ${CUSTOM_JAX} -eq 1 ]] ; then
   echo "Installing custom jax and jaxlib"
   docker build --network host -f ./maxtext_custom_wheels.Dockerfile -t ${LOCAL_IMAGE_NAME} .

docker_upload_runner.sh

@@ -30,9 +30,10 @@ export PROJECT=$(gcloud config get-value project)
 
 # Set environment variables
 for ARGUMENT in "$@"; do
-    IFS='=' read -r KEY VALUE <<< "$ARGUMENT"
+    IFS='=' read -r RAW_KEY VALUE <<< "$ARGUMENT"
+    KEY=$(echo "$RAW_KEY" | tr '[:lower:]' '[:upper:]')
     export "$KEY"="$VALUE"
-    echo "$KEY"="$VALUE"
+    echo "$KEY=$VALUE"
 done
 
 export LOCAL_IMAGE_NAME_RUNNER=${LOCAL_IMAGE_NAME}__runner
@@ -44,9 +45,36 @@ if [[ ! -v CLOUD_IMAGE_NAME ]]; then
   exit 1
 fi
 
+# Check for dangling symbolic links (target does not exist).
+DANGLING_LINKS=$(find -L . -type l)
+if [ -n "$DANGLING_LINKS" ]; then
+  echo "ERROR: Found dangling symbolic links in the build context:"
+  echo "$DANGLING_LINKS"
+  echo "These can cause 'failed to compute cache key' errors during 'docker build'."
+  echo "Please remove or fix them before building the Docker image."
+  exit 1
+fi
+
+# Check for absolute symbolic links, which Docker can't follow outside the build context.
+ABSOLUTE_LINKS=$(find . -type l -lname '/*')
+if [ -n "$ABSOLUTE_LINKS" ]; then
+  echo "ERROR: Found symbolic links with absolute paths in the build context:"
+  echo "$ABSOLUTE_LINKS"
+  echo "Docker cannot follow absolute paths outside of the build context, which can cause 'failed to compute cache key' errors."
+  echo "Please remove these links or convert them to relative paths before building the Docker image."
+  exit 1
+fi
+
 # Download other test assets from GCS into "${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/MaxText}"/test_assets
-if ! gcloud storage cp gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT:-${MAXTEXT_REPO_ROOT:-$PWD}/test_assets}"; then
-  echo "WARNING: Failed to download test assets from GCS. These files are only used for end-to-end tests; you may not have access to the bucket."
+# if ! gcloud storage cp gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT:-${MAXTEXT_REPO_ROOT:-$PWD}/test_assets}"; then
+#   echo "WARNING: Failed to download test assets from GCS. These files are only used for end-to-end tests; you may not have access to the bucket."
+# fi
+
+# Check if the base image exists locally
+if ! docker image inspect "${LOCAL_IMAGE_NAME}" &> /dev/null; then
+  echo "ERROR: Base image '${LOCAL_IMAGE_NAME}' not found locally."
+  echo "Please build it first by running 'bash docker_build_dependency_image.sh'."
+  exit 1
 fi
 
 docker build --build-arg BASEIMAGE=${LOCAL_IMAGE_NAME} -f ./maxtext_runner.Dockerfile -t ${LOCAL_IMAGE_NAME_RUNNER} .

docs/tutorials/grpo with pathways.md

@@ -0,0 +1,70 @@
+<!--
+ Copyright 2024 Google LLC
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      https://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+ -->
+
+# Try GRPO with Pathways!
+
+This tutorial demonstrates step-by-step instructions for setting up the environment and then training the Llama3.1 70B-IT model on the GSM8K math reasoning benchmark using Group Relative Policy Optimization (GRPO). GRPO can enhance your model's problem-solving skills on mathematical word problems, coding problems, etc.
+
+GRPO is an RL algorithm designed to enhance the reasoning abilities of LLMs. It is a variant of Proximal Policy Optimization (PPO) that reduces memory usage by eliminating the need for a separate value function model. GRPO works by generating multiple responses for a given prompt, evaluating these responses using a reward model, and then calculating a relative advantage based on the group's performance to update the policy.
+
+We use Tunix as the library for GRPO. 
+And we use vLLM as the library for efficient model inference and generation.
+
+Furthermore, we use Pathways for [orchestration](https://cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro). Using Pathways, you can also run GRPO in a disaggregated mode where the trainer and the samplers are running on separate mesh. Try out the following recipe `v5p-64`. You can submit jobs to a Pathways enabled GKE cluster.
+ 
+## Build and Upload MaxText Docker Image with Tunix, vLLM, tpu-commons dependencies
+Run the following bash script to create a docker image with all the dependencies of MaxText, Tunix, vLLM and tpu-commons installed.
+
+In addition to MaxText dependencies, 
+
+1. It installs `pip install keyring keyrings.google-artifactregistry-auth` which enables pip to authenticate with Google Artifact Registry automatically.
+2. Next, it installs `vLLM` for Jax and TPUs from the artifact registry `https://us-python.pkg.dev/cloud-tpu-images/maxtext-rl/simple/`
+3. Then, it installs `tpu-commons` from the same artifact registry.
+
+
+`tpu_commons` is the TPU backend for vLLM. You will need both libraries to run vLLM on tpus.
+We use the scheduler code from vLLM, and the model runner code from `tpu_commons`
+
+```
+bash docker_build_dependency_image.sh MODE=grpo
+```
+
+You can also use `bash docker_build_dependency_image.sh MODE=grpo-experimental` to try out new features via experimental dependencies such as improved pathwaysutils resharding API
+
+
+
+### Upload the dependency docker image along with MaxText code
+```
+bash docker_upload_runner.sh CLOUD_IMAGE_NAME=path/to/gcr.io
+```
+
+### Submit your jobs
+
+Please use a pathways enabled cluster, and you can submit the script `maxtext/src/MaxText/examples/grpo_llama3_1_70b_demo_pw.py` via XPK
+```
+xpk workload create-pathways --workload $WORKLOAD \
+--docker-image path/to/gcr.io:latest --cluster $TPU_CLUSTER \
+--tpu-type=$TPU_TYPE --num-slices=1  --zone=$ZONE \
+--project=$PROJECT_ID --priority=high \
+--command "HF_TOKEN=$HF_TOKEN TF_CPP_MIN_LOG_LEVEL=0 JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE='1' python  src/MaxText/examples/grpo_llama3_1_70b_demo_pw.py"
+```
+
+The overview of the demo script ~/maxtext/src/MaxText/examples/grpo_llama3_1_70b_demo_pw.py` is as follows:
+
+1. We load a policy model and a reference model. Both are copies of `Llama3.1-70b-Instruct`.
+2. Evaluate the policy model's performance on GSM8K math reasoning benchmark.
+3. Train the policy model using GRPO with potentially different meshes for trainer and rollout dependending on the parameters `TRAINER_DEVICES_FRACTION` and `SAMPLER_DEVICES_FRACTION`. If we set both of these to `1.0`, the entire (same) mesh will be used for both trainer and rollout. If we set say `TRAINER_DEVICES_FRACTION=0.5` and `SAMPLER_DEVICES_FRACTION=0.5`, the first half of the devices will be used for trainer and the second half will be used for rollout
+4. Evaluate the policy model's performance on GSM8K math reasoning benchmark after the post-training with GRPO.

docs/tutorials/grpo.md

@@ -32,7 +32,7 @@ In this tutorial we use a single host TPUVM such as `v6e-8/v5p-8`. Let's get sta
 bash setup.sh
 ```
 
-### Activate your virtual environment
+### Activate your virtual environment (Skip if you have already done this for running `bash setup.sh` )
 ```
 # Replace with your virtual environment name if not using this default name
 venv_name="maxtext_venv"
@@ -60,7 +60,7 @@ We use the scheduler code from vLLM, and the model runner code from `tpu_commons
 
 Finally, run the script
 
-`python ~/maxtext/MaxText/examples/grpo_llama3_demo.py`
+`python ~/maxtext/src/MaxText/examples/grpo_llama3_1_8b_demo.py`
 
 The overview of the demo script is as follows:
 

maxtext_grpo_dependencies.Dockerfile

@@ -0,0 +1,59 @@
+# Copyright 2023–2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+ARG BASEIMAGE
+FROM ${BASEIMAGE}
+ARG MODE
+
+ENV MODE=$MODE
+
+RUN echo "Installing GRPO dependencies (vLLM, tpu-common, tunix) with MODE=${MODE}"
+
+
+# Uninstall existing jax to avoid conflicts
+RUN pip uninstall -y jax jaxlib libtpu
+
+RUN pip install aiohttp==3.12.15
+
+# Install Python packages that enable pip to authenticate with Google Artifact Registry automatically.
+RUN pip install keyring keyrings.google-artifactregistry-auth
+
+RUN pip install numba==0.61.2
+
+# Install vLLM for Jax and TPUs from the artifact registry
+RUN VLLM_TARGET_DEVICE="tpu" pip install --no-cache-dir --pre \
+    --index-url https://us-python.pkg.dev/cloud-tpu-images/maxtext-rl/simple/ \
+    --extra-index-url https://pypi.org/simple/ \
+    --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \
+    --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
+    --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \
+    --find-links https://storage.googleapis.com/libtpu-wheels/index.html \
+    --find-links https://storage.googleapis.com/libtpu-releases/index.html \
+    --find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
+    --find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html \
+    vllm==0.11.1rc1.dev292+g1b86bd8e1.tpu
+
+# Install tpu-commons from the artifact registry
+RUN pip install --no-cache-dir --pre \
+    --index-url https://us-python.pkg.dev/cloud-tpu-images/maxtext-rl/simple/ \
+    --extra-index-url https://pypi.org/simple/ \
+    --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ \
+    --find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html \
+    tpu-commons==0.1.2
+
+RUN if [ "$MODE" = "grpo-experimental" ]; then \
+    pip uninstall -y jax jaxlib libtpu && \
+    pip install --pre -U jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ && \
+    pip install -U --pre libtpu -f https://storage.googleapis.com/jax-releases/libtpu_releases.html; \
+    fi

maxtext_jax_ai_image.Dockerfile

@@ -49,7 +49,7 @@ RUN if [ "$DEVICE" = "tpu" ] && [ "$JAX_STABLE_STACK_BASEIMAGE" = "us-docker.pkg
 
 # Install google-tunix for TPU devices, skip for GPU
 RUN if [ "$DEVICE" = "tpu" ]; then \
-        python3 -m pip install 'google-tunix>=0.1.0'; \
+        python3 -m pip install 'google-tunix>=0.1.2'; \
         # TODO: Once tunix stopped pinning jax 0.7.1, we should remove our 0.7.0 version pin (b/450286600)
         python3 -m pip install 'jax==0.7.0' 'jaxlib==0.7.0'; \
   fi

setup.sh

@@ -184,7 +184,7 @@ if [[ "$MODE" == "stable" || ! -v MODE ]]; then
 
         # TODO: Once tunix has support for GPUs, move it from here to requirements.txt
         echo "Installing google-tunix for stable TPU environment"
-        python3 -m uv pip install 'google-tunix>=0.1.0'
+        python3 -m uv pip install 'google-tunix>=0.1.2'
         echo "Installing stable jax, jaxlib for tpu"
         if [[ -n "$JAX_VERSION" ]]; then
             echo "Installing stable jax, jaxlib, libtpu version ${JAX_VERSION}"