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=<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\'
\ 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=<your Hugging Face access 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=<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
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=<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'
-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=<your Hugging Face access 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