Commit b4ea221
committed
Fix Qwen2.5-VL MROPE implementation — 9 bugs causing wrong bounding box output
When porting Qwen2.5-VL grounding inference to native Swift via mlx-swift-lm,
the model consistently produced wrong bbox_2d coordinates (the "panels" it
detected were hallucinated by hundreds of pixels). After weeks of debugging,
found 9 distinct bugs. With all fixes applied, Swift output matches the Python
reference (mlx-vlm) at 0px delta on all 8 bbox edges.
Full writeup + diagnostics: https://dev.to/nivdvir/building-a-real-time-screen-reader-on-macos-that-actually-works-471
The 9 bugs:
1. MROPE section selection (split-select vs slice-replace)
Multi-Resolution Rotary Position Embedding assigns different frequency
bands to temporal (T), height (H), and width (W). The Swift impl split the
frequency tensor via modulo indexing (i % 3), interleaving frequencies.
Python starts with temporal freqs and overwrites H/W slices in place:
[T_0-15, H_16-39, W_40-63]. Wrong layout destroys attention patterns.
2. invFreq registered as a Module weight
invFreq is a computed constant, not a learned weight. Declaring it as a
property on a Module subclass exposed it to MLX's weight loader, which
either threw keyNotFound or silently overwrote with garbage. Fix: wrap in
a non-Module class (InvFreqBox) to hide from reflection.
3. rope_deltas unused during autoregressive generation
After the prefill pass, cached position IDs were cleared but rope_deltas
were never applied to subsequent tokens. Correct computation:
positionIds = cache_offset + rope_deltas + arange(seqLen). Without deltas,
position embeddings drifted with each generated token.
4. MROPE state not reset between successive images
Cached position IDs and rope deltas from one inference persisted into the
next. Processing a new image meant position embeddings started from the
previous image's offset. Progressively worse results on 2nd, 3rd images.
5. Image resize using 1800px max instead of 1280px
Swift code resized to max 1800px (2688 visual tokens). Python reference
uses 1280px max (1305 visual tokens). The model was trained on 1280px.
1800px pushed visual token positions outside the training distribution.
6. Chat template ordering (text vs image token placement)
Swift message generator placed text before the image token in the content
array. Python puts image first: <|vision_start|><|image_pad|><|vision_end|>PROMPT.
Ordering matters: text tokens attending to positions where image features
have not yet been injected produces wrong attention patterns.
7. Vision attention mask ignored — THE ROOT CAUSE
The vision encoder's self-attention uses a mask for windowed attention
(each patch only attends to patches within its window). Swift passed
mask: .none to scaledDotProductAttention instead of mask: .array(floatMask).
Result: every patch attended globally to every other patch, destroying
spatial locality the model relies on for precise coordinate prediction.
This single bug was most responsible for bbox inaccuracy.
Tests against Python reference (mlx-vlm 0.1.31, Qwen2.5-VL-7B-Instruct-4bit):
- Before fixes: bounding boxes off by 200-800px, inconsistent across runs
- After fixes: 0px delta on all 8 bbox edges (x1,y1,x2,y2 for 2 panels)
Filed in context of upstream issue ml-explore#221 (Upstreaming improvements from
fast-moving forks). Happy to split into smaller PRs if maintainers prefer.
Related bugs ml-explore#6 (prompt format) and ml-explore#7 (maxTokens) live in consumer code,
not mlx-swift-lm, so they're out of scope for this PR.1 parent 8c9dd63 commit b4ea221
1 file changed
Lines changed: 410 additions & 41 deletions
0 commit comments