Skip to content

fix(qwen2_5_vl): f32 tower norm and a correct window inverse - #1601

Merged
inureyes merged 2 commits into
mainfrom
fix/issue-1596-qwen2-5-vl-window-unreorder
Sep 2, 2026
Merged

inureyes merged 2 commits into
mainfrom
fix/issue-1596-qwen2-5-vl-window-unreorder

Conversation

@inureyes

@inureyes inureyes commented Sep 2, 2026 •

Copy link
Copy Markdown
Member

Summary

Qwen2.5-VL dropped an object from every image description because the vision tower's RMSNorm reduction overflowed float16 and erased whole token rows. A second, independent defect returned the merged tokens in the wrong order at any image size whose merged grid is not a multiple of the 4-token attention window. Both are fixed here, and the family gets its first content-asserting test.

The stage diff

Taps (a) through (h) against a transformers float32 oracle (Qwen/Qwen2.5-VL-3B-Instruct, greedy, prompt "What shapes and colors are in this image? Answer briefly."), on the 448x448 fixture and a 336x336 rendering of the same scene. rel is mean absolute difference over the oracle's mean absolute value.

Tap 448x448 before 336x336 before 448x448 after 336x336 after
(a) processor pixel_values, grid_thw, image tokens rel 6.8e-8, (1, 32, 32), 256 rel 6.8e-8, (1, 24, 24), 144 unchanged unchanged
(b) patch_embed [N, 1280] rel 2.2e-4 rel 2.2e-4 rel 2.2e-4 rel 2.2e-4
(c) window_index, cu_window_seqlens, cu_seqlens integer-identical integer-identical integer-identical integer-identical
(c) rot_pos_emb [N, 40] exact, 0.0 exact, 0.0 exact, 0.0 exact, 0.0
(d) vision blocks 0 to 14 rel 4.6e-4 rising to 3.3e-3 same same same
(d) block 15 rel 3.5e-3 rel 3.2e-3 rel 3.0e-3 rel 2.8e-3
(d) block 16 rel 3.8e-2, max 9.4e1 rel 3.6e-2 rel 3.3e-3 rel 3.0e-3
(d) block 17 rel 2.8e-1, max 6.4e3 rel 2.7e-1 rel 4.8e-3 rel 5.2e-3
(d) block 31 rel 8.0e-1, max 1.9e4 rel 7.9e-1 rel 4.6e-3 rel 4.2e-3
(e) merger, window order rel 9.9e-1 rel 9.9e-1 rel 3.3e-3 rel 3.2e-3
(e) merger, raster order (after the un-reorder) rel 9.9e-1 rel 8.5e-1 rel 3.3e-3 rel 3.2e-3
(f) merged embeddings at the image slots not comparable not comparable rel 3.3e-3, 256 slots rel 3.2e-3, 144 slots
(g) MRoPE position ids [3, 1, seq], rope_deltas integer-identical, -240 integer-identical, -132 integer-identical integer-identical
(h) greedy generation "The image contains a red square and a green triangle." "A red circle and a blue triangle." "square, circle, triangle, red, blue, green", token-identical to the oracle names all three shapes with their own colors

First diverging stage: (d) block 16. The relative error is flat at 3.3e-3 through block 14, holds at block 15, then jumps an order of magnitude at block 16 and becomes catastrophic at block 17. Stages (a), (b), (c) and (g) never diverge, so the processor, the patch embedding, the window bookkeeping, the rotary table and the MRoPE indices were all already correct.

The float16 versus float32 tower experiment

Running the same tower with float32 activations (float16 weights, MLX promotion) on the bf16 export brought every block to rel 8e-6 against the oracle, flat from block 0 to block 31, with max|.| agreeing to 23410.50 versus 23410.75 at block 31. That is what settles the question the issue asks: precision, not structure.

