Fix Phi-3 bridge support for explicit head_dim

tests/integration/model_bridge/test_phi3_tiny.py

@@ -0,0 +1,131 @@
+"""Tiny-random Phi-3 bridge tests for configs with an explicit ``head_dim``.
+
+Valid Phi-3 configs may set ``head_dim`` independently of
+``hidden_size // num_attention_heads`` (the fused QKV width then differs from
+``3 * hidden_size``). The adapter previously re-derived ``d_head`` from
+``d_model // n_heads`` and crashed splitting the fused matrix during bridge
+construction. Uses tiny programmatic configs with real weights — no network
+access or weight downloads. Fixture pattern mirrors
+tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py.
+"""
+
+import pytest
+import torch
+from transformers import AutoConfig, AutoModelForCausalLM
+
+from transformer_lens.model_bridge.bridge import TransformerBridge
+from transformer_lens.model_bridge.sources import build_bridge_config_from_hf
+from transformer_lens.model_bridge.supported_architectures.phi3 import (
+    Phi3ArchitectureAdapter,
+)
+
+FP32_TOL = 1e-5
+
+# head_dim is decoupled from hidden_size // num_attention_heads in both:
+# MHA derives 32 // 3 = 10 vs explicit 8; GQA derives 32 // 4 = 8 vs explicit 6.
+GEOMETRIES = {
+    "mha": dict(num_attention_heads=3, num_key_value_heads=3, head_dim=8),
+    "gqa": dict(num_attention_heads=4, num_key_value_heads=2, head_dim=6),
+}
+
+
+class _MockTokenizer:
+    """Stand-in to satisfy TransformerBridge(tokenizer=...)."""
+
+    pass
+
+
+@pytest.fixture(scope="module", params=sorted(GEOMETRIES), ids=sorted(GEOMETRIES))
+def geometry(request):
+    """One attention geometry (MHA or GQA), each with a decoupled head_dim."""
+    return GEOMETRIES[request.param]
+
+
+@pytest.fixture(scope="module")
+def hf_model(geometry):
+    hf_config = AutoConfig.for_model(
+        "phi3",
+        vocab_size=100,
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        max_position_embeddings=64,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        **geometry,
+    )
+    torch.manual_seed(0)
+    model = AutoModelForCausalLM.from_config(hf_config, attn_implementation="eager")
+    return model.eval()
+
+
+@pytest.fixture(scope="module")
+def tokens():
+    # Ids stay below the tiny vocab_size=100.
+    return torch.tensor([[1, 5, 7, 42, 9]])
+
+
+@pytest.fixture(scope="module")
+def hf_logits(hf_model, tokens):
+    """Reference logits, captured before the bridge wraps (and mutates) the HF module.
+
+    This keeps them an independent ground truth — helpers.assert_bridge_matches_hf
+    compares against the already-wrapped model, so it can't provide that.
+    """
+    with torch.no_grad():
+        return hf_model(tokens).logits
+
+
+@pytest.fixture(scope="module")
+def bridge(hf_model, hf_logits):
+    """TransformerBridge over the tiny HF model.
+
+    The config goes through build_bridge_config_from_hf — the same translation
+    path boot_transformers uses, including the explicit-head_dim mapping under
+    test. Depends on hf_logits so the reference capture happens before wrapping.
+    """
+    bridge_config = build_bridge_config_from_hf(
+        hf_model.config, "Phi3ForCausalLM", "phi3-tiny", torch.float32
+    )
+    adapter = Phi3ArchitectureAdapter(bridge_config)
+    return TransformerBridge(model=hf_model, adapter=adapter, tokenizer=_MockTokenizer())
+
+
+def test_mapped_d_head_honours_explicit_head_dim(bridge, geometry) -> None:
+    assert bridge.cfg.d_head == geometry["head_dim"]
+
+
+def test_split_qkv_weights_match_fused_slices(bridge, geometry) -> None:
+    """Q/K/V weights must be bitwise slices of the original fused qkv_proj."""
+    head_dim = geometry["head_dim"]
+    q_size = bridge.cfg.n_heads * head_dim
+    kv_size = (bridge.cfg.n_key_value_heads or bridge.cfg.n_heads) * head_dim
+    for block in bridge.blocks:
+        fused = block.attn.qkv.weight
+        assert fused.shape[0] == q_size + 2 * kv_size
+        assert torch.equal(block.attn.q.weight, fused[:q_size])
+        assert torch.equal(block.attn.k.weight, fused[q_size : q_size + kv_size])
+        assert torch.equal(block.attn.v.weight, fused[q_size + kv_size :])
+
+
+def test_forward_matches_hf_eager(bridge, hf_logits, tokens) -> None:
+    with torch.no_grad():
+        bridge_logits = bridge(tokens)
+    max_diff = (bridge_logits - hf_logits).abs().max().item()
+    assert max_diff < FP32_TOL, (
+        f"bridge vs HF eager drift={max_diff:.2e} exceeds {FP32_TOL:.0e} "
+        f"for explicit-head_dim Phi-3"
+    )
+
+
+def test_run_with_cache_head_shapes(bridge, tokens, geometry) -> None:
+    """Cached per-head activations must use the explicit head_dim."""
+    head_dim = geometry["head_dim"]
+    _, cache = bridge.run_with_cache(tokens)
+    n_kv = bridge.cfg.n_key_value_heads or bridge.cfg.n_heads
+    batch, seq = tokens.shape
+    z = cache["blocks.0.attn.hook_z"]
+    assert z.shape == (batch, seq, bridge.cfg.n_heads, head_dim)
+    k = cache["blocks.0.attn.k.hook_out"]
+    assert k.shape[-2:] == (n_kv, head_dim)

