diff --git a/docs/tutorials/posttraining/rl_gemma4_e4b.md b/docs/tutorials/posttraining/rl_gemma4_e4b.md index 6a01a955ac..95688bc927 100644 --- a/docs/tutorials/posttraining/rl_gemma4_e4b.md +++ b/docs/tutorials/posttraining/rl_gemma4_e4b.md @@ -17,7 +17,7 @@ # Reinforcement Learning with gemma4-e4b on Multi-Host TPUs This tutorial provides step-by-step instructions for setting up the environment -and training the gemma4-e4b model with GRPO on the [OpenMathInstruct-2 dataset](https://huggingface.co/datasets/nvidia/OpenMathInstruct-2) on a Cloud TPU v6e (Trillium) GKE cluster using a `v6e-32` (4x8) slice. +and training the gemma4-e4b model with GRPO on the [OpenMathInstruct-2 dataset](https://huggingface.co/datasets/nvidia/OpenMathInstruct-2) on a Cloud TPU v6e (Trillium) GKE cluster using a `v6e-32` (4x8) slice with Cluster Toolkit. ## Prerequisites @@ -26,8 +26,8 @@ Before starting, ensure you have: - Access to a Google Cloud Project with TPU quotas. - A Hugging Face account with an access token for downloading models (the `google/gemma-4-E4B` and `google/gemma-4-E4B-it` repositories are gated; request access before proceeding). - Permissions for Google Artifact Registry (Artifact Registry Writer role). -- Prerequisites for XPK installed (follow [official documentation](https://github.com/AI-Hypercomputer/xpk/blob/main/docs/installation.md#1-prerequisites)). -- A Pathways-ready GKE cluster (see [create GKE cluster](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster)). +- Cluster Toolkit installed and configured. Follow [Running MaxText with Cluster Toolkit](../../run_maxtext/run_maxtext_via_cluster_toolkit.md) for `gcluster` setup. +- A GKE cluster configured for Cluster Toolkit, including healthy Kueue and JobSet components. - **Docker** installed and configured for sudoless use. Follow the steps to [configure sudoless Docker](https://docs.docker.com/engine/install/linux-postinstall/). ## Setup Environment Variables @@ -49,6 +49,16 @@ export ZONE= # e.g., 'us-central1' or 'us-central1-a' # Use a GCS bucket you own to store logs and checkpoints. export BASE_OUTPUT_DIRECTORY= # e.g., gs://my-bucket/maxtext-runs + +# An arbitrary string to identify this specific run. +export RUN_NAME="rl-gemma4-$(date +%Y%m%d-%H%M%S)" + +# The Docker image you pushed in the prerequisite step +export CLOUD_IMAGE_NAME= +export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/${CLOUD_IMAGE_NAME?}" + +# Your Hugging Face access token. +export HF_TOKEN= ``` ## Authenticate with Hugging Face @@ -96,36 +106,71 @@ Both files are already included under `src/maxtext/examples/chat_templates/`, so ### Build and Upload MaxText Docker Image -For instructions on building and uploading the MaxText Docker image with post-training dependencies, please refer to the [official documentation](build-docker). +For instructions on building and uploading the MaxText Docker image with post-training dependencies, please refer to the [official documentation](../build_maxtext.md). ### Submit your workload ```bash -# The Docker image you pushed in the previous step -export CLOUD_IMAGE_NAME= -export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/${CLOUD_IMAGE_NAME?}" - -# Run the RL training script on your cluster -run_tutorial maxtext/trainers/post_train/rl/scripts/run_gemma4_e4b_rl.sh +export COMPUTE_TYPE= +export TOPOLOGY= + +gcloud config set project ${PROJECT_ID?} +gcloud container clusters get-credentials ${CLUSTER_NAME?} \ + --location ${ZONE?} \ + --project ${PROJECT_ID?} +gcluster job config set project ${PROJECT_ID?} +gcluster job config set cluster ${CLUSTER_NAME?} +gcluster job config set location ${ZONE?} + +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name=${RUN_NAME?} \ + --pathways \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --num-slices=1 \ + --pathways-gcs-location=${BASE_OUTPUT_DIRECTORY?} \ + --command="python3 -m maxtext.trainers.post_train.rl.train_rl \ + model_name=gemma4-e4b \ + tokenizer_path=google/gemma-4-E4B \ + load_parameters_path=${MAXTEXT_CKPT_PATH?} \ + run_name=${RUN_NAME?} \ + base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ + hf_access_token=${HF_TOKEN?} \ + scan_layers=False \ + data_template_path=maxtext/examples/chat_templates/openmathinstruct2_rl.json \ + chat_template_path=maxtext/examples/chat_templates/gemma-3-27b-chat_template.json \ + dataset_name=nvidia/OpenMathInstruct-2 \ + hf_train_files=hf://datasets/nvidia/OpenMathInstruct-2/data/train_1M-*.parquet \ + train_split=train_1M \ + chips_per_vm=4 \ + ici_tensor_parallelism=2 \ + rollout_tensor_parallelism=2 \ + enable_single_controller=True" ``` -> **Note:** The `run_gemma4_e4b_rl.sh` script pins the Pathways component images to specific versions via the xpk `--server-image` and `--proxy-server-image` flags (set through the `PATHWAYS_SERVER_IMAGE` and `PATHWAYS_PROXY_SERVER_IMAGE` variables at the top of the script). The `--server-image` is used for both the Pathways resource-manager server and the workers (the reference config uses the same image for both). Update these variables if you need a different Pathways release. - ### Monitor your workload -To monitor your job's progress, you can use `kubectl` to check the `Jobset` status and stream logs directly from the pods. +To monitor your job's progress, you can use `gcluster` or `kubectl` to check the `JobSet` status and stream logs directly: ```bash -kubectl get jobset -n default ${WORKLOAD_NAME} +# Check job status with Cluster Toolkit +gcluster job list + +# Stream logs with Cluster Toolkit (specify --main-only=false for Pathways workloads) +gcluster job logs ${RUN_NAME?} --main-only=false + +# Alternatively, check JobSet status with kubectl +kubectl get jobset -l gcluster.google.com/workload=${RUN_NAME?} -# List pods to find the specific name -kubectl get pods | grep ${WORKLOAD_NAME} +# List pods (use jobset-name to select both head and worker pods in Pathways) +kubectl get pods -l jobset.sigs.k8s.io/jobset-name=${RUN_NAME?} -# stream the logs from the running pod (replace with the name you found) -kubectl logs -f +# Stream logs with kubectl +kubectl logs -f -l jobset.sigs.k8s.io/jobset-name=${RUN_NAME?} --all-containers=true ``` -Alternatively, after running the bash script, you will also get a link to the Google Cloud Console to view your workload logs. Follow the link to view logs and monitor your workload's progress in the Cloud Console. +Alternatively, `gcluster job submit` provides a link to the Google Cloud Console to view your workload logs. Follow the link to view logs and monitor your workload's progress in the Cloud Console. ### Monitor RL Metrics diff --git a/docs/tutorials/posttraining/rl_gptoss_20b.md b/docs/tutorials/posttraining/rl_gptoss_20b.md index 05f78fa3d2..e19c89e57c 100644 --- a/docs/tutorials/posttraining/rl_gptoss_20b.md +++ b/docs/tutorials/posttraining/rl_gptoss_20b.md @@ -17,7 +17,7 @@ # Reinforcement Learning with GPT-OSS 20B on Multi-Host TPUs This tutorial provides step-by-step instructions for setting up the environment -and training the GPT-OSS 20B model on the [GSM8K dataset](https://huggingface.co/datasets/openai/gsm8k) on a GKE cluster with `v5p-64` nodes. +and training the GPT-OSS 20B model on the [GSM8K dataset](https://huggingface.co/datasets/openai/gsm8k) on a GKE cluster with `v5p-64` nodes using Cluster Toolkit. ## Prerequisites @@ -26,8 +26,8 @@ Before starting, ensure you have: - Access to a Google Cloud Project with TPU quotas. - A Hugging Face account with an access token for downloading models. - Permissions for Google Artifact Registry (Artifact Registry Writer role). -- Prerequisites for XPK installed (follow [official documentation](https://github.com/AI-Hypercomputer/xpk/blob/main/docs/installation.md#1-prerequisites)). -- A Pathways-ready GKE cluster (see [create GKE cluster](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster)). +- Cluster Toolkit installed and configured. Follow [Running MaxText with Cluster Toolkit](../../run_maxtext/run_maxtext_via_cluster_toolkit.md) for `gcluster` setup. +- A GKE cluster configured for Cluster Toolkit, including healthy Kueue and JobSet components. - **Docker** installed and configured for sudoless use. Follow the steps to [configure sudoless Docker](https://docs.docker.com/engine/install/linux-postinstall/). ## Setup Environment Variables @@ -49,6 +49,19 @@ export ZONE= # e.g., 'us-central1' or 'us-central1-a' # Use a GCS bucket you own to store logs and checkpoints. export BASE_OUTPUT_DIRECTORY= # e.g., gs://my-bucket/maxtext-runs + +# An arbitrary string to identify this specific run. +export RUN_NAME="rl-gptoss-$(date +%Y%m%d-%H%M%S)" + +# The Docker image you pushed in the prerequisite step +export CLOUD_IMAGE_NAME= +export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/${CLOUD_IMAGE_NAME?}" + +# Your Hugging Face access token. +export HF_TOKEN= + +# Tokenizer path for GPT-OSS 20B +export TOKENIZER_PATH="unsloth/gpt-oss-20b-BF16" ``` ## Authenticate with Hugging Face @@ -84,34 +97,65 @@ export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/ ### Build and Upload MaxText Docker Image -For instructions on building and uploading the MaxText Docker image with post-training dependencies, please refer to the [official documentation](build-docker). +For instructions on building and uploading the MaxText Docker image with post-training dependencies, please refer to the [official documentation](../build_maxtext.md). ### Submit your workload ```bash -# The Docker image you pushed in the previous step -export CLOUD_IMAGE_NAME= -export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/${CLOUD_IMAGE_NAME?}" - -# Run the RL training script on your cluster -run_tutorial maxtext/trainers/post_train/rl/scripts/run_gptoss_20b_rl.sh +export COMPUTE_TYPE= +export TOPOLOGY= + +gcloud config set project ${PROJECT_ID?} +gcloud container clusters get-credentials ${CLUSTER_NAME?} \ + --location ${ZONE?} \ + --project ${PROJECT_ID?} +gcluster job config set project ${PROJECT_ID?} +gcluster job config set cluster ${CLUSTER_NAME?} +gcluster job config set location ${ZONE?} + +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name=${RUN_NAME?} \ + --pathways \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --num-slices=1 \ + --pathways-gcs-location=${BASE_OUTPUT_DIRECTORY?} \ + --command="python3 -m maxtext.trainers.post_train.rl.train_rl \ + model_name=gpt-oss-20b \ + tokenizer_path=${TOKENIZER_PATH?} \ + load_parameters_path=${MAXTEXT_CKPT_PATH?} \ + run_name=${RUN_NAME?} \ + base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ + hf_access_token=${HF_TOKEN?} \ + chat_template_path=maxtext/examples/chat_templates/gpt_oss_rl.json \ + chips_per_vm=4 \ + rollout_tensor_parallelism=8 \ + enable_single_controller=True" ``` ### Monitor your workload -To monitor your job's progress, you can use `kubectl` to check the `Jobset` status and stream logs directly from the pods. +To monitor your job's progress, you can use `gcluster` or `kubectl` to check the `JobSet` status and stream logs directly: ```bash -kubectl get jobset -n default ${WORKLOAD_NAME} +# Check job status with Cluster Toolkit +gcluster job list + +# Stream logs with Cluster Toolkit (specify --main-only=false for Pathways workloads) +gcluster job logs ${RUN_NAME?} --main-only=false + +# Alternatively, check JobSet status with kubectl +kubectl get jobset -l gcluster.google.com/workload=${RUN_NAME?} -# List pods to find the specific name -kubectl get pods | grep ${WORKLOAD_NAME} +# List pods (use jobset-name to select both head and worker pods in Pathways) +kubectl get pods -l jobset.sigs.k8s.io/jobset-name=${RUN_NAME?} -# stream the logs from the running pod (replace with the name you found) -kubectl logs -f +# Stream logs with kubectl +kubectl logs -f -l jobset.sigs.k8s.io/jobset-name=${RUN_NAME?} --all-containers=true ``` -Alternatively, after running the bash script, you will also get a link to the Google Cloud Console to view your workload logs. Follow the link to view logs and monitor your workload's progress in the Cloud Console. +Alternatively, `gcluster job submit` provides a link to the Google Cloud Console to view your workload logs. Follow the link to view logs and monitor your workload's progress in the Cloud Console. ### Monitor RL Metrics diff --git a/docs/tutorials/posttraining/rl_on_multi_host.md b/docs/tutorials/posttraining/rl_on_multi_host.md index 828ec7b1f7..c484492a9d 100644 --- a/docs/tutorials/posttraining/rl_on_multi_host.md +++ b/docs/tutorials/posttraining/rl_on_multi_host.md @@ -61,9 +61,9 @@ Before starting, ensure you have: - **Artifact Registry Writer** (`roles/artifactregistry.writer`) to upload Docker images. - **Storage Admin** (`roles/storage.admin`) or **Storage Object Admin** (`roles/storage.objectAdmin`) combined with **Storage Legacy Bucket Reader** (`roles/storage.legacyBucketReader`) on your GCS bucket to read/write checkpoints and logs. (Note: A bucket-level read permission like `storage.buckets.get` is required by JAX/TensorStore to verify bucket existence and metadata; using `roles/storage.objectAdmin` alone will cause a misleading "bucket not found" error). - A Hugging Face account with an access token for downloading models. -- Prerequisites for XPK installed (follow [official documentation](https://github.com/AI-Hypercomputer/xpk/blob/main/docs/installation.md#1-prerequisites)). +- Cluster Toolkit installed and configured. Follow [Running MaxText with Cluster Toolkit](../../run_maxtext/run_maxtext_via_cluster_toolkit.md) for `gcluster` setup. - **Important:** Modern GKE clusters require the GKE auth plugin. If you encounter `gke-gcloud-auth-plugin not found` when running `kubectl` commands, you must install it locally (e.g., `sudo apt-get install google-cloud-cli-gke-gcloud-auth-plugin` for `apt` installations, or `gcloud components install gke-gcloud-auth-plugin` for standalone archive installations). -- A Pathways-ready GKE cluster (see [create GKE cluster](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster)). +- A GKE cluster configured for Cluster Toolkit, with healthy Kueue and JobSet components. - **Docker** installed and configured for sudoless use. Follow the steps to [configure sudoless Docker](https://docs.docker.com/engine/install/linux-postinstall/). ## Build and upload MaxText Docker image @@ -93,9 +93,8 @@ export HF_TOKEN= export BASE_OUTPUT_DIRECTORY= # e.g., gs://my-bucket/maxtext-runs # An arbitrary string to identify this specific run. -# We recommend to include the model, user, and timestamp. -# Note: Kubernetes requires workload names to be valid DNS labels (lowercase, no underscores or periods). -export RUN_NAME= +# Note: Workload names cannot exceed 28 characters and must be valid DNS labels (lowercase alphanumeric and hyphens). +export RUN_NAME="rl-$(date +%m%d%H%M%S)" # The directory containing the MaxText-compatible model checkpoint. # If you are converting from a Hugging Face checkpoint, see: @@ -161,60 +160,86 @@ export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/ See the **Troubleshooting** section for concise instructions on how to retry or resume a failed workload. -Ensure you have a Pathways-ready GKE cluster (as mentioned in Prerequisites) and -submit the `train_rl.py` script via XPK. +Configure `kubectl` and `gcluster` for the target cluster before submitting: -> **Note:** XPK v0.14.0+ automatically discovers your cluster's location from -> GCP. You don't need to specify `--zone` in the commands below. If using an -> older XPK version, add `--zone=` to the workload commands. +```bash +gcloud config set project ${PROJECT_ID?} +gcloud container clusters get-credentials ${GKE_CLUSTER?} \ + --location ${ZONE?} \ + --project ${PROJECT_ID?} +gcluster job config set project ${PROJECT_ID?} +gcluster job config set cluster ${GKE_CLUSTER?} +gcluster job config set location ${ZONE?} +``` + +Set the Cluster Toolkit placement values. `TOPOLOGY` must match the TPU slice +available on the cluster; for example, verify the supported topology before +using a four-slice v6e cluster. + +```bash +export COMPUTE_TYPE= +export TOPOLOGY= +``` ### Submit GRPO workload ```bash -xpk workload create-pathways --workload ${RUN_NAME?} \ ---docker-image ${DOCKER_IMAGE?} --cluster ${GKE_CLUSTER?} \ ---tpu-type=${TPU_TYPE?} --num-slices=1 \ ---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' \ -python3 -m maxtext.trainers.post_train.rl.train_rl \ +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name=${RUN_NAME?}-grpo \ + --pathways \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --num-slices=1 \ + --pathways-gcs-location=${BASE_OUTPUT_DIRECTORY?} \ + --command="python3 -m maxtext.trainers.post_train.rl.train_rl \ model_name=${MODEL?} \ load_parameters_path=${MAXTEXT_CKPT_PATH?} \ - run_name=${RUN_NAME?} \ + run_name=${RUN_NAME?}-grpo \ base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ rollout_tensor_parallelism=8 \ - hf_access_token=${HF_TOKEN?}" + hf_access_token=${HF_TOKEN?} \ + enable_single_controller=True" ``` ### Submit GSPO workload ```bash -xpk workload create-pathways --workload ${RUN_NAME?} \ ---docker-image ${DOCKER_IMAGE?} --cluster ${GKE_CLUSTER?} \ ---tpu-type=${TPU_TYPE?} --num-slices=1 \ ---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' \ -python3 -m maxtext.trainers.post_train.rl.train_rl \ +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name=${RUN_NAME?}-gspo \ + --pathways \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --num-slices=1 \ + --pathways-gcs-location=${BASE_OUTPUT_DIRECTORY?} \ + --command="python3 -m maxtext.trainers.post_train.rl.train_rl \ model_name=${MODEL?} \ load_parameters_path=${MAXTEXT_CKPT_PATH?} \ - run_name=${RUN_NAME?} \ + run_name=${RUN_NAME?}-gspo \ base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ rollout_tensor_parallelism=8 \ hf_access_token=${HF_TOKEN?} \ - loss_algo=gspo-token" + loss_algo=gspo-token \ + enable_single_controller=True" ``` -## Managing Workloads +## Monitor and clean up -- **Monitor workload status**: Check Pathways job status: `kubectl get pathwaysjob`. Check pod status: `kubectl get pods`. -- **Delete a workload**: To remove a failed or unwanted Pathways job, use XPK: - ```bash - xpk workload delete \ - --workload ${RUN_NAME?} \ - --cluster ${GKE_CLUSTER?} \ - --project ${PROJECT_ID?} - ``` - In case the job still lingers on, you can use - `kubectl get pods` to obtain the name of the pod and then run: `kubectl delete pod `. +```bash +gcluster job list +# Note: For Pathways workloads (> 5 pods), specify --main-only=false to retrieve logs from all pods: +gcluster job logs ${RUN_NAME?} --main-only=false +gcluster job cancel ${RUN_NAME?} +``` + +You can also inspect the Kubernetes resources directly: + +```bash +kubectl get jobset -l gcluster.google.com/workload=${RUN_NAME?} +# In Pathways workloads, use the jobset-name label to select all pods (both pathways-head and worker pods): +kubectl get pods -l jobset.sigs.k8s.io/jobset-name=${RUN_NAME?} +``` ## Troubleshooting @@ -230,9 +255,13 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \ - **Solution**: Explicitly pass at least one of them in your training command (e.g., `rollout_tensor_parallelism=8` as shown in the example commands above). - **Workload retry / resume**: - **Retry (fresh run)**: Use a unique run name to avoid overwriting - outputs: `export RUN_NAME=${RUN_NAME?}-retry1 export MAXTEXT_CKPT_PATH=${BASE_OUTPUT_DIRECTORY?}/${RUN_NAME?}/0/items`. Then - submit the XPK workload. If "workload already exists" error occurs, pick - a new name or list jobs: `kubectl get pathwaysjob`. + outputs: + ```bash + export MAXTEXT_CKPT_PATH=${BASE_OUTPUT_DIRECTORY?}/${RUN_NAME?}/0/items + export RUN_NAME=${RUN_NAME?}-retry1 + ``` + Then submit the Cluster Toolkit workload. If a "workload already exists" error occurs, pick + a new name or cancel the previous job (`gcluster job cancel ${RUN_NAME}`). - **Resume from checkpoint**: Keep the same `RUN_NAME` and set the checkpoint path: `export load_parameters_path=${MAXTEXT_CKPT_PATH?}/checkpoint-0000`. Then submit the workload again. @@ -241,4 +270,4 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \ For more detailed troubleshooting, refer to the [MaxText documentation](../../index.md) and -[XPK documentation](https://github.com/AI-Hypercomputer/xpk). +[Cluster Toolkit guide](../../run_maxtext/run_maxtext_via_cluster_toolkit.md). diff --git a/docs/tutorials/posttraining/rl_qwen3_30b.md b/docs/tutorials/posttraining/rl_qwen3_30b.md index 00de74f6f6..cb8b45b011 100644 --- a/docs/tutorials/posttraining/rl_qwen3_30b.md +++ b/docs/tutorials/posttraining/rl_qwen3_30b.md @@ -17,7 +17,7 @@ # Reinforcement Learning with Qwen3-30b-a3b-base on Multi-Host TPUs This tutorial provides step-by-step instructions for setting up the environment -and training the Qwen3-30b-a3b-base model on the [OpenMathInstruct-2 dataset](https://huggingface.co/datasets/nvidia/OpenMathInstruct-2) on Ironwood GKE cluster with `tpu7x-128` nodes. +and training the Qwen3-30b-a3b-base model on the [OpenMathInstruct-2 dataset](https://huggingface.co/datasets/nvidia/OpenMathInstruct-2) on Ironwood GKE cluster with `tpu7x-128` nodes using Cluster Toolkit. ## Prerequisites @@ -26,8 +26,8 @@ Before starting, ensure you have: - Access to a Google Cloud Project with TPU quotas. - A Hugging Face account with an access token for downloading models. - Permissions for Google Artifact Registry (Artifact Registry Writer role). -- Prerequisites for XPK installed (follow [official documentation](https://github.com/AI-Hypercomputer/xpk/blob/main/docs/installation.md#1-prerequisites)). -- A Pathways-ready GKE cluster (see [create GKE cluster](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster)). +- Cluster Toolkit installed and configured. Follow [Running MaxText with Cluster Toolkit](../../run_maxtext/run_maxtext_via_cluster_toolkit.md) for `gcluster` setup. +- A GKE cluster configured for Cluster Toolkit, including healthy Kueue and JobSet components. - **Docker** installed and configured for sudoless use. Follow the steps to [configure sudoless Docker](https://docs.docker.com/engine/install/linux-postinstall/). ## Setup Environment Variables @@ -49,6 +49,19 @@ export ZONE= # e.g., 'us-central1' or 'us-central1-a' # Use a GCS bucket you own to store logs and checkpoints. export BASE_OUTPUT_DIRECTORY= # e.g., gs://my-bucket/maxtext-runs + +# An arbitrary string to identify this specific run. +export RUN_NAME="rl-qwen3-$(date +%Y%m%d-%H%M%S)" + +# The Docker image you pushed in the prerequisite step +export CLOUD_IMAGE_NAME= +export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/${CLOUD_IMAGE_NAME?}" + +# Your Hugging Face access token. +export HF_TOKEN= + +# Tokenizer path for Qwen3 +export TOKENIZER_PATH="Qwen/Qwen2.5-32B" ``` ## Authenticate with Hugging Face @@ -84,34 +97,68 @@ export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/ ### Build and Upload MaxText Docker Image -For instructions on building and uploading the MaxText Docker image with post-training dependencies, please refer to the [official documentation](build-docker). +For instructions on building and uploading the MaxText Docker image with post-training dependencies, please refer to the [official documentation](../build_maxtext.md). ### Submit your workload ```bash -# The Docker image you pushed in the previous step -export CLOUD_IMAGE_NAME= -export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/${CLOUD_IMAGE_NAME?}" - -# Run the RL training script on your cluster -run_tutorial maxtext/trainers/post_train/rl/scripts/run_qwen3_30b_rl.sh +export COMPUTE_TYPE= +export TOPOLOGY= + +gcloud config set project ${PROJECT_ID?} +gcloud container clusters get-credentials ${CLUSTER_NAME?} \ + --location ${ZONE?} \ + --project ${PROJECT_ID?} +gcluster job config set project ${PROJECT_ID?} +gcluster job config set cluster ${CLUSTER_NAME?} +gcluster job config set location ${ZONE?} + +gcluster job submit \ + --image=${DOCKER_IMAGE?} \ + --name=${RUN_NAME?} \ + --pathways \ + --compute-type=${COMPUTE_TYPE?} \ + --topology=${TOPOLOGY?} \ + --num-slices=1 \ + --pathways-gcs-location=${BASE_OUTPUT_DIRECTORY?} \ + --command="python3 -m maxtext.trainers.post_train.rl.train_rl \ + model_name=qwen3-30b-a3b-base \ + tokenizer_path=${TOKENIZER_PATH?} \ + load_parameters_path=${MAXTEXT_CKPT_PATH?} \ + run_name=${RUN_NAME?} \ + base_output_directory=${BASE_OUTPUT_DIRECTORY?} \ + hf_access_token=${HF_TOKEN?} \ + dataset_name=nvidia/OpenMathInstruct-2 \ + hf_train_files=hf://datasets/nvidia/OpenMathInstruct-2/data/train_1M-*.parquet \ + train_split=train_1M \ + chips_per_vm=8 \ + rollout_data_parallelism=16 \ + rollout_tensor_parallelism=4 \ + enable_single_controller=True" ``` ### Monitor your workload -To monitor your job's progress, you can use `kubectl` to check the `Jobset` status and stream logs directly from the pods. +To monitor your job's progress, you can use `gcluster` or `kubectl` to check the `JobSet` status and stream logs directly: ```bash -kubectl get jobset -n default ${WORKLOAD_NAME} +# Check job status with Cluster Toolkit +gcluster job list + +# Stream logs with Cluster Toolkit (specify --main-only=false for Pathways workloads) +gcluster job logs ${RUN_NAME?} --main-only=false + +# Alternatively, check JobSet status with kubectl +kubectl get jobset -l gcluster.google.com/workload=${RUN_NAME?} -# List pods to find the specific name -kubectl get pods | grep ${WORKLOAD_NAME} +# List pods (use jobset-name to list both head and worker pods in Pathways) +kubectl get pods -l jobset.sigs.k8s.io/jobset-name=${RUN_NAME?} -# stream the logs from the running pod (replace with the name you found) -kubectl logs -f +# Stream logs with kubectl +kubectl logs -f -l jobset.sigs.k8s.io/jobset-name=${RUN_NAME?} --all-containers=true ``` -Alternatively, after running the bash script, you will also get a link to the Google Cloud Console to view your workload logs. Follow the link to view logs and monitor your workload's progress in the Cloud Console. +Alternatively, `gcluster job submit` provides a link to the Google Cloud Console to view your workload logs. Follow the link to view logs and monitor your workload's progress in the Cloud Console. ### Monitor RL Metrics diff --git a/src/maxtext/experimental/rl/README.md b/src/maxtext/experimental/rl/README.md index 30941bc449..a6c0cf7e4e 100644 --- a/src/maxtext/experimental/rl/README.md +++ b/src/maxtext/experimental/rl/README.md @@ -20,7 +20,7 @@ This directory contains code and documentation for **GRPO**, a reinforcement lea * Follow the instructions to create a GKE cluster with Pathways support: [Create a GKE Cluster with Pathways](https://cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/create-gke-cluster) * This involves setting up node pools for both training and inference. 3. **Pathways Container Images:** Your GCP service account must be allowlisted to access the Pathways container images. Contact your Google Cloud representative for assistance. -4. **XPK:** (Recommended) Install and configure XPK, the Accelerated Processing Kit, to simplify the management of your GKE cluster and Pathways environment. +4. **Cluster Management Tools:** For standard GKE workloads, [Cluster Toolkit](../../../../docs/run_maxtext/run_maxtext_via_cluster_toolkit.md) is recommended. For legacy Pathways-specific orchestration, install and configure XPK to manage your GKE cluster and Pathways environment. 5. **HF_TOKEN:** You need to have a Hugging Face token to access the models. ### Setup @@ -47,7 +47,10 @@ MODEL=llama3.3-70b \ TOKENIZER=meta-llama/Llama-3.3-70B-Instruct \ NUM_SAMPLERS=8 \ DEVICES_PER_SAMPLER=8 \ +TRAINING_SUBSLICE=4,16 \ +INFERENCE_SUBSLICE=4,16 \ TRAINING_PER_DEVICE_BATCH_SIZE=1 \ INFERENCE_PER_DEVICE_BATCH_SIZE=8 \ STEPS=20 \ bash tests/end_to_end/tpu/test_grpo.sh +```