The overflowing op is the RMSNorm reduction, and only that one. The tower carries massive activations: one channel of the residual stream reaches 5.4e2 by block 15 and 2.3e4 by block 31. rms_norm reduces sum(x * x) over the 1280-wide feature axis, and in float16 that sum saturates at 65504, so a row whose energy crosses the ceiling reduces to infinity and the norm returns rsqrt(inf) = 0. Measured directly: 15 of 1024 tower rows came out of block 16's norm1 as exact zeros, 21 at block 17, and the threshold sat between a row peak of 213.75 (kept) and 239.5 (erased), which is where the row's sum of squares crosses 65504. Nothing else in the block overflows: norm outputs stay under 17, and the projections accumulate in float32.

So the fix is scoped to the reduction, not the tower. VisionRMSNorm::forward promotes a float16 input to float32 for the norm and returns the tower's own dtype. Cost, A/B'd inside one binary so no build variance leaks in, min over 30 tower forwards: 144.525 ms with the promotion against 144.201 ms without at 448x448, and 95.185 ms against 95.196 ms at 336x336. That is inside run-to-run noise, and decode is untouched because the tower runs once per image at prefill. Quantized scales and biases are not involved: the mlx-community 4-bit conversion ships the vision tower unquantized, which is why the 4-bit and bf16 paths produced byte-identical tower outputs and the same sentence.

The window un-reorder

forward_with_grid sorted (value, index) pairs and wrote reverse_indices[orig_idx] = rank. For a permutation of 0..N the element with value v always lands at rank v, so that loop reproduced window_index itself rather than argsort(window_index), which is what upstream builds (https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/qwen2_5_vl/vision.py). Applying a permutation twice restores the original order only when it is an involution, which a 16-wide merged grid satisfies, so the 448x448 case in the report was unaffected and the defect stayed invisible there.

Measured from the dumped merger output: 0 of 256 merged tokens misplaced at 448x448, 120 of 144 at 336x336. Tap (e) at 336x336 is rel 8.5e-1 with the old construction and rel 3.2e-3 with argsort. invert_window_index now builds the inverse directly, bounds-checked so a malformed grid cannot panic.

What changed

  • src/vision/encoders/qwen2_5_vl.rs: VisionRMSNorm::forward reduces a float16 input in float32; invert_window_index replaces the incorrect inverse construction; get_window_index's body is extracted as the free function window_index_for_grid so the permutation is testable from a grid alone.
  • src/vision/encoders/qwen2_5_vl_tests.rs: four checkpoint-free tests for the window inverse (the minimal non-involution [2, 0, 1], a 24x24 grid, the 32x32 involution grid that hid the defect, and a padded 37x37 plus two-image grid) and one for the float16 RMSNorm reduction. All five fail on the pre-fix code.
  • tests/qwen2_5_vl_parity.rs (new): three #[ignore] real-checkpoint tests that assert on content, not token count. They skip with a printed reason when neither models/qwen2.5-vl-3b-4bit nor models/mlx/qwen2.5-vl-3b-4bit exists. All three fail on the pre-fix code.
  • tests/fixtures/test_image_shapes.png, tests/fixtures/test_image_shapes_336.png, tests/fixtures/generate_test_image_shapes.py (new): a three-object fixture at both a size whose window permutation is an involution and one where it is not, reproducible from the committed script.
  • TECHNICAL_REPORTS/1601-qwen2-5-vl-f16-norm-overflow-20260903.{en,ko}.md (new): the full localization record, the decisions and the alternatives they rule out, and the parts that remain unverified.

Why the existing tests could not catch this

The repository's only image fixture is a 224x224 solid orange square, and every VLM test that uses it asserts only that tokens were produced. Before this PR that fixture was described as "a person wearing a white shirt and black pants, standing in front of a white wall with a black and white patterned rug on the floor". It now reads "The image is a solid orange color." The erased tower rows left the language model with no grounded content and it invented a scene, which a token-count assertion passes. tests/qwen2_5_vl_parity.rs pins that fixture's description too.

Validation

All on an Apple M1 Ultra, metal,accelerate, greedy (-t 0), 48 tokens.

