Skip to content

feat(models): port FastVLM (FastViTHD vision + Qwen2 text)#661

Merged
inureyes merged 2 commits into
mainfrom
feat/issue-538-fastvlm
Jul 5, 2026
Merged

feat(models): port FastVLM (FastViTHD vision + Qwen2 text)#661
inureyes merged 2 commits into
mainfrom
feat/issue-538-fastvlm

Conversation

@inureyes

@inureyes inureyes commented Jul 5, 2026

Copy link
Copy Markdown
Member

Summary

Add FastVLM (Apple, llava_qwen2 / fastvlm): a FastViTHD hybrid convolutional/transformer vision encoder feeding a stock Qwen2 text decoder through an mlp2x_gelu MLP projector.

  • Vision tower (src/vision/encoders/fastvlm.rs): the FastViTHD tower runs entirely on channels-last maps. A 3-block conv stem, three RepMixer stages (depthwise token mixing plus a BatchNorm ConvFFN with a layer scale), two attention stages (channel LayerNorm, head_dim 32 self-attention over the flattened map, dual layer scales), large-kernel PatchEmbed downsamples, RepCPE position encoders, and a squeeze-excite conv_exp head. The final (16, 16, 3072) map is flattened to 256 tokens. Reuses Conv2dLayer (grouped conv) and inference BatchNorm from existing modules.
  • Projector: reuses MLPProjector (mm_projector.0 / .2 renamed to linear_1 / linear_2); rejects any mm_projector_type other than mlp2x_gelu.
  • Processor (src/vision/processors/fastvlm.rs): pad-to-square, bicubic resize to 1024, rescale, per-channel normalize, channels-first (B, 3, 1024, 1024).
  • Prompt / runtime: the <image> placeholder is the fixed -200 sentinel (not a vocabulary token). A dedicated runtime arm re-tokenizes the rendered prompt, splices one sentinel per image, expands each to 256 tokens, preprocesses images, and merges through the generic VisionModule (LLaVA scatter). The -200 id survives the embedding lookup (mx::take) and is overwritten by the scatter.
  • Loader (src/loading/vlm_fastvlm.rs): text reuses Qwen2 unchanged (qkv biases auto-load from weight presence). Handles both the genuine (apple/FastVLM-0.5B, channels-first convs, (C,1,1) layer scales, patch_embed.<n>) and converted (mlx-community/FastVLM-0.5B-bf16, already channels-last) layouts; the conv layout is decided once from an unambiguous depthwise probe. Drops the classification head, BatchNorm counters, and the tied lm_head.

Validation

Real checkpoint mlx-community/FastVLM-0.5B-bf16:

  • Text-only: "What is the capital of France?" produced "Paris."
  • Shapes image (256 image tokens): produced "three distinct geometric shapes: a red circle, a blue square, and a green triangle. The red circle is positioned in the upper left corner ... the blue square is located in the upper right corner. The green triangle is situated in the lower center." Colors, shapes, and positions are all correct.

6 co-located unit tests and 4 real-model parity tests (tests/fastvlm_parity.rs) pass. Release build (metal,accelerate), cargo fmt, and cargo clippy --all-targets -- -D warnings are clean.

Out of scope

Quantized FastVLM conversions, LoRA on the vision tower, non-mlp2x_gelu projectors, the classification head, video, and dynamic resolutions (the model is fixed 1024x1024 pad-to-square).

Drive-by

Ignore in (cxx unsound advisory). mlxcel does not use that macro, so the path is unreachable; the advisory was published after the last merge and fails cargo-deny on every PR until ignored.

Closes #538

Add FastVLM (Apple, llava_qwen2 / fastvlm): a FastViTHD hybrid convolutional/transformer vision encoder feeding a stock Qwen2 text decoder through an mlp2x_gelu MLP projector. The FastViTHD tower (src/vision/encoders/fastvlm.rs) runs entirely on channels-last maps: a 3-block conv stem, three RepMixer stages (depthwise token mixing plus a BatchNorm ConvFFN with a layer scale), two attention stages (channel LayerNorm, head_dim 32 self-attention over the flattened map, dual layer scales), large-kernel PatchEmbed downsamples, RepCPE position encoders, and a squeeze-excite conv_exp head; the final (16, 16, 3072) map is flattened to 256 tokens and projected to the decoder width. It reuses Conv2dLayer (grouped conv) and inference BatchNorm from existing modules. The <image> placeholder is the fixed -200 sentinel, spliced one per image and expanded to 256 tokens, then scattered through the generic VisionModule (LLaVA merge); -200 survives the embedding lookup and is overwritten by the scatter. Text reuses Qwen2 unchanged (qkv biases auto-load from weight presence). The loader accepts both the genuine (apple/FastVLM-0.5B, channels-first convs) and converted (mlx-community/FastVLM-0.5B-bf16, already channels-last) weight layouts, deciding the conv layout once from an unambiguous depthwise probe, and drops the classification head, BatchNorm counters, and the tied lm_head. Validated on mlx-community/FastVLM-0.5B-bf16: text 'Paris', and a shapes image described as a red circle, a blue square, and a green triangle with correct positions. 6 unit tests and 4 real-model parity tests pass.
@inureyes inureyes added type:enhancement New features, capabilities, or significant additions area:models Model architectures, weights, loading, metadata priority:low Low priority labels Jul 5, 2026
The cxx advisory RUSTSEC-2026-0202 flags the let_cxx_string! macro for an uninitialized value on a C++ exception. mlxcel never uses that macro (the FFI bridge passes owned Rust types, not stack CxxString values), so the unsound path is unreachable. This unblocks cargo-deny, which fails on every PR since the advisory was published. Remove once an upstream cxx release ships the fix.
@inureyes inureyes merged commit bb4ac36 into main Jul 5, 2026
5 checks passed
@inureyes inureyes deleted the feat/issue-538-fastvlm branch July 5, 2026 12:04
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 priority:low Low priority type:enhancement New features, capabilities, or significant additions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat(models): port FastVLM (FastViTHD vision + text)

1 participant