Add vision and audio support for Gemma 3n - #474
Merged
Conversation
Gemma 3n expresses intermediate_size as a per-layer list ([8192]*30 for E2B, [16384]*35 for E4B) rather than a scalar. from_transformers copied it through raw, so the list reached nn.Parameter/ir.Shape as a weight dimension and raised TypeError during build. Wrap the intermediate_size extraction in _as_int (already used for hidden_size directly above), which collapses a uniform list to its first element and leaves scalars untouched. Every shipped Gemma 3n checkpoint keeps the list uniform, so the first element is the correct MLP width. Add regression tests for both the scalar pass-through and the list case. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Gemma 3n's trailing decoder layers borrow K,V from an earlier layer instead of projecting their own. HF (modeling_gemma3n.py) has layers at or after ``num_hidden_layers - num_kv_shared_layers`` reuse the K,V of the last preceding layer of the *same* attention type, so sliding and full-attention layers borrow from different sources. For E4B (35 layers, 15 shared) that means: layers 20-23,25-28,30-33 (sliding) -> reuse layer 18's K,V layers 24,29,34 (full) -> reuse layer 19's K,V mobius computed k_proj/v_proj unconditionally on every layer. The checkpoint does ship k_proj/v_proj/k_norm for all 35 layers, so nothing failed to load — the export silently computed KV that HF never uses and emitted a 35-entry KV cache where HF effectively has 20. Mirrors the existing Gemma4TextAttention implementation: - ``Gemma3nConfig.num_kv_shared_layers``, populated from the HF config. - ``Gemma3nAttention`` computes ``is_kv_shared_layer`` / ``kv_shared_layer_index`` / ``provides_shared_kv`` and, for shared layers, drops the k_proj/v_proj/k_norm submodules the parent built so they emit no initializers. Source layers publish their post-attention K,V into a ``shared_kv_states`` dict that shared layers read from, converting BNSH back to 3D. The opset-24 ``Attention`` op's present outputs have no ORT shape inference, so the known BNSH shape is pinned to keep the downstream o_proj MatMul inferable. - ``Gemma3nTextModel`` emits present entries only for cache-owning layers and expands the shorter ``past_key_values`` list back to one slot per layer; ``kv_cache_layer_count()`` reports the real count, and ``CausalLMTask`` consults that hook when sizing the cache inputs (mirroring the existing ``static_kv_cache_specs`` precedent). - ``preprocess_weights`` drops the shared layers' K/V tensors, which HF does not construct either. Sharing was previously untestable: synthetic_parity_test.py pinned ``num_kv_shared_layers: 0`` for both gemma3n entries. The mixed- layer_types entries keep an explicit 0 (with TINY_LAYERS=2 there is no same-type source layer to borrow from), and the all-full-attention gemma3n_text entry now sets 1, so HF-vs-mobius parity exercises the shared path. TestBuildGemma3nKvSharing covers the cache I/O, absent initializers, per-type source mapping, state_dict pruning, and the every-layer-shared rejection. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ORT-GenAI uses ``decoder.num_hidden_layers`` to decide how many
``past_key_values.%d.{key,value}`` bindings to create, so it must match
the exported graph rather than the architecture. ``from_config`` read it
straight off the model config, which is wrong for KV-layer-sharing
models: their trailing layers borrow K,V from an earlier layer and own
no cache entry.
This is a pre-existing bug affecting the shipping gemma4 export, not
just gemma3n — a 4-layer / 2-shared gemma4 build emits 2 KV cache pairs
but wrote ``num_hidden_layers: 4``. Cached google/gemma-4-E2B-it has 35
layers with 20 shared.
``_write_genai_config`` already introspects the decoder graph for its
input names and GQA support, so count the ``past_key_values.{i}.key``
inputs there too and pass the result through. Static-cache exports use
``key_cache.{i}`` instead and fall back to the config value.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Gemma 3n sparsifies the MLP gate branch on its early layers. HF
(Gemma3nTextMLP._gaussian_topk) treats each row of gate_proj as a
Gaussian sample, derives the value below which `sparsity` of the mass
falls, and clamps everything under it to zero *before* the activation:
cutoff = mean(x) + std(x, unbiased=False) * Phi^-1(sparsity)
x = relu(x - cutoff)
E4B ships activation_sparsity_pattern = 0.95 for layers 0-9 and 0.0 for
layers 10-34, so nearly a third of the decoder was running a materially
different MLP than the checkpoint was trained with. mobius used the
plain components.MLP everywhere and never read the field —
`grep activation_sparsity src/mobius/` returned nothing.
Phi^-1 depends only on the config, so it folds into a Python float at
build time via statistics.NormalDist().inv_cdf() — no ONNX erfinv
needed. It matches torch.distributions.Normal(0, 1).icdf() to ~4e-7,
well inside float32.
- Gemma3nConfig.activation_sparsity_pattern, populated from the HF
config (verified: 35 floats for E4B, non-zero on layers 0-9).
- Gemma3nMLP overrides MLP.forward to insert the cutoff between
gate_proj and act_fn when sparsity is non-zero, and rejects a pattern
that is out of range or too short to cover every layer. Zero-sparsity
layers keep the plain path and emit no extra ops.
The tiny gemma3n configs now set [0.95, 0.0], covering both branches of
the fork, so the existing HF-vs-mobius parity harness exercises it:
without the cutoff, cosine similarity against HF drops from 0.996 to
0.78 (max_abs_diff 0.46). gemma3n_test.py additionally compares the
ONNX subgraph against a verbatim port of HF's _gaussian_topk under ORT,
and checks 0.95 sparsity keeps ~5% of a Gaussian row.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Groundwork for the Gemma 3n multimodal export: extract the MobileNet-V5 vision tower and USM Conformer audio tower configs so the component builders have something to read. No graph changes yet — the registry still maps "gemma3n" to the text-only decoder. Verified against google/gemma-3n-E4B-it: every field below matches the checkpoint, resolved identically whether from_transformers() is handed the outer config or the text sub-config plus parent. - VisionConfig gains vocab_offset/vocab_size (the image soft-token id range [262144, 262272) the multimodal embedder rebases into its own 128-entry table), architecture (the timm spec name — MobileNet-V5 has no per-layer config, so the name is the only handle on the block layout), do_pooling, and rms_norm_eps. - Gemma3nAudioConfig mirrors HF's field names verbatim (conf_* for the Conformer stack, sscp_* for the strided conv subsampler), so the hook is a copy rather than a remapping. - Gemma3nMultiModalConfig extends Gemma3nConfig (the decoder is unchanged) with the fixed soft-token counts: 256 for a 768x768 image (16x16 grid), 188 for audio. It also lifts audio_token_id to the top level, which no generic extractor populates for gemma3n. - The vision hook pins image_size=768: HF ships no image_size, MobileNet-V5 has no dynamic-resolution path, and the checkpoint's processor resizes to a fixed 768x768. Pinning it here keeps the graph's pixel_values shape and the generated preprocessing pipeline in sync. Both hooks are bare (unfiltered) because build() may dispatch on "gemma3n_text" while the gemma3n parent holds the sub-configs. The predicate therefore reads the parent model_type off parent_config, not off the `parent_config or config` composite — the composite form made the hooks fire for *any* dispatched model_type whenever the config itself was a gemma3n one, which the existing cross-contamination guards in _extractors_test.py caught. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The Gemma 3n vision tower is MobileNet-V5-300m, not SigLIP: HF maps gemma3n_vision to TimmWrapperModel, so there is no HF-side module to port and timm is deliberately not a mobius dependency. This adds a native onnxscript implementation driven by an explicit block-spec table transcribed from the mobilenetv5_300m_enc checkpoint (84 blocks across 4 stages: EdgeResidual, UIB, and multi-query attention 2D), plus the Multi-Scale Fusion Adapter that fuses the last two stages (640 + 1280 -> 1920) into 256 soft tokens. Two things the checkpoint forces, both easy to get wrong: * The 116 tensors timm names `bn.weight` are RmsNorm2d, not BatchNorm — there is no bn.bias, running_mean, running_var or num_batches_tracked anywhere in the checkpoint. BatchNorm2d would demand four initializers that do not exist, so RmsNorm2d (channel-axis RMS, scale-only) is added to _conv.py. * timm builds this tower with TensorFlow-style SAME padding, which is asymmetric when the total pad is odd (extra pixel on the end). Symmetric k//2 padding shifts every strided feature map by a pixel. Conv2d/Conv2dNoBias therefore now accept an ONNX-order [top, left, bottom, right] 4-tuple for `padding` alongside the existing int. Folding SAME padding into the Conv `pads` attribute (rather than emitting Pad nodes) keeps static shape inference intact through the whole tower, and keeps every conv going through Module.__call__ — which is what qualifies parameter names and registers initializers. Validated against timm with identical weights: max abs diff 8.1e-06 at 256x256 and 5.2e-06 at 768x768, cosine 1.000000 at both. Parameter names and shapes are an exact 548/548 match against google/gemma-3n-E4B-it. The 47 new tests cover the SAME-padding amounts, RmsNorm2d's reduction axis and scale broadcast, the block-spec table, the weight-name contract, per-block ONNX execution, and the resolution flow. Whole-tower assertions stay at the graph level: 300M parameters is past protobuf's 2 GB serialization ceiling once weights are attached. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Gemma3nMultimodalEmbedder bridges both modality towers into the text embedding space. Gemma 3n uses one per modality (model.embed_vision.* and model.embed_audio.*, 4 tensors each), and HF drives each through two distinct paths: the soft path embeds continuous encoder features (256 vision soft tokens, 188 audio), while the hard path embeds the reserved placeholder token ids the processor splices into the prompt, via a 128-row table offset by vocab_offset. Both paths are dispatched from a single forward, matching HF's signature — reaching into a submodule's parameters from outside Module.__call__ would leave them unqualified and unregistered, so the branch is resolved at graph-build time instead. One instance serves both paths in one graph, which the audio path needs. embedding_post_projection_norm is scale-free: HF builds it with with_scale=False and the checkpoint ships no tensor for it. gemma4 already had a private class for exactly this, so it is promoted to components as ScaleFreeRMSNorm and gemma4's four call sites now use the shared one rather than duplicating it. Verified against HuggingFace's Gemma3nMultimodalEmbedder: both paths match to 4.8e-07. Weight names are an exact match for the checkpoint, so preprocess_weights needs no renaming hook. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Phase 4 of gemma-3n-E4B-it multimodal export: the 269 tensors under model.audio_tower. — a two-stage strided conv subsampler (SSCP) plus 12 conformer blocks. There is no output projection; that lives in model.embed_audio (Gemma3nMultimodalEmbedder, added in the previous commit). The conformer feed-forward and light-conv1d sub-blocks are Gemma 4's, identical in both weight names and arithmetic, so Gemma4FeedForward and Gemma4LightConv1d gain a linear_cls hook and are reused directly. Gemma 3n's checkpoint ships no activation-clipping bounds, so the ClippableLinear default would demand four initializers per projection that do not exist. Three things genuinely diverge and are implemented here: * _CumulativeGroupNorm — group statistics accumulated over the time axis via CumSum, where Gemma 4 uses a plain per-frame LayerNorm. HF's masking support in this layer is vestigial (it builds an all-ones mask unconditionally), so the element count is exactly (t+1)*F*C and HF's zero-count guard and trailing mask multiply are both no-ops. * Reverse-causal time padding (0, kernel_h-1) plus fixed frequency padding (1, 1) in the SSCP convolutions, folded into the ONNX Conv pads attribute so shape inference still propagates. * Attention: an explicit relative-position projection, no key rescaling, and a 13-key window where Gemma4Attention admits 12 — so that class cannot be subclassed. The logit cap lands after the rel-pos bias and before the mask, which also rules out the ONNX Attention op's native softcap. HF's chunked attention is flattened to full T×T attention, which is exactly equivalent offline: working the block mask through, query i attends exactly the keys j with -R <= i-j <= L. Mask polarity follows the mobius convention (True = valid), the negation of HF's audio_mel_mask. Every learned tensor keeps its HF name; the only extra initializers are the derived relative_position_embedding.sin_emb constants. Tests diff the whole encoder against HF over five sequence lengths, three context configurations (including a lookahead window and a zero-length history), a batch with differing valid lengths, an interior mask hole, and a gradient clamp small enough to actually bind — plus the 269-tensor name/shape contract, the 1024-wide projection the frequency arithmetic has to produce, initializer qualification, and mask polarity. 18 of 18 seeded mutations of the component are caught. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Splits Gemma 3n into the four graphs the ORT-GenAI multimodal pipeline needs, following the gemma3/gemma4 pattern: - _Gemma3nDecoderModel takes inputs_embeds + per_layer_inputs, keeping the 4.7 GB embed_tokens_per_layer table out of the decoder while preserving the token-embedding term HF's project_per_layer_inputs adds. - _Gemma3nVisionEncoderModel runs MobileNet-V5 -> reshape/transpose -> scale -> embed_vision soft path (mirrors get_image_features). - _Gemma3nAudioEncoderModel runs the USM conformer -> embed_audio soft path. - _Gemma3nEmbeddingModel does the scaled token lookup, the per-layer tables, and image/audio fusion, reusing gemma3's zero-pad Gather guard so a decode step's empty [0, hidden] feature tensor cannot gather OOB. - Gemma3nMultiModalModel.preprocess_weights routes the HF checkpoint's five prefixes to the right components, synthesizes lm_head from the tied token embedding (E4B ships no lm_head key), and duplicates the embedder weights into both the tower and the embedding graph. Also threads final_logit_softcapping (30.0 in every published Gemma 3n config) through Gemma3nConfig and applies it in both the text-only and multimodal LM heads, and makes Gemma3nTextModel.forward accept a precomputed per_layer_inputs. _compute_per_layer_inputs now raises when input_ids is absent instead of silently dropping the token-embedding term and its 1/sqrt(2) scale. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Wires the four sub-models into a package via Gemma3nTask, closely following Gemma4Task (which already solves the same 4-way split with optional audio): - vision input pixel_values [batch, 3, image_size, image_size] — fixed NCHW, since MobileNet-V5 has no dynamic-resolution path - audio inputs input_features [B, T, feat] + input_features_mask [B, T] - declare_component_presence / declare_optional_input so a checkpoint without an audio tower still produces a legal 3-model package The "gemma3n" registry key now maps to Gemma3nMultiModalModel instead of Gemma3nCausalLMModel. That key previously produced a silently vision- and audio-blind text decoder: Gemma3nCausalLMModel.preprocess_weights deletes every vision_tower./audio_tower. key, so exporting google/gemma-3n-E4B-it yielded a package that could not accept images or audio at all. "gemma3n_text" keeps the text-only decoder, and _TEXT_ONLY_MODEL_TYPE maps gemma3n -> gemma3n_text so text_only=True still works. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Extends the gemma3 fixed-resize branch to cover gemma3n: MobileNet-V5 takes the same fixed NCHW pixel_values contract as gemma3's SigLIP, so the generic VLM branch's smart_resize (variable HxW, no Permute3D) would produce a tensor the tower cannot accept. Two behavioural differences from gemma3, both verified against the E4B checkpoint's SiglipImageProcessorFast: - image_size defaults to 768 (gemma3 uses 896). - Normalize is *omitted*, giving a 5-step pipeline. The processor sets do_normalize=False, so the tower is trained on [0, 1] pixels; emitting Normalize with mean/std 0.5 would map them to [-1, 1] and silently degrade every caption. An HF processor that reports do_normalize=True still gets the 6-step pipeline. Also maps the unwrapped model_type "gemma3n_text" to ORT model type "gemma3n" for VLM packages. Deliberately not aliased to "gemma3": the package threads per_layer_inputs (and optional audio) that gemma3's ORT pipeline does not bind. onnxruntime-genai 0.14.1 has no gemma3n pipeline yet, so this currently trips the existing "could not determine ORT-GenAI model type" warning rather than mis-wiring the graph. Tests cover the 5-step pipeline, the HF do_normalize override, the gemma3n_text -> gemma3n mapping, and a regression guard that gemma3 still gets its Normalize. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
- Move the "gemma3n" test config from ALL_CAUSAL_LM_CONFIGS to VL_CONFIGS so the parametrized vision-language suites pick it up, and add a vl_overrides() helper for direct construction. The tiny vision config respects MobileNet-V5's image_size rules (size % 32 == 0 and (size // 16) % 16 == 0, so 256 is the smallest legal value). - Add test_gemma3n_multimodal_graph and the ..._without_audio variant to build_graph_test.py, asserting the 4- and 3-model package shapes and the pixel_values / input_features / per_layer_inputs graph I/O. - Drop the three now-unreachable "gemma3n" entries from synthetic_parity_test.py (_ATOL_OVERRIDES, _HF_EXTRA_CONFIG, _HF_MODEL_TYPE_OVERRIDES). None of those dicts has a staleness guard, so they were dead rather than failing. The gemma3n_text entries still cover text-only parity. - Relabel testdata/cases/causal-lm/gemma3n.yaml as gemma3n_text: that case asserts a text-generation package, which "gemma3n" no longer builds. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The two existing gemma3n_text parity entries each disable half of the feature interaction: one has mixed layer_types with num_kv_shared_layers=0, the other has sharing on but a single layer type. Neither reaches the configuration the real E4B checkpoint uses (35 mixed layers, num_kv_shared_layers=15), where HF picks a *different* KV source layer per attention type -- a sliding shared layer borrows from the last pre-cutoff sliding layer, a full-attention one from the last pre-cutoff full layer. Add a 4-layer entry with layer_types [sliding, full, sliding, full] and num_kv_shared_layers=2, so prev_layers is [sliding, full] and layer 2 must borrow from layer 0 while layer 3 borrows from layer 1. Also keeps activation sparsity on for a subset of layers. Verified load-bearing by mutation: replacing the per-type lookup in Gemma3nAttention with a plain `len(prev_layers) - 1` (always borrow from the last pre-cutoff layer, ignoring type) fails this entry at cosine=0.881 with argmax_match=False, while both pre-existing entries still pass. Getting this wrong yields a loadable graph with plausible-looking attention, so the shape-level tests do not catch it either. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The gemma3n audio tower was exported but nothing referenced it: no feature extraction config was written, and the speech section of genai_config.json carried neither a config_filename nor an input mapping the runtime accepts. Three pieces: - _write_audio_processor_config gains a gemma3n branch. It reuses gemma4's Gemma4LogMel op but not its attributes -- Gemma3nAudioFeatureExtractor is a different filterbank (32 ms frames, 125-7600 Hz, 0.97 preemphasis, FFT overdrive on, 1e-5 mel floor). The upstream frame/hop are in samples and are converted to the milliseconds the op takes. - The speech section now names that file. ORT-GenAI rejects a speech section that sets filename without config_filename outright, so leaving it off turns a missing-file error into a load-time throw. - The speech input mapping is hardcoded rather than introspected. Unlike the decoder and vision sections, model.speech.inputs keys are a closed set the runtime defines (audio_embeds / attention_mask / audio_sizes / audio_projection_mode); an identity map from graph input names is rejected with 'model:speech:inputs: Unknown value "input_features"'. Also point the vision section at processor_config.json, the file gemma3n actually writes -- it shares gemma3's fixed-resize branch, not gemma4's, so with_vision's image_processor.json default named a file that never exists. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Three defects in the processor_config.json that every mobius VLM export writes. All three are silent: the pipeline is well-formed, the model loads, and the captions are merely wrong. 1. ConvertRGB. ort-extensions' convert_to_rgb unconditionally swaps R and B -- it exists to fix up a BGR decode. DecodeImage(color_space="RGB") already emits RGB, so chaining the two handed the encoder BGR. Drop the step from every pipeline. Measured against HF's own processor on a 1920x1242 JPEG, this alone put the pixel tensor 22% (relative L2) away from the reference. 2. Resize interpolation. ort-extensions defaults to CUBIC; HF image processors overwhelmingly ship resample=2, PIL BILINEAR. The Resize step now carries an explicit interpolation derived from image_processor.resample, mapped through the PIL constants. PIL BOX/HAMMING have no ort-extensions counterpart, so those log a warning and fall back rather than crash. 3. size handling. transformers >= 5 returns a SizeDict from image_processor.size, and SizeDict is not a dict subclass -- so the isinstance(size, dict) guard silently discarded HF's real size and left whatever default the caller hardcoded. Normalise through _size_mapping() instead. With all three fixed the pixel tensor lands 0.01% from HF's reference, down from 22%. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
KV-shared layers (those at or after first_kv_shared_layer_idx) borrow the
already-RoPE'd K,V of the last non-shared layer of the same type. Because
those tensors already span past + current, the layer passes them whole and
supplies no past_key/past_value -- so at decode q_len=1 while kv_len=total.
The opset-24 Attention spec aligns its built-in causal mask UPPER-LEFT
("the attention masking has the form of the upper left causal bias due to the
alignment"), which pins that lone query to key 0 and collapses those 15 layers
onto the BOS token. ORT <= 1.27 aligned it bottom-right -- non-conforming, and
it hid this entirely; 1.28 conforms and the decode output degenerates.
attention_bias from create_attention_bias already bakes causal + sliding +
padding keyed on absolute cumsum positions, so it carries causality on its own
at any q_len/kv_len ratio. Emit is_causal=0 on shared layers and let the
explicit mask do the work, and raise if a shared layer is ever built without
one. Non-shared layers keep is_causal=1: they pass past_key/past_value, which
takes the op's cache path (bottom-right aligned in every ORT version), and at
prefill q_len == kv_len makes the two alignments identical.
The regression test is a decode step against a full re-prefill. This is
prefill-vs-decode, so the existing synthetic parity sweep -- prefill only --
could not have caught it.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ArchitectureConfig.from_transformers gates attn_qk_norm on a model-type allowlist that gemma3n was never added to, so Gemma3nAttention built neither q_norm nor k_norm and the weights were dropped on the floor. They are load-bearing rather than cosmetic. HF hardcodes scaling = 1.0 instead of head_dim**-0.5 precisely because those norms leave Q and K at unit RMS; without them the logits are scaled by |q||k| and softmax degenerates toward a one-hot argmax. Adding them takes the text-path synthetic parity max diff from ~0.094 to ~0.026, so tighten the tolerance to 0.05 to keep the fix from silently regressing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The causal-LM case pointed at google/gemma-3n-E2B-pt and carried skip_reason "Model not yet published on HuggingFace." No -pt checkpoint was ever published -- only the -it releases, which are multimodal. Follow the gemma-4-e2b precedent and drive the text-only path from an -it checkpoint via model_type "gemma3n_text", and drop the skip. Add the E4B multimodal case alongside it, so the 4-model split (decoder + vision + audio + embedding) is covered too. _TEST_MODEL_IDS gets the same treatment for both keys. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
There was a problem hiding this comment.
lintrunner found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.
justinchuby
reviewed
Aug 10, 2026
Contributor
There was a problem hiding this comment.
Pull request overview
Enables full multimodal export for Gemma 3n (google/gemma-3n-E4B-it) by splitting the package into dedicated decoder / vision / audio / embedding sub-models, and updates the surrounding config + export plumbing so downstream runtimes (notably ORT-GenAI) can wire the graphs correctly (including KV-layer-sharing cache contracts and processor config files).
Changes:
- Add a Gemma 3n multimodal task and model/package contract (3- or 4-model split depending on audio).
- Implement Gemma 3n’s modality-specific components (MobileNet-V5 vision tower, multimodal embedder, USM audio tower coverage) plus shared ScaleFreeRMSNorm.
- Update testing + ORT-GenAI export/config generation to match new I/O contracts (processor pipelines, KV cache layer counts, model_type routing).
Reviewed changes
Copilot reviewed 37 out of 37 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/synthetic_parity_test.py | Removes gemma3n multimodal entry from causal-LM parity suite; tightens gemma3n_text tolerance and unpins KV-sharing config escape hatch. |
| tests/build_graph_test.py | Adds multimodal Gemma3n graph-build tests and adjusts KV-cache assertions for KV-sharing models. |
| tests/_test_configs.py | Adds Gemma3n multimodal tiny config in VL_CONFIGS, expands gemma3n_text coverage for KV sharing + sparsity, and introduces vl_overrides(). |
| testdata/cases/vision-language/gemma-3n-e4b-it.yaml | New VLM case for Gemma 3n E4B-it multimodal export. |
| testdata/cases/causal-lm/gemma3n.yaml | Switches causal-lm case to gemma3n_text driven from an -it checkpoint. |
| src/mobius/tasks/_gemma3n.py | New Gemma3nTask implementing decoder/vision/audio/embedding graph split and contracts. |
| src/mobius/tasks/_causal_lm.py | Uses kv_cache_layer_count() hook to size KV-cache inputs for KV-sharing decoders. |
| src/mobius/tasks/init.py | Exposes Gemma3nTask and registers task key gemma3n. |
| src/mobius/models/gemma4.py | Replaces private scale-free RMSNorm implementation with shared ScaleFreeRMSNorm. |
| src/mobius/models/gemma4_test.py | Updates scale-free RMSNorm unit test to import from components. |
| src/mobius/models/gemma3n_test.py | New tests for Gemma3n activation sparsity, KV-sharing masking behavior, and multimodal weight renaming/coverage. |
| src/mobius/models/init.py | Exports Gemma3nMultiModalModel. |
| src/mobius/integrations/ort_genai/genai_config.py | Adds num_kv_cache_layers override for ORT-GenAI KV binding count; clarifies audio config requirements. |
| src/mobius/integrations/ort_genai/genai_config_test.py | Adds coverage for num_kv_cache_layers override behavior. |
| src/mobius/integrations/ort_genai/auto_export.py | Fixes/extends processor config writing (no ConvertRGB, HF resample→Resize interpolation, SizeDict handling), adds Gemma3n processor branches, and counts KV-cache layers from graph inputs. |
| src/mobius/integrations/ort_genai/auto_export_test.py | Updates vision pipeline expectations, adds Gemma3n-specific processor/config coverage, and tests KV-cache layer counting helper. |
| src/mobius/components/_rms_norm.py | Introduces shared ScaleFreeRMSNorm. |
| src/mobius/components/_mobilenetv5.py | New MobileNet-V5-300m encoder implementation for Gemma3n vision tower. |
| src/mobius/components/_mobilenetv5_test.py | New unit tests validating MobileNet-V5 spec table, naming/shape contracts, SAME padding, and ONNX execution. |
| src/mobius/components/_gemma4_audio.py | Parameterizes certain Gemma4 audio blocks to allow Gemma3n to use non-clippable linear projections. |
| src/mobius/components/_gemma3n_embedder.py | New Gemma3n multimodal embedder component (soft/hard paths). |
| src/mobius/components/_gemma3n_embedder_test.py | New unit tests for embedder checkpoint contract and HF numerical parity. |
| src/mobius/components/_gemma3n_audio_test.py | New unit tests for Gemma3n USM audio encoder contract, masking, and HF numerical parity. |
| src/mobius/components/_conv.py | Extends Conv2d padding handling to accept ONNX-order 4-tuples for asymmetric padding; adds RmsNorm2d. |
| src/mobius/components/init.py | Exports Gemma3n components, MobileNetV5Encoder, RmsNorm2d, and ScaleFreeRMSNorm. |
| src/mobius/_registry.py | Splits gemma3n vs gemma3n_text registrations; adds multimodal mapping and text-only aliasing rules; updates default model IDs. |
| src/mobius/_configs/per_model/_gemma3n_vision.py | New gemma3n vision extractor hook for MobileNet-V5 tower fields + fixed image size. |
| src/mobius/_configs/per_model/_gemma3n_audio.py | New gemma3n audio extractor hook producing Gemma3nAudioConfig. |
| src/mobius/_configs/per_model/init.py | Registers new gemma3n per-model hooks. |
| src/mobius/_configs/_sub_configs.py | Extends VisionConfig for gemma3n fields and adds Gemma3nAudioConfig dataclass. |
| src/mobius/_configs/_extractors_test.py | Adds tests ensuring gemma3n vision/audio hooks fire correctly and don’t leak to unrelated model types. |
| src/mobius/_configs/_base.py | Fixes intermediate_size coercion, adds gemma3n model types to special-case handling, and adds gemma3n KV-sharing + sparsity fields/config class. |
| src/mobius/_configs/_base_test.py | New tests for ArchitectureConfig.from_transformers intermediate_size list→scalar coercion. |
| src/mobius/_configs/init.py | Exposes Gemma3nMultiModalConfig and Gemma3nAudioConfig. |
| src/mobius/init.py | Re-exports Gemma3nMultiModalConfig. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
justinchuby
reviewed
Aug 10, 2026
Nothing behavioural. `ruff check` flagged 20 violations across the new gemma3n files and `ruff format` disagreed on 11 of them: - RUF012: `Gemma3nTask.model_roles` needs `ClassVar[dict[str, str]]`, the same annotation every other `ModelTask` already carries. - RUF002/RUF003: the `x` in dimension comments (`768x768`, `[B, T, 128]`) was a U+00D7 multiplication sign. - D205: three docstrings in `gemma3n_test.py` ran the summary into the body without a blank line. - E501 / I001 in `_mobilenetv5_test.py`, plus `timm's` -> `Timm's` for the docstring capitalisation rule. - PT011-adjacent float asserts in `_extractors_test.py` and `gemma3n_test.py` now go through `pytest.approx`, matching `_attention_test.py` and `_moe_test.py`. `ruff check --config=pyproject.toml` and `ruff format --check` are both clean on the 35 files this branch touches, and the branch test set is unchanged at 1632 passed / 44 skipped. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Both cases have existed since this branch added them, but neither had a
golden, and a missing golden makes `tests/e2e_golden_test.py` *skip* --
so the two jobs that would actually catch a wrong export have been silently
doing nothing.
Generated with `scripts/generate_golden.py --force`, float32 on CPU, per
the dtype each case declares:
- `causal-lm/gemma3n` from `google/gemma-3n-E2B-it` via `model_type`
`gemma3n_text` (30 layers, `num_kv_shared_layers=10`), 20 new tokens.
- `vision-language/gemma-3n-e4b-it` from `google/gemma-3n-E4B-it`,
`pipeline-cat-chonk.jpeg` at 768x768, 30 new tokens. HF captions it
"A medium shot captures a Pallas's cat, also known as a manul, walking
across a snowy surface."
All four now-unskipped tests pass locally against the f32 export:
L4 prefill argmax text-generation/gemma3n PASSED 201s
L4 prefill argmax image-text-to-text/gemma-3n-e4b-it PASSED
L5 generation text-generation/gemma3n PASSED 370s
L5 generation image-text-to-text/gemma-3n-e4b-it PASSED 738s
L5 reproduces every golden token, which is the first end-to-end evidence
that KV layer sharing and activation sparsity are right on a real
checkpoint -- the synthetic parity harness only ever saw tiny configs.
Note for CI: the L4/L5 workflows pass `--timeout=300`. These ran 201-738s
on CPU here; the A10 pool should be faster, but the multimodal L5 in
particular may need a longer per-test timeout.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
# Conflicts: # src/mobius/_registry.py
Merging main brought in onnxruntime#475's test for the packed-Qwen image pipeline, which asserts a ConvertRGB step this branch had already removed from every VLM pipeline. Git merged the file cleanly because the conflict is semantic, so all six unit-test combos went red on one assertion. The removal stands: ort-extensions' convert_to_rgb swaps R and B unconditionally -- it exists to fix up a BGR decode -- so chaining it after DecodeImage(color_space="RGB") hands the encoder BGR. With it gone the pixel tensor lands 0.01% from HF's reference instead of 22%. So this updates the newer expectation rather than restoring the transform, and says why in the docstring so it does not come back. The pipeline is now DecodeImage, Resize, Rescale, Normalize, PatchImage, which moves the qwen2_5_vl flag onto index 3 and PatchImage's attrs to index 4. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Review feedback on onnxruntime#474: _resolve_pads accepted any non-int padding verbatim, so a PyTorch-style (h, w) pair became a 2-element ONNX pads attribute. That is not caught here -- it surfaces later, either as a shape-inference error far from the caller or, worse, as a silent numerical mismatch once the values land on the wrong edges. The sequence form is deliberately ONNX order [top, left, bottom, right], so a 2-tuple cannot be padded out to a 4-tuple without guessing which edges the caller meant. Raise instead, and say what to pass. Co-Authored-By: Claude Opus 5 <noreply@anthropic.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.
What this does
Makes
google/gemma-3n-E4B-itexport as a real 4-model multimodal package(
decoder+vision_encoder+audio_encoder+embedding) instead of atext-only decoder.
Before this, both
gemma3nandgemma3n_textmapped toGemma3nCausalLMModel, soGemma3nCausalLMModel.preprocess_weightsdeletedevery
vision_tower.,audio_tower., andmulti_modal_projector.key. Theexport succeeded and produced a package that could not accept images or audio at
all.
gemma3nnow maps to a newGemma3nMultiModalModelwithtask="gemma3n",matching how
gemma3/gemma3_textare already split.Two pre-existing text-path bugs fixed first
These affect the already-shipping text export and are independent of the
modality work, so they land as their own commits at the base of the branch. Both
silently produced wrong numbers rather than crashing.
num_kv_shared_layers=15, so layers >= 20 reuse the KV of the last non-sharedlayer of the same type (sliding -> layer 18, full -> layer 19). Attention
always recomputed
k_proj/v_proj. The checkpoint does shipk_projfor all35 layers, so nothing failed to load -- the graph just computed KV that HF
never uses, and the exported cache had 35 entries where HF effectively has 20.
_gaussian_topk) was missing (08a14a0). For E4Blayers 0-9
activation_sparsity_patternis 0.95. HF computescutoff = mean + std * icdf(0.95)over the last dim and appliesrelu(gate_proj - cutoff)before the activation.icdf(0.95)is acompile-time constant, so no ONNX erf/ppf is needed.
Both were previously invisible because
tests/synthetic_parity_test.pypinnednum_kv_shared_layers: 0forgemma3nand never set a sparsity pattern. Thoseescape hatches are removed here, so the existing HF-vs-mobius parity harness now
covers both.
Architecture notes
Three places where gemma3n does not match the obvious existing component, and
reusing it would have been wrong:
gemma3n_visiontoTimmWrapperModel, and timm is not a dependency. Newcomponents/_mobilenetv5.pyimplements the stem, EdgeResidual / UIB /UIB+MQA blocks, and the Multi-Scale Fusion Adapter, driven by a hard-coded
per-stage spec table -- there is no config source for the block layout.
bn.weighttensors are timmRmsNorm2d, not BatchNorm. Thecheckpoint ships zero
bn.bias/running_mean/running_var/num_batches_tracked. Using the existingcomponents.BatchNorm2dwould havedemanded four initializers that do not exist.
_gemma4_audio.pyexcept for the norm. gemma3n'sSSCP conv blocks use
Gemma3nAudioCumulativeGroupNorm-- group norm computedcumulatively over the time axis -- where gemma4 uses plain
LayerNormNoBias.That needs
CumSumover T. gemma3n also adds an explicitrelative_position_embedding.pos_projthat gemma4 lacks.components/_rms_norm.pygains a shared scale-free RMSNorm, promoted out ofmodels/gemma4.pyrather than duplicated.