Test hook shape coverage

docs/source/content/model_structure.md

@@ -0,0 +1,154 @@
+# TransformerBridge Model Structure
+
+This page describes the structure exposed by TransformerBridge, the canonical hook names to use, and the expected tensor shapes at each hook point.
+
+## Overview
+
+TransformerBridge wraps a Hugging Face model behind a consistent TransformerLens interface. It relies on:
+- An ArchitectureAdapter that understands the HF module graph and provides a mapping to bridge components
+- Generalized components (Embedding, Attention, MLP, Normalization, Block) exposing uniform hook points
+- A light aliasing layer for backwards compatibility with legacy TransformerLens hook names
+
+Construct a bridge from a HF model id:
+
+```python
+from transformer_lens.model_bridge.bridge import TransformerBridge
+from transformer_lens.model_bridge.sources import transformers as bridge_sources  # registers boot
+
+bridge = TransformerBridge.boot_transformers("gpt2", device="cpu")
+```
+
+You can then call the familiar APIs: `to_tokens`, `to_string`, `generate`, `run_with_hooks`, `run_with_cache`.
+
+## Top-Level Components
+
+Typical decoder-only models expose these top-level components (names vary by architecture):
+- `embed`: token embedding
+- `pos_embed` (if applicable) or rotary embeddings inside attention
+- `blocks`: list-like container of transformer blocks
+- `ln_final` (if applicable): final normalization
+- `unembed`: output projection to vocabulary logits
+
+Each `blocks.{i}` is a `BlockBridge` with subcomponents:
+- `ln1`: normalization before attention
+- `attn`: attention module
+- `ln2`: normalization before MLP
+- `mlp`: MLP module
+
+## Canonical Hook Names
+
+Use these canonical (non-aliased) names when adding hooks or reading from the cache.
+
+### Embedding
+- `embed.hook_in`: token ids (batch, pos)
+- `embed.hook_out`: embeddings (batch, pos, d_model)
+  - *Legacy alias: `hook_embed`*
+- `pos_embed.hook_in` / `pos_embed.hook_out`: same shapes as above
+  - *Legacy alias: `hook_pos_embed`*
+
+### Residual stream
+- `blocks.{i}.hook_in`: residual stream into block (batch, pos, d_model)
+  - *Legacy alias: `blocks.{i}.hook_resid_pre`*
+- `blocks.{i}.hook_out`: residual stream out of block (batch, pos, d_model)
+  - *Legacy alias: `blocks.{i}.hook_resid_post`*
+- `blocks.{i}.attn.hook_out`: residual stream after attention (batch, pos, d_model)
+  - *Legacy alias: `blocks.{i}.hook_resid_mid`*
+
+### Attention
+- `blocks.{i}.attn.hook_in`: (batch, pos, d_model)
+  - *Legacy alias: `blocks.{i}.hook_attn_in`*
+- `blocks.{i}.attn.hook_out`: (batch, pos, d_model)
+  - *Legacy alias: `blocks.{i}.hook_attn_out`*
+- `blocks.{i}.attn.hook_hidden_states`: primary output for caching (batch, pos, d_model)
+  - *Legacy alias: `blocks.{i}.attn.hook_result`*
+- `blocks.{i}.attn.hook_attn_scores`: raw attention scores before softmax (batch, n_heads, pos, pos)
+- `blocks.{i}.attn.hook_pattern`: attention pattern after softmax and NaN handling (n_heads, pos, pos)
+  - *Legacy alias: `blocks.{i}.attn.hook_attention_weights`*
+- When present, sub-projections: `blocks.{i}.attn.q/k/v/o.hook_in` / `.hook_out` (commonly (batch, pos, d_model))
+  - *Legacy aliases: `blocks.{i}.hook_q_input`, `blocks.{i}.hook_k_input`, `blocks.{i}.hook_v_input`, `blocks.{i}.hook_q`, `blocks.{i}.hook_k`, `blocks.{i}.hook_v`*
+
+#### Individual Q/K/V Hooks
+All attention bridges provide access to individual Q, K, V activations through `HookPointWrapper` properties:
+
+- `blocks.{i}.attn.q.hook_in` / `blocks.{i}.attn.q.hook_out`: Q projection hooks (batch, pos, n_heads, d_head)
+- `blocks.{i}.attn.k.hook_in` / `blocks.{i}.attn.k.hook_out`: K projection hooks (batch, pos, n_heads, d_head)
+- `blocks.{i}.attn.v.hook_in` / `blocks.{i}.attn.v.hook_out`: V projection hooks (batch, pos, n_heads, d_head)
+
+#### Joint QKV Attention (GPT-2 style)
+For models using fused QKV projections (like GPT-2), the `JointQKVAttentionBridge` provides additional hooks:
+
+- `blocks.{i}.attn.qkv.hook_in`: input to QKV projection (batch, pos, d_model)
+- `blocks.{i}.attn.qkv.hook_out`: output from QKV projection (batch, pos, 3*d_model)
+- `blocks.{i}.attn.qkv.q_hook_in`: input to Q projection (batch, pos, d_model)
+- `blocks.{i}.attn.qkv.q_hook_out`: output from Q projection (batch, pos, n_heads, d_head)
+- `blocks.{i}.attn.qkv.k_hook_in`: input to K projection (batch, pos, d_model)
+- `blocks.{i}.attn.qkv.k_hook_out`: output from K projection (batch, pos, n_heads, d_head)
+- `blocks.{i}.attn.qkv.v_hook_in`: input to V projection (batch, pos, d_model)
+- `blocks.{i}.attn.qkv.v_hook_out`: output from V projection (batch, pos, n_heads, d_head)
+
+### MLP
+- `blocks.{i}.mlp.hook_in`: (batch, pos, d_model)
+  - *Legacy alias: `blocks.{i}.hook_mlp_in`*
+- `blocks.{i}.mlp.hook_pre`: (batch, pos, d_mlp)
+  - *Legacy alias: `blocks.{i}.hook_mlp_in` (via `mlp.in.hook_out`)*
+- `blocks.{i}.mlp.hook_out`: (batch, pos, d_model)
+  - *Legacy alias: `blocks.{i}.hook_mlp_out`*
+
+### Normalization
+- `blocks.{i}.ln1.hook_in` / `.hook_out`: (batch, pos, d_model)
+  - *Legacy aliases for `.hook_out`: `blocks.{i}.ln1.hook_normalized`, `blocks.{i}.ln1.hook_scale`*
+- Similarly for `ln2`
+  - *Legacy aliases for `.hook_out`: `blocks.{i}.ln2.hook_normalized`, `blocks.{i}.ln2.hook_scale`*
+
+### Unembedding / Logits
+- `unembed.hook_in`: (batch, pos, d_model)
+- `unembed.hook_out`: (batch, pos, d_vocab)
+
+## Shapes at a Glance
+
+- Residual stream and hidden states: (batch, pos, d_model)
+- Attention scores: (batch, n_heads, pos, pos)
+- Attention patterns: (n_heads, pos, pos) - after batch dimension removal
+- QKV projections: (batch, pos, n_heads, d_head)
+- MLP pre-activation: (batch, pos, d_mlp)
+- Embeddings: (batch, pos, d_model)
+- Unembedding logits: (batch, pos, d_vocab)
+- LayerNorm normalized / scale: (batch, pos, d_model)
+
+These shapes are exercised in the multi-model shape test: `tests/integration/test_hook_shape_compatibility.py`.
+
+## Booting from Hugging Face
+
+`TransformerBridge.boot_transformers(model_id, ...)`:
+- Loads the HF config/model/tokenizer
+- Selects the appropriate ArchitectureAdapter
+- Maps HF config fields to TransformerLens config (e.g., `d_model`, `n_heads`, `n_layers`, `d_mlp`, `d_vocab`, `n_ctx`, ...)
+- Constructs the bridge and registers all hook points
+
+## Fused QKV Attention
+
+Some architectures use a fused QKV projection (like GPT-2). The bridge's `JointQKVAttentionBridge` provides access to individual Q, K, V activations through the `QKVBridge` submodule. This allows for:
+
+1. **Individual Q/K/V hooking**: You can hook into `blocks.{i}.attn.qkv.q_hook_out`, `k_hook_out`, or `v_hook_out` to modify individual attention heads
+2. **Attention pattern creation**: The bridge automatically creates attention patterns from the attention scores and applies them through `hook_pattern`
+3. **Compatibility with legacy code**: Legacy hook names like `blocks.{i}.hook_v` are aliased to the appropriate QKV hooks
+
+The canonical attention hooks (`attn.hook_in/out`, `attn.hook_pattern`, etc.) retain the shapes listed above, while the QKV-specific hooks provide access to the individual attention components.
+
+## Aliases and Backwards Compatibility
+
+A minimal alias layer exists to ease migration from older TransformerLens names (e.g., `blocks.{i}.hook_resid_pre` → `blocks.{i}.hook_in`). New code should prefer the canonical names documented here.
+
+## Example: Caching and Inspecting Hooks
+
+```python
+prompt = "Hello world"
+logits, cache = bridge.run_with_cache(prompt)
+
+# List some attention-related hooks on the first block
+for k in cache.keys():
+    if k.startswith("blocks.0.attn"):
+        print(k, cache[k].shape)
+```
+
+For larger examples and a multi-model shape check, see `tests/integration/test_hook_shape_compatibility.py`.

