feat(test): add modular E2E TPU testing pipeline scripts for GPTOSS-20B

Author: RexBearIU

docs/superpowers/plans/2026-07-21-gptoss-20b-e2e.md

@@ -0,0 +1,392 @@
+# GPTOSS-20B E2E TPU Testing Implementation Plan
+
+> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
+
+**Goal:** Create a modular suite of E2E TPU testing scripts for GPTOSS-20B, excluding multimodal capabilities and aligning with the structure used by Gemma3-4B.
+
+**Architecture:** Split the monolithic `test_gpt_oss.sh` into modular scripts: `test_gpt_oss_to_mt.sh` (which installs common dependencies), `test_gpt_oss.sh` (pre-training baseline), `test_gpt_oss_sft.sh` (supervised fine-tuning), `test_gpt_oss_rl.sh` (reinforcement learning), and `test_gpt_oss_to_hf.sh` (reverse checkpoint conversion running last).
+
+**Tech Stack:** Bash, Python, JAX, MaxText, PyTorch (for conversion).
+
+## Global Constraints
+- Target directory: `tests/end_to_end/tpu/gpt_oss/20b/`
+- No multimodal functionality (multimodal excluded).
+- Dependencies must be installed in `test_gpt_oss_to_mt.sh` first, so subsequent scripts do not repeat them.
+- Reverse conversion to Hugging Face (`test_gpt_oss_to_hf.sh`) must be the final pipeline step.
+
+---
+
+### Task 1: Checkpoint Conversion to MaxText & Logit Verification
+
+**Files:**
+- Create: `tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh`
+
+- [ ] **Step 1: Create and write the file**
+Write the following contents to `tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh`:
+```bash
+#!/bin/bash
+
+# Converts GPTOSS-20B HuggingFace checkpoint to MaxText format and validates logit correctness.
+
+set -ex
+
+run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+export MODEL_NAME='gpt-oss-20b'
+export TOKENIZER_PATH='openai/gpt-oss-20b'
+
+if [ -z "${BASE_OUTPUT_PATH}" ]; then
+  export BASE_OUTPUT_PATH=gs://runner-maxtext-logs/${MODEL_NAME}
+fi
+BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/}
+echo "Using BASE_OUTPUT_PATH = ${BASE_OUTPUT_PATH}"
+
+# Install dependencies (CPU) required for checkpoint conversion & logit check
+python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+
+if [ -z "${CKPT_DISK_LOCATION}" ]; then
+  export CKPT_BUCKET=gs://maxtext-model-checkpoints/gpt-oss-20b/hf-bf16
+  gcloud storage cp -r ${CKPT_BUCKET} /tmp
+  export CKPT_DISK_LOCATION=/tmp/hf-bf16
+fi
+
+# 1. Convert to scanned checkpoint (for training)
+JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.standalone_scripts.convert_gpt_oss_ckpt \
+    --base-model-path ${CKPT_DISK_LOCATION} \
+    --maxtext-model-path ${BASE_OUTPUT_PATH}/scanned/${run_id} \
+    --model-size ${MODEL_NAME}
+
+SCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/scanned/${run_id}/0/items
+echo "Scanned checkpoint path: ${SCANNED_CKPT_PATH}"
+
+# 2. Convert to unscanned checkpoint (for inference)
+JAX_PLATFORMS=cpu python3 -m maxtext.checkpoint_conversion.standalone_scripts.convert_gpt_oss_unscanned_ckpt \
+    --base-model-path ${CKPT_DISK_LOCATION} \
+    --maxtext-model-path ${BASE_OUTPUT_PATH}/unscanned/${run_id} \
+    --model-size ${MODEL_NAME}
+
+UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/${run_id}/0/items
+echo "Unscanned checkpoint path: ${UNSCANNED_CKPT_PATH}"
+
+# 3. Logit correctness check
+python3 -m tests.utils.forward_pass_logit_checker \
+    "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \
+    base_output_directory=${BASE_OUTPUT_PATH} \
+    run_name=forward_logits_check \
+    model_name=${MODEL_NAME} \
+    load_parameters_path=${UNSCANNED_CKPT_PATH} \
+    scan_layers=false \
+    attention=dot_product \
+    sparse_matmul=True \
+    megablox=True \
+    per_device_batch_size=1 \
+    max_target_length=4 \
+    max_prefill_predict_length=4 \
+    dtype=float32 \
+    --atol=0.1 --rtol=0.1 \
+    --max_kl_div=3e-4
+```
+
+- [ ] **Step 2: Make the script executable**
+Run: `chmod +x tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh`
+
+- [ ] **Step 3: Run baseline logit check to verify conversion fails/passes properly**
+Run: `bash tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh`
+Expected: Completion of dependencies installation and starting conversion. (If CPU OOMs, note it as an issue).
+
+- [ ] **Step 4: Commit**
+```bash
+git add tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh
+git commit -m "test: add GPTOSS-20B HF to MT checkpoint conversion test script"
+```
+
+---
+
+### Task 2: Refactor Baseline Pre-training Script
+
+**Files:**
+- Modify: `tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh`
+
+- [ ] **Step 1: Refactor the file contents**
+Replace the content of `tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh` with the following simplified pre-training flow:
+```bash
+#!/bin/bash
+
+# Validates the GPTOSS-20B pre-training pipeline starting from converted MaxText checkpoint.
+
+set -ex
+
+run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+export MODEL_NAME='gpt-oss-20b'
+export TOKENIZER_PATH='openai/gpt-oss-20b'
+
+if [ -z "${BASE_OUTPUT_PATH}" ]; then
+  export BASE_OUTPUT_PATH=gs://runner-maxtext-logs/${MODEL_NAME}
+fi
+BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/}
+
+export SCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/scanned/${run_id}/0/items
+export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/${run_id}/0/items
+
+# 1. Run Pre-training using synthetic dataset with Megablocks
+python3 -m maxtext.trainers.pre_train.train \
+    "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \
+    base_output_directory=${BASE_OUTPUT_PATH}/train \
+    run_name=${run_id} \
+    model_name=${MODEL_NAME} \
+    tokenizer_type=huggingface \
+    tokenizer_path=${TOKENIZER_PATH} \
+    dataset_type=synthetic \
+    enable_checkpointing=true \
+    async_checkpointing=false \
+    load_parameters_path=${SCANNED_CKPT_PATH} \
+    attention=flash \
+    sparse_matmul=True \
+    megablox=True \
+    dtype=bfloat16 \
+    weight_dtype=bfloat16 \
+    per_device_batch_size=4 \
+    steps=5 \
+    max_target_length=1024 \
+    ici_fsdp_parallelism=4 \
+    gcs_metrics=true
+
+# 2. Run Verification Decoding from the converted checkpoint
+python3 -m maxtext.inference.decode \
+    "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \
+    base_output_directory=${BASE_OUTPUT_PATH} \
+    run_name=decode \
+    model_name=${MODEL_NAME} \
+    tokenizer_type=huggingface \
+    tokenizer_path=${TOKENIZER_PATH} \
+    load_parameters_path=${UNSCANNED_CKPT_PATH} \
+    scan_layers=False \
+    attention=dot_product \
+    sparse_matmul=True \
+    megablox=True \
+    dtype=bfloat16 \
+    weight_dtype=bfloat16 \
+    per_device_batch_size=1 \
+    max_prefill_predict_length=64 \
+    max_target_length=128 \
+    prompt="I love to" \
+    ici_fsdp_parallelism=1 \
+    ici_tensor_parallelism=4
+```
+
+- [ ] **Step 2: Commit**
+```bash
+git add tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh
+git commit -m "test: refactor GPTOSS-20B pre-training baseline test script"
+```
+
+---
+
+### Task 3: Supervised Fine-Tuning Pipeline
+
+**Files:**
+- Create: `tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh`
+
+- [ ] **Step 1: Create and write the file**
+Write the following contents to `tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh`:
+```bash
+#!/bin/bash
+
+# Validates the GPTOSS-20B Supervised Fine-Tuning (SFT) pipeline.
+
+set -ex
+
+run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+export MODEL_NAME='gpt-oss-20b'
+export TOKENIZER_PATH='openai/gpt-oss-20b'
+
+if [ -z "${BASE_OUTPUT_PATH}" ]; then
+  export BASE_OUTPUT_PATH=gs://runner-maxtext-logs/${MODEL_NAME}
+fi
+BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/}
+
+export SCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/scanned/${run_id}/0/items
+
+# 1. Run Supervised Fine-Tuning
+python3 -m maxtext.trainers.post_train.sft.train_sft_native \
+    "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs/post_train}"/sft.yml \
+    base_output_directory=${BASE_OUTPUT_PATH}/sft \
+    run_name=${run_id} \
+    model_name=${MODEL_NAME} \
+    tokenizer_type=huggingface \
+    tokenizer_path=${TOKENIZER_PATH} \
+    dataset_type=hf \
+    enable_checkpointing=true \
+    async_checkpointing=false \
+    load_parameters_path=${SCANNED_CKPT_PATH} \
+    scan_layers=True \
+    attention=flash \
+    sparse_matmul=True \
+    megablox=True \
+    dtype=bfloat16 \
+    weight_dtype=bfloat16 \
+    per_device_batch_size=4 \
+    steps=5 \
+    max_target_length=1024 \
+    ici_fsdp_parallelism=1 \
+    ici_expert_parallelism=4 \
+    gcs_metrics=true \
+    abort_on_nan_loss=false
+
+# 2. Run Decoding on the newly produced SFT checkpoint
+python3 -m maxtext.inference.decode \
+    "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \
+    base_output_directory=${BASE_OUTPUT_PATH} \
+    run_name=decode_sft \
+    model_name=${MODEL_NAME} \
+    tokenizer_type=huggingface \
+    tokenizer_path=${TOKENIZER_PATH} \
+    load_parameters_path=${BASE_OUTPUT_PATH}/sft/${run_id}/checkpoints/4/items \
+    scan_layers=True \
+    attention=dot_product \
+    sparse_matmul=True \
+    megablox=True \
+    dtype=bfloat16 \
+    weight_dtype=bfloat16 \
+    per_device_batch_size=1 \
+    max_prefill_predict_length=64 \
+    max_target_length=128 \
+    prompt="I love to" \
+    ici_fsdp_parallelism=1 \
+    ici_tensor_parallelism=4
+```
+
+- [ ] **Step 2: Make the script executable**
+Run: `chmod +x tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh`
+
+- [ ] **Step 3: Commit**
+```bash
+git add tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh
+git commit -m "test: add GPTOSS-20B SFT pipeline test script"
+```
+
+---
+
+### Task 4: Reinforcement Learning Pipeline
+
+**Files:**
+- Create: `tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh`
+
+- [ ] **Step 1: Create and write the file**
+Write the following contents to `tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh`:
+```bash
+#!/bin/bash
+
+# Validates the GPTOSS-20B Reinforcement Learning (RL) pipeline using GRPO.
+
+set -ex
+
+run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+export MODEL_NAME='gpt-oss-20b'
+export TOKENIZER_PATH='openai/gpt-oss-20b'
+
+if [ -z "${BASE_OUTPUT_PATH}" ]; then
+  export BASE_OUTPUT_PATH=gs://runner-maxtext-logs/${MODEL_NAME}
+fi
+BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/}
+
+export SCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/scanned/${run_id}/0/items
+
+# 1. Run GRPO Reinforcement Learning
+python3 -m maxtext.trainers.post_train.rl.train_rl \
+    base_output_directory=${BASE_OUTPUT_PATH}/rl \
+    load_parameters_path=${SCANNED_CKPT_PATH} \
+    run_name=${run_id} \
+    rl.loss_algo='grpo' \
+    scan_layers=true \
+    num_batches=5 \
+    batch_size=1 \
+    num_test_batches=5 \
+    model_name=${MODEL_NAME} \
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False \
+    rollout_tensor_parallelism=1 \
+    vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
+    vllm_additional_config='{"maxtext_config": {"model_name": "gpt-oss-20b", "log_config": "false"}}'
+
+# 2. Run Verification Decoding on the newly produced actor checkpoint
+python3 -m maxtext.inference.decode \
+    "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \
+    base_output_directory=${BASE_OUTPUT_PATH} \
+    run_name=decode_rl \
+    model_name=${MODEL_NAME} \
+    tokenizer_type=huggingface \
+    tokenizer_path=${TOKENIZER_PATH} \
+    load_parameters_path=${BASE_OUTPUT_PATH}/rl/${run_id}/checkpoints/actor/4/items \
+    scan_layers=True \
+    attention=dot_product \
+    sparse_matmul=True \
+    megablox=True \
+    dtype=bfloat16 \
+    weight_dtype=bfloat16 \
+    per_device_batch_size=1 \
+    max_prefill_predict_length=64 \
+    max_target_length=128 \
+    prompt="I love to" \
+    ici_fsdp_parallelism=1 \
+    ici_tensor_parallelism=4
+```
+
+- [ ] **Step 2: Make the script executable**
+Run: `chmod +x tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh`
+
+- [ ] **Step 3: Commit**
+```bash
+git add tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh
+git commit -m "test: add GPTOSS-20B RL pipeline test script"
+```
+
+---
+
+### Task 5: Checkpoint Conversion to Hugging Face
+
+**Files:**
+- Create: `tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_hf.sh`
+
+- [ ] **Step 1: Create and write the file**
+Write the following contents to `tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_hf.sh`:
+```bash
+#!/bin/bash
+
+# Converts a MaxText checkpoint to a Hugging Face model checkpoint for GPTOSS-20B.
+
+set -ex
+
+run_id=$1
+CKPT_PATH=$2
+SCAN_LAYERS=${3:-false}
+
+export MODEL_NAME='gpt-oss-20b'
+
+if [ -z "${BASE_OUTPUT_PATH}" ]; then
+  export BASE_OUTPUT_PATH=gs://runner-maxtext-logs/${MODEL_NAME}
+fi
+BASE_OUTPUT_PATH=${BASE_OUTPUT_PATH%/}
+
+if [ "${SCAN_LAYERS,,}" = "true" ]; then
+    scan_status="scanned"
+else
+    scan_status="unscanned"
+fi
+
+python3 -m maxtext.checkpoint_conversion.to_huggingface \
+    model_name=${MODEL_NAME} \
+    tokenizer_type="huggingface" \
+    load_parameters_path=${CKPT_PATH} \
+    base_output_directory=${BASE_OUTPUT_PATH}/to_huggingface/${scan_status}/${run_id} \
+    use_multimodal=false \
+    scan_layers=$SCAN_LAYERS
+```
+
+- [ ] **Step 2: Make the script executable**
+Run: `chmod +x tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_hf.sh`
+
+- [ ] **Step 3: Commit**
+```bash
+git add tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_hf.sh
+git commit -m "test: add GPTOSS-20B to HF checkpoint conversion test script"
+```