Add RL tutorial and run script for GPT-OSS 20B on Multi-Host TPUs

susanbao · susanbao · commit becbf01765c3 · 2026-07-13T02:04:49.000Z
diff --git a/docs/tutorials/post_training_index.md b/docs/tutorials/post_training_index.md
@@ -37,6 +37,7 @@ MaxText was co-designed with key Google led innovations to provide a unified pos
   - [RL on Single-Host TPUs](./posttraining/rl.md)
   - [RL on Multi-Host TPUs](./posttraining/rl_on_multi_host.md)
   - [RL with Qwen3-30b-a3b-base](./posttraining/rl_qwen3_30b.md)
+  - [RL with GPT-OSS 20B](./posttraining/rl_gptoss_20b.md)
 
 ## Step by step RL
 
@@ -75,6 +76,7 @@ posttraining/dpo.md
 posttraining/rl.md
 posttraining/rl_on_multi_host.md
 posttraining/rl_qwen3_30b.md
+posttraining/rl_gptoss_20b.md
 posttraining/knowledge_distillation.md
 posttraining/lora.md
 posttraining/lora_on_multi_host.md
diff --git a/docs/tutorials/posttraining/rl_gptoss_20b.md b/docs/tutorials/posttraining/rl_gptoss_20b.md
@@ -0,0 +1,201 @@
+<!--
+ Copyright 2023-2026 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.
+ -->
+
+# 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 an Viperfish 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/).
+
+## 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).
+
+## 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=<PROJECT_ID>
+
+# The name of your GKE cluster.
+export CLUSTER_NAME=<CLUSTER_NAME>
+
+# The GCP location of your GKE cluster.
+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
+
+# 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?}"
+```
+
+## Clone MaxText Repository
+
+If you haven't already, clone the MaxText repository to your local machine:
+
+```bash
+git clone https://github.com/AI-Hypercomputer/maxtext.git
+cd maxtext
+```
+
+## 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=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items
+```
+
+### Option 2: Converting from a Hugging Face checkpoint
+
+> **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.
+
+```bash
+# Create and activate a virtual environment
+uv venv --python 3.12 --seed tpu_venv
+source tpu_venv/bin/activate
+uv pip install -e .[tpu] --resolution=lowest
+
+# Install torch for conversion
+python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+
+# Download Hugging Face weights to a local directory
+huggingface-cli download openai/gpt-oss-20b --local-dir ./gpt-oss-20b-hf
+
+# Set up the MaxText checkpoint path environment variable where the converted checkpoint will be saved
+export MAXTEXT_CKPT_PATH="${BASE_OUTPUT_DIRECTORY}/checkpoints/gpt-oss-20b-converted/0/items"
+
+# Run the conversion script to convert the Hugging Face checkpoint to MaxText format
+python3 -m maxtext.checkpoint_conversion.standalone_scripts.convert_gpt_oss_ckpt \
+    --base-model-path ./gpt-oss-20b-hf \
+    --maxtext-model-path ${MAXTEXT_CKPT_PATH} \
+    --model-size gpt-oss-20b \
+    --use-ocdbt=False --use-zarr3=False
+
+# Deactivate the virtual environment
+deactivate
+rm -rf tpu_venv
+```
+
+## Run RL Workload
+
+### Submit your workload
+
+```bash
+# Create and activate a virtual environment
+uv venv --python 3.12 --seed runner_venv
+source runner_venv/bin/activate
+uv pip install -e .[runner] --resolution=lowest
+
+# Run the RL training script on your cluster
+bash scripts/run_gptoss_20b_rl.sh
+
+# Deactivate the virtual environment
+deactivate
+rm -rf runner_venv
+```
+
+### 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 <POD_NAME> with the name you found)
+kubectl logs -f <POD_NAME>
+```
+
+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 `src/maxtext/configs/post_train/rl.yml`. Note that this flag is already set to `True` by default in the [scripts/run_gptoss_20b_rl.sh](../../../scripts/run_gptoss_20b_rl.sh) script 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
+
+After training, you may want to convert your MaxText checkpoint back to Hugging Face format. Use the following command to perform the conversion:
+
+```bash
+# Create and activate a virtual environment
+uv venv --python 3.12 --seed tpu_venv
+source tpu_venv/bin/activate
+uv pip install -e .[tpu] --resolution=lowest
+
+# Install torch for conversion
+python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+
+# Define the output path for the converted checkpoint
+export HF_CKPT_OUTPUT_DIR="${BASE_OUTPUT_DIRECTORY}/checkpoints/gpt-oss-20b-hf-converted"
+
+# Run the conversion script to convert the MaxText checkpoint back to Hugging Face format 
+python3 -m maxtext.checkpoint_conversion.to_huggingface \
+    src/maxtext/configs/base.yml \
+    model_name=gpt-oss-20b \
+    load_parameters_path="${MAXTEXT_CKPT_PATH}" \
+    base_output_directory="${HF_CKPT_OUTPUT_DIR}" \
+    scan_layers=True \
+    weight_dtype=bfloat16 hardware=cpu skip_jax_distributed_system=True
+
+# Deactivate the virtual environment
+deactivate
+rm -rf tpu_venv
+```
diff --git a/scripts/run_gptoss_20b_rl.sh b/scripts/run_gptoss_20b_rl.sh
@@ -0,0 +1,134 @@
+#!/bin/bash
+
+# This script launches a Reinforcement Learning (RL) training workload for the
+# GPT-OSS 20B model on a GKE cluster using XPK.
+
+set -e
+
+# --- Environment Setup ---
+if ! pip show xpk &> /dev/null; then
+    echo "xpk not found in the environment. Please install it by running:"
+    echo "uv pip install -e .[runner] --resolution=lowest"
+    exit 1
+fi
+
+# --- Environment Variables ---
+export PROJECT_ID="${PROJECT_ID:-}" # GCP project ID where the Ironwood cluster is deployed
+export CLUSTER_NAME="${CLUSTER_NAME:-}" # Name of your Ironwood cluster
+export ZONE="${ZONE:-}" # Zone where your Ironwood cluster is deployed
+export BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-}" # GCS bucket path for outputs (e.g., gs://my-bucket/outputs)
+export DOCKER_IMAGE="${DOCKER_IMAGE:-}" # Full path to the Docker image you pushed (e.g., gcr.io/my-project/my-image:tag)
+export MAXTEXT_CKPT_PATH="${MAXTEXT_CKPT_PATH:-}" # GCS path of the MaxText checkpoint you want to fine-tune from (e.g., gs://my-bucket/checkpoints/maxtext-ckpt)
+export TPU_TYPE="v5p-64"
+export WORKLOAD_NAME="rl-$(date +%Y%m%d-%H%M)"
+
+# --- Variable Validation ---
+if [ -z "$PROJECT_ID" ]; then
+    echo "Error: PROJECT_ID is not set. Please set it in the script or as an environment variable."
+    exit 1
+fi
+if [ -z "$CLUSTER_NAME" ]; then
+    echo "Error: CLUSTER_NAME is not set. Please set it in the script or as an environment variable."
+    exit 1
+fi
+if [ -z "$ZONE" ]; then
+    echo "Error: ZONE is not set. Please set it in the script or as an environment variable."
+    exit 1
+fi
+if [ -z "$BASE_OUTPUT_DIRECTORY" ]; then
+    echo "Error: BASE_OUTPUT_DIRECTORY is not set. Please set it in the script or as an environment variable."
+    exit 1
+fi
+if [ -z "$DOCKER_IMAGE" ]; then
+    echo "Error: DOCKER_IMAGE is not set. Please set it in the script or as an environment variable."
+    exit 1
+fi
+
+if [ -z "$MAXTEXT_CKPT_PATH" ]; then
+    echo "MAXTEXT_CKPT_PATH is not set. Please set it in the script or as an environment variable."
+    exit 1
+fi
+
+# XLA Flags
+XLA_FLAGS="--xla_tpu_dvfs_p_state=7 \
+--xla_tpu_scoped_vmem_limit_kib=65536 \
+--xla_tpu_num_sparse_cores_for_gather_offloading=1 \
+--xla_tpu_bf16_emission_mode=NATIVE_EMISSION \
+--xla_tpu_enable_sparse_core_reduce_scatter_v2=true \
+--xla_tpu_enable_sparse_core_collective_offload_all_gather=true \
+--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \
+--xla_tpu_use_tc_device_shape_on_sc=True \
+--xla_sc_disable_megacore_partitioning=True \
+--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \
+--xla_enable_async_all_gather=true \
+--xla_tpu_prefer_async_allgather_to_allreduce=true \
+--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \
+--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \
+--xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \
+--xla_tpu_use_single_sparse_core_for_all_gather_offload=true \
+--xla_tpu_enable_concurrent_sparse_core_offloading=true \
+--xla_tpu_enable_offloading_gather_to_sparsecore=true \
+--xla_tpu_sparse_core_all_gather_latency_multiplier=1 \
+--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \
+--xla_tpu_enable_sparse_core_collective_aggregator=true \
+--xla_tpu_enable_latency_hiding_layer_scheduler=true \
+--xla_tpu_scheduler_percent_shared_memory_limit=150 \
+--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \
+--xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \
+--xla_tpu_pcie_bandwidth_multiplier=0.03 \
+--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \
+--xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \
+--xla_tpu_enable_3d_reduce_scatter_decomposer=false"
+
+MAXTEXT_COMMAND="MODEL_IMPL_TYPE=flax_nnx \
+  GRPC_ENABLE_FORK_SUPPORT=0 \
+  JAX_RANDOM_WEIGHTS=1 \
+  VLLM_ENABLE_V1_MULTIPROCESSING=0 \
+  SKIP_JAX_PRECOMPILE=1 \
+  NEW_MODEL_DESIGN=0 \
+  TPU_MIN_LOG_LEVEL=0 \
+  TF_CPP_MIN_LOG_LEVEL=0 \
+  TPU_STDERR_LOG_LEVEL=0 \
+  JAX_PLATFORMS=proxy,cpu \
+  JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 \
+  ENABLE_PATHWAYS_PERSISTENCE=1 \
+  python3 -m maxtext.trainers.post_train.rl.train_rl src/maxtext/configs/post_train/rl.yml  \
+  model_name=gpt-oss-20b  \
+  tokenizer_path=unsloth/gpt-oss-20b-BF16   \
+  load_parameters_path=${MAXTEXT_CKPT_PATH}  \
+  run_name=${WORKLOAD_NAME}  \
+  base_output_directory=${BASE_OUTPUT_DIRECTORY}  \
+  checkpoint_storage_use_ocdbt=False \
+  checkpoint_storage_use_zarr3=False \
+  rollout_data_parallelism=2 \
+  rollout_tensor_parallelism=8 \
+  hbm_utilization_vllm=0.8 \
+  batch_size=8 profiler=xplane \
+  profiler_steps=2 \
+  num_batches=500 \
+  base_emb_dim=2880 \
+  vocab_size=201088 \
+  batch_size=16 \
+  train_micro_batch_size=16 \
+  rollout_micro_batch_size=16 \
+  enable_dp_attention=False \
+  async_scheduling=True \
+  chips_per_vm=4 \
+  chat_template_path=maxtext/examples/chat_templates/gpt_oss_rl.json \
+  enable_tunix_perf_metrics=True \
+"
+
+# Workload Creation
+xpk workload create-pathways \
+  --cluster=$CLUSTER_NAME \
+  --project=$PROJECT_ID \
+  --zone=$ZONE \
+  --priority=medium \
+  --max-restarts=0 \
+  --tpu-type=$TPU_TYPE \
+  --num-slices=1 \
+  --docker-image="${DOCKER_IMAGE}" \
+  --workload="${WORKLOAD_NAME}" \
+  --custom-pathways-proxy-server-args='${XLA_FLAGS}' \
+  --command="${MAXTEXT_COMMAND}"
+
