add Qwen3 VL tests, update logit checker, and golden logit generator#4330
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
subawocit
requested review from
A9isha,
NicoGrande,
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
darisoy,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
igorts-git,
jiangjy1982,
khatwanimohit,
richjames0,
shralex,
suexu1025,
vipannalla and
xibinliu
as code owners
hengtaoguo
left a comment
hengtaoguo
left a comment
There was a problem hiding this comment.
Thanks for the great work!
Looking at the "Text and image input" logs in the PR description:
- Please try to improve the output logs readability.
- We don't need to print out the full token_ids and kl_divergence for the entire prompt, since prompt with image has relatively longer sequence length and it will flush the terminal.
- Print diff values in a consistent format. E.g. I saw both
2.2498674e-09and0.0008294873405247927. Maybe let's consolidate to two decimal representation like2.24e-09 - Could you also print "top 10 tokens" for this one prompt's first generation, similar to the text-only input?
| 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 |
There was a problem hiding this comment.
Could you remind me why this new branch is needed?
There was a problem hiding this comment.
Supported Gemma and Llama models' val_np has a dummy batch dimension of size 1, and val_np[0] was used to strip this batch dimension. But in Qwen3-VL models, the val_np is [num_patches, patch_dim], so val_np[0] would delete the entire image except for the first image patch. If-else statement was added to prevent this.
|
|
||
| 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. |
There was a problem hiding this comment.
Could you also provide a command example here to generate the hf logits?
| 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 |
There was a problem hiding this comment.
Maybe merge some lines together to make this example command more condensed?
e.g.
maxtext/benchmarks/mmlu/mmlu_eval.py
Lines 29 to 34 in b9aee47
| src/maxtext/configs/base.yml \ | ||
| tokenizer_path=<your-bucket>/gemma-2-2b \ | ||
| load_parameters_path=<your-bucket>/gemma-2-2b/maxtext_ckpt/0/items \ | ||
| model_name=gemma-2-2b \ |
There was a problem hiding this comment.
Could you double check the model_name in maxtext? It looks like gemma2-2b is the correct spelling. Or we could use qwen3-0.6b as it's newer and smaller, also belong to the same model family as qwen3-vl.
There was a problem hiding this comment.
Fixed - the code now uses qwen3-0.6b
| else: | ||
| pixel_values = np.transpose(pixel_values, (1, 2, 0)) |
There was a problem hiding this comment.
Could you help me understand why this new else branch for transpose is needed? IIUC we are only updating for the qwen3 vision family.
There was a problem hiding this comment.
Good catch, these two lines are unnecessary, I just fixed the code
| 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"], |
There was a problem hiding this comment.
Make sure this change doesn't impact any other text-only models?
There was a problem hiding this comment.
Checked with Gemma 3 4B (text only input) and Gemma 2 2B. KL were all below threshold and no error detected
| --output-path=${GOLDEN_LOGITS_PATH} \ | ||
| --prompts="${TEST_PROMPT}" \ | ||
| --image-paths=${TEST_IMAGE} \ | ||
| --hf-model-path=${LOCAL_PATH} \ |
There was a problem hiding this comment.
Have you tried directly passing --hf-model-path=Qwen/Qwen3-VL-2B-Instruct to compare against the original HF model?
There was a problem hiding this comment.
Good call - I just did and results look reasonable: max KL divergence = 0.0008294873405247927 at index 346, the corresponding token id is 419
There was a problem hiding this comment.
Great, I would recommend you to directly compare against HF model at this line instead of using LOCAL_PATH from the HF roundtrip conversion model.
As for the to_huggingface command, let's move it towards the end of this file. WDYT?
There was a problem hiding this comment.
SG - I just updated the file
subawocit
force-pushed
the
update-logit-checker
branch
2 times, most recently
from
0579310 to
356b09d
Compare
|
To improve output readability, several code blocks were modified in
Test ResultUpdated Text and image input (USE_MULTIMODAL=true) using Qwen3-VL-2BQwen3 VL 2B |
hengtaoguo
left a comment
hengtaoguo
left a comment
There was a problem hiding this comment.
Thanks for the quick fix!
| --output-path=${GOLDEN_LOGITS_PATH} \ | ||
| --prompts="${TEST_PROMPT}" \ | ||
| --image-paths=${TEST_IMAGE} \ | ||
| --hf-model-path=${LOCAL_PATH} \ |
There was a problem hiding this comment.
Great, I would recommend you to directly compare against HF model at this line instead of using LOCAL_PATH from the HF roundtrip conversion model.
As for the to_huggingface command, let's move it towards the end of this file. WDYT?
subawocit
force-pushed
the
update-logit-checker
branch
from
e2d5f81 to
c9f6111
Compare
shuningjin
left a comment
shuningjin
left a comment
There was a problem hiding this comment.
Thank you for the tool enhancement, along with comprehensive doc and testing! Overall looks good. Some minor comments.
| image_path=${TEST_IMAGE} \ | ||
| --max_kl_div=0.1 \ | ||
| --golden_logits_path=${GOLDEN_LOGITS_PATH} \ | ||
| override_model_config=true |
There was a problem hiding this comment.
nit: Do we need override_model_config=true here (same for below)? I didn't see you change config in commands.
There was a problem hiding this comment.
Depends on how model's YAML config is initialized. In qwen3-vl models, we set use_multimodal: true in the config, and if we define USE_MULTIMODAL=false in .sh test file directly, we'll get this error:
ValueError: Keys ['use_multimodal'] are overridden by both model config and CLI/kwargs with different values.This is not allowed, unless setting override_model_config=True.
I added override_model_config=true to handle this.
subawocit
force-pushed
the
update-logit-checker
branch
2 times, most recently
from
314e5e8 to
8d550d2
Compare
shuningjin
left a comment
shuningjin
left a comment
There was a problem hiding this comment.
Great work! (let's squash commit before merge)
subawocit
force-pushed
the
update-logit-checker
branch
from
ef2af73 to
144fb46
Compare
add Qwen3 VL tests, update logit checker, and golden logit generator Fix pyink formatting Fix trailing whitespace and pyink list formatting Fix final pyink indentation and comma Fix pyink indentation improve example usage; fix code minor usage example change minor update to test file; improve log output format update to test file and description pyink update
subawocit
force-pushed
the
update-logit-checker
branch
from
8d84a47 to
d81f8fa
Compare
3 participants
Description
This PR adds logit checking and end-to-end integration validation support for Qwen3-VL models. It extends both the golden logit generation utility and the forward-pass logit checker to handle multimodal components for Qwen3-VL models. An end-to-end shell script was also added to compare MaxText logits with Hugging Face golden logits.
generate_hf_golden_logits.pyQwen3VLForConditionalGenerationmodel class for"qwen3-vl"model identifiers.forward_pass_logit_checker.py(channels, time * temporal_patch_size, height * patch_size, width * patch_size).processor_qwen3_omni.get_rope_indexwhenuse_mropeis enabled for Qwen3 models.vision_start,vision_end,image_pad,video_pad) during KL divergence calculation. Since the model is not trained to predict next token from raw vision placeholders, their logits can be random and vary numerically. Masking them prevents visual placeholders from skewing the max KL divergence score, and focuses the comparison on actual language-generation logits.test_qwen3_vl_2b_to_hf_e2e.shtests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.shfollowing test_gemma4_to_hf.sh:--max_kl_div=0.1).Tests
Define variables: HF_TOKEN, MODEL_BUCKET, LOCAL_PATH, USE_MULTIMODAL, USE_SCAN_LAYERS in
tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.shTest Results
Text-only input (USE_MULTIMODAL=false)
Text and image input (USE_MULTIMODAL=true)
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.