Case Before After transformers oracle
448x448, bf16 export The image contains a red square and a green triangle. square, circle, triangle, red, blue, green square, circle, triangle, red, blue, green
448x448, 4-bit The image contains a red square and a green triangle. The image contains a red square, a blue circle, and a green triangle. as above
336x336, bf16 export The image contains a red circle and a blue triangle. rectangle, circle, triangle, red, blue, green. rectangle, circle, triangle; red, blue, green
336x336, 4-bit A red circle and a blue triangle. The image contains a red square, a blue circle, and a green triangle. as above
1036x1036 (37x37 merged, padded windows), 4-bit Green triangle and blue square. Red square, blue circle, green triangle. Red square, blue circle, green triangle.
Two images (448x448 plus 224x448), 4-bit four invented images, the first repeated names all four real objects across both grid_thw [[1, 32, 32], [1, 32, 16]] matches
tests/fixtures/test_image.png, 4-bit a person wearing a white shirt and black pants ... The image is a solid orange color. n/a
qwen2-vl-2b-4bit on the same image rectangle, circle, triangle rectangle, circle, triangle unchanged, as expected for a separate encoder file

Test plan

  • cargo test --profile test-fast --features metal,accelerate --lib vision::encoders::qwen2_5_vl (9 passed)
  • cargo test --profile test-fast --features metal,accelerate --lib vision::processors::qwen2_vl (4 passed), --lib models::colqwen2_5 (6 passed), --lib loading::vlm (197 passed)
  • cargo test --profile test-fast --features metal,accelerate --test qwen2_5_vl_parity -- --ignored (3 passed; 3 failed before the fix)
  • cargo test --profile test-fast --features metal,accelerate --test vlm_concurrency qwen2_5_vl -- --ignored
  • cargo clippy --profile test-fast --lib --tests --features metal,accelerate -- -D warnings
  • cargo fmt --all -- --check
  • python3 tests/fixtures/generate_test_image_shapes.py /tmp/regen.png && cmp /tmp/regen.png tests/fixtures/test_image_shapes.png (byte-identical, Pillow 12.3.0)

Scope notes

mlxcel_core::rms_norm is shared by every family in the workspace, so the reduction dtype is not changed there: the promotion lives in Qwen2.5-VL's own VisionRMSNorm, which is private to src/vision/encoders/qwen2_5_vl.rs and reaches ColQwen2.5 only through the same encoder. Whether other float16 families hit the same ceiling is a separate question this PR does not answer.

The identical un-reorder construction in src/vision/encoders/youtu_vl_window.rs:147-159 (reverse_window_indices, used by youtu_vl.rs:425) is a different family with its own checkpoint and validation path, and is filed as #1600 rather than fixed blind here.

Two divergences the issue lists as out of scope were confirmed inert at these sizes and left alone: the processor resamples with Lanczos3 where Qwen2VLImageProcessor uses bicubic, and the loader builds the processor from code defaults rather than preprocessor_config.json. Tap (a) is exact at 448x448 and 336x336 because both are multiples of 28, so smart_resize performs no resample, and this checkpoint's preprocessor_config.json carries the same min_pixels and max_pixels as the code defaults.

Closes #1596

Qwen2.5-VL described a synthetic three-shape image as "The image contains a red square and a green triangle.", dropping the blue circle, on both the mlx-community 4-bit conversion and the raw bf16 export of Qwen/Qwen2.5-VL-3B-Instruct, while transformers named all three objects for the same PNG and prompt under greedy decoding.

A stage-by-stage numerical diff against a transformers float32 oracle localized it to the vision tower's RMSNorm. The processor output, the patch grid, the window index, both cu_seqlens lists and the rotary table were bit-identical or at the float32 rounding floor. The per-block relative error held flat at 3.3e-3 through block 14, then jumped an order of magnitude at block 16 and became catastrophic at block 17 (relative 0.28, max absolute 6.4e3). Running the same tower in float32 brought every block back to 8e-6 relative, which is what identified the failure as precision rather than structure.

