Add E2E TPU testing scripts for GPTOSS-20B

tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh

@@ -1,58 +1,80 @@
 #!/bin/bash
 
-# Validates the GPTOSS-20B pre-training pipeline starting from converted MaxText checkpoint.
+# Validates the GPT-OSS-20b pre-training pipeline using a pre-converted MaxText checkpoint.
+
+# The flow of this script is as follows:
+# 1. Run inference on the pre-converted checkpoint.
+# 2. Run pre-training starting from the pre-converted checkpoint.
+# 3. Run inference on the checkpoint produced by the pre-training run.
+
+# Usage:
+# export HF_TOKEN=<your Hugging Face access token>
+# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S)
+# bash test_gpt_oss_to_mt.sh $RUN_ID
+# bash test_gpt_oss.sh $RUN_ID
 
 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'
+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%/}
+# Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you have the scanned and unscanned checkpoints stored
+BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}
+UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items
 
-export SCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/scanned/${run_id}/0/items
-export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/${run_id}/0/items
+# Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
+DATASET_PATH=gs://maxtext-dataset
 
-export SPARSE_MATMUL="True"
-export MEGABLOX="True"
-export PRETRAIN_ATTENTION="flash"
+# Step 1: Run inference on the original checkpoint converted from Hugging Face
+    python3 -m maxtext.inference.decode \
+    model_name=${MODEL_NAME} \
+    tokenizer_type="huggingface" \
+    load_parameters_path=${UNSCANNED_CKPT_PATH} \
+    per_device_batch_size=1 \
+    run_name=${run_id} \
+    max_prefill_predict_length=8 \
+    max_target_length=16 \
+    steps=1 \
+    async_checkpointing=false \
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False \
+    scan_layers=false \
+    prompt='I love to' \
+    attention=\'dot_product\'
 
-# 1. Run Pre-training using synthetic dataset with Megablocks
+# Step 2: Run Pre-training on the converted checkpoint
+# We can also run training by using the scanned converted checkpoint
+# Note that scanned checkpoint helps with efficient training
 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 \
+    base_output_directory=${BASE_OUTPUT_DIRECTORY}/train \
+    dataset_path=${DATASET_PATH} \
+    tokenizer_type="huggingface" \
+    load_parameters_path=${UNSCANNED_CKPT_PATH} \
+    per_device_batch_size=1 \
     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=${PRETRAIN_ATTENTION} \
-    sparse_matmul=${SPARSE_MATMUL} \
-    megablox=${MEGABLOX} \
-    dtype=bfloat16 \
-    weight_dtype=bfloat16 \
-    per_device_batch_size=4 \
-    steps=5 \
     max_target_length=1024 \
-    ici_fsdp_parallelism=4 \
-    gcs_metrics=true
+    steps=5 \
+    weight_dtype=bfloat16 \
+    async_checkpointing=false \
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False \
+    model_name=${MODEL_NAME} \
+    scan_layers=false \
+    use_multimodal=false
 
-# 2. Run Verification Decoding from the converted checkpoint
-python3 -m maxtext.inference.decode \
-    base_output_directory=${BASE_OUTPUT_PATH} \
-    run_name=decode \
+# Step 3: Run inference on the checkpoint generated from the previous run
+    python3 -m maxtext.inference.decode \
     model_name=${MODEL_NAME} \
-    tokenizer_path=${TOKENIZER_PATH} \
-    load_parameters_path=${BASE_OUTPUT_PATH}/train/${run_id}/checkpoints/4/items \
-    scan_layers=True \
-    attention=dot_product \
-    sparse_matmul=${SPARSE_MATMUL} \
-    megablox=${MEGABLOX} \
-    prompt="I love to" \
-    ici_tensor_parallelism=4
+    tokenizer_type="huggingface" \
+    load_parameters_path=${BASE_OUTPUT_DIRECTORY}/train/${run_id}/checkpoints/4/items \
+    per_device_batch_size=1 \
+    run_name=${run_id} \
+    max_prefill_predict_length=8 \
+    max_target_length=16 \
+    steps=1 \
+    async_checkpointing=false \
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False \
+    scan_layers=false \
+    prompt='I love to' \
+    attention=\'dot_product\'

tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_lora.sh

@@ -1,87 +1,74 @@
 #!/bin/bash
 
-# Validates the GPT-OSS-20B LoRA pipeline using a pre-converted MaxText checkpoint.
+# Validates the GPT-OSS-20b LoRA pipeline using a pre-converted MaxText checkpoint.
 
 # The flow of this script is as follows:
