Refactor RL tutorials

PiperOrigin-RevId: 948616994

Author: SurbhiJainUSC

docs/guides/checkpointing_solutions/convert_checkpoint.md

@@ -86,6 +86,8 @@ You can find your converted checkpoint files under `${BASE_OUTPUT_DIRECTORY}/0/i
 - `--hf_model_path` (Optional): Specifies a customized remote directory or local directory containing the model weights. If unspecified, we use the [default Hugging Face repository ID](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/utils/globals.py) (e.g., openai/gpt-oss-20b). This is necessary for locally dequantized models like GPT-OSS or DeepSeek.
 - `--save_dtype` (Optional): Specifies the data type of saved model weights. Default to `bfloat16` to save memory.
 
+(maxtext-to-hf)=
+
 ## MaxText to Hugging Face
 
 Use the `to_huggingface.py` script to convert a MaxText checkpoint into the Hugging Face format. This is useful for sharing your models or integrating them with the Hugging Face ecosystem.

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 an Viperfish 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.
 
 ## Prerequisites
 
@@ -30,10 +30,6 @@ Before starting, ensure you have:
 - 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
@@ -53,19 +49,6 @@ 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
@@ -89,51 +72,29 @@ export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/
 
 ### 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.
+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
-# 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
+export MAXTEXT_CKPT_PATH=<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
-# 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
+# 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
-bash scripts/run_gptoss_20b_rl.sh
-
-# Deactivate the virtual environment
-deactivate
-rm -rf runner_venv
+run_tutorial maxtext/trainers/post_train/rl/scripts/run_gptoss_20b_rl.sh
 ```
 
 ### Monitor your workload
@@ -156,7 +117,7 @@ Alternatively, after running the bash script, you will also get a link to the Go
 
 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.
+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:
 
@@ -172,30 +133,4 @@ For a complete list of collected metrics, see the [Tunix Metrics Documentation](
 
 ## 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
-```
+Refer to [MaxText to Hugging Face](maxtext-to-hf) to convert a MaxText checkpoint back to Hugging Face format.

docs/tutorials/posttraining/rl_qwen3_30b.md

@@ -30,10 +30,6 @@ Before starting, ensure you have:
 - 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
@@ -53,19 +49,6 @@ 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
@@ -89,42 +72,29 @@ export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/
 
 ### Option 2: Converting from a Hugging Face checkpoint
 
-> **Note:** Converting the 30B model requires approximately 62 GB of free disk space to download its safetensors. Please verify you have sufficient space before running the conversion script.
+Refer to [Hugging Face to MaxText](hf-to-maxtext) to convert a Hugging Face checkpoint to MaxText format. You can find an example script to convert `qwen3-30b-a3b-base`model to MaxText format [here](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_mt.sh). After conversion finishes, set `MAXTEXT_CKPT_PATH` to the converted MaxText checkpoint path.
 
 ```bash
-# Optional: If you run out of disk space when downloading Hugging Face safetensors,
-# customize your "HF_HOME" to redirect the cache to a larger or mounted disk (e.g., on a TPU VM).
-# export HF_HOME="/dev/shm/huggingface_tmp"
+export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items
+```
 
-# 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
+> **Note:** Converting the 30B model requires approximately 62 GB of free disk space to download its safetensors. Please verify you have sufficient space before running the conversion script.
 
-# Run the conversion script to convert the Hugging Face checkpoint to MaxText format
-bash scripts/run_qwen3_30b_hf_to_maxtext.sh
+## Run RL Workload
 
-# Deactivate the virtual environment
-deactivate
-rm -rf tpu_venv
-```
+### Build and Upload MaxText Docker Image
 
-## Run RL Workload
+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
-# 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
+# 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
-bash scripts/run_qwen3_30b_rl.sh
-
-# Deactivate the virtual environment
-deactivate
-rm -rf runner_venv
+run_tutorial maxtext/trainers/post_train/rl/scripts/run_qwen3_30b_rl.sh
 ```
 
 ### Monitor your workload
@@ -147,7 +117,7 @@ Alternatively, after running the bash script, you will also get a link to the Go
 
 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_qwen3_30b_rl.sh](../../../scripts/run_qwen3_30b_rl.sh) script for this tutorial workload. When enabled, Tunix automatically collects and uploads these metrics to TensorBoard.
+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:
 
@@ -163,18 +133,6 @@ For a complete list of collected metrics, see the [Tunix Metrics Documentation](
 
 ## Convert Checkpoint to Hugging Face Format
 
-After training, you may want to convert your MaxText checkpoint back to Hugging Face format. Use the following script 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
-
-# Run the conversion script to convert the MaxText checkpoint back to Hugging Face format 
-bash scripts/run_qwen3_30b_maxtext_to_hf.sh
+## Convert Checkpoint to Hugging Face Format
 
-# Deactivate the virtual environment
-deactivate
-rm -rf tpu_venv
-```
+Refer to [MaxText to Hugging Face](maxtext-to-hf) to convert a MaxText checkpoint back to Hugging Face format. You can find an example script to convert `qwen3-30b-a3b-base`model to Hugging Face format [here](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_hf.sh).

pyproject.toml

@@ -42,11 +42,10 @@ allow-direct-references = true
 [tool.hatch.build.targets.wheel]
 packages = ["src/maxtext", "src/dependencies"]
 
-
-
 [project.scripts]
 install_tpu_pre_train_extra_deps = "dependencies.scripts.install_pre_train_extra_deps:main"
 install_cuda12_pre_train_extra_deps = "dependencies.scripts.install_pre_train_extra_deps:main"
 install_tpu_post_train_extra_deps = "dependencies.scripts.install_post_train_extra_deps:main"
 build_maxtext_docker_image = "dependencies.scripts.build_maxtext_docker_image:main"
 upload_maxtext_docker_image = "dependencies.scripts.upload_maxtext_docker_image:main"
+run_tutorial = "dependencies.scripts.run_tutorial:main"