fix(qwen2_5_vl): image description drops an object that transformers reports (blue circle in a three-shape image) #1596

Problem / Background

On the Qwen2.5-VL path, a synthetic three-object image (red filled square, blue filled circle, green filled triangle on a light gray background) is described as "The image contains a red square and a green triangle." The blue circle is dropped. transformers names all three objects for the same PNG and prompt under greedy decoding. Measured 2026-09-02 on main at 5b57a10 (the #1423 fix, PR #1582) and later, Apple M1 Ultra, release build with metal,accelerate.

The omission is not a quantization artifact. The raw HuggingFace export (Qwen/Qwen2.5-VL-3B-Instruct, bf16 weights, loaded through the loader fixed in #1423) produces the identical sentence as the mlx-community 4-bit conversion, so the raw and converted paths agree with each other and both disagree with transformers. It also predates #1423: the 4-bit output was byte-identical before and after PR #1582.

The repository's only image fixture, tests/fixtures/test_image.png, is a solid orange square. Every VLM test that uses it (tests/vlm_concurrency.rs, tests/smolvlm_parity.rs, tests/hunyuan_vl_parity.rs, the other tests/*_parity.rs, and src/commands/generate_tests.rs:47) can only catch "no output" or "non-finite output". None of them can catch an omitted object, so this defect class has no coverage today.

Evidence

Prompt for all three rows: What shapes and colors are in this image? Answer briefly. Same 448x448 PNG for all three rows.

Current Behavior

The image enters through Qwen2VLProcessor (src/vision/processors/qwen2_vl.rs). smart_resize (:248-279) rounds to multiples of 28 under the code defaults DEFAULT_MIN_PIXELS = 3136 and DEFAULT_MAX_PIXELS = 12845056 (:44-52); the loader builds the processor from those defaults and does not read preprocessor_config.json (src/loading/vlm.rs:775-786), which is harmless for this checkpoint because its preprocessor_config.json carries the same two values. A 448x448 input therefore stays 448x448 (no resample), gives grid_thw = (1, 32, 32), 1024 patches of 3 * 2 * 14 * 14 = 1176 features (preprocess_with_grid, :294-311, emits 2048 rows of 588 that PatchEmbed::forward regroups), and 256 <|image_pad|> tokens (id 151655) between <|vision_start|> (151652) and <|vision_end|>, inserted by insert_qwen_vl_image_tokens (src/multimodal/qwen_vl.rs:634). Normalization uses the OpenAI CLIP mean/std (:176-177), matching the checkpoint.

Qwen25VLModel::get_input_embeddings_with_cache (src/vision/qwen2_5_vl.rs:68-133) casts pixel_values to the text embedding dtype (:92-94), runs Qwen25VLVisionEncoder::forward_with_grid (src/vision/encoders/qwen2_5_vl.rs:755-826), scatters the result over the image-token slots with merge::merge_llava_any (:112-117), and computes MRoPE positions with compute_qwen_vl_mrope_position_ids (src/multimodal/qwen_vl.rs:828) plus rope_deltas = max_pos + 1 - seq_len (:120-130).

Inside the tower: PatchEmbed::forward (:235-257, a matmul against the Conv3d filter flattened in [out, T, C, H, W] order), rot_pos_emb (:560-607, built on VisionRotaryEmbedding::forward from src/vision/encoders/qwen2_vl.rs:204-218 and applied by apply_rotary_pos_emb_vision, qwen2_vl.rs:357), get_window_index (:614-732, vit_merger_window_size = 112 / 2 / 14 = 4, so a 16x16 merged grid becomes 16 windows of 16 merged tokens), 32 VisionBlocks with full attention at fullatt_block_indexes = [7, 15, 23, 31] and windowed attention elsewhere via per-segment attention_from_ptr (VisionAttention::forward, :289-380), then PatchMerger::forward (:489-495, RMSNorm, reshape to 5120, GELU MLP to 2048) and the un-reorder at :805-823.

On Apple Silicon, finish_vlm_weights_common (src/loading/vlm.rs:290-313) converts every bf16 tensor of a non-quantized checkpoint to f16 through convert_bf16_weights (src/models/sanitize.rs:1872), so the bf16 export runs its whole vision tower in f16. The 4-bit checkpoint runs the tower through the quantized linears that inherit_qwen_vision_quantization selects (src/loading/vlm_qwen.rs:255) with bf16 scales.

Reproduction

Generate the image (Pillow):

from PIL import Image, ImageDraw
img = Image.new("RGB", (448, 448), (245, 245, 245))
d = ImageDraw.Draw(img)
d.rectangle([40, 60, 200, 220], fill=(220, 30, 30))
d.ellipse([250, 80, 410, 240], fill=(30, 60, 220))
d.polygon([(90, 400), (230, 260), (370, 400)], fill=(30, 170, 60))
img.save("shapes.png")

mlxcel (greedy is the default, -t 0):

cargo build --release --features metal,accelerate
PROMPT='What shapes and colors are in this image? Answer briefly.'
./target/release/mlxcel generate -m models/mlx/qwen2.5-vl-3b-4bit --image shapes.png -p "$PROMPT" -n 48 -t 0
./target/release/mlxcel generate -m models/mlx/qwen2.5-vl-3b-hf --image shapes.png -p "$PROMPT" -n 48 -t 0

Both print The image contains a red square and a green triangle. on M1 Ultra at 5b57a10.

transformers oracle (bf16, MPS, greedy):

import torch
from PIL import Image
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
model_id = "Qwen/Qwen2.5-VL-3B-Instruct"
processor = AutoProcessor.from_pretrained(model_id)
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(model_id, torch_dtype=torch.bfloat16).to("mps")
messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "What shapes and colors are in this image? Answer briefly."}]}]
text = processor.apply_chat_template(messages, add_generation_prompt=True)
inputs = processor(text=[text], images=[Image.open("shapes.png")], return_tensors="pt").to("mps")
out = model.generate(**inputs, max_new_tokens=48, do_sample=False)
print(processor.batch_decode(out[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0])

Prints square, circle, triangle, red, blue, green.

Proposed Solution

1. Reproduce with the snippet above on both local checkpoints and confirm the transformers output on the same PNG before changing anything. Record the exact strings.

2. Compare the vision-side pipeline against transformers stage by stage on this image and stop at the first stage that diverges. Tap points on the mlxcel side, with the transformers counterpart and the expected shape, are: (a) processor output, preprocess_with_grid versus processor.image_processor(images=[img])["pixel_values"], expected [1024, 1176] after the PatchEmbed regroup, grid_thw = (1, 32, 32), and 256 image tokens in the tokenized prompt on both sides; (b) PatchEmbed::forward versus model.visual.patch_embed, [1024, 1280]; (c) rot_pos_emb versus model.visual.rot_pos_emb(grid_thw), [1024, 40], plus get_window_index and both cu_seqlens lists versus model.visual.get_window_index(grid_thw), which must be integer-identical; (d) the output of every VisionBlock (forward hooks on model.visual.blocks[i] on the transformers side), reporting mean absolute difference per block in the window-reordered token order; (e) PatchMerger::forward versus model.visual.merger, [256, 2048], after undoing the window order on both sides; (f) the merged input embeddings at the 256 image slots; (g) compute_qwen_vl_mrope_position_ids versus get_rope_index(input_ids, image_grid_thw), [3, 1, seq_len], integer-identical, and rope_deltas; (h) the prefill logits at the last position, whose argmax must agree before greedy decode can. Run the transformers side in f32 on CPU for the tap comparison so its own bf16 rounding does not mask a real gap; a mean absolute difference that jumps by an order of magnitude between two consecutive stages is the divergence.

3. While mapping stage (e), the un-reorder in forward_with_grid (src/vision/encoders/qwen2_5_vl.rs:805-823) was checked algebraically and by simulation: it sorts (value, index) pairs and writes reverse_indices[orig_idx] = rank, which for a permutation of 0..N reproduces window_index itself rather than argsort(window_index) as upstream does (https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/qwen2_5_vl/vision.py reverse_indices = mx.argsort(window_index, axis=0)). Applying take(h, window_index) twice is only the identity when the permutation is an involution, which happens to hold for a 16x16 merged grid with 4x4 windows (this issue's 448x448 image and any image whose merged width is 16), so it cannot explain the omission reported here. For other sizes it scrambles the merged tokens: simulating the same get_window_index gives 120 of 144 tokens misplaced at 336x336, 360 of 400 at 560x560, and 352 of 384 at 448x672. The correct construction is reverse_indices[rank] = orig_idx. Fix it in this issue because the stage comparison in step 2 must also be run at a non-involution size (336x336) to be meaningful, and confirm the effect on model output at that size. The identical construction in src/vision/encoders/youtu_vl_window.rs:147-159 (reverse_window_indices, used by youtu_vl.rs:425) is a separate family and goes to a follow-up issue, linked from here, rather than into this PR.

4. Check whether the f16 tower explains the 448x448 omission: keep every vision_tower.* tensor in f32 for one forward (a keep predicate through convert_bf16_weights_with_keep, src/models/sanitize.rs:1882, then astype to f32, plus a matching f32 cast of pv at src/vision/qwen2_5_vl.rs:93) on the bf16 export, and compare the merger output and the generated sentence against the f16 run and against transformers. If f32 restores the circle, locate the overflowing op with the per-block diffs from step 2(d) and upcast only that op (RMSNorm input, attention scores, or the SwiGLU intermediate), not the whole tower, and keep the quantized checkpoint's bf16 scales untouched per the Apple Silicon precision rules. If f32 does not restore it, the defect is structural and the stage diff from step 2 is the answer.

5. Add a content-rich fixture and a parity test. Fixture: tests/fixtures/test_image_shapes.png, the exact 448x448 image from the snippet, committed together with tests/fixtures/generate_test_image_shapes.py (same convention as tests/fixtures/generate_inkling_dmel_numpy.py) so it is reproducible; a flat-color PNG of this kind is a few KB. Test: new tests/qwen2_5_vl_parity.rs with mod common;, resolving the checkpoint with repo_model_dir("qwen2.5-vl-3b-4bit") exactly as tests/vlm_concurrency.rs:41 does and falling back to repo_model_dir("mlx/qwen2.5-vl-3b-4bit"), marked #[ignore = "requires the local qwen2.5-vl-3b-4bit checkpoint"] (precedent tests/chat_template_kwargs.rs:122) and returning early with an eprintln! skip message when neither directory exists so -- --ignored stays green on a machine without the model. The test spawns the built binary through common::repo_binary_path("mlxcel") with generate -m <dir> --image tests/fixtures/test_image_shapes.png -p "<prompt>" -n 48 -t 0, so it exercises the same loader, processor, tower, merge and decode path as the CLI, and asserts case-insensitively that stdout contains all of square, circle, triangle, red, blue, green. Also add a checkpoint-free unit test in src/vision/encoders/qwen2_5_vl_tests.rs for the extracted reverse-index helper: input [2, 0, 1] yields [1, 2, 0], and for the (1, 24, 24) grid the composed permutation is the identity.

Scope

In scope: src/vision/encoders/qwen2_5_vl.rs (tower, un-reorder, any precision fix found), src/vision/qwen2_5_vl.rs (pixel dtype cast if the fix lands there), src/vision/processors/qwen2_vl.rs only if stage (a) diverges, tests/fixtures/test_image_shapes.png, tests/fixtures/generate_test_image_shapes.py, tests/qwen2_5_vl_parity.rs, src/vision/encoders/qwen2_5_vl_tests.rs, and docs/supported-models.md only if the fix changes a documented limitation.

Out of scope: quantized-weight numerics (the raw bf16 export shows the same omission); the Youtu-VL copy of the reverse-index construction (src/vision/encoders/youtu_vl_window.rs:147-159, file a follow-up and link it here); the Lanczos3 versus bicubic resample mismatch (normalize_resized_frame, src/vision/processors/qwen2_vl.rs:452, versus the Qwen2VLImageProcessor default BICUBIC), which cannot act at 448x448 because no resample happens there, and which belongs in its own issue if stage (a) diverges at a non-multiple-of-28 size; reading min_pixels / max_pixels from preprocessor_config.json, which is a no-op for this checkpoint.

Implementation Notes

• Reuse: the stage taps on the mlxcel side already exist as public or crate-visible functions (Qwen2VLProcessor::preprocess_with_grid, Qwen25VLVisionEncoder::forward_with_grid, compute_qwen_vl_mrope_position_ids); add an output_hidden_states-style per-block capture only if it is needed for the diff and keep it out of the hot path. Use transformers as the oracle for this issue (mlx-vlm shares the same un-reorder and window logic and would mask a shared mistake); docs/adding-models.md (Choosing and Using Reference Implementations) covers when each reference applies.
• Shared code: apply_rotary_pos_emb_vision, VisionRotaryEmbedding (src/vision/encoders/qwen2_vl.rs), Qwen2VLProcessor, and src/multimodal/qwen_vl.rs are used by Qwen2-VL, Qwen3-VL, and ColQwen2.5 as well. If the fix touches any of them, update the // Used by: comment per docs/code-guidelines.md and smoke-test qwen2-vl-2b-4bit on the same PNG.
• Constraints: quantized checkpoints keep bf16 scales and biases; do not promote them. Any f32 upcast must be scoped to the offending op and measured for decode and prefill cost on the 4-bit checkpoint (scripts/bench_decode.sh), since a whole-tower f32 run is a diagnostic, not a fix.
• Edge cases: the un-reorder fix must be validated with two images in one prompt (448x448 plus 224x448, where window_index_id offsets accumulate) and with a padded window grid (an image whose merged grid is not a multiple of 4, for example 1036x1036 giving 37x37), both against transformers.
• Error handling: no new failure paths; the parity test skips with a printed reason when the checkpoint is absent and panics with the full stdout when an object is missing, as tests/vlm_concurrency.rs:271-281 does for the empty-output case.

Acceptance Criteria

• mlxcel generate -m models/mlx/qwen2.5-vl-3b-4bit --image shapes.png -p "What shapes and colors are in this image? Answer briefly." -n 48 -t 0 names the square, the circle, and the triangle with their colors, and the same command on models/mlx/qwen2.5-vl-3b-hf does too.
• The stage-by-stage comparison is recorded in the PR: for each tap (a) through (h), the mean absolute difference or integer equality against transformers on the 448x448 image and on a 336x336 image, and the first diverging stage is named.
• The f16-versus-f32 tower experiment result on the bf16 export is recorded in the PR, with the conclusion (precision or structural) and, if precision, the single op that was upcast.
• The un-reorder in forward_with_grid equals argsort(window_index); the checkpoint-free unit test in src/vision/encoders/qwen2_5_vl_tests.rs covers [2, 0, 1] and the (1, 24, 24) grid; a 336x336 image gives the same description on mlxcel and transformers.
• A follow-up issue for src/vision/encoders/youtu_vl_window.rs:147-159 is filed and linked here.
• tests/fixtures/test_image_shapes.png and tests/fixtures/generate_test_image_shapes.py exist, the PNG is byte-identical to the script's output, and tests/qwen2_5_vl_parity.rs passes under -- --ignored with the checkpoint present and skips with a printed reason without it.
• Integrated into the real code flow: the fix is on the path used by mlxcel generate --image and by mlxcel-server chat completions with an image part, not in a test-only helper; tests/vlm_concurrency.rs::qwen2_5_vl_single_row_bench_prefill_succeeds_with_m5_warmup still passes.
• qwen2-vl-2b-4bit on the same PNG is unchanged or improved if any shared file was touched.
• cargo test --workspace --profile test-fast --features metal,accelerate, cargo clippy --workspace --all-targets --features metal,accelerate -- -D warnings, and cargo fmt --all -- --check are green. (narrow scope run here: the module suites, --test qwen2_5_vl_parity, --test vlm_concurrency qwen2_5_vl, cargo clippy --profile test-fast --lib --tests --features metal,accelerate -- -D warnings and cargo fmt --all -- --check are all green; the full workspace gate runs centrally)

Verification

cargo build --release --features metal,accelerate
PROMPT='What shapes and colors are in this image? Answer briefly.'
./target/release/mlxcel generate -m models/mlx/qwen2.5-vl-3b-4bit --image tests/fixtures/test_image_shapes.png -p "$PROMPT" -n 48 -t 0
./target/release/mlxcel generate -m models/mlx/qwen2.5-vl-3b-hf --image tests/fixtures/test_image_shapes.png -p "$PROMPT" -n 48 -t 0
python3 tests/fixtures/generate_test_image_shapes.py /tmp/regen.png && cmp /tmp/regen.png tests/fixtures/test_image_shapes.png
cargo test --profile test-fast --features metal,accelerate -p mlxcel qwen2_5_vl
cargo test --profile test-fast --features metal,accelerate --test qwen2_5_vl_parity -- --ignored --nocapture
cargo test --profile test-fast --features metal,accelerate --test vlm_concurrency qwen2_5_vl
cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets --features metal,accelerate -- -D warnings
cargo fmt --all -- --check

Pass looks like: both generate runs print a sentence containing square, circle, and triangle with red, blue, and green; cmp prints nothing; the ignored parity test reports 1 passed; the workspace gate, clippy, and fmt exit 0. Manual step: run the transformers snippet from the Reproduction section once on the final fixture and paste its output next to the mlxcel output in the PR.

Technical Considerations

Related: #1423 and PR #1582 (raw HF Conv3d patch-embed layout, the loader this issue's bf16 measurement depends on), #294 (earlier qwen2.5-vl-3b-4bit M5 warmup abort, unrelated to output content), #866 (XLA windowed vision attention for Qwen2.5-VL, which will need the corrected un-reorder). Upstream reference for the tower: https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/qwen2_5_vl/vision.py.