Mistral: use fork-capable attention bridge (enable hook_attn_in / set_use_attn_in)

tests/integration/model_bridge/test_bridge_integration.py

@@ -704,22 +704,18 @@ def test_AttentionBridge_preserves_fp_input_when_first_param_is_quantized():
     producing gibberish logits on every quantized model.
 
     Fakes a "quantized first parameter" by replacing q_proj.weight with a
-    uint8 tensor, then runs a forward and asserts the input the original
-    component receives is still floating-point.
+    uint8 tensor, then runs a forward and asserts the hidden states reaching
+    q_proj are still floating-point.
     """
     from transformer_lens.model_bridge.generalized_components.attention import (
         AttentionBridge,
     )
 
-    # Use tiny Mistral — it's a plain AttentionBridge (not JointQKV).
     bridge: TransformerBridge = TransformerBridge.boot_transformers(  # type: ignore
         "trl-internal-testing/tiny-MistralForCausalLM-0.2", device="cpu"
     )
 
     attn_bridge = bridge.blocks[0].attn  # type: ignore[attr-defined]
-    assert (
-        type(attn_bridge).__name__ == "AttentionBridge"
-    ), f"Expected plain AttentionBridge, got {type(attn_bridge).__name__}"
     assert isinstance(attn_bridge, AttentionBridge)
 
     original = attn_bridge.original_component
@@ -734,42 +730,32 @@ def test_AttentionBridge_preserves_fp_input_when_first_param_is_quantized():
         next(original.parameters()).dtype == torch.uint8
     ), "Test setup: first param should be uint8 to trigger the bug condition"
 
-    # Capture what dtype reaches the original component's forward.
+    # Capture what dtype reaches q_proj. The uint8 weight makes the matmul
+    # itself fail, which is fine: the pre-hook already recorded the dtype.
     received_dtype: list = []
-    orig_forward = original.forward
-
-    def capture(*args, **kwargs):
-        if "hidden_states" in kwargs:
-            received_dtype.append(kwargs["hidden_states"].dtype)
-        elif args:
-            received_dtype.append(args[0].dtype)
-        # Don't actually run forward — fake-quantized weight would error.
-        # Return a shape-compatible dummy. HF Mistral attention returns a tuple.
-        bsz, seq, d_model = (kwargs.get("hidden_states", args[0] if args else None)).shape
-        n_heads = bridge.cfg.n_heads  # type: ignore[attr-defined]
-        return (
-            torch.zeros(bsz, seq, d_model, dtype=torch.float32),
-            torch.zeros(bsz, n_heads, seq, seq, dtype=torch.float32),
-        )
 
-    original.forward = capture  # type: ignore[method-assign]
+    def capture(module, inputs):
+        if inputs and isinstance(inputs[0], torch.Tensor):
+            received_dtype.append(inputs[0].dtype)
+
+    handle = original.q_proj.register_forward_pre_hook(capture)
     try:
         test_input = torch.tensor([[1, 2, 3, 4, 5]])
         with torch.no_grad():
             try:
                 bridge(test_input)
             except Exception:
-                pass  # downstream may fail; we only care what reached attn forward
+                pass  # downstream may fail; we only care what reached q_proj
     finally:
-        original.forward = orig_forward  # type: ignore[method-assign]
+        handle.remove()
         original.q_proj.weight = fp_weight
 
-    assert len(received_dtype) > 0, "Original attention forward never called"
+    assert len(received_dtype) > 0, "q_proj was never called"
     for dt in received_dtype:
         assert dt.is_floating_point, (
-            f"Bridge passed dtype={dt} to original attention forward, but it must be "
-            f"floating point. Regression of the AttentionBridge dtype-cast bug — "
-            f"target_dtype must skip non-fp (quantized-storage) parameters."
+            f"Bridge passed dtype={dt} into q_proj, but it must be floating point. "
+            f"Regression of the AttentionBridge dtype-cast bug: target_dtype must "
+            f"skip non-fp (quantized-storage) parameters."
         )
 
 

tests/unit/model_bridge/supported_architectures/test_mistral_adapter.py

@@ -112,12 +112,13 @@ def test_block_submodule_types(self, adapter: MistralArchitectureAdapter) -> Non
         assert isinstance(blocks.submodules["attn"], AttentionBridge)
         assert isinstance(blocks.submodules["mlp"], GatedMLPBridge)
 
-    def test_attn_is_not_position_embeddings_subclass(
-        self, adapter: MistralArchitectureAdapter
-    ) -> None:
-        """Mistral uses plain AttentionBridge, not PositionEmbeddingsAttentionBridge."""
+    def test_attn_is_position_embeddings_bridge(self, adapter: MistralArchitectureAdapter) -> None:
+        """Mistral uses PositionEmbeddingsAttentionBridge, like Qwen2.
+
+        That bridge is what gives Mistral hook_attn_in and the q/k/v forks.
+        """
         attn = adapter.component_mapping["blocks"].submodules["attn"]
-        assert not isinstance(attn, PositionEmbeddingsAttentionBridge)
+        assert isinstance(attn, PositionEmbeddingsAttentionBridge)
 
     def test_block_submodule_hf_paths(self, adapter: MistralArchitectureAdapter) -> None:
         blocks = adapter.component_mapping["blocks"]

tests/unit/model_bridge/test_mistral_attn_in_fork.py

@@ -0,0 +1,56 @@
+"""Mistral adapter uses a fork-capable attention bridge (per-receiver ``hook_attn_in``).
+
+Mistral has separate q/k/v/o projections with RoPE + GQA, identical in structure to Qwen2, but
+its adapter previously used the plain ``AttentionBridge`` (which delegates q/k/v to HF and exposes
+no fork point), so ``set_use_attn_in`` raised and ``blocks.{i}.hook_attn_in`` did not exist. It now
+uses ``PositionEmbeddingsAttentionBridge`` like Qwen2, enabling the input fork the circuit-analysis
+tooling relies on. These tests boot the tiny random Mistral so they stay cheap.
+"""
+from __future__ import annotations
+
+import pytest
+import torch
+
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.model_bridge.generalized_components.position_embeddings_attention import (
+    PositionEmbeddingsAttentionBridge,
+)
+
+_MODEL = "hf-internal-testing/tiny-random-MistralForCausalLM"
+
+
+@pytest.fixture(scope="module")
+def bridge() -> TransformerBridge:
+    return TransformerBridge.boot_transformers(_MODEL, device="cpu")
+
+
+@pytest.mark.slow
+def test_mistral_uses_fork_capable_attention(bridge: TransformerBridge) -> None:
+    assert isinstance(bridge.blocks[0].attn, PositionEmbeddingsAttentionBridge)
+
+
+@pytest.mark.slow
+def test_mistral_hook_attn_in_shape_and_intervenes(bridge: TransformerBridge) -> None:
+    toks = torch.arange(1, 7).unsqueeze(0)
+    baseline = bridge(toks)
+
+    bridge.set_use_attn_in(True)  # previously raised on the plain AttentionBridge
+
+    captured: dict = {}
+
+    def _grab(tensor: torch.Tensor, hook: object) -> torch.Tensor:
+        captured["shape"] = tuple(tensor.shape)
+        return tensor
+
+    bridge.run_with_hooks(toks, fwd_hooks=[("blocks.0.hook_attn_in", _grab)])
+    # [batch, pos, n_heads, d_model]
+    assert captured["shape"] == (1, 6, bridge.cfg.n_heads, bridge.cfg.d_model)
+
+    def _zero_head0(tensor: torch.Tensor, hook: object) -> torch.Tensor:
+        tensor = tensor.clone()
+        tensor[:, :, 0, :] = 0.0
+        return tensor
+
+    ablated = bridge.run_with_hooks(toks, fwd_hooks=[("blocks.0.hook_attn_in", _zero_head0)])
+    # zeroing a single head's forked input must actually change the output
+    assert (ablated - baseline).abs().max().item() > 1e-6

transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py

@@ -274,12 +274,17 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
         # Apply input hook
         hidden_states = self.hook_in(hidden_states)
 
-        # Match dtype of HF module
+        # Match dtype of HF module. Skip non-fp params: quantized weights (bnb
+        # uint8/int8, GPTQ/AWQ int32, HQQ, torchao) are stored in integer dtypes
+        # and dequantized internally during matmul. The compute dtype must come
+        # from a fp parameter; casting fp inputs to an integer storage dtype
+        # destroys precision.
         target_dtype = None
-        try:
-            target_dtype = next(hf_attn.parameters()).dtype
-        except StopIteration:
-            pass
+        for p in hf_attn.parameters():
+            if not p.dtype.is_floating_point:
+                continue
+            target_dtype = p.dtype
+            break
         if target_dtype is not None and hidden_states.is_floating_point():
             hidden_states = hidden_states.to(dtype=target_dtype)
 

transformer_lens/model_bridge/supported_architectures/mistral.py

@@ -4,11 +4,11 @@
 
 from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
 from transformer_lens.model_bridge.generalized_components import (
-    AttentionBridge,
     BlockBridge,
     EmbeddingBridge,
     GatedMLPBridge,
     LinearBridge,
+    PositionEmbeddingsAttentionBridge,
     RMSNormalizationBridge,
     RotaryEmbeddingBridge,
     UnembeddingBridge,
@@ -49,10 +49,11 @@ def __init__(self, cfg: Any) -> None:
             "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg),
             "blocks": BlockBridge(
                 name="model.layers",
+                config=self.cfg,
                 submodules={
                     "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
                     "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg),
-                    "attn": AttentionBridge(
+                    "attn": PositionEmbeddingsAttentionBridge(
                         name="self_attn",
                         config=self.cfg,
                         requires_position_embeddings=True,