Skip to content

Move Gemma4 per-layer embeddings to embedding model (Alternative A) - #296

Merged
justinchuby merged 4 commits into
mainfrom
gemma4-perlayer-fused
May 18, 2026
Merged

Move Gemma4 per-layer embeddings to embedding model (Alternative A)#296
justinchuby merged 4 commits into
mainfrom
gemma4-perlayer-fused

Conversation

@justinchuby

@justinchuby justinchuby commented May 8, 2026

Copy link
Copy Markdown
Member

Summary

Replace L separate [V, D] per-layer embedding tables with a single fused [V, L*D] table, matching HuggingFace's original weight layout.

Background

The per-layer split (nn.ModuleList of L separate embeddings) was a workaround for the ORT CUDA Gather int32 overflow bug. Each [V, D] table had only 67M elements (V=262144, D=256), staying under the 2.1B int32 limit.

That ORT bug is now fixed (merged in ORT 1.27), so the fused [V, L*D] table (2.35B elements for E2B, 2.82B for E4B) works correctly with int64 indexing.

Changes

Model (__init__):

  • Replace nn.ModuleList of L Gemma3TextScaledWordEmbedding(V, D) with single Gemma3TextScaledWordEmbedding(V, L*D)
  • embed_scale remains sqrt(D) (not sqrt(L*D))

Forward pass (_compute_per_layer_inputs):

  • Single Gather on fused table → [B, S, L*D]
  • Reshape to [B, S, L, D]
  • Slice per layer (same as projection slicing)

Weight loading (preprocess_weights):

  • Remove split logic in Gemma4CausalLMModel (was splitting [V, L*D] → L × [V, D])
  • Remove split logic in Gemma4MultimodalCausalLMModel (same)
  • HF weight model.embed_tokens_per_layer.weight maps directly to ONNX parameter

Testing

  • L1: All 15 Gemma4 tests pass
  • Full L1: 2736 passed, 0 failed
  • Net code change: -39 lines, +25 lines

@justinchuby
justinchuby requested review from a team and Copilot May 8, 2026 20:54
@github-actions

github-actions Bot commented May 8, 2026 •
edited
Loading

Copy link
Copy Markdown

🏗️ Architecture Diff

Comparing 47da8cb → 3881873

Model Sub-model Changes Status
bert (feature-extraction) model 0 ⚪
falcon model 0 ⚪
gemma2 model 0 ⚪
gemma4 (gemma4) decoder 122 🔴
gemma4 (gemma4) embedding 0 ⚪
gemma4 (gemma4) vision_encoder 0 ⚪
gemma4_text model 0 ⚪
gpt2 model 0 ⚪
llama model 0 ⚪
llama (static-cache) model 0 ⚪
mamba (ssm-text-generation) model 0 ⚪
phi3 model 0 ⚪
phi3 (static-cache) model 0 ⚪
qwen model 0 ⚪
qwen (static-cache) model 0 ⚪
qwen2 model 0 ⚪
qwen2 (static-cache) model 0 ⚪
qwen2_moe model 0 ⚪
qwen2_moe (static-cache) model 0 ⚪
qwen3 model 0 ⚪
qwen3 (static-cache) model 0 ⚪
qwen3_5_moe (hybrid-text-generation) model 0 ⚪
qwen3_5_text (hybrid-text-generation) model 0 ⚪
qwen3_5_vl (hybrid-qwen-vl) decoder 0 ⚪
qwen3_5_vl (hybrid-qwen-vl) embedding 0 ⚪
qwen3_5_vl (hybrid-qwen-vl) vision_encoder 0 ⚪
qwen3_moe model 0 ⚪
qwen3_moe (static-cache) model 0 ⚪
qwen3_next (hybrid-text-generation) model 0 ⚪
t5 (seq2seq) decoder 0 ⚪
t5 (seq2seq) encoder 0 ⚪
whisper (speech-to-text) decoder 0 ⚪
whisper (speech-to-text) encoder 0 ⚪
gemma4 (gemma4) / decoder — 122 change(s)

