# 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. ## Prerequisites 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)). - **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 Set up the following environment variables to configure your training run. Replace placeholders with your actual values. ```bash # Your GCP project ID. # If you've already set it in your local config, you can retrieve it via: # gcloud config get-value project export PROJECT_ID= # The name of your GKE cluster. export CLUSTER_NAME= # The GCP location of your GKE cluster. 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 ``` ## Authenticate with Hugging Face To download the `gpt-oss-20b` model checkpoint from Hugging Face, you need to authenticate using your Hugging Face account credentials. Run the following command and follow the prompts to log in: ```bash hf auth login ``` ## Get Your MaxText Compatible Model Checkpoint ### Option 1: Using an existing MaxText checkpoint If you already have a MaxText-compatible model checkpoint, simply set the following environment variable and move on to the next section. ```bash export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/0/items ``` ### Option 2: Converting from a Hugging Face checkpoint Refer to [Hugging Face to MaxText](hf-to-maxtext) to convert a Hugging Face checkpoint to MaxText format. After conversion finishes, set `MAXTEXT_CKPT_PATH` to the converted MaxText checkpoint path. ```bash export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/0/items ``` > **Note:** Converting the 20B model requires approximately 40 GB of free disk space to download its safetensors. Please verify you have sufficient space before running the conversion. ## Run RL Workload ### 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_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 ``` ### 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. ```bash kubectl get jobset -n default ${WORKLOAD_NAME} # List pods to find the specific name kubectl get pods | grep ${WORKLOAD_NAME} # stream the logs from the running pod (replace with the name you found) kubectl logs -f ``` 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. ### Monitor RL Metrics During RL training, you can monitor key metrics to track model convergence, reward trends, and hardware performance. To enable Tunix-managed metrics measurement, set `enable_tunix_perf_metrics` to `true` in RL configurations. Note that this flag is already set to `True` by default for this tutorial workload. When enabled, Tunix automatically collects and uploads these metrics to TensorBoard. For a complete list of collected metrics, see the [Tunix Metrics Documentation](https://tunix.readthedocs.io/en/latest/metrics.html). Key metrics to monitor include: - **Model Quality & Reward Metrics:** - `rewards/mean`: The average reward across the batch (crucial for tracking learning progress). - `score/mean`: The average raw score from the reward model before applying the KL penalty. - **Rollout & Generation Metrics:** - `rollout_time`: How long each rollout step takes. - `completions/mean_length`: The average token length of generated completions. - `actor_dequeue_time`: The time spent waiting for data from the rollout workers (relevant when async rollout is enabled). - **Performance & Efficiency Metrics:** - `step_time_sec`: The execution time for a single training step. ## Convert Checkpoint to Hugging Face Format Refer to [MaxText to Hugging Face](maxtext-to-hf) to convert a MaxText checkpoint back to Hugging Face format.