Fix GPU integration/L4/L5 test failures: TF32, transformers 5.x vision renames, gemma3 multimodal (#350)

## Summary

Ran the full integration + L4 (golden) + L5 (generation) suite on GPU
(`MOBIUS_TEST_DEVICE=cuda`, 8× H200) to find regressions and
long-standing
bugs. Triaged 88 failures by root cause. **PR #338 (gemma4 bidirectional
overlay + GQA-cap removal) introduced zero regressions** — every
spot-checked
failure reproduces identically on the parent commit `f972184`.

This PR lands the verified, high-impact fixes. Remaining buckets (deep
numeric/decode bugs, ref-API drift, env/external) are tracked separately
and
will follow.

## Fixes

### 1. Disable TF32 on GPU (`tests/conftest.py`)
On Ampere+/Hopper the ORT CUDA EP uses TF32 for fp32 matmuls by default,
while
the PyTorch reference computes in true fp32. The ~1e-2 logit discrepancy
spuriously failed **~35** fp32 numeric-parity tests (rtol/atol 1e-3).
Set
`NVIDIA_TF32_OVERRIDE=0` in conftest before any CUDA library initializes
(uses
`setdefault` so users can opt back in). Verified: with the env var
unset,
`gpt2` and `qwen2.5-0.5b` L4 now pass.

### 2. transformers ≥5.x flattened ViT/CLIP weight names
(`src/mobius/models/vit.py`)
transformers 5.x flattened the ViT state dict to `layers.N.*` with
consolidated
`attention.{q,k,v,o}_proj` / `mlp.fc1/fc2` names. The legacy rename map
no
longer matched, leaving graph initializers unfilled (ORT load failure).
Added
an **additive** new-naming branch (legacy 5.0–5.9 path preserved) and
aligned
the in-test torch reference modules with the mobius graph param names.

### 3. gemma3 multimodal vision encoder (`src/mobius/models/gemma3.py`)
Two bugs broke the gemma-3 image-text-to-text pipeline:
- The full-VLM `preprocess_weights` prefixed vision weights but didn't
rename
the vision MLP `fc1/fc2` → `up_proj/down_proj`, so FCMLP initializers
were
  never filled (ORT load failure).
- The vision encoder returned rank-3 `(batch, tokens, hidden)`, but the
embedding sub-model declares `image_features` as rank-2 and gathers
along
  axis 0. Squeeze the batch dim to honor the 2-D contract (matching the
  `PixtralVLTask` precedent and the ort-genai runtime).

With both, **gemma-3-4b-it L4 golden passes** on CUDA.

## Verification
- TF32 fix: `gpt2`, `qwen2.5-0.5b` L4 pass with env unset.
- ViT/CLIP parity tests pass.
- gemma-3-4b-it L4 golden passes; 25 gemma3 build/task unit tests pass.
- `lintrunner` clean on all changed files.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---

## Architecture-diff notes (review follow-up)

Two entries in the `arch-diff-bot` comment (`e005489 -> 2f13c7d`)
deserve explanation:

### falcon: `Gelu` -> `Sigmoid + Mul`
The bloom-GELU fix changed falcon's MLP from a hardcoded
`activation="gelu"` to `activation=config.hidden_act or "gelu"` so
`FCMLP` can be shared with Bloom (which needs `gelu_pytorch_tanh`). For
a **real** falcon config this is a no-op:
`ArchitectureConfig.hidden_act` resolves to `"gelu"` via
`config.activation` (HF `FalconConfig` has no `hidden_act` and defaults
`activation="gelu"`). The arch-diff only showed SiLU (`Sigmoid+Mul`)
because the **synthetic test config** used the generic `_base_config`
default `hidden_act="silu"`. Fixed by setting `hidden_act="gelu"` in
falcon's synthetic config so it matches real Falcon.

### `RotaryEmbedding: num_heads: 2 -> 4`
This is a **false positive** from the diff tool's positional node
matching, *caused by* the SiLU regression above -- not a real change.
Each falcon layer emits two `RotaryEmbedding` nodes: Q-rotary
(`num_heads=4`) and K-rotary (`num_heads=2`, GQA `kv_heads=2`).
Inserting `Sigmoid+Mul` per layer shifted all subsequent node indices by
+2, so the tool aligned base's K-rotary (`num_heads=2`) against head's
Q-rotary (`num_heads=4`). Per-head counts are unchanged at both base and
head. Restoring `Gelu` removes the extra nodes, realigns indices, and
makes this spurious entry disappear.

