From 2414d0bc797baf27c8299bf66e1956ca7328c0b2 Mon Sep 17 00:00:00 2001 From: Jacky Fang Date: Wed, 22 Jul 2026 01:24:57 +0000 Subject: [PATCH 1/2] feat(test): add modular E2E TPU testing pipeline scripts for GPTOSS-20B --- .../tpu/gpt_oss/20b/test_gpt_oss.sh | 107 ++++++++---------- .../tpu/gpt_oss/20b/test_gpt_oss_lora.sh | 87 ++++++++++++++ .../tpu/gpt_oss/20b/test_gpt_oss_rl.sh | 55 +++++++++ .../tpu/gpt_oss/20b/test_gpt_oss_sft.sh | 60 ++++++++++ .../tpu/gpt_oss/20b/test_gpt_oss_to_hf.sh | 30 +++++ .../tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh | 60 ++++++++++ 6 files changed, 338 insertions(+), 61 deletions(-) create mode 100755 tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_lora.sh create mode 100755 tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh create mode 100755 tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh create mode 100755 tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_hf.sh create mode 100755 tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh index f11c117d7e..dabd89cdab 100644 --- a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh +++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh @@ -1,73 +1,58 @@ #!/bin/bash -# This file is documentation for how to get started with gpt-oss-20b on v5p-8. - -# The flow of this file is as follows: -# 1. Convert the HuggingFace checkpoint (bf16) to MaxText-compatible checkpoint (bf16): -# Scanned format is better for training; unscanned format is better for decoding. -# 2. Run logit check, pre-training, fine-tuning, and decoding. - -# Example Usage: export HF_TOKEN=; export BASE_OUTPUT_PATH=; bash test_gpt_oss.sh - -# The golden logit can be generated by: -# python3 -m tests.assets.logits_generation.generate_hf_golden_logits --model-id=openai/gpt-oss-20b --output-path=golden_data_gpt-oss-20b.jsonl --prompts='I love to;Today is a;What is the' --hf-model-path=$local_bf16_path +# 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 - # Non-Googlers please remember to point `BASE_OUTPUT_PATH` to GCS buckets that you own, this script uses internal buckets for testing. - export BASE_OUTPUT_PATH=gs://runner-maxtext-logs/$(date +%Y-%m-%d-%H-%M) - echo "BASE_OUTPUT_PATH is not set" + 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} - -# Installing torch for checkpoint conversion and forward_pass_logit_checker.py -python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu - -# Step 1: Checkpoint conversion -# Assume HF checkpoints are uploaded to GCS bucket at CKPT_BUCKET -# Non-Googlers please remember to point `CKPT_BUCKET` to GCS buckets that you own -# Copying the HF checkpoint into a local directory `/tmp` -- you are free to use a different directory -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.1 Convert checkpoint to `scanned` format, more suitable 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 --model-size ${MODEL_NAME} - -# 1.2 Convert checkpoint to `unscanned` format, more suitable for decoding -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 --model-size ${MODEL_NAME} - -# Step 2: -# We define the checkpoint paths. This way it is easier to use these paths in the `train.py` and `decode.py` commands -export SCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/scanned/0/items -export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/0/items -# Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data -export DATASET_PATH=gs://maxtext-dataset - -export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=81920' - -# Test whether the forward pass logits match the golden logits -# default golden_logits_path=/deps/tests/assets/golden_logits/golden_data_{MODEL_NAME}.jsonl, copied from gs://maxtext-test-assets/golden_data_${MODEL_NAME}.jsonl -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 - -# Run pre-training - megablox implementation -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} run_name=megablox_pre_training model_name=${MODEL_NAME} tokenizer_type=huggingface tokenizer_path=${TOKENIZER_PATH} dataset_type=synthetic enable_checkpointing=false 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 - -# Run fine-tuning - megablox implementation -# TODO: remove `abort_on_nan_loss=false` after b/497864549 -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} run_name=megablox_fine_tuning model_name=${MODEL_NAME} tokenizer_type=huggingface tokenizer_path=${TOKENIZER_PATH} dataset_path=${DATASET_PATH} 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 - -# Run supervised fine-tuning - megablox implementation -# TODO: remove `abort_on_nan_loss=false` after b/497864549 -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} run_name=megablox_supervised_fine_tuning 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 -# Run decoding - megablox implementation -# Note decode requires the access token for huggingface tokenizer even if the model is not gated -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} hf_access_token=${HF_TOKEN} 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 +export SCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/scanned/${run_id}/0/items +export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/${run_id}/0/items + +export SPARSE_MATMUL="True" +export MEGABLOX="True" +export PRETRAIN_ATTENTION="flash" + +# 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=${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 + +# 2. Run Verification Decoding from the converted checkpoint +python3 -m maxtext.inference.decode \ + base_output_directory=${BASE_OUTPUT_PATH} \ + run_name=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 diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_lora.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_lora.sh new file mode 100755 index 0000000000..70dbedd9a1 --- /dev/null +++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_lora.sh @@ -0,0 +1,87 @@ +#!/bin/bash + +# 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. + +# Usage: +# export HF_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)} +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 + +SPARSE_MATMUL="True" +MEGABLOX="True" +LORA_ATTENTION="flash" + +# Step 1: Run LoRA SFT 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 \ + 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 \ + steps=5 \ + max_target_length=1024 \ + ici_fsdp_parallelism=1 \ + ici_expert_parallelism=4 \ + gcs_metrics=true \ + abort_on_nan_loss=false \ + 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 + +# 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 \ + 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_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} \ + scan_layers=true \ + enable_nnx=True \ + pure_nnx_decoder=True diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh new file mode 100755 index 0000000000..51d8c405f6 --- /dev/null +++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh @@ -0,0 +1,55 @@ +#!/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 + +export SPARSE_MATMUL="True" +export MEGABLOX="True" +export ATTENTION="flash" +export VLLM_ADDITIONAL_CONFIG='{"maxtext_config": {"model_name": "gpt-oss-20b", "log_config": "false"}}' + +# 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 \ + attention=${ATTENTION} \ + sparse_matmul=${SPARSE_MATMUL} \ + megablox=${MEGABLOX} \ + vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ + vllm_additional_config="${VLLM_ADDITIONAL_CONFIG}" + +# 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 \ + 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 diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh new file mode 100755 index 0000000000..2f6c39a74f --- /dev/null +++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh @@ -0,0 +1,60 @@ +#!/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 + +export SPARSE_MATMUL="True" +export MEGABLOX="True" +export SFT_ATTENTION="flash" + +# 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=${SFT_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=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 \ + base_output_directory=${BASE_OUTPUT_PATH} \ + run_name=decode_sft \ + model_name=${MODEL_NAME} \ + tokenizer_path=${TOKENIZER_PATH} \ + load_parameters_path=${BASE_OUTPUT_PATH}/sft/${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 diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_hf.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_hf.sh new file mode 100755 index 0000000000..14d046c908 --- /dev/null +++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_hf.sh @@ -0,0 +1,30 @@ +#!/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 diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh new file mode 100755 index 0000000000..955d518cfb --- /dev/null +++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh @@ -0,0 +1,60 @@ +#!/bin/bash + +# Converts GPTOSS-20B HuggingFace checkpoint to MaxText format and validates logit correctness. + +set -ex + +export PYTHONPATH=src + +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}" + +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 +if [ ! -f /tmp/golden_data_gpt-oss-20b.jsonl ]; then + gcloud storage cp gs://maxtext-test-assets/golden_data_gpt-oss-20b.jsonl /tmp/golden_data_gpt-oss-20b.jsonl +fi + +SPARSE_MATMUL="True" +MEGABLOX="True" + +python3 -m tests.utils.forward_pass_logit_checker \ + base_output_directory=${BASE_OUTPUT_PATH} \ + model_name=${MODEL_NAME} \ + load_parameters_path=${UNSCANNED_CKPT_PATH} \ + scan_layers=false \ + attention=dot_product \ + sparse_matmul=${SPARSE_MATMUL} \ + megablox=${MEGABLOX} \ + --golden_logits_path=/tmp/golden_data_gpt-oss-20b.jsonl \ + --max_kl_div=0.01 From f93aa68bd86670bd453aad44cf84a5fbcc23f2ba Mon Sep 17 00:00:00 2001 From: Emmalien Date: Wed, 29 Jul 2026 09:30:54 +0000 Subject: [PATCH 2/2] fix (#4641) --- .../tpu/gpt_oss/20b/test_gpt_oss.sh | 108 +++++++++++------- .../tpu/gpt_oss/20b/test_gpt_oss_lora.sh | 99 +++++++--------- .../tpu/gpt_oss/20b/test_gpt_oss_rl.sh | 77 ++++++++----- .../tpu/gpt_oss/20b/test_gpt_oss_sft.sh | 100 ++++++++-------- .../tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh | 100 +++++++++------- 5 files changed, 267 insertions(+), 217 deletions(-) diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh index dabd89cdab..aea09c2721 100644 --- a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss.sh +++ b/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= +# 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\' \ No newline at end of file diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_lora.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_lora.sh index 70dbedd9a1..01433cc10b 100755 --- a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_lora.sh +++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_lora.sh @@ -1,11 +1,12 @@ #!/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= @@ -13,75 +14,61 @@ # 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} diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh index 51d8c405f6..b0843e1c0c 100755 --- a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_rl.sh +++ b/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= +# 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 diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh index 2f6c39a74f..912a546007 100755 --- a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh +++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_sft.sh @@ -1,60 +1,68 @@ #!/bin/bash -# Validates the GPTOSS-20B Supervised Fine-Tuning (SFT) pipeline. +# Validates the GPT-OSS-20b SFT 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 SFT starting from the pre-converted checkpoint. +# 3. Run inference on the checkpoint produced by the SFT run. + +# Usage: +# export HF_TOKEN= +# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S) +# bash test_gpt_oss_to_mt.sh $RUN_ID +# bash test_gpt_oss_sft.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 SFT_ATTENTION="flash" - -# 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} \ +# Step 1: Run inference on the original checkpoint converted from Hugging Face +python3 -m maxtext.inference.vllm_decode \ model_name=${MODEL_NAME} \ - tokenizer_type=huggingface \ - tokenizer_path=${TOKENIZER_PATH} \ - dataset_type=hf \ - enable_checkpointing=true \ - async_checkpointing=false \ + load_parameters_path=${UNSCANNED_CKPT_PATH} \ + tokenizer_path='unsloth/gpt-oss-20b-BF16' \ + 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} \ + ici_tensor_parallelism=8 + + +# Step 2: Run SFT on the converted checkpoint +python3 -m maxtext.trainers.post_train.sft.train_sft \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/sft \ load_parameters_path=${SCANNED_CKPT_PATH} \ - scan_layers=True \ - attention=${SFT_ATTENTION} \ - sparse_matmul=${SPARSE_MATMUL} \ - megablox=${MEGABLOX} \ - dtype=bfloat16 \ - weight_dtype=bfloat16 \ - per_device_batch_size=4 \ + 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 - -# 2. Run Decoding on the newly produced SFT checkpoint -python3 -m maxtext.inference.decode \ - base_output_directory=${BASE_OUTPUT_PATH} \ - run_name=decode_sft \ + scan_layers=true \ + model_name=${MODEL_NAME} \ + tokenizer_path='unsloth/gpt-oss-20b-BF16' \ + enable_single_controller=${use_pathways} \ + checkpoint_storage_use_zarr3=False \ + checkpoint_storage_use_ocdbt=False + +# 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}/sft/${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 + load_parameters_path=${BASE_OUTPUT_DIRECTORY}/sft/${run_id}/checkpoints/5/model_params \ + tokenizer_path='unsloth/gpt-oss-20b-BF16' \ + vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ + hbm_utilization_vllm=0.7 \ + prompt="Suggest some famous landmarks in London." \ + use_chat_template=True \ + scan_layers=true \ + enable_single_controller=${use_pathways} \ + ici_tensor_parallelism=8 diff --git a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh index 955d518cfb..28d6227356 100755 --- a/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh +++ b/tests/end_to_end/tpu/gpt_oss/20b/test_gpt_oss_to_mt.sh @@ -1,60 +1,78 @@ #!/bin/bash -# Converts GPTOSS-20B HuggingFace checkpoint to MaxText format and validates logit correctness. +# Converts GPT-OSS-20b HuggingFace checkpoint to MaxText format and validates logit correctness. + +# The flow of this script is as follows: +# 1. Install PyTorch (CPU) required for checkpoint conversion. +# 2. Convert the HuggingFace checkpoint to MaxText format in both unscanned and scanned formats. +# 3. Run a forward pass logits check to verify the converted checkpoint matches the original HF model. + +# Usage: +# export HF_TOKEN= +# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S) +# bash test_gpt_oss_to_mt.sh $RUN_ID - to convert the checkpoint and run logit check set -ex export PYTHONPATH=src 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}" +MODEL_NAME='gpt-oss-20b' -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 +# Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you want to store scanned and unscanned checkpoints +BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}/to_maxtext -# 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} +# Step 1: Install torch +python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu -SCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/scanned/${run_id}/0/items -echo "Scanned checkpoint path: ${SCANNED_CKPT_PATH}" +# Step 2: Convert the checkpoint from Hugging Face to make it compatible with MaxText -# 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} +# Step 2.a: Convert to unscanned checkpoint (for inference) +python3 -m maxtext.checkpoint_conversion.to_maxtext \ + model_name=${MODEL_NAME} \ + --hf_model_path="unsloth/gpt-oss-20b-BF16" \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/unscanned/${run_id} \ + use_multimodal=false \ + scan_layers=false \ + hardware=cpu \ + skip_jax_distributed_system=True \ + checkpoint_storage_use_zarr3=False \ + checkpoint_storage_use_ocdbt=False \ + attention=\'dot_product\' -UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/${run_id}/0/items +UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/unscanned/${run_id}/0/items echo "Unscanned checkpoint path: ${UNSCANNED_CKPT_PATH}" -# 3. Logit correctness check +# Step 2.b: Convert to scanned checkpoint (for training) +python3 -m maxtext.checkpoint_conversion.to_maxtext \ + model_name=${MODEL_NAME} \ + --hf_model_path="unsloth/gpt-oss-20b-BF16" \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/scanned/${run_id} \ + use_multimodal=false \ + scan_layers=true \ + hardware=cpu \ + skip_jax_distributed_system=True \ + checkpoint_storage_use_zarr3=False \ + checkpoint_storage_use_ocdbt=False \ + attention=\'dot_product\' + +SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/scanned/${run_id}/0/items +echo "Scanned checkpoint path: ${SCANNED_CKPT_PATH}" + +# Step 3: Test whether the forward pass logits match the original HF model +# to get higher precision (eg. float32) run on CPU with `JAX_PLATFORMS=cpu` +# ToDo: improve forward_pass_logit_checker to test multi-modal prompt if [ ! -f /tmp/golden_data_gpt-oss-20b.jsonl ]; then gcloud storage cp gs://maxtext-test-assets/golden_data_gpt-oss-20b.jsonl /tmp/golden_data_gpt-oss-20b.jsonl fi -SPARSE_MATMUL="True" -MEGABLOX="True" - -python3 -m tests.utils.forward_pass_logit_checker \ - base_output_directory=${BASE_OUTPUT_PATH} \ - model_name=${MODEL_NAME} \ - load_parameters_path=${UNSCANNED_CKPT_PATH} \ - scan_layers=false \ - attention=dot_product \ - sparse_matmul=${SPARSE_MATMUL} \ - megablox=${MEGABLOX} \ - --golden_logits_path=/tmp/golden_data_gpt-oss-20b.jsonl \ - --max_kl_div=0.01 + python3 -m tests.utils.forward_pass_logit_checker \ + load_parameters_path=${UNSCANNED_CKPT_PATH} \ + model_name=${MODEL_NAME} \ + use_multimodal=false \ + scan_layers=false \ + global_batch_size_to_train_on=1 \ + per_device_batch_size=1 \ + max_target_length=512 \ + --golden_logits_path=/tmp/golden_data_gpt-oss-20b.jsonl \ + --max_kl_div=0.01