Op summary: 127 → 127 nodes

No op-sequence changes.

Connectivity changes:

  • node[0] Gather: input_ids [8, 2] → [7, 2]
  • node[1] Gather: input_ids [9, 2] → [8, 2]
  • node[2] Gather: input_ids [10, 2] → [9, 2]
  • node[3] Gather: input_ids [11, 2] → [10, 2]
  • node[4] CumSum: input_ids [1, 12] → [1, 11]
  • node[5] Unsqueeze: input_ids [55, 13] → [54, 12]
  • node[6] Shape: input_ids [3] → [0]
  • node[8] Sub: input_ids [58, 57] → [57, 56]
  • node[9] Slice: input_ids [55, 59, 58, 13] → [54, 58, 57, 12]
  • node[10] Unsqueeze: input_ids [60, 14] → [59, 13]
  • node[11] GreaterOrEqual: input_ids [61, 56] → [60, 55]
  • node[12] Unsqueeze: input_ids [1, 13] → [1, 12]
  • node[13] Cast: input_ids [63] → [62]
  • node[14] And: input_ids [64, 62] → [63, 61]
  • node[15] Where: input_ids [65, 15, 16] → [64, 14, 15]
  • node[16] Cast: input_ids [66] → [65]
  • node[17] Unsqueeze: input_ids [67, 13] → [66, 12]
  • node[18] CumSum: input_ids [1, 12] → [1, 11]
  • node[19] Unsqueeze: input_ids [69, 13] → [68, 12]
  • node[20] Shape: input_ids [3] → [0]
  • node[22] Sub: input_ids [72, 71] → [71, 70]
  • node[23] Slice: input_ids [69, 73, 72, 13] → [68, 72, 71, 12]
  • node[24] Unsqueeze: input_ids [74, 14] → [73, 13]
  • node[25] GreaterOrEqual: input_ids [75, 70] → [74, 69]
  • node[26] Unsqueeze: input_ids [1, 13] → [1, 12]
  • node[27] Cast: input_ids [77] → [76]
  • node[28] And: input_ids [78, 76] → [77, 75]
  • node[29] Where: input_ids [79, 15, 16] → [78, 14, 15]
  • node[30] Cast: input_ids [80] → [79]
  • node[31] Unsqueeze: input_ids [81, 13] → [80, 12]
  • node[32] RMSNormalization: input_ids [0, 18] → [0, 17]
  • node[33] Transpose: input_ids [19] → [18]
  • node[34] MatMul: input_ids [83, 84] → [82, 83]
  • node[35] Reshape: input_ids [85, 20] → [84, 19]
  • node[36] RMSNormalization: input_ids [86, 21] → [85, 20]
  • node[37] Reshape: input_ids [87, 22] → [86, 21]
  • node[38] RotaryEmbedding: input_ids [88, 51, 52] → [87, 50, 51]
  • node[39] Transpose: input_ids [23] → [22]
  • node[40] MatMul: input_ids [83, 90] → [82, 89]
  • node[41] Reshape: input_ids [91, 20] → [90, 19]
  • node[42] RMSNormalization: input_ids [92, 24] → [91, 23]
  • node[43] Reshape: input_ids [93, 22] → [92, 21]
  • node[44] RotaryEmbedding: input_ids [94, 51, 52] → [93, 50, 51]
  • node[45] Transpose: input_ids [25] → [24]
  • node[46] MatMul: input_ids [83, 96] → [82, 95]
  • node[48] Reshape: input_ids [97, 98] → [96, 97]
  • node[49] Cast: input_ids [99] → [98]
  • node[50] Mul: input_ids [100, 100] → [99, 99]
  • node[51] ReduceMean: input_ids [101, 26] → [100, 25]
  • node[53] Add: input_ids [102, 103] → [101, 102]
  • node[54] Sqrt: input_ids [104] → [103]
  • node[55] Div: input_ids [100, 105] → [99, 104]
  • node[56] CastLike: input_ids [106, 99] → [105, 98]
  • node[57] Reshape: input_ids [107, 22] → [106, 21]
  • node[58] Attention: input_ids [89, 95, 108, 68, 4, 5] → [88, 94, 107, 67, 3, 4]
  • node[59] Transpose: input_ids [27] → [26]
  • node[60] MatMul: input_ids [109, 112] → [108, 111]
  • node[61] RMSNormalization: input_ids [113, 28] → [112, 27]
  • node[62] Add: input_ids [0, 114] → [0, 113]
  • node[63] RMSNormalization: input_ids [115, 29] → [114, 28]
  • node[64] Transpose: input_ids [30] → [29]
  • node[65] MatMul: input_ids [116, 117] → [115, 116]
  • node[66] Sigmoid: input_ids [118] → [117]
  • node[67] Mul: input_ids [118, 119] → [117, 118]
  • node[68] Transpose: input_ids [31] → [30]
  • node[69] MatMul: input_ids [116, 121] → [115, 120]
  • node[70] Mul: input_ids [120, 122] → [119, 121]
  • node[71] Transpose: input_ids [32] → [31]
  • node[72] MatMul: input_ids [123, 124] → [122, 123]
  • node[73] RMSNormalization: input_ids [125, 33] → [124, 32]
  • node[74] Add: input_ids [115, 126] → [114, 125]
  • node[75] Mul: input_ids [127, 17] → [126, 16]
  • node[76] RMSNormalization: input_ids [128, 35] → [127, 34]
  • node[77] Transpose: input_ids [36] → [35]
  • node[78] MatMul: input_ids [129, 130] → [128, 129]
  • node[79] Reshape: input_ids [131, 20] → [130, 19]
  • node[80] RMSNormalization: input_ids [132, 37] → [131, 36]
  • node[81] Reshape: input_ids [133, 22] → [132, 21]
  • node[82] RotaryEmbedding: input_ids [134, 53, 54] → [133, 52, 53]
  • node[83] Transpose: input_ids [38] → [37]
  • node[84] MatMul: input_ids [129, 136] → [128, 135]
  • node[85] Reshape: input_ids [137, 20] → [136, 19]
  • node[86] RMSNormalization: input_ids [138, 39] → [137, 38]
  • node[87] Reshape: input_ids [139, 22] → [138, 21]
  • node[88] RotaryEmbedding: input_ids [140, 53, 54] → [139, 52, 53]
  • node[89] Transpose: input_ids [40] → [39]
  • node[90] MatMul: input_ids [129, 142] → [128, 141]
  • node[92] Reshape: input_ids [143, 144] → [142, 143]
  • node[93] Cast: input_ids [145] → [144]
  • node[94] Mul: input_ids [146, 146] → [145, 145]
  • node[95] ReduceMean: input_ids [147, 26] → [146, 25]
  • node[97] Add: input_ids [148, 149] → [147, 148]
  • node[98] Sqrt: input_ids [150] → [149]
  • node[99] Div: input_ids [146, 151] → [145, 150]
  • node[100] CastLike: input_ids [152, 145] → [151, 144]
  • node[101] Reshape: input_ids [153, 22] → [152, 21]
  • node[102] Attention: input_ids [135, 141, 154, 82, 6, 7] → [134, 140, 153, 81, 5, 6]
  • node[103] Transpose: input_ids [41] → [40]
  • node[104] MatMul: input_ids [155, 158] → [154, 157]
  • node[105] RMSNormalization: input_ids [159, 42] → [158, 41]
  • node[106] Add: input_ids [128, 160] → [127, 159]
  • node[107] RMSNormalization: input_ids [161, 43] → [160, 42]
  • node[108] Transpose: input_ids [44] → [43]
  • node[109] MatMul: input_ids [162, 163] → [161, 162]
  • node[110] Sigmoid: input_ids [164] → [163]
  • node[111] Mul: input_ids [164, 165] → [163, 164]
  • node[112] Transpose: input_ids [45] → [44]
  • node[113] MatMul: input_ids [162, 167] → [161, 166]
  • node[114] Mul: input_ids [166, 168] → [165, 167]
  • node[115] Transpose: input_ids [46] → [45]
  • node[116] MatMul: input_ids [169, 170] → [168, 169]
  • node[117] RMSNormalization: input_ids [171, 47] → [170, 46]
  • node[118] Add: input_ids [161, 172] → [160, 171]
  • node[119] Mul: input_ids [173, 34] → [172, 33]
  • node[120] RMSNormalization: input_ids [174, 48] → [173, 47]
  • node[121] Transpose: input_ids [49] → [48]
  • node[122] MatMul: input_ids [175, 176] → [174, 175]
  • node[123] CastLike: input_ids [50, 177] → [49, 176]
  • node[124] Div: input_ids [177, 178] → [176, 177]
  • node[125] Tanh: input_ids [179] → [178]
  • node[126] Mul: input_ids [180, 178] → [179, 177]

