From 1c20e9c59920080801c71f1e5b334fa25971f886 Mon Sep 17 00:00:00 2001 From: Copybara Date: Thu, 23 Jul 2026 18:28:02 +0000 Subject: [PATCH] Add RL GRPO gemma4-26b tutorial and fix doc conversion links --- docs/tutorials/post_training_index.md | 4 + docs/tutorials/posttraining/rl.md | 2 +- docs/tutorials/posttraining/rl_gemma4_26b.md | 152 +++++++++++++++ docs/tutorials/posttraining/rl_gemma4_e4b.md | 6 +- docs/tutorials/posttraining/rl_gptoss_20b.md | 6 +- docs/tutorials/posttraining/rl_qwen3_30b.md | 6 +- .../rl/scripts/run_gemma4_26b_rl.sh | 178 ++++++++++++++++++ .../rl/scripts/run_gemma4_e4b_rl.sh | 2 +- 8 files changed, 345 insertions(+), 11 deletions(-) create mode 100644 docs/tutorials/posttraining/rl_gemma4_26b.md create mode 100644 src/maxtext/trainers/post_train/rl/scripts/run_gemma4_26b_rl.sh diff --git a/docs/tutorials/post_training_index.md b/docs/tutorials/post_training_index.md index 03460338d3..32087a37f8 100644 --- a/docs/tutorials/post_training_index.md +++ b/docs/tutorials/post_training_index.md @@ -38,6 +38,8 @@ MaxText was co-designed with key Google led innovations to provide a unified pos - [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) + - [RL with gemma4-e4b](./posttraining/rl_gemma4_e4b.md) + - [RL with gemma4-26b](./posttraining/rl_gemma4_26b.md) ## Step by step RL @@ -77,6 +79,8 @@ posttraining/rl.md posttraining/rl_on_multi_host.md posttraining/rl_qwen3_30b.md posttraining/rl_gptoss_20b.md +posttraining/rl_gemma4_e4b.md +posttraining/rl_gemma4_26b.md posttraining/knowledge_distillation.md posttraining/lora.md posttraining/lora_on_multi_host.md diff --git a/docs/tutorials/posttraining/rl.md b/docs/tutorials/posttraining/rl.md index 42109c0a4a..f277d75469 100644 --- a/docs/tutorials/posttraining/rl.md +++ b/docs/tutorials/posttraining/rl.md @@ -97,7 +97,7 @@ export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/ ### 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. +Refer to [Hugging Face to MaxText](../../guides/checkpointing_solutions/convert_checkpoint.md#hugging-face-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 diff --git a/docs/tutorials/posttraining/rl_gemma4_26b.md b/docs/tutorials/posttraining/rl_gemma4_26b.md new file mode 100644 index 0000000000..4ff9e6c043 --- /dev/null +++ b/docs/tutorials/posttraining/rl_gemma4_26b.md @@ -0,0 +1,152 @@ + + +# Reinforcement Learning with gemma4-26b on Multi-Host TPUs + +This tutorial provides step-by-step instructions for setting up the environment +and training the gemma4-26b (26B-A4B MoE) model with GRPO on the [OpenMathInstruct-2 dataset](https://huggingface.co/datasets/nvidia/OpenMathInstruct-2) on a Cloud TPU v7x (Ironwood) GKE cluster using a `tpu7x-4x4x4` (64-chip) slice. + +## 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 (the `google/gemma-4-26B-A4B` and `google/gemma-4-26B-A4B-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)). +- **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 `gemma4-26b` 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](../../guides/checkpointing_solutions/convert_checkpoint.md#hugging-face-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 (Gemma 4 26B specifics):** +> +> - Run the conversion with `scan_layers=False` so the resulting unscanned checkpoint matches the `scan_layers=False` setting used for the RL run below. +> - This RL recipe fine-tunes the **base** model (`google/gemma-4-26B-A4B`), not the instruction-tuned default. Pass `--hf_model_path=google/gemma-4-26B-A4B` explicitly when running `to_maxtext` — otherwise MaxText defaults to `HF_IDS[gemma4-26b]` = `google/gemma-4-26b-a4b-it`. + +## Chat Template Configuration + +Unlike an instruction-tuned tokenizer, the `google/gemma-4-26B-A4B` base tokenizer does **not** ship with a chat template, so the RL run must supply one explicitly. The [run_gemma4_26b_rl.sh](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_26b_rl.sh) script points at two files bundled with the repo (and therefore baked into your Docker image): + +- `data_template_path=maxtext/examples/chat_templates/openmathinstruct2_rl.json` — a stripped-down data template that injects only the system prompt and question (no turn markers). This avoids the doubled `` delimiters that would occur if the default `gsm8k_rl.json` (which bakes literal Gemma turn markers into the message content) were combined with the tokenizer's Jinja chat template. +- `chat_template_path=maxtext/examples/chat_templates/gemma-3-27b-chat_template.json` — the Gemma 3 chat template, wrapped in a JSON object with a `chat_template` key. It is applied by the tokenizer as the single source of truth for turn-marker formatting. + +Both files are already included under `src/maxtext/examples/chat_templates/`, so no additional setup is required. If you want to customize the prompt formatting, edit these files (or point the config at your own) before building the Docker image. + +## 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_gemma4_26b_rl.sh +``` + +> **Note:** The `run_gemma4_26b_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. + +```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](../../guides/checkpointing_solutions/convert_checkpoint.md#maxtext-to-hugging-face) to convert a MaxText checkpoint back to Hugging Face format. You can find an example script to convert the `gemma4-26b` model to Hugging Face format [here](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/gemma4/26b/test_gemma4_to_hf.sh). + +> **Note (Gemma 4 26B specifics):** Because this recipe fine-tunes the **base** model, pass `--hf_model_path=google/gemma-4-26B-A4B` to `to_huggingface` so the exported checkpoint bundles the base tokenizer. Without it, `to_huggingface` sources the tokenizer from `HF_IDS[gemma4-26b]` = `google/gemma-4-26b-a4b-it` (the instruction-tuned model). Also keep `scan_layers=False`, matching the RL run that produced the checkpoint. diff --git a/docs/tutorials/posttraining/rl_gemma4_e4b.md b/docs/tutorials/posttraining/rl_gemma4_e4b.md index 37cfd2f23b..c5a5fffffa 100644 --- a/docs/tutorials/posttraining/rl_gemma4_e4b.md +++ b/docs/tutorials/posttraining/rl_gemma4_e4b.md @@ -72,7 +72,7 @@ export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/ ### 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. +Refer to [Hugging Face to MaxText](../../guides/checkpointing_solutions/convert_checkpoint.md#hugging-face-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 @@ -96,7 +96,7 @@ 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_maxtext.md). +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 @@ -147,6 +147,6 @@ For a complete list of collected metrics, see the [Tunix Metrics Documentation]( ## 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. +Refer to [MaxText to Hugging Face](../../guides/checkpointing_solutions/convert_checkpoint.md#maxtext-to-hugging-face) to convert a MaxText checkpoint back to Hugging Face format. > **Note (Gemma 4 E4B specifics):** Because this recipe fine-tunes the **base** model, pass `--hf_model_path=google/gemma-4-E4B` to `to_huggingface` so the exported checkpoint bundles the base tokenizer. Without it, `to_huggingface` sources the tokenizer from `HF_IDS[gemma4-e4b]` = `google/gemma-4-E4B-it` (the instruction-tuned model). Also keep `scan_layers=False`, since the `gemma4_small` decoder block is not compatible with scanned layers. diff --git a/docs/tutorials/posttraining/rl_gptoss_20b.md b/docs/tutorials/posttraining/rl_gptoss_20b.md index 0869456056..be7f9bccd4 100644 --- a/docs/tutorials/posttraining/rl_gptoss_20b.md +++ b/docs/tutorials/posttraining/rl_gptoss_20b.md @@ -72,7 +72,7 @@ export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/ ### 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. +Refer to [Hugging Face to MaxText](../../guides/checkpointing_solutions/convert_checkpoint.md#hugging-face-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 @@ -84,7 +84,7 @@ 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_maxtext.md). +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 @@ -133,4 +133,4 @@ For a complete list of collected metrics, see the [Tunix Metrics Documentation]( ## 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. +Refer to [MaxText to Hugging Face](../../guides/checkpointing_solutions/convert_checkpoint.md#maxtext-to-hugging-face) to convert a MaxText checkpoint back to Hugging Face format. diff --git a/docs/tutorials/posttraining/rl_qwen3_30b.md b/docs/tutorials/posttraining/rl_qwen3_30b.md index 16a050d80e..d06356e391 100644 --- a/docs/tutorials/posttraining/rl_qwen3_30b.md +++ b/docs/tutorials/posttraining/rl_qwen3_30b.md @@ -72,7 +72,7 @@ export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/ ### 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. 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. +Refer to [Hugging Face to MaxText](../../guides/checkpointing_solutions/convert_checkpoint.md#hugging-face-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 export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/0/items @@ -84,7 +84,7 @@ 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_maxtext.md). +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 @@ -135,4 +135,4 @@ For a complete list of collected metrics, see the [Tunix Metrics Documentation]( ## 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. 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). +Refer to [MaxText to Hugging Face](../../guides/checkpointing_solutions/convert_checkpoint.md#maxtext-to-hugging-face) 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). diff --git a/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_26b_rl.sh b/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_26b_rl.sh new file mode 100644 index 0000000000..48487bffcc --- /dev/null +++ b/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_26b_rl.sh @@ -0,0 +1,178 @@ +#!/bin/bash + +# This script launches a Reinforcement Learning (RL) training workload for the +# gemma4-26b 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 v7x cluster is deployed +export CLUSTER_NAME="${CLUSTER_NAME:-}" # Name of your v7x (Ironwood) cluster +export ZONE="${ZONE:-}" # Zone where your v7x 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="tpu7x-4x4x4" # v7x (Ironwood) 4x4x4 = 64-chip slice (tpu7x-standard-4t, 4 chips/VM) +export WORKLOAD_NAME="rl-$(date +%Y%m%d-%H%M)" + +# Pathways component images, pinned to specific versions (mirrors the reference config). +# Note: --server-image is used for BOTH the Pathways resource-manager server and the +# workers; in the reference config the worker and pathwaysServer images are identical. +export PATHWAYS_SERVER_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server@sha256:6f5031086c3d603f275763c933e85e587e9332be1642fff515e9ecb52374950b" +export PATHWAYS_PROXY_SERVER_IMAGE="us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server@sha256:8ffcc48a7bbd093a9071e0842ed7055d09d140d75f2b6c0a98bdb76cac77c44e" + +# --- 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 (tuned for v7x / Ironwood) +XLA_FLAGS="--xla_tpu_scoped_vmem_limit_kib=65536 \ +--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_single_sparse_core_for_all_gather_offload=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_enable_all_gather_offload_tracing=true \ +--xla_tpu_use_tc_device_shape_on_sc=True \ +--xla_sc_disable_megacore_partitioning=True \ +--xla_enable_async_all_gather=true \ +--xla_tpu_prefer_async_allgather_to_allreduce=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_enable_sparse_core_collective_aggregator=true" + +# MaxText command +MAXTEXT_COMMAND="JAX_RANDOM_WEIGHTS=1 \ +VLLM_ENABLE_V1_MULTIPROCESSING=0 \ +NEW_MODEL_DESIGN=1 \ +TF_CPP_MIN_LOG_LEVEL=0 \ +ENABLE_PJRT_COMPATIBILITY=true \ +ENABLE_PATHWAYS_PERSISTENCE=1 \ +TPU_BACKEND_TYPE=jax \ +VLLM_LOGGING_LEVEL=WARNING \ +JAX_PLATFORMS=proxy,cpu \ +JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 \ +NUM_SLICES=1 \ +VLLM_ENGINE_READY_TIMEOUT_S=7200 \ +RPA_D_BLOCK_SIZES=1,1024,1,512 \ +TOKENIZERS_PARALLELISM=False \ +python3 -m maxtext.trainers.post_train.rl.train_rl \ +model_name=gemma4-26b \ +tokenizer_type=huggingface \ +tokenizer_path=google/gemma-4-26B-A4B \ +run_name=$WORKLOAD_NAME \ +async_scheduling=True \ +base_output_directory=$BASE_OUTPUT_DIRECTORY \ +chips_per_vm=4 \ +data_template_path=maxtext/examples/chat_templates/openmathinstruct2_rl.json \ +chat_template_path=maxtext/examples/chat_templates/gemma-3-27b-chat_template.json \ +remat_policy=full \ +decoder_layer_input=offload \ +megablox=True \ +sparse_matmul=True \ +use_custom_sort_vjp=True \ +attention=flash \ +use_tokamax_splash=True \ +sa_use_fused_bwd_kernel=True \ +sa_block_q=1024 \ +sa_block_kv=1024 \ +sa_block_kv_compute=512 \ +sa_block_q_dkv=2048 \ +sa_block_kv_dkv=2048 \ +sa_block_kv_dkv_compute=256 \ +mu_dtype=bfloat16 \ +grad_dtype=bfloat16 \ +use_iota_embed=True \ +batch_size=128 \ +num_batches=500 \ +learning_rate_schedule_steps=500 \ +num_test_batches=5 \ +rl.num_generations=8 \ +rl.num_iterations=1 \ +rl.grpo_beta=0.05 \ +rl.grpo_epsilon=0.2 \ +debug=True \ +gradient_clipping_threshold=1.0 \ +decode_sampling_temperature=0.8 \ +decode_sampling_top_k=50 \ +decode_sampling_nucleus_p=0.95 \ +hf_name=default \ +dataset_name=nvidia/OpenMathInstruct-2 \ +hf_train_files=hf://datasets/nvidia/OpenMathInstruct-2/data/train_1M-*.parquet \ +train_split=train_1M \ +eval_dataset_name=nvidia/OpenMathInstruct-2 \ +eval_split=test \ +learning_rate=1e-6 \ +max_prefill_predict_length=512 \ +max_target_length=8192 \ +kv_cache_buffer=512 \ +max_num_seqs=256 \ +max_num_batched_tokens=8192 \ +rollout_data_parallelism=8 \ +rollout_tensor_parallelism=1 \ +rollout_expert_parallelism=8 \ +rl.reshard_chunk_size=50 \ +enable_dp_attention=False \ +hbm_utilization_vllm=0.8 \ +scan_layers=False \ +allow_split_physical_axes=True \ +enable_tunix_perf_metrics=True \ +checkpoint_period=25 \ +max_num_checkpoints_to_keep=30 \ +enable_checkpointing=True \ +train_micro_batch_size=8 \ +rollout_micro_batch_size=128 \ +load_parameters_path=$MAXTEXT_CKPT_PATH \ +vllm_hf_overrides='{architectures: [\"MaxTextForCausalLM\"]}' \ +vllm_additional_config='{\"maxtext_config\": {\"model_name\": \"gemma4-26b\", \"model_call_mode\": \"inference\", \"enable_dp_attention\": false, \"allow_split_physical_axes\": true, \"log_config\": false, \"weight_dtype\": \"bfloat16\", \"prefuse_moe_weights\": 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}" \ + --server-image="${PATHWAYS_SERVER_IMAGE}" \ + --proxy-server-image="${PATHWAYS_PROXY_SERVER_IMAGE}" \ + --custom-pathways-proxy-server-args='${XLA_FLAGS}' \ + --command="${MAXTEXT_COMMAND}" diff --git a/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_e4b_rl.sh b/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_e4b_rl.sh index 4ec84d7b46..40cdafa34f 100644 --- a/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_e4b_rl.sh +++ b/src/maxtext/trainers/post_train/rl/scripts/run_gemma4_e4b_rl.sh @@ -127,7 +127,7 @@ rl.num_generations=8 \ rl.num_iterations=1 \ rl.grpo_beta=0.05 \ rl.grpo_epsilon=0.2 \ -debug.rl=True \ +debug=True \ gradient_clipping_threshold=1.0 \ decode_sampling_temperature=0.8 \ decode_sampling_top_k=50 \