tests/unit/model_bridge/supported_architectures/test_phi3_adapter.py

@@ -49,11 +49,12 @@ def _make_cfg(
     d_mlp: int = D_MLP,
     d_vocab: int = D_VOCAB,
     n_ctx: int = N_CTX,
+    d_head: int | None = None,
 ) -> TransformerBridgeConfig:
     """Return a minimal TransformerBridgeConfig for Phi-3 adapter tests."""
     return TransformerBridgeConfig(
         d_model=d_model,
-        d_head=d_model // n_heads,
+        d_head=d_head if d_head is not None else d_model // n_heads,
         n_layers=n_layers,
         n_ctx=n_ctx,
         n_heads=n_heads,
@@ -302,3 +303,100 @@ def test_fold_skipped_when_not_requested(self, adapter: Phi3ArchitectureAdapter)
         out = adapter.preprocess_weights(sd)
         assert torch.allclose(out["blocks.0.ln1.weight"], torch.full((D_MODEL,), 2.0))
         assert torch.allclose(out["blocks.0.attn.q.weight"], torch.ones(N_HEADS * D_HEAD, D_MODEL))
+
+
+# ---------------------------------------------------------------------------
+# Explicit head_dim: d_head decoupled from d_model // n_heads
+# ---------------------------------------------------------------------------
+
+
+class _FakeAttention(torch.nn.Module):
+    """Minimal stand-in for HF Phi3Attention exposing only qkv_proj."""
+
+    def __init__(self, d_model: int, qkv_rows: int, bias: bool) -> None:
+        super().__init__()
+        self.qkv_proj = torch.nn.Linear(d_model, qkv_rows, bias=bias)
+        with torch.no_grad():
+            self.qkv_proj.weight.copy_(
+                torch.arange(qkv_rows * d_model, dtype=torch.float32).reshape(qkv_rows, d_model)
+            )
+            if bias:
+                self.qkv_proj.bias.copy_(torch.arange(qkv_rows, dtype=torch.float32))
+
+
+class TestPhi3ExplicitHeadDim:
+    """HF configs may set head_dim explicitly, decoupled from d_model // n_heads.
+
+    Geometry is chosen so d_model // n_heads never equals d_head — the old
+    derived-d_head formula cannot accidentally produce the right split sizes.
+    """
+
+    # MHA: d_model=16, n_heads=3 → derived would be 5, explicit d_head=4
+    MHA = dict(d_model=16, n_heads=3, n_kv_heads=3, d_head=4)
+    # GQA: d_model=16, n_heads=4 → derived would be 4, explicit d_head=3
+    GQA = dict(d_model=16, n_heads=4, n_kv_heads=2, d_head=3)
+
+    def _split_sizes(self, adapter: Phi3ArchitectureAdapter) -> list[int]:
+        conv = adapter.weight_processing_conversions["blocks.{i}.attn.q"]
+        assert isinstance(conv.tensor_conversion, _SizedSplitConversion)
+        return conv.tensor_conversion.sizes
+
+    def test_mha_conversion_split_sizes(self) -> None:
+        """Q/K/V each get n_heads * d_head = 12 rows, not 3 * (16 // 3) = 15."""
+        adapter = Phi3ArchitectureAdapter(_make_cfg(**self.MHA))
+        assert self._split_sizes(adapter) == [12, 12, 12]
+
+    def test_gqa_conversion_split_sizes(self) -> None:
+        """Q gets n_heads * d_head = 12 rows; K/V get n_kv_heads * d_head = 6."""
+        adapter = Phi3ArchitectureAdapter(_make_cfg(**self.GQA))
+        assert self._split_sizes(adapter) == [12, 6, 6]
+
+    def test_conversion_split_sizes_consistent_across_qkv(self) -> None:
+        """All three conversions must share the same size list."""
+        adapter = Phi3ArchitectureAdapter(_make_cfg(**self.GQA))
+        for key in ["blocks.{i}.attn.q", "blocks.{i}.attn.k", "blocks.{i}.attn.v"]:
+            conv = adapter.weight_processing_conversions[key]
+            assert isinstance(conv.tensor_conversion, _SizedSplitConversion)
+            assert conv.tensor_conversion.sizes == [12, 6, 6]
+
+    def test_split_phi3_qkv_mha_with_bias(self) -> None:
+        """Live split of a fused 36-row qkv_proj into 12/12/12 with biases."""
+        adapter = Phi3ArchitectureAdapter(_make_cfg(**self.MHA))
+        fake = _FakeAttention(d_model=16, qkv_rows=36, bias=True)
+        q, k, v = adapter._split_phi3_qkv(fake)
+        assert q.weight.shape == (12, 16)
+        assert k.weight.shape == (12, 16)
+        assert v.weight.shape == (12, 16)
+        fused_w = fake.qkv_proj.weight
+        fused_b = fake.qkv_proj.bias
+        assert torch.equal(q.weight, fused_w[:12])
+        assert torch.equal(k.weight, fused_w[12:24])
+        assert torch.equal(v.weight, fused_w[24:])
+        assert torch.equal(q.bias, fused_b[:12])
+        assert torch.equal(k.bias, fused_b[12:24])
+        assert torch.equal(v.bias, fused_b[24:])
+
+    def test_split_phi3_qkv_gqa_without_bias(self) -> None:
+        """Live split of a fused 24-row GQA qkv_proj into 12/6/6 without biases."""
+        adapter = Phi3ArchitectureAdapter(_make_cfg(**self.GQA))
+        fake = _FakeAttention(d_model=16, qkv_rows=24, bias=False)
+        q, k, v = adapter._split_phi3_qkv(fake)
+        assert q.weight.shape == (12, 16)
+        assert k.weight.shape == (6, 16)
+        assert v.weight.shape == (6, 16)
+        fused_w = fake.qkv_proj.weight
+        assert torch.equal(q.weight, fused_w[:12])
+        assert torch.equal(k.weight, fused_w[12:18])
+        assert torch.equal(v.weight, fused_w[18:])
+        assert q.bias is None
+        assert k.bias is None
+        assert v.bias is None
+
+    def test_derived_geometry_unchanged(self) -> None:
+        """Ordinary configs (d_head == d_model // n_heads) keep the same sizes."""
+        adapter = Phi3ArchitectureAdapter(_make_cfg())
+        assert self._split_sizes(adapter) == [
+            N_HEADS * D_HEAD,
+            N_KV_HEADS * D_HEAD,
+            N_KV_HEADS * D_HEAD,
+        ]

tests/unit/model_bridge/test_config_mapping.py

@@ -0,0 +1,41 @@
+"""Unit tests for map_default_transformer_lens_config field precedence.
+
+Style mirrors tests/unit/model_bridge/test_output_logits_contract.py: plain
+SimpleNamespace HF-config stand-ins, no network access.
+"""
+
+from types import SimpleNamespace
+
+from transformer_lens.model_bridge.sources.transformers import (
+    map_default_transformer_lens_config,
+)
+
+
+def _hf_config(**overrides: object) -> SimpleNamespace:
+    values: dict[str, object] = {
+        "hidden_size": 16,
+        "num_attention_heads": 4,
+        "num_hidden_layers": 2,
+        "vocab_size": 32,
+        "max_position_embeddings": 64,
+        "intermediate_size": 32,
+    }
+    values.update(overrides)
+    return SimpleNamespace(**values)
+
+
+def test_explicit_head_dim_wins_over_derived() -> None:
+    """head_dim need not equal hidden_size // num_attention_heads (16 // 4 = 4)."""
+    mapped = map_default_transformer_lens_config(_hf_config(head_dim=8))
+    assert mapped.d_head == 8
+
+
+def test_d_head_falls_back_to_derived_when_head_dim_absent() -> None:
+    mapped = map_default_transformer_lens_config(_hf_config())
+    assert mapped.d_head == 16 // 4
+
+
+def test_d_head_falls_back_to_derived_when_head_dim_is_none() -> None:
+    """HF configs may carry head_dim=None; treat it the same as absent."""
+    mapped = map_default_transformer_lens_config(_hf_config(head_dim=None))
+    assert mapped.d_head == 16 // 4

transformer_lens/model_bridge/supported_architectures/phi3.py

@@ -61,8 +61,10 @@ def __init__(self, cfg: Any) -> None:
         # LN folding is handled in preprocess_weights() instead.
         self.supports_fold_ln = False
 
-        # GQA: Q has n_heads * d_head, K/V have n_kv_heads * d_head
-        d_head = cfg.d_model // cfg.n_heads
+        # GQA: Q has n_heads * d_head, K/V have n_kv_heads * d_head.
+        # cfg.d_head honours an explicit HF head_dim, which need not equal
+        # d_model // n_heads.
+        d_head = cfg.d_head
         n_kv_heads = cfg.n_key_value_heads or cfg.n_heads
         q_size = cfg.n_heads * d_head
         kv_size = n_kv_heads * d_head
@@ -167,8 +169,10 @@ def _split_phi3_qkv(
         qkv_weight = original_attention_component.qkv_proj.weight
         d_model = qkv_weight.shape[1]
 
-        # GQA: Q has n_heads * d_head, K/V have n_kv_heads * d_head each
-        d_head = self.cfg.d_model // self.cfg.n_heads
+        # GQA: Q has n_heads * d_head, K/V have n_kv_heads * d_head each.
+        # cfg.d_head honours an explicit HF head_dim, which need not equal
+        # d_model // n_heads.
+        d_head = self.cfg.d_head
         n_kv_heads = self.cfg.n_key_value_heads or self.cfg.n_heads
         q_size = self.cfg.n_heads * d_head
         kv_size = n_kv_heads * d_head