Interface changes:

  • input count 8 → 7; input[3]: dtype INT64 → FLOAT; shape ['?', '?'] → ['?', '2', '?', '16']

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

@github-actions

github-actions Bot commented May 8, 2026 •
edited
Loading

Copy link
Copy Markdown

Performance Comparison

Comparing 47da8cb → 3881873

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 66 66 +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 53 53 +0.0% ⚪
llama model_size_bytes 425 KB 425 KB +0.0% ⚪
llama num_nodes 61 61 +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 59 59 +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 61 61 +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 408 408 +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 May 8, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 91.17647% with 6 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/mobius/models/gemma4.py 88.46% 2 Missing and 4 partials ⚠️

📢 Thoughts on this report? Let us know!

Copilot AI reviewed May 8, 2026
View reviewed changes

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR simplifies Gemma4’s optional per-layer input embedding path by switching from L separate [V, D] embedding tables to a single fused [V, L*D] table, aligning Mobius’ parameter layout with HuggingFace’s embed_tokens_per_layer.weight and removing weight-splitting logic.

Changes:

  • Replace per-layer nn.ModuleList embeddings with a single fused Gemma3TextScaledWordEmbedding(vocab_per_layer, L*D).
  • Update _compute_per_layer_inputs() to do one fused Gather and reshape to [B, S, L, D], then slice per-layer embeddings.
  • Remove fused-weight splitting logic in preprocess_weights() for both text-only and multimodal Gemma4 models.

