diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3.sh new file mode 100644 index 0000000000..9fa9e9a7b1 --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3.sh @@ -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= +# 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\ \\' \ + 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\ \\' \ + 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 \ No newline at end of file diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_multimodal_sft.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_multimodal_sft.sh new file mode 100644 index 0000000000..3ffef50584 --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_multimodal_sft.sh @@ -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= +# 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\ \\' \ + 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\ \\' \ + image_path=\'tests/assets/test_image.jpg\' \ + attention=\'dot_product\' \ + skip_jax_distributed_system=True \ No newline at end of file diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_rl.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_rl.sh new file mode 100644 index 0000000000..704f771328 --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_rl.sh @@ -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= +# 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} \ No newline at end of file diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_sft.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_sft.sh new file mode 100644 index 0000000000..b404cbc7db --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_sft.sh @@ -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= +# 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} \ No newline at end of file diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_to_hf.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_to_hf.sh new file mode 100644 index 0000000000..5b8a33a3a1 --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_to_hf.sh @@ -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 \ No newline at end of file diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_to_mt.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_to_mt.sh new file mode 100644 index 0000000000..420e31b03b --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_to_mt.sh @@ -0,0 +1,108 @@ +#!/bin/bash + +# Converts Qwen3-VL-2B 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_qwen3_to_mt.sh $RUN_ID + +set -ex + +run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)} +MODEL_NAME='qwen3-vl-2b' +HF_GOLDEN_MODEL=Qwen/Qwen3-VL-2B-Instruct + +# To convert the multimodal model, make sure the use_multimodal is set to be true + +BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}/to_maxtext + +# Step 1: Install torch +python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu +python3 -m pip install decord + +# Step 2.a: Convert to unscanned checkpoint (for inference) +python3 -m maxtext.checkpoint_conversion.to_maxtext \ + model_name=${MODEL_NAME} \ + 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 \ + override_model_config=True \ + --lazy_load_tensors=False \ + --eager_load_method='safetensors' + +UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/unscanned_multimodal/${run_id}/0/items +echo "Unscanned checkpoint path: ${UNSCANNED_CKPT_PATH}" + +# Step 2.b: Convert to scanned multimodal checkpoint (for multimodal training) +python3 -m maxtext.checkpoint_conversion.to_maxtext \ + model_name=${MODEL_NAME} \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/unscanned/${run_id} \ + use_multimodal=true \ + scan_layers=false \ + hardware=cpu \ + skip_jax_distributed_system=True \ + checkpoint_storage_use_zarr3=False \ + checkpoint_storage_use_ocdbt=False \ + --lazy_load_tensors=False \ + --eager_load_method='safetensors' + +MULTIMODAL_UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/unscanned_multimodal/${run_id}/0/items +echo "Multimodal Scanned checkpoint path: ${MULTIMODAL_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 +TEST_PROMPT='Describe this image' +TEST_IMAGE='tests/assets/test_image.jpg' +export GOLDEN_LOGITS_PATH=/tmp/golden_qwen3_vl_2b_vision.jsonl + +python3 -m tests.assets.logits_generation.generate_hf_golden_logits \ + --model-id=${HF_GOLDEN_MODEL} \ + --output-path=${GOLDEN_LOGITS_PATH} \ + --prompts="${TEST_PROMPT}" \ + --image-paths=${TEST_IMAGE} \ + --hf-model-path=${HF_GOLDEN_MODEL} \ + --apply-chat-template \ + --output-format=json + +echo "=== Running MaxText Multimodal Forward Pass Logit Checker ===" +python3 -m tests.utils.forward_pass_logit_checker \ + tokenizer_path=${HF_GOLDEN_MODEL} \ + load_parameters_path=${MULTIMODAL_UNSCANNED_CKPT_PATH} \ + model_name=${MODEL_NAME} \ + use_multimodal=true \ + scan_layers=false \ + dtype=float32 \ + matmul_precision=highest \ + per_device_batch_size=1 \ + attention=dot_product \ + prompt="${TEST_PROMPT}" \ + image_path=${TEST_IMAGE} \ + --max_kl_div=0.1 \ + --golden_logits_path=${GOLDEN_LOGITS_PATH} \ + override_model_config=true + +echo "=== Running MaxText Forward Pass Logit Checker ===" +python3 -m tests.utils.forward_pass_logit_checker \ + tokenizer_path=${HF_GOLDEN_MODEL} \ + load_parameters_path=${UNSCANNED_CKPT_PATH} \ + model_name=${MODEL_NAME} \ + use_multimodal=false \ + scan_layers=false \ + per_device_batch_size=1 \ + dtype=float32 \ + --max_kl_div=0.1 \ + --run_hf_model=true \ + --hf_model_path=${HF_GOLDEN_MODEL} \ + override_model_config=true +