From db9a47030064955dc7f7c2c02b9320eb4021d732 Mon Sep 17 00:00:00 2001 From: justinchuby Date: Thu, 30 Jul 2026 16:04:24 +0000 Subject: [PATCH] Fix L4/L5 golden regressions across 7 broken models An L4/L5 golden-test regression sweep found seven models failing at build or run time. Root-cause and fix each: - diffllama: create_padding_mask now returns a 4D bool [B,1,S,Sk] mask, but DiffLlamaAttention still unsqueezed it again, producing a rank-5 bias that broke ORT Split. Add the (already-4D) pad bias directly. - gemma-4 audio encoder (e2b + e4b): (1) op.Range got a 1-D [1] limit from op.Shape; squeeze to a scalar at both Range sites. (2) the mask-downsample Slice uses an INT64_MAX "to-the-end" sentinel with step 2; onnx-shape- inference < 0.3.1 did not normalize the sentinel for strided slices and derived a 2^62 extent that overflowed the downstream Reshape at load. Bump the dependency to onnx-shape-inference>=0.3.1 (which normalizes a positive- sentinel end to ceil(dim/step)) and keep the natural [INT64_MAX] slice. - gemma-4 assistant drafter: the generic Attention path built the sliding- window bias from the full buffer-width attention_mask (kv_len + 1) while K is only the shared KV (kv_len), so ORT rejected an inconsistent total_sequence_length. Slice the mask to kv_len first. - qwen3-asr / qwen3-asr-en: thinker_config ships as a plain dict, so the text_config-unwrap branch in build() missed it and the decoder config kept hidden_size=0. Convert the dict via _dict_to_pretrained_config before unwrapping text_config. - nemotron-h: build() without trust_remote_code mis-detected layer_types as linear_attention (TypeError). Add trust_remote_code and a clear None-guard in linear_attention_dims. The build now succeeds with correct mamba2 layer types, but the HF reference needs trust_remote_code custom modeling plus mamba-ssm CUDA kernels that are unavailable in CI golden generation (same as jamba-v0_1 / granitemoehybrid-tiny), and the previously committed golden was generated without trust_remote_code (wrong native config). Documented skip_reason added until a valid golden can be produced. - sensevoice-small: audio-ctc L4/L5 is not wired into the e2e harness (custom LFR-fbank input_features + language_id, FunASR reference, no audio-ctc golden generator). Documented skip_reason added. Verified on GPU (CUDA EP): diffllama, qwen3-asr, qwen3-asr-en, gemma-4-e2b-it-audio, gemma-4-e4b-it-audio, gemma-4-e2b-it-assistant all pass L4 and L5; nemotron-h and sensevoice-small skip. Regression-checked gemma-4-e2b, llama-3_2-1b, qwen3-0_6b (unchanged, still pass). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby --- pyproject.toml | 2 +- src/mobius/_builder.py | 9 ++++++++- src/mobius/components/_gemma4_assistant.py | 12 +++++++++++- src/mobius/components/_gemma4_audio.py | 14 +++++++++++--- src/mobius/models/diffllama.py | 9 ++++++--- src/mobius/tasks/_cache_utils.py | 18 ++++++++++++++++++ .../cases/causal-lm/nemotron-h-nano-4b.yaml | 2 ++ testdata/cases/speech/sensevoice-small.yaml | 6 ++++++ 8 files changed, 63 insertions(+), 9 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c00f53b4..78963652 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ dependencies = [ "huggingface_hub", "numpy>=1.24.0", "onnx_ir>=0.1.0", - "onnx-shape-inference>=0.1.9", + "onnx-shape-inference>=0.3.1", "onnxscript>=0.7.1", "safetensors", "torch>=2.1.0", diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index 4638ecd4..17bec3ea 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -447,6 +447,7 @@ def build( from mobius._config_resolver import ( _config_from_hf, _default_task_for_model, + _dict_to_pretrained_config, _try_load_config_json, ) from mobius._diffusers_builder import build_diffusers_pipeline @@ -481,7 +482,13 @@ def build( hf_config = hf_config.talker_config elif hasattr(hf_config, "thinker_config"): thinker = hf_config.thinker_config - if hasattr(thinker, "text_config"): + # Some checkpoints (e.g. Qwen3-ASR) ship ``thinker_config`` as a plain + # dict rather than a nested ``PretrainedConfig``. Convert it so the + # decoder ``text_config`` (and its scalar fields such as hidden_size) + # is reachable via attribute access. + if isinstance(thinker, dict): + thinker = _dict_to_pretrained_config(thinker) + if getattr(thinker, "text_config", None) is not None: hf_config = thinker.text_config elif hasattr(hf_config, "decoder_config") and model_type == "qwen3_tts_tokenizer_12hz": # Codec tokenizer: use decoder_config as the primary config source diff --git a/src/mobius/components/_gemma4_assistant.py b/src/mobius/components/_gemma4_assistant.py index 2706ae80..c66cad91 100644 --- a/src/mobius/components/_gemma4_assistant.py +++ b/src/mobius/components/_gemma4_assistant.py @@ -205,10 +205,20 @@ def forward( if self.sliding_window and attention_mask is not None: from mobius._build_context import get_build_dtype + # The shared KV holds ``kv_len`` real positions, but ``attention_mask`` + # is the target's full buffer-width mask (kv_len + 1 for the slot of + # the token being predicted — GQA consumes that extra slot via + # ``total_seq_len`` while the generic path passes only the shared KV + # as K). Slice the mask to ``kv_len`` so the additive bias width + # matches the Attention op's total_sequence_length (= K's seq dim); + # otherwise ORT rejects the model with an inconsistent + # total_sequence_length error. + kv_len = op.Shape(shared_key, start=2, end=3) # [1] = shared KV seq len + mask_kv = op.Slice(attention_mask, [0], kv_len, [1], [1]) # [B, kv_len] attn_bias = create_attention_bias( op, input_ids=q, - attention_mask=attention_mask, + attention_mask=mask_kv, sliding_window=self.sliding_window, dtype=get_build_dtype(), ) diff --git a/src/mobius/components/_gemma4_audio.py b/src/mobius/components/_gemma4_audio.py index 33c54a45..6717df85 100644 --- a/src/mobius/components/_gemma4_audio.py +++ b/src/mobius/components/_gemma4_audio.py @@ -211,7 +211,10 @@ def _mask_and_downsample( # Broadcast mask [B, T] → [B, 1, T, 1] over C and F dims mask_4d = op.Unsqueeze(mask, [1, 3]) # [B, 1, T, 1] x = op.Mul(x, op.CastLike(mask_4d, x)) - # Downsample mask by conv stride (2): mask[:, ::2] + # Downsample mask by conv stride (2): mask[:, ::2]. The INT64_MAX + # ``ends`` sentinel means "to the end"; onnx-shape-inference (>=0.3.1) + # normalizes it to the axis length for strided slices, yielding a + # ceil(T/2) extent instead of a 2^62 sentinel. mask = op.Slice( mask, [0], # starts @@ -477,7 +480,10 @@ def _build_causal_window_mask(self, op: OpBuilder, seq_len: ir.Value) -> ir.Valu """Build [1, 1, T, T] causal sliding-window attention bias.""" zero = op.Constant(value_int=0) one = op.Constant(value_int=1) - positions = op.Range(zero, seq_len, one) # [T] int64 + # ``seq_len`` is a 1-D [1] tensor from op.Shape; Range requires a + # scalar limit, so squeeze the singleton axis. + seq_len_scalar = op.Squeeze(seq_len, [0]) + positions = op.Range(zero, seq_len_scalar, one) # [T] int64 q_pos = op.Unsqueeze(positions, [1]) # [T, 1] k_pos = op.Unsqueeze(positions, [0]) # [1, T] diff = op.Sub(q_pos, k_pos) # [T, T]: i - j @@ -570,7 +576,9 @@ def forward( # pos_embed is ordered [ctx_left-1, ..., 0], so distance d maps to index ctx_left-1-d zero_i = op.Constant(value_int=0) one_i = op.Constant(value_int=1) - positions = op.Range(zero_i, seq_len, one_i) # [T] + # Range requires a scalar limit; ``seq_len`` is a 1-D [1] shape tensor. + seq_len_scalar = op.Squeeze(seq_len, [0]) + positions = op.Range(zero_i, seq_len_scalar, one_i) # [T] q_pos_i = op.Unsqueeze(positions, [1]) # [T, 1] k_pos_i = op.Unsqueeze(positions, [0]) # [1, T] diff_i = op.Sub(q_pos_i, k_pos_i) # [T, T]: i-j diff --git a/src/mobius/models/diffllama.py b/src/mobius/models/diffllama.py index 72da50b0..b5f72f5f 100644 --- a/src/mobius/models/diffllama.py +++ b/src/mobius/models/diffllama.py @@ -194,12 +194,15 @@ def forward( # Unsqueeze to [1, 1, S, Sk] and add to scores [B, H, S, Sk] attn_scores = op.Add(attn_scores, op.Unsqueeze(causal_bias, [0, 1])) - # Padding mask (3D bool [B, S, Sk]) → additive bias [B, 1, S, Sk] + # Padding mask: ``create_padding_mask`` yields a 4D bool tensor + # [B, 1, S, Sk] (the explicit singleton head dim prevents the batch + # axis being misread as q_num_heads for batch > 1). Convert to an + # additive bias and add directly — it broadcasts across the head axis. if attention_bias is not None: pad_zero = op.CastLike(0.0, q_4d) pad_neg_inf = op.CastLike(float("-inf"), q_4d) - pad_bias = op.Where(attention_bias, pad_zero, pad_neg_inf) - attn_scores = op.Add(attn_scores, op.Unsqueeze(pad_bias, [1])) + pad_bias = op.Where(attention_bias, pad_zero, pad_neg_inf) # [B, 1, S, Sk] + attn_scores = op.Add(attn_scores, pad_bias) # Softmax and weighted sum with doubled V attn_weights = op.Softmax(attn_scores, axis=-1) # [B, H, S, Sk] diff --git a/src/mobius/tasks/_cache_utils.py b/src/mobius/tasks/_cache_utils.py index d8c156b5..c69787e4 100644 --- a/src/mobius/tasks/_cache_utils.py +++ b/src/mobius/tasks/_cache_utils.py @@ -49,6 +49,24 @@ def linear_attention_dims(config: BaseModelConfig) -> LinearAttentionDims: head_k_dim = config.linear_key_head_dim head_v_dim = config.linear_value_head_dim conv_kernel = config.linear_conv_kernel_dim + missing = [ + name + for name, value in ( + ("linear_num_key_heads", num_k_heads), + ("linear_num_value_heads", num_v_heads), + ("linear_key_head_dim", head_k_dim), + ("linear_value_head_dim", head_v_dim), + ("linear_conv_kernel_dim", conv_kernel), + ) + if value is None + ] + if missing: + raise ValueError( + "Linear-attention cache inputs requested but the following config " + f"fields are unset: {', '.join(missing)}. This usually means the " + "architecture's layer_types were mis-detected as 'linear_attention' " + "(e.g. a hybrid model built without trust_remote_code)." + ) key_dim = head_k_dim * num_k_heads value_dim = head_v_dim * num_v_heads conv_dim = key_dim * 2 + value_dim diff --git a/testdata/cases/causal-lm/nemotron-h-nano-4b.yaml b/testdata/cases/causal-lm/nemotron-h-nano-4b.yaml index 9e4c6cba..6ce598ea 100644 --- a/testdata/cases/causal-lm/nemotron-h-nano-4b.yaml +++ b/testdata/cases/causal-lm/nemotron-h-nano-4b.yaml @@ -3,6 +3,8 @@ model_type: "nemotron_h" revision: "main" task_type: "text-generation" dtype: "float32" +trust_remote_code: true +skip_reason: "Requires trust_remote_code custom modeling + mamba-ssm CUDA kernels for the HF reference, which are unavailable in CI golden generation (same limitation as jamba-v0_1 and granitemoehybrid-tiny). The mobius build now succeeds (correct mamba2 layer_types via trust_remote_code), but a valid golden cannot be regenerated here; the previously committed golden was produced without trust_remote_code and reflects the wrong native `linear_attention` config." inputs: prompts: diff --git a/testdata/cases/speech/sensevoice-small.yaml b/testdata/cases/speech/sensevoice-small.yaml index e70b5e42..64a0c7c6 100644 --- a/testdata/cases/speech/sensevoice-small.yaml +++ b/testdata/cases/speech/sensevoice-small.yaml @@ -10,6 +10,12 @@ inputs: level: "L4+L5" +skip_reason: > + audio-ctc L4/L5 is not wired into the e2e harness: the encoder needs + custom LFR-fbank ``input_features`` (560-dim) plus a ``language_id``, + and the reference is FunASR (no transformers class), so + generate_golden.py has no audio-ctc generator. Tracked as a follow-up. + notes: > SenseVoiceSmall CTC encoder-only ASR. Single model: SANM encoder + CTC head. Audio input is LFR fbank features (batch, seq_len, 560) with language ID.