diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py index 265f8b0f..5ee0af9b 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -2639,9 +2639,19 @@ class _Gemma4AudioEncoderModel(nn.Module): Inputs: - ``input_features [B, T, input_size]``: mel-spectrogram features + - ``input_features_mask [B, T]``: BOOL mask, ``True`` for valid frames Output: - - ``audio_features [B, T//4, text_hidden_size]``: projected audio tokens + - ``audio_features [num_valid_frames, text_hidden_size]``: projected audio + tokens with padding frames stripped so output rows align 1:1 with audio + placeholder tokens. This mirrors HF, which selects + ``audio_features[audio_mask]`` (a rank-2 gather) before scattering the + rows into the language-model embeddings. Emitting rank-2 here keeps the + producer rank aligned with the ``embedding`` sub-model's rank-2 + ``audio_features`` consumer. + - ``downsampled_mask [B, T//4]``: BOOL mask over the subsampled time axis, + returned for the ``audio_features_mask`` diagnostic output. ``None`` when + no ``input_features_mask`` was supplied. """ def __init__(self, config: Gemma4Config): @@ -2701,7 +2711,23 @@ def forward( rms = op.Sqrt(op.Add(mean_sq, eps)) audio_features = op.CastLike(op.Div(x_f32, rms), audio_features) # → projector → [B, T//4, text_hidden_size] - return self.projector(op, audio_features), downsampled_mask + audio_features = self.projector(op, audio_features) + if downsampled_mask is None: + # No validity mask (e.g. eager unit call): emit the rank-3 tensor + # unchanged. Export always supplies a mask, so the rank-2 path below + # is the one that reaches the ``embedding`` consumer. + return audio_features, None + # Select valid frames so the output rows align 1:1 with audio placeholder + # tokens, matching HF's ``audio_features[audio_mask]``. Flatten the batch + # and subsampled-time axes, then compress by the downsampled mask to get + # rank-2 ``[num_valid_frames, text_hidden_size]`` — the rank the + # ``embedding`` sub-model consumes. + flat_features = op.Reshape( + audio_features, op.Constant(value_ints=[-1, self.config.hidden_size]) + ) + flat_mask = op.Reshape(downsampled_mask, op.Constant(value_ints=[-1])) + selected = _dtype_safe_compress(op, flat_features, flat_mask, axis=0) + return selected, downsampled_mask # --------------------------------------------------------------------------- diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py index 66e24e19..c06c9241 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -284,6 +284,86 @@ def test_no_block_sequence_ids_when_causal(self): assert "input_ids" not in dec_inputs +def _tiny_gemma4_audio_config(**overrides) -> Gemma4Config: + """Minimal gemma4 config with a Conformer audio encoder (E2B-style).""" + from mobius._configs import Gemma4AudioConfig + + base = dict( + image_token_id=255, + audio=Gemma4AudioConfig( + input_size=32, + hidden_size=16, + num_layers=1, + subsampling_conv_channels=[8, 8], + output_proj_dims=16, + audio_token_id=254, + ), + ) + base.update(overrides) + return _tiny_gemma4_config(**base) + + +class TestGemma4AudioEncoderRank: + """Lock the audio dataflow-edge rank contract (regression for #570). + + The exported ``audio_encoder`` must emit ``audio_features`` at the SAME rank + the ``embedding`` sub-model consumes it (rank-2 + ``[num_valid_frames, text_hidden_size]``), mirroring HF's + ``audio_features[audio_mask]`` selection. A rank-3 producer against a + rank-2 consumer is rejected at onnx-genai pipeline admission. + """ + + def test_audio_features_producer_consumer_ranks_match(self): + from mobius.models.gemma4 import Gemma4Model + from mobius.tasks._gemma4 import Gemma4Task + + config = _tiny_gemma4_audio_config() + pkg = Gemma4Task().build(Gemma4Model(config), config) + + audio_out = {o.name: o for o in pkg["audio_encoder"].graph.outputs} + emb_in = {i.name: i for i in pkg["embedding"].graph.inputs} + + producer = audio_out["audio_features"] + consumer = emb_in["audio_features"] + assert producer.shape is not None + assert consumer.shape is not None + producer_rank = len(producer.shape) + consumer_rank = len(consumer.shape) + + # Non-vacuous: assert the concrete rank-2 contract on BOTH ends and that + # they agree, not merely that export ran. + assert producer_rank == 2, ( + f"audio_encoder audio_features must be rank-2 " + f"[num_valid_frames, hidden]; got rank {producer_rank} {producer.shape}" + ) + assert consumer_rank == 2, ( + f"embedding audio_features input must be rank-2; " + f"got rank {consumer_rank} {consumer.shape}" + ) + assert producer_rank == consumer_rank + # Trailing dim must be the text hidden size on both ends. + assert producer.shape[1] == config.hidden_size + assert consumer.shape[1] == config.hidden_size + + def test_audio_encoder_selects_valid_frames_with_compress(self): + """The rank reduction is a real mask-select (HF ``audio_features[audio_mask]``). + + Guards against a "fix" that reshapes B*T//4 into rank-2 without dropping + padding frames, which would misalign rows with placeholder tokens. + """ + from mobius.models.gemma4 import Gemma4Model + from mobius.tasks._gemma4 import Gemma4Task + + config = _tiny_gemma4_audio_config() + pkg = Gemma4Task().build(Gemma4Model(config), config) + + audio_graph = pkg["audio_encoder"].graph + assert any(n.op_type == "Compress" for n in audio_graph), ( + "audio_encoder must Compress by the downsampled validity mask so " + "output rows align 1:1 with audio placeholder tokens" + ) + + def _tiny_gemma4_unified_config(**overrides) -> Gemma4Config: """Minimal gemma4_unified config (dense decoder + encoder-free embedders).""" from mobius._configs import Gemma4AudioConfig, VisionConfig diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index a8233a79..de121f97 100644 --- a/src/mobius/tasks/_gemma4.py +++ b/src/mobius/tasks/_gemma4.py @@ -666,19 +666,9 @@ def _build_audio( ) if downsampled_mask is None: raise ValueError("Gemma4 audio encoder must return a downsampled validity mask") - flattened_features = op.Reshape( - audio_features, - op.Constant(value_ints=[-1, config.hidden_size]), - ) - flattened_mask = op.Reshape( - downsampled_mask, - op.Constant(value_ints=[-1]), - ) - flattened_features_f32 = op.Cast(flattened_features, to=ir.DataType.FLOAT) - audio_features = op.CastLike( - op.Compress(flattened_features_f32, flattened_mask, axis=0), - flattened_features, - ) + # The audio encoder module already strips padding frames and flattens the + # batch axis, so ``audio_features`` is rank-2 ``[num_valid, hidden_size]`` + # here — the rank the ``embedding`` sub-model consumes. builder.add_output(audio_features, "audio_features") builder.add_output(downsampled_mask, "audio_features_mask")