-# 1. Run LoRA starting from the pre-converted checkpoint.
-# 2. Run inference on the checkpoint produced by the LoRA run.
-# 3. Convert the checkpoint produced by the LoRA run back to HuggingFace format.
+# 1. Run inference on the pre-converted checkpoint.
+# 2. Run LoRA starting from the pre-converted checkpoint.
+# 3. Run inference on the checkpoint produced by the LoRA run.
+# 4. Convert the checkpoint produced by the LoRA run back to HuggingFace format.
 
 # Usage:
 # export HF_TOKEN=<your Hugging Face access token>
 # export RUN_ID=$(date +%Y-%m-%d-%H-%M)
 # bash test_gpt_oss_to_mt.sh $RUN_ID
 # bash test_gpt_oss_lora.sh $RUN_ID
 
+
 set -ex
 
 run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+use_pathways=${2:-false}
 export MODEL_NAME='gpt-oss-20b'
-export TOKENIZER_PATH='openai/gpt-oss-20b'
 
-export BASE_OUTPUT_PATH=gs://runner-maxtext-logs/gpt-oss-20b
-SCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/scanned/${run_id}/0/items
+# Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you have the scanned and unscanned checkpoints stored
+BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}
+UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items
+SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned/${run_id}/0/items
 
-SPARSE_MATMUL="True"
-MEGABLOX="True"
-LORA_ATTENTION="flash"
+# Step 1: Run inference on the original checkpoint converted from Hugging Face
+python3 -m maxtext.inference.vllm_decode \
+    model_name=${MODEL_NAME} \
+    load_parameters_path=${UNSCANNED_CKPT_PATH} \
+    vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
+    hbm_utilization_vllm=0.7 \
+    prompt="Suggest some famous landmarks in London." \
+    use_chat_template=True \
+    scan_layers=false \
+    enable_single_controller=${use_pathways} \
+    prefuse_moe_weights=True \
+    ici_tensor_parallelism=8
 
-# Step 1: Run LoRA SFT on the converted checkpoint
+# Step 2: Run LoRA on the converted checkpoint
 python3 -m maxtext.trainers.post_train.sft.train_sft \
-    base_output_directory=${BASE_OUTPUT_PATH}/lora \
-    run_name=${run_id} \
-    model_name=${MODEL_NAME} \
-    tokenizer_path=${TOKENIZER_PATH} \
-    dataset_name=lucasmccabe-lmt/math_alpaca \
-    dataset_type=hf \
-    enable_checkpointing=true \
-    async_checkpointing=false \
+    base_output_directory=${BASE_OUTPUT_DIRECTORY}/lora \
     load_parameters_path=${SCANNED_CKPT_PATH} \
-    scan_layers=True \
-    attention=${LORA_ATTENTION} \
-    sparse_matmul=${SPARSE_MATMUL} \
-    megablox=${MEGABLOX} \
-    dtype=bfloat16 \
-    weight_dtype=bfloat16 \
-    per_device_batch_size=4 \
+    tokenizer_path='unsloth/gpt-oss-20b-BF16' \
+    per_device_batch_size=1 \
+    run_name=${run_id} \
     steps=5 \
-    max_target_length=1024 \
-    ici_fsdp_parallelism=1 \
-    ici_expert_parallelism=4 \
-    gcs_metrics=true \
-    abort_on_nan_loss=false \
+    scan_layers=true \
+    model_name=${MODEL_NAME} \
+    learning_rate=3e-6 \
     lora.enable_lora=True \
     lora.lora_rank=16 \
     lora.lora_alpha=32.0 \
-    enable_nnx=True \
-    pure_nnx_decoder=True \
-    enable_single_controller=True \
-    checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False
+    enable_single_controller=${use_pathways} \
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False
 
-# Step 2: Run inference decoding on the checkpoint generated from the previous run
-python3 -m maxtext.inference.decode \
-    base_output_directory=${BASE_OUTPUT_PATH} \
-    run_name=decode_lora \
+# Step 3: Run inference on the checkpoint generated from the previous run
+python3 -m maxtext.inference.vllm_decode \
+    --use_tunix=True \
     model_name=${MODEL_NAME} \
-    tokenizer_path=${TOKENIZER_PATH} \
     load_parameters_path=${SCANNED_CKPT_PATH} \
     lora.enable_lora=True \
-    lora.lora_restore_path=${BASE_OUTPUT_PATH}/lora/${run_id}/checkpoints/4/items \
+    lora.lora_restore_path=${BASE_OUTPUT_DIRECTORY}/lora/${run_id}/checkpoints/5/model_params \
     lora.lora_rank=16 \
     lora.lora_alpha=32.0 \