tests/integration/test_hook_shape_compatibility.py

@@ -0,0 +1,238 @@
+import gc
+import os
+from typing import Iterable, Tuple
+
+import pytest
+import torch
+
+
+def _to_list(keys: Iterable[str]) -> list[str]:
+    return list(keys) if not isinstance(keys, list) else keys
+
+
+# Test models selected for architectural diversity while respecting memory constraints.
+# Note: Python doesn't release model memory between parameterized tests, causing accumulation.
+# Models are ordered by size to minimize peak memory usage.
+PUBLIC_HF_MODELS = [
+    "sshleifer/tiny-gpt2",
+    "EleutherAI/pythia-70m",
+    "facebook/opt-125m",
+]
+
+# Extended model list for authenticated CI runs
+FULL_HF_MODELS = [
+    "sshleifer/tiny-gpt2",
+    "EleutherAI/pythia-70m",
+    "roneneldan/TinyStories-33M",
+    "facebook/opt-125m",
+    "gpt2",
+]
+
+
+def _select_model_ids_from_acceptance_lists() -> list[str]:
+    return FULL_HF_MODELS if os.environ.get("HF_TOKEN", "") else PUBLIC_HF_MODELS
+
+
+# Allow overriding via env, comma-separated HF ids
+DEFAULT_IDS = ",".join(_select_model_ids_from_acceptance_lists())
+MODELS_ENV = os.getenv("TL_HOOK_SHAPE_MODELS", DEFAULT_IDS)
+MODEL_NAMES = [m.strip() for m in MODELS_ENV.split(",") if m.strip()]
+
+
+def _expected_shape_for_name(
+    name: str,
+    *,
+    batch: int,
+    pos: int,
+    d_model: int,
+    d_vocab: int | None,
+    n_heads: int | None,
+    d_head: int | None,
+    d_mlp: int | None,
+) -> Tuple[int, ...] | None:
+    # Canonical TransformerBridge hook names only (no legacy aliases)
+
+    # Unembedding (check before embedding to avoid matching "embed" in "unembed")
+    if name.endswith("unembed.hook_in"):
+        return (batch, pos, d_model)
+    if name.endswith("unembed.hook_out") and d_vocab is not None:
+        return (batch, pos, d_vocab)
+
+    # Embedding components
+    if name.endswith("embed.hook_in") or name.endswith("pos_embed.hook_in"):
+        return (batch, pos)
+    if name.endswith("embed.hook_out") or name.endswith("pos_embed.hook_out"):
+        return (batch, pos, d_model)
+
+    # Block IO
+    if ".hook_in" in name and ".attn." not in name and ".mlp." not in name and ".ln" not in name:
+        # blocks.{i}.hook_in
+        return (batch, pos, d_model)
+    if ".hook_out" in name and ".attn." not in name and ".mlp." not in name and ".ln" not in name:
+        # blocks.{i}.hook_out
+        return (batch, pos, d_model)
+
+    # Attention module (canonical TB names)
+    if name.endswith("attn.hook_in") or name.endswith("attn.hook_out"):
+        return (batch, pos, d_model)
+    if name.endswith("attn.hook_hidden_states"):
+        return (batch, pos, d_model)
+    if name.endswith("attn.hook_attention_weights") and n_heads is not None:
+        return (batch, n_heads, pos, pos)
+    if name.endswith("attn.hook_attn_scores") and n_heads is not None:
+        return (batch, n_heads, pos, pos)
+    if name.endswith("attn.hook_pattern") and n_heads is not None:
+        return (batch, n_heads, pos, pos)
+
+    # Attention subprojections: q/k/v/o
+    # Note: q/k/v hooks can be either:
+    # - (batch, pos, n_heads, d_head) for models with split heads (GPT-2, Pythia, etc.)
+    # - (batch, pos, d_model) for models without split heads (GPT-Neo, etc.)
+    # Both are valid depending on the architecture
+    if name.endswith("attn.o.hook_in"):
+        return (batch, pos, d_model)
+    if name.endswith("attn.o.hook_out"):
+        return (batch, pos, d_model)
+
+    # LayerNorms within blocks
+    if ".ln" in name and name.endswith("hook_in"):
+        return (batch, pos, d_model)
+    if ".ln" in name and name.endswith("hook_out"):
+        return (batch, pos, d_model)
+    if name.endswith("hook_normalized"):
+        return (batch, pos, d_model)
+    if name.endswith("hook_scale"):
+        # LayerNorm scale is (batch, pos, 1) for broadcasting
+        return (batch, pos, 1)
+
+    # MLP module
+    if name.endswith("mlp.hook_in") or name.endswith("mlp.hook_out"):
+        return (batch, pos, d_model)
+    if name.endswith("mlp.hook_pre") and d_mlp is not None:
+        return (batch, pos, d_mlp)
+    # MLP submodules: input and out projections
+    if name.endswith("mlp.input.hook_in") or name.endswith("mlp.out.hook_out"):
+        return (batch, pos, d_model)
+    if (
+        name.endswith("mlp.input.hook_out") or name.endswith("mlp.out.hook_in")
+    ) and d_mlp is not None:
+        return (batch, pos, d_mlp)
+
+    return None
+
+
+@pytest.mark.parametrize("model_name", MODEL_NAMES)
+def test_transformer_bridge_hook_shapes(model_name: str):
+    # Ensure boot method is registered
+    from transformer_lens.model_bridge.bridge import TransformerBridge
+    from transformer_lens.model_bridge.sources import (  # noqa: F401
+        transformers as bridge_sources,
+    )
+
+    bridge = TransformerBridge.boot_transformers(model_name, device="cpu")
+
+    prompt = "Hello world"
+    tokens = bridge.to_tokens(prompt, move_to_device=False)
+    batch, pos = int(tokens.shape[0]), int(tokens.shape[1])
+
+    cfg = bridge.cfg
+    d_model = int(getattr(cfg, "d_model"))
+    # Use actual vocab size from weights (may differ from config due to padding)
+    d_vocab = None
+    if hasattr(bridge, "unembed") and hasattr(bridge.unembed, "weight"):
+        d_vocab = int(bridge.unembed.weight.shape[0])
+    elif hasattr(cfg, "d_vocab"):
+        d_vocab = int(getattr(cfg, "d_vocab", 0))
+
+    n_heads = int(getattr(cfg, "n_heads", 0)) if hasattr(cfg, "n_heads") else None
+    d_head = int(getattr(cfg, "d_head", 0)) if hasattr(cfg, "d_head") else None
+    d_mlp = int(getattr(cfg, "d_mlp", 0)) if hasattr(cfg, "d_mlp") else None
+    if n_heads == 0:
+        n_heads = None
+    if d_head == 0:
+        d_head = None
+    if d_mlp == 0:
+        d_mlp = None
+
+    _, cache = bridge.run_with_cache(tokens, device="cpu")
+    keys = sorted(_to_list(cache.keys()))
+
+    # OPT reshapes to (batch*seq, d_model) internally for efficiency
+    is_opt_model = "opt" in model_name.lower()
+
+    mismatches: list[tuple[str, Tuple[int, ...], Tuple[int, ...]]] = []
+    checked = 0
+    for name in keys:
+        # Special handling for q/k/v hooks which can have two valid shapes
+        is_qkv_hook = any(
+            name.endswith(suf)
+            for suf in (
+                "attn.q.hook_in",
+                "attn.k.hook_in",
+                "attn.v.hook_in",
+                "attn.q.hook_out",
+                "attn.k.hook_out",
+                "attn.v.hook_out",
+            )
+        )
+
+        if is_qkv_hook:
+            tensor = cache[name]
+            assert isinstance(tensor, torch.Tensor), f"Non-tensor cached for {name}"
+            got = tuple(tensor.shape)
+            # Valid shapes: (batch, pos, n_heads, d_head) or (batch, pos, d_model)
+            valid_shapes = []
+            if n_heads is not None and d_head is not None:
+                valid_shapes.append((batch, pos, n_heads, d_head))
+            valid_shapes.append((batch, pos, d_model))
+
+            if got not in valid_shapes:
+                exp_str = " or ".join(str(s) for s in valid_shapes)
+                mismatches.append((name, exp_str, got))  # type: ignore
+            checked += 1
+            continue
+
+        # Rotary embeddings have architecture-specific partial dimensions
+        if "rotary" in name.lower():
+            checked += 1
+            continue
+
+        exp = _expected_shape_for_name(
+            name,
+            batch=batch,
+            pos=pos,
+            d_model=d_model,
+            d_vocab=d_vocab,
+            n_heads=n_heads,
+            d_head=d_head,
+            d_mlp=d_mlp,
+        )
+        if exp is None:
+            continue
+        tensor = cache[name]
+        assert isinstance(tensor, torch.Tensor), f"Non-tensor cached for {name}"
+        got = tuple(tensor.shape)
+
+        # OPT flattens batch and sequence dimensions for MLP/LayerNorm
+        if is_opt_model and got != exp:
+            is_flattened_hook = (".ln" in name and ".hook" in name and ".attn.ln" not in name) or (
+                ".mlp." in name and "hook" in name
+            )
+
+            if is_flattened_hook and len(exp) == 3 and len(got) == 2:
+                if got == (batch * pos, exp[2]):
+                    checked += 1
+                    continue
+
+        if got != exp:
+            mismatches.append((name, exp, got))
+        checked += 1
+
+    assert checked > 0, "No hooks were checked; update expected mapping or model filter"
+    msg = "\n".join(f"{n}: expected {e}, got {g}" for n, e, g in mismatches[:20])
+    assert not mismatches, f"Found {len(mismatches)} shape mismatches. Examples:\n{msg}"
+
+    # Clean up to reduce memory usage during parameterized test runs
+    del bridge, cache, tokens
+    torch.cuda.empty_cache() if torch.cuda.is_available() else None
+    gc.collect()