From 7c525010d8848216d8ad4333246b96a95a637cd9 Mon Sep 17 00:00:00 2001 From: Yuchen Hou Date: Wed, 1 Jul 2026 22:42:55 +0000 Subject: [PATCH 1/6] add Qwen3 VL tests, update logit checker, and golden logit generator --- .../generate_hf_golden_logits.py | 8 +- .../qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh | 129 ++++++++++++++++++ tests/utils/forward_pass_logit_checker.py | 103 +++++++++++++- 3 files changed, 234 insertions(+), 6 deletions(-) create mode 100644 tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh diff --git a/tests/assets/logits_generation/generate_hf_golden_logits.py b/tests/assets/logits_generation/generate_hf_golden_logits.py index c57d58c380..6fcb7898b0 100644 --- a/tests/assets/logits_generation/generate_hf_golden_logits.py +++ b/tests/assets/logits_generation/generate_hf_golden_logits.py @@ -80,6 +80,9 @@ def save_golden_logits( from transformers import Llama4ForConditionalGeneration # pylint: disable=import-outside-toplevel model_class = Llama4ForConditionalGeneration + elif "qwen3-vl" in model_id.lower(): + from transformers import Qwen3VLForConditionalGeneration # pylint: disable=import-outside-toplevel + model_class = Qwen3VLForConditionalGeneration else: from transformers import AutoModelForCausalLM # pylint: disable=import-outside-toplevel @@ -151,7 +154,10 @@ def save_golden_logits( for key, value in inputs.items(): new_key = "tokens" if key == "input_ids" else key val_np = value.cpu().numpy() - data_to_save[new_key] = val_np[0] if val_np.ndim > 0 else val_np + if key == "pixel_values" and "qwen3-vl" in model_id.lower(): + data_to_save[new_key] = val_np + else: + data_to_save[new_key] = val_np[0] if val_np.ndim > 0 else val_np data_to_save["logits"] = logits[0] print(f"Token length is {len(data_to_save['tokens'])} for prompt: {prompt_text}") diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh new file mode 100644 index 0000000000..2d38ae0751 --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh @@ -0,0 +1,129 @@ +#!/bin/bash + +# This script is both an end-to-end test and documentation for converting a +# Qwen3-VL-2B MaxText checkpoint to Hugging Face format. Can be run on a v4-8. + +# The flow of this script is as follows: +# 1. Convert an original Hugging Face model checkpoint to MaxText format. +# 2. Convert the resulting MaxText checkpoint back to Hugging Face format. +# 3. Run a forward pass check to compare the logits and KL divergence between +# the MaxText checkpoint and the Hugging Face checkpoint. + +# Pre-requisites: +# 1. Set HF_TOKEN environment variable to your Hugging Face access token. +# export HF_TOKEN= +# 2. Configure USE_MULTIMODAL (true for multimodal, false for text-only). +# 3. Configure USE_SCAN_LAYERS (true if checkpoint was trained with scanned layers, false otherwise). + +set -ex + + +MODEL_NAME='qwen3-vl-2b' +export MODEL_VARIATION='vl_2b' +export HF_MODEL=Qwen/Qwen3-VL-2B-Instruct + +idx=$(date +%Y-%m-%d-%H-%M) + +# Set USE_SCAN_LAYERS=true if the checkpoint was trained with scanned layers +USE_SCAN_LAYERS=false +if ${USE_SCAN_LAYERS}; then export CHECKPOINT_TYPE=scanned; else export CHECKPOINT_TYPE=unscanned; fi + +USE_MULTIMODAL=true + + +export HF_TOKEN= + +export MODEL_BUCKET= + +# Path to Maxtext converted to HF checkpoint +export LOCAL_PATH=/hf/${MODEL_NAME}/${idx} + + +# Installing torch for deps in forward_pass_logit_checker.py +python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu +python3 -m pip install decord + +# Check point conversion +python3 -m maxtext.checkpoint_conversion.to_maxtext \ + "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \ + model_name=${MODEL_NAME} \ + base_output_directory=${MODEL_BUCKET}/${MODEL_NAME}/${CHECKPOINT_TYPE}/${idx} \ + scan_layers=false \ + hf_access_token=${HF_TOKEN} \ + weight_dtype=bfloat16 \ + hardware=cpu \ + skip_jax_distributed_system=True \ + checkpoint_storage_use_ocdbt=False \ + checkpoint_storage_use_zarr3=False \ + --eager_load_method=safetensors \ + --lazy_load_tensors=False + +# Path to MaxText checkpoint +export CKPT_PATH=${MODEL_BUCKET}/${MODEL_NAME}/${CHECKPOINT_TYPE}/${idx}/0/items + +python3 -m maxtext.checkpoint_conversion.to_huggingface \ + "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \ + model_name=${MODEL_NAME} \ + hf_access_token=${HF_TOKEN} \ + load_parameters_path=${CKPT_PATH} \ + base_output_directory=${LOCAL_PATH} \ + use_multimodal=${USE_MULTIMODAL} \ + scan_layers=${USE_SCAN_LAYERS} \ + override_model_config=true + +# Run forward pass logit checker to validate the converted checkpoint. +if [ "${USE_MULTIMODAL}" == true ]; then + 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_MODEL} \ + --output-path=${GOLDEN_LOGITS_PATH} \ + --prompts="${TEST_PROMPT}" \ + --image-paths=${TEST_IMAGE} \ + --hf-model-path=${LOCAL_PATH} \ + --apply-chat-template \ + --output-format=json + + echo "=== Running MaxText Forward Pass Logit Checker ===" + python3 -m tests.utils.forward_pass_logit_checker \ + "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \ + tokenizer_path=${HF_MODEL} \ + load_parameters_path=${CKPT_PATH} \ + model_name=${MODEL_NAME} \ + use_multimodal=${USE_MULTIMODAL} \ + scan_layers=${USE_SCAN_LAYERS} \ + dtype=float32 \ + wi_tile_fwd_embed_dim=512 \ + wi_tile_fwd_mlp_dim=512 \ + wo_tile_fwd_embed_dim=512 \ + wo_tile_fwd_mlp_dim=512 \ + 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 +else + echo "=== Running MaxText Forward Pass Logit Checker ===" + python3 -m tests.utils.forward_pass_logit_checker \ + "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \ + tokenizer_path=${HF_MODEL} \ + load_parameters_path=${CKPT_PATH} \ + model_name=${MODEL_NAME} \ + use_multimodal=${USE_MULTIMODAL} \ + scan_layers=${USE_SCAN_LAYERS} \ + per_device_batch_size=1 \ + dtype=float32 \ + wi_tile_fwd_embed_dim=512 \ + wi_tile_fwd_mlp_dim=512 \ + wo_tile_fwd_embed_dim=512 \ + wo_tile_fwd_mlp_dim=512 \ + --max_kl_div=0.1 \ + --run_hf_model=true \ + --hf_model_path=${LOCAL_PATH} \ + override_model_config=true +fi diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py index 76038e745b..d1a7066aad 100644 --- a/tests/utils/forward_pass_logit_checker.py +++ b/tests/utils/forward_pass_logit_checker.py @@ -34,7 +34,46 @@ # For example: # tests/assets/logits_generation/golden_llama2-7b_export.ipynb -"""Check if the logits generated by a model's src/maxtext/HF implementation matches golden logits for the same inputs""" +"""Check if the logits generated by a model's src/maxtext/HF implementation matches golden logits for the same inputs. + +Usage: + +1. For multimodal models (comparing MaxText against a pre-generated HF golden logits file): + +python3 -m tests.utils.forward_pass_logit_checker \ + src/maxtext/configs/base.yml \ + tokenizer_path=/qwen3-vl-2b \ + load_parameters_path=/qwen3-vl-2b/maxtext_ckpt/0/items \ + model_name=qwen3-vl-2b \ + use_multimodal=true \ + scan_layers=false \ + dtype=float32 \ + per_device_batch_size=1 \ + prompt="Describe this image" \ + image_path=tests/assets/test_image.jpg \ + --max_kl_div=0.1 \ + --golden_logits_path=golden_qwen2-vl-7b_vision.jsonl \ + override_model_config=true + +Note: For multimodal models, running the HuggingFace model on-the-fly inside this script is not supported. +You must pre-generate the HuggingFace golden logits file first using +tests/assets/logits_generation/generate_hf_golden_logits.py. + +2. For text-only models (running HuggingFace model on-the-fly to compare against MaxText): + +python3 -m tests.utils.forward_pass_logit_checker \ + src/maxtext/configs/base.yml \ + tokenizer_path=/gemma-2-2b \ + load_parameters_path=/gemma-2-2b/maxtext_ckpt/0/items \ + model_name=gemma-2-2b \ + use_multimodal=false \ + per_device_batch_size=1 \ + dtype=float32 \ + --max_kl_div=0.1 \ + --run_hf_model=true \ + --hf_model_path=/path/to/local/hf/gemma-2-2b \ + override_model_config=true +""" import argparse import functools @@ -211,6 +250,18 @@ def get_data(golden_data_point, config): pixel_values = np.transpose(pixel_values, (1, 2, 0)) elif model_prefix in ["llama4"]: pixel_values = pixel_values[None, :] + elif model_prefix in ["qwen3"]: + grid_thw = np.asarray(golden_data_point["image_grid_thw"]) + if grid_thw.ndim == 2: + grid_t, grid_h, grid_w = grid_thw[0] + else: + grid_t, grid_h, grid_w = grid_thw + tps = config.temporal_patch_size_for_vit + p = config.patch_size_for_vit + c = config.num_channels_for_vit + pixel_values = np.reshape(pixel_values, (c, int(grid_t * tps), int(grid_h * p), int(grid_w * p))) + else: + pixel_values = np.transpose(pixel_values, (1, 2, 0)) pixel_values = np.stack([pixel_values for _ in range(config.global_batch_size_to_train_on)]) else: pixel_values = None @@ -238,9 +289,30 @@ def get_data(golden_data_point, config): decoder_segment_ids = np.zeros(s, dtype=np.int32) decoder_segment_ids[:, :seq_len] = DECODING_ACTIVE_SEQUENCE_INDICATOR - decoder_positions = np.stack( - [np.arange(config.max_target_length, dtype=np.int32) for _ in range(config.global_batch_size_to_train_on)] - ) + + # For Qwen3-VL model, compute RoPE embeddings to generate 3D position IDs + if model_prefix in ["qwen3"] and config.use_mrope: + from maxtext.multimodal import processor_qwen3_omni # pylint: disable=import-outside-toplevel + image_grid_thw = np.atleast_2d(golden_data_point["image_grid_thw"]) if "image_grid_thw" in golden_data_point else None + video_grid_thw = np.atleast_2d(golden_data_point["video_grid_thw"]) if "video_grid_thw" in golden_data_point else None + + attention_mask = np.zeros(s, dtype=np.int32) + attention_mask[:, :seq_len] = 1 + + position_ids, _ = processor_qwen3_omni.get_rope_index( + input_ids=ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, + spatial_merge_size=config.spatial_merge_size_for_vit, + position_id_per_seconds=config.position_id_per_seconds, + config=config, + ) + decoder_positions = position_ids.astype(np.int32) + else: # For non-Qwen3-VL models, keep 1D position IDs + decoder_positions = np.stack( + [np.arange(config.max_target_length, dtype=np.int32) for _ in range(config.global_batch_size_to_train_on)] + ) return ids, decoder_segment_ids, decoder_positions, logits, seq_len, pixel_values @@ -371,6 +443,21 @@ def main(config, test_args): # pylint: disable=W0621 max_logging.log(f"{model_probabilities[1]=}") kl_div = jax.numpy.sum(jax.scipy.special.kl_div(golden_probabilities, model_probabilities), axis=-1) + + # Mask out vision placeholder tokens for KL calculation + ignore_token_ids = [] + if "qwen3" in config.model_name.lower(): + from maxtext.multimodal.processor_qwen3_omni import QwenTokens # pylint: disable=import-outside-toplevel + qwen_tokens = QwenTokens(config) + ignore_token_ids = [qwen_tokens.vision_start, qwen_tokens.vision_end, qwen_tokens.image_pad, qwen_tokens.video_pad] + + if ignore_token_ids: + slice_ids = ids[0, start_index:token_size] + mask = jnp.ones_like(slice_ids, dtype=jnp.bool_) + for ignore_id in ignore_token_ids: + mask = mask & (slice_ids != ignore_id) + kl_div = jnp.where(mask, kl_div, 0.0) + max_kl_div_val = jax.numpy.max(kl_div) max_kl_div_idx = jax.numpy.argmax(kl_div) max_logging.log( @@ -430,8 +517,14 @@ def main(config, test_args): # pylint: disable=W0621 # Default to bfloat16 if dtype is unrecognized torch_dtype = dtype_mapping.get(config.dtype.name.lower(), torch.bfloat16) max_logging.log(f"Loading HF model with dtype: {torch_dtype} (derived from config.dtype: {config.dtype})") + + if "qwen3-vl" in config.model_name.lower(): + from transformers import Qwen3VLForConditionalGeneration # pylint: disable=import-outside-toplevel + model_class = Qwen3VLForConditionalGeneration + else: + model_class = AutoModelForCausalLM - hf_model = AutoModelForCausalLM.from_pretrained( + hf_model = model_class.from_pretrained( test_args.hf_model_path, dtype=torch_dtype, token=hf_token, trust_remote_code=test_args.trust_remote_code ) hf_lora_path = config.hf_lora_adapter_path From 310e248f01e4b7ba2c6f3e8f976c2a3d62a7b70b Mon Sep 17 00:00:00 2001 From: Yuchen Hou Date: Wed, 1 Jul 2026 23:53:51 +0000 Subject: [PATCH 2/6] add Qwen3 VL tests, update logit checker, and golden logit generator --- tests/utils/forward_pass_logit_checker.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py index d1a7066aad..76ce8feeb7 100644 --- a/tests/utils/forward_pass_logit_checker.py +++ b/tests/utils/forward_pass_logit_checker.py @@ -468,7 +468,7 @@ def main(config, test_args): # pylint: disable=W0621 if jax.process_index() == 0 and test_args.output_logits_path: data_to_save = { - "prompt": golden_data[golden_data_index]["prompt"], + "prompt": golden_data_point["prompt"], "tokens": ids[0, :seq_len].tolist(), "logits": full_train_logits[0].tolist(), } @@ -517,7 +517,7 @@ def main(config, test_args): # pylint: disable=W0621 # Default to bfloat16 if dtype is unrecognized torch_dtype = dtype_mapping.get(config.dtype.name.lower(), torch.bfloat16) max_logging.log(f"Loading HF model with dtype: {torch_dtype} (derived from config.dtype: {config.dtype})") - + if "qwen3-vl" in config.model_name.lower(): from transformers import Qwen3VLForConditionalGeneration # pylint: disable=import-outside-toplevel model_class = Qwen3VLForConditionalGeneration From 6be9015b0f1e17a0fb931ba353bce96b7e81219a Mon Sep 17 00:00:00 2001 From: Yuchen Hou Date: Thu, 2 Jul 2026 00:09:25 +0000 Subject: [PATCH 3/6] Fix pyink formatting --- .../logits_generation/generate_hf_golden_logits.py | 1 + tests/utils/forward_pass_logit_checker.py | 10 +++++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/assets/logits_generation/generate_hf_golden_logits.py b/tests/assets/logits_generation/generate_hf_golden_logits.py index 6fcb7898b0..5b59214d62 100644 --- a/tests/assets/logits_generation/generate_hf_golden_logits.py +++ b/tests/assets/logits_generation/generate_hf_golden_logits.py @@ -82,6 +82,7 @@ def save_golden_logits( model_class = Llama4ForConditionalGeneration elif "qwen3-vl" in model_id.lower(): from transformers import Qwen3VLForConditionalGeneration # pylint: disable=import-outside-toplevel + model_class = Qwen3VLForConditionalGeneration else: from transformers import AutoModelForCausalLM # pylint: disable=import-outside-toplevel diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py index 76ce8feeb7..d8dca64cfd 100644 --- a/tests/utils/forward_pass_logit_checker.py +++ b/tests/utils/forward_pass_logit_checker.py @@ -293,6 +293,7 @@ def get_data(golden_data_point, config): # For Qwen3-VL model, compute RoPE embeddings to generate 3D position IDs if model_prefix in ["qwen3"] and config.use_mrope: from maxtext.multimodal import processor_qwen3_omni # pylint: disable=import-outside-toplevel + image_grid_thw = np.atleast_2d(golden_data_point["image_grid_thw"]) if "image_grid_thw" in golden_data_point else None video_grid_thw = np.atleast_2d(golden_data_point["video_grid_thw"]) if "video_grid_thw" in golden_data_point else None @@ -448,8 +449,14 @@ def main(config, test_args): # pylint: disable=W0621 ignore_token_ids = [] if "qwen3" in config.model_name.lower(): from maxtext.multimodal.processor_qwen3_omni import QwenTokens # pylint: disable=import-outside-toplevel + qwen_tokens = QwenTokens(config) - ignore_token_ids = [qwen_tokens.vision_start, qwen_tokens.vision_end, qwen_tokens.image_pad, qwen_tokens.video_pad] + ignore_token_ids = [ + qwen_tokens.vision_start, + qwen_tokens.vision_end, + qwen_tokens.image_pad, + qwen_tokens.video_pad + ] if ignore_token_ids: slice_ids = ids[0, start_index:token_size] @@ -520,6 +527,7 @@ def main(config, test_args): # pylint: disable=W0621 if "qwen3-vl" in config.model_name.lower(): from transformers import Qwen3VLForConditionalGeneration # pylint: disable=import-outside-toplevel + model_class = Qwen3VLForConditionalGeneration else: model_class = AutoModelForCausalLM From 372def3257b3a9e81f709ae887fe6717f2dbdedd Mon Sep 17 00:00:00 2001 From: Yuchen Hou Date: Thu, 2 Jul 2026 00:14:27 +0000 Subject: [PATCH 4/6] Fix trailing whitespace and pyink list formatting --- tests/utils/forward_pass_logit_checker.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py index d8dca64cfd..4f407ddd75 100644 --- a/tests/utils/forward_pass_logit_checker.py +++ b/tests/utils/forward_pass_logit_checker.py @@ -452,9 +452,9 @@ def main(config, test_args): # pylint: disable=W0621 qwen_tokens = QwenTokens(config) ignore_token_ids = [ - qwen_tokens.vision_start, - qwen_tokens.vision_end, - qwen_tokens.image_pad, + qwen_tokens.vision_start, + qwen_tokens.vision_end, + qwen_tokens.image_pad, qwen_tokens.video_pad ] From dbdf879c9bc1032202aa782db7812ef4c3a37838 Mon Sep 17 00:00:00 2001 From: Yuchen Hou Date: Thu, 2 Jul 2026 00:17:22 +0000 Subject: [PATCH 5/6] Fix final pyink indentation and comma --- tests/utils/forward_pass_logit_checker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py index 4f407ddd75..4f2d2bfdba 100644 --- a/tests/utils/forward_pass_logit_checker.py +++ b/tests/utils/forward_pass_logit_checker.py @@ -455,7 +455,7 @@ def main(config, test_args): # pylint: disable=W0621 qwen_tokens.vision_start, qwen_tokens.vision_end, qwen_tokens.image_pad, - qwen_tokens.video_pad + qwen_tokens.video_pad, ] if ignore_token_ids: From 22cb2b0d900d1e39cf0ccfaf98a0677cb91e047b Mon Sep 17 00:00:00 2001 From: Yuchen Hou Date: Thu, 2 Jul 2026 00:21:09 +0000 Subject: [PATCH 6/6] Fix pyink indentation --- tests/utils/forward_pass_logit_checker.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py index 4f2d2bfdba..acf22978bb 100644 --- a/tests/utils/forward_pass_logit_checker.py +++ b/tests/utils/forward_pass_logit_checker.py @@ -452,10 +452,10 @@ def main(config, test_args): # pylint: disable=W0621 qwen_tokens = QwenTokens(config) ignore_token_ids = [ - qwen_tokens.vision_start, - qwen_tokens.vision_end, - qwen_tokens.image_pad, - qwen_tokens.video_pad, + qwen_tokens.vision_start, + qwen_tokens.vision_end, + qwen_tokens.image_pad, + qwen_tokens.video_pad, ] if ignore_token_ids: