AI-Hypercomputer · chiajunglien · Jul 31, 2026
@@ -1,73 +1,80 @@
 #!/bin/bash
 
-# This file is documentation for how to get started with gpt-oss-20b on v5p-8.
+# Validates the GPT-OSS-20b pre-training pipeline using a pre-converted MaxText checkpoint.
 
-# 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.
+# 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.
 
-# Example Usage: export HF_TOKEN=<huggingface_access_token>; export BASE_OUTPUT_PATH=<GCS_bucket_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
+# 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
 
-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"
-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
+run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+MODEL_NAME='gpt-oss-20b'
 
-# 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
+# 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
 
-# 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
+DATASET_PATH=gs://maxtext-dataset
+
+# 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\'
+
+# 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 \
+    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} \
+    max_target_length=1024 \
+    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
+
+# Step 3: Run inference on the checkpoint generated from the previous run
+    python3 -m maxtext.inference.decode \
+    model_name=${MODEL_NAME} \
+    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\'
@@ -0,0 +1,74 @@
+#!/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 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'
+
+# 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
+
+# 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 2: Run LoRA on the converted checkpoint
+python3 -m maxtext.trainers.post_train.sft.train_sft \
+    base_output_directory=${BASE_OUTPUT_DIRECTORY}/lora \
+    load_parameters_path=${SCANNED_CKPT_PATH} \
+    tokenizer_path='unsloth/gpt-oss-20b-BF16' \
+    per_device_batch_size=1 \
+    run_name=${run_id} \
+    steps=5 \
+    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_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 \
+    --use_tunix=True \
+    model_name=${MODEL_NAME} \
+    load_parameters_path=${SCANNED_CKPT_PATH} \
+    lora.enable_lora=True \
+    lora.lora_restore_path=${BASE_OUTPUT_DIRECTORY}/lora/${run_id}/checkpoints/5/model_params \
+    lora.lora_rank=16 \
+    lora.lora_alpha=32.0 \
+    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_single_controller=${use_pathways}
@@ -0,0 +1,70 @@
+#!/bin/bash
+
+# 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'
+
+# 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
+
+
+# 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 2: Run RL on the converted checkpoint
+python3 -m maxtext.trainers.post_train.rl.train_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=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=4 \
+    vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
+    vllm_additional_config='{"maxtext_config": {"model_name": "gpt-oss-20b", "log_config": "false", "prefuse_moe_weights": "true"}}'
+
+# Step 3: Run inference on the checkpoint generated from the previous run
+python3 -m maxtext.inference.vllm_decode \
+    model_name=${MODEL_NAME} \
+    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
@@ -0,0 +1,68 @@
+#!/bin/bash
+
+# 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=<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_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'
+
+# 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
+
+
+# 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} \
+    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} \
+    per_device_batch_size=1 \
+    run_name=${run_id} \
+    steps=5 \
+    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} \
+    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
@@ -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