---------

Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: justinchuby

examples/phi4mm_multimodal.py

@@ -16,7 +16,8 @@
 The model is split into 4 ONNX graphs, each running in its own
 ONNX Runtime session:
 
-    - **Vision**:    ``pixel_values`` + ``image_sizes`` → ``image_features``
+    - **Vision**:    ``pixel_values`` + ``image_sizes`` +
+      ``image_attention_mask`` → ``image_features``
       (SigLIP encoder + projection)
     - **Speech**:    ``audio_features`` + ``audio_sizes`` +
       ``audio_projection_mode`` → ``audio_features``
@@ -111,12 +112,16 @@ def prepare_vision_feeds(
     """Prepare feeds for the **vision** session.
 
     Returns:
-        ``{"pixel_values": [1, 3, H, W], "image_sizes": [1, 2]}``
-        ready for the vision ONNX model.
+        ``{"pixel_values": [crops, 3, H, W], "image_sizes": [1, 2],
+        "image_attention_mask": [crops, P, P]}`` ready for the vision
+        ONNX model (single image).
     """
-    pixel_values = _load_image(image_path, processor)
-    image_sizes = np.array([[pixel_values.shape[-2], pixel_values.shape[-1]]], dtype=np.int64)
-    return {"pixel_values": pixel_values, "image_sizes": image_sizes}
+    pixel_values, image_sizes, image_attention_mask = _load_image(image_path, processor)
+    return {
+        "pixel_values": pixel_values,
+        "image_sizes": image_sizes,
+        "image_attention_mask": image_attention_mask,
+    }
 
 
 def prepare_speech_feeds(
@@ -175,6 +180,8 @@ def prepare_decoder_feeds(
     inputs_embeds: np.ndarray,
     past_seq_len: int,
     past_kv: dict[str, np.ndarray],
+    vision_gate: np.ndarray,
+    speech_gate: np.ndarray,
 ) -> dict[str, np.ndarray]:
     """Prepare feeds for the **decoder** session.
 
@@ -185,6 +192,10 @@ def prepare_decoder_feeds(
         inputs_embeds: ``[batch, cur_seq_len, hidden_size]`` float32.
         past_seq_len: Number of tokens already in the KV cache.
         past_kv: Dict of ``past_key_values.{i}.key/value`` arrays.
+        vision_gate: Scalar LoRA gate emitted by the embedding model
+            (1.0 if any image token is present, else 0.0).
+        speech_gate: Scalar LoRA gate emitted by the embedding model
+            (1.0 if audio present and no image, else 0.0).
 
     Returns:
         Complete feeds dict for the decoder ONNX model.
@@ -197,6 +208,11 @@ def prepare_decoder_feeds(
         "inputs_embeds": inputs_embeds,
         "attention_mask": np.ones((batch_size, total_seq_len), dtype=np.int64),
         "position_ids": np.arange(past_seq_len, total_seq_len, dtype=np.int64)[np.newaxis, :],
+        # The decoder declares vision_gate/speech_gate as required scalar
+        # inputs; they are produced by the embedding model and select the
+        # active per-modality LoRA adapter.
+        "vision_gate": vision_gate,
+        "speech_gate": speech_gate,
         **past_kv,
     }
 
@@ -279,17 +295,27 @@ def _build_input_ids_vision_audio(
 # ---------------------------------------------------------------------------
 
 
-def _load_image(image_path: str, processor) -> np.ndarray:
+def _load_image(image_path: str, processor) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Load and preprocess an image using the HuggingFace processor.
 
+    Phi4MM's image processor performs an HD multi-crop transform and emits
+    ``input_image_embeds`` (the crop pixel tensor), ``image_sizes`` and an
+    ``image_attention_mask`` marking valid (non-padding) patches per crop.
+
     Returns:
-        ``pixel_values`` as ``[1, 3, H, W]`` float32 numpy array.
+        ``(pixel_values [crops, 3, H, W], image_sizes [1, 2],
+        image_attention_mask [crops, P, P])`` as float32/int64 arrays for a
+        single image.
     """
     from PIL import Image
 
     img = Image.open(image_path).convert("RGB")
     processed = processor.image_processor(images=img, return_tensors="np")
-    return processed["pixel_values"].astype(np.float32)
+    # Single image -> index 0 of the leading num_images dimension.
+    pixel_values = processed["input_image_embeds"][0].astype(np.float32)
+    image_sizes = processed["image_sizes"][0:1].astype(np.int64)
+    image_attention_mask = processed["image_attention_mask"][0].astype(np.float32)
+    return pixel_values, image_sizes, image_attention_mask
 
 
 def _load_audio(
@@ -417,9 +443,15 @@ def generate(
         )
         embed_out = embedding_session.run(embed_feeds)
         inputs_embeds = embed_out["inputs_embeds"]
+        # The embedding model also emits the per-modality LoRA gates derived
+        # from input_ids; thread them into the decoder, which requires them.
+        vision_gate = embed_out["vision_gate"]
+        speech_gate = embed_out["speech_gate"]
 
         # --- Decoder session ---
-        decoder_feeds = prepare_decoder_feeds(inputs_embeds, past_seq_len, past_kv)
+        decoder_feeds = prepare_decoder_feeds(
+            inputs_embeds, past_seq_len, past_kv, vision_gate, speech_gate
+        )
         outputs = decoder_session.run(decoder_feeds)
 
         logits = outputs["logits"]

examples/phi4mm_ort_genai.py

@@ -110,6 +110,8 @@ def _write_genai_config(config, output_dir: str) -> None:
                     "inputs_embeds": "inputs_embeds",
                     "attention_mask": "attention_mask",
                     "position_ids": "position_ids",
+                    "vision_gate": "vision_gate",
+                    "speech_gate": "speech_gate",
                     "past_key_names": "past_key_values.%d.key",
                     "past_value_names": "past_key_values.%d.value",
                 },
@@ -132,6 +134,8 @@ def _write_genai_config(config, output_dir: str) -> None:
                 },
                 "outputs": {
                     "inputs_embeds": "inputs_embeds",
+                    "vision_gate": "vision_gate",
+                    "speech_gate": "speech_gate",
                 },
             },
             "vision": {
@@ -144,6 +148,7 @@ def _write_genai_config(config, output_dir: str) -> None:
                 "inputs": {
                     "pixel_values": "pixel_values",
                     "image_sizes": "image_sizes",
+                    "image_attention_mask": "image_attention_mask",
                 },
                 "outputs": {
                     "image_features": "image_features",

scripts/generate_golden.py

@@ -951,6 +951,79 @@ def _generate_image_classification(case: TestCase, json_path: Path, device: str)
     )
 
 
+def _detection_forced_size(model_id: str, trust_remote_code: bool) -> dict | None:
+    """Return a fixed ``{height, width}`` size for object-detection export.
+
+    mobius exports object-detection models (e.g. YOLOS) at a *fixed* input
+    resolution taken from ``config.image_size`` (position embeddings are not
+    interpolated for arbitrary sizes).  To keep the golden reference and the
+    ONNX forward pass on the same footing, the HF image processor must be
+    forced to emit exactly that resolution instead of its default
+    aspect-preserving resize.
+    """
+    import transformers
+
+    config = transformers.AutoConfig.from_pretrained(
+        model_id, trust_remote_code=trust_remote_code
+    )
+    image_size = getattr(config, "image_size", None)
+    if isinstance(image_size, (list, tuple)) and len(image_size) == 2:
+        return {"height": int(image_size[0]), "width": int(image_size[1])}
+    if isinstance(image_size, int):
+        return {"height": image_size, "width": image_size}
+    return None
+
+
+def _generate_object_detection(case: TestCase, json_path: Path, device: str) -> None:
+    """Generate golden data for object detection (e.g. YOLOS).
+
+    The model emits per-query class ``logits`` of shape
+    ``[batch, num_queries, num_labels + 1]``.  The golden top-K is taken over
+    the *last* query's class-logit vector to match ``compare_golden``, which
+    slices ``logits[:, -1, :]``.  The image processor is forced to the model's
+    fixed export resolution so the golden and ONNX forward pass agree.
+    """
+    import torch
+    import transformers
+    from PIL import Image
+
+    from mobius._testing.golden import save_golden_ref
+
+    processor = transformers.AutoImageProcessor.from_pretrained(
+        case.model_id, trust_remote_code=case.trust_remote_code
+    )
+    model = transformers.AutoModelForObjectDetection.from_pretrained(
+        case.model_id,
+        torch_dtype=torch.float32,
+        device_map=device,
+        trust_remote_code=case.trust_remote_code,
+    ).eval()
+
+    image = Image.open(Path("testdata") / case.images[0])
+    forced_size = _detection_forced_size(case.model_id, case.trust_remote_code)
+    proc_kwargs = {"images": image, "return_tensors": "pt"}
+    if forced_size is not None:
+        proc_kwargs["size"] = forced_size
+    processed = processor(**proc_kwargs)
+    pixel_values = processed["pixel_values"].to(device)
+
+    with torch.no_grad():
+        outputs = model(pixel_values=pixel_values)
+    # logits: [batch, num_queries, num_labels + 1] -> last query's class vector
+    last_logits = outputs.logits[0, -1, :].cpu().numpy()
+    golden = _extract_logits_golden(last_logits)
+
+    save_golden_ref(
+        json_path,
+        top1_id=golden["top1_id"],
+        top2_id=golden["top2_id"],
+        top10_ids=golden["top10_ids"],
+        top10_logits=golden["top10_logits"],
+        logits_summary=golden["logits_summary"],
+        input_ids=np.array([[0]], dtype=np.int64),  # placeholder (no text input)
+    )
+
+
 # ---- Phi4MM multimodal generator ----
 
 
@@ -1129,10 +1202,13 @@ def _generate_phi4mm_multimodal(case: TestCase, json_path: Path, device: str) ->
         processor = transformers.AutoProcessor.from_pretrained(
             case.model_id, trust_remote_code=True
         )
-        # Load in bfloat16 to reduce memory (14B model)
+        # Load in float32 to match the f32 runtime used by the L4/L5 tests.
+        # bf16 goldens produced flat/tied logit distributions that caused
+        # argmax instability against the f32 model output (see phi4mm L4
+        # false-failures: exact top-2 ties, top-10 spans <3 logits).
         model = transformers.AutoModelForCausalLM.from_pretrained(
             case.model_id,
-            torch_dtype=torch.bfloat16,
+            torch_dtype=torch.float32,
             device_map=device,
             trust_remote_code=True,
             _attn_implementation="eager",
@@ -1284,7 +1360,7 @@ def _generate_phi4mm_multimodal(case: TestCase, json_path: Path, device: str) ->
     "depth-estimation": _generate_image_classification,
     "image-segmentation": _generate_image_classification,
     "image-to-image": _generate_image_classification,
-    "object-detection": _generate_image_classification,
+    "object-detection": _generate_object_detection,
     "phi4mm-multimodal": _generate_phi4mm_multimodal,
 }
 

src/mobius/_configs/_base.py

@@ -1370,13 +1370,32 @@ class YolosConfig(EncoderConfig):
     """
 
     num_detection_tokens: int = 100
+    # YOLOS uses a rectangular input resolution (e.g. [800, 1333] for
+    # yolos-tiny). The base ArchitectureConfig collapses ``image_size`` to a
+    # single int (height), which would size the learned position embeddings
+    # for a square image and mismatch the pretrained weights. Preserve both
+    # dimensions here so the patch grid (and position-embedding length) is
+    # computed correctly.
+    image_height: int = 800
+    image_width: int = 1333
 
     @classmethod
     def from_transformers(cls, config, parent_config=None) -> YolosConfig:
         base = ArchitectureConfig.from_transformers(config, parent_config)
+        raw_image_size = getattr(config, "image_size", [base.image_size, base.image_size])
+        if isinstance(raw_image_size, dict):
+            height = int(raw_image_size.get("height", base.image_size))
+            width = int(raw_image_size.get("width", base.image_size))
+        elif isinstance(raw_image_size, (list, tuple)):
+            height = int(raw_image_size[0])
+            width = int(raw_image_size[-1])
+        else:
+            height = width = int(raw_image_size)
         return cls(
             **_shallow_fields(base),
             num_detection_tokens=getattr(config, "num_detection_tokens", 100),
+            image_height=height,
+            image_width=width,
         )
 
 

src/mobius/_testing/generation.py

@@ -88,6 +88,10 @@ def generate(
                 shape = self.session.get_input_shape(name) or []
                 static = [d if isinstance(d, int) and d > 0 else batch_size for d in shape]
                 past_kv[name] = np.zeros(static, dtype=np.float32)
+            elif ltype in ("mlp", "moe"):
+                # Pure feed-forward layers (e.g. NemotronH hybrid) carry no
+                # attention KV and no recurrent state — nothing to initialize.
+                continue
             else:
                 past_kv[f"past_key_values.{i}.key"] = np.zeros(
                     (batch_size, num_kv_heads, 0, head_dim), dtype=np.float32
@@ -154,6 +158,9 @@ def generate(
                     dst = f"past_key_values.{i}.conv_state"
                     if src in outputs:
                         past_kv[dst] = outputs[src]
+                elif ltype in ("mlp", "moe"):
+                    # Pure feed-forward layers have no cache state to carry.
+                    continue
                 else:
                     past_kv[f"past_key_values.{i}.key"] = outputs[f"present.{i}.key"]
                     past_kv[f"past_key_values.{i}.value"] = outputs[f"present.{i}.value"]

src/mobius/_testing/parity.py

@@ -65,6 +65,18 @@ class ParityReport:
     "int4": 1.0,
 }
 
+# Minimum top-10 overlap (count of shared tokens) at/above which an argmax
+# mismatch is treated as a tie-break (AMBIGUOUS) rather than a divergence
+# (FAIL).  A count of >=9 means at least 9 of the 10 highest tokens agree with
+# the reference.  NOTE: this MUST be a count, not a Jaccard ratio: for two
+# size-10 sets sharing k tokens the Jaccard is k/(20-k), so 9/10 overlap is
+# only 0.818 (a >=0.9 ratio would require identical sets, k=10, defeating the
+# "9 of 10 agree" intent).  The branch additionally requires the predicted
+# token to be within the golden top-10 AND the golden argmax to remain in the
+# ONNX top-2, so it only rescues genuine #1/#2 tie-break swaps (e.g. CUDA
+# float32 accumulation noise), not a low-ranked token promoted to #1.
+_AMBIGUOUS_TOP10_OVERLAP: int = 9
+
 # Per-dtype default tolerances for L3 synthetic parity.
 DEFAULT_TOLERANCES: dict[str, tuple[float, float]] = {
     # (atol, rtol)
@@ -271,14 +283,15 @@ def compare_golden(
 
     argmax_match = onnx_top1 == golden_top1_id
 
-    # Top-10 Jaccard
+    # Top-10 Jaccard (reported) and raw overlap count (gate).
     onnx_top10 = set(np.argsort(onnx_last_f64)[-10:].tolist())
     golden_top10 = set(golden_top10_ids)
-    jaccard = (
-        len(onnx_top10 & golden_top10) / len(onnx_top10 | golden_top10)
-        if golden_top10
-        else 0.0
-    )
+    overlap = len(onnx_top10 & golden_top10)
+    jaccard = overlap / len(onnx_top10 | golden_top10) if golden_top10 else 0.0
+    # ONNX top-2 token ids, used to confirm an argmax mismatch is a #1/#2
+    # tie-break swap (golden argmax still ranked #1 or #2 by ONNX) rather than
+    # a low-ranked token being promoted to #1.
+    onnx_top2_ids = set(np.argsort(onnx_last_f64)[-2:].tolist())
 
     # Gate decision
     if argmax_match:
@@ -291,6 +304,26 @@ def compare_golden(
             f"but matches top2={golden_top2_id} and near-tie detected "
             f"(gap={abs(top1_logit - top2_logit):.4f} < margin={margin})"
         )
+    elif (
+        overlap >= _AMBIGUOUS_TOP10_OVERLAP
+        and onnx_top1 in golden_top10
+        and golden_top1_id in onnx_top2_ids
+    ):
+        # >=9 of the golden top-10 tokens agree, the predicted token is itself a
+        # golden top-10 token, AND the golden argmax is still ranked #1 or #2 by
+        # ONNX — i.e. the argmax difference is a near-tie #1/#2 swap, not a
+        # divergence.  This catches CUDA float32 near-ties whose absolute logit
+        # gap exceeds the per-dtype ``near_tie`` margin (CUDA accumulation noise
+        # > CPU) yet whose ranking otherwise matches the reference.  Requiring
+        # the golden argmax to stay in the ONNX top-2 prevents masking a real
+        # divergence where a low-ranked token is promoted to #1 with a large gap.
+        result = ParityResult.AMBIGUOUS
+        message = (
+            f"L4 AMBIGUOUS: argmax={onnx_top1} != golden_top1={golden_top1_id}, "
+            f"but {overlap}/10 of the golden top-10 agree, argmax is within the "
+            f"golden top-10, and golden_top1 remains in the ONNX top-2 "
+            f"— tie-break, not divergence"
+        )
     else:
         result = ParityResult.FAIL
         message = (