-    scan_layers=True \
-    attention=dot_product \
-    sparse_matmul=${SPARSE_MATMUL} \
-    megablox=${MEGABLOX} \
-    prompt="I love to" \
-    ici_tensor_parallelism=4
-
-# Step 3: Convert the checkpoint from MaxText format to Hugging Face format
-python3 -m maxtext.checkpoint_conversion.to_huggingface \
-    model_name=${MODEL_NAME} \
-    load_parameters_path=${SCANNED_CKPT_PATH} \
-    lora.lora_restore_path=${BASE_OUTPUT_PATH}/lora/${run_id}/checkpoints/4/items \
-    base_output_directory=${BASE_OUTPUT_PATH}/to_huggingface/unscanned/${run_id} \
+    vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
+    hbm_utilization_vllm=0.6 \
+    prompt="Suggest some famous landmarks in London." \
+    use_chat_template=True \
     scan_layers=true \
-    enable_nnx=True \
-    pure_nnx_decoder=True
+    enable_single_controller=${use_pathways}

tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh

@@ -1,55 +1,70 @@
 #!/bin/bash
 
-# Validates the GPTOSS-20B Reinforcement Learning (RL) pipeline using GRPO.
+# Validates the GPT-OSS-20b RL pipeline using a pre-converted MaxText checkpoint.
+
+# The flow of this script is as follows:
+# 1. Run inference on the pre-converted checkpoint.
+# 2. Run RL starting from the pre-converted checkpoint.
+# 3. Run inference on the checkpoint produced by the RL run.
+
+# Usage:
+# export HF_TOKEN=<your Hugging Face access token>
+# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S)
+# bash test_gpt_oss_to_mt.sh $RUN_ID
+# bash test_gpt_oss_rl.sh $RUN_ID
 
 set -ex
 
 run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+use_pathways=${2:-false}
 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%/}
+# Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you have the scanned and unscanned checkpoints stored
+BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}
+UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items
+SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned/${run_id}/0/items
 
-export SCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/scanned/${run_id}/0/items
 
-export SPARSE_MATMUL="True"
-export MEGABLOX="True"
-export ATTENTION="flash"
-export VLLM_ADDITIONAL_CONFIG='{"maxtext_config": {"model_name": "gpt-oss-20b", "log_config": "false"}}'
+# Step 1: Run inference on the original checkpoint converted from Hugging Face
+python3 -m maxtext.inference.vllm_decode \
+    model_name=${MODEL_NAME} \
+    load_parameters_path=${UNSCANNED_CKPT_PATH} \
+    vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
+    hbm_utilization_vllm=0.7 \
+    prompt="Suggest some famous landmarks in London." \
+    use_chat_template=True \
+    scan_layers=false \
+    enable_single_controller=${use_pathways} \
+    prefuse_moe_weights=True \
+    ici_tensor_parallelism=8
 
-# 1. Run GRPO Reinforcement Learning
+# Step 2: Run RL on the converted checkpoint
 python3 -m maxtext.trainers.post_train.rl.train_rl \
-    base_output_directory=${BASE_OUTPUT_PATH}/rl \
+    base_output_directory=${BASE_OUTPUT_DIRECTORY}/rl \
     load_parameters_path=${SCANNED_CKPT_PATH} \
     run_name=${run_id} \
     rl.loss_algo='grpo' \
     scan_layers=true \
     num_batches=5 \
-    batch_size=1 \
+    batch_size=16 \
     num_test_batches=5 \
     model_name=${MODEL_NAME} \
+    enable_single_controller=${use_pathways} \
     checkpoint_storage_use_zarr3=False \
     checkpoint_storage_use_ocdbt=False \
-    rollout_tensor_parallelism=1 \
-    attention=${ATTENTION} \
-    sparse_matmul=${SPARSE_MATMUL} \
-    megablox=${MEGABLOX} \
+    rollout_tensor_parallelism=4 \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
-    vllm_additional_config="${VLLM_ADDITIONAL_CONFIG}"
+    vllm_additional_config='{"maxtext_config": {"model_name": "gpt-oss-20b", "log_config": "false", "prefuse_moe_weights": "true"}}'
 
-# 2. Run Verification Decoding on the newly produced actor checkpoint
-python3 -m maxtext.inference.decode \
-    base_output_directory=${BASE_OUTPUT_PATH} \
-    run_name=decode_rl \
+# Step 3: Run inference on the checkpoint generated from the previous run
+python3 -m maxtext.inference.vllm_decode \
     model_name=${MODEL_NAME} \
-    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=${SPARSE_MATMUL} \
-    megablox=${MEGABLOX} \
-    prompt="I love to" \
-    ici_tensor_parallelism=4
+    load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/5/model_params \
+    vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
+    hbm_utilization_vllm=0.85 \
+    prompt='Suggest some famous landmarks in London.' \
+    use_chat_template=True \
+    scan_layers=true \
+    enable_single_controller=${use_pathways} \
+    prefuse_moe_weights=True \
+    ici_tensor_parallelism=8