Skip to content

Fix gemma3 multimodal export: fixed-size vision resize + empty-decode Gather guard - #431

Merged
justinchuby merged 5 commits into
mainfrom
fix/gemma3-multimodal-export
Jul 30, 2026
Merged

Fix gemma3 multimodal export: fixed-size vision resize + empty-decode Gather guard#431
justinchuby merged 5 commits into
mainfrom
fix/gemma3-multimodal-export

Conversation

@NishantN2005

Copy link
Copy Markdown
Contributor

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-it and verified end-to-end on CPU.

Fix 1 — Vision processor config emits the wrong image resize

Issue. _write_vision_processor_config had no Gemma3 branch, so Gemma3 fell through to the generic-VLM path. That path emits a smart_resize (variable HxW, driven by min_pixels/max_pixels) pipeline with no Permute3D step — leaving a variable-size HWC tensor. But Gemma3's SigLIP vision encoder expects a fixed NCHW pixel_values tensor of shape [batch, 3, image_size, image_size] (896×896). The mismatch causes inference to fail on a pixel_values shape error.

Fix. Add a dedicated Gemma3 branch that writes processor_config.json with 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 trailing Permute3D so 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 time config.model_type is "gemma3_text", not "gemma3". The stable discriminator is therefore the vision sub-config type — we match model_type ∈ {gemma3, gemma3_text} or vision_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 by Where(image_mask, gathered, text_embeds). On decode steps there are no image tokens, so image_features is an empty [0, hidden] tensor. ONNX evaluates the Gather before the Where can 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_features with a single zero row before the Gather so index 0 is always valid, mirroring the existing pattern in components.InputMixer. The pad row is never selected in a correct prefill (the Where mask 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 correct inputs_embeds shape. Reproduces the exact crash when the pad is reverted.
  • auto_export_test.py::test_gemma3_vision_config — asserts the Gemma3 branch writes processor_config.json with the 6-step pipeline, smart_resize == 0, no min_pixels/max_pixels, and Permute3D dims [2,0,1].

Verified end-to-end on CPU.

🤖 Generated with Claude Code

@NishantN2005
NishantN2005 requested a review from a team July 27, 2026 18:22
@github-actions

github-actions Bot commented Jul 27, 2026

Copy link
Copy Markdown

Performance Comparison

Comparing 2227f5d0760cac

Model Metric Baseline Current Delta
bert (feature-extraction) model_size_bytes 359 KB 359 KB +0.0%
bert (feature-extraction) num_nodes 60 60 +0.0%
falcon model_size_bytes 364 KB 364 KB +0.0%
falcon num_nodes 68 68 +0.0%
gemma2 model_size_bytes 428 KB 428 KB +0.0%
gemma2 num_nodes 107 107 +0.0%
gpt2 model_size_bytes 388 KB 388 KB +0.0%
gpt2 num_nodes 54 54 +0.0%
llama model_size_bytes 425 KB 425 KB +0.0%
llama num_nodes 62 62 +0.0%
llama (static-cache) model_size_bytes 425 KB 425 KB +0.0%
llama (static-cache) num_nodes 58 58 +0.0%
mamba (ssm-text-generation) model_size_bytes 296 KB 296 KB +0.0%
mamba (ssm-text-generation) num_nodes 98 98 +0.0%
phi3 model_size_bytes 421 KB 421 KB +0.0%
phi3 num_nodes 60 60 +0.0%
phi3 (static-cache) model_size_bytes 421 KB 421 KB +0.0%
phi3 (static-cache) num_nodes 56 56 +0.0%
qwen2 model_size_bytes 425 KB 425 KB +0.0%
qwen2 num_nodes 62 62 +0.0%
qwen2 (static-cache) model_size_bytes 425 KB 425 KB +0.0%
qwen2 (static-cache) num_nodes 58 58 +0.0%
qwen3_5_moe (hybrid-text-generation) model_size_bytes 506 KB 506 KB +0.0%
qwen3_5_moe (hybrid-text-generation) num_nodes 275 275 +0.0%
qwen3_5_text (hybrid-text-generation) model_size_bytes 458 KB 458 KB +0.0%
qwen3_5_text (hybrid-text-generation) num_nodes 129 129 +0.0%
qwen3_5_vl (hybrid-qwen-vl) model_size_bytes 977 KB 977 KB +0.0%
qwen3_5_vl (hybrid-qwen-vl) num_nodes 413 413 +0.0%
t5 (seq2seq) model_size_bytes 836 KB 836 KB +0.0%
t5 (seq2seq) num_nodes 166 166 +0.0%
whisper (speech-to-text) model_size_bytes 1008 KB 1008 KB +0.0%
whisper (speech-to-text) num_nodes 128 128 +0.0%

No performance regressions.

@codecov

codecov Bot commented Jul 27, 2026

Copy link
Copy Markdown

The author of this PR, NishantN2005, is not an activated member of this organization on Codecov.
Please activate this user on Codecov to display this PR comment.
Coverage data is still being uploaded to Codecov.io for purposes of overall coverage calculations.
Please don't hesitate to email us at support@codecov.io with any questions.

@github-actions

github-actions Bot commented Jul 27, 2026

Copy link
Copy Markdown

🏗️ Architecture Diff

Comparing 2227f5d0760cac

Model Sub-model Changes Status

No architecture changes detected.


Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed)

justinchuby and others added 2 commits July 29, 2026 17:18
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>
@CLAassistant

CLAassistant commented Jul 29, 2026

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

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
justinchuby force-pushed the fix/gemma3-multimodal-export branch from b86f369 to f4ae468 Compare July 30, 2026 03:52
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
justinchuby force-pushed the fix/gemma3-multimodal-export branch from f4ae468 to 0760cac Compare July 30, 2026 03:57
@justinchuby
justinchuby merged commit 63cd4cb into main Jul 30, 2026
24 checks passed
@justinchuby
justinchuby deleted the fix/gemma3-multimodal-export branch July 30, 2026 04:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants