Add Pixtral/Ministral3 VLM support: dynamic vision, FP8 dequant, PatchMerger fix, MHA - #130
Merged
Merged
Conversation
justinchuby
reviewed
Apr 8, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
Adds Pixtral/Ministral-3 VLM enablement across the export pipeline by (1) introducing a Pixtral-specific VLM task with dynamic vision input resolution, (2) supporting direct loading of FP8 HuggingFace checkpoints via dequantization during weight download, and (3) adding Ministral3/Mistral4 YaRN RoPE llama_4_attn_scale handling and applying it in attention.
Changes:
- Add
PixtralVLTaskand registerpixtral/mistral3to use it for dynamic-resolution vision export. - Extend
YarnRopeto optionally return(cos, sin, attn_scale)and applyattn_scaleto queries inAttention. - Add
_dequantize_fp8_weights()to dequantize FP8 checkpoint weights at download time, with new unit tests.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
src/mobius/tasks/_vision_language_3model.py |
Adds PixtralVLTask overriding the vision model input shape to symbolic H/W for dynamic resolution. |
src/mobius/tasks/__init__.py |
Exports PixtralVLTask and registers a new task name (pixtral-vl). |
src/mobius/components/_rotary_embedding.py |
Updates RoPE APIs/docs to allow optional attn_scale; implements YaRN llama_4_scaling_beta scaling output. |
src/mobius/components/_rotary_embedding_test.py |
Adds unit tests for YaRN returning 2-tuple vs 3-tuple and RoPE application compatibility. |
src/mobius/components/_attention.py |
Applies llama_4_attn_scale to query states when present in position_embeddings. |
src/mobius/_weight_loading.py |
Adds FP8 dequantization during weight download and filters auxiliary FP8 scale tensors. |
src/mobius/_weight_loading_test.py |
Adds unit tests for FP8 dequantization behavior and key derivation. |
src/mobius/_registry.py |
Routes pixtral/mistral3 model types to the new pixtral-vl task. |
titaiwangms
force-pushed
the
fix-pixtral-dynamic-vision
branch
from
April 8, 2026 23:33
30a5634 to
fdb1bc2
Compare
Performance Comparison
|
justinchuby
reviewed
Apr 8, 2026
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
titaiwangms
force-pushed
the
fix-pixtral-dynamic-vision
branch
from
April 8, 2026 23:38
fdb1bc2 to
52387fc
Compare
titaiwangms
marked this pull request as ready for review
April 9, 2026 21:35
titaiwangms
force-pushed
the
fix-pixtral-dynamic-vision
branch
from
April 9, 2026 21:37
c4ac85c to
c714125
Compare
|
The author of this PR, titaiwangms, is not an activated member of this organization on Codecov. |
justinchuby
reviewed
Apr 9, 2026
…_scale - Dynamic Pixtral vision export via PixtralVLTask - FP8 dequantization for HF checkpoints - llama_4_attn_scale for long-context attention - FP16 dtype mismatch fix - LlavaEmbedding zero-padding for text-only mode - Add vision parity test and golden test case - 8 unit tests Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
titaiwangms
force-pushed
the
fix-pixtral-dynamic-vision
branch
from
April 10, 2026 00:04
0f1129b to
2897b81
Compare
justinchuby
reviewed
Apr 10, 2026
Member
|
Only comments are LazyTensor + Referece json files, which can also be a follow up if needed |
justinchuby
approved these changes
Apr 10, 2026
justinchuby
reviewed
Apr 10, 2026
…sion output Three fixes for the Pixtral vision model: 1. Mistral3PatchMerger transpose perm [0,1,3,2,4,5] (patch-major) changed to [0,1,3,5,2,4] (dim-major) to match HuggingFace's F.unfold element ordering. The learned merging_layer weights expect HF's ordering — wrong permutation caused cosine sim 0.007 vs HF. 2. PixtralAttention: use com.microsoft.MultiHeadAttention for all EPs except onnx-standard (which falls back to standard Attention since MHA has no function body for InlinePass decomposition). 3. PixtralVLTask._build_vision: Squeeze(axis=0) on vision output to produce rank-2 [num_patches, hidden] — ort-genai expects rank-2. Also addresses PR review comments: - _weight_loading: always return new dict from _dequantize_fp8_weights, update docstring to list all removed suffixes (.input_scale) - _attention: move attn_scale Mul before RoPE to preserve RotaryEmbedding -> Attention pattern for GQA rewrite rule Results: - Cosine sim vs HuggingFace: 0.999973 (was 0.007) - E2E model-mm.py: fish.jpg and challenge.jpg produce correct descriptions - Olive-recipe pipeline verified (optimize.py export + E2E) - All 1261 unit tests pass, lintrunner clean Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
titaiwangms
changed the title
Member
|
@titaiwangms if you have a Olive-recipe e2e, please feel free to add it to examples/olive/ as a demo. Thanks! |
Member
|
@titaiwangms CI failing: https://github.com/onnxruntime/mobius/actions/runs/24267296245/job/70864946872?pr=157 is there a way to run the new ort test without saving model to disk? |
Contributor
Author
let me fix it. |
This was referenced Apr 10, 2026
titaiwangms
added a commit
that referenced
this pull request
Apr 30, 2026
…161) ## Summary Add an end-to-end olive-recipe demo for Ministral-3-3B VLM and enhance the mobius genai integration with Pixtral processor config support and tokenizer class remapping. ## Changes ### Core: genai integration enhancements (`src/mobius/integrations/ort_genai/`) - **`_write_processor_config`**: Enhanced to generate full ORT-extensions image transform pipeline for VL models. Dispatches by model type (Pixtral: longest_edge resize; generic: smart_resize). Derives normalization params from HF processor config. - **`_fix_tokenizer_config`**: New function with `_TOKENIZER_CLASS_REMAP` dict to remap unsupported tokenizer classes (e.g. `TokenizersBackend` → `LlamaTokenizer`). Called automatically from `write_ort_genai_config()`. - 4 new tests (31 total passing) ### Example: olive-recipe demo (`examples/olive/ministral-3-3b-vlm/`) - **optimize.py**: Pure mobius export (`build()` → `save()` → `write_ort_genai_config()`) with optional Olive quantization via `--olive-config` - **inference.py**: ORT GenAI multimodal inference (text-only, image+text, interactive) - **eval.py**: AI2D benchmark evaluation (ONNX vs PyTorch comparison) - Olive configs for CPU (INT4) and CUDA (FP16) quantization - README with setup, export, inference, and evaluation instructions ### Skills - New `olive-recipe` skill documenting the mobius + Olive hybrid pipeline pattern - Updated `ort-genai-config` skill with Pixtral processor config and tokenizer remap info ## References - [Issue #158](#158) - [PR #130](#130) (Pixtral/Ministral3 VLM support) - [olive-recipes PR #352](microsoft/olive-recipes#352) - [onnxruntime-genai PR #2077](microsoft/onnxruntime-genai#2077) Closes #158 --------- Signed-off-by: Ti-Tai Wang <titaiwang@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
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 Pixtral/Ministral3 VLM support with 3-model split (vision/embedding/decoder) for onnxruntime-genai.
Key Changes
Vision encoder (
_pixtral_vision.py)PixtralVisionTower: Conv2d patch embedding → RMSNorm → 24-layer transformer with 2D RoPEPixtralRoPE2D: Precomputed 2D rotary position embeddings over spatial gridPixtralAttention: Bidirectional MHA with 2D RoPE (com.microsoft.MultiHeadAttention;op.Attentionfallback for onnx-standard EP)Mistral3PatchMerger: Spatial 2×2 patch merging matching HFF.unfolddim-major orderingMistral3MultiModalProjector: norm → merge → GELU MLP projectionPatchMerger fix (cosine sim 0.007 → 0.999973)
The original transpose permutation
[0,1,3,2,4,5]produced patch-major ordering, but HFF.unfoldproduces dim-major[D, ms_h, ms_w]. The learnedmerging_layerweights expect HF ordering. Fixed to[0,1,3,5,2,4].FP8 weight dequantization (
_weight_loading.py)fp8.to(bf16) * scaleAttention scale optimization (
_attention.py)attn_scaleMul before RoPE (scaling commutes with rotation) to preserveRotaryEmbedding → Attentionpattern for GQA rewrite rule matchingTask and infrastructure
PixtralVLTask: Dynamic H×W vision input with Squeeze for rank-2 outputpixtral-vltask in_vision_language_3model.pyoptimize.pyintegration verifiedTesting
model-mm.pywith fish.jpg and challenge.jpg: correct descriptionstest_patch_merger_matches_hf_unfold_ordering: regression test against HF F.unfoldFiles Changed
src/mobius/components/_pixtral_vision.py— PatchMerger fix + MHAsrc/mobius/components/_pixtral_vision_test.py— F.unfold ordering testsrc/mobius/tasks/_vision_language_3model.py— Squeeze vision outputsrc/mobius/_weight_loading.py— Docstring + mutation fixsrc/mobius/components/_attention.py— attn_scale before RoPE