Fix Gemma4 CUDA EP support and audio encoder accuracy - #172
Merged
Conversation
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
Updates the Gemma4 multimodal example CLI to better support selecting an execution target, by adding device/provider flags and wiring them into model build + ONNX Runtime session creation.
Changes:
- Add
--deviceto select the ONNX Runtime inference provider used byOnnxModelSession. - Add
--epand pass it tomobius.build(..., execution_provider=...)to generate EP-specific graphs. - Extend
--dtypechoices to includebf16(plus some formatting-only line wraps).
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Three fixes to make examples/gemma4_multimodal.py --device cuda work: 1. Opset 24→23 lowering for CUDA EP: ORT 1.24.4 CUDA EP lacks kernels for standard opset-24 ops (Squeeze, Reshape, etc.). Added automatic opset lowering in OnnxModelSession with a feature flag (MOBIUS_ORT_LOWER_OPSET_FOR_EP, default on). 2. ClippableLinear for audio encoder: HF Gemma4ClippableLinear has learned input/output activation clipping buffers (NOT ±inf as previously assumed). Added ClippableLinear component and replaced plain Linear in FeedForward, Attention, and LightConv1d modules. This reduced audio encoder max diff from 52.68 to 0.0003. 3. Audio HF comparison: Fixed prepare_audio_feeds to use the processor's Gemma4AudioFeatureExtractor (instead of Whisper's). Added _hf_generate_audio() and _hf_generate_vision_audio() for --compare-hf support across all modality modes. Also fixed off-by-one in _build_causal_window_mask (context_left → context_left - 1) matching HF's sliding_window=(context_left-1, 0). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
justinchuby
changed the title
justinchuby
marked this pull request as ready for review
April 16, 2026 22:05
- --ep defaults to None, falls back to --device at runtime - Added 'cpu' to --ep choices - Updated help strings: --device for inference, --ep for build - Follows nemotron_3_nano example convention Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
commented
Apr 16, 2026
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
ORT ≤1.24.x CUDA Gather kernel uses int32 for element offset computation. When the data tensor exceeds INT32_MAX elements (e.g. Gemma4 embed_tokens_per_layer [262144, 8960] = 2.35B), token IDs above row 239674 cause illegal memory access. Split oversized Gather ops into N shards along axis 0, each staying under INT32_MAX elements. Route indices to the correct shard with Less/Where/Sub ops. Only applies to axis=0, rank-2, initializer-backed Gathers when running on CUDA. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Replace the post-hoc _split_large_gathers() IR mutation in ort_inference.py with flag-gated sharding in the Embedding component itself. When ort_shard_large_gathers is enabled (default) and num_embeddings * embedding_dim > INT32_MAX, Embedding.__init__ creates per-shard weight parameters and forward() emits a sharded Gather subgraph using Less/Where/Sub routing. Added Embedding.shard_weight_dict() helper for preprocess_weights to split HF weight tensors into per-shard slices. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Instead of one [V, L*D] embedding table (262144×8960 = 2.35B elements) that overflows ORT CUDA Gather int32 arithmetic, use L separate [V, D] tables (262144×256 = 67M each). Each layer does its own small Gather directly — no post-Gather reshape or axis-2 slicing needed. This eliminates the ORT int32 overflow without any sharding workaround, Where/Less/Sub routing, or runtime IR mutation. The ort_shard_large_gathers flag remains as a generic safety net for other models. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
The per-layer ModuleList approach eliminates oversized Gather ops at the model architecture level, making the generic sharding workaround unnecessary. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Slice is a direct memory copy without index indirection, more natural for constant-index axis extraction. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Coverage for E2B, E4B, 26B-A4B, and 31B across text-generation, image-text-to-text, and speech-language tasks. All cases are skip_reason-gated (requires HF authentication). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
The HF Gemma4 vision encoder uses Gemma4ClippableLinear (with learned
input/output activation clamping) for all linear layers in both the
self-attention (q/k/v/o_proj) and MLP (gate/up/down_proj). Our ONNX
implementation was using plain Linear, missing the clamping. This caused
a max diff of 3.92 (vs max value 5.34) in vision encoder output, leading
to wrong generation output ('cat' instead of 'raccoon').
Changes:
- Use ClippableLinear in Gemma4VisionSelfAttention for q/k/v/o projections
- Pass ClippableLinear as linear_class to MLP in Gemma4VisionEncoderLayer
- Export ClippableLinear from components/__init__.py
- Add audio boundary markers (AUDIO_OPEN/CLOSE tokens) in example script
After fix: vision max diff = 0.000071, all modalities match HF exactly.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
…nerator - Unskip E2B test cases (3 cases: text, vision, audio) - Generate golden JSON files for all 4 Gemma4 sizes: - E2B: causal-lm, vision-language, speech (unskipped) - E4B: causal-lm, vision-language, speech (skipped) - 26B-A4B: causal-lm, vision-language (skipped) - 31B: causal-lm, vision-language (skipped) - Add _generate_speech_language() to generate_golden.py for Gemma4-style multimodal audio models (not Whisper-style) - Total: 20 golden files (10 L4 + 10 L5 generation) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Skills updated with knowledge gained from Gemma4 multimodal work: - reusable-components: Add ClippableLinear component docs, _gemma4_audio.py to directory listing - multimodal-models: Expand to cover audio models, add Gemma4 architecture section (ClippableLinear, audio boundary markers, per-layer embeddings, feature extractor) - debugging-vl-pipeline: Add 5 new failure modes (ClippableLinear, audio markers, CUDA Gather int32 overflow, opset 24 lowering, wrong feature extractor), expand scope to audio modality - phi4mm-component-parity: Add Gemma4 as applicable model, add ClippableLinear and boundary token failure modes - quality-checklist: Add speech-language golden generation note, CUDA EP testing checklist item Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Gemma4Processor accepts `audio=` (singular), not `audios=` (plural). The wrong kwarg was silently ignored, causing the model to receive no audio features and produce 'Please provide the audio...' output. Regenerated golden files for E2B and E4B speech tasks with actual audio input — output now correctly transcribes the test audio. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Generated L4 golden reference files for the two audio-feature-extraction models that had no skip_reason and were missing golden data. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Unskip whisper-tiny: speech-to-text generator already handles audio input_features correctly. Regenerated L4+L5 golden files. - Fix Qwen3-ASR model IDs: Qwen/Qwen3-ASR → Qwen/Qwen3-ASR-0.6B, Qwen/Qwen3-ASR-2B-Instruct → Qwen/Qwen3-ASR-1.7B (actual HF IDs) - Qwen3-ASR still skipped: model_type 'qwen3_asr' not recognized by installed transformers version Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Add _get_test_device_kwargs() reading MOBIUS_TEST_DEVICE/MOBIUS_TEST_EP env vars, pass to all OnnxModelSession calls - Add _make_empty_kv_cache() using ORT declared shapes for per-layer KV cache dimensions (needed for Gemma4 shared KV layers) - Add _run_speech_to_text_prefill() for whisper-style models - Add _run_text_only_multimodel_prefill() for multi-model text-gen - Add _run_speech_language_prefill() for Gemma4 audio - Fix pixel_position_ids naming (HF image_position_ids -> ONNX pixel_) - Provide empty modality tensors for unused inputs in embedding models - Pass input_ids to Gemma4 decoder alongside inputs_embeds - Remove nonexistent qwen3-asr-2b test case - Skip unispeech models (no safetensors weights) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
The previous ID (Qwen/Qwen3-ASR-2B-Instruct) does not exist on HuggingFace. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Unskip qwen3-asr.yaml (trust_remote_code: true, L4 only) - Generate L4 golden from Qwen/Qwen3-ASR-0.6B with patched qwen_asr package (multiple transformers 5.6 compat fixes needed) - L5 generation skipped: qwen_asr package incompatible with transformers generate() in 5.6.0.dev0 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Generated using transformers 4.57.6 + qwen_asr 0.0.6 with the proper
Qwen3ASRProcessor chat template. L4 top1=11528 ('language'), L5 outputs
correct English transcription of the LibriSpeech audio sample.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
Add _try_register_qwen3_asr() to register the qwen_asr config with transformers AutoConfig before loading. This is needed because the HF repo lacks auto_map and trust_remote_code alone won't resolve it. Refactor _generate_speech_language into three functions: - _load_speech_language_model: loads model/processor, detects architecture - _prepare_speech_language_inputs: builds inputs for Gemma4 or Qwen3-ASR - _generate_speech_language: orchestrates the full pipeline The Qwen3-ASR path uses Qwen3ASRProcessor.apply_chat_template() for proper input construction (not manual token assembly). Note: Qwen3-ASR generation requires qwen_asr pip package. The package (v0.0.6) is incompatible with transformers>=5.0; use transformers<=4.57.6. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
- Fix providers kwarg: pass as list not string (was iterating chars) - Fix feature extractor fallback: use WhisperFeatureExtractor when AutoProcessor returns a tokenizer (Qwen3-ASR HF repo lacks auto_map) - Fix audio_encoder key: use 'audio_encoder' when 'audio' not in pkg - Fix audio token count mismatch: adjust input_ids placeholders to match ONNX audio encoder's actual output length - Fix MRoPE position_ids: handle 3D shape (dims, batch, seq_len) - Fix golden JSON format: convert to hex-encoded floats Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Simplified help text for the --ep argument. Signed-off-by: Justin Chu <justinchuby@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.
Gemma4 multimodal CUDA support & audio encoder fixes
Changes
CUDA EP compatibility
ort_lower_opset_for_epfeature flag (default: on).[V, L*D]table (262144×8960 = 2.35B elements) to L separate[V, D]tables (262144×256 = 67M each). Avoids ORT CUDA Gather int32 overflow (onnxruntime#28107) without any sharding workaround — each layer does its own small Gather directly.ort_shard_large_gathersflag +Embedding.shard_weight_dict()safety net for other models that may hit the same int32 limit.Audio encoder accuracy
ClippableLinearcomponent with learned input/output clamping buffers (HFGemma4ClippableLinear). Fixes audio encoder divergence: max diff 52.68 → 0.0003._build_causal_window_mask:context_left → context_left - 1to match HF behavior.Example script (
examples/gemma4_multimodal.py)Gemma4AudioFeatureExtractor(notWhisperFeatureExtractor)--compare-hfsupport for audio and vision+audio modes--ep/--deviceflags with established CLI patternTesting
--compare-hf --device cuda: max diff ~0.0003 (float32)