Fix gemma3 multimodal export: fixed-size vision resize + empty-decode Gather guard - #431
Merged
Merged
Conversation
Performance Comparison
|
|
The author of this PR, NishantN2005, is not an activated member of this organization on Codecov. |
Run lintrunner after merging origin/main and commit the formatting it applied to the Gemma3 multimodal export changes. No functional changes beyond formatting. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Restrict the Gemma3 fixed-resize image processor path to Gemma3 model types so other SigLIP-based VLMs keep their generic preprocessing. Also map config-mode Gemma3 VLM exports back to the multimodal ORT GenAI model type and add regression coverage for both cases. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
fix/gemma3-multimodal-export
branch
from
July 30, 2026 03:52
b86f369 to
f4ae468
Compare
The L4 (prefill argmax) and L5 (30-token generation) golden cases for google/gemma-3-4b-it now run end-to-end on the mobius ONNX pipeline. L5 exercises the empty-image-features decode path this PR fixes: before the _Gemma3EmbeddingModel zero-row Gather pad, L5 crashed on the first decode step (Gather indices out of bounds); it now runs to completion. L4 matches the golden argmax exactly and tokens 0-19 match token-for-token (ORT 1.28, CPU and CUDA both deterministically 20/30). At token 20 the ONNX pipeline picks the reference's rank-2 token (golden 32219 logit 31.039 vs onnx 131371 logit 30.812; gap 0.227) and the greedy paths fork. HF-torch with current transformers reproduces this golden exactly (30/30), confirming the golden is current and the divergence is expected fp32 accumulation over 20 autoregressive decode steps at a near-tie, not a regression. Exact 30/30 is not achievable, so set min_token_match_ratio=0.5 to tolerate the near-tie fork while still catching real regressions (a crash or garbled output falls far below 0.5). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
fix/gemma3-multimodal-export
branch
from
July 30, 2026 03:57
f4ae468 to
0760cac
Compare
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
Fixes two independent bugs in the Gemma3 multimodal (3-model split) export path. Both produce a model that builds successfully but fails or silently misbehaves when run through ORT-GenAI. Found while exporting
gemma-3-4b-itand verified end-to-end on CPU.Fix 1 — Vision processor config emits the wrong image resize
Issue.
_write_vision_processor_confighad no Gemma3 branch, so Gemma3 fell through to the generic-VLM path. That path emits asmart_resize(variable HxW, driven bymin_pixels/max_pixels) pipeline with noPermute3Dstep — leaving a variable-size HWC tensor. But Gemma3's SigLIP vision encoder expects a fixed NCHWpixel_valuestensor of shape[batch, 3, image_size, image_size](896×896). The mismatch causes inference to fail on apixel_valuesshape error.Fix. Add a dedicated Gemma3 branch that writes
processor_config.jsonwith a 6-step pipeline:DecodeImage → ConvertRGB → Resize[fixed] → Rescale → Normalize → Permute3D([2,0,1]). It uses a fixed-size resize (smart_resize: 0) and a trailingPermute3Dso the HWC→CHW layout and fixed dimensions match the encoder's contract. Mean/std/rescale/size are pulled from the HF processor when available, falling back to Gemma3 defaults (image_size=896, mean/std=0.5).Why match on the vision sub-config.
build()unwraps the composite HF config to its text sub-config, so at export timeconfig.model_typeis"gemma3_text", not"gemma3". The stable discriminator is therefore the vision sub-config type — we matchmodel_type ∈ {gemma3, gemma3_text}orvision_model_type == "siglip_vision_model".Fix 2 — Embedding graph Gather goes out of bounds on decode steps
Issue. The embedding sub-model fuses vision features into token embeddings with
Gather(image_features, indices)followed byWhere(image_mask, gathered, text_embeds). On decode steps there are no image tokens, soimage_featuresis an empty[0, hidden]tensor. ONNX evaluates theGatherbefore theWherecan mask it, and gathering into a zero-row tensor throws:indices element out of data bounds, idx=0 must be within the inclusive range [0,-1].Fix. Pad
image_featureswith a single zero row before the Gather so index 0 is always valid, mirroring the existing pattern incomponents.InputMixer. The pad row is never selected in a correct prefill (theWheremask discards it), so outputs are unchanged in the normal case — it only keeps the Gather in-bounds when the modality is absent.Tests
tests/integration_test.py::test_gemma3_embedding_runs_with_empty_image_features— builds the 3-model split with random weights and runs the embedding session with an empty[0, hidden]image_features, asserting the correctinputs_embedsshape. Reproduces the exact crash when the pad is reverted.auto_export_test.py::test_gemma3_vision_config— asserts the Gemma3 branch writesprocessor_config.jsonwith the 6-step pipeline,smart_resize == 0, nomin_pixels/max_pixels, andPermute3Ddims[2,0,1].Verified end-to-end on CPU.
🤖 Generated with Claude Code