Docs migration from XPK to CTK(3/5)

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=<ZONE> # e.g., 'us-central1' or 'us-central1-a'
 
 # Use a GCS bucket you own to store logs and checkpoints.
 export BASE_OUTPUT_DIRECTORY=<GCS_BUCKET> # 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=<IMAGE_NAME>
+export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/${CLOUD_IMAGE_NAME?}"
+
+# Your Hugging Face access token.
+export HF_TOKEN=<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=<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=<CLUSTER_TOOLKIT_COMPUTE_TYPE>
+export TOPOLOGY=<TPU_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 <POD_NAME> with the name you found)
-kubectl logs -f <POD_NAME>
+# 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
 

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=<ZONE> # e.g., 'us-central1' or 'us-central1-a'
 
 # Use a GCS bucket you own to store logs and checkpoints.
 export BASE_OUTPUT_DIRECTORY=<GCS_BUCKET> # 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=<IMAGE_NAME>
+export DOCKER_IMAGE="gcr.io/${PROJECT_ID?}/${CLOUD_IMAGE_NAME?}"
+
+# Your Hugging Face access token.
+export HF_TOKEN=<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=<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=<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=<CLUSTER_TOOLKIT_COMPUTE_TYPE>
+export TOPOLOGY=<TPU_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 <POD_NAME> with the name you found)
-kubectl logs -f <POD_NAME>
+# 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