Comment thread src/mobius/models/gemma4.py Outdated
@justinchuby
Move Gemma4 per-layer embeddings to embedding model (Alternative A)
6bbb2cd 
Move per-layer input embedding computation from the decoder to the
embedding sub-model. This eliminates the decoder's dependency on
input_ids for VLM 3-model split, simplifying the runtime interface.

Architecture change (VLM 3-model split):
- Embedding model now computes per_layer_inputs [B, S, L*D] alongside
  inputs_embeds, using a fused [V, L*D] table (ORT#28107 fixed)
- Decoder accepts per_layer_inputs as a graph input instead of
  computing them internally from input_ids
- Decoder no longer needs input_ids in its graph signature

Text-only single-model path (Gemma4CausalLMModel) is unchanged:
Gemma4TextModel retains per-layer weights and _compute_per_layer_inputs
for models that have input_ids available directly.

Weight routing (Gemma4MultimodalCausalLMModel.preprocess_weights):
- embed_tokens_per_layer, per_layer_model_projection,
  per_layer_projection_norm weights now route to embedding.* instead
  of decoder.model.*

Changes:
- gemma4.py: Add per-layer components to Gemma4EmbeddingModel, update
  _Gemma4DecoderModel to accept per_layer_inputs, keep dual path in
  Gemma4TextModel for text-only vs VLM split
- _gemma4.py: Update _build_decoder (add per_layer_inputs input,
  remove input_ids), update _build_embedding (add per_layer_inputs
  output)
- auto_export_test.py: Update Gemma4 genai config tests

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
@justinchuby
justinchuby force-pushed the gemma4-perlayer-fused branch from 5e828ba to 6bbb2cd Compare May 8, 2026 21:23
@justinchuby justinchuby changed the title Simplify Gemma4 per-layer embeddings: use fused [V, L*D] table Move Gemma4 per-layer embeddings to embedding model (Alternative A) May 8, 2026
@justinchuby
Fix lint: lowercase local variable name in preprocess_weights
7d30095 
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
@justinchuby justinchuby mentioned this pull request May 8, 2026
Merged
@justinchuby
justinchuby requested a review from Copilot May 8, 2026 22:22
Copilot AI reviewed May 8, 2026
View reviewed changes

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

Comment thread src/mobius/integrations/ort_genai/auto_export_test.py
Comment thread src/mobius/integrations/ort_genai/auto_export_test.py
@justinchuby
Address review: ORT version comment, drop input_ids from decoder mock…
3ef5750 
…/assertions

- Add ORT >= 1.27 requirement comment for fused [V, L*D] Gather
- Remove input_ids from mock Gemma4 decoder inputs (decoder no longer
  needs input_ids in VLM split)
- Add 'input_ids not in decoder_inputs' assertions in both mock and
  real-model genai config tests

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
@justinchuby
justinchuby enabled auto-merge (squash) May 9, 2026 01:43
@justinchuby

justinchuby commented May 9, 2026

Copy link
Copy Markdown
Member Author 
FAILED tests/e2e_golden_test.py::TestL4CheckpointVerified::test_prefill_argmax_matches_golden[text-generation/gemma-4-e2b] - ValueError: Required inputs (['per_layer_inputs']) are missing from input feed (['inputs_embeds', 'past_key_values.0.key', 'past_key_values.0.value', 'past_key_values.1.key', 'past_key_values.1.value', 'past_key_values.2.key', 'past_key_values.2.value', 'past_key_values.3.key', 'past_key_values.3.value', 'past_key_values.4.key', 'past_key_values.4.value', 'past_key_values.5.key', 'past_key_values.5.value', 'past_key_values.6.key', 'past_key_values.6.value', 'past_key_values.7.key', 'past_key_values.7.value', 'past_key_values.8.key', 'past_key_values.8.value', 'past_key_values.9.key', 'past_key_values.9.value', 'past_key_values.10.key', 'past_key_values.10.value', 'past_key_values.11.key', 'past_key_values.11.value', 'past_key_values.12.key', 'past_key_values.12.value', 'past_key_values.13.key', 'past_key_values.13.value', 'past_key_values.14.key', 'past_key_values.14.value', 'attention_mask', 'position_ids']).
FAILED tests/e2e_golden_test.py::TestL4CheckpointVerified::test_prefill_argmax_matches_golden[image-text-to-text/gemma-4-e2b-it] - ValueError: Required inputs (['per_layer_inputs']) are missing from input feed (['inputs_embeds', 'past_key_values.0.key', 'past_key_values.0.value', 'past_key_values.1.key', 'past_key_values.1.value', 'past_key_values.2.key', 'past_key_values.2.value', 'past_key_values.3.key', 'past_key_values.3.value', 'past_key_values.4.key', 'past_key_values.4.value', 'past_key_values.5.key', 'past_key_values.5.value', 'past_key_values.6.key', 'past_key_values.6.value', 'past_key_values.7.key', 'past_key_values.7.value', 'past_key_values.8.key', 'past_key_values.8.value', 'past_key_values.9.key', 'past_key_values.9.value', 'past_key_values.10.key', 'past_key_values.10.value', 'past_key_values.11.key', 'past_key_values.11.value', 'past_key_values.12.key', 'past_key_values.12.value', 'past_key_values.13.key', 'past_key_values.13.value', 'past_key_values.14.key', 'past_key_values.14.value', 'attention_mask', 'position_ids']).
FAILED tests/e2e_golden_test.py::TestL4CheckpointVerified::test_prefill_argmax_matches_golden[speech-language/gemma-4-e2b-it-audio] - onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Unexpected input data type. Actual: (tensor(float)) , expected: (tensor(bool))
FAILED tests/e2e_golden_test.py::TestL4CheckpointVerified::test_prefill_argmax_matches_golden[text-generation/gemma-4-e4b] - Failed: Timeout (>300.0s) from pytest-timeout.
FAILED tests/e2e_golden_test.py::TestL4CheckpointVerified::test_prefill_argmax_matches_golden[image-text-to-text/gemma-4-e4b-it] - ValueError: Required inputs (['per_layer_inputs']) are missing from input feed (['inputs_embeds', 'past_key_values.0.key', 'past_key_values.0.value', 'past_key_values.1.key', 'past_key_values.1.value', 'past_key_values.2.key', 'past_key_values.2.value', 'past_key_values.3.key', 'past_key_values.3.value', 'past_key_values.4.key', 'past_key_values.4.value', 'past_key_values.5.key', 'past_key_values.5.value', 'past_key_values.6.key', 'past_key_values.6.value', 'past_key_values.7.key', 'past_key_values.7.value', 'past_key_values.8.key', 'past_key_values.8.value', 'past_key_values.9.key', 'past_key_values.9.value', 'past_key_values.10.key', 'past_key_values.10.value', 'past_key_values.11.key', 'past_key_values.11.value', 'past_key_values.12.key', 'past_key_values.12.value', 'past_key_values.13.key', 'past_key_values.13.value', 'past_key_values.14.key', 'past_key_values.14.value', 'past_key_values.15.key', 'past_key_values.15.value', 'past_key_values.16.key', 'past_key_values.16.value', 'past_key_values.17.key', 'past_key_values.17.value', 'past_key_values.18.key', 'past_key_values.18.value', 'past_key_values.19.key', 'past_key_values.19.value', 'past_key_values.20.key', 'past_key_values.20.value', 'past_key_values.21.key', 'past_key_values.21.value', 'past_key_values.22.key', 'past_key_values.22.value', 'past_key_values.23.key', 'past_key_values.23.value', 'attention_mask', 'position_ids']).
FAILED tests/e2e_golden_test.py::TestL4CheckpointVerified::test_prefill_argmax_matches_golden[speech-language/gemma-4-e4b-it-audio] - onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Unexpected input data type. Actual: (tensor(float)) , expected: (tensor(bool))

@justinchuby
justinchuby merged commit 114a1bc into main May 18, 2026
21 of 23 checks passed
@justinchuby
justinchuby deleted the gemma4-perlayer-fused branch May 18, 2026 18:22
apsonawane
apsonawane reviewed May 18, 2026
View reviewed changes
Comment thread src/mobius/models/gemma4.py
proj_i = op.Squeeze(op.Slice(proj, starts=[i], ends=[i + 1], axes=[2]), [2])
pad = op.Constant(value_int=0)
masked_ids = input_ids
if self._image_token_id:

@apsonawane apsonawane May 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

image_token_ids_mask and image_token_id is similar? Store the same values? Even audio_token_ids_mask and audio_token_id?

This comment was marked as resolved.

Sign in to view

@justinchuby justinchuby mentioned this pull request May 19, 2026
Closed
Copilot AI mentioned this pull request May 23, 2026
Merged
@justinchuby justinchuby mentioned this pull request May 27, 2026
Merged
justinchuby added a commit that referenced this pull request May 27, 2026
@justinchuby
Fix L4/L5 e2e harness for Gemma4 per_layer_inputs + BOOL audio mask (#…
3ce0444 
…318)

## Problem

L4 / L5 e2e tests have been failing on `main` for every Gemma4 case
since
#296 (Move Gemma4 per-layer embeddings to embedding model). Two distinct
root causes:

### 1. `per_layer_inputs` missing from decoder feed

After #296 the embedding sub-model emits a second output
`per_layer_inputs` and the decoder accepts it as a required input. The
e2e harness only wired the first embedding output (`inputs_embeds`)
into the decoder, so every multi-model Gemma4 path failed with:

```
ValueError: Required inputs (['per_layer_inputs']) are missing from input feed
(['inputs_embeds', 'past_key_values.0.key', ..., 'attention_mask', 'position_ids'])
```

Affected: text-only on multi-model (`text-generation/gemma-4-e2b`),
VL prefill (`image-text-to-text/gemma-4-e2b-it`, `…-e4b-it`), VL
generation (L5), speech-language prefill+generation.

### 2. `input_features_mask` dtype mismatch

The Gemma4 audio encoder (`_gemma4.py:367`) declares
`input_features_mask` as `tensor(bool)`, but the harness unconditionally
cast every feature-extractor output to `np.float32`, producing:

```
InvalidArgument: Unexpected input data type. Actual: (tensor(float)), expected: (tensor(bool))
```

Even when the fallback all-True mask path was taken (line 1267), the
bool numpy array crashed `ort_easy`'s DLPack-first conversion path
because DLPack has no native bool type code.

## Fix

### `tests/e2e_golden_test.py`

In every decoder-feed setup (5 sites: text-only multi-model prefill,
VL prefill, VL generation, speech-language prefill, speech-language
generation), wire any extra embedding outputs through to the decoder by
name. For models without `per_layer_inputs` the extra loop iteration is
a no-op:

```python
elif name in emb_out:
    dec_feeds[name] = emb_out[name]
```

For both audio-encoder setup sites (L4 + L5 paths), honor the session's
declared input dtype:

```python
target_dtype = audio_session.get_input_dtype(name) or np.float32
audio_feeds[name] = audio_processed[name].astype(target_dtype)
```

And use the session-declared dtype for the fallback all-True mask too.

### `src/mobius/_testing/ort_inference.py`

Route bool numpy arrays through `OrtValue.ortvalue_from_numpy` directly
in `_numpy_to_ort_value`, bypassing `ort_easy`'s DLPack-first path
(which has no bool type code).

## Verification (locally on H200)

| Test | Before | After |
|---|---|---|
| L4 `text-generation/gemma-4-e2b` | FAIL (missing per_layer_inputs) |
**PASS** |
| L4 `image-text-to-text/gemma-4-e2b-it` | FAIL (missing
per_layer_inputs) | **PASS** |
| L4 `speech-language/gemma-4-e2b-it-audio` | FAIL (bool dtype mismatch)
| **PASS** |
| L5 `image-text-to-text/gemma-4-e2b-it` | FAIL | **PASS** |
| L5 `speech-language/gemma-4-e2b-it-audio` | FAIL | **PASS** |

5 passed, 1 skipped (L5 text-gen is gated by integration markers under
fast runs), 0 failed. `ruff check` + `ruff format --check` both pass.

L1 + L3 + non-gemma e2e suites are unaffected (the embedding-output
loop is a generic no-op for models that don't emit extra outputs;
audio-mask changes only trigger when the session declares a bool input).

## Why this didn't get caught earlier

The CI matrix only runs L4 / L5 in the affected-models lane, and #296
was tested standalone before this lane started exercising the multi-
model Gemma4 path through the e2e harness. The bool-mask issue is even
older — it was masked by the harness's unconditional
`astype(np.float32)`
until the audio_encoder was upgraded to take a real BOOL mask.

Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
@justinchuby justinchuby mentioned this pull request May 27, 2026
Merged
justinchuby added a commit that referenced this pull request May 27, 2026
@justinchuby
examples/gemma4_multimodal.py: pass per_layer_inputs from embedding t…
7f19c93 
…o decoder

After #296 ("Move Gemma4 per-layer embeddings to embedding model") the
Gemma4 embedding sub-model emits a second output ``per_layer_inputs``
([B, S, num_layers * per_layer_dim]) that the decoder consumes on every
step (prefill + decode). The example wasn't updated to forward this
output, so any Gemma4 build with hidden_size_per_layer_input > 0
(e.g. google/gemma-4-E2B-it) blows up at the first decoder.run() with:

  ValueError: Required inputs (['per_layer_inputs']) are missing from
  input feed (['inputs_embeds', ..., 'attention_mask', 'position_ids'])

Fix:

  * prepare_decoder_feeds(): accept an optional per_layer_inputs
    argument and add it to the feeds dict when present. Builds without
    per-layer inputs (hidden_size_per_layer_input == 0, e.g. larger
    Gemma 4 variants) keep working since the kwarg is optional.

  * generate(): pull embed_out.get("per_layer_inputs") and pass it to
    prepare_decoder_feeds on every step.

Verified locally: 'python examples/gemma4_multimodal.py --mode text
--prompt "What is 2+2?"' now generates "2 + 2 = **4**".

Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
justinchuby added a commit that referenced this pull request May 28, 2026
@justinchuby
examples/gemma4_multimodal.py: pass per_layer_inputs to decoder (#320)
ccf96d6 
## Summary

After #296 ("Move Gemma4 per-layer embeddings to embedding model") the
Gemma4 embedding sub-model emits a second output `per_layer_inputs`
(`[B, S, num_layers * per_layer_dim]`) that the decoder consumes on
every step. `examples/gemma4_multimodal.py` wasn't updated to forward
this output, so any Gemma 4 build with `hidden_size_per_layer_input > 0`
(e.g. `google/gemma-4-E2B-it`) crashes at the first `decoder.run()`:

```
ValueError: Required inputs (['per_layer_inputs']) are missing from input feed
(['inputs_embeds', ..., 'attention_mask', 'position_ids'])
```

## Fix

- `prepare_decoder_feeds()`: accept an optional `per_layer_inputs`
argument and add it to the feeds dict when present. Builds without
per-layer inputs (`hidden_size_per_layer_input == 0`, e.g. larger Gemma
4 variants) keep working — the kwarg is optional.
- `generate()`: pull `embed_out.get("per_layer_inputs")` and pass it to
`prepare_decoder_feeds` on every step.

## Verification

```
$ python examples/gemma4_multimodal.py --mode text --prompt "What is 2+2?"
...
📝  TEXT-ONLY GENERATION
================================================================
Prompt: What is 2+2?
----------------------------------------------------------------
2 + 2 = **4**
```

## Related

This is the example-side counterpart to #318 which fixed the same gap in
the L4/L5 e2e test harness.

---------

Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
feich-ms added a commit that referenced this pull request Jul 1, 2026
@feich-ms @claude
Revert _gemma4.py to main (internal per-layer path was pre-PR#296 code)
e4173a3 
The internal input_ids path and compute_per_layer_inputs kwarg depended
on gemma4.py changes that were already reverted when we restored the
fused-table architecture from PR #296. Revert to main's external
per_layer_inputs path to keep the two files consistent.

Co-Authored-By: Claude <noreply@anthropic.com>
feich-ms added a commit that referenced this pull request Jul 8, 2026
@feich-ms @claude
Revert _gemma4.py to main (internal per-layer path was pre-PR#296 code)
dc56190 
The internal input_ids path and compute_per_layer_inputs kwarg depended
on gemma4.py changes that were already reverted when we restored the
fused-table architecture from PR #296. Revert to main's external
per_layer_inputs path to keep the two files consistent.

Co-Authored-By: Claude <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@apsonawane apsonawane apsonawane left review comments
Copilot code review Copilot
Copilot review effort
Applies to this pull request for everyone.Learn more about Copilot code review.
Copilot left review comments
@kunal-vaishnavi kunal-vaishnavi kunal-vaishnavi approved these changes
@titaiwangms titaiwangms Awaiting requested review from titaiwangms

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@justinchuby @apsonawane @kunal-vaishnavi