AI-Hypercomputer · chiajunglien · Jul 31, 2026
@@ -0,0 +1,120 @@
+#!/bin/bash
+
+# Validates the Qwen3-VL-2B 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_qwen3_to_mt.sh $RUN_ID
+# bash test_qwen3.sh $RUN_ID
+
+
+set -ex
+
+run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+MODEL_NAME='qwen3-vl-2b'
+
+USE_MULTIMODAL=false
+
+# 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
+
+# Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
+DATASET_PATH=gs://maxtext-dataset
+
+# Step 1: Run inference on the original checkpoint converted from Hugging Face
+if [ ${USE_MULTIMODAL} == true ]; then
+    python3 -m maxtext.inference.decode \
+    model_name=${MODEL_NAME} \
+    tokenizer_type="huggingface" \
+    load_parameters_path=${UNSCANNED_CKPT_PATH} \
+    per_device_batch_size=1 r\
+    un_name=${run_id} \
+    max_prefill_predict_length=272 \
+    max_target_length=300 \
+    steps=1 \
+    async_checkpointing=false \
+    scan_layers=false \
+    use_multimodal=true \
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False \
+    prompt=\'Describe\ image\ \<start_of_image\>\' \
+    image_path=\'tests/assets/test_image.jpg\' \
+    attention=\'dot_product\'
+else
+    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\'
+fi
+
+# 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=8192 \
+    steps=5 async_checkpointing=false \
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False \
+    model_name=${MODEL_NAME} \
+    scan_layers=false \
+    use_multimodal=${USE_MULTIMODAL}
+
+# Step 3: Run inference on the checkpoint generated from the previous run
+if [ ${USE_MULTIMODAL} == true ]; then
+    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=272 \
+    max_target_length=300 \
+    steps=1 \
+    async_checkpointing=false \
+    scan_layers=false \
+    use_multimodal=true \
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False \
+    prompt=\'Describe\ image\ \<start_of_image\>\' \
+    image_path=\'tests/assets/test_image.jpg\' \
+    attention=\'dot_product\'
+else
+    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\'
+fi
@@ -0,0 +1,88 @@
+#!/bin/bash
+
+# Validates the Qwen3-VL-2B SFT multimodal 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 of Qwen3-VL-2B on ChartQA dataset with the 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_qwen3_to_mt.sh $RUN_ID true
+# bash test_qwen3_multimodal_sft.sh $RUN_ID
+
+# Note: You can stop at any step if you just want to run part of the flow.
+
+set -ex
+
+run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+use_pathways=${2:-false}
+MODEL_NAME='qwen3-vl-2b'
+
+# 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}
+MULTIMODAL_UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned_multimodal/${run_id}/0/items
+
+# Step 1: Install google-jetstream
+python3 -m pip install google-jetstream@https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip --no-deps
+
+# Step 2: Run inference on the original checkpoint converted from Hugging Face
+python3 -m maxtext.inference.decode \
+    model_name=${MODEL_NAME} \
+    load_parameters_path=${MULTIMODAL_SCANNED_CKPT_PATH} \
+    per_device_batch_size=1 \
+    run_name=${run_id} \
+    max_prefill_predict_length=4096 \
+    max_target_length=4200 \
+    steps=1 \
+    async_checkpointing=false \
+    scan_layers=false \
+    use_multimodal=true \
+    tokenizer_type=huggingface \
+    prompt=\'Describe\ image\ \<start_of_image\>\' \
+    image_path=\'tests/assets/test_image.jpg\' \
+    attention=\'dot_product\' \
+    skip_jax_distributed_system=True
+
+# Step 3: Run SFT on the MaxText checkpoint on ChartQA dataset
+python -m maxtext.trainers.post_train.sft.train_sft_native "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/post_train/sft-vision-chartqa.yml \
+    run_name=${run_id} \
+    model_name=${MODEL_NAME} \
+    per_device_batch_size=1 \
+    max_prefill_predict_length=1024 \
+    max_target_length=2048 \
+    steps=5 \
+    scan_layers=true \
+    async_checkpointing=False \
+    attention=\'dot_product\' \
+    dataset_type=hf hf_path=parquet \
+    hf_train_files=gs://aireenmei-multipod/dataset/hf/chartqa/train-* \
+    base_output_directory=${BASE_OUTPUT_DIRECTORY}/multimodal/sft \
+    load_parameters_path=${MULTIMODAL_UNSCANNED_CKPT_PATH} \
+    dtype=bfloat16 \
+    weight_dtype=bfloat16 \
+    sharding_tolerance=0.05 \
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False \
+    enable_single_controller=${use_pathways} \
+    grain_worker_count=0
+
+# Step 4: Run inference on the checkpoint generated from the previous run
+python3 -m maxtext.inference.decode \
+    model_name=${MODEL_NAME} \
+    load_parameters_path=${BASE_OUTPUT_DIRECTORY}/multimodal/sft/${run_id}/checkpoints/4/items \
+    per_device_batch_size=1 \
+    run_name=${run_id} \
+    max_prefill_predict_length=4096 \
+    max_target_length=4200 \
+    steps=1 \
+    async_checkpointing=false \
+    scan_layers=true \
+    use_multimodal=true \
+    tokenizer_type=huggingface \
+    prompt=\'Describe\ image\ \<start_of_image\>\' \
+    image_path=\'tests/assets/test_image.jpg\' \
+    attention=\'dot_product\' \
+    skip_jax_distributed_system=True
@@ -0,0 +1,67 @@
+#!/bin/bash
+
+# Validates the Qwen3-VL-2B 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_qwen3_to_mt.sh $RUN_ID
+# bash test_qwen3_rl.sh $RUN_ID
+
+
+set -ex
+
+run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+use_pathways=${2:-false}
+MODEL_NAME='qwen3-vl-2b'
+
+# 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
+
+# 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.5 \
+    prompt='Suggest some famous landmarks in London.' \
+    use_chat_template=True \
+    scan_layers=false \
+    enable_single_controller=${use_pathways}
+
+# 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=${UNSCANNED_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} \
+    enable_single_controller=${use_pathways} \
+    checkpoint_storage_use_zarr3=False \
+    checkpoint_storage_use_ocdbt=False \
+    rollout_tensor_parallelism=1 \
+    use_standalone_converter=true \
+    vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
+    vllm_additional_config='{"maxtext_config": {"model_name": "qwen3-vl-2b", "log_config": "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}/rl/${run_id}/checkpoints/actor/5/model_params \
+    vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
+    hbm_utilization_vllm=0.5 \
+    prompt='Suggest some famous landmarks in London.' \
+    use_chat_template=True \
+    scan_layers=false \
+    enable_single_controller=${use_pathways}
@@ -0,0 +1,60 @@
+#!/bin/bash
+
+# Validates the Qwen3-VL-2B 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_qwen3_to_mt.sh $RUN_ID
+# bash test_qwen3_sft.sh $RUN_ID
+
+
+set -ex
+
+run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)}
+use_pathways=${2:-false}
+MODEL_NAME='qwen3-vl-2b'
+
+# 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
+
+# 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.5 \
+    prompt="Suggest some famous landmarks in London." \
+    use_chat_template=True \
+    scan_layers=false \
+    enable_single_controller=${use_pathways}
+
+# 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=${UNSCANNED_CKPT_PATH} \
+    per_device_batch_size=1 \
+    run_name=${run_id} \
+    steps=5 \
+    scan_layers=false \
+    model_name=${MODEL_NAME} \
+    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 \
+    vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
+    hbm_utilization_vllm=0.5 \
+    prompt="Suggest some famous landmarks in London." \
+    use_chat_template=True \
+    scan_layers=false \
+    enable_single_controller=${use_pathways}
@@ -0,0 +1,31 @@
+#!/bin/bash
+
+# Converts a MaxText checkpoint to a Hugging Face model checkpoint.
+
+# Usage:
+# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S)
+# bash test_gemma4_to_hf.sh $RUN_ID $CHECKPOINT_PATH $USE_MULTIMODAL $SCAN_LAYERS
+
+set -ex
+
+run_id=$1
+CKPT_PATH=$2
+USE_MULTIMODAL=${3:-false}
+SCAN_LAYERS=${4:-false}
+
+MODEL_NAME='qwen3-vl-2b'
+BASE_OUTPUT_DIRECTORY="gs://runner-maxtext-logs/${MODEL_NAME}"
+
+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_DIRECTORY}/to_huggingface/${scan_status}/${run_id} \
+    use_multimodal=${USE_MULTIMODAL} \
+    scan_layers=$SCAN_LAYERS