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

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/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,