Repository navigation
fix(qwen2_5_vl): f32 tower norm and a correct window inverse - #1601
Merged
Merged
Conversation
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
11 of 17 tasks
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
6 tasks done
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
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.relis mean absolute difference over the oracle's mean absolute value.(1, 32, 32), 256(1, 24, 24), 144[N, 1280][N, 40][3, 1, seq], rope_deltasFirst 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_normreducessum(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 returnsrsqrt(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::forwardpromotes 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_gridsorted(value, index)pairs and wrotereverse_indices[orig_idx] = rank. For a permutation of0..Nthe element with valuevalways lands at rankv, so that loop reproducedwindow_indexitself rather thanargsort(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_indexnow builds the inverse directly, bounds-checked so a malformed grid cannot panic.What changed
src/vision/encoders/qwen2_5_vl.rs:VisionRMSNorm::forwardreduces a float16 input in float32;invert_window_indexreplaces the incorrect inverse construction;get_window_index's body is extracted as the free functionwindow_index_for_gridso 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 neithermodels/qwen2.5-vl-3b-4bitnormodels/mlx/qwen2.5-vl-3b-4bitexists. 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.rspins that fixture's description too.Validation
All on an Apple M1 Ultra,
metal,accelerate, greedy (-t 0), 48 tokens.[[1, 32, 32], [1, 32, 16]]matchestests/fixtures/test_image.png, 4-bitqwen2-vl-2b-4biton the same imageTest 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 -- --ignoredcargo clippy --profile test-fast --lib --tests --features metal,accelerate -- -D warningscargo fmt --all -- --checkpython3 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_normis shared by every family in the workspace, so the reduction dtype is not changed there: the promotion lives in Qwen2.5-VL's ownVisionRMSNorm, which is private tosrc/vision/encoders/qwen2_5_vl.rsand 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 byyoutu_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
Qwen2VLImageProcessoruses bicubic, and the loader builds the processor from code defaults rather thanpreprocessor_config.json. Tap (a) is exact at 448x448 and 336x336 because both are multiples of 28, sosmart_resizeperforms no resample, and this checkpoint'spreprocessor_config.jsoncarries the samemin_pixelsandmax_pixelsas the code defaults.Closes #1596