The tower carries massive activations: one channel of the residual stream reaches 5.4e2 by block 15 and 2.3e4 by block 31. `rms_norm` reduces `sum(x * x)` over the 1280-wide feature axis, and in float16 that sum saturates at 65504, so a row whose energy crosses the ceiling reduces to infinity and the norm returns `rsqrt(inf) = 0`. Fifteen of 1024 tower tokens were erased at block 16 and 21 at block 17. `VisionRMSNorm::forward` now promotes a float16 input to float32 for the reduction and returns the tower's own dtype; the surrounding projections and attention still run in float16. An A/B in one binary put the tower at 144.5 ms with the promotion against 144.2 ms without at 448x448, and 95.185 ms against 95.196 ms at 336x336, so the cost is inside run-to-run noise. Decode is untouched: the tower runs once per image at prefill.

The un-reorder after the patch merger was separately wrong. It sorted `(value, index)` pairs and wrote `reverse_indices[orig_idx] = rank`, which for a permutation of `0..N` reproduces `window_index` itself instead of `argsort(window_index)`. Applying the permutation twice restores raster order only when it is an involution, which a 16-wide merged grid happens to satisfy, so the 448x448 case in the report was unaffected while 120 of 144 merged tokens were misplaced at 336x336 and 352 of 384 at 448x672. `invert_window_index` now builds the inverse directly, and `get_window_index` is extracted as `window_index_for_grid` so the permutation can be exercised on a grid alone.

The float16 overflow was not confined to hard images. The repository's own fixture, a 224x224 solid orange square, was described as "a person wearing a white shirt and black pants, standing in front of a white wall"; it now reads "The image is a solid orange color." Every VLM test that drives that fixture asserts only that tokens were produced, so nothing caught it.

Validation, all on an M1 Ultra with the mlx-community 4-bit conversion and the raw bf16 export:

- 448x448: the bf16 export now emits "square, circle, triangle, red, blue, green", token-identical to the transformers greedy output; the 4-bit conversion emits "The image contains a red square, a blue circle, and a green triangle."
- 336x336 (12x12 merged grid, not an involution): was "A red circle and a blue triangle"; now names all three shapes with their own colors, against a transformers oracle of "rectangle, circle, triangle; red, blue, green".
- 1036x1036 (37x37 merged grid, padded windows): was "Green triangle and blue square."; now "Red square, blue circle, green triangle.", byte-identical to the transformers oracle.
- Two images in one prompt (448x448 plus 224x448): was four invented images; now names all four real objects across both.
- qwen2-vl-2b-4bit on the same image is byte-identical before and after, as expected for a family with its own encoder file.

Tests: four checkpoint-free unit tests for the window inverse (including the involution case that hid the defect, a padded 37x37 grid, and a two-image grid) and one for the float16 RMSNorm reduction; a new `tests/qwen2_5_vl_parity.rs` with three ignored real-checkpoint tests that assert on content rather than token count; and `tests/fixtures/test_image_shapes.png` plus its 336x336 rendering, both reproducible from `tests/fixtures/generate_test_image_shapes.py`. All nine unit tests and all three parity tests fail on the pre-fix code.

`mlxcel_core::rms_norm` is shared by every family, so the reduction dtype is not changed there. The identical un-reorder construction in `src/vision/encoders/youtu_vl_window.rs` is a separate family and is filed as #1600.

Closes #1596
@inureyes inureyes added type:bug Bug fixes, error corrections, or issue resolutions priority:medium Medium priority area:models Model architectures, weights, loading, metadata modeltype:vlm Vision-language model status:review Under review labels Sep 2, 2026
@inureyes inureyes added status:done Completed and removed status:review Under review labels Sep 2, 2026
@inureyes
inureyes merged commit 089eeef into main Sep 2, 2026
13 checks passed
@inureyes
inureyes deleted the fix/issue-1596-qwen2-5-vl-window-unreorder branch September 2, 2026 20:05
rapsealk added a commit to rapsealk/mlxcel that referenced this pull request Sep 3, 2026
…on itself (lablup#1600)

`reverse_window_indices` sorted `(value, index)` pairs and wrote the rank
back by index. For a permutation of `0..N` the element with value `v` always
sorts to rank `v`, so that reproduced `window_index` rather than its inverse,
and the encoder un-reordered the merger output by applying the window
permutation a second time. That is only correct when the permutation is an
involution, which needs the window count per edge to equal the window width.
With Youtu-VL's 8x8 merged-token window that is a single 256x256 image or a
2048x2048 one; at 512x512, 192 of the 256 merged tokens reached the language
model out of place, and at 1024x1024, 960 of 1024.

The helper now builds `inverse[window_index[rank]] = rank` directly,
bounds-checked so a malformed grid cannot panic, the same construction lablup#1601
landed for Qwen2.5-VL. Four checkpoint-free tests at the upstream defaults pin
the minimal `[2, 0, 1]` case, the 512x512 and 384x384 grids with their
misplaced counts, the single-window identity case, and a padded two-image
batch; all four fail on the old construction.

Not done here: the real-checkpoint before/after the issue asks for. No
Youtu-VL checkpoint or Metal toolchain is available on this host; the PR body
says which image sizes to use.

Verified on a CPU-only build: `cargo fmt --all -- --check`, `cargo clippy
--profile test-fast -p mlxcel --all-targets -- -D warnings`, and `cargo test
--profile test-fast -p mlxcel --lib vision::encoders::youtu_vl` (10 passed).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0134qsVMvsvs9LpZgCVH5m3Y
rapsealk added a commit to rapsealk/mlxcel that referenced this pull request Sep 3, 2026
…ndow index

The bounds guard in `reverse_window_indices` exists so a malformed grid
cannot index out of range; the value it leaves behind is not a contract.
Asserting `[0, 7, 1]` maps to `[0, 2, 0]` promised that corruption, so the
assertion is gone. The guard itself matches the Qwen2.5-VL helper from lablup#1601.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0134qsVMvsvs9LpZgCVH5m3Y
inureyes pushed a commit that referenced this pull request Sep 3, 2026
## Summary

`reverse_window_indices` in the Youtu-VL vision encoder returned `window_index` itself instead of its inverse, so the merger output was un-reordered by applying the window permutation a second time. That is only correct when the permutation is an involution. The helper now builds `argsort(window_index)` directly, the same construction #1601 landed for Qwen2.5-VL, with checkpoint-free tests that pin the affected grids.

## Related issues

Closes #1600. Same defect as #1596 / #1601 (Qwen2.5-VL), whose corrected helper this copies.

## Type of change

- [ ] `feat` — new user-visible feature
- [x] `fix` — bug fix
- [ ] `perf` — performance improvement (include before/after numbers in the PR body)
- [ ] `refactor` — internal restructuring without behavior change
- [ ] `chore` — build, CI, dependencies, release infrastructure
- [ ] `docs` — documentation only
- [ ] `test` — tests only

## Which grids are affected

A Youtu-VL window is 8x8 merged tokens (`window_size` 256 over `patch_size` 16 times `spatial_merge_size` 2), and `build_processor` caps the processor at `vision_config.num_patches` (4096), so the patch grid tops out at 64x64 and the merged grid at 32x32. Enumerating every merged grid the loader can actually produce, the permutation is an involution only when the whole image fits in one window, that is when it resizes to 256x256 or smaller:

| image | resized | patch grid | merged grid | windows | merged tokens misplaced by the old un-reorder |
|---|---|---|---|---|---|
| 224x224 | 224x224 | 14x14 | 7x7 | 1x1 | 0 of 49 (the case that hid the defect) |
| 256x256 | 256x256 | 16x16 | 8x8 | 1x1 | 0 of 64 |
| 336x336 | 352x352 | 22x22 | 11x11 | 2x2 padded | 99 of 121 |
| 384x384 | 384x384 | 24x24 | 12x12 | 2x2 padded | 120 of 144 |
| 448x448 | 448x448 | 28x28 | 14x14 | 2x2 padded | 160 of 196 |
| 512x512 | 512x512 | 32x32 | 16x16 | 2x2 | 192 of 256 |
| 768x768 | 768x768 | 48x48 | 24x24 | 3x3 | 528 of 576 |
| 1024x1024 and larger | 1024x1024 | 64x64 | 32x32 | 4x4 | 960 of 1024 |

Every merged grid from 9x9 through 32x32 is a non-involution; there is no second safe size above one window. A grid whose window count per edge equals the window width would be one, but that is a 64x64 merged grid, which needs 16384 patches and is past the 4096 cap, so it is unreachable.

For Qwen2.5-VL the 4-wide window makes the common 16x16 merged grid (448x448) an involution, which is why #1596 was invisible at the default fixture size. Youtu-VL is exposed at every size above a single window instead.

## What changed

- `src/vision/encoders/youtu_vl_window.rs`: `reverse_window_indices` builds `inverse[window_index[rank]] = rank` directly, bounds-checked so a malformed grid cannot panic. The rustdoc records why the old construction degenerated. Name and signature unchanged, so the single call site is untouched.
- `src/vision/encoders/youtu_vl.rs`: the call-site comment no longer claims the old code was equivalent to `argsort`.
- `src/vision/encoders/youtu_vl_tests.rs`: four checkpoint-free tests, all failing on the pre-fix code: the minimal non-involution `[2, 0, 1]` (plus an involution), the 512x512 grid with its misplaced count pinned, the single-window identity case, and a two-image batch with padding.
- `CHANGELOG.md`: entry under Unreleased / Fixed, scoped to the token order at that stage and pointing at #1610 for the reason the family still misdescribes images.

## Real-checkpoint validation, and why it does not show a difference

Acceptance criterion 3 asked for a Youtu-VL run at a non-involution size before and after. It was run on `models/mlx/youtu-vl-4b-instruct` (`tencent/Youtu-VL-4B-Instruct`) on an M1 Ultra, greedy (`-t 0 -n 48`), release build with `metal,accelerate`, at 448x448 and at 336x336, both non-involution per the table. **The output is byte-identical before and after this change, and wrong in both.**

| fixture | before | after |
|---|---|---|
| `test_image_shapes.png` (448, 160 of 196 misplaced) | "The image contains a single, solid black circle on a white background." | identical |
| `test_image_shapes_336.png` (352, 99 of 121 misplaced) | "The image contains a single, solid black circle on a white background." | identical |
| `test_image.png` (224, one window, 0 misplaced) | "The image is completely black and contains no visible content, objects, text, or details." | identical |

The third row is the diagnostic one. 224x224 is a single window, so its permutation is the identity and this change provably cannot alter it, yet the description of a solid **orange** square is already wrong. So a second defect upstream of the window reordering dominates the vision path at every size, and no window-inverse fix can show a content difference until it is fixed. That defect is filed separately as #1610 and is out of scope here per this issue's own Scope section.

This PR is therefore justified on the arithmetic, which is checked directly by the new unit tests rather than inferred from generation quality: the pre-fix construction is reproduced in the tests with its exact misplaced counts, and the new inverse round-trips to the identity on every grid tried.

## Test plan

- [x] `cargo fmt --all -- --check`
- [x] `cargo clippy --workspace --all-targets --features metal,accelerate -- -D warnings`
- [x] `cargo test --workspace --profile test-fast --features metal,accelerate`
- [x] `cargo deny check` (PR CI)
- [x] Real checkpoint run at two non-involution sizes, recorded above

## Notes for reviewers

The math is not in doubt: the tests reproduce the old construction's exact misplaced counts and prove the new inverse round-trips to the identity on every grid tried. What the real-checkpoint run adds is the knowledge that this family has a second, larger defect, which is why the changelog entry claims a corrected token order at that stage rather than a corrected image description.

## Checklist

- [x] PR title uses a Conventional Commits prefix (`feat:`, `fix:`, etc.)
- [x] One logical change per PR (split unrelated changes)
- [x] Updated `docs/` if user-facing behavior or supported models changed (none needed)
- [x] Updated `// Used by: ...` comments on any shared function I modified (unchanged: the module header already carries it)
- [x] No secrets, credentials, or `.env` files committed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:models Model architectures, weights, loading, metadata modeltype:vlm Vision-language model priority:medium Medium priority status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant