feat(models): port FastVLM (FastViTHD vision + Qwen2 text)#661
Merged
Conversation
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.
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.
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
Add FastVLM (Apple,
llava_qwen2/fastvlm): a FastViTHD hybrid convolutional/transformer vision encoder feeding a stock Qwen2 text decoder through anmlp2x_geluMLP projector.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_dim32 self-attention over the flattened map, dual layer scales), large-kernel PatchEmbed downsamples, RepCPE position encoders, and a squeeze-exciteconv_exphead. The final(16, 16, 3072)map is flattened to 256 tokens. ReusesConv2dLayer(grouped conv) and inferenceBatchNormfrom existing modules.MLPProjector(mm_projector.0/.2renamed tolinear_1/linear_2); rejects anymm_projector_typeother thanmlp2x_gelu.src/vision/processors/fastvlm.rs): pad-to-square, bicubic resize to 1024, rescale, per-channel normalize, channels-first(B, 3, 1024, 1024).<image>placeholder is the fixed-200sentinel (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 genericVisionModule(LLaVA scatter). The-200id survives the embedding lookup (mx::take) and is overwritten by the scatter.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 tiedlm_head.Validation
Real checkpoint
mlx-community/FastVLM-0.5B-bf16:6 co-located unit tests and 4 real-model parity tests (
tests/fastvlm_parity.rs) pass. Release build (metal,accelerate),cargo fmt, andcargo clippy --all-targets -- -D warningsare clean.Out of scope
Quantized FastVLM conversions, LoRA on the vision tower, non-
mlp2x_geluprojectors, 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