From 2414d0bc797baf27c8299bf66e1956ca7328c0b2 Mon Sep 17 00:00:00 2001 From: Jacky Fang Date: Wed, 22 Jul 2026 01:24:57 +0000 Subject: [PATCH] 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