feat(models): port ERNIE-4.5 MoE VL (DFNRope ViT + modality-split MoE + 3D MRoPE)#662
Merged
Conversation
… + 3D MRoPE)
Add ERNIE-4.5 MoE VL (ernie4_5_moe_vl), Baidu's vision-language MoE. The text decoder (src/models/ernie4_5_moe_vl.rs) extends the ERNIE-4.5 MoE family with two structures new to the tree: modality-split expert banks (separate text and multimodal routers plus stacked expert banks selected per token by token type, with an e_score_correction_bias that shifts expert selection but never the mixing weights, and a fused shared-experts SwiGLU added for every token) and interleaved 3D MRoPE ([T, H, W] axes assigned per frequency index with H on even and W on odd indices below mrope_section[0]+[1] and T above, rotating adjacent element pairs in f32; a unit test proves the text-only degeneration to fast_rope traditional). The DFNRope ViT (src/vision/encoders/ernie4_5_vl.rs) reuses the Qwen2-VL packed cu_seqlens attention and 2D merge-window RoPE helpers with a plain 588-wide linear patch embedding and quick_gelu MLP; the variable-resolution resampler (src/vision/connectors/ernie4_5_vl.rs) folds 2x2 spatial windows and even/odd temporal pairs before projecting to the decoder width. The processor emits 588-wide patch rows in merge-window order (load-bearing for the resampler fold), and the runtime arm reuses the Qwen2-VL placeholder insertion and MRoPE state plumbing. Also fixes the SentencePiece tokenizer for non-special added tokens: ids outside the SentencePiece vocab that are declared in added_tokens_decoder with special: false (e.g. ERNIE's image placeholder ids) previously made decode error "Out of range", which the CLI swallowed into empty output; they now decode to their content strings. Validated against mlx-community/ERNIE-4.5-VL-28B-A3B-Thinking-4bit: text reasoning ("Paris"), a shapes image described with correct colors, shapes, and positions, and a red/blue split image localized correctly ("left side is red, right side is blue"). 10 unit tests and 4 real-model parity tests pass (the parity suite runs with --test-threads=1 since two tests each load the 16 GB checkpoint).
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
Adds ERNIE-4.5 MoE VL (
ernie4_5_moe_vl), Baidu's vision-language MoE, with two structures new to the tree.src/models/ernie4_5_moe_vl.rs): each MoE layer holds independent text (gate/switch_mlp) and multimodal (gate_1/switch_mlp_1) expert banks. Routing per bank is softmax in f32, top-k selected overprobs + e_score_correction_bias(the bias shifts selection only), mixing weights are the uncorrected probs renormalized bymax(sum, moe_norm_min). Token-type dispatch takes the text bank for text tokens and the multimodal bank at image-placeholder positions; text-only prompts and every decode step skip the multimodal bank entirely. A fused shared-experts SwiGLU is added for all tokens. Config MoE fields parse in int-or-2-list form.[T, H, W]axes assigned per frequency index (H on even, W on odd indices belowmrope_section[0]+[1]= 44; T above), rotating adjacent element pairs in f32. A unit test proves the exact degeneration tofast_rope(traditional)for scalar positions. Position-id construction and per-sequence delta tracking reuseqwen_mrope_state::MRopeState(PaddleOCR-VL plumbing).src/vision/encoders/ernie4_5_vl.rs): reuses the Qwen2-VL packedcu_seqlensattention, merge-window 2D RoPE, and rotation helpers; a plain 588-wide linear patch embedding (no conv reorder, no bias), quick_gelu MLP, trailing LayerNorm, no in-encoder merger.src/vision/connectors/ernie4_5_vl.rs): 2x2 spatial fold over consecutive merge-window rows, even/odd temporal pair fold (a single image duplicates its frame), GELU MLP stacks with LayerNorms, then a Linear to the decoder width plus RMSNorm.src/vision/processors/ernie4_5_vl.rs): Qwen2-VL-family smart resize (factor 28, CLIP normalization), 588-wide rows in merge-window order (load-bearing for the resampler fold, unit-tested),grid_thw = (1, gh, gw).insert_qwen_vl_image_tokens(ERNIE's<|IMAGE_START|>/<|IMAGE_END|>frame the<|IMAGE_PLACEHOLDER|>run, end = start + 1) and the QwenVlm preparation summary; no new summary variant.Tokenizer fix (required by this family)
decodeon the SentencePiece path errored "Out of range" for ids outside the SentencePiece vocab thatadded_tokens_decoderdeclares withspecial: false(ERNIE's image ids: the sp vocab ends at 100295,<|IMAGE_PLACEHOLDER|>is id 100295 withspecial: false). The CLI swallowed that error into empty output, so image runs displayed nothing while generating correct tokens. Non-special added tokens now decode to their content strings (they are real text per HF semantics and are never skipped).Validation
Real checkpoint
mlx-community/ERNIE-4.5-VL-28B-A3B-Thinking-4bit(28B-A3B; text and resampler 4-bit, vision tower bf16):10 co-located unit tests (MRoPE degeneration, selection-only correction bias, MoE layer rule, int-or-2-list parsing, merge-window row order) and 4 real-model parity tests pass; the parity suite documents
--test-threads=1since two tests each load the 16 GB checkpoint. Release build,cargo fmt, andcargo clippy --all-targets -- -D warningsare clean.Out of scope
Video inputs, the
moe_use_token_type_biastraining path, server per-sequence MRoPE binding beyond the PaddleOCR-VL fallback-slot scope, and grounding post-processing.Closes #535