diff --git a/AGENTS.md b/AGENTS.md
index 10fbef1a2..12f81bdb8 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -93,6 +93,7 @@ Python: **>=3.10, <4.0**. CI tests 3.10, 3.11, 3.12. Format/type/docstring check
| [transformer_lens/hook_points.py](transformer_lens/hook_points.py) | `HookPoint` class and `LensHandle` |
| [transformer_lens/supported_models.py](transformer_lens/supported_models.py) | **HT-only** registry (`OFFICIAL_MODEL_NAMES`, `MODEL_ALIASES`) |
| [transformer_lens/tools/model_registry/](transformer_lens/tools/model_registry/) | Bridge-side registry + `verify_models.py` benchmark suite |
+| [transformer_lens/tools/analysis/](transformer_lens/tools/analysis/) | High-level single-call analyses over the cache (e.g. `direct_logit_attribution`); works with both HT and Bridge |
| [transformer_lens/patching.py](transformer_lens/patching.py), [evals.py](transformer_lens/evals.py) | Activation patching, IOI, ROME, etc. |
| [tests/unit/](tests/unit/), [tests/integration/](tests/integration/), [tests/acceptance/](tests/acceptance/), [tests/benchmarks/](tests/benchmarks/), [tests/mps/](tests/mps/) | Test tiers |
| [demos/](demos/) | Jupyter notebooks; a subset runs in CI under `nbval` with sanitization from [demos/doc_sanitize.cfg](demos/doc_sanitize.cfg) |
diff --git a/README.md b/README.md
index d909863c6..74a4c3258 100644
--- a/README.md
+++ b/README.md
@@ -202,7 +202,7 @@ exploratory research!
### Creator's Note (Neel Nanda)
-I (Neel Nanda) used to work for the [Anthropic interpretability team](transformer-circuits.pub), and
+I (Neel Nanda) used to work for the [Anthropic interpretability team](https://transformer-circuits.pub), and
I wrote this library because after I left and tried doing independent research, I got extremely
frustrated by the state of open source tooling. There's a lot of excellent infrastructure like
HuggingFace and DeepSpeed to _use_ or _train_ models, but very little to dig into their internals
diff --git a/demos/Exploratory_Analysis_Demo.ipynb b/demos/Exploratory_Analysis_Demo.ipynb
index 097f25c10..b0aa0c945 100644
--- a/demos/Exploratory_Analysis_Demo.ipynb
+++ b/demos/Exploratory_Analysis_Demo.ipynb
@@ -66,6 +66,7 @@
"metadata": {},
"outputs": [],
"source": [
+ "import os\n",
"\n",
"# Detect if we're running in Google Colab\n",
"try:\n",
@@ -79,8 +80,6 @@
"if IN_COLAB:\n",
" %pip install transformer_lens\n",
" %pip install circuitsvis\n",
- " # Install a faster Node version\n",
- " !curl -fsSL https://deb.nodesource.com/setup_16.x | sudo -E bash -; sudo apt-get install -y nodejs # noqa\n",
"\n",
"# Hot reload in development mode & not running on the CD\n",
"if not IN_COLAB:\n",
@@ -88,7 +87,9 @@
" ip = get_ipython()\n",
" if not ip.extension_manager.loaded:\n",
" ip.extension_manager.load('autoreload')\n",
- " %autoreload 2\n"
+ " %autoreload 2\n",
+ "\n",
+ "IN_GITHUB = os.getenv(\"GITHUB_ACTIONS\") == \"true\""
]
},
{
@@ -166,13 +167,16 @@
"metadata": {},
"outputs": [],
"source": [
- "def imshow(tensor, **kwargs):\n",
- " px.imshow(\n",
+ "def imshow(tensor, show=True, **kwargs):\n",
+ " fig = px.imshow(\n",
" utils.to_numpy(tensor),\n",
" color_continuous_midpoint=0.0,\n",
" color_continuous_scale=\"RdBu\",\n",
" **kwargs,\n",
- " ).show()\n",
+ " )\n",
+ " if show:\n",
+ " fig.show()\n",
+ " return fig\n",
"\n",
"\n",
"def line(tensor, **kwargs):\n",
@@ -255,6 +259,7 @@
" fold_ln=True,\n",
" refactor_factored_attn_matrices=True,\n",
")\n",
+ "model.set_use_attn_result(True)\n",
"\n",
"# Get the default device used\n",
"device: torch.device = utils.get_device()"
@@ -1389,6 +1394,230 @@
")"
]
},
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Direct Path Patching\n",
+ "\n",
+ "Decomposing attention heads helped us understand whether their outputs are affected more by their computed attention pattern or values, but we've looked at each head in isolation. We can try taking this one step further by attempting to find how one attention head's output affects another's attention pattern and values. This could start revealing the circuitry of how information is moved across layers!\n",
+ "\n",
+ "One mechanism we can use is Direct Path Patching. Path patching is like activation patching, but rather than patching an activation we patch the *effect* of one activation on another, later activation. *Direct* path patching, specifically, means patching the linear part of that effect, i.e., the part that passes through the residual stream from the first activation to the later one, without going through attention heads or MLPs (see this [ARENA notebook](https://colab.research.google.com/drive/1KgrEwvCKdX-8DQ1uSiIuxwIiwzJuQ3Gw#scrollTo=b32Qdk-Gl6mU) for an alternative definition which does include MLPs).\n",
+ "\n",
+ "For example, to do direct path patching from the output of attention head `A` to the query of attention head `B` in a later layer `L`, we'd add a hook saying:\n",
+ "```py\n",
+ "patched_B_query = corrupted_B_query + (clean_A_output - corrupted_A_output) @ W_Q(B) / corrupted_ln1_scale(L)\n",
+ "```\n",
+ "\n",
+ "\n",
+ "Why the corrupted ln1?
\n",
+ "Since all earlier components in the run are corrupt, we apply the cached layer norm from the corrupted run rather than that of the clean one to get a better approximation of the linear effect (see \"Ignoring LayerNorm\" above for why dividing by the cached LayerNorm makes sense in the first place).\n",
+ " \n",
+ "
\n",
+ "\n",
+ "We'll look at the effects of direct path patching from outputs of attention heads to query, key and value vectors of following attention heads. As before, we patch the effects across all token positions, although it's possible to further zoom in to specific token positions (notably the last one and the second subject one).\n",
+ "\n",
+ "\n",
+ "More on token positions
\n",
+ "Note that direct path patching can only meaningfully be done for activations in the same token position; any path between two activations in two different token positions must go through attention, hence it is not a direct path. Still, patching a direct path from the output of an attention head in a specific token position to either the key or value vectors of a later attention head attending to the same token position could affect that head's output in other querying token positions.\n",
+ " \n",
+ "
\n",
+ "\n",
+ "First, the implementation:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def patch_direct_path_A_output_to_B_attn_act(\n",
+ " corrupted_attn_act: Float[torch.Tensor, \"batch pos head_index d_head\"],\n",
+ " hook,\n",
+ " attn_act_type,\n",
+ " A_layer,\n",
+ " A_head_index,\n",
+ " B_layer,\n",
+ " B_head_index,\n",
+ " clean_cache,\n",
+ " corrupted_cache,\n",
+ "):\n",
+ " A_output_act_name = utils.get_act_name(\"result\", A_layer, \"attn\")\n",
+ " clean_A_output = clean_cache[A_output_act_name][:, :, A_head_index, :]\n",
+ " corrupted_A_output = corrupted_cache[A_output_act_name][:, :, A_head_index, :]\n",
+ " corrupted_B_act = corrupted_attn_act[:, :, B_head_index, :]\n",
+ " W_B = _get_attention_weights(model.blocks[B_layer].attn, attn_act_type)[B_head_index]\n",
+ " corrupted_ln1_scale = corrupted_cache[utils.get_act_name(\"scale\", B_layer, \"ln1\")]\n",
+ "\n",
+ " patched_B_act = corrupted_B_act + (clean_A_output - corrupted_A_output) @ W_B / corrupted_ln1_scale\n",
+ "\n",
+ " patched_attn_act = corrupted_attn_act.clone()\n",
+ " patched_attn_act[:, :, B_head_index, :] = patched_B_act\n",
+ " return patched_attn_act\n",
+ "\n",
+ "\n",
+ "def direct_path_output_to_attn_act_diffs(output_layer, attn_act_layer, attn_act_type):\n",
+ " A_to_B_diffs = torch.zeros(model.cfg.n_heads, model.cfg.n_heads, device=device, dtype=torch.float32)\n",
+ " for A in range(model.cfg.n_heads):\n",
+ " for B in range(model.cfg.n_heads):\n",
+ " patch_hook = (\n",
+ " utils.get_act_name(attn_act_type, attn_act_layer, \"attn\"),\n",
+ " partial(\n",
+ " patch_direct_path_A_output_to_B_attn_act,\n",
+ " attn_act_type=attn_act_type,\n",
+ " A_layer=output_layer,\n",
+ " A_head_index=A,\n",
+ " B_layer=attn_act_layer,\n",
+ " B_head_index=B,\n",
+ " clean_cache=cache,\n",
+ " corrupted_cache=corrupted_cache,\n",
+ " ),\n",
+ " )\n",
+ " # It would have been great to set start_layer=B and the inputs to the cached corrupted residual stream at that layer;\n",
+ " # unfortunately, this is not supported by TransformerBridge. Legacy HookedTransformer supports it but is slower, so\n",
+ " # it doesn't pay off to switch.\n",
+ " patched_logits = model.run_with_hooks(\n",
+ " corrupted_tokens,\n",
+ " fwd_hooks=[patch_hook],\n",
+ " return_type=\"logits\",\n",
+ " )\n",
+ " patched_logit_diff = logits_to_ave_logit_diff(patched_logits, answer_tokens)\n",
+ " A_to_B_diffs[A, B] = normalize_patched_logit_diff(patched_logit_diff)\n",
+ " return A_to_B_diffs\n",
+ "\n",
+ "\n",
+ "def _get_attention_weights(attn, attn_act_type):\n",
+ " if attn_act_type == \"q\":\n",
+ " return attn.W_Q\n",
+ " if attn_act_type == \"k\":\n",
+ " return attn.W_K\n",
+ " if attn_act_type == \"v\":\n",
+ " return attn.W_V\n",
+ " raise ValueError"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def show_direct_path_output_to_attn_act_diffs(diffs, output_layer, attn_act_layer, attn_act_type, zrange=0.3):\n",
+ " attn_act_name = _get_attention_act_type_name(attn_act_type)\n",
+ " fig = imshow(\n",
+ " diffs,\n",
+ " show=False,\n",
+ " title=f\"Direct Path Patch Logit Difference, L{output_layer} Output to L{attn_act_layer} {attn_act_name}\",\n",
+ " labels={\"x\": f\"{attn_act_name} Head\", \"y\": \"Output Head\"},\n",
+ " zmin=-zrange,\n",
+ " zmax=zrange,\n",
+ " )\n",
+ " fig.update_traces(\n",
+ " hovertemplate=f\"L{output_layer}H%{{y}} -> L{attn_act_layer}H%{{x}}
%{{z:.3f}}\"\n",
+ " )\n",
+ " fig.show()\n",
+ "\n",
+ "\n",
+ "def _get_attention_act_type_name(attn_act_type):\n",
+ " if attn_act_type == \"q\":\n",
+ " return \"Query\"\n",
+ " if attn_act_type == \"k\":\n",
+ " return \"Key\"\n",
+ " if attn_act_type == \"v\":\n",
+ " return \"Value\"\n",
+ " raise ValueError"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Testing that it works:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def test_direct_path_patching():\n",
+ " output_layer, value_layer = (5, 8)\n",
+ " diffs = direct_path_output_to_attn_act_diffs(output_layer, value_layer, \"v\")\n",
+ " show_direct_path_output_to_attn_act_diffs(diffs, output_layer, value_layer, \"v\")\n",
+ "\n",
+ "\n",
+ "test_direct_path_patching()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "We can see that L5H5's output directly affects L8H6's attention value vector. We'd previously seen that the output of either head is significant to correct prediction, but now we can further conclude that they perform (at least some of) their function together rather than in isolation!\n",
+ "\n",
+ "Moving on to queries: The previous section showed how patching the attention patterns of late-layer heads affects their outputs significantly. We could hypothesize that the computed query vectors of these heads are different in the corrupted vs. clean runs because of information that the mid-layer heads are generating. If that's the case and the effect is direct (passes through the residual stream), then direct path patching might find it! The time to iterate pairs of layers increases quadratically with the number of layers, so with this hypothesis in mind we can focus on just the mid-late layers range, and we can also limit looking up to two layers ahead."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "if not IN_GITHUB:\n",
+ " for output_layer in range(7, model.cfg.n_layers - 1):\n",
+ " for query_layer in range(output_layer + 1, min(output_layer + 3, model.cfg.n_layers)):\n",
+ " diffs = direct_path_output_to_attn_act_diffs(output_layer, query_layer, \"q\")\n",
+ " show_direct_path_output_to_attn_act_diffs(diffs, output_layer, query_layer, \"q\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Two diffs that stand out are that patching L8H6's output direct path to L9H9's query improves next token prediction, whereas patching L9H9's output direct paths to both L10H7's and L11H10's queries worsens it. Previously, we'd already seen that patching L8H6 and L9H9 outputs improves prediction and that patching L10H7 and L11H10 outputs worsens it, and that the computed attention patterns for L9H9, L10H7 and L11H10 play a significant role; we can now further conclude that L8H6's output directly affects L9H9's attention pattern significantly, and that L9H9's output directly (negatively) affects L10H7's and L11H10's attention patterns significantly, via the computed query vectors.\n",
+ "\n",
+ "What about keys? We can guess that outputs of earlier attention heads set up keys on token positions to determine what later heads attend to:"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "if not IN_GITHUB:\n",
+ " for output_layer in range(5, 8):\n",
+ " for key_layer in range(8, model.cfg.n_layers):\n",
+ " diffs = direct_path_output_to_attn_act_diffs(output_layer, key_layer, \"k\")\n",
+ " show_direct_path_output_to_attn_act_diffs(diffs, output_layer, key_layer, \"k\", zrange=0.03)"
+ ]
+ },
+ {
+ "attachments": {},
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "We do find direct effects of outputs on keys, albeit at a much lower scale than on queries and values. The most significant ones are from L6H9 and L7H1 on L8H6.\n",
+ "\n",
+ "Note that it's possible stronger effects of outputs on keys exist which are not revealed, even on the layers that we did test: there might be some indirect effects, or possibly additional direct effects which are masked by some other model behavior, e.g., a later corrupted component could remove such an effect.\n",
+ "\n",
+ "Combining the results of the above experiments, we get an outline of some circuitry:\n",
+ "\n",
+ "```\n",
+ " 5.5 6.9,7.1\n",
+ " | |\n",
+ " v | | k\n",
+ " \\ /\n",
+ " 8.6 --> 9.9 --> 10.7,11.10\n",
+ " q q\n",
+ "```\n",
+ "\n",
+ "cf. the paper diagram below."
+ ]
+ },
{
"attachments": {},
"cell_type": "markdown",
@@ -1471,7 +1700,7 @@
"Breaking down their categories:\n",
"\n",
"* Early: The duplicate token heads, previous token heads and induction heads. These serve the purpose of detecting that the second subject is duplicated and which earlier name is the duplicate.\n",
- " * We found a direct duplicate token head which behaves exactly as expected, L3H0. Heads L5H0 and L6H9 are induction heads, which explains why they don't attend directly to the earlier copy of John!\n",
+ " * We found a direct duplicate token head which behaves exactly as expected, L3H0. Heads L5H5 and L6H9 are induction heads, which explains why they don't attend directly to the earlier copy of John!\n",
" * Note that the duplicate token heads and induction heads do not compose with each other - both directly add to the S-Inhibition heads. The diagram is somewhat misleading.\n",
"* Middle: They call these S-Inhibition heads - they copy the information about the duplicate token from the second subject to the to token, and their output is used to *inhibit* the attention paid from the name movers to the first subject copy. We found all these heads, and had a decent guess for what they did.\n",
" * In either case they attend to the second subject, so the patch that mattered was their value vectors!\n",
diff --git a/demos/Grokking_Demo.ipynb b/demos/Grokking_Demo.ipynb
index 2fab2d43d..2566cca49 100644
--- a/demos/Grokking_Demo.ipynb
+++ b/demos/Grokking_Demo.ipynb
@@ -3492,22 +3492,9 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
- "outputs": [
- {
- "ename": "RuntimeError",
- "evalue": "Size does not match at dimension 0 expected index [12769, 1] to be smaller than self [113, 113] apart from dimension 1",
- "output_type": "error",
- "traceback": [
- "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
- "\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)",
- "\u001b[0;32m/tmp/ipykernel_1215793/3004607503.py\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mall_logits\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabels\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
- "\u001b[0;32m/tmp/ipykernel_1215793/4096650173.py\u001b[0m in \u001b[0;36mloss_fn\u001b[0;34m(logits, labels)\u001b[0m\n\u001b[1;32m 4\u001b[0m \u001b[0mlogits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlogits\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat64\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 5\u001b[0m \u001b[0mlog_probs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlogits\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlog_softmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 6\u001b[0;31m \u001b[0mcorrect_log_probs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlog_probs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgather\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindex\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mlabels\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 7\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0;34m-\u001b[0m\u001b[0mcorrect_log_probs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmean\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 8\u001b[0m \u001b[0mtrain_logits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrain_data\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
- "\u001b[0;31mRuntimeError\u001b[0m: Size does not match at dimension 0 expected index [12769, 1] to be smaller than self [113, 113] apart from dimension 1"
- ]
- }
- ],
+ "outputs": [],
"source": [
- "print(loss_fn(all_logits, labels)) # This bugged on models not fully trained "
+ "print(loss_fn(original_logits, labels))"
]
},
{
diff --git a/demos/Main_Demo.ipynb b/demos/Main_Demo.ipynb
index 70ffabbea..87c01d86d 100644
--- a/demos/Main_Demo.ipynb
+++ b/demos/Main_Demo.ipynb
@@ -18,7 +18,7 @@
"\n",
"To use this notebook, go to Runtime > Change Runtime Type and select GPU as the hardware accelerator.\n",
"\n",
- "This is a reference notebook covering the main features of the [TransformerLens](https://github.com/TransformerLensOrg/TransformerLens) library for mechanistic interpretability. See [Callum McDougall's tutorial](https://transformerlens-intro.streamlit.app/TransformerLens_&_induction_circuits) for a more structured and gentler introduction to the library"
+ "This is a reference notebook covering the main features of the [TransformerLens](https://github.com/TransformerLensOrg/TransformerLens) library for mechanistic interpretability. See this [ARENA chapter](https://learn.arena.education/chapter1_transformer_interp/) for a more structured and gentler introduction to the library."
]
},
{
@@ -65,8 +65,6 @@
"if IN_COLAB:\n",
" %pip install transformer_lens\n",
" %pip install circuitsvis\n",
- " # Install a faster Node version\n",
- " !curl -fsSL https://deb.nodesource.com/setup_16.x | sudo -E bash -; sudo apt-get install -y nodejs # noqa\n",
"\n",
"# Hot reload in development mode & not running on the CD\n",
"if not IN_COLAB:\n",
@@ -181,7 +179,7 @@
"\n",
"\n",
"def line(tensor, renderer=None, xaxis=\"\", yaxis=\"\", **kwargs):\n",
- " px.line(utils.to_numpy(tensor), labels={\"x\": xaxis, \"y\": yaxis}, **kwargs).show(renderer)\n",
+ " px.line(y=utils.to_numpy(tensor), labels={\"x\": xaxis, \"y\": yaxis}, **kwargs).show(renderer)\n",
"\n",
"\n",
"def scatter(x, y, xaxis=\"\", yaxis=\"\", caxis=\"\", renderer=None, **kwargs):\n",
@@ -217,7 +215,7 @@
"source": [
"## Loading and Running Models\n",
"\n",
- "TransformerLens comes loaded with >40 open source GPT-style models. You can load any of them in with `HookedTransformer.from_pretrained(MODEL_NAME)`. For this demo notebook we'll look at GPT-2 Small, an 80M parameter model, see the Available Models section for info on the rest."
+ "TransformerLens supports 9,000+ models across 50+ architecture families. You can load any of them in with `TransformerBridge.boot_transformers(MODEL_NAME)`. For this demo notebook we'll look at GPT-2 Small, an 80M parameter model, see the Available Models section for info on the rest."
]
},
{
@@ -236,8 +234,8 @@
"outputs": [],
"source": [
"# NBVAL_IGNORE_OUTPUT\n",
- "model = TransformerBridge.boot_transformers(\"gpt2\", device=device)\n",
- "model.enable_compatibility_mode(disable_warnings=True)"
+ "model = TransformerBridge.boot_transformers(\"openai-community/gpt2\", device=device)\n",
+ "model.enable_compatibility_mode(disable_warnings=True) # for legacy HookedTransformer-equivalent numerics"
]
},
{
@@ -245,11 +243,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "To try the model out, let's find the loss on this text! Models can be run on a single string or a tensor of tokens (shape: [batch, position], all integers), and the possible return types are: \n",
- "* \"logits\" (shape [batch, position, d_vocab], floats), \n",
- "* \"loss\" (the cross-entropy loss when predicting the next token), \n",
- "* \"both\" (a tuple of (logits, loss)) \n",
- "* None (run the model, but don't calculate the logits - this is faster when we only want to use intermediate activations)"
+ "To try the model out, let's find the loss on this text!"
]
},
{
@@ -258,15 +252,27 @@
"metadata": {},
"outputs": [],
"source": [
- "model_description_text = \"\"\"## Loading Models\n",
+ "model_description_text = \"\"\"## Loading and Running Models\n",
"\n",
- "HookedTransformer comes loaded with >40 open source GPT-style models. You can load any of them in with `HookedTransformer.from_pretrained(MODEL_NAME)`. See my explainer for documentation of all supported models, and this table for hyper-parameters and the name used to load them. Each model is loaded into the consistent HookedTransformer architecture, designed to be clean, consistent and interpretability-friendly. \n",
+ "TransformerLens supports 9,000+ models across 50+ architecture families. You can load any of them in with `TransformerBridge.boot_transformers(MODEL_NAME)`. For this demo notebook we'll look at GPT-2 Small, an 80M parameter model, see the Available Models section for info on the rest.\n",
"\n",
- "For this demo notebook we'll look at GPT-2 Small, an 80M parameter model. To try the model the model out, let's find the loss on this paragraph!\"\"\"\n",
+ "To try the model out, let's find the loss on this text!\"\"\"\n",
"loss = model(model_description_text, return_type=\"loss\")\n",
"print(\"Model loss:\", loss)"
]
},
+ {
+ "attachments": {},
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Models can be run on a single string or a tensor of tokens (shape: [batch, position], all integers), and the possible return types are: \n",
+ "* \"logits\" (shape [batch, position, d_vocab], floats), \n",
+ "* \"loss\" (the cross-entropy loss when predicting the next token), \n",
+ "* \"both\" (a tuple of (logits, loss)) \n",
+ "* None (run the model, but don't calculate the logits - this is faster when we only want to use intermediate activations)"
+ ]
+ },
{
"attachments": {},
"cell_type": "markdown",
@@ -279,7 +285,7 @@
"On `remove_batch_dim`
\n",
"\n",
"Every activation inside the model begins with a batch dimension. Here, because we only entered a single batch dimension, that dimension is always length 1 and kinda annoying, so passing in the `remove_batch_dim=True` keyword removes it. `gpt2_cache_no_batch_dim = gpt2_cache.remove_batch_dim()` would have achieved the same effect.\n",
- ""
+ "Technical details
\n",
"\n",
"* We attach the hook to the attention pattern activation. There's one big pattern activation per layer, stacked across all heads, so we need to do some tensor manipulation to get a per-head score. \n",
- "* Hook functions can access global state, so we make a big tensor to store the induction head score for each head, and then we just add the score for each head to the appropriate position in the tensor. \n",
+ "* Hook functions can access global state, so we make a big tensor to store the induction head score for each head, and then we just store the score for each head in the appropriate position in the tensor. \n",
"* To get a single hook function that works for each layer, we use the `hook.layer()` method to get the layer index (internally this is just inferred from the hook names).\n",
"* As we want to add this to *every* activation pattern hook point, rather than giving the string for an activation name, this time we give a **name filter**. This is a Boolean function on hook point names, and it adds the hook function to every hook point where the function evaluates as true. \n",
" * `run_with_hooks` allows us to enter a list of (act_name, hook_function) pairs to all be added at once, so we could also have done this by inputting a list with a hook for each layer.\n",
@@ -712,11 +711,9 @@
" hook: HookPoint,\n",
"):\n",
" display(\n",
- " cv.attention.attention_patterns(\n",
+ " cv.attention.attention_pattern(\n",
" tokens=model.to_str_tokens(repeated_random_sequence),\n",
- " attention=pattern[0, induction_head_index, :, :][\n",
- " None, :, :\n",
- " ], # Add a dummy axis, as CircuitsVis expects 3D patterns.\n",
+ " attention=pattern[0, induction_head_index, :, :],\n",
" )\n",
" )\n",
"\n",
@@ -740,9 +737,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "TransformerLens comes with over 40 open source models available, all of which can be loaded into a consistent(-ish) architecture by just changing the name in `from_pretrained`. The open source models available are [documented here](https://dynalist.io/d/n2ZWtnoYHrU1s4vnFSAQ519J#z=jHj79Pj58cgJKdq4t-ygK-4h), and a set of interpretability friendly models I've trained are [documented here](https://dynalist.io/d/n2ZWtnoYHrU1s4vnFSAQ519J#z=NCJ6zH_Okw_mUYAwGnMKsj2m), including a set of toy language models (tiny one to four layer models) and a set of [SoLU models](https://dynalist.io/d/n2ZWtnoYHrU1s4vnFSAQ519J#z=FZ5W6GGcy6OitPEaO733JLqf) up to GPT-2 Medium size (300M parameters). You can see [a table of the official alias and hyper-parameters of available models here](https://github.com/TransformerLensOrg/TransformerLens/blob/main/transformer_lens/model_properties_table.md).\n",
- "\n",
- "**Note:** TransformerLens does not currently support multi-GPU models (which you want for models above eg 7B parameters), but this feature is coming soon!"
+ "TransformerLens supports 9,000+ models across 50+ architecture families, all of which can be loaded into a consistent(-ish) interface by just changing the name in `boot_transformers`. The available models are [documented here](https://transformerlensorg.github.io/TransformerLens/generated/transformer_bridge_models.html) with some notable ones [documented here](https://dynalist.io/d/n2ZWtnoYHrU1s4vnFSAQ519J#z=jHj79Pj58cgJKdq4t-ygK-4h), and a set of interpretability friendly models I've trained are [documented here](https://dynalist.io/d/n2ZWtnoYHrU1s4vnFSAQ519J#z=NCJ6zH_Okw_mUYAwGnMKsj2m), including a set of toy language models (tiny one to four layer models) and a set of [SoLU models](https://dynalist.io/d/n2ZWtnoYHrU1s4vnFSAQ519J#z=FZ5W6GGcy6OitPEaO733JLqf) up to GPT-2 Medium size (300M parameters)."
]
},
{
@@ -814,22 +809,22 @@
"metadata": {},
"source": [
"\n",
- "### An overview of the important open source models in the library\n",
+ "### An overview of the important open source models in the library (last updated 2023)\n",
"\n",
"* **GPT-2** - the classic generative pre-trained models from OpenAI\n",
" * Sizes Small (85M), Medium (300M), Large (700M) and XL (1.5B).\n",
- " * Trained on ~22B tokens of internet text. ([Open source replication](https://huggingface.co/datasets/openwebtext))\n",
+ " * Trained on ~22B tokens of internet text. ([Open source replication](https://huggingface.co/datasets/Skylion007/openwebtext))\n",
"* **GPT-Neo** - Eleuther's replication of GPT-2\n",
" * Sizes 125M, 1.3B, 2.7B\n",
" * Trained on 300B(ish?) tokens of [the Pile](https://pile.eleuther.ai/) a large and diverse dataset including a bunch of code (and weird stuff)\n",
- "* **[OPT](https://ai.facebook.com/blog/democratizing-access-to-large-scale-language-models-with-opt-175b/)** - Meta AI's series of open source models\n",
+ "* **[OPT](https://ai.meta.com/blog/democratizing-access-to-large-scale-language-models-with-opt-175b/)** - Meta AI's series of open source models\n",
" * Trained on 180B tokens of diverse text.\n",
" * 125M, 1.3B, 2.7B, 6.7B, 13B, 30B, 66B\n",
"* **GPT-J** - Eleuther's 6B parameter model, trained on the Pile\n",
"* **GPT-NeoX** - Eleuther's 20B parameter model, trained on the Pile\n",
"* **StableLM** - Stability AI's 3B and 7B models, with and without chat and instruction fine-tuning\n",
"* **Stanford CRFM models** - a replication of GPT-2 Small and GPT-2 Medium, trained on 5 different random seeds.\n",
- " * Notably, 600 checkpoints were taken during training per model, and these are available in the library with eg `HookedTransformer.from_pretrained(\"stanford-gpt2-small-a\", checkpoint_index=265)`.\n",
+ " * Notably, 600 checkpoints were taken during training per model, and these are available in the library with eg `TransformerBridge.boot_transformers(\"stanford-crfm/alias-gpt2-small-x21\", checkpoint_index=265)`.\n",
"- **BERT** - Google's bidirectional encoder-only transformer.\n",
" - Size Base (108M), trained on English Wikipedia and BooksCorpus.\n",
" \n",
@@ -845,7 +840,7 @@
"\n",
"(Feel free to [reach out](mailto:neelnanda27@gmail.com) if you want more details on any of these models)\n",
"\n",
- "Each of these models has about ~200 checkpoints taken during training that can also be loaded from TransformerLens, with the `checkpoint_index` argument to `from_pretrained`.\n",
+ "Each of these models has about ~200 checkpoints taken during training that can also be loaded from TransformerLens, with the `checkpoint_index` argument to `boot_transformers`.\n",
"\n",
"Note that all models are trained with a Beginning of Sequence token, and will likely break if given inputs without that! \n",
"\n",
@@ -877,7 +872,7 @@
" * **[Exploratory Analysis Demo](https://neelnanda.io/exploratory-analysis-demo)**, a demonstration of my standard toolkit for how to use TransformerLens to explore a mysterious behaviour in a language model.\n",
" * [Interpretability in the Wild](https://github.com/redwoodresearch/Easy-Transformer) a codebase from Arthur Conmy and Alex Variengien at Redwood research using this library to do a detailed and rigorous reverse engineering of the Indirect Object Identification circuit, to accompany their paper\n",
" * Note - this was based on an earlier version of this library, called EasyTransformer. It's pretty similar, but several breaking changes have been made since. \n",
- " * A [recorded walkthrough](https://www.youtube.com/watch?v=yo4QvDn-vsU) of me doing research with TransformerLens on whether a tiny model can re-derive positional information, with [an accompanying Colab](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/No_Position_Experiment.ipynb)\n",
+ " * A [recorded walkthrough](https://www.youtube.com/watch?v=yo4QvDn-vsU) of me doing research with TransformerLens on whether a tiny model can re-derive positional information, with [an accompanying Colab](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/No_Position_Experiment.ipynb)\n",
"* [Neuroscope](https://neuroscope.io), a website showing the text in the dataset that most activates each neuron in some selected models. Good to explore to get a sense for what kind of features the model tends to represent, and as a \"wiki\" to get some info\n",
" * A tutorial on how to make an [Interactive Neuroscope](https://github.com/TransformerLensOrg/TransformerLens/blob/main/Hacky-Interactive-Lexoscope.ipynb), where you type in text and see the neuron activations over the text update live."
]
@@ -886,7 +881,9 @@
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Transformer architecture\n",
+ "## HookedTransformer architecture\n",
+ "\n",
+ "**Note:** HookedTransformer is deprecated as of TransformerLens 3.0.\n",
"\n",
"HookedTransformer is a somewhat adapted GPT-2 architecture, but is computationally identical. The most significant changes are to the internal structure of the attention heads: \n",
"* The weights (W_K, W_Q, W_V) mapping the residual stream to queries, keys and values are 3 separate matrices, rather than big concatenated one.\n",
@@ -959,9 +956,9 @@
"source": [
"### Activation + Hook Names\n",
"\n",
- "Lets get out a list of the activation/hook names in the model and their shapes. In practice, I recommend using the `utils.get_act_name` function to get the names, but this is a useful fallback, and necessary to eg write a name filter function.\n",
+ "Let's get out a list of the activation/hook names in the model and their shapes. In practice, I recommend using the `utils.get_act_name` function to get the names, but this is a useful fallback, and necessary to eg write a name filter function.\n",
"\n",
- "Let's do this by entering in a short, 10 token prompt, and add a hook function to each activations to print its name and shape. To avoid spam, let's just add this to activations in the first block or not in a block.\n",
+ "Let's do this by entering in a short, 10 token prompt, and add a hook function to each activation to print its name and shape. To avoid spam, let's just add this to activations in the first block or not in a block.\n",
"\n",
"Note 1: Each LayerNorm has a hook for the scale factor (ie the standard deviation of the input activations for each token position & batch element) and for the normalized output (ie the input activation with mean 0 and standard deviation 1, but *before* applying scaling or translating with learned weights). LayerNorm is applied every time a layer reads from the residual stream: `ln1` is the LayerNorm before the attention layer in a block, `ln2` the one before the MLP layer, and `ln_final` is the LayerNorm before the unembed. \n",
"\n",
@@ -1084,7 +1081,7 @@
"source": [
"# Features\n",
"\n",
- "An overview of some other important features of the library. I recommend checking out the [Exploratory Analysis Demo](https://colab.research.google.com/github/TransformerLensOrg/Easy-Transformer/blob/main/Exploratory_Analysis_Demo.ipynb) for some other important features not mentioned here, and for a demo of what using the library in practice looks like."
+ "An overview of some other important features of the library. I recommend checking out the [Exploratory Analysis Demo](https://colab.research.google.com/github/TransformerLensOrg/TransformerLens/blob/main/demos/Exploratory_Analysis_Demo.ipynb) for some other important features not mentioned here, and for a demo of what using the library in practice looks like."
]
},
{
@@ -1096,22 +1093,14 @@
"\n",
"**Tokenization** is one of the most annoying features of studying language models. We want language models to be able to take in arbitrary text as input, but the transformer architecture needs the inputs to be elements of a fixed, finite vocabulary. The solution to this is **tokens**, a fixed vocabulary of \"sub-words\", that any natural language can be broken down into with a **tokenizer**. This is invertible, and we can recover the original text, called **de-tokenization**. \n",
"\n",
- "TransformerLens comes with a range of utility functions to deal with tokenization. Different models can have different tokenizers, so these are all methods on the model.\n",
- "\n",
- "get_token_position, to_tokens, to_string, to_str_tokens, prepend_bos, to_single_token"
+ "TransformerLens comes with a range of utility functions to deal with tokenization. Different models can have different tokenizers, so these are all methods on the model: `get_token_position`, `to_tokens`, `to_string`, `to_str_tokens`, `prepend_bos`, `to_single_token`"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "The first thing you need to figure out is *how* things are tokenized. `model.to_str_tokens` splits a string into the tokens *as a list of substrings*, and so lets you explore what the text looks like. To demonstrate this, let's use it on this paragraph.\n",
- "\n",
- "Some observations - there are a lot of arbitrary-ish details in here!\n",
- "* The tokenizer splits on spaces, so no token contains two words.\n",
- "* Tokens include the preceding space, and whether the first token is a capital letter. `how` and ` how` are different tokens!\n",
- "* Common words are single tokens, even if fairly long (` paragraph`) while uncommon words are split into multiple tokens (` token|ized`).\n",
- "* Tokens *mostly* split on punctuation characters (eg `*` and `.`), but eg `'s` is a single token."
+ "The first thing you need to figure out is *how* things are tokenized. `model.to_str_tokens` splits a string into the tokens *as a list of substrings*, and so lets you explore what the text looks like. To demonstrate this, let's use it on this paragraph."
]
},
{
@@ -1125,6 +1114,17 @@
"print(example_text_str_tokens)"
]
},
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Some observations - there are a lot of arbitrary-ish details in here!\n",
+ "* The tokenizer splits on spaces, so no token contains two words.\n",
+ "* Tokens include the preceding space, and whether the first token is a capital letter. `|how|` and `| how|` are different tokens!\n",
+ "* Common words are single tokens, even if fairly long (` paragraph`) while uncommon words are split into multiple tokens (` token|ized`).\n",
+ "* Tokens *mostly* split on punctuation characters (eg `*` and `.`), but eg `'s` is a single token."
+ ]
+ },
{
"cell_type": "markdown",
"metadata": {},
@@ -1322,8 +1322,6 @@
"\n",
"Further, *some* models are trained to need a BOS token (OPT and my interpretability-friendly models are, GPT-2 and GPT-Neo are not). But despite GPT-2 not being trained with this, empirically it seems to make interpretability easier.\n",
"\n",
- "(However, if you want to change the default behaviour to *not* prepending a BOS token, pass `default_prepend_bos=False` when you instantiate the model, e.g., `model = HookedTransformer.from_pretrained('gpt2', default_prepend_bos=False)`.)\n",
- "\n",
"For example, the model can get much worse at Indirect Object Identification without a BOS (and with a name as the first token):"
]
},
@@ -1616,7 +1614,7 @@
" utils.to_numpy(full_OV_copying_score),\n",
" xaxis=\"Head\",\n",
" yaxis=\"Layer\",\n",
- " title=\"OV Copying Score for each head in GPT-2 Small\",\n",
+ " title=\"Full OV Copying Score for each head in GPT-2 Small\",\n",
" zmax=1.0,\n",
" zmin=-1.0,\n",
")"
@@ -1645,19 +1643,6 @@
")"
]
},
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "print(f\"Token 256 - the most common pair of ASCII characters: |{model.to_string(256)}|\")\n",
- "# Squeeze means to remove dimensions of length 1.\n",
- "# Here, that removes the dummy batch dimension so it's a rank 1 tensor and returns a string\n",
- "# Rank 2 tensors map to a list of strings\n",
- "print(f\"De-Tokenizing the example tokens: {model.to_string(example_text_tokens.squeeze())}\")"
- ]
- },
{
"attachments": {},
"cell_type": "markdown",
@@ -1684,7 +1669,6 @@
" \"(CNN) President Barack Obama caught in embarrassing new scandal\\n\",\n",
" max_new_tokens=50,\n",
" temperature=0.7,\n",
- " prepend_bos=True,\n",
")"
]
},
@@ -1806,7 +1790,7 @@
"metadata": {},
"source": [
"\n",
- "We can also use hooks to intervene on activations - eg, we can set the intermediate value in layer 2 to zero to change the output to -5\n"
+ "We can also use hooks to intervene on activations - eg, we can set the intermediate value in layer 2 to zero to change the output to -4\n"
]
},
{
@@ -1821,7 +1805,7 @@
"\n",
"\n",
"print(\n",
- " \"Output after intervening on layer2.hook_scaled\",\n",
+ " \"Output after intervening on layer2.hook_square\",\n",
" model.run_with_hooks(\n",
" torch.tensor(5.0), fwd_hooks=[(\"layer2.hook_square\", set_to_zero_hook)]\n",
" ).item(),\n",
@@ -1834,9 +1818,11 @@
"source": [
"## Loading Pre-Trained Checkpoints\n",
"\n",
+ "**Note:** `TransformerBridge.boot_transformers` mainly works with HuggingFace revisions, and is only compatible with checkpoints for a few model families. For other models with saved checkpoints, keep using `HookedTransformer.from_pretrained`.\n",
+ "\n",
"There are a lot of interesting questions combining mechanistic interpretability and training dynamics - analysing model capabilities and the underlying circuits that make them possible, and how these change as we train the model. \n",
"\n",
- "TransformerLens supports these by having several model families with checkpoints throughout training. `HookedTransformer.from_pretrained` can load a checkpoint of a model with the `checkpoint_index` (the label 0 to `num_checkpoints-1`) or `checkpoint_value` (the step or token number, depending on how the checkpoints were labelled)."
+ "TransformerLens supports these by having several model families with checkpoints throughout training. `TransformerBridge.boot_transformers` can load a checkpoint of a model with the `checkpoint_index` (the label 0 to `num_checkpoints-1`) or `checkpoint_value` (the step or token number, depending on how the checkpoints were labelled)."
]
},
{
@@ -1934,7 +1920,7 @@
"from transformer_lens import evals\n",
"\n",
"# We use the two layer model with SoLU activations, chosen fairly arbitrarily as being both small (so fast to download and keep in memory) and pretty good at the induction task.\n",
- "model_name = \"solu-2l\"\n",
+ "model_name = \"NeelNanda/SoLU_2L512W_C4_Code\"\n",
"# We can load a model from a checkpoint by specifying the checkpoint_index, -1 means the final checkpoint\n",
"checkpoint_indices = [10, 25, 35, 60, -1]\n",
"checkpointed_models = []\n",
@@ -1978,7 +1964,7 @@
"source": [
"We can plot this, and see there's a sharp shift from ~200-500M tokens trained on (note the log scale on the x axis). Interestingly, this is notably earlier than the phase transition in the paper, I'm not sure what's up with that.\n",
"\n",
- "(To contextualise the numbers, the tokens in the random sequence are uniformly chosen from the first 20,000 tokens (out of ~48,000 total), so random performance is at least $\\ln(20000)\\approx 10$. A naive strategy like \"randomly choose a token that's already appeared in the first half of the sequence (384 elements)\" would get $\\ln(384)\\approx 5.95$, so the model is doing pretty well here.)"
+ "(To contextualise the numbers, the tokens in the random sequence are uniformly chosen from the first 20,000 tokens (out of ~48,000 total), so random performance is at least $\\ln(20000)\\approx 10$. A naive strategy like \"randomly choose a token that's already appeared in the first half of the sequence (384 elements)\" would get $\\ln(384)\\approx 6$, so the model is doing pretty well here.)"
]
},
{
@@ -1987,15 +1973,16 @@
"metadata": {},
"outputs": [],
"source": [
- "line(\n",
- " induction_losses,\n",
- " x=tokens_trained_on,\n",
- " xaxis=\"Tokens Trained On\",\n",
- " yaxis=\"Induction Loss\",\n",
- " title=\"Induction Loss over training: solu-2l\",\n",
- " markers=True,\n",
- " log_x=True,\n",
- ")"
+ "if not IN_GITHUB:\n",
+ " line(\n",
+ " induction_losses,\n",
+ " x=tokens_trained_on,\n",
+ " xaxis=\"Tokens Trained On\",\n",
+ " yaxis=\"Induction Loss\",\n",
+ " title=\"Induction Loss over training: solu-2l\",\n",
+ " markers=True,\n",
+ " log_x=True,\n",
+ " )"
]
}
],
diff --git a/docs/source/_static/adapter-template.py b/docs/source/_static/adapter-template.py
index 7b14928b0..61b40da00 100644
--- a/docs/source/_static/adapter-template.py
+++ b/docs/source/_static/adapter-template.py
@@ -64,10 +64,6 @@ def __init__(self, cfg: Any) -> None:
self.cfg.attn_only = False # True only for attention-only models (rare)
self.cfg.uses_rms_norm = True # Should match normalization_type
- # TODO: Set the epsilon attribute name used by this model's normalization
- # Check the HF model's norm layer to find the correct attribute name
- self.cfg.eps_attr = "variance_epsilon" # or "layer_norm_eps", "rms_norm_eps", etc.
-
# TODO: Handle GQA if applicable
# If the model uses Grouped Query Attention (n_key_value_heads < n_heads):
if hasattr(cfg, "n_key_value_heads") and cfg.n_key_value_heads is not None:
diff --git a/docs/source/content/adapter_development/adapter-creation-guide.md b/docs/source/content/adapter_development/adapter-creation-guide.md
index 2a39b3bb2..886ccad14 100644
--- a/docs/source/content/adapter_development/adapter-creation-guide.md
+++ b/docs/source/content/adapter_development/adapter-creation-guide.md
@@ -104,7 +104,6 @@ Set these on `self.cfg` in `__init__` *before* building the component mapping (t
| `gated_mlp` | `bool` | MLP has gate projection (SwiGLU) |
| `attn_only` | `bool` | Model has no MLP layers (rare) |
| `uses_rms_norm` | `bool` | Should match `normalization_type == "RMS"` |
-| `eps_attr` | `str` | HF attribute name for norm epsilon |
For GQA models, also forward `n_key_value_heads`:
@@ -284,7 +283,6 @@ Both must be clean. Don't paper over mypy errors with `# type: ignore` — fix t
## Common pitfalls
-- **Wrong `eps_attr` name.** Models that look identical use different attribute names (`variance_epsilon`, `rms_norm_eps`, `eps`). Read the norm class.
- **Forgetting `n_key_value_heads`.** Without it, GQA models silently reshape weights as if they were MHA — verification fails with cryptic shape errors.
- **Missing registration.** Adapter exists but the factory can't find it. Update both `__init__.py` and `architecture_adapter_factory.py`.
- **Skipping `setup_component_testing` for RoPE.** Rotary embeddings need to be wired through to each attention bridge or component testing produces nonsense.
diff --git a/docs/source/content/adapter_development/adapter-specification.md b/docs/source/content/adapter_development/adapter-specification.md
index 17798b6b6..f177a3926 100644
--- a/docs/source/content/adapter_development/adapter-specification.md
+++ b/docs/source/content/adapter_development/adapter-specification.md
@@ -44,7 +44,6 @@ Set these on `self.cfg` in `__init__` before building the component mapping:
| `gated_mlp` | `bool` | Whether MLP uses gate projection | Llama=True, GPT2=False |
| `attn_only` | `bool` | Whether model has no MLP layers | Usually False |
| `uses_rms_norm` | `bool` | Redundant with normalization_type but needed | Match normalization_type |
-| `eps_attr` | `str` | Attribute name for norm epsilon | `"variance_epsilon"`, `"layer_norm_eps"` |
### GQA (Grouped Query Attention)
diff --git a/docs/source/content/adapter_development/adapter-unit-test-guide.md b/docs/source/content/adapter_development/adapter-unit-test-guide.md
index 0113c11fa..767503879 100644
--- a/docs/source/content/adapter_development/adapter-unit-test-guide.md
+++ b/docs/source/content/adapter_development/adapter-unit-test-guide.md
@@ -22,7 +22,7 @@ Organize around the three things an adapter decides (config, component mapping,
| Area | Worth asserting | Skip |
| --- | --- | --- |
| **Component mapping** | The HF module paths and bridge **types** for this arch — especially non-standard ones (`transformer.wte`, `model.tok_embeddings`, `out_proj`, `fc_in`, `EncDecAttention`); the distinctive bridge (`JointQKVAttentionBridge`, `ParallelBlockBridge`, `SymbolicBridge`, `MoEBridge`, `SigLIP`); the exact submodule **set** (e.g. attention has `q_norm`/`k_norm`, or block has no `ln2`). | — |
-| **Config quirks** | Propagation that drives *behavior*: `n_key_value_heads` (GQA) through the adapter's own branch, custom `eps_attr` value, softcap / `logit_scale` coercion + `None`-fallback, `rmsnorm_uses_offset`, `parallel_attn_mlp`, `uses_combined_qkv`, `supports_fold_ln=False` when a fused projection forces it, multimodal/`gated_q_proj` flags. | A flag whose only effect is the literal you set (see "config-literal" below). |
+| **Config quirks** | Propagation that drives *behavior*: `n_key_value_heads` (GQA) through the adapter's own branch, softcap / `logit_scale` coercion + `None`-fallback, `rmsnorm_uses_offset`, `parallel_attn_mlp`, `uses_combined_qkv`, `supports_fold_ln=False` when a fused projection forces it, multimodal/`gated_q_proj` flags. | A flag whose only effect is the literal you set (see "config-literal" below). |
| **Weight conversions** | Logic the **adapter** implements: a fused-QKV split's numerical partition (which rows are Q vs K vs V — e.g. GPT-2 thirds, CodeGen's `[Q,V,K]` `mp_num` ordering, Baichuan/InternLM2 interleaved layouts), a manual LayerNorm fold (values folded, weight reset to ones, dtype preserved), the exact conversion **key set** (no stray norm/bias entries). | The einops rearrange itself (see "dependency test"). |
| **Overrides** | Each branch of `setup_component_testing` / `preprocess_weights` / `prepare_model` / `prepare_loading` you wrote — the happy path *and* the defensive `hasattr`/`None` guards, the no-op-when-absent path, the rejection guard. | Overrides you didn't write. |
| **Behavioral hook shapes** | Where the adapter's config drives reshaping: GQA `hook_k`/`hook_v` at `n_key_value_heads`, MQA single KV head, hybrid layers where attn hooks are **absent** on linear-attention layers. | Generic `(batch, seq, d_model)` output shape (it's the shared bridge's contract, not yours). |
diff --git a/docs/source/content/contributing.md b/docs/source/content/contributing.md
index fc3092ad8..0a6a09e52 100644
--- a/docs/source/content/contributing.md
+++ b/docs/source/content/contributing.md
@@ -6,6 +6,25 @@
The HookedTransformer **acceptance test suite is currently quarantined** due to a CI test-pollution issue (see `tests/QUARANTINES.md` in the repo). Changes that touch HookedTransformer internals therefore land essentially untested at the acceptance level — extra manual care is required until the suite is re-enabled.
```
+## Contributing with AI coding agents
+
+This repo ships first-class guidance for AI coding agents (Claude Code, Cursor, Copilot, Codex, and the like). If you contribute with one, point it at these files — they encode the same conventions documented on this page, in a form an agent reads in-context.
+
+The primary source of information is [`AGENTS.md`](https://github.com/TransformerLensOrg/TransformerLens/blob/main/AGENTS.md) at the repo root. Vendor-specific entry points all defer to it. Several subdirectories carry their own `AGENTS.md` with narrower rules (`tests/`, `transformer_lens/model_bridge/supported_architectures/`, and `transformer_lens/tools/model_registry/`). An agent working in one of those areas should read the local file as well as the root one.
+
+Claude Code users also get slash commands in `.claude/commands/` that wrap the common workflows. Most have plain `make`/`uv` equivalent, so non-Claude agents (and humans) can run the same thing by hand, but there are a handful that are Claude-specific at the moment.
+
+| Claude shortcut | Manual equivalent |
+|---|---|
+| `/test-unit` | `make unit-test` |
+| `/test-all` | `make test` (long) |
+| `/format` | `make format && uv run mypy .` |
+| `/typecheck` | `uv run mypy .` |
+| `/build-docs` | `set -a; source .env; set +a && uv run build-docs` |
+| `/verify-model ` | `uv run python -m transformer_lens.tools.model_registry.verify_models --model ` (dry-run first) |
+| `/add-model-support ` | Follow the adapter-authoring workflow under [Creating Architecture Adapters](#creating-architecture-adapters) |
+| `/task-complete` | `make format && uv run mypy . && make test-pr` (loop until clean) |
+
## Setup
### DevContainer
@@ -285,7 +304,7 @@ Two test layers:
### Common adapter gotchas
-- **HF raw config attributes are invisible to TL-side consumers unless explicitly propagated to `self.cfg`.** Walk the HF `config.json` and mirror any non-standard knobs (`final_logit_softcapping`, `attn_logit_softcapping`, `query_pre_attn_scalar`, `sliding_window`, `layer_types`, custom `eps_attr` names) onto `self.cfg` so weight processing and forward passes can see them.
+- **HF raw config attributes are invisible to TL-side consumers unless explicitly propagated to `self.cfg`.** Walk the HF `config.json` and mirror any non-standard knobs (`final_logit_softcapping`, `attn_logit_softcapping`, `query_pre_attn_scalar`, `sliding_window`, `layer_types`) onto `self.cfg` so weight processing and forward passes can see them.
- **Some config attrs need both surface-on-cfg AND fold-into-weight** via a `preprocess_weights()` override. The trigger: a numerical operation HF's forward applies natively must also be baked into the raw weights, or `bridge.enable_compatibility_mode()` (which calls `process_weights` on raw weights) produces wrong results. Concrete examples in-tree: Cohere `logit_scale` → `unembed.weight`; Gemma embedding scale (`√d_model`) → `embed.weight`. Skip the fold and Phase 3 / Phase 4 of `verify_models` will silently degrade.
- **Tokenizer policy is per-model, not per-architecture.** Sibling models in the same family routinely differ — the chat-instruct variant may prepend BOS where the base does not, padding side can flip, EOS handling can differ. It's worth re-checking `default_prepend_bos`, padding side, and EOS handling against the specific target rather than copying them from a starter adapter. `tokenizer_config.json` is not always reliable on its own — some architectures (Cohere is a notable example) declare `add_bos_token=False` but HF's `__call__` prepends BOS anyway. The most reliable check is to invoke the tokenizer directly:
diff --git a/docs/source/content/debugging_numerical_divergence.md b/docs/source/content/debugging_numerical_divergence.md
index 01568a3ce..23f1c67a0 100644
--- a/docs/source/content/debugging_numerical_divergence.md
+++ b/docs/source/content/debugging_numerical_divergence.md
@@ -40,7 +40,6 @@ The first hop where they disagree localizes the bug.
| Off by a constant scale in residual | Final-RMS-norm offset missing | `cfg.rmsnorm_uses_offset = True` + `ArithmeticTensorConversion(ADDITION, 1.0)` |
| Logits flat / saturated at extremes | Missing logit softcap | `cfg.output_logits_soft_cap` from HF's `final_logit_softcapping` |
| Attention pattern collapses to argmax | Missing attention-score softcap | `cfg.attn_scores_soft_cap` from HF's `attn_logit_softcapping` |
-| Off by `eps` magnitudes in norm | Wrong RMSNorm eps attribute name | `cfg.eps_attr` (Llama uses `"variance_epsilon"`, most others use `"eps"`) |
| First MLP off; gate matches | Forgot gated-MLP wiring | `GatedMLPBridge` with `{gate, in, out}` submodules — not `MLPBridge` |
| Bias-related drift | Adapter assumes biases that don't exist (Llama / RMSNorm) | `ProcessWeights._safe_get_tensor` handles `None`; check the weight-processing conversions are bias-aware |
| Drift only in compatibility mode | Hook semantic carve-out missing for post-norm or MLA | See [compatibility_mode.md](compatibility_mode.md) §"Hook semantic parity" |
diff --git a/pyproject.toml b/pyproject.toml
index a98ab985c..0039daf42 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -17,16 +17,13 @@
"rich>=12.6.0",
"sentencepiece",
"torch>=2.6",
- # Pin to the torchvision that pairs with the project's torch 2.7.x so adding
- # it for multimodal processors doesn't force a repo-wide torch upgrade.
- "torchvision>=0.22,<0.23",
"tqdm>=4.64.1",
"transformers-stream-generator>=0.0.5,<0.1",
"transformers>=5.4.0",
"typeguard>=4.2,<5",
"typing-extensions",
"wandb>=0.13.5",
-]
+ ]
description="An implementation of transformers tailored for mechanistic interpretability."
license={text="MIT"}
name="transformer-lens"
@@ -38,7 +35,7 @@
[project.optional-dependencies]
# chardet<6 works around a `requests<=2.32` compatibility-check warning that fires
# whenever chardet>=6 is installed. Remove the pin when psf/requests bumps the cap.
- evals=["lm-eval>=0.4", "chardet<6"]
+ evals=["chardet<6", "lm-eval>=0.4"]
lit=["lit-nlp>=1.3"]
[project.scripts]
@@ -79,9 +76,11 @@
jupyter=["ipywidgets>=8.1.1", "jupyterlab>=3.5.0"]
# Vision/audio deps for multimodal adapters (Gemma 3n's vision tower needs timm; image
# processors need torchvision). In default-groups so contributor/CI test runs install them,
- # while downstream users get a light core install. torchvision is pinned to the release
- # pairing with torch 2.7.x so it doesn't force a repo-wide torch upgrade.
- multimodal=["timm>=1.0.27", "torchvision>=0.22,<0.23"]
+ # while downstream users get a light core install. No torchvision upper bound: torchvision
+ # exact-pins its matching torch, so the resolver co-selects a compatible
+ # (Python, torch, torchvision) set (e.g. py3.12→0.22/torch2.7, py3.13→0.23+/torch2.8+).
+ # Capping it would just reintroduce the Python-3.13 block on the next torch/torchvision bump.
+ multimodal=["timm>=1.0.27", "torchvision>=0.22"]
quantization=["bitsandbytes>=0.46.1", "optimum-quanto>=0.2.7"]
[tool.uv]
@@ -102,9 +101,7 @@
"ignore:distutils Version classes are deprecated:DeprecationWarning",
"ignore:pkg_resources is deprecated as an API:DeprecationWarning",
]
- markers=[
- "slow: marks tests as slow (deselect with '-m \"not slow\"')",
- ]
+ markers=["slow: marks tests as slow (deselect with '-m \"not slow\"')"]
pythonpath=["."]
testpaths=["tests", "transformer_lens"] # Only test these directories
diff --git a/tests/QUARANTINES.md b/tests/QUARANTINES.md
index 5cf3e9fb2..5a587ea38 100644
--- a/tests/QUARANTINES.md
+++ b/tests/QUARANTINES.md
@@ -66,6 +66,7 @@ Rule ([AGENTS.md §10](../AGENTS.md#10-hard-rules)): **never add `xfail` / `skip
| Path | Reason | Issue |
|---|---|---|
| [`unit/model_bridge/test_bridge_generate_no_tokenizer.py`:30,128](unit/model_bridge/test_bridge_generate_no_tokenizer.py) | `skipif(_MACOS_ARM64)` — KV-cache NaN | Upstream PyTorch/HF on M-series Macs |
+| [`integration/model_bridge/test_bridge_generate_stopping_criteria.py`](integration/model_bridge/test_bridge_generate_stopping_criteria.py) | `skipif(_MACOS_ARM64)`, KV-cache NaN (one `use_past_kv_cache=True` test) | Upstream PyTorch/HF on M-series Macs |
**Un-skip:** when upstream resolves. Don't bypass — produces NaN logits.
diff --git a/tests/integration/model_bridge/test_bridge_generate_stopping_criteria.py b/tests/integration/model_bridge/test_bridge_generate_stopping_criteria.py
new file mode 100644
index 000000000..92256c4d9
--- /dev/null
+++ b/tests/integration/model_bridge/test_bridge_generate_stopping_criteria.py
@@ -0,0 +1,456 @@
+"""Tests for stop_strings and stopping_criteria on TransformerBridge.generate().
+
+Coverage:
+- stop_strings (single, list) halts generation early at the match.
+- stopping_criteria (bare StoppingCriteria, list, StoppingCriteriaList) halts.
+- The three signals (EOS, stop_strings, stopping_criteria) are independent: in
+ particular stop_at_eos=False still stops on a stop string (the early-exit and
+ finished-row padding must not be gated on stop_at_eos).
+- Defaults (both None) are a byte-for-byte no-op (back-compat guard).
+- Batched generation stops correctly.
+- Error contracts: stop_strings without a tokenizer raises ValueError, a
+ stopping_criteria callable needs no tokenizer, unsupported input paths raise
+ NotImplementedError, and a bad stopping_criteria type raises TypeError.
+- generate_stream honors the same parameters (shared _generate_tokens loop).
+
+Uses distilgpt2 (CI-cached). Greedy (do_sample=False) for determinism, and
+use_past_kv_cache=False so the tests stay robust on macOS-arm64 CI where the
+cached-eager-attention path can NaN (issue #1322).
+"""
+
+import platform
+
+import pytest
+import torch
+from transformers import StoppingCriteria, StoppingCriteriaList
+
+_MACOS_ARM64 = platform.system() == "Darwin" and platform.machine() == "arm64"
+
+# Common kwargs for the greedy, macOS-safe, token-returning generate calls below.
+_GEN = dict(do_sample=False, use_past_kv_cache=False, return_type="tokens", verbose=False)
+
+
+@pytest.fixture()
+def bridge(distilgpt2_bridge):
+ """Alias the shared session fixture for concise test signatures."""
+ return distilgpt2_bridge
+
+
+@pytest.fixture()
+def tokenizerless_bridge():
+ """A private distilgpt2 bridge with its tokenizer removed.
+
+ Function-scoped (a fresh boot per test) so removing the tokenizer never
+ contaminates the shared session fixture.
+ """
+ from transformer_lens.model_bridge import TransformerBridge
+
+ b = TransformerBridge.boot_transformers("distilgpt2", device="cpu")
+ b.tokenizer = None
+ return b
+
+
+@pytest.fixture(scope="module")
+def bridge_with_pad():
+ """A private distilgpt2 bridge with eos set as the pad token, for batched tests.
+
+ A dedicated boot (not the shared session fixture) so setting the pad token does not
+ leak into other tests.
+ """
+ from transformer_lens.model_bridge import TransformerBridge
+
+ b = TransformerBridge.boot_transformers("distilgpt2", device="cpu")
+ if b.tokenizer.pad_token is None:
+ b.tokenizer.pad_token = b.tokenizer.eos_token
+ return b
+
+
+class _StopAfterTotalLen(StoppingCriteria):
+ """Stop once the running sequence reaches a fixed total length.
+
+ Length-based so the stop step is deterministic and independent of what the
+ model emits, which makes for exact assertions.
+ """
+
+ def __init__(self, total_len: int):
+ self.total_len = total_len
+
+ def __call__(self, input_ids, scores, **kwargs):
+ return input_ids.shape[-1] >= self.total_len
+
+
+def _decode(bridge, out_row):
+ return bridge.tokenizer.decode(out_row, skip_special_tokens=True)
+
+
+def test_stop_string_halts_early(bridge):
+ """A single stop string stops generation before max_new_tokens and appears in the output."""
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+
+ out = bridge.generate(tokens, max_new_tokens=40, stop_strings=".", **_GEN)
+
+ assert out.shape[1] > prompt_len, "should have generated at least one token"
+ assert out.shape[1] < prompt_len + 40, "should have stopped before the token budget"
+ assert "." in _decode(bridge, out[0]), "the stop string should be present in the output"
+
+
+def test_stop_string_list(bridge):
+ """A list of stop strings stops on whichever appears first, deterministically."""
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+ stops = [".", ",", " the"]
+
+ out1 = bridge.generate(tokens, max_new_tokens=40, stop_strings=stops, **_GEN)
+ out2 = bridge.generate(tokens, max_new_tokens=40, stop_strings=stops, **_GEN)
+
+ assert out1.shape[1] < prompt_len + 40
+ assert any(s in _decode(bridge, out1[0]) for s in stops)
+ assert torch.equal(out1, out2), "greedy generation with stop strings must be deterministic"
+
+
+def test_custom_stopping_criteria_exact_stop(bridge):
+ """A length-based StoppingCriteria stops at exactly the expected step."""
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+
+ out = bridge.generate(
+ tokens,
+ max_new_tokens=20,
+ stopping_criteria=_StopAfterTotalLen(prompt_len + 3),
+ **_GEN,
+ )
+
+ assert out.shape[1] == prompt_len + 3
+
+
+@pytest.mark.parametrize("wrap", ["bare", "list", "criteria_list"])
+def test_stopping_criteria_accepted_forms(bridge, wrap):
+ """stopping_criteria accepts a bare StoppingCriteria, a list, or a StoppingCriteriaList."""
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+ criterion = _StopAfterTotalLen(prompt_len + 3)
+
+ if wrap == "bare":
+ stopping_criteria = criterion
+ elif wrap == "list":
+ stopping_criteria = [criterion]
+ else:
+ stopping_criteria = StoppingCriteriaList([criterion])
+
+ out = bridge.generate(tokens, max_new_tokens=20, stopping_criteria=stopping_criteria, **_GEN)
+
+ assert out.shape[1] == prompt_len + 3
+
+
+def test_stop_at_eos_false_still_stops_on_string(bridge):
+ """With stop_at_eos=False, a stop string still halts generation.
+
+ The early-exit and finished-row padding are shared with the EOS path but must
+ not be gated on stop_at_eos, otherwise a stop string could never end generation
+ when EOS stopping is off.
+ """
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+
+ stopped = bridge.generate(
+ tokens, max_new_tokens=40, stop_at_eos=False, stop_strings=".", **_GEN
+ )
+ baseline = bridge.generate(tokens, max_new_tokens=40, stop_at_eos=False, **_GEN)
+
+ assert baseline.shape[1] == prompt_len + 40, "no stop signal should run to the full budget"
+ assert stopped.shape[1] < prompt_len + 40, "the stop string should stop before the budget"
+ assert "." in _decode(bridge, stopped[0])
+
+
+def test_no_stopping_is_noop(bridge):
+ """Passing both new params as None is byte-identical to not passing them."""
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+
+ out_explicit_none = bridge.generate(
+ tokens,
+ max_new_tokens=5,
+ stop_at_eos=False,
+ stop_strings=None,
+ stopping_criteria=None,
+ **_GEN,
+ )
+ out_baseline = bridge.generate(tokens, max_new_tokens=5, stop_at_eos=False, **_GEN)
+
+ assert torch.equal(out_explicit_none, out_baseline)
+ assert out_explicit_none.shape[1] == prompt_len + 5, "nothing should stop generation"
+
+
+def test_empty_stop_strings_is_noop(bridge):
+ """An empty stop_strings list is a no-op rather than an error."""
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+
+ out = bridge.generate(tokens, max_new_tokens=4, stop_at_eos=False, stop_strings=[], **_GEN)
+
+ assert out.shape[1] == prompt_len + 4
+
+
+def test_stopping_criteria_takes_precedence_over_unmet_stop_string(bridge):
+ """When a stop string never appears, a criterion still stops generation."""
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+
+ out = bridge.generate(
+ tokens,
+ max_new_tokens=20,
+ stop_strings="zzzzzqqqqq", # not emitted by greedy distilgpt2
+ stopping_criteria=_StopAfterTotalLen(prompt_len + 3),
+ **_GEN,
+ )
+
+ assert out.shape[1] == prompt_len + 3
+
+
+def test_batched_generation_stops(bridge_with_pad):
+ """Batched generation stops on stop strings, deterministically, for every row."""
+ bridge = bridge_with_pad
+ prompts = ["The capital of France is", "Hello, my name is"]
+ tokens = bridge.to_tokens(prompts, prepend_bos=False, padding_side="left")
+ prompt_len = tokens.shape[1]
+
+ out1 = bridge.generate(tokens, max_new_tokens=40, stop_strings=".", **_GEN)
+ out2 = bridge.generate(tokens, max_new_tokens=40, stop_strings=".", **_GEN)
+
+ assert out1.shape[0] == 2, "batch dimension preserved"
+ assert out1.shape[1] < prompt_len + 40, "the batch stopped before the budget"
+ for i in range(2):
+ assert "." in _decode(bridge, out1[i]), f"row {i} should contain its stop string"
+ assert torch.equal(out1, out2), "batched greedy generation must be deterministic"
+
+
+@pytest.mark.skipif(_MACOS_ARM64, reason="Upstream macOS-arm64 KV-cache NaN, see issue #1322.")
+def test_stop_string_with_kv_cache(bridge):
+ """stop_strings also works on the default KV-cache path (not only the no-cache path)."""
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+
+ out = bridge.generate(
+ tokens,
+ max_new_tokens=40,
+ stop_strings=".",
+ do_sample=False,
+ use_past_kv_cache=True,
+ return_type="tokens",
+ verbose=False,
+ )
+
+ assert out.shape[1] < prompt_len + 40
+ assert "." in _decode(bridge, out[0])
+
+
+def test_stop_strings_requires_tokenizer(tokenizerless_bridge):
+ """stop_strings without a tokenizer raises a clear ValueError, but a criterion does not."""
+ b = tokenizerless_bridge
+ tokens = torch.tensor([[15496, 11, 314, 1101, 257]], dtype=torch.long)
+
+ # stop_at_eos=False isolates the stop_strings tokenizer requirement from the
+ # pre-existing eos-needs-a-tokenizer assertion.
+ with pytest.raises(ValueError, match="tokenizer"):
+ b.generate(tokens, max_new_tokens=3, stop_at_eos=False, stop_strings=".", **_GEN)
+
+ # A token-based stopping_criteria needs no tokenizer and must still work.
+ prompt_len = tokens.shape[1]
+ out = b.generate(
+ tokens,
+ max_new_tokens=20,
+ stop_at_eos=False,
+ stopping_criteria=_StopAfterTotalLen(prompt_len + 3),
+ **_GEN,
+ )
+ assert out.shape[1] == prompt_len + 3
+
+
+def test_invalid_stopping_criteria_type_raises(bridge):
+ """A stopping_criteria of the wrong type raises TypeError."""
+ tokens = bridge.to_tokens("The quick brown")
+ with pytest.raises(TypeError, match="StoppingCriteria"):
+ bridge.generate(tokens, max_new_tokens=3, stopping_criteria=123, **_GEN)
+
+
+def test_unsupported_input_path_raises(bridge):
+ """stop_strings with an inputs_embeds input raises NotImplementedError pointing to hf_generate."""
+ embeds = bridge.original_model.get_input_embeddings()(torch.tensor([[15496, 11, 314]])).float()
+ with pytest.raises(NotImplementedError, match="hf_generate"):
+ bridge.generate(embeds, max_new_tokens=3, stop_strings=".", verbose=False)
+
+
+def test_generate_stream_honors_stop_strings(bridge):
+ """generate_stream stops on a stop string too (shares the _generate_tokens loop)."""
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+
+ chunks = list(
+ bridge.generate_stream(
+ tokens,
+ max_new_tokens=40,
+ stop_strings=".",
+ do_sample=False,
+ use_past_kv_cache=False,
+ return_type="tokens",
+ verbose=False,
+ )
+ )
+ # First chunk includes the prompt, later chunks are deltas.
+ # Concatenating gives the full prompt + generated sequence.
+ full = torch.cat(chunks, dim=1)
+
+ assert full.shape[1] < prompt_len + 40, "stream should have stopped before the budget"
+ assert "." in _decode(bridge, full[0])
+
+
+def test_stop_string_with_output_logits(bridge):
+ """output_logits=True returns one logits entry per generated token, even on early stop."""
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+
+ out = bridge.generate(tokens, max_new_tokens=40, stop_strings=".", output_logits=True, **_GEN)
+ n_generated = out.sequences.shape[1] - prompt_len
+
+ assert n_generated < 40, "should have stopped before the budget"
+ assert len(out.logits) == n_generated, "one logits tensor per generated token"
+ assert out.logits[0].shape[0] == 1, "batch dimension on the per-step logits"
+
+
+def test_stop_string_return_type_str(bridge):
+ """return_type='str' decodes the early-stopped sequence, which contains the stop string."""
+ out = bridge.generate(
+ "The quick brown",
+ max_new_tokens=40,
+ stop_strings=".",
+ do_sample=False,
+ use_past_kv_cache=False,
+ return_type="str",
+ verbose=False,
+ )
+ assert isinstance(out, str)
+ assert "." in out
+
+
+def test_malformed_stopping_criteria_shape_raises(bridge):
+ """A criterion returning a non per-row shape gets a clear ValueError, not an opaque error."""
+
+ class _BadShape(StoppingCriteria):
+ def __call__(self, input_ids, scores, **kwargs):
+ return torch.zeros(input_ids.shape[0], 1, dtype=torch.bool) # [batch, 1], not [batch]
+
+ tokens = bridge.to_tokens("The quick brown")
+ with pytest.raises(ValueError, match="per-row bool"):
+ bridge.generate(tokens, max_new_tokens=3, stopping_criteria=_BadShape(), **_GEN)
+
+
+def test_batched_stopping_criteria_without_pad_token_raises(tokenizerless_bridge):
+ """Batched stopping_criteria with no pad/eos id and stop_at_eos=False raises a clear error.
+
+ Without a padding token, finished rows could not be frozen while the rest of the batch
+ continues, so generate refuses rather than emitting token-0 padding.
+ """
+ b = tokenizerless_bridge
+ batched = torch.tensor([[15496, 11, 314], [15496, 11, 314]], dtype=torch.long)
+ with pytest.raises(ValueError, match="padding token"):
+ b.generate(
+ batched,
+ max_new_tokens=5,
+ stop_at_eos=False,
+ stopping_criteria=_StopAfterTotalLen(5),
+ **_GEN,
+ )
+
+
+# Unsupported-path guards. These exercise the NotImplementedError branches in
+# generate() without booting a heavy enc-dec / multimodal / SSM model, by flipping the
+# exact flag each guard reads on the existing distilgpt2 bridge. The guard fires before
+# any architecture-specific generation code runs, so a stand-in model is sufficient.
+
+
+def test_stop_strings_encoder_decoder_raises(bridge, monkeypatch):
+ """stop_strings on an encoder-decoder model raises NotImplementedError."""
+ monkeypatch.setattr(bridge.original_model.config, "is_encoder_decoder", True)
+ tokens = bridge.to_tokens("The quick brown")
+ with pytest.raises(NotImplementedError, match="encoder-decoder"):
+ bridge.generate(tokens, max_new_tokens=3, stop_strings=".", verbose=False)
+
+
+def test_stop_strings_multimodal_raises(bridge):
+ """stop_strings with a multimodal input (pixel_values) raises NotImplementedError."""
+ tokens = bridge.to_tokens("The quick brown")
+ with pytest.raises(NotImplementedError, match="multimodal"):
+ bridge.generate(
+ tokens,
+ max_new_tokens=3,
+ stop_strings=".",
+ pixel_values=torch.zeros(1, 3, 8, 8),
+ verbose=False,
+ )
+
+
+def test_stop_strings_stateful_fallback_raises(bridge, monkeypatch):
+ """A stateful/SSM model with use_past_kv_cache=False raises, pointing to the cached path.
+
+ With use_past_kv_cache=False the stateful cache is off, so generate() would fall back to
+ hf_generate() and drop the kwargs. The error points to use_past_kv_cache=True, which keeps
+ generation on the hooked loop where stopping is applied.
+ """
+ monkeypatch.setattr(bridge.cfg, "is_stateful", True)
+ tokens = bridge.to_tokens("The quick brown")
+ with pytest.raises(NotImplementedError, match="use_past_kv_cache=True"):
+ bridge.generate(
+ tokens, max_new_tokens=3, stop_strings=".", use_past_kv_cache=False, verbose=False
+ )
+
+
+def test_generate_stream_stop_at_eos_false_stops_on_string(bridge):
+ """generate_stream with stop_at_eos=False still stops on a stop string."""
+ tokens = bridge.to_tokens("The quick brown")
+ prompt_len = tokens.shape[1]
+
+ chunks = list(
+ bridge.generate_stream(
+ tokens,
+ max_new_tokens=40,
+ stop_at_eos=False,
+ stop_strings=".",
+ do_sample=False,
+ use_past_kv_cache=False,
+ return_type="tokens",
+ verbose=False,
+ )
+ )
+ full = torch.cat(chunks, dim=1)
+
+ assert full.shape[1] < prompt_len + 40, "stream should have stopped before the budget"
+ assert "." in _decode(bridge, full[0])
+
+
+def test_generate_stream_batched_without_pad_token_raises(tokenizerless_bridge):
+ """Batched generate_stream stopping_criteria with no pad token and stop_at_eos=False raises."""
+ b = tokenizerless_bridge
+ batched = torch.tensor([[15496, 11, 314], [15496, 11, 314]], dtype=torch.long)
+ stream = b.generate_stream(
+ batched,
+ max_new_tokens=5,
+ stop_at_eos=False,
+ stopping_criteria=_StopAfterTotalLen(5),
+ do_sample=False,
+ use_past_kv_cache=False,
+ return_type="tokens",
+ verbose=False,
+ )
+ with pytest.raises(ValueError, match="padding token"):
+ list(stream)
+
+
+def test_stop_at_eos_false_with_pad_token(bridge_with_pad):
+ """Covers the pad_token_id arm of the no-eos padding fallback (a pad token is present)."""
+ tokens = bridge_with_pad.to_tokens("The quick brown")
+ out = bridge_with_pad.generate(
+ tokens, max_new_tokens=40, stop_at_eos=False, stop_strings=".", **_GEN
+ )
+ assert out.shape[1] < tokens.shape[1] + 40
+ assert "." in _decode(bridge_with_pad, out[0])
diff --git a/tests/integration/model_bridge/test_direct_logit_attribution.py b/tests/integration/model_bridge/test_direct_logit_attribution.py
new file mode 100644
index 000000000..5c175737e
--- /dev/null
+++ b/tests/integration/model_bridge/test_direct_logit_attribution.py
@@ -0,0 +1,167 @@
+"""Integration tests for the Direct Logit Attribution tool.
+
+DLA decomposes the part of a logit that comes from the residual stream via the
+unembedding *direction* ``W_U[:, token]``. The unembedding *bias* ``b_U`` is a
+per-token constant that no component produces, so the exact correctness
+invariant is::
+
+ sum(component DLA for token) + b_U[token] == logit[token]
+
+and, for a difference of two tokens, the two bias terms do **not** generally
+cancel (gpt2's folded ``ln_final`` bias makes them differ), so::
+
+ sum(component DLA, correct vs incorrect) == logit_diff - (b_U[c] - b_U[i])
+
+We assert these for ``HookedTransformer`` and for ``TransformerBridge``
+(compatibility mode) — the latter is the reason issue #1263 exists.
+
+These tests load gpt2 (cached), so they live in ``integration/`` per
+``tests/AGENTS.md``.
+"""
+
+import pytest
+
+PROMPT = "The Eiffel Tower is in the city of"
+CORRECT = " Paris"
+INCORRECT = " London"
+
+
+def _refs(model):
+ """Reference values: (logit_correct, logit_incorrect, b_U[c], b_U[i])."""
+ logits = model(PROMPT)
+ if logits.ndim == 2: # some Bridge configs may drop the batch dim
+ logits = logits[None]
+ c = model.to_single_token(CORRECT)
+ i = model.to_single_token(INCORRECT)
+ return (
+ logits[0, -1, c].item(),
+ logits[0, -1, i].item(),
+ model.b_U[c].item(),
+ model.b_U[i].item(),
+ )
+
+
+def _assert_complete_decomposition(model, unit):
+ """sum(DLA) reconstructs the logit / logit-diff up to the b_U constant."""
+ from transformer_lens.tools.analysis import direct_logit_attribution
+
+ logit_c, logit_i, bu_c, bu_i = _refs(model)
+
+ diff = direct_logit_attribution(
+ model, PROMPT, answer_tokens=CORRECT, incorrect_tokens=INCORRECT, unit=unit
+ )
+ single = direct_logit_attribution(model, PROMPT, answer_tokens=CORRECT, unit=unit)
+
+ # accumulated_resid ("layer") is cumulative: the last entry is the full
+ # residual stream, so it (not the column sum) is the reconstruction.
+ diff_total = diff.attribution[-1].sum() if unit == "layer" else diff.attribution.sum()
+ single_total = single.attribution[-1].sum() if unit == "layer" else single.attribution.sum()
+
+ assert diff_total.item() == pytest.approx((logit_c - logit_i) - (bu_c - bu_i), abs=1e-2)
+ assert single_total.item() == pytest.approx(logit_c - bu_c, abs=1e-2)
+
+
+@pytest.fixture(scope="module")
+def gpt2_ht():
+ from transformer_lens import HookedTransformer
+
+ return HookedTransformer.from_pretrained("gpt2", device="cpu")
+
+
+class TestDirectLogitAttributionHooked:
+ """Correctness on HookedTransformer (the reference numerics)."""
+
+ @pytest.mark.parametrize("unit", ["component", "layer", "head"])
+ def test_decomposition_reconstructs_logit(self, gpt2_ht, unit):
+ _assert_complete_decomposition(gpt2_ht, unit)
+
+ def test_component_labels_and_shape(self, gpt2_ht):
+ from transformer_lens.tools.analysis import direct_logit_attribution
+
+ res = direct_logit_attribution(
+ gpt2_ht, PROMPT, answer_tokens=CORRECT, incorrect_tokens=INCORRECT, unit="component"
+ )
+ assert res.unit == "component"
+ assert res.attribution.shape[0] == len(res.labels)
+ # Embedding term(s) plus each layer's attn_out and mlp_out.
+ assert "embed" in res.labels
+ assert sum(label.endswith("_attn_out") for label in res.labels) == gpt2_ht.cfg.n_layers
+ assert sum(label.endswith("_mlp_out") for label in res.labels) == gpt2_ht.cfg.n_layers
+
+ def test_head_labels_include_remainder(self, gpt2_ht):
+ from transformer_lens.tools.analysis import direct_logit_attribution
+
+ res = direct_logit_attribution(gpt2_ht, PROMPT, answer_tokens=CORRECT, unit="head")
+ assert len(res.labels) == gpt2_ht.cfg.n_layers * gpt2_ht.cfg.n_heads + 1
+ assert res.labels[-1] == "remainder"
+
+ def test_reuses_precomputed_cache(self, gpt2_ht):
+ from transformer_lens.tools.analysis import direct_logit_attribution
+
+ logit_c, _, bu_c, _ = _refs(gpt2_ht)
+ _, cache = gpt2_ht.run_with_cache(PROMPT)
+ res = direct_logit_attribution(gpt2_ht, answer_tokens=CORRECT, cache=cache)
+ assert res.attribution.sum().item() == pytest.approx(logit_c - bu_c, abs=1e-2)
+
+ def test_pos_none_keeps_position_axis(self, gpt2_ht):
+ from transformer_lens.tools.analysis import direct_logit_attribution
+
+ n_tokens = gpt2_ht.to_tokens(PROMPT).shape[1]
+ res = direct_logit_attribution(
+ gpt2_ht, PROMPT, answer_tokens=CORRECT, unit="component", pos=None
+ )
+ assert res.attribution.ndim == 3 # [component, batch, pos]
+ assert res.attribution.shape[-1] == n_tokens
+
+ def test_top_returns_sorted_pairs(self, gpt2_ht):
+ from transformer_lens.tools.analysis import direct_logit_attribution
+
+ res = direct_logit_attribution(
+ gpt2_ht, PROMPT, answer_tokens=CORRECT, incorrect_tokens=INCORRECT, unit="head"
+ )
+ top = res.top(3)
+ assert len(top) == 3
+ values = [v for _, v in top]
+ assert values == sorted(values, reverse=True)
+
+
+class TestDirectLogitAttributionBridge:
+ """The point of #1263: DLA must work on TransformerBridge."""
+
+ @pytest.mark.parametrize("unit", ["component", "head"])
+ def test_decomposition_reconstructs_logit(self, gpt2_bridge_compat, unit):
+ _assert_complete_decomposition(gpt2_bridge_compat, unit)
+
+
+class TestDirectLogitAttributionBridgeGuards:
+ """Guards required for correctness on TransformerBridge (from PR #1316).
+
+ Without compatibility mode the projection direction is wrong on a Bridge and
+ DLA would silently return incorrect numbers — verify the explicit refusal.
+ """
+
+ def test_non_compat_bridge_raises(self, gpt2_bridge):
+ from transformer_lens.tools.analysis import direct_logit_attribution
+
+ with pytest.raises(ValueError, match="compatibility mode"):
+ direct_logit_attribution(gpt2_bridge, PROMPT, answer_tokens=CORRECT)
+
+
+class TestDirectLogitAttributionValidation:
+ def test_invalid_unit_raises(self, gpt2_ht):
+ from transformer_lens.tools.analysis import direct_logit_attribution
+
+ with pytest.raises(ValueError, match="unit must be one of"):
+ direct_logit_attribution(gpt2_ht, PROMPT, answer_tokens=CORRECT, unit="neuron")
+
+ def test_missing_answer_tokens_raises(self, gpt2_ht):
+ from transformer_lens.tools.analysis import direct_logit_attribution
+
+ with pytest.raises(ValueError, match="answer_tokens is required"):
+ direct_logit_attribution(gpt2_ht, PROMPT)
+
+ def test_missing_input_and_cache_raises(self, gpt2_ht):
+ from transformer_lens.tools.analysis import direct_logit_attribution
+
+ with pytest.raises(ValueError, match="either `input`"):
+ direct_logit_attribution(gpt2_ht, answer_tokens=CORRECT)
diff --git a/tests/unit/model_bridge/compatibility/test_svd_interpreter.py b/tests/unit/model_bridge/compatibility/test_svd_interpreter.py
index b2b0de900..e73a087d6 100644
--- a/tests/unit/model_bridge/compatibility/test_svd_interpreter.py
+++ b/tests/unit/model_bridge/compatibility/test_svd_interpreter.py
@@ -5,7 +5,7 @@
from transformer_lens import SVDInterpreter
from transformer_lens.model_bridge import TransformerBridge
-MODEL = "gpt2" # Use a model that works with TransformerBridge
+MODEL = "Intel/tiny-random-gpt2" # Use a model that works with TransformerBridge
VECTOR_TYPES = ["OV", "w_in", "w_out"]
ATOL = 2e-4 # Absolute tolerance - how far does a float have to be before we consider it no longer equal?
@@ -17,14 +17,7 @@ def model():
@pytest.fixture(scope="module")
def second_model():
- # Use a different model architecture if available, otherwise same model
- # Note: If gpt2-medium fails to load, tests that need different models will be skipped
- try:
- return TransformerBridge.boot_transformers("gpt2-medium", device="cpu")
- except Exception:
- # Fallback to same model if gpt2-medium is not available
- # The test will skip if both models end up being the same
- return TransformerBridge.boot_transformers(MODEL, device="cpu")
+ return TransformerBridge.boot_transformers("hyper-accel/tiny-random-gpt2", device="cpu")
def test_svd_interpreter_returns_meaningful_values(model):
@@ -56,10 +49,6 @@ def test_svd_interpreter_returns_meaningful_values(model):
def test_svd_interpreter_returns_different_answers_for_different_layers(model):
- # Only test if model has multiple layers
- if model.cfg.n_layers < 2:
- pytest.skip("Model only has one layer")
-
svd_interpreter = SVDInterpreter(model)
# Layer 0 results
@@ -91,10 +80,6 @@ def test_svd_interpreter_returns_different_answers_for_different_layers(model):
def test_svd_interpreter_returns_different_answers_for_different_models(model, second_model):
- # Skip if both models are the same (check model name/config, not just object ID)
- if id(model) == id(second_model) or model.cfg.model_name == second_model.cfg.model_name:
- pytest.skip("Same model used for both fixtures")
-
# Get results from first model
svd_interpreter_1 = SVDInterpreter(model)
ov_1 = svd_interpreter_1.get_singular_vectors(
diff --git a/tests/unit/model_bridge/supported_architectures/test_baichuan_adapter.py b/tests/unit/model_bridge/supported_architectures/test_baichuan_adapter.py
index 457367718..fadd43afb 100644
--- a/tests/unit/model_bridge/supported_architectures/test_baichuan_adapter.py
+++ b/tests/unit/model_bridge/supported_architectures/test_baichuan_adapter.py
@@ -74,9 +74,6 @@ def _make_w_pack_component(d_model: int) -> Any:
class TestBaichuanAdapterConfig:
- def test_eps_attr(self, adapter: BaichuanArchitectureAdapter) -> None:
- assert adapter.cfg.eps_attr == "variance_epsilon"
-
def test_supports_fold_ln_false(self, adapter: BaichuanArchitectureAdapter) -> None:
assert adapter.supports_fold_ln is False
diff --git a/tests/unit/model_bridge/supported_architectures/test_cohere_adapter.py b/tests/unit/model_bridge/supported_architectures/test_cohere_adapter.py
index 87eaf8c9c..de8ff8673 100644
--- a/tests/unit/model_bridge/supported_architectures/test_cohere_adapter.py
+++ b/tests/unit/model_bridge/supported_architectures/test_cohere_adapter.py
@@ -69,10 +69,6 @@ def test_uses_rms_norm_is_false(self, adapter: CohereArchitectureAdapter) -> Non
# CohereLayerNorm subtracts the mean — NOT RMSNorm.
assert adapter.cfg.uses_rms_norm is False
- def test_eps_attr_is_variance_epsilon(self, adapter: CohereArchitectureAdapter) -> None:
- # CohereLayerNorm stores epsilon as self.variance_epsilon.
- assert adapter.cfg.eps_attr == "variance_epsilon"
-
def test_parallel_attn_mlp_is_true(self, adapter: CohereArchitectureAdapter) -> None:
# Single input_layernorm; attn and MLP run in parallel on same normed input.
assert adapter.cfg.parallel_attn_mlp is True
diff --git a/tests/unit/model_bridge/supported_architectures/test_gpt_bigcode_adapter.py b/tests/unit/model_bridge/supported_architectures/test_gpt_bigcode_adapter.py
index 572fa3e13..ce0e7bac8 100644
--- a/tests/unit/model_bridge/supported_architectures/test_gpt_bigcode_adapter.py
+++ b/tests/unit/model_bridge/supported_architectures/test_gpt_bigcode_adapter.py
@@ -265,10 +265,6 @@ def test_only_qkvo_conversion_keys(self, adapter: GPTBigCodeArchitectureAdapter)
def test_uses_rms_norm_false(self, adapter: GPTBigCodeArchitectureAdapter) -> None:
assert adapter.cfg.uses_rms_norm is False
- def test_eps_attr(self, adapter: GPTBigCodeArchitectureAdapter) -> None:
- # GPT-2 family eps (not RMS variance_epsilon).
- assert adapter.cfg.eps_attr == "layer_norm_epsilon"
-
# ---------------------------------------------------------------------------
# MQAQKVConversionRule tests
diff --git a/tests/unit/model_bridge/supported_architectures/test_gpt_oss_adapter.py b/tests/unit/model_bridge/supported_architectures/test_gpt_oss_adapter.py
index a555774f4..996baef5f 100644
--- a/tests/unit/model_bridge/supported_architectures/test_gpt_oss_adapter.py
+++ b/tests/unit/model_bridge/supported_architectures/test_gpt_oss_adapter.py
@@ -158,15 +158,6 @@ def __init__(self, cfg: TransformerBridgeConfig) -> None:
self.o_proj = nn.Linear(cfg.n_heads * cfg.d_head, cfg.d_model, bias=False)
-class TestGPTOSSAdapterConfig:
- """Adapter-owned config defaults that downstream bridge code relies on."""
-
- def test_eps_attr_is_variance_epsilon(self, adapter: GPTOSSArchitectureAdapter) -> None:
- """GPT-OSS uses HF's `variance_epsilon` attribute name on RMSNorm modules,
- not the default `eps`. Downstream norm-folding reads this attribute."""
- assert adapter.cfg.eps_attr == "variance_epsilon"
-
-
class TestGPTOSSWeightConversions:
"""GPT-OSS uses the standard QKVO weight conversions (no biases), with GQA head counts."""
@@ -370,29 +361,23 @@ def _hook(x: Any, hook: Any) -> Any:
# Identity RoPE inputs keep this test focused on hook reshaping, not rotation math.
cos = ones(1, self.SEQ, self.D_HEAD)
sin = zeros(1, self.SEQ, self.D_HEAD)
- out = attn_bridge(hidden, position_embeddings=(cos, sin))
- # The attention bridge may return either a bare tensor or an (output, ...) tuple.
- out_tensor = out[0] if isinstance(out, tuple) else out
+ attn_bridge(hidden, position_embeddings=(cos, sin))
- return captured["q"], captured["k"], captured["v"], out_tensor
+ return captured["q"], captured["k"], captured["v"]
def test_hook_q_uses_n_heads(
self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
) -> None:
- q, _, _, _ = self._run_and_capture(wired_attn_bridge)
+ q, _, _ = self._run_and_capture(wired_attn_bridge)
assert q.shape == (self.BATCH, self.SEQ, self.N_HEADS, self.D_HEAD)
def test_hook_kv_use_n_kv_heads(
self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
) -> None:
- _, k, v, _ = self._run_and_capture(wired_attn_bridge)
+ _, k, v = self._run_and_capture(wired_attn_bridge)
assert k.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
assert v.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
- def test_attn_output_shape(self, wired_attn_bridge: PositionEmbeddingsAttentionBridge) -> None:
- _, _, _, out = self._run_and_capture(wired_attn_bridge)
- assert out.shape == (self.BATCH, self.SEQ, self.D_MODEL)
-
class TestGPTOSSSetupHookCompatibility:
"""setup_hook_compatibility wires the bridge model's rotary_emb onto every
diff --git a/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py b/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py
new file mode 100644
index 000000000..06202464f
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_granite_adapter.py
@@ -0,0 +1,266 @@
+"""Unit tests for GraniteArchitectureAdapter and GraniteMoeArchitectureAdapter.
+
+Tests cover:
+- Component mapping structure (bridge types and HF module names)
+- Weight conversion key set
+- Config flags set by the adapter
+"""
+
+import pytest
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge.generalized_components import (
+ BlockBridge,
+ EmbeddingBridge,
+ GatedMLPBridge,
+ LinearBridge,
+ MoEBridge,
+ PositionEmbeddingsAttentionBridge,
+ RMSNormalizationBridge,
+ RotaryEmbeddingBridge,
+ UnembeddingBridge,
+)
+from transformer_lens.model_bridge.supported_architectures.granite import (
+ GraniteArchitectureAdapter,
+)
+from transformer_lens.model_bridge.supported_architectures.granite_moe import (
+ GraniteMoeArchitectureAdapter,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+N_HEADS = 8
+N_KV_HEADS = 2
+D_MODEL = 64
+D_MLP = 256
+N_LAYERS = 2
+N_CTX = 256
+D_VOCAB = 1000
+
+
+def _make_cfg(
+ n_heads: int = N_HEADS,
+ n_kv_heads: int = N_KV_HEADS,
+ d_model: int = D_MODEL,
+ n_layers: int = N_LAYERS,
+ d_mlp: int = D_MLP,
+ d_vocab: int = D_VOCAB,
+ n_ctx: int = N_CTX,
+) -> TransformerBridgeConfig:
+ """Return a minimal TransformerBridgeConfig for Granite adapter tests."""
+ return TransformerBridgeConfig(
+ d_model=d_model,
+ d_head=d_model // n_heads,
+ n_layers=n_layers,
+ n_ctx=n_ctx,
+ n_heads=n_heads,
+ d_vocab=d_vocab,
+ d_mlp=d_mlp,
+ n_key_value_heads=n_kv_heads,
+ default_prepend_bos=False,
+ architecture="GraniteForCausalLM",
+ )
+
+
+@pytest.fixture
+def cfg() -> TransformerBridgeConfig:
+ return _make_cfg()
+
+
+@pytest.fixture
+def adapter(cfg: TransformerBridgeConfig) -> GraniteArchitectureAdapter:
+ return GraniteArchitectureAdapter(cfg)
+
+
+@pytest.fixture
+def moe_cfg() -> TransformerBridgeConfig:
+ return _make_cfg()
+
+
+@pytest.fixture
+def moe_adapter(moe_cfg: TransformerBridgeConfig) -> GraniteMoeArchitectureAdapter:
+ return GraniteMoeArchitectureAdapter(moe_cfg)
+
+
+# ---------------------------------------------------------------------------
+# Config flag tests
+# ---------------------------------------------------------------------------
+
+
+class TestGraniteAdapterConfig:
+ """Tests that the adapter sets the correct config flags."""
+
+ def test_normalization_type(self, adapter: GraniteArchitectureAdapter) -> None:
+ assert adapter.cfg.normalization_type == "RMS"
+
+ def test_positional_embedding_type(self, adapter: GraniteArchitectureAdapter) -> None:
+ assert adapter.cfg.positional_embedding_type == "rotary"
+
+ def test_final_rms(self, adapter: GraniteArchitectureAdapter) -> None:
+ """Granite uses RMSNorm as the final norm (final_rms=True)."""
+ assert adapter.cfg.final_rms is True
+
+ def test_gated_mlp(self, adapter: GraniteArchitectureAdapter) -> None:
+ assert adapter.cfg.gated_mlp is True
+
+ def test_default_prepend_bos_false(self, adapter: GraniteArchitectureAdapter) -> None:
+ """Granite models do not prepend BOS by default."""
+ assert adapter.cfg.default_prepend_bos is False
+
+ def test_n_key_value_heads_propagated(self, adapter: GraniteArchitectureAdapter) -> None:
+ assert adapter.cfg.n_key_value_heads == N_KV_HEADS
+
+
+# ---------------------------------------------------------------------------
+# Component mapping tests — dense Granite
+# ---------------------------------------------------------------------------
+
+
+class TestGraniteAdapterComponentMapping:
+ """Tests that component_mapping has the correct bridge types and HF module names."""
+
+ def test_top_level_keys(self, adapter: GraniteArchitectureAdapter) -> None:
+ assert set(adapter.component_mapping.keys()) == {
+ "embed",
+ "rotary_emb",
+ "blocks",
+ "ln_final",
+ "unembed",
+ }
+
+ def test_bridge_types(self, adapter: GraniteArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ assert isinstance(mapping["embed"], EmbeddingBridge)
+ assert isinstance(mapping["rotary_emb"], RotaryEmbeddingBridge)
+ assert isinstance(mapping["blocks"], BlockBridge)
+ assert isinstance(mapping["ln_final"], RMSNormalizationBridge)
+ assert isinstance(mapping["unembed"], UnembeddingBridge)
+
+ def test_top_level_hf_paths(self, adapter: GraniteArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ assert mapping["embed"].name == "model.embed_tokens"
+ assert mapping["rotary_emb"].name == "model.rotary_emb"
+ assert mapping["blocks"].name == "model.layers"
+ assert mapping["ln_final"].name == "model.norm"
+ assert mapping["unembed"].name == "lm_head"
+
+ def test_block_submodule_keys(self, adapter: GraniteArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert set(blocks.submodules.keys()) == {"ln1", "ln2", "attn", "mlp"}
+
+ def test_block_bridge_types(self, adapter: GraniteArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert isinstance(blocks.submodules["ln1"], RMSNormalizationBridge)
+ assert isinstance(blocks.submodules["ln2"], RMSNormalizationBridge)
+ assert isinstance(blocks.submodules["attn"], PositionEmbeddingsAttentionBridge)
+ assert isinstance(blocks.submodules["mlp"], GatedMLPBridge)
+
+ def test_block_hf_paths(self, adapter: GraniteArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["ln1"].name == "input_layernorm"
+ assert blocks.submodules["ln2"].name == "post_attention_layernorm"
+ assert blocks.submodules["attn"].name == "self_attn"
+ assert blocks.submodules["mlp"].name == "mlp"
+
+ def test_attention_submodule_keys(self, adapter: GraniteArchitectureAdapter) -> None:
+ """Granite uses separate Q, K, V, O projections."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert set(attn.submodules.keys()) == {"q", "k", "v", "o"}
+
+ def test_attention_hf_paths(self, adapter: GraniteArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.submodules["q"].name == "q_proj"
+ assert attn.submodules["k"].name == "k_proj"
+ assert attn.submodules["v"].name == "v_proj"
+ assert attn.submodules["o"].name == "o_proj"
+
+ def test_mlp_submodule_keys(self, adapter: GraniteArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert set(mlp.submodules.keys()) == {"gate", "in", "out"}
+
+ def test_mlp_hf_paths(self, adapter: GraniteArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert mlp.submodules["gate"].name == "gate_proj"
+ assert mlp.submodules["in"].name == "up_proj"
+ assert mlp.submodules["out"].name == "down_proj"
+
+ def test_attention_linear_bridge_types(self, adapter: GraniteArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ for submodule in attn.submodules.values():
+ assert isinstance(submodule, LinearBridge)
+
+ def test_mlp_linear_bridge_types(self, adapter: GraniteArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ for submodule in mlp.submodules.values():
+ assert isinstance(submodule, LinearBridge)
+
+
+# ---------------------------------------------------------------------------
+# Weight conversion key tests — dense Granite
+# ---------------------------------------------------------------------------
+
+
+class TestGraniteAdapterWeightConversions:
+ """Tests that weight_processing_conversions has exactly the expected keys."""
+
+ def test_exact_conversion_key_set(self, adapter: GraniteArchitectureAdapter) -> None:
+ assert set(adapter.weight_processing_conversions.keys()) == {
+ "blocks.{i}.attn.q.weight",
+ "blocks.{i}.attn.k.weight",
+ "blocks.{i}.attn.v.weight",
+ "blocks.{i}.attn.o.weight",
+ }
+
+
+# ---------------------------------------------------------------------------
+# GraniteMoe component mapping tests
+# ---------------------------------------------------------------------------
+
+
+class TestGraniteMoeAdapterComponentMapping:
+ """GraniteMoe replaces dense MLP with MoE; everything else is identical to Granite."""
+
+ def test_top_level_keys(self, moe_adapter: GraniteMoeArchitectureAdapter) -> None:
+ assert set(moe_adapter.component_mapping.keys()) == {
+ "embed",
+ "rotary_emb",
+ "blocks",
+ "ln_final",
+ "unembed",
+ }
+
+ def test_mlp_is_moe_bridge(self, moe_adapter: GraniteMoeArchitectureAdapter) -> None:
+ mlp = moe_adapter.component_mapping["blocks"].submodules["mlp"]
+ assert isinstance(mlp, MoEBridge)
+
+ def test_moe_hf_path(self, moe_adapter: GraniteMoeArchitectureAdapter) -> None:
+ mlp = moe_adapter.component_mapping["blocks"].submodules["mlp"]
+ assert mlp.name == "block_sparse_moe"
+
+ def test_non_mlp_components_match_dense(
+ self,
+ adapter: GraniteArchitectureAdapter,
+ moe_adapter: GraniteMoeArchitectureAdapter,
+ ) -> None:
+ """Embed, rotary_emb, ln_final, unembed, and attention are shared with dense Granite."""
+ for key in ("embed", "rotary_emb", "ln_final", "unembed"):
+ assert type(moe_adapter.component_mapping[key]) is type(adapter.component_mapping[key])
+ assert moe_adapter.component_mapping[key].name == adapter.component_mapping[key].name
+
+ def test_attention_unchanged_in_moe(self, moe_adapter: GraniteMoeArchitectureAdapter) -> None:
+ attn = moe_adapter.component_mapping["blocks"].submodules["attn"]
+ assert isinstance(attn, PositionEmbeddingsAttentionBridge)
+ assert set(attn.submodules.keys()) == {"q", "k", "v", "o"}
+
+ def test_moe_config_flags_match_dense(
+ self,
+ adapter: GraniteArchitectureAdapter,
+ moe_adapter: GraniteMoeArchitectureAdapter,
+ ) -> None:
+ """MoE variant inherits the same config flags as dense Granite."""
+ assert moe_adapter.cfg.normalization_type == adapter.cfg.normalization_type
+ assert moe_adapter.cfg.positional_embedding_type == adapter.cfg.positional_embedding_type
+ assert moe_adapter.cfg.final_rms == adapter.cfg.final_rms
+ assert moe_adapter.cfg.default_prepend_bos == adapter.cfg.default_prepend_bos
diff --git a/tests/unit/model_bridge/supported_architectures/test_internlm2_adapter.py b/tests/unit/model_bridge/supported_architectures/test_internlm2_adapter.py
index 059f2fc73..f3861eece 100644
--- a/tests/unit/model_bridge/supported_architectures/test_internlm2_adapter.py
+++ b/tests/unit/model_bridge/supported_architectures/test_internlm2_adapter.py
@@ -94,9 +94,6 @@ def _fill_interleaved(
class TestInternLM2AdapterConfig:
"""Adapter sets all required config attributes."""
- def test_eps_attr(self, adapter: InternLM2ArchitectureAdapter) -> None:
- assert adapter.cfg.eps_attr == "variance_epsilon"
-
def test_supports_fold_ln_false(self, adapter: InternLM2ArchitectureAdapter) -> None:
# fold_ln silently skips attn when wqkv is fused in bridge state dict.
assert adapter.supports_fold_ln is False
diff --git a/tests/unit/model_bridge/supported_architectures/test_llama_adapter.py b/tests/unit/model_bridge/supported_architectures/test_llama_adapter.py
new file mode 100644
index 000000000..a1168a591
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_llama_adapter.py
@@ -0,0 +1,300 @@
+"""Unit tests for LlamaArchitectureAdapter.
+
+Tests cover:
+- Component mapping structure (bridge types and HF module names)
+- Weight conversion key set and rearrange patterns
+- GQA: n_key_value_heads propagates to K/V conversions only
+- setup_component_testing rotary embedding wiring
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from transformer_lens.config.transformer_bridge_config import TransformerBridgeConfig
+from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
+from transformer_lens.conversion_utils.param_processing_conversion import (
+ ParamProcessingConversion,
+)
+from transformer_lens.model_bridge.generalized_components import (
+ BlockBridge,
+ EmbeddingBridge,
+ GatedMLPBridge,
+ LinearBridge,
+ PositionEmbeddingsAttentionBridge,
+ RMSNormalizationBridge,
+ RotaryEmbeddingBridge,
+ UnembeddingBridge,
+)
+from transformer_lens.model_bridge.supported_architectures.llama import (
+ LlamaArchitectureAdapter,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+
+def _make_cfg(
+ n_heads: int = 32,
+ d_model: int = 4096,
+ n_layers: int = 32,
+ d_vocab: int = 32000,
+ n_ctx: int = 4096,
+ **overrides,
+) -> TransformerBridgeConfig:
+ """Return a minimal TransformerBridgeConfig for LLaMA adapter tests."""
+ cfg = TransformerBridgeConfig(
+ d_model=d_model,
+ d_head=d_model // n_heads,
+ n_heads=n_heads,
+ n_layers=n_layers,
+ n_ctx=n_ctx,
+ d_vocab=d_vocab,
+ architecture="LlamaForCausalLM",
+ )
+ for k, v in overrides.items():
+ setattr(cfg, k, v)
+ return cfg
+
+
+@pytest.fixture(scope="module")
+def adapter() -> LlamaArchitectureAdapter:
+ return LlamaArchitectureAdapter(_make_cfg())
+
+
+# ---------------------------------------------------------------------------
+# Component mapping — top-level key set and bridge types
+# ---------------------------------------------------------------------------
+
+
+class TestLlamaComponentMapping:
+ """Component mapping has the correct slots, bridge types, and HF module paths."""
+
+ def test_top_level_keys(self, adapter: LlamaArchitectureAdapter) -> None:
+ assert set(adapter.component_mapping.keys()) == {
+ "embed",
+ "rotary_emb",
+ "blocks",
+ "ln_final",
+ "unembed",
+ }
+
+ def test_no_pos_embed_key(self, adapter: LlamaArchitectureAdapter) -> None:
+ """LLaMA uses rotary embeddings — no learned positional embedding component."""
+ assert "pos_embed" not in adapter.component_mapping
+
+ def test_bridge_types(self, adapter: LlamaArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ assert isinstance(mapping["embed"], EmbeddingBridge)
+ assert isinstance(mapping["rotary_emb"], RotaryEmbeddingBridge)
+ assert isinstance(mapping["blocks"], BlockBridge)
+ assert isinstance(mapping["ln_final"], RMSNormalizationBridge)
+ assert isinstance(mapping["unembed"], UnembeddingBridge)
+
+ def test_top_level_hf_paths(self, adapter: LlamaArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ assert mapping["embed"].name == "model.embed_tokens"
+ assert mapping["rotary_emb"].name == "model.rotary_emb"
+ assert mapping["blocks"].name == "model.layers"
+ assert mapping["ln_final"].name == "model.norm"
+ assert mapping["unembed"].name == "lm_head"
+
+ def test_block_submodule_keys(self, adapter: LlamaArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert set(blocks.submodules.keys()) == {"ln1", "ln2", "attn", "mlp"}
+
+ def test_block_submodule_types(self, adapter: LlamaArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert isinstance(blocks.submodules["ln1"], RMSNormalizationBridge)
+ assert isinstance(blocks.submodules["ln2"], RMSNormalizationBridge)
+ assert isinstance(blocks.submodules["attn"], PositionEmbeddingsAttentionBridge)
+ assert isinstance(blocks.submodules["mlp"], GatedMLPBridge)
+
+ def test_block_submodule_hf_paths(self, adapter: LlamaArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["ln1"].name == "input_layernorm"
+ assert blocks.submodules["ln2"].name == "post_attention_layernorm"
+ assert blocks.submodules["attn"].name == "self_attn"
+ assert blocks.submodules["mlp"].name == "mlp"
+
+ def test_attn_requires_mask_and_position_embeddings(
+ self, adapter: LlamaArchitectureAdapter
+ ) -> None:
+ """LLaMA RoPE attention requires both attention mask and position embeddings."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.requires_attention_mask is True
+ assert attn.requires_position_embeddings is True
+
+ def test_attn_qkvo_submodule_paths(self, adapter: LlamaArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert set(attn.submodules.keys()) == {"q", "k", "v", "o"}
+ assert attn.submodules["q"].name == "q_proj"
+ assert attn.submodules["k"].name == "k_proj"
+ assert attn.submodules["v"].name == "v_proj"
+ assert attn.submodules["o"].name == "o_proj"
+
+ def test_attn_qkvo_are_linear_bridges(self, adapter: LlamaArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ for sub in attn.submodules.values():
+ assert isinstance(sub, LinearBridge)
+
+ def test_mlp_submodule_paths(self, adapter: LlamaArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert set(mlp.submodules.keys()) == {"gate", "in", "out"}
+ assert mlp.submodules["gate"].name == "gate_proj"
+ assert mlp.submodules["in"].name == "up_proj"
+ assert mlp.submodules["out"].name == "down_proj"
+
+
+# ---------------------------------------------------------------------------
+# Anti-drift config flags
+# ---------------------------------------------------------------------------
+
+
+class TestLlamaAdapterConfig:
+ """Anti-drift flags that must not silently regress."""
+
+ def test_final_rms_is_true(self, adapter: LlamaArchitectureAdapter) -> None:
+ assert adapter.cfg.final_rms is True
+
+ def test_uses_rms_norm_is_true(self, adapter: LlamaArchitectureAdapter) -> None:
+ assert adapter.cfg.uses_rms_norm is True
+
+ def test_gated_mlp_is_true(self, adapter: LlamaArchitectureAdapter) -> None:
+ """LLaMA uses a gated SwiGLU MLP — must not silently revert to vanilla MLP."""
+ assert adapter.cfg.gated_mlp is True
+
+
+# ---------------------------------------------------------------------------
+# Weight processing conversions
+# ---------------------------------------------------------------------------
+
+
+class TestLlamaWeightConversions:
+ """weight_processing_conversions has exactly the expected QKVO keys."""
+
+ def test_exact_conversion_key_set(self, adapter: LlamaArchitectureAdapter) -> None:
+ assert set(adapter.weight_processing_conversions.keys()) == {
+ "blocks.{i}.attn.q.weight",
+ "blocks.{i}.attn.k.weight",
+ "blocks.{i}.attn.v.weight",
+ "blocks.{i}.attn.o.weight",
+ }
+
+ def test_qkv_conversions_use_split_heads_pattern(
+ self, adapter: LlamaArchitectureAdapter
+ ) -> None:
+ """'(n h) m -> n m h' splits [n_heads*d_head, d_model] → [n, d_model, d_head]."""
+ for slot in ("q", "k", "v"):
+ conv = adapter.weight_processing_conversions[f"blocks.{{i}}.attn.{slot}.weight"]
+ assert isinstance(conv, ParamProcessingConversion)
+ assert isinstance(conv.tensor_conversion, RearrangeTensorConversion)
+ assert conv.tensor_conversion.pattern == "(n h) m -> n m h"
+
+ def test_o_conversion_uses_merge_heads_pattern(self, adapter: LlamaArchitectureAdapter) -> None:
+ """'m (n h) -> n h m' moves n to the front for the output projection."""
+ conv = adapter.weight_processing_conversions["blocks.{i}.attn.o.weight"]
+ assert isinstance(conv, ParamProcessingConversion)
+ assert isinstance(conv.tensor_conversion, RearrangeTensorConversion)
+ assert conv.tensor_conversion.pattern == "m (n h) -> n h m"
+
+ def test_no_bias_conversion_keys(self, adapter: LlamaArchitectureAdapter) -> None:
+ """LLaMA has no attention or MLP biases — no b_Q/b_K/b_V/b_O conversions."""
+ keys = set(adapter.weight_processing_conversions.keys())
+ assert not any("bias" in k or ".b_" in k for k in keys)
+
+ def test_no_norm_conversion_keys(self, adapter: LlamaArchitectureAdapter) -> None:
+ """RMSNorm has no bias offset — no ln1/ln2/ln_final conversion entries."""
+ keys = set(adapter.weight_processing_conversions.keys())
+ assert not any("ln" in k for k in keys)
+
+
+# ---------------------------------------------------------------------------
+# GQA support — LLaMA 3.1 / 3.2 / 3.3
+# ---------------------------------------------------------------------------
+
+
+class TestLlamaGQASupport:
+ """n_key_value_heads must propagate to K/V conversions and leave Q/O unchanged."""
+
+ def test_no_gqa_defaults_to_n_heads(self) -> None:
+ """Without n_key_value_heads, K/V use n_heads (MHA mode)."""
+ adapter = LlamaArchitectureAdapter(_make_cfg(n_heads=32))
+ k_conv = adapter.weight_processing_conversions["blocks.{i}.attn.k.weight"]
+ assert k_conv.tensor_conversion.axes_lengths["n"] == 32
+
+ def test_gqa_propagates_to_kv_conversions(self) -> None:
+ """With 8 KV heads (LLaMA-3 style), K/V conversions must use n=8."""
+ adapter = LlamaArchitectureAdapter(_make_cfg(n_heads=32, n_key_value_heads=8))
+ for slot in ("k", "v"):
+ conv = adapter.weight_processing_conversions[f"blocks.{{i}}.attn.{slot}.weight"]
+ assert conv.tensor_conversion.axes_lengths["n"] == 8
+
+ def test_gqa_does_not_affect_q_conversion(self) -> None:
+ """Q always uses full n_heads regardless of GQA."""
+ adapter = LlamaArchitectureAdapter(_make_cfg(n_heads=32, n_key_value_heads=8))
+ q_conv = adapter.weight_processing_conversions["blocks.{i}.attn.q.weight"]
+ assert q_conv.tensor_conversion.axes_lengths["n"] == 32
+
+ def test_gqa_does_not_affect_o_conversion(self) -> None:
+ """O projection always uses n_heads; GQA only affects K/V."""
+ adapter = LlamaArchitectureAdapter(_make_cfg(n_heads=32, n_key_value_heads=8))
+ o_conv = adapter.weight_processing_conversions["blocks.{i}.attn.o.weight"]
+ assert o_conv.tensor_conversion.axes_lengths["n"] == 32
+
+
+# ---------------------------------------------------------------------------
+# setup_component_testing — rotary embedding wiring
+# ---------------------------------------------------------------------------
+
+
+class _DummyAttn:
+ def __init__(self) -> None:
+ self.rotary_emb = None
+
+ def set_rotary_emb(self, rotary_emb: object) -> None:
+ self.rotary_emb = rotary_emb
+
+
+class _DummyBlock:
+ def __init__(self, has_attn: bool = True) -> None:
+ if has_attn:
+ self.attn = _DummyAttn()
+
+
+class _DummyBridgeModel:
+ def __init__(self, blocks: list) -> None:
+ self.blocks = blocks
+
+
+def _fake_hf_model(rotary_emb: object) -> SimpleNamespace:
+ return SimpleNamespace(model=SimpleNamespace(rotary_emb=rotary_emb))
+
+
+class TestLlamaSetupComponentTesting:
+ """setup_component_testing wires rotary_emb onto every block's attention bridge."""
+
+ def test_sets_rotary_emb_on_all_blocks(self) -> None:
+ adapter = LlamaArchitectureAdapter(_make_cfg())
+ rotary_emb = object()
+ bridge_model = _DummyBridgeModel([_DummyBlock(), _DummyBlock(), _DummyBlock()])
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model)
+
+ for block in bridge_model.blocks:
+ assert block.attn.rotary_emb is rotary_emb
+
+ def test_skips_blocks_without_attn(self) -> None:
+ adapter = LlamaArchitectureAdapter(_make_cfg())
+ rotary_emb = object()
+ bridge_model = _DummyBridgeModel([_DummyBlock(), _DummyBlock(has_attn=False)])
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model)
+
+ assert bridge_model.blocks[0].attn.rotary_emb is rotary_emb
+
+ def test_no_bridge_model_does_not_raise(self) -> None:
+ """setup_component_testing without a bridge_model must not raise."""
+ adapter = LlamaArchitectureAdapter(_make_cfg())
+ adapter.setup_component_testing(_fake_hf_model(object()))
diff --git a/tests/unit/model_bridge/supported_architectures/test_llava_adapter.py b/tests/unit/model_bridge/supported_architectures/test_llava_adapter.py
index 0344c0e76..d494b6f11 100644
--- a/tests/unit/model_bridge/supported_architectures/test_llava_adapter.py
+++ b/tests/unit/model_bridge/supported_architectures/test_llava_adapter.py
@@ -63,9 +63,6 @@ def adapter(self):
def test_is_multimodal(self, adapter):
assert adapter.cfg.is_multimodal is True
- def test_eps_attr(self, adapter):
- assert adapter.cfg.eps_attr == "variance_epsilon"
-
def test_vision_config_extracted(self, adapter):
assert adapter.cfg.vision_hidden_size == 1024
assert adapter.cfg.vision_num_layers == 24
diff --git a/tests/unit/model_bridge/supported_architectures/test_llava_onevision_adapter.py b/tests/unit/model_bridge/supported_architectures/test_llava_onevision_adapter.py
new file mode 100644
index 000000000..290203ad1
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_llava_onevision_adapter.py
@@ -0,0 +1,133 @@
+"""Unit tests for LlavaOnevisionArchitectureAdapter.
+
+LlavaOnevisionArchitectureAdapter inherits its config, component mapping, and
+weight conversions from LlavaArchitectureAdapter (covered by test_llava_adapter.py).
+This suite pins the subclass contract and the prepare_model weight-tying override.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from transformer_lens.config.transformer_bridge_config import TransformerBridgeConfig
+from transformer_lens.model_bridge.supported_architectures.llava import (
+ LlavaArchitectureAdapter,
+)
+from transformer_lens.model_bridge.supported_architectures.llava_onevision import (
+ LlavaOnevisionArchitectureAdapter,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+
+def _make_cfg(vision_model_type: str = "clip_vision_model") -> TransformerBridgeConfig:
+ """Return a minimal TransformerBridgeConfig for LLaVA-OneVision tests."""
+ cfg = TransformerBridgeConfig(
+ d_model=64,
+ d_head=16,
+ n_heads=4,
+ n_layers=2,
+ n_ctx=512,
+ d_vocab=1000,
+ architecture="LlavaOnevisionForConditionalGeneration",
+ )
+ cfg.vision_config = SimpleNamespace(
+ model_type=vision_model_type,
+ hidden_size=256,
+ num_hidden_layers=4,
+ num_attention_heads=4,
+ )
+ return cfg
+
+
+@pytest.fixture
+def adapter() -> LlavaOnevisionArchitectureAdapter:
+ return LlavaOnevisionArchitectureAdapter(_make_cfg())
+
+
+# ---------------------------------------------------------------------------
+# Inheritance tests
+# ---------------------------------------------------------------------------
+
+
+class TestLlavaOnevisionInheritance:
+ """LlavaOnevisionArchitectureAdapter must be a LlavaArchitectureAdapter subclass."""
+
+ def test_subclass_of_llava(self) -> None:
+ assert issubclass(LlavaOnevisionArchitectureAdapter, LlavaArchitectureAdapter)
+
+ def test_instance_of_llava(self, adapter: LlavaOnevisionArchitectureAdapter) -> None:
+ assert isinstance(adapter, LlavaArchitectureAdapter)
+
+
+# ---------------------------------------------------------------------------
+# prepare_model weight-tying tests
+# ---------------------------------------------------------------------------
+
+
+class TestLlavaOnevisionPrepareModel:
+ """prepare_model fixes weight tying when text_config and top-level config disagree."""
+
+ def _make_hf_model(
+ self, tie_word_embeddings_text: bool, tie_word_embeddings_top: bool
+ ) -> MagicMock:
+ """Build a minimal mock HF model for prepare_model testing."""
+ embed = MagicMock()
+ embed.weight = "original_weight"
+
+ language_model = MagicMock()
+ language_model.embed_tokens = embed
+
+ model = MagicMock()
+ model.language_model = language_model
+
+ lm_head = MagicMock()
+ lm_head.weight = "random_weight"
+
+ text_config = SimpleNamespace(tie_word_embeddings=tie_word_embeddings_text)
+ config = SimpleNamespace(
+ tie_word_embeddings=tie_word_embeddings_top,
+ text_config=text_config,
+ )
+
+ hf_model = MagicMock()
+ hf_model.model = model
+ hf_model.lm_head = lm_head
+ hf_model.config = config
+ return hf_model
+
+ def test_ties_weights_when_text_config_says_tied_but_top_level_says_not(
+ self, adapter: LlavaOnevisionArchitectureAdapter
+ ) -> None:
+ """lm_head.weight should be set to embed.weight when text_config disagrees."""
+ hf_model = self._make_hf_model(tie_word_embeddings_text=True, tie_word_embeddings_top=False)
+ adapter.prepare_model(hf_model)
+ assert hf_model.lm_head.weight == "original_weight"
+
+ def test_no_tying_when_both_agree_tied(
+ self, adapter: LlavaOnevisionArchitectureAdapter
+ ) -> None:
+ """No weight override when top-level config already says tied."""
+ hf_model = self._make_hf_model(tie_word_embeddings_text=True, tie_word_embeddings_top=True)
+ original_weight = hf_model.lm_head.weight
+ adapter.prepare_model(hf_model)
+ assert hf_model.lm_head.weight == original_weight
+
+ def test_no_tying_when_text_config_says_not_tied(
+ self, adapter: LlavaOnevisionArchitectureAdapter
+ ) -> None:
+ """No weight override when text_config says not tied."""
+ hf_model = self._make_hf_model(
+ tie_word_embeddings_text=False, tie_word_embeddings_top=False
+ )
+ original_weight = hf_model.lm_head.weight
+ adapter.prepare_model(hf_model)
+ assert hf_model.lm_head.weight == original_weight
+
+ def test_no_op_when_no_lm_head(self, adapter: LlavaOnevisionArchitectureAdapter) -> None:
+ """prepare_model is a no-op when lm_head is absent."""
+ hf_model = MagicMock(spec=[]) # no attributes at all
+ adapter.prepare_model(hf_model) # must not raise
diff --git a/tests/unit/model_bridge/supported_architectures/test_mistral_adapter.py b/tests/unit/model_bridge/supported_architectures/test_mistral_adapter.py
new file mode 100644
index 000000000..1d423e744
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_mistral_adapter.py
@@ -0,0 +1,258 @@
+"""Unit tests for MistralArchitectureAdapter.
+
+Tests cover:
+- Component mapping structure (bridge types and HF module names)
+- Weight conversion key set and rearrange patterns
+- GQA: n_key_value_heads propagates to K/V conversions only
+- Anti-drift config flags
+"""
+
+import pytest
+
+from transformer_lens.config.transformer_bridge_config import TransformerBridgeConfig
+from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
+from transformer_lens.conversion_utils.param_processing_conversion import (
+ ParamProcessingConversion,
+)
+from transformer_lens.model_bridge.generalized_components import (
+ AttentionBridge,
+ BlockBridge,
+ EmbeddingBridge,
+ GatedMLPBridge,
+ LinearBridge,
+ PositionEmbeddingsAttentionBridge,
+ RMSNormalizationBridge,
+ RotaryEmbeddingBridge,
+ UnembeddingBridge,
+)
+from transformer_lens.model_bridge.supported_architectures.mistral import (
+ MistralArchitectureAdapter,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+
+def _make_cfg(
+ n_heads: int = 32,
+ d_model: int = 4096,
+ n_layers: int = 32,
+ d_vocab: int = 32000,
+ n_ctx: int = 4096,
+ n_key_value_heads: int | None = 8,
+ **overrides,
+) -> TransformerBridgeConfig:
+ """Return a minimal TransformerBridgeConfig for Mistral adapter tests."""
+ cfg = TransformerBridgeConfig(
+ d_model=d_model,
+ d_head=d_model // n_heads,
+ n_heads=n_heads,
+ n_layers=n_layers,
+ n_ctx=n_ctx,
+ d_vocab=d_vocab,
+ n_key_value_heads=n_key_value_heads,
+ architecture="MistralForCausalLM",
+ )
+ for k, v in overrides.items():
+ setattr(cfg, k, v)
+ return cfg
+
+
+@pytest.fixture(scope="module")
+def adapter() -> MistralArchitectureAdapter:
+ return MistralArchitectureAdapter(_make_cfg())
+
+
+# ---------------------------------------------------------------------------
+# Component mapping — top-level key set and bridge types
+# ---------------------------------------------------------------------------
+
+
+class TestMistralComponentMapping:
+ """Component mapping has the correct slots, bridge types, and HF module paths."""
+
+ def test_top_level_keys(self, adapter: MistralArchitectureAdapter) -> None:
+ assert set(adapter.component_mapping.keys()) == {
+ "embed",
+ "rotary_emb",
+ "blocks",
+ "ln_final",
+ "unembed",
+ }
+
+ def test_no_pos_embed_key(self, adapter: MistralArchitectureAdapter) -> None:
+ """Mistral uses rotary embeddings — no learned positional embedding component."""
+ assert "pos_embed" not in adapter.component_mapping
+
+ def test_bridge_types(self, adapter: MistralArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ assert isinstance(mapping["embed"], EmbeddingBridge)
+ assert isinstance(mapping["rotary_emb"], RotaryEmbeddingBridge)
+ assert isinstance(mapping["blocks"], BlockBridge)
+ assert isinstance(mapping["ln_final"], RMSNormalizationBridge)
+ assert isinstance(mapping["unembed"], UnembeddingBridge)
+
+ def test_top_level_hf_paths(self, adapter: MistralArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ assert mapping["embed"].name == "model.embed_tokens"
+ assert mapping["rotary_emb"].name == "model.rotary_emb"
+ assert mapping["blocks"].name == "model.layers"
+ assert mapping["ln_final"].name == "model.norm"
+ assert mapping["unembed"].name == "lm_head"
+
+ def test_block_submodule_keys(self, adapter: MistralArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert set(blocks.submodules.keys()) == {"ln1", "ln2", "attn", "mlp"}
+
+ def test_block_submodule_types(self, adapter: MistralArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert isinstance(blocks.submodules["ln1"], RMSNormalizationBridge)
+ assert isinstance(blocks.submodules["ln2"], RMSNormalizationBridge)
+ 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."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert not isinstance(attn, PositionEmbeddingsAttentionBridge)
+
+ def test_block_submodule_hf_paths(self, adapter: MistralArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["ln1"].name == "input_layernorm"
+ assert blocks.submodules["ln2"].name == "post_attention_layernorm"
+ assert blocks.submodules["attn"].name == "self_attn"
+ assert blocks.submodules["mlp"].name == "mlp"
+
+ def test_attn_requires_mask_and_position_embeddings(
+ self, adapter: MistralArchitectureAdapter
+ ) -> None:
+ """Mistral RoPE attention requires both attention mask and position embeddings."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.requires_attention_mask is True
+ assert attn.requires_position_embeddings is True
+
+ def test_attn_qkvo_submodule_paths(self, adapter: MistralArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert set(attn.submodules.keys()) == {"q", "k", "v", "o"}
+ assert attn.submodules["q"].name == "q_proj"
+ assert attn.submodules["k"].name == "k_proj"
+ assert attn.submodules["v"].name == "v_proj"
+ assert attn.submodules["o"].name == "o_proj"
+
+ def test_attn_qkvo_are_linear_bridges(self, adapter: MistralArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ for sub in attn.submodules.values():
+ assert isinstance(sub, LinearBridge)
+
+ def test_mlp_submodule_paths(self, adapter: MistralArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert set(mlp.submodules.keys()) == {"gate", "in", "out"}
+ assert mlp.submodules["gate"].name == "gate_proj"
+ assert mlp.submodules["in"].name == "up_proj"
+ assert mlp.submodules["out"].name == "down_proj"
+
+
+# ---------------------------------------------------------------------------
+# Anti-drift config flags
+# ---------------------------------------------------------------------------
+
+
+class TestMistralAdapterConfig:
+ """Anti-drift flags that must not silently regress."""
+
+ def test_final_rms_is_false(self, adapter: MistralArchitectureAdapter) -> None:
+ """Mistral does not use final RMSNorm — final_rms must remain False."""
+ assert adapter.cfg.final_rms is False
+
+ def test_uses_rms_norm_is_true(self, adapter: MistralArchitectureAdapter) -> None:
+ assert adapter.cfg.uses_rms_norm is True
+
+ def test_gated_mlp_is_true(self, adapter: MistralArchitectureAdapter) -> None:
+ """Mistral uses a gated SwiGLU MLP — must not silently revert to vanilla MLP."""
+ assert adapter.cfg.gated_mlp is True
+
+ def test_attn_only_is_false(self, adapter: MistralArchitectureAdapter) -> None:
+ assert adapter.cfg.attn_only is False
+
+
+# ---------------------------------------------------------------------------
+# Weight processing conversions
+# ---------------------------------------------------------------------------
+
+
+class TestMistralWeightConversions:
+ """weight_processing_conversions has exactly the expected QKVO keys."""
+
+ def test_exact_conversion_key_set(self, adapter: MistralArchitectureAdapter) -> None:
+ assert set(adapter.weight_processing_conversions.keys()) == {
+ "blocks.{i}.attn.q.weight",
+ "blocks.{i}.attn.k.weight",
+ "blocks.{i}.attn.v.weight",
+ "blocks.{i}.attn.o.weight",
+ }
+
+ def test_qkv_conversions_use_split_heads_pattern(
+ self, adapter: MistralArchitectureAdapter
+ ) -> None:
+ """'(n h) m -> n m h' splits [n_heads*d_head, d_model] → [n, d_model, d_head]."""
+ for slot in ("q", "k", "v"):
+ conv = adapter.weight_processing_conversions[f"blocks.{{i}}.attn.{slot}.weight"]
+ assert isinstance(conv, ParamProcessingConversion)
+ assert isinstance(conv.tensor_conversion, RearrangeTensorConversion)
+ assert conv.tensor_conversion.pattern == "(n h) m -> n m h"
+
+ def test_o_conversion_uses_merge_heads_pattern(
+ self, adapter: MistralArchitectureAdapter
+ ) -> None:
+ """'m (n h) -> n h m' moves n to the front for the output projection."""
+ conv = adapter.weight_processing_conversions["blocks.{i}.attn.o.weight"]
+ assert isinstance(conv, ParamProcessingConversion)
+ assert isinstance(conv.tensor_conversion, RearrangeTensorConversion)
+ assert conv.tensor_conversion.pattern == "m (n h) -> n h m"
+
+ def test_no_bias_conversion_keys(self, adapter: MistralArchitectureAdapter) -> None:
+ """Mistral has no attention biases — no bias conversion entries."""
+ keys = set(adapter.weight_processing_conversions.keys())
+ assert not any("bias" in k or ".b_" in k for k in keys)
+
+ def test_no_norm_conversion_keys(self, adapter: MistralArchitectureAdapter) -> None:
+ """RMSNorm has no bias offset — no ln1/ln2/ln_final conversion entries."""
+ keys = set(adapter.weight_processing_conversions.keys())
+ assert not any("ln" in k for k in keys)
+
+
+# ---------------------------------------------------------------------------
+# GQA support
+# ---------------------------------------------------------------------------
+
+
+class TestMistralGQASupport:
+ """n_key_value_heads must propagate to K/V conversions and leave Q/O unchanged."""
+
+ def test_no_kv_heads_falls_back_to_n_heads(self) -> None:
+ """Without n_key_value_heads, K/V fall back to n_heads."""
+ adapter = MistralArchitectureAdapter(_make_cfg(n_heads=32, n_key_value_heads=None))
+ k_conv = adapter.weight_processing_conversions["blocks.{i}.attn.k.weight"]
+ assert k_conv.tensor_conversion.axes_lengths["n"] == 32
+
+ def test_gqa_propagates_to_kv_conversions(self) -> None:
+ """With 8 KV heads, K/V conversions must use n=8."""
+ adapter = MistralArchitectureAdapter(_make_cfg(n_heads=32, n_key_value_heads=8))
+ for slot in ("k", "v"):
+ conv = adapter.weight_processing_conversions[f"blocks.{{i}}.attn.{slot}.weight"]
+ assert conv.tensor_conversion.axes_lengths["n"] == 8
+
+ def test_gqa_does_not_affect_q_conversion(self) -> None:
+ """Q always uses full n_heads regardless of GQA."""
+ adapter = MistralArchitectureAdapter(_make_cfg(n_heads=32, n_key_value_heads=8))
+ q_conv = adapter.weight_processing_conversions["blocks.{i}.attn.q.weight"]
+ assert q_conv.tensor_conversion.axes_lengths["n"] == 32
+
+ def test_gqa_does_not_affect_o_conversion(self) -> None:
+ """O projection always uses n_heads; GQA only affects K/V."""
+ adapter = MistralArchitectureAdapter(_make_cfg(n_heads=32, n_key_value_heads=8))
+ o_conv = adapter.weight_processing_conversions["blocks.{i}.attn.o.weight"]
+ assert o_conv.tensor_conversion.axes_lengths["n"] == 32
diff --git a/tests/unit/model_bridge/supported_architectures/test_mixtral_adapter.py b/tests/unit/model_bridge/supported_architectures/test_mixtral_adapter.py
index 999ca92b0..9b08dab09 100644
--- a/tests/unit/model_bridge/supported_architectures/test_mixtral_adapter.py
+++ b/tests/unit/model_bridge/supported_architectures/test_mixtral_adapter.py
@@ -392,29 +392,23 @@ def _hook(x: Any, hook: Any) -> Any:
# Identity RoPE inputs keep this test focused on hook reshaping, not rotation math.
cos = ones(1, self.SEQ, self.D_HEAD)
sin = zeros(1, self.SEQ, self.D_HEAD)
- out = attn_bridge(hidden, position_embeddings=(cos, sin))
- # The attention bridge may return either a bare tensor or an (output, ...) tuple.
- out_tensor = out[0] if isinstance(out, tuple) else out
+ attn_bridge(hidden, position_embeddings=(cos, sin))
- return captured["q"], captured["k"], captured["v"], out_tensor
+ return captured["q"], captured["k"], captured["v"]
def test_hook_q_uses_n_heads(
self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
) -> None:
- q, _, _, _ = self._run_and_capture(wired_attn_bridge)
+ q, _, _ = self._run_and_capture(wired_attn_bridge)
assert q.shape == (self.BATCH, self.SEQ, self.N_HEADS, self.D_HEAD)
def test_hook_kv_use_n_kv_heads(
self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
) -> None:
- _, k, v, _ = self._run_and_capture(wired_attn_bridge)
+ _, k, v = self._run_and_capture(wired_attn_bridge)
assert k.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
assert v.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
- def test_attn_output_shape(self, wired_attn_bridge: PositionEmbeddingsAttentionBridge) -> None:
- _, _, _, out = self._run_and_capture(wired_attn_bridge)
- assert out.shape == (self.BATCH, self.SEQ, self.D_MODEL)
-
class TestMixtralSetupComponentTesting:
"""setup_component_testing wires the shared rotary embedding and forces eager attention."""
diff --git a/tests/unit/model_bridge/supported_architectures/test_neo_adapter.py b/tests/unit/model_bridge/supported_architectures/test_neo_adapter.py
new file mode 100644
index 000000000..5b7d4aad3
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_neo_adapter.py
@@ -0,0 +1,225 @@
+"""Unit tests for NeoArchitectureAdapter.
+
+Tests cover:
+- Component mapping structure (bridge types and HF module names)
+- Weight conversion key set
+- NeoLinearTransposeConversion numerical correctness
+"""
+
+import pytest
+import torch
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge.generalized_components import (
+ AttentionBridge,
+ BlockBridge,
+ EmbeddingBridge,
+ LinearBridge,
+ MLPBridge,
+ NormalizationBridge,
+ PosEmbedBridge,
+ UnembeddingBridge,
+)
+from transformer_lens.model_bridge.supported_architectures.neo import (
+ NeoArchitectureAdapter,
+ NeoLinearTransposeConversion,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+
+def _make_cfg(
+ n_heads: int = 4,
+ d_model: int = 64,
+ n_layers: int = 2,
+ d_mlp: int = 256,
+ d_vocab: int = 1000,
+ n_ctx: int = 512,
+) -> TransformerBridgeConfig:
+ """Return a minimal TransformerBridgeConfig for Neo adapter tests."""
+ return TransformerBridgeConfig(
+ d_model=d_model,
+ d_head=d_model // n_heads,
+ n_layers=n_layers,
+ n_ctx=n_ctx,
+ n_heads=n_heads,
+ d_vocab=d_vocab,
+ d_mlp=d_mlp,
+ default_prepend_bos=True,
+ architecture="GPTNeoForCausalLM",
+ )
+
+
+@pytest.fixture
+def cfg() -> TransformerBridgeConfig:
+ return _make_cfg()
+
+
+@pytest.fixture
+def adapter(cfg: TransformerBridgeConfig) -> NeoArchitectureAdapter:
+ return NeoArchitectureAdapter(cfg)
+
+
+# ---------------------------------------------------------------------------
+# Component mapping structure tests
+# ---------------------------------------------------------------------------
+
+
+class TestNeoAdapterComponentMapping:
+ """Tests that component_mapping has the correct bridge types and HF module names."""
+
+ def test_top_level_keys(self, adapter: NeoArchitectureAdapter) -> None:
+ assert set(adapter.component_mapping.keys()) == {
+ "embed",
+ "pos_embed",
+ "blocks",
+ "ln_final",
+ "unembed",
+ }
+
+ def test_bridge_types(self, adapter: NeoArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ assert isinstance(mapping["embed"], EmbeddingBridge)
+ assert isinstance(mapping["pos_embed"], PosEmbedBridge)
+ assert isinstance(mapping["blocks"], BlockBridge)
+ assert isinstance(mapping["ln_final"], NormalizationBridge)
+ assert isinstance(mapping["unembed"], UnembeddingBridge)
+
+ def test_top_level_hf_paths(self, adapter: NeoArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ assert mapping["embed"].name == "transformer.wte"
+ assert mapping["pos_embed"].name == "transformer.wpe"
+ assert mapping["blocks"].name == "transformer.h"
+ assert mapping["ln_final"].name == "transformer.ln_f"
+ assert mapping["unembed"].name == "lm_head"
+
+ def test_no_rotary_emb_key(self, adapter: NeoArchitectureAdapter) -> None:
+ """Neo uses standard positional embeddings — no rotary embedding component."""
+ assert "rotary_emb" not in adapter.component_mapping
+
+ def test_block_submodule_keys(self, adapter: NeoArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert set(blocks.submodules.keys()) == {"ln1", "attn", "ln2", "mlp"}
+
+ def test_attention_submodule_keys(self, adapter: NeoArchitectureAdapter) -> None:
+ """Neo uses separate Q, K, V, O projections (no combined QKV matrix)."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert set(attn.submodules.keys()) == {"q", "k", "v", "o"}
+
+ def test_mlp_submodule_keys(self, adapter: NeoArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert set(mlp.submodules.keys()) == {"in", "out"}
+
+ def test_block_bridge_types(self, adapter: NeoArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert isinstance(blocks.submodules["ln1"], NormalizationBridge)
+ assert isinstance(blocks.submodules["attn"], AttentionBridge)
+ assert isinstance(blocks.submodules["ln2"], NormalizationBridge)
+ assert isinstance(blocks.submodules["mlp"], MLPBridge)
+
+ def test_attention_hf_paths(self, adapter: NeoArchitectureAdapter) -> None:
+ """Neo's attention is nested as attn.attention in HuggingFace."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.name == "attn.attention"
+ assert attn.submodules["q"].name == "q_proj"
+ assert attn.submodules["k"].name == "k_proj"
+ assert attn.submodules["v"].name == "v_proj"
+ assert attn.submodules["o"].name == "out_proj"
+
+ def test_block_hf_paths(self, adapter: NeoArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["ln1"].name == "ln_1"
+ assert blocks.submodules["ln2"].name == "ln_2"
+ assert blocks.submodules["mlp"].name == "mlp"
+ assert blocks.submodules["mlp"].submodules["in"].name == "c_fc"
+ assert blocks.submodules["mlp"].submodules["out"].name == "c_proj"
+
+ def test_linear_submodule_bridge_types(self, adapter: NeoArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ for submodule in [*attn.submodules.values(), *mlp.submodules.values()]:
+ assert isinstance(submodule, LinearBridge)
+
+
+# ---------------------------------------------------------------------------
+# Weight processing conversion tests
+# ---------------------------------------------------------------------------
+
+
+class TestNeoAdapterWeightConversions:
+ """Tests that weight_processing_conversions has exactly the expected keys.
+
+ Neo uses standard PyTorch Linear layers whose weights are stored as
+ [out_features, in_features], requiring a transpose to Conv1D format for
+ attention heads and an optional einops rearrangement for head dimensions.
+ """
+
+ def test_exact_conversion_key_set(self, adapter: NeoArchitectureAdapter) -> None:
+ assert set(adapter.weight_processing_conversions.keys()) == {
+ "blocks.{i}.attn.q.weight",
+ "blocks.{i}.attn.k.weight",
+ "blocks.{i}.attn.v.weight",
+ "blocks.{i}.attn.o.weight",
+ "blocks.{i}.mlp.in.weight",
+ "blocks.{i}.mlp.out.weight",
+ "blocks.{i}.attn.q.bias",
+ "blocks.{i}.attn.k.bias",
+ "blocks.{i}.attn.v.bias",
+ }
+
+
+# ---------------------------------------------------------------------------
+# NeoLinearTransposeConversion — numerical correctness tests
+# ---------------------------------------------------------------------------
+
+
+class TestNeoLinearTransposeConversion:
+ """Numerical correctness of Neo's Linear weight transposition."""
+
+ D_MODEL, N_HEADS, D_HEAD = 64, 4, 16 # D_MODEL = N_HEADS * D_HEAD
+
+ def test_transpose_only_roundtrips(self) -> None:
+ """A weight transposed and reverted should recover the original."""
+ torch.manual_seed(0)
+ conv = NeoLinearTransposeConversion()
+ original = torch.randn(self.D_MODEL, self.D_MODEL)
+ reverted = conv.revert(conv.handle_conversion(original))
+ assert reverted.shape == original.shape
+ assert torch.allclose(original, reverted)
+
+ def test_transpose_changes_shape(self) -> None:
+ """handle_conversion transposes [out, in] -> [in, out]."""
+ w = torch.zeros(128, 64) # [out_features, in_features]
+ out = NeoLinearTransposeConversion().handle_conversion(w)
+ assert out.shape == (64, 128)
+
+ def test_transpose_with_rearrange_q_weight(self) -> None:
+ """Q/K/V weight: [d_model, n*d_head] -> transpose -> rearrange to [n, d_model, d_head]."""
+ conv = NeoLinearTransposeConversion("d_model (n h) -> n d_model h", n=self.N_HEADS)
+ w = torch.randn(self.D_MODEL, self.N_HEADS * self.D_HEAD)
+ out = conv.handle_conversion(w)
+ assert out.shape == (self.N_HEADS, self.D_MODEL, self.D_HEAD)
+
+ def test_transpose_with_rearrange_o_weight(self) -> None:
+ """O weight: [n*d_head, d_model] -> transpose -> rearrange to [n, d_head, d_model]."""
+ conv = NeoLinearTransposeConversion("(n h) d_model -> n h d_model", n=self.N_HEADS)
+ w = torch.randn(self.N_HEADS * self.D_HEAD, self.D_MODEL)
+ out = conv.handle_conversion(w)
+ assert out.shape == (self.N_HEADS, self.D_HEAD, self.D_MODEL)
+
+ def test_rearrange_roundtrip(self) -> None:
+ """handle_conversion -> revert recovers the original weight for Q projection."""
+ torch.manual_seed(1)
+ conv = NeoLinearTransposeConversion("d_model (n h) -> n d_model h", n=self.N_HEADS)
+ original = torch.randn(self.D_MODEL, self.N_HEADS * self.D_HEAD)
+ recovered = conv.revert(conv.handle_conversion(original))
+ assert recovered.shape == original.shape
+ assert torch.allclose(original, recovered, atol=1e-6)
+
+ def test_values_preserved_after_transpose(self) -> None:
+ """Values should be identical after transpose (not just shape)."""
+ w = torch.arange(12, dtype=torch.float).reshape(3, 4)
+ out = NeoLinearTransposeConversion().handle_conversion(w)
+ assert torch.allclose(out, w.T)
diff --git a/tests/unit/model_bridge/supported_architectures/test_neox_adapter.py b/tests/unit/model_bridge/supported_architectures/test_neox_adapter.py
new file mode 100644
index 000000000..6bf4324cc
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_neox_adapter.py
@@ -0,0 +1,296 @@
+"""Unit tests for NeoxArchitectureAdapter.
+
+Tests cover:
+- Component mapping structure (bridge types and HF module names)
+- Weight conversion key set and shared source keys
+- setup_component_testing rotary embedding wiring
+- split_qkv_matrix interleaved QKV partition correctness
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge.generalized_components import (
+ EmbeddingBridge,
+ JointQKVPositionEmbeddingsAttentionBridge,
+ LinearBridge,
+ MLPBridge,
+ NormalizationBridge,
+ ParallelBlockBridge,
+ RotaryEmbeddingBridge,
+ UnembeddingBridge,
+)
+from transformer_lens.model_bridge.supported_architectures.neox import (
+ NeoxArchitectureAdapter,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+
+def _make_cfg(
+ n_heads: int = 4,
+ d_model: int = 64,
+ n_layers: int = 2,
+ d_mlp: int = 256,
+ d_vocab: int = 1000,
+ n_ctx: int = 512,
+) -> TransformerBridgeConfig:
+ """Return a minimal TransformerBridgeConfig for NeoX adapter tests."""
+ return TransformerBridgeConfig(
+ d_model=d_model,
+ d_head=d_model // n_heads,
+ n_layers=n_layers,
+ n_ctx=n_ctx,
+ n_heads=n_heads,
+ d_vocab=d_vocab,
+ d_mlp=d_mlp,
+ default_prepend_bos=False,
+ architecture="GPTNeoXForCausalLM",
+ )
+
+
+@pytest.fixture
+def cfg() -> TransformerBridgeConfig:
+ return _make_cfg()
+
+
+@pytest.fixture
+def adapter(cfg: TransformerBridgeConfig) -> NeoxArchitectureAdapter:
+ return NeoxArchitectureAdapter(cfg)
+
+
+def _fake_hf_model(rotary_emb: object) -> SimpleNamespace:
+ return SimpleNamespace(gpt_neox=SimpleNamespace(rotary_emb=rotary_emb))
+
+
+class DummyAttention:
+ def __init__(self) -> None:
+ self.rotary_emb = None
+
+ def set_rotary_emb(self, rotary_emb: object) -> None:
+ self.rotary_emb = rotary_emb
+
+
+class DummyBlock:
+ def __init__(self, has_attention: bool = True) -> None:
+ if has_attention:
+ self.attn = DummyAttention()
+
+
+class DummyBridgeModel:
+ def __init__(self, blocks: list[DummyBlock]) -> None:
+ self.blocks = blocks
+
+
+# ---------------------------------------------------------------------------
+# Component mapping structure tests
+# ---------------------------------------------------------------------------
+
+
+class TestNeoxAdapterComponentMapping:
+ """Tests that component_mapping has the correct bridge types and HF module names."""
+
+ def test_top_level_keys(self, adapter: NeoxArchitectureAdapter) -> None:
+ assert set(adapter.component_mapping.keys()) == {
+ "embed",
+ "rotary_emb",
+ "blocks",
+ "ln_final",
+ "unembed",
+ }
+
+ def test_bridge_types(self, adapter: NeoxArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ blocks = mapping["blocks"]
+ assert isinstance(mapping["embed"], EmbeddingBridge)
+ assert isinstance(mapping["rotary_emb"], RotaryEmbeddingBridge)
+ assert isinstance(blocks, ParallelBlockBridge)
+ assert isinstance(mapping["ln_final"], NormalizationBridge)
+ assert isinstance(mapping["unembed"], UnembeddingBridge)
+
+ def test_top_level_hf_paths(self, adapter: NeoxArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ assert mapping["embed"].name == "gpt_neox.embed_in"
+ assert mapping["rotary_emb"].name == "gpt_neox.rotary_emb"
+ assert mapping["blocks"].name == "gpt_neox.layers"
+ assert mapping["ln_final"].name == "gpt_neox.final_layer_norm"
+ assert mapping["unembed"].name == "embed_out"
+
+ def test_no_pos_embed_key(self, adapter: NeoxArchitectureAdapter) -> None:
+ """NeoX uses rotary embeddings — no learned positional embedding component."""
+ assert "pos_embed" not in adapter.component_mapping
+
+ def test_block_submodule_keys(self, adapter: NeoxArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert set(blocks.submodules.keys()) == {"ln1", "ln2", "attn", "mlp"}
+
+ def test_attention_submodule_keys(self, adapter: NeoxArchitectureAdapter) -> None:
+ """NeoX uses a combined QKV projection alongside derived q/k/v split bridges."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert set(attn.submodules.keys()) == {"qkv", "q", "k", "v", "o"}
+
+ def test_mlp_submodule_keys(self, adapter: NeoxArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert set(mlp.submodules.keys()) == {"in", "out"}
+
+ def test_block_bridge_types(self, adapter: NeoxArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert isinstance(blocks.submodules["ln1"], NormalizationBridge)
+ assert isinstance(blocks.submodules["ln2"], NormalizationBridge)
+ assert isinstance(blocks.submodules["attn"], JointQKVPositionEmbeddingsAttentionBridge)
+ assert isinstance(blocks.submodules["mlp"], MLPBridge)
+
+ def test_attention_hf_paths(self, adapter: NeoxArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.name == "attention"
+ assert attn.submodules["qkv"].name == "query_key_value"
+ assert attn.submodules["o"].name == "dense"
+
+ def test_attn_requires_attention_mask(self, adapter: NeoxArchitectureAdapter) -> None:
+ """GPTNeoX requires an explicit attention mask."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.requires_attention_mask is True
+
+ def test_block_hf_paths(self, adapter: NeoxArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["ln1"].name == "input_layernorm"
+ assert blocks.submodules["ln2"].name == "post_attention_layernorm"
+ assert blocks.submodules["mlp"].name == "mlp"
+ assert blocks.submodules["mlp"].submodules["in"].name == "dense_h_to_4h"
+ assert blocks.submodules["mlp"].submodules["out"].name == "dense_4h_to_h"
+
+ def test_linear_submodule_bridge_types(self, adapter: NeoxArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ for submodule in [*attn.submodules.values(), *mlp.submodules.values()]:
+ assert isinstance(submodule, LinearBridge)
+
+
+# ---------------------------------------------------------------------------
+# Weight processing conversion tests
+# ---------------------------------------------------------------------------
+
+
+class TestNeoxAdapterWeightConversions:
+ """Tests that weight_processing_conversions has the expected key set and source keys.
+
+ NeoX stores Q, K, V (weights and biases) in a single interleaved matrix
+ gpt_neox.layers.{i}.attention.query_key_value.{weight,bias}.
+ All three projections share the same source key — each conversion extracts
+ its slice via SplitTensorConversion.
+ """
+
+ def test_exact_conversion_key_set(self, adapter: NeoxArchitectureAdapter) -> None:
+ assert set(adapter.weight_processing_conversions.keys()) == {
+ "blocks.{i}.attn.q",
+ "blocks.{i}.attn.k",
+ "blocks.{i}.attn.v",
+ "blocks.{i}.attn.b_Q",
+ "blocks.{i}.attn.b_K",
+ "blocks.{i}.attn.b_V",
+ "blocks.{i}.attn.o",
+ }
+
+ def test_qkv_weights_share_source_key(self, adapter: NeoxArchitectureAdapter) -> None:
+ """Q, K, V weights all come from the same interleaved QKV matrix."""
+ expected = "gpt_neox.layers.{i}.attention.query_key_value.weight"
+ for key in ("blocks.{i}.attn.q", "blocks.{i}.attn.k", "blocks.{i}.attn.v"):
+ assert adapter.weight_processing_conversions[key].source_key == expected
+
+ def test_qkv_biases_share_source_key(self, adapter: NeoxArchitectureAdapter) -> None:
+ """Q, K, V biases all come from the same interleaved QKV bias vector."""
+ expected = "gpt_neox.layers.{i}.attention.query_key_value.bias"
+ for key in ("blocks.{i}.attn.b_Q", "blocks.{i}.attn.b_K", "blocks.{i}.attn.b_V"):
+ assert adapter.weight_processing_conversions[key].source_key == expected
+
+ def test_o_projection_source_key(self, adapter: NeoxArchitectureAdapter) -> None:
+ expected = "gpt_neox.layers.{i}.attention.dense.weight"
+ assert adapter.weight_processing_conversions["blocks.{i}.attn.o"].source_key == expected
+
+
+# ---------------------------------------------------------------------------
+# setup_component_testing — rotary embedding wiring
+# ---------------------------------------------------------------------------
+
+
+class TestNeoxSetupComponentTesting:
+ """setup_component_testing must wire NeoX's rotary embedding into attention bridges."""
+
+ def test_sets_rotary_emb_on_bridge_model_blocks(self, adapter: NeoxArchitectureAdapter) -> None:
+ rotary_emb = object()
+ bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(), DummyBlock()])
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model)
+
+ for block in bridge_model.blocks:
+ assert block.attn.rotary_emb is rotary_emb
+
+ def test_skips_bridge_blocks_without_attention(self, adapter: NeoxArchitectureAdapter) -> None:
+ rotary_emb = object()
+ bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(has_attention=False)])
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model)
+
+ assert bridge_model.blocks[0].attn.rotary_emb is rotary_emb
+
+ def test_no_bridge_model_does_not_raise(self, adapter: NeoxArchitectureAdapter) -> None:
+ """setup_component_testing without a bridge_model should not raise."""
+ adapter.setup_component_testing(_fake_hf_model(object()))
+
+
+# ---------------------------------------------------------------------------
+# split_qkv_matrix — interleaved QKV partition
+# ---------------------------------------------------------------------------
+
+
+class TestNeoxSplitQkvMatrix:
+ """split_qkv_matrix correctly partitions the interleaved [Q_h0,K_h0,V_h0,...] format."""
+
+ def _make_attention_component(self, n_heads: int, d_head: int, d_model: int) -> SimpleNamespace:
+ """Build a mock attention layer with a known interleaved QKV weight and bias."""
+ total_out = n_heads * 3 * d_head
+ qkv_weight = torch.arange(total_out * d_model, dtype=torch.float32).reshape(
+ total_out, d_model
+ )
+ qkv_bias = torch.arange(total_out, dtype=torch.float32)
+ qkv_linear = SimpleNamespace(weight=qkv_weight, bias=qkv_bias)
+ return SimpleNamespace(query_key_value=qkv_linear)
+
+ def test_weight_partition_matches_interleaved_format(self) -> None:
+ """Q/K/V weight slices must follow the NeoX per-head interleaved layout."""
+ n_heads, d_head, d_model = 4, 16, 64
+ adapter = NeoxArchitectureAdapter(_make_cfg(n_heads=n_heads, d_model=d_model))
+ component = self._make_attention_component(n_heads, d_head, d_model)
+
+ W = component.query_key_value.weight.view(n_heads, 3 * d_head, d_model)
+ expected_q = W[:, :d_head, :].reshape(n_heads * d_head, d_model)
+ expected_k = W[:, d_head : 2 * d_head, :].reshape(n_heads * d_head, d_model)
+ expected_v = W[:, 2 * d_head :, :].reshape(n_heads * d_head, d_model)
+
+ q_lin, k_lin, v_lin = adapter.split_qkv_matrix(component)
+
+ assert torch.equal(q_lin.weight.data, expected_q)
+ assert torch.equal(k_lin.weight.data, expected_k)
+ assert torch.equal(v_lin.weight.data, expected_v)
+
+ def test_bias_partition_matches_interleaved_format(self) -> None:
+ """Q/K/V bias slices must follow the NeoX per-head interleaved layout."""
+ n_heads, d_head, d_model = 4, 16, 64
+ adapter = NeoxArchitectureAdapter(_make_cfg(n_heads=n_heads, d_model=d_model))
+ component = self._make_attention_component(n_heads, d_head, d_model)
+
+ b = component.query_key_value.bias.view(n_heads, 3 * d_head)
+ expected_bq = b[:, :d_head].reshape(n_heads * d_head)
+ expected_bk = b[:, d_head : 2 * d_head].reshape(n_heads * d_head)
+ expected_bv = b[:, 2 * d_head :].reshape(n_heads * d_head)
+
+ q_lin, k_lin, v_lin = adapter.split_qkv_matrix(component)
+
+ assert torch.equal(q_lin.bias.data, expected_bq)
+ assert torch.equal(k_lin.bias.data, expected_bk)
+ assert torch.equal(v_lin.bias.data, expected_bv)
diff --git a/tests/unit/model_bridge/supported_architectures/test_olmo2_adapter.py b/tests/unit/model_bridge/supported_architectures/test_olmo2_adapter.py
new file mode 100644
index 000000000..e9c7a6e46
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_olmo2_adapter.py
@@ -0,0 +1,465 @@
+"""Unit tests for the Olmo2ArchitectureAdapter (download-free, tiny programmatic configs
+plus small synthetic tensors and a fake attention module, no real checkpoints).
+
+Covered:
+- Anti-drift config: supports_fold_ln=False (post-norm cannot fold).
+- Weight conversions: the adapter uses the base QKVO helper, so only the exact
+ key set is asserted at the adapter layer.
+- Component-mapping structure, bridge types, and HF module paths.
+- Post-norm block wiring: ln1 maps to post_attention_layernorm and ln2 maps to
+ post_feedforward_layernorm. This is the central arch-specific decision.
+- Q/K-norm submodules under attention with the correct HF names.
+- hook_alias_overrides: hook_resid_mid points at mlp.hook_in (the true post-attn
+ pre-mlp residual under post-norm), overriding BlockBridge's default ln2.hook_in.
+- GQA forward hook shapes for Q and K/V with Q/K-norm wired into the fake attention.
+- setup_component_testing: rotary wiring on template and bridge-model attentions,
+ plus forcing eager attention on the HF model and per-layer self_attn configs.
+- Architecture guards.
+"""
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch.nn as nn
+from torch import ones, randn, zeros
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge.generalized_components import (
+ BlockBridge,
+ EmbeddingBridge,
+ GatedMLPBridge,
+ LinearBridge,
+ PositionEmbeddingsAttentionBridge,
+ RMSNormalizationBridge,
+ RotaryEmbeddingBridge,
+ UnembeddingBridge,
+)
+from transformer_lens.model_bridge.supported_architectures.olmo2 import (
+ Olmo2ArchitectureAdapter,
+)
+
+
+@pytest.fixture(scope="class")
+def cfg() -> TransformerBridgeConfig:
+ return TransformerBridgeConfig(
+ d_model=64,
+ d_head=16,
+ n_layers=2,
+ n_ctx=128,
+ n_heads=4,
+ n_key_value_heads=2,
+ d_vocab=256,
+ architecture="Olmo2ForCausalLM",
+ )
+
+
+@pytest.fixture(scope="class")
+def adapter(cfg: TransformerBridgeConfig) -> Olmo2ArchitectureAdapter:
+ return Olmo2ArchitectureAdapter(cfg)
+
+
+def _cfg(*, n_key_value_heads: int | None = 2) -> TransformerBridgeConfig:
+ # Keep dimensions tiny so adapter tests do not need HF downloads or real checkpoints.
+ return TransformerBridgeConfig(
+ d_model=64,
+ d_head=16,
+ n_layers=2,
+ n_ctx=128,
+ n_heads=4,
+ n_key_value_heads=n_key_value_heads,
+ d_vocab=256,
+ architecture="Olmo2ForCausalLM",
+ )
+
+
+def _mapping(adapter: Olmo2ArchitectureAdapter) -> dict:
+ """Narrow component_mapping (Optional on the base class) to a non-None dict."""
+ mapping = adapter.component_mapping
+ assert mapping is not None
+ return mapping
+
+
+def _conversions(adapter: Olmo2ArchitectureAdapter) -> dict:
+ """weight_processing_conversions is Optional on the base class, assert it is populated."""
+ conversions = adapter.weight_processing_conversions
+ assert conversions is not None
+ return conversions
+
+
+def _fake_hf_model(rotary_emb: object) -> SimpleNamespace:
+ """Minimal HF model exposing only model.rotary_emb (no config, no layers)."""
+ return SimpleNamespace(model=SimpleNamespace(rotary_emb=rotary_emb))
+
+
+def _fake_hf_model_with_eager_targets(rotary_emb: object) -> SimpleNamespace:
+ """HF model whose top-level and per-layer attention implementation start non-eager."""
+ layers = [
+ SimpleNamespace(
+ self_attn=SimpleNamespace(config=SimpleNamespace(_attn_implementation="sdpa"))
+ )
+ for _ in range(2)
+ ]
+ return SimpleNamespace(
+ config=SimpleNamespace(_attn_implementation="sdpa"),
+ model=SimpleNamespace(rotary_emb=rotary_emb, layers=layers),
+ )
+
+
+class DummyAttention:
+ def __init__(self) -> None:
+ self.rotary_emb = None
+
+ def set_rotary_emb(self, rotary_emb: object) -> None:
+ self.rotary_emb = rotary_emb
+
+
+class DummyBlock:
+ def __init__(self, has_attention: bool = True) -> None:
+ if has_attention:
+ self.attn = DummyAttention()
+
+
+class DummyBridgeModel:
+ def __init__(self, blocks: list[DummyBlock]) -> None:
+ self.blocks = blocks
+
+
+class FakeOlmo2Attention(nn.Module):
+ """Minimal OLMo-2-style attention module for adapter hook-shape tests.
+
+ OLMo 2 has no attention bias and applies RMSNorm to the flattened Q and K
+ projections (pre-reshape phase): q_norm over n_heads * head_dim and k_norm
+ over n_key_value_heads * head_dim. Matches HF's Olmo2Attention shape.
+ """
+
+ def __init__(self, cfg: TransformerBridgeConfig) -> None:
+ super().__init__()
+ # PositionEmbeddingsAttentionBridge reads these HF-style attributes during forward.
+ self.head_dim = cfg.d_head
+ self.num_key_value_groups = cfg.n_heads // (cfg.n_key_value_heads or cfg.n_heads)
+ self.scaling = cfg.d_head**-0.5
+ self.attention_dropout = 0.0
+
+ n_kv = cfg.n_key_value_heads or cfg.n_heads
+ kv_width = n_kv * cfg.d_head
+ self.q_proj = nn.Linear(cfg.d_model, cfg.n_heads * cfg.d_head, bias=False)
+ self.k_proj = nn.Linear(cfg.d_model, kv_width, bias=False)
+ self.v_proj = nn.Linear(cfg.d_model, kv_width, bias=False)
+ self.o_proj = nn.Linear(cfg.n_heads * cfg.d_head, cfg.d_model, bias=False)
+ # Pre-reshape RMSNorm over the flattened head dimension.
+ self.q_norm = nn.RMSNorm(cfg.n_heads * cfg.d_head)
+ self.k_norm = nn.RMSNorm(kv_width)
+
+
+class TestOlmo2AdapterConfig:
+ """Anti-drift config: post-norm forces supports_fold_ln=False because folding
+ a norm that runs AFTER attention/MLP would corrupt the weights."""
+
+ def test_supports_fold_ln_is_false(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ """OLMo 2 is post-norm: RMSNorm applies after attention/MLP, not before.
+ Folding LN into QKV/MLP weights would be incorrect."""
+ assert adapter.supports_fold_ln is False
+
+
+class TestOlmo2WeightConversions:
+ """The adapter uses `self._qkvo_weight_conversions()` from the base helper with no
+ overrides. Per the unit-test guide, the rearrange patterns and the GQA n_kv_heads
+ axis are the base helper's responsibility and are covered by base-class tests.
+ The adapter-owned decision here is the exact set of conversion keys: four QKVO
+ weights, no biases, no extras."""
+
+ def test_conversion_keys_are_exactly_qkvo_weights(
+ self, adapter: Olmo2ArchitectureAdapter
+ ) -> None:
+ assert set(_conversions(adapter).keys()) == {
+ "blocks.{i}.attn.q.weight",
+ "blocks.{i}.attn.k.weight",
+ "blocks.{i}.attn.v.weight",
+ "blocks.{i}.attn.o.weight",
+ }
+
+
+class TestOlmo2ComponentMapping:
+ """Structure of the component mapping: required keys and HF module paths."""
+
+ def test_has_required_top_level_keys(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ mapping = _mapping(adapter)
+ for key in ("embed", "rotary_emb", "blocks", "ln_final", "unembed"):
+ assert key in mapping, f"Missing top-level key: {key!r}"
+
+ def test_top_level_hf_paths(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ mapping = _mapping(adapter)
+ assert mapping["embed"].name == "model.embed_tokens"
+ assert mapping["rotary_emb"].name == "model.rotary_emb"
+ assert mapping["blocks"].name == "model.layers"
+ assert mapping["ln_final"].name == "model.norm"
+ assert mapping["unembed"].name == "lm_head"
+
+
+class TestOlmo2ComponentTypes:
+ """Bridge classes selected for each component slot."""
+
+ def test_rotary_emb_is_rotary_bridge(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ assert isinstance(_mapping(adapter)["rotary_emb"], RotaryEmbeddingBridge)
+
+ def test_blocks_is_block_bridge(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ assert isinstance(_mapping(adapter)["blocks"], BlockBridge)
+
+ def test_ln_final_is_rms_normalization_bridge(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ assert isinstance(_mapping(adapter)["ln_final"], RMSNormalizationBridge)
+
+ def test_block_attn_is_position_embeddings_bridge(
+ self, adapter: Olmo2ArchitectureAdapter
+ ) -> None:
+ block = _mapping(adapter)["blocks"]
+ assert isinstance(block.submodules["attn"], PositionEmbeddingsAttentionBridge)
+
+ def test_block_mlp_is_gated_mlp_bridge(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ block = _mapping(adapter)["blocks"]
+ assert isinstance(block.submodules["mlp"], GatedMLPBridge)
+
+ def test_block_norms_are_rms(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ block = _mapping(adapter)["blocks"]
+ assert isinstance(block.submodules["ln1"], RMSNormalizationBridge)
+ assert isinstance(block.submodules["ln2"], RMSNormalizationBridge)
+
+ def test_embed_and_unembed_are_correct_bridge_types(
+ self, adapter: Olmo2ArchitectureAdapter
+ ) -> None:
+ mapping = _mapping(adapter)
+ assert isinstance(mapping["embed"], EmbeddingBridge)
+ assert isinstance(mapping["unembed"], UnembeddingBridge)
+
+ def test_attn_q_k_v_o_are_linear_bridges(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ attn = _mapping(adapter)["blocks"].submodules["attn"]
+ for slot in ("q", "k", "v", "o"):
+ assert isinstance(attn.submodules[slot], LinearBridge)
+
+ def test_mlp_gate_in_out_are_linear_bridges(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ mlp = _mapping(adapter)["blocks"].submodules["mlp"]
+ for slot in ("gate", "in", "out"):
+ assert isinstance(mlp.submodules[slot], LinearBridge)
+
+
+class TestOlmo2PostNormBlockWiring:
+ """OLMo 2 is post-norm: RMSNorm applies AFTER attention (ln1) and AFTER MLP (ln2).
+ The HF module names diverge from the pre-norm Llama family default, where ln1
+ would be `input_layernorm` and ln2 would be `post_attention_layernorm`. This is
+ the central arch-specific decision and the single test most likely to catch a
+ porting regression."""
+
+ def test_ln1_maps_to_post_attention_layernorm(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ block = _mapping(adapter)["blocks"]
+ assert block.submodules["ln1"].name == "post_attention_layernorm"
+
+ def test_ln2_maps_to_post_feedforward_layernorm(
+ self, adapter: Olmo2ArchitectureAdapter
+ ) -> None:
+ block = _mapping(adapter)["blocks"]
+ assert block.submodules["ln2"].name == "post_feedforward_layernorm"
+
+ def test_block_submodule_set(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ """Exact block submodule set: sequential transformer with both ln1 and ln2."""
+ block = _mapping(adapter)["blocks"]
+ assert set(block.submodules.keys()) == {"ln1", "ln2", "attn", "mlp"}
+
+
+class TestOlmo2QKNormStructure:
+ """OLMo 2 applies a pre-reshape RMSNorm to the flattened Q and K projection outputs
+ (`q_norm(q_proj(x))`, then reshape into heads). The adapter exposes those norms as
+ `q_norm` and `k_norm` submodules under the attention bridge, with HF names matching."""
+
+ def test_attn_submodule_set(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ attn = _mapping(adapter)["blocks"].submodules["attn"]
+ assert set(attn.submodules.keys()) == {"q", "k", "v", "o", "q_norm", "k_norm"}
+
+ def test_q_norm_is_rms_normalization_bridge(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ attn = _mapping(adapter)["blocks"].submodules["attn"]
+ assert isinstance(attn.submodules["q_norm"], RMSNormalizationBridge)
+
+ def test_k_norm_is_rms_normalization_bridge(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ attn = _mapping(adapter)["blocks"].submodules["attn"]
+ assert isinstance(attn.submodules["k_norm"], RMSNormalizationBridge)
+
+ def test_q_norm_uses_q_norm_hf_name(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ attn = _mapping(adapter)["blocks"].submodules["attn"]
+ assert attn.submodules["q_norm"].name == "q_norm"
+
+ def test_k_norm_uses_k_norm_hf_name(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ attn = _mapping(adapter)["blocks"].submodules["attn"]
+ assert attn.submodules["k_norm"].name == "k_norm"
+
+
+class TestOlmo2HookAliasOverrides:
+ """Under post-norm, `ln2.hook_in` no longer captures the residual between
+ attention and MLP (ln2 applies AFTER MLP, so ln2.hook_in is the MLP output).
+ The true post-attn pre-mlp residual is at `mlp.hook_in`. The adapter overrides
+ `hook_resid_mid` accordingly. BlockBridge's default would point at the wrong
+ tensor without this override."""
+
+ def test_hook_resid_mid_points_at_mlp_hook_in(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ block = _mapping(adapter)["blocks"]
+ assert block.hook_aliases["hook_resid_mid"] == "mlp.hook_in"
+
+ def test_override_differs_from_blockbridge_default(
+ self, adapter: Olmo2ArchitectureAdapter
+ ) -> None:
+ """Asserts the override is load-bearing: the BlockBridge class default for
+ hook_resid_mid is `ln2.hook_in` (correct for pre-norm), which would be wrong
+ here. Catches a regression where the override gets dropped."""
+ block = _mapping(adapter)["blocks"]
+ assert block.hook_aliases["hook_resid_mid"] != BlockBridge.hook_aliases["hook_resid_mid"]
+
+
+class TestOlmo2GQAHookShapes:
+ """Wire a fake attention module into the bridge and verify GQA hook shapes.
+
+ Q must surface n_heads while K/V surface n_key_value_heads, which is the whole
+ point of grouped-query attention. The fake carries OLMo 2's pre-reshape Q/K norms
+ so the bridge takes its Q/K-norm code path.
+ """
+
+ N_HEADS = 4
+ N_KV_HEADS = 2
+ D_MODEL = 64
+ D_HEAD = D_MODEL // N_HEADS
+ BATCH = 2
+ SEQ = 8
+
+ @pytest.fixture
+ def wired_attn_bridge(self) -> PositionEmbeddingsAttentionBridge:
+ adapter = Olmo2ArchitectureAdapter(_cfg(n_key_value_heads=self.N_KV_HEADS))
+ fake_attn = FakeOlmo2Attention(adapter.cfg)
+ attn_bridge = _mapping(adapter)["blocks"].submodules["attn"]
+ assert isinstance(attn_bridge, PositionEmbeddingsAttentionBridge)
+ attn_bridge.set_original_component(fake_attn)
+ # A full TransformerBridge build materializes these child bridge modules for us.
+ # This unit test wires them by hand so it can stay download-free.
+ for name, original in {
+ "q": fake_attn.q_proj,
+ "k": fake_attn.k_proj,
+ "v": fake_attn.v_proj,
+ "o": fake_attn.o_proj,
+ "q_norm": fake_attn.q_norm,
+ "k_norm": fake_attn.k_norm,
+ }.items():
+ submodule = attn_bridge.submodules[name]
+ submodule.set_original_component(original)
+ attn_bridge.add_module(name, submodule)
+ attn_bridge.setup_hook_compatibility()
+ return attn_bridge
+
+ def _run_and_capture(self, attn_bridge: PositionEmbeddingsAttentionBridge) -> tuple:
+ captured: dict = {}
+
+ def _capture(name: str) -> Any:
+ def _hook(x: Any, hook: Any) -> Any:
+ captured[name] = x.detach()
+ return x
+
+ return _hook
+
+ attn_bridge.q.hook_out.add_hook(_capture("q"))
+ attn_bridge.k.hook_out.add_hook(_capture("k"))
+ attn_bridge.v.hook_out.add_hook(_capture("v"))
+
+ hidden = randn(self.BATCH, self.SEQ, self.D_MODEL)
+ # Identity RoPE inputs keep this test focused on hook reshaping, not rotation math.
+ cos = ones(1, self.SEQ, self.D_HEAD)
+ sin = zeros(1, self.SEQ, self.D_HEAD)
+ attn_bridge(hidden, position_embeddings=(cos, sin))
+
+ return captured["q"], captured["k"], captured["v"]
+
+ def test_hook_q_uses_n_heads(
+ self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
+ ) -> None:
+ q, _, _ = self._run_and_capture(wired_attn_bridge)
+ assert q.shape == (self.BATCH, self.SEQ, self.N_HEADS, self.D_HEAD)
+
+ def test_hook_kv_use_n_kv_heads(
+ self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
+ ) -> None:
+ _, k, v = self._run_and_capture(wired_attn_bridge)
+ assert k.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
+ assert v.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
+
+
+class TestOlmo2SetupComponentTesting:
+ """setup_component_testing wires the shared rotary embedding onto the template
+ attention bridge and onto each bridge-model block's attention. It also forces
+ eager attention on the HF model (top-level config and per-layer self_attn.config)
+ for numerical parity with the bridge's eager-mode reimplementation."""
+
+ def test_sets_rotary_emb_on_template_attention(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ rotary_emb = object()
+ attn_template = adapter.get_generalized_component("blocks.0.attn")
+ assert isinstance(attn_template, PositionEmbeddingsAttentionBridge)
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb))
+
+ assert attn_template._rotary_emb is rotary_emb
+
+ def test_sets_rotary_emb_on_each_bridge_model_attention(
+ self, adapter: Olmo2ArchitectureAdapter
+ ) -> None:
+ rotary_emb = object()
+ bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(), DummyBlock()])
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model)
+
+ for block in bridge_model.blocks:
+ assert block.attn.rotary_emb is rotary_emb
+
+ def test_skips_bridge_blocks_without_attention(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ rotary_emb = object()
+ bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(has_attention=False)])
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model)
+
+ assert bridge_model.blocks[0].attn.rotary_emb is rotary_emb
+
+ def test_forces_eager_attention_implementation(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ """Bridge attention only matches HF under eager attention, so it is forced on
+ at both the top-level config and on each per-layer self_attn.config."""
+ hf_model = _fake_hf_model_with_eager_targets(object())
+
+ adapter.setup_component_testing(hf_model)
+
+ assert hf_model.config._attn_implementation == "eager"
+ for layer in hf_model.model.layers:
+ assert layer.self_attn.config._attn_implementation == "eager"
+
+ def test_tolerates_minimal_hf_model_without_config_or_layers(
+ self, adapter: Olmo2ArchitectureAdapter
+ ) -> None:
+ """The defensive hasattr branches must not raise when config/layers are absent."""
+ rotary_emb = object()
+ # _fake_hf_model exposes only model.rotary_emb (no config, no layers).
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb))
+
+ attn_template = adapter.get_generalized_component("blocks.0.attn")
+ assert isinstance(attn_template, PositionEmbeddingsAttentionBridge)
+ assert attn_template._rotary_emb is rotary_emb
+
+
+class TestOlmo2ArchitectureGuards:
+ """Guards against drift from OLMo-2 conventions."""
+
+ def test_no_normalization_weights_in_conversions(
+ self, adapter: Olmo2ArchitectureAdapter
+ ) -> None:
+ """Post-norm prevents LN fold (supports_fold_ln=False), so no norm weights
+ appear in the conversion map."""
+ for key in _conversions(adapter):
+ assert "ln1" not in key
+ assert "ln2" not in key
+ assert "ln_final" not in key
+ assert "q_norm" not in key
+ assert "k_norm" not in key
+
+ def test_no_bias_conversions(self, adapter: Olmo2ArchitectureAdapter) -> None:
+ """OLMo 2 has no biases on any projection."""
+ for key in _conversions(adapter):
+ assert not key.endswith(".bias")
diff --git a/tests/unit/model_bridge/supported_architectures/test_olmoe_adapter.py b/tests/unit/model_bridge/supported_architectures/test_olmoe_adapter.py
index 5cbc7ba4e..466ea77e0 100644
--- a/tests/unit/model_bridge/supported_architectures/test_olmoe_adapter.py
+++ b/tests/unit/model_bridge/supported_architectures/test_olmoe_adapter.py
@@ -17,7 +17,7 @@
import pytest
import torch.nn as nn
-from torch import equal, ones, randn, zeros
+from torch import ones, randn, zeros
from transformer_lens.config import TransformerBridgeConfig
from transformer_lens.conversion_utils.conversion_steps.rearrange_tensor_conversion import (
@@ -225,50 +225,6 @@ def test_gqa_does_not_affect_q_or_o(self, adapter: OlmoeArchitectureAdapter) ->
assert _rearrange(adapter, "blocks.{i}.attn.o.weight").axes_lengths["n"] == 4
-class TestOlmoeWeightConversionRoundTrips:
- """Run the rearrange conversions on synthetic HF-shaped tensors.
-
- The pattern/axis assertions above only check metadata. These confirm the
- conversions actually reshape realistic weight tensors into the split-head
- layout and revert losslessly (a rearrange operation is a pure permutation,
- so the round-trip must be exactly equal).
- """
-
- N_HEADS = 4
- N_KV_HEADS = 2
- D_HEAD = 16
- D_MODEL = 64
-
- @pytest.fixture
- def adapter(self) -> OlmoeArchitectureAdapter:
- return OlmoeArchitectureAdapter(_cfg(n_key_value_heads=self.N_KV_HEADS))
-
- def _roundtrip(self, adapter: OlmoeArchitectureAdapter, key: str, tensor: Any) -> tuple:
- conv = _param_conversion(adapter, key)
- converted = conv.convert({key: tensor}, key)
- reverted = conv.revert(converted)
- return converted, reverted
-
- def test_q_weight_splits_into_n_heads(self, adapter: OlmoeArchitectureAdapter) -> None:
- w = randn(self.N_HEADS * self.D_HEAD, self.D_MODEL)
- converted, reverted = self._roundtrip(adapter, "blocks.{i}.attn.q.weight", w)
- assert converted.shape == (self.N_HEADS, self.D_MODEL, self.D_HEAD)
- assert equal(reverted, w)
-
- def test_kv_weight_splits_into_n_kv_heads(self, adapter: OlmoeArchitectureAdapter) -> None:
- for slot in ("k", "v"):
- w = randn(self.N_KV_HEADS * self.D_HEAD, self.D_MODEL)
- converted, reverted = self._roundtrip(adapter, f"blocks.{{i}}.attn.{slot}.weight", w)
- assert converted.shape == (self.N_KV_HEADS, self.D_MODEL, self.D_HEAD)
- assert equal(reverted, w)
-
- def test_o_weight_merges_heads(self, adapter: OlmoeArchitectureAdapter) -> None:
- w = randn(self.D_MODEL, self.N_HEADS * self.D_HEAD)
- converted, reverted = self._roundtrip(adapter, "blocks.{i}.attn.o.weight", w)
- assert converted.shape == (self.N_HEADS, self.D_HEAD, self.D_MODEL)
- assert equal(reverted, w)
-
-
class TestOlmoeComponentMapping:
"""Structure of the component mapping: required keys and submodules."""
@@ -438,29 +394,23 @@ def _hook(x: Any, hook: Any) -> Any:
# Identity RoPE inputs keep this test focused on hook reshaping, not rotation math.
cos = ones(1, self.SEQ, self.D_HEAD)
sin = zeros(1, self.SEQ, self.D_HEAD)
- out = attn_bridge(hidden, position_embeddings=(cos, sin))
- # The attention bridge may return either a bare tensor or an (output, ...) tuple.
- out_tensor = out[0] if isinstance(out, tuple) else out
+ attn_bridge(hidden, position_embeddings=(cos, sin))
- return captured["q"], captured["k"], captured["v"], out_tensor
+ return captured["q"], captured["k"], captured["v"]
def test_hook_q_uses_n_heads(
self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
) -> None:
- q, _, _, _ = self._run_and_capture(wired_attn_bridge)
+ q, _, _ = self._run_and_capture(wired_attn_bridge)
assert q.shape == (self.BATCH, self.SEQ, self.N_HEADS, self.D_HEAD)
def test_hook_kv_use_n_kv_heads(
self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
) -> None:
- _, k, v, _ = self._run_and_capture(wired_attn_bridge)
+ _, k, v = self._run_and_capture(wired_attn_bridge)
assert k.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
assert v.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
- def test_attn_output_shape(self, wired_attn_bridge: PositionEmbeddingsAttentionBridge) -> None:
- _, _, _, out = self._run_and_capture(wired_attn_bridge)
- assert out.shape == (self.BATCH, self.SEQ, self.D_MODEL)
-
class TestOlmoeSetupComponentTesting:
"""setup_component_testing wires the shared rotary embedding and forces eager attention."""
diff --git a/tests/unit/model_bridge/supported_architectures/test_openelm_adapter.py b/tests/unit/model_bridge/supported_architectures/test_openelm_adapter.py
new file mode 100644
index 000000000..0fba94511
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_openelm_adapter.py
@@ -0,0 +1,207 @@
+"""Unit tests for OpenElmArchitectureAdapter.
+
+Tests cover:
+- Config attribute validation (all required attributes are set correctly)
+- Component mapping structure (correct bridge types and HF module names)
+- Weight conversion keys (empty for OpenELM — native attention handles all variants)
+"""
+
+import pytest
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge.generalized_components import (
+ AttentionBridge,
+ BlockBridge,
+ EmbeddingBridge,
+ LinearBridge,
+ MLPBridge,
+ RMSNormalizationBridge,
+ UnembeddingBridge,
+)
+from transformer_lens.model_bridge.supported_architectures.openelm import (
+ OpenElmArchitectureAdapter,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+
+def _make_cfg(
+ n_heads: int = 4,
+ d_model: int = 64,
+ n_layers: int = 2,
+ d_mlp: int = 256,
+ d_vocab: int = 1000,
+ n_ctx: int = 512,
+) -> TransformerBridgeConfig:
+ """Return a minimal TransformerBridgeConfig for OpenELM adapter tests."""
+ return TransformerBridgeConfig(
+ d_model=d_model,
+ d_head=d_model // n_heads,
+ n_layers=n_layers,
+ n_ctx=n_ctx,
+ n_heads=n_heads,
+ d_vocab=d_vocab,
+ d_mlp=d_mlp,
+ default_prepend_bos=True,
+ architecture="OpenELMForCausalLM",
+ )
+
+
+@pytest.fixture
+def cfg() -> TransformerBridgeConfig:
+ return _make_cfg()
+
+
+@pytest.fixture
+def adapter(cfg: TransformerBridgeConfig) -> OpenElmArchitectureAdapter:
+ return OpenElmArchitectureAdapter(cfg)
+
+
+# ---------------------------------------------------------------------------
+# Config attribute tests
+# ---------------------------------------------------------------------------
+
+
+class TestOpenElmAdapterConfig:
+ """Anti-drift flags and non-obvious config choices for OpenElmArchitectureAdapter."""
+
+ def test_final_rms_is_true(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert adapter.cfg.final_rms is True
+
+ def test_uses_rms_norm_is_true(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert adapter.cfg.uses_rms_norm is True
+
+ def test_tokenizer_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ """OpenELM has no bundled tokenizer — uses LLaMA-2 tokenizer as proxy."""
+ assert adapter.cfg.tokenizer_name == "NousResearch/Llama-2-7b-hf"
+
+
+# ---------------------------------------------------------------------------
+# Component mapping structure tests
+# ---------------------------------------------------------------------------
+
+
+class TestOpenElmAdapterComponentMapping:
+ """Tests that component_mapping has the correct bridge types and HF module names."""
+
+ # -- Top-level keys --
+
+ def test_embed_is_embedding_bridge(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert isinstance(adapter.component_mapping["embed"], EmbeddingBridge)
+
+ def test_embed_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert adapter.component_mapping["embed"].name == "transformer.token_embeddings"
+
+ def test_no_pos_embed_key(self, adapter: OpenElmArchitectureAdapter) -> None:
+ """OpenELM uses per-layer rotary embeddings — no shared positional embedding."""
+ assert "pos_embed" not in adapter.component_mapping
+
+ def test_no_rotary_emb_key(self, adapter: OpenElmArchitectureAdapter) -> None:
+ """OpenELM RoPE is embedded per-layer in attention, not a top-level component."""
+ assert "rotary_emb" not in adapter.component_mapping
+
+ def test_blocks_is_block_bridge(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert isinstance(adapter.component_mapping["blocks"], BlockBridge)
+
+ def test_blocks_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert adapter.component_mapping["blocks"].name == "transformer.layers"
+
+ def test_ln_final_is_rms_normalization_bridge(
+ self, adapter: OpenElmArchitectureAdapter
+ ) -> None:
+ assert isinstance(adapter.component_mapping["ln_final"], RMSNormalizationBridge)
+
+ def test_ln_final_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert adapter.component_mapping["ln_final"].name == "transformer.norm"
+
+ def test_unembed_is_unembedding_bridge(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert isinstance(adapter.component_mapping["unembed"], UnembeddingBridge)
+
+ def test_unembed_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert adapter.component_mapping["unembed"].name == "lm_head"
+
+ # -- Block submodules --
+
+ def test_blocks_ln1_is_rms_normalization_bridge(
+ self, adapter: OpenElmArchitectureAdapter
+ ) -> None:
+ assert isinstance(
+ adapter.component_mapping["blocks"].submodules["ln1"], RMSNormalizationBridge
+ )
+
+ def test_blocks_ln1_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert adapter.component_mapping["blocks"].submodules["ln1"].name == "attn_norm"
+
+ def test_blocks_ln2_is_rms_normalization_bridge(
+ self, adapter: OpenElmArchitectureAdapter
+ ) -> None:
+ assert isinstance(
+ adapter.component_mapping["blocks"].submodules["ln2"], RMSNormalizationBridge
+ )
+
+ def test_blocks_ln2_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert adapter.component_mapping["blocks"].submodules["ln2"].name == "ffn_norm"
+
+ def test_attn_is_attention_bridge(self, adapter: OpenElmArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert isinstance(blocks.submodules["attn"], AttentionBridge)
+
+ def test_attn_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["attn"].name == "attn"
+
+ def test_attn_requires_attention_mask(self, adapter: OpenElmArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.requires_attention_mask is True
+
+ def test_attn_qkv_is_linear_bridge(self, adapter: OpenElmArchitectureAdapter) -> None:
+ """OpenELM uses a combined QKV projection (not separate q/k/v)."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert isinstance(attn.submodules["qkv"], LinearBridge)
+
+ def test_attn_qkv_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.submodules["qkv"].name == "qkv_proj"
+
+ def test_attn_o_is_linear_bridge(self, adapter: OpenElmArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert isinstance(attn.submodules["o"], LinearBridge)
+
+ def test_attn_o_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.submodules["o"].name == "out_proj"
+
+ def test_mlp_is_mlp_bridge(self, adapter: OpenElmArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert isinstance(blocks.submodules["mlp"], MLPBridge)
+
+ def test_mlp_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ """OpenELM names its MLP submodule 'ffn' (feedforward network)."""
+ assert adapter.component_mapping["blocks"].submodules["mlp"].name == "ffn"
+
+ def test_mlp_in_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert mlp.submodules["in"].name == "proj_1"
+
+ def test_mlp_out_name(self, adapter: OpenElmArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert mlp.submodules["out"].name == "proj_2"
+
+
+# ---------------------------------------------------------------------------
+# Weight processing conversion tests
+# ---------------------------------------------------------------------------
+
+
+class TestOpenElmAdapterWeightConversions:
+ """Tests that weight_processing_conversions is empty for OpenELM.
+
+ OpenELM uses per-layer varying head counts and FFN dimensions handled
+ entirely by native HuggingFace attention — no static weight rearrangements
+ are needed at the bridge level.
+ """
+
+ def test_no_weight_processing_conversions(self, adapter: OpenElmArchitectureAdapter) -> None:
+ assert len(adapter.weight_processing_conversions) == 0
diff --git a/tests/unit/model_bridge/supported_architectures/test_phi3_adapter.py b/tests/unit/model_bridge/supported_architectures/test_phi3_adapter.py
new file mode 100644
index 000000000..b66a07f06
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_phi3_adapter.py
@@ -0,0 +1,333 @@
+"""Unit tests for Phi3ArchitectureAdapter.
+
+Tests cover:
+- Component mapping structure (bridge types and HF module names)
+- Weight conversion key set
+- _SizedSplitConversion numerical correctness
+- Config flags set by the adapter
+- preprocess_weights LN folding
+"""
+
+import pytest
+import torch
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge.generalized_components import (
+ BlockBridge,
+ EmbeddingBridge,
+ JointGateUpMLPBridge,
+ JointQKVPositionEmbeddingsAttentionBridge,
+ LinearBridge,
+ RMSNormalizationBridge,
+ RotaryEmbeddingBridge,
+ UnembeddingBridge,
+)
+from transformer_lens.model_bridge.supported_architectures.phi3 import (
+ Phi3ArchitectureAdapter,
+ _SizedSplitConversion,
+)
+
+# ---------------------------------------------------------------------------
+# Helpers / fixtures
+# ---------------------------------------------------------------------------
+
+N_HEADS = 4
+N_KV_HEADS = 2
+D_MODEL = 64
+D_HEAD = D_MODEL // N_HEADS # 16
+D_MLP = 128
+N_LAYERS = 2
+N_CTX = 128
+D_VOCAB = 500
+
+
+def _make_cfg(
+ n_heads: int = N_HEADS,
+ n_kv_heads: int = N_KV_HEADS,
+ d_model: int = D_MODEL,
+ n_layers: int = N_LAYERS,
+ d_mlp: int = D_MLP,
+ d_vocab: int = D_VOCAB,
+ n_ctx: int = N_CTX,
+) -> TransformerBridgeConfig:
+ """Return a minimal TransformerBridgeConfig for Phi-3 adapter tests."""
+ return TransformerBridgeConfig(
+ d_model=d_model,
+ d_head=d_model // n_heads,
+ n_layers=n_layers,
+ n_ctx=n_ctx,
+ n_heads=n_heads,
+ d_vocab=d_vocab,
+ d_mlp=d_mlp,
+ n_key_value_heads=n_kv_heads,
+ default_prepend_bos=True,
+ architecture="Phi3ForCausalLM",
+ )
+
+
+@pytest.fixture
+def cfg() -> TransformerBridgeConfig:
+ return _make_cfg()
+
+
+@pytest.fixture
+def adapter(cfg: TransformerBridgeConfig) -> Phi3ArchitectureAdapter:
+ return Phi3ArchitectureAdapter(cfg)
+
+
+# ---------------------------------------------------------------------------
+# Config flag tests
+# ---------------------------------------------------------------------------
+
+
+class TestPhi3AdapterConfig:
+ """Tests that the adapter sets the correct config flags."""
+
+ def test_normalization_type(self, adapter: Phi3ArchitectureAdapter) -> None:
+ assert adapter.cfg.normalization_type == "RMS"
+
+ def test_positional_embedding_type(self, adapter: Phi3ArchitectureAdapter) -> None:
+ assert adapter.cfg.positional_embedding_type == "rotary"
+
+ def test_gated_mlp(self, adapter: Phi3ArchitectureAdapter) -> None:
+ assert adapter.cfg.gated_mlp is True
+
+ def test_final_rms(self, adapter: Phi3ArchitectureAdapter) -> None:
+ assert adapter.cfg.final_rms is False
+
+ def test_supports_fold_ln_false(self, adapter: Phi3ArchitectureAdapter) -> None:
+ """Standard fold_ln is disabled — handled in preprocess_weights instead."""
+ assert adapter.supports_fold_ln is False
+
+
+# ---------------------------------------------------------------------------
+# Component mapping tests
+# ---------------------------------------------------------------------------
+
+
+class TestPhi3AdapterComponentMapping:
+ """Tests that component_mapping has the correct bridge types and HF module names."""
+
+ def test_top_level_keys(self, adapter: Phi3ArchitectureAdapter) -> None:
+ assert set(adapter.component_mapping.keys()) == {
+ "embed",
+ "rotary_emb",
+ "blocks",
+ "ln_final",
+ "unembed",
+ }
+
+ def test_bridge_types(self, adapter: Phi3ArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ assert isinstance(mapping["embed"], EmbeddingBridge)
+ assert isinstance(mapping["rotary_emb"], RotaryEmbeddingBridge)
+ assert isinstance(mapping["blocks"], BlockBridge)
+ assert isinstance(mapping["ln_final"], RMSNormalizationBridge)
+ assert isinstance(mapping["unembed"], UnembeddingBridge)
+
+ def test_top_level_hf_paths(self, adapter: Phi3ArchitectureAdapter) -> None:
+ mapping = adapter.component_mapping
+ assert mapping["embed"].name == "model.embed_tokens"
+ assert mapping["rotary_emb"].name == "model.rotary_emb"
+ assert mapping["blocks"].name == "model.layers"
+ assert mapping["ln_final"].name == "model.norm"
+ assert mapping["unembed"].name == "lm_head"
+
+ def test_block_submodule_keys(self, adapter: Phi3ArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert set(blocks.submodules.keys()) == {"ln1", "ln2", "attn", "mlp"}
+
+ def test_block_bridge_types(self, adapter: Phi3ArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert isinstance(blocks.submodules["ln1"], RMSNormalizationBridge)
+ assert isinstance(blocks.submodules["ln2"], RMSNormalizationBridge)
+ assert isinstance(blocks.submodules["attn"], JointQKVPositionEmbeddingsAttentionBridge)
+ assert isinstance(blocks.submodules["mlp"], JointGateUpMLPBridge)
+
+ def test_block_hf_paths(self, adapter: Phi3ArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["ln1"].name == "input_layernorm"
+ assert blocks.submodules["ln2"].name == "post_attention_layernorm"
+ assert blocks.submodules["attn"].name == "self_attn"
+ assert blocks.submodules["mlp"].name == "mlp"
+
+ def test_attention_submodule_keys(self, adapter: Phi3ArchitectureAdapter) -> None:
+ """Phi-3 uses a fused qkv_proj with a separate o_proj."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert set(attn.submodules.keys()) == {"qkv", "q", "k", "v", "o"}
+
+ def test_attention_hf_paths(self, adapter: Phi3ArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.submodules["qkv"].name == "qkv_proj"
+ assert attn.submodules["o"].name == "o_proj"
+
+ def test_mlp_submodule_keys(self, adapter: Phi3ArchitectureAdapter) -> None:
+ """Phi-3 MLP exposes only the output projection; gate/up come from fused gate_up_proj."""
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert set(mlp.submodules.keys()) == {"gate", "in", "out"}
+
+ def test_mlp_hf_paths(self, adapter: Phi3ArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert mlp.submodules["out"].name == "down_proj"
+
+ def test_linear_submodule_bridge_types(self, adapter: Phi3ArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ for submodule in [*attn.submodules.values(), *mlp.submodules.values()]:
+ assert isinstance(submodule, LinearBridge)
+
+
+# ---------------------------------------------------------------------------
+# Weight conversion key tests
+# ---------------------------------------------------------------------------
+
+
+class TestPhi3AdapterWeightConversions:
+ """Tests that weight_processing_conversions has exactly the expected keys."""
+
+ def test_exact_conversion_key_set(self, adapter: Phi3ArchitectureAdapter) -> None:
+ assert set(adapter.weight_processing_conversions.keys()) == {
+ "blocks.{i}.attn.q",
+ "blocks.{i}.attn.k",
+ "blocks.{i}.attn.v",
+ "blocks.{i}.attn.o",
+ "blocks.{i}.mlp.in",
+ "blocks.{i}.mlp.gate",
+ }
+
+ def test_qkv_source_key(self, adapter: Phi3ArchitectureAdapter) -> None:
+ """Q, K, V all source from the same fused qkv_proj weight."""
+ for key in ["blocks.{i}.attn.q", "blocks.{i}.attn.k", "blocks.{i}.attn.v"]:
+ conv = adapter.weight_processing_conversions[key]
+ assert conv.source_key == "model.layers.{i}.self_attn.qkv_proj.weight"
+
+ def test_mlp_source_key(self, adapter: Phi3ArchitectureAdapter) -> None:
+ """Gate and up projections both source from fused gate_up_proj."""
+ for key in ["blocks.{i}.mlp.in", "blocks.{i}.mlp.gate"]:
+ conv = adapter.weight_processing_conversions[key]
+ assert conv.source_key == "model.layers.{i}.mlp.gate_up_proj.weight"
+
+
+# ---------------------------------------------------------------------------
+# _SizedSplitConversion numerical correctness tests
+# ---------------------------------------------------------------------------
+
+
+class TestSizedSplitConversion:
+ """Numerical correctness of Phi-3's GQA split conversion."""
+
+ def test_extracts_q_slice(self) -> None:
+ """Index 0 should return the first (Q) chunk."""
+ q_size, kv_size = 8, 4
+ sizes = [q_size, kv_size, kv_size]
+ tensor = torch.arange(float(q_size + 2 * kv_size)).unsqueeze(1) # [16, 1]
+ conv = _SizedSplitConversion(sizes=sizes, index=0)
+ out = conv.handle_conversion(tensor)
+ assert out.shape[0] == q_size
+ assert torch.allclose(out, tensor[:q_size])
+
+ def test_extracts_k_slice(self) -> None:
+ """Index 1 should return the second (K) chunk."""
+ q_size, kv_size = 8, 4
+ sizes = [q_size, kv_size, kv_size]
+ tensor = torch.arange(float(q_size + 2 * kv_size)).unsqueeze(1)
+ conv = _SizedSplitConversion(sizes=sizes, index=1)
+ out = conv.handle_conversion(tensor)
+ assert out.shape[0] == kv_size
+ assert torch.allclose(out, tensor[q_size : q_size + kv_size])
+
+ def test_extracts_v_slice(self) -> None:
+ """Index 2 should return the third (V) chunk."""
+ q_size, kv_size = 8, 4
+ sizes = [q_size, kv_size, kv_size]
+ tensor = torch.arange(float(q_size + 2 * kv_size)).unsqueeze(1)
+ conv = _SizedSplitConversion(sizes=sizes, index=2)
+ out = conv.handle_conversion(tensor)
+ assert out.shape[0] == kv_size
+ assert torch.allclose(out, tensor[q_size + kv_size :])
+
+ def test_dim_1_split(self) -> None:
+ """Splitting along dim=1 returns the correct column slice."""
+ sizes = [3, 5]
+ tensor = torch.ones(4, 8)
+ conv = _SizedSplitConversion(sizes=sizes, index=1, dim=1)
+ out = conv.handle_conversion(tensor)
+ assert out.shape == (4, 5)
+
+ def test_values_are_correct_not_just_shape(self) -> None:
+ """Returned slice should contain the correct values, not just the right shape."""
+ tensor = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
+ conv = _SizedSplitConversion(sizes=[2, 2, 2], index=1)
+ out = conv.handle_conversion(tensor)
+ expected = torch.tensor([3.0, 4.0])
+ assert torch.allclose(out, expected)
+
+ def test_three_slices_reconstruct_original(self) -> None:
+ """Concatenating all three slices should recover the original tensor."""
+ torch.manual_seed(42)
+ q_size, kv_size = 16, 8
+ sizes = [q_size, kv_size, kv_size]
+ tensor = torch.randn(q_size + 2 * kv_size, 32)
+ parts = [_SizedSplitConversion(sizes, i).handle_conversion(tensor) for i in range(3)]
+ assert torch.allclose(torch.cat(parts, dim=0), tensor)
+
+
+# ---------------------------------------------------------------------------
+# preprocess_weights: LN folding
+# ---------------------------------------------------------------------------
+
+
+class TestPhi3PreprocessWeights:
+ """Tests that preprocess_weights correctly folds RMS-norm scales."""
+
+ def _make_state_dict(self, n_layers: int = 2, d_model: int = D_MODEL, d_mlp: int = D_MLP):
+ """Build a minimal state dict matching what weight_processing would see."""
+ sd = {}
+ for i in range(n_layers):
+ sd[f"blocks.{i}.ln1.weight"] = torch.full((d_model,), 2.0)
+ sd[f"blocks.{i}.ln2.weight"] = torch.full((d_model,), 3.0)
+ sd[f"blocks.{i}.attn.q.weight"] = torch.ones(N_HEADS * D_HEAD, d_model)
+ sd[f"blocks.{i}.attn.k.weight"] = torch.ones(N_KV_HEADS * D_HEAD, d_model)
+ sd[f"blocks.{i}.attn.v.weight"] = torch.ones(N_KV_HEADS * D_HEAD, d_model)
+ sd[f"blocks.{i}.mlp.gate.weight"] = torch.ones(d_mlp, d_model)
+ sd[f"blocks.{i}.mlp.in.weight"] = torch.ones(d_mlp, d_model)
+ sd["ln_final.weight"] = torch.full((d_model,), 4.0)
+ sd["unembed.weight"] = torch.ones(D_VOCAB, d_model)
+ return sd
+
+ def test_ln1_folded_into_qkv(self, adapter: Phi3ArchitectureAdapter) -> None:
+ """ln1 scale should be multiplied into Q/K/V weights."""
+ sd = self._make_state_dict()
+ adapter._fold_ln_requested = True
+ out = adapter.preprocess_weights(sd)
+ # ln1.weight was 2.0, QKV weights were 1.0 → expect 2.0
+ for key in ["blocks.0.attn.q.weight", "blocks.0.attn.k.weight", "blocks.0.attn.v.weight"]:
+ assert torch.allclose(out[key], torch.full_like(out[key], 2.0)), key
+
+ def test_ln1_set_to_ones_after_fold(self, adapter: Phi3ArchitectureAdapter) -> None:
+ sd = self._make_state_dict()
+ adapter._fold_ln_requested = True
+ out = adapter.preprocess_weights(sd)
+ assert torch.allclose(out["blocks.0.ln1.weight"], torch.ones(D_MODEL))
+
+ def test_ln2_folded_into_mlp(self, adapter: Phi3ArchitectureAdapter) -> None:
+ """ln2 scale should be multiplied into gate and up projection weights."""
+ sd = self._make_state_dict()
+ adapter._fold_ln_requested = True
+ out = adapter.preprocess_weights(sd)
+ for key in ["blocks.0.mlp.gate.weight", "blocks.0.mlp.in.weight"]:
+ assert torch.allclose(out[key], torch.full_like(out[key], 3.0)), key
+
+ def test_ln2_set_to_ones_after_fold(self, adapter: Phi3ArchitectureAdapter) -> None:
+ sd = self._make_state_dict()
+ adapter._fold_ln_requested = True
+ out = adapter.preprocess_weights(sd)
+ assert torch.allclose(out["blocks.0.ln2.weight"], torch.ones(D_MODEL))
+
+ def test_fold_skipped_when_not_requested(self, adapter: Phi3ArchitectureAdapter) -> None:
+ """When _fold_ln_requested=False the state dict is returned unchanged."""
+ sd = self._make_state_dict()
+ adapter._fold_ln_requested = False
+ 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))
diff --git a/tests/unit/model_bridge/supported_architectures/test_phi_adapter.py b/tests/unit/model_bridge/supported_architectures/test_phi_adapter.py
new file mode 100644
index 000000000..0a70eb7d2
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_phi_adapter.py
@@ -0,0 +1,344 @@
+"""Unit tests for PhiArchitectureAdapter.
+
+Tests cover:
+- Config attribute validation
+- Component mapping structure
+- Weight conversion keys and rearrange patterns
+- Architecture guards
+- Setup component tests
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
+from transformer_lens.conversion_utils.param_processing_conversion import (
+ ParamProcessingConversion,
+)
+from transformer_lens.model_bridge.generalized_components import (
+ EmbeddingBridge,
+ LinearBridge,
+ MLPBridge,
+ NormalizationBridge,
+ ParallelBlockBridge,
+ PositionEmbeddingsAttentionBridge,
+ RotaryEmbeddingBridge,
+ UnembeddingBridge,
+)
+from transformer_lens.model_bridge.supported_architectures.phi import (
+ PhiArchitectureAdapter,
+)
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+def _make_cfg(
+ n_heads: int = 4,
+ d_model: int = 64,
+ n_layers: int = 2,
+ d_mlp: int = 256,
+ d_vocab: int = 1000,
+ n_ctx: int = 512,
+) -> TransformerBridgeConfig:
+ """Return a minimal TransformerBridgeConfig for Phi adapter tests."""
+ return TransformerBridgeConfig(
+ d_model=d_model,
+ d_head=d_model // n_heads,
+ n_layers=n_layers,
+ n_ctx=n_ctx,
+ n_heads=n_heads,
+ d_vocab=d_vocab,
+ d_mlp=d_mlp,
+ architecture="PhiForCausalLM",
+ )
+
+
+@pytest.fixture
+def cfg() -> TransformerBridgeConfig:
+ return _make_cfg()
+
+
+@pytest.fixture
+def adapter(cfg: TransformerBridgeConfig) -> PhiArchitectureAdapter:
+ return PhiArchitectureAdapter(cfg)
+
+
+def _fake_hf_model(rotary_emb: object) -> SimpleNamespace:
+ return SimpleNamespace(model=SimpleNamespace(rotary_emb=rotary_emb))
+
+
+class DummyAttention:
+ def __init__(self) -> None:
+ self.rotary_emb = None
+
+ def set_rotary_emb(self, rotary_emb: object) -> None:
+ self.rotary_emb = rotary_emb
+
+
+class DummyBlock:
+ def __init__(self, has_attention: bool = True) -> None:
+ if has_attention:
+ self.attn = DummyAttention()
+
+
+class DummyBridgeModel:
+ def __init__(self, blocks: list[DummyBlock]) -> None:
+ self.blocks = blocks
+
+
+# ---------------------------------------------------------------------------
+# Config attribute tests
+# ---------------------------------------------------------------------------
+
+
+class TestPhiAdapterConfig:
+ """Adapter must set all required config flags to the values Phi expects."""
+
+ def test_use_fast_is_false(self, adapter: PhiArchitectureAdapter) -> None:
+ """Do not use the rust based HF tokenizer. Uses python based version instead"""
+ assert adapter.cfg.use_fast is False
+
+ def test_default_prepend_bos_is_false(self, adapter: PhiArchitectureAdapter) -> None:
+ """Phi was not trained with BOS token so TransformerLens should not append it"""
+ assert adapter.cfg.default_prepend_bos is False
+
+
+# ---------------------------------------------------------------------------
+# Component mapping structure tests
+# ---------------------------------------------------------------------------
+
+
+class TestPhiAdapterComponentMapping:
+ """Component mapping must have the correct bridge types and HF module names."""
+
+ # -- Top-level keys --
+
+ def test_embed_is_embedding_bridge(self, adapter: PhiArchitectureAdapter) -> None:
+ assert isinstance(adapter.component_mapping["embed"], EmbeddingBridge)
+
+ def test_embed_name(self, adapter: PhiArchitectureAdapter) -> None:
+ assert adapter.component_mapping["embed"].name == "model.embed_tokens"
+
+ def test_rotary_emb_is_rotary_embedding_bridge(self, adapter: PhiArchitectureAdapter) -> None:
+ assert isinstance(adapter.component_mapping["rotary_emb"], RotaryEmbeddingBridge)
+
+ def test_rotary_emb_name(self, adapter: PhiArchitectureAdapter) -> None:
+ assert adapter.component_mapping["rotary_emb"].name == "model.rotary_emb"
+
+ def test_blocks_is_parallel_block_bridge(self, adapter: PhiArchitectureAdapter) -> None:
+ """Parallel attn+MLP requires ParallelBlockBridge, not sequential BlockBridge."""
+ assert isinstance(adapter.component_mapping["blocks"], ParallelBlockBridge)
+
+ def test_blocks_name(self, adapter: PhiArchitectureAdapter) -> None:
+ assert adapter.component_mapping["blocks"].name == "model.layers"
+
+ def test_ln_final_is_normalization_bridge(self, adapter: PhiArchitectureAdapter) -> None:
+ assert isinstance(adapter.component_mapping["ln_final"], NormalizationBridge)
+
+ def test_ln_final_name(self, adapter: PhiArchitectureAdapter) -> None:
+ assert adapter.component_mapping["ln_final"].name == "model.final_layernorm"
+
+ def test_ln_final_use_native_layernorm_autograd_is_true(
+ self, adapter: PhiArchitectureAdapter
+ ) -> None:
+ assert adapter.component_mapping["ln_final"].use_native_layernorm_autograd is True
+
+ def test_unembed_is_unembedding_bridge(self, adapter: PhiArchitectureAdapter) -> None:
+ assert isinstance(adapter.component_mapping["unembed"], UnembeddingBridge)
+
+ def test_unembed_name(self, adapter: PhiArchitectureAdapter) -> None:
+ assert adapter.component_mapping["unembed"].name == "lm_head"
+
+ # -- Block submodules --
+
+ def test_blocks_ln1_is_normalization_bridge(self, adapter: PhiArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert isinstance(blocks.submodules["ln1"], NormalizationBridge)
+
+ def test_blocks_ln1_name(self, adapter: PhiArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["ln1"].name == "input_layernorm"
+
+ def test_blocks_ln1_use_native_layernorm_autograd_is_true(
+ self, adapter: PhiArchitectureAdapter
+ ) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["ln1"].use_native_layernorm_autograd is True
+
+ def test_attn_is_position_embeddings_attention_bridge(
+ self, adapter: PhiArchitectureAdapter
+ ) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert isinstance(blocks.submodules["attn"], PositionEmbeddingsAttentionBridge)
+
+ def test_attn_name(self, adapter: PhiArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["attn"].name == "self_attn"
+
+ def test_attn_requires_attention_mask_is_true(self, adapter: PhiArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["attn"].requires_attention_mask is True
+
+ def test_attn_requires_position_embeddings_is_true(
+ self, adapter: PhiArchitectureAdapter
+ ) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["attn"].requires_position_embeddings is True
+
+ def test_mlp_is_mlp_bridge(self, adapter: PhiArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert isinstance(blocks.submodules["mlp"], MLPBridge)
+
+ def test_mlp_name(self, adapter: PhiArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert blocks.submodules["mlp"].name == "mlp"
+
+ # -- Attention submodules --
+
+ @pytest.mark.parametrize("slot", ["q", "k", "v", "o"])
+ def test_attn_submodule_is_linear_bridge(
+ self, adapter: PhiArchitectureAdapter, slot: str
+ ) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert isinstance(attn.submodules[slot], LinearBridge)
+
+ @pytest.mark.parametrize(
+ "slot, hf_name",
+ [("q", "q_proj"), ("k", "k_proj"), ("v", "v_proj"), ("o", "dense")],
+ )
+ def test_attn_submodule_name(
+ self, adapter: PhiArchitectureAdapter, slot: str, hf_name: str
+ ) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.submodules[slot].name == hf_name
+
+ # -- MLP submodules --
+
+ def test_mlp_in_is_linear_bridge(self, adapter: PhiArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert isinstance(mlp.submodules["in"], LinearBridge)
+
+ def test_mlp_in_name(self, adapter: PhiArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert mlp.submodules["in"].name == "fc1"
+
+ def test_mlp_out_is_linear_bridge(self, adapter: PhiArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert isinstance(mlp.submodules["out"], LinearBridge)
+
+ def test_mlp_out_name(self, adapter: PhiArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert mlp.submodules["out"].name == "fc2"
+
+
+# ---------------------------------------------------------------------------
+# Weight processing conversion tests
+# ---------------------------------------------------------------------------
+
+
+class TestPhiAdapterWeightConversions:
+ """Adapter must define exactly the four QKVO weight and three QKV bias conversions."""
+
+ def test_conversion_keys_present(self, adapter: PhiArchitectureAdapter) -> None:
+ """Phi has 4 weight matrices (QKVO) and 3 bias vectors (QKV) per layer"""
+ assert adapter.weight_processing_conversions.keys() == {
+ "blocks.{i}.attn.q.weight",
+ "blocks.{i}.attn.k.weight",
+ "blocks.{i}.attn.v.weight",
+ "blocks.{i}.attn.q.bias",
+ "blocks.{i}.attn.k.bias",
+ "blocks.{i}.attn.v.bias",
+ "blocks.{i}.attn.o.weight",
+ }
+
+ @pytest.mark.parametrize("slot", ["q", "k", "v"])
+ def test_qkv_weight_uses_split_heads_pattern(
+ self, adapter: PhiArchitectureAdapter, slot: str
+ ) -> None:
+ conv = adapter.weight_processing_conversions[f"blocks.{{i}}.attn.{slot}.weight"]
+ assert isinstance(conv, ParamProcessingConversion)
+ assert isinstance(conv.tensor_conversion, RearrangeTensorConversion)
+ assert conv.tensor_conversion.pattern == "(n h) m -> n m h"
+ assert conv.tensor_conversion.axes_lengths["n"] == adapter.cfg.n_heads
+
+ @pytest.mark.parametrize("slot", ["q", "k", "v"])
+ def test_qkv_bias_uses_split_heads_pattern(
+ self, adapter: PhiArchitectureAdapter, slot: str
+ ) -> None:
+ conv = adapter.weight_processing_conversions[f"blocks.{{i}}.attn.{slot}.bias"]
+ assert isinstance(conv, ParamProcessingConversion)
+ assert isinstance(conv.tensor_conversion, RearrangeTensorConversion)
+ assert conv.tensor_conversion.pattern == "(n h) -> n h"
+ assert conv.tensor_conversion.axes_lengths["n"] == adapter.cfg.n_heads
+
+ def test_o_uses_merge_heads_pattern(self, adapter: PhiArchitectureAdapter) -> None:
+ conv = adapter.weight_processing_conversions["blocks.{i}.attn.o.weight"]
+ assert isinstance(conv, ParamProcessingConversion)
+ assert isinstance(conv.tensor_conversion, RearrangeTensorConversion)
+ assert conv.tensor_conversion.pattern == "m (n h) -> n h m"
+ assert conv.tensor_conversion.axes_lengths["n"] == adapter.cfg.n_heads
+
+
+# ---------------------------------------------------------------------------
+# Architecture guards
+# ---------------------------------------------------------------------------
+
+
+class TestPhiArchitectureGuards:
+ """Guard against accidental introduction of features Phi does not have."""
+
+ def test_no_pos_embed_component(self, adapter: PhiArchitectureAdapter) -> None:
+ """Phi uses rotary embeddings, so there is no learned positional embedding."""
+ assert "pos_embed" not in adapter.component_mapping
+
+ def test_no_ln2_in_blocks(self, adapter: PhiArchitectureAdapter) -> None:
+ """Parallel attn+MLP shares a single ln_1; no ln2 exists."""
+ blocks = adapter.component_mapping["blocks"]
+ assert "ln2" not in blocks.submodules
+
+ def test_no_gate_in_mlp(self, adapter: PhiArchitectureAdapter) -> None:
+ """Phi uses a standard non-gated MLP; no gate submodule."""
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert "gate" not in mlp.submodules
+
+
+# ---------------------------------------------------------------------------
+# Setup component testing tests
+# ---------------------------------------------------------------------------
+
+
+class TestPhiSetupComponentTesting:
+ """setup_component_testing must wire Phi's shared rotary embedding into attention bridges."""
+
+ def test_sets_rotary_emb_on_template_attention(self, adapter: PhiArchitectureAdapter) -> None:
+ rotary_emb = object()
+ attn_template = adapter.get_generalized_component("blocks.0.attn")
+ assert isinstance(attn_template, PositionEmbeddingsAttentionBridge)
+ assert attn_template._rotary_emb is None
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb))
+
+ assert attn_template._rotary_emb is rotary_emb
+
+ def test_sets_rotary_emb_on_each_bridge_model_attention(
+ self, adapter: PhiArchitectureAdapter
+ ) -> None:
+ rotary_emb = object()
+ bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(), DummyBlock()])
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model)
+
+ for block in bridge_model.blocks:
+ assert block.attn.rotary_emb is rotary_emb
+
+ def test_skips_bridge_blocks_without_attention(self, adapter: PhiArchitectureAdapter) -> None:
+ rotary_emb = object()
+ bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(has_attention=False)])
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model)
+
+ assert bridge_model.blocks[0].attn.rotary_emb is rotary_emb
diff --git a/tests/unit/model_bridge/supported_architectures/test_qwen_adapter.py b/tests/unit/model_bridge/supported_architectures/test_qwen_adapter.py
new file mode 100644
index 000000000..42b83f52f
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_qwen_adapter.py
@@ -0,0 +1,342 @@
+"""Unit tests for QwenArchitectureAdapter
+
+Tests cover:
+- Config attributes the adapter sets (RMSNorm, rotary, gated MLP)
+- Component mapping: TL canonical names, Qwen HF module paths, and bridge types
+- The weight processing conversion keys {q/k/v/o} and their HF source weights
+- _split_qkv_matrix: Qwen fuses Q, K, V into one `c_attn` matrix, this splits it back
+- Factory registration: "QwenForCausalLM" resolves to this adapter
+
+These are pure unit tests. We build the adapter from a tiny mock config and inspect
+the Python object. Nothing is downloaded and no real Qwen checkpoint is loaded
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.factories.architecture_adapter_factory import (
+ ArchitectureAdapterFactory,
+)
+from transformer_lens.model_bridge.generalized_components import (
+ BlockBridge,
+ EmbeddingBridge,
+ GatedMLPBridge,
+ JointQKVAttentionBridge,
+ LinearBridge,
+ NormalizationBridge,
+ UnembeddingBridge,
+)
+from transformer_lens.model_bridge.supported_architectures.qwen import (
+ QwenArchitectureAdapter,
+)
+
+
+# Helpers
+def _make_cfg(
+ n_heads: int = 4,
+ d_model: int = 64,
+ n_layers: int = 2,
+ d_mlp: int = 256,
+ d_vocab: int = 1000,
+ n_ctx: int = 512,
+) -> TransformerBridgeConfig:
+ """Return a minimal TransformerBridgeConfig for Qwen adapter tests
+
+ Dimensions are kept tiny so the tests stay fast and need no HF download.
+ d_head is derived so that d_model == n_heads * d_head.
+ """
+ return TransformerBridgeConfig(
+ d_model=d_model,
+ d_head=d_model // n_heads,
+ n_layers=n_layers,
+ n_ctx=n_ctx,
+ n_heads=n_heads,
+ d_vocab=d_vocab,
+ d_mlp=d_mlp,
+ default_prepend_bos=False,
+ architecture="QwenForCausalLM",
+ )
+
+
+@pytest.fixture
+def cfg() -> TransformerBridgeConfig:
+ return _make_cfg()
+
+
+@pytest.fixture
+def adapter(cfg: TransformerBridgeConfig) -> QwenArchitectureAdapter:
+ return QwenArchitectureAdapter(cfg)
+
+
+# Config attribute tests
+class TestQwenAdapterConfig:
+ """The adapter sets these flags so downstream weight processing behaves correctly."""
+
+ # Qwen uses RMSNorm, not LayerNorm
+ def test_normalization_type_is_rms(self, adapter: QwenArchitectureAdapter) -> None:
+ assert adapter.cfg.normalization_type == "RMS"
+
+ def test_positional_embedding_type_is_rotary(self, adapter: QwenArchitectureAdapter) -> None:
+ assert adapter.cfg.positional_embedding_type == "rotary"
+
+ def test_final_rms_is_true(self, adapter: QwenArchitectureAdapter) -> None:
+ assert adapter.cfg.final_rms is True
+
+ def test_gated_mlp_is_true(self, adapter: QwenArchitectureAdapter) -> None:
+ """Qwen's MLP has a gate branch, so the adapter flags gated_mlp."""
+ assert adapter.cfg.gated_mlp is True
+
+ def test_attn_only_is_false(self, adapter: QwenArchitectureAdapter) -> None:
+ assert adapter.cfg.attn_only is False
+
+
+# Component mapping tests
+class TestQwenComponentMapping:
+ """The adapter contract: TL canonical names mapped to Qwen HF module paths."""
+
+ # Top-level keys
+
+ def test_top_level_keys(self, adapter: QwenArchitectureAdapter) -> None:
+ assert set(adapter.component_mapping.keys()) == {
+ "embed",
+ "blocks",
+ "ln_final",
+ "unembed",
+ }
+
+ def test_embed_is_embedding_bridge(self, adapter: QwenArchitectureAdapter) -> None:
+ assert isinstance(adapter.component_mapping["embed"], EmbeddingBridge)
+
+ def test_embed_name(self, adapter: QwenArchitectureAdapter) -> None:
+ assert adapter.component_mapping["embed"].name == "transformer.wte"
+
+ def test_blocks_is_block_bridge(self, adapter: QwenArchitectureAdapter) -> None:
+ assert isinstance(adapter.component_mapping["blocks"], BlockBridge)
+
+ def test_blocks_name(self, adapter: QwenArchitectureAdapter) -> None:
+ assert adapter.component_mapping["blocks"].name == "transformer.h"
+
+ def test_ln_final_is_normalization_bridge(self, adapter: QwenArchitectureAdapter) -> None:
+ assert isinstance(adapter.component_mapping["ln_final"], NormalizationBridge)
+
+ def test_ln_final_name(self, adapter: QwenArchitectureAdapter) -> None:
+ assert adapter.component_mapping["ln_final"].name == "transformer.ln_f"
+
+ def test_unembed_is_unembedding_bridge(self, adapter: QwenArchitectureAdapter) -> None:
+ assert isinstance(adapter.component_mapping["unembed"], UnembeddingBridge)
+
+ def test_unembed_name(self, adapter: QwenArchitectureAdapter) -> None:
+ assert adapter.component_mapping["unembed"].name == "lm_head"
+
+ # Block submodules
+
+ def test_block_submodule_keys(self, adapter: QwenArchitectureAdapter) -> None:
+ blocks = adapter.component_mapping["blocks"]
+ assert set(blocks.submodules.keys()) == {"ln1", "attn", "ln2", "mlp"}
+
+ def test_ln1_is_normalization_bridge_named_ln_1(self, adapter: QwenArchitectureAdapter) -> None:
+ ln1 = adapter.component_mapping["blocks"].submodules["ln1"]
+ assert isinstance(ln1, NormalizationBridge)
+ assert ln1.name == "ln_1"
+
+ def test_ln2_is_normalization_bridge_named_ln_2(self, adapter: QwenArchitectureAdapter) -> None:
+ ln2 = adapter.component_mapping["blocks"].submodules["ln2"]
+ assert isinstance(ln2, NormalizationBridge)
+ assert ln2.name == "ln_2"
+
+ # Attention submodules
+
+ def test_attn_is_joint_qkv_attention_bridge(self, adapter: QwenArchitectureAdapter) -> None:
+ """Qwen fuses Q/K/V into one matrix, so attention uses a JointQKVAttentionBridge."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert isinstance(attn, JointQKVAttentionBridge)
+
+ def test_attn_name(self, adapter: QwenArchitectureAdapter) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert attn.name == "attn"
+
+ def test_attn_submodule_keys(self, adapter: QwenArchitectureAdapter) -> None:
+ """Adapter only defines qkv and o, but JointQKVAttentionBridge already splits
+ the fused qkv back into q/k/v, so we include those in the adapter test here."""
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ assert set(attn.submodules.keys()) == {"qkv", "o", "q", "k", "v"}
+
+ def test_attn_qkv_is_linear_bridge_named_c_attn(self, adapter: QwenArchitectureAdapter) -> None:
+ qkv = adapter.component_mapping["blocks"].submodules["attn"].submodules["qkv"]
+ assert isinstance(qkv, LinearBridge)
+ assert qkv.name == "c_attn"
+
+ def test_attn_o_is_linear_bridge_named_c_proj(self, adapter: QwenArchitectureAdapter) -> None:
+ o = adapter.component_mapping["blocks"].submodules["attn"].submodules["o"]
+ assert isinstance(o, LinearBridge)
+ assert o.name == "c_proj"
+
+ # MLP submodule
+
+ def test_mlp_is_gated_mlp_bridge(self, adapter: QwenArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert isinstance(mlp, GatedMLPBridge)
+
+ def test_mlp_name(self, adapter: QwenArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert mlp.name == "mlp"
+
+ def test_mlp_submodule_keys(self, adapter: QwenArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert set(mlp.submodules.keys()) == {"gate", "in", "out"}
+
+ def test_mlp_hf_paths(self, adapter: QwenArchitectureAdapter) -> None:
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ assert mlp.submodules["gate"].name == "w1"
+ assert mlp.submodules["in"].name == "w2"
+ assert mlp.submodules["out"].name == "c_proj"
+
+ def test_all_linear_submodules_are_linear_bridges(
+ self, adapter: QwenArchitectureAdapter
+ ) -> None:
+ attn = adapter.component_mapping["blocks"].submodules["attn"]
+ mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+ for submodule in [*attn.submodules.values(), *mlp.submodules.values()]:
+ assert isinstance(submodule, LinearBridge)
+
+
+# Weight processing conversion tests
+
+
+class TestQwenWeightConversions:
+ """weight_processing_conversions reshape Qwen's fused weights into TL head-split form."""
+
+ @pytest.mark.parametrize(
+ "key",
+ [
+ "blocks.{i}.attn.q",
+ "blocks.{i}.attn.k",
+ "blocks.{i}.attn.v",
+ "blocks.{i}.attn.o",
+ ],
+ )
+ def test_conversion_key_present(self, adapter: QwenArchitectureAdapter, key: str) -> None:
+ assert key in adapter.weight_processing_conversions
+
+ def test_exactly_four_conversion_keys(self, adapter: QwenArchitectureAdapter) -> None:
+ assert len(adapter.weight_processing_conversions) == 4
+
+ @pytest.mark.parametrize("key", ["q", "k", "v"])
+ def test_qkv_conversions_read_from_c_attn(
+ self, adapter: QwenArchitectureAdapter, key: str
+ ) -> None:
+ """Q, K, V all come from the single fused c_attn weight."""
+ conversion = adapter.weight_processing_conversions[f"blocks.{{i}}.attn.{key}"]
+ assert conversion.source_key == "transformer.h.{i}.attn.c_attn.weight"
+
+ def test_o_conversion_reads_from_c_proj(self, adapter: QwenArchitectureAdapter) -> None:
+ conversion = adapter.weight_processing_conversions["blocks.{i}.attn.o"]
+ assert conversion.source_key == "transformer.h.{i}.attn.c_proj.weight"
+
+
+# _split_qkv_matrix — numerical correctness tests
+
+
+class MockQwenAttention(nn.Module):
+ """Stand-in for Qwen's HF attention module.
+
+ _split_qkv_matrix only looks at a single attribute, `c_attn`, so this is all
+ we need. c_attn is a fused linear that produces Q, K, V stacked together (3x wide).
+ """
+
+ def __init__(self, c_attn: nn.Linear) -> None:
+ super().__init__()
+ self.c_attn = c_attn
+
+
+class TestQwenQKVSplit:
+ """Qwen stores Q/K/V in one c_attn matrix; _split_qkv_matrix slices it back into three for use."""
+
+ D_MODEL = 64
+
+ def test_returns_three_linears_of_right_shape(self, adapter: QwenArchitectureAdapter) -> None:
+ d_model = self.D_MODEL
+ mock = MockQwenAttention(nn.Linear(d_model, 3 * d_model, bias=True))
+
+ q, k, v = adapter._split_qkv_matrix(mock)
+
+ for proj in (q, k, v):
+ assert isinstance(proj, nn.Linear)
+ assert proj.weight.shape == (d_model, d_model)
+ assert proj.bias.shape == (d_model,)
+
+ def test_thirds_land_in_correct_projection(self, adapter: QwenArchitectureAdapter) -> None:
+ """Standard Linear layout: rows [0:d], [d:2d], [2d:3d] become Q, K, V."""
+ d_model = self.D_MODEL
+ c_attn = nn.Linear(d_model, 3 * d_model, bias=False)
+ with torch.no_grad():
+ # Fill each third with a recognizable constant: Q=1, K=2, V=3
+ c_attn.weight[:d_model] = 1.0
+ c_attn.weight[d_model : 2 * d_model] = 2.0
+ c_attn.weight[2 * d_model :] = 3.0
+
+ q, k, v = adapter._split_qkv_matrix(MockQwenAttention(c_attn))
+
+ assert torch.all(q.weight == 1.0)
+ assert torch.all(k.weight == 2.0)
+ assert torch.all(v.weight == 3.0)
+
+ def test_split_matches_fused_output(self, adapter: QwenArchitectureAdapter) -> None:
+ """Guarantee that the split projections reproduce the fused layer's output."""
+ torch.manual_seed(0)
+ d_model = self.D_MODEL
+ c_attn = nn.Linear(d_model, 3 * d_model, bias=True) # random weights and biases for testing
+ mock = MockQwenAttention(c_attn)
+
+ q, k, v = adapter._split_qkv_matrix(mock)
+
+ x = torch.randn(2, 8, d_model)
+ fused = c_attn(x)
+ assert torch.allclose(q(x), fused[..., :d_model], atol=1e-6)
+ assert torch.allclose(k(x), fused[..., d_model : 2 * d_model], atol=1e-6)
+ assert torch.allclose(v(x), fused[..., 2 * d_model :], atol=1e-6)
+
+ def test_no_bias_gives_zero_biases(self, adapter: QwenArchitectureAdapter) -> None:
+ """When c_attn has no bias, the split projections get zero biases."""
+ d_model = self.D_MODEL
+ mock = MockQwenAttention(nn.Linear(d_model, 3 * d_model, bias=False))
+
+ q, k, v = adapter._split_qkv_matrix(mock)
+
+ for proj in (q, k, v):
+ assert torch.all(proj.bias == 0.0)
+
+ def test_conv1d_style_layout_is_transposed(self, adapter: QwenArchitectureAdapter) -> None:
+ """Conv1D-style storage has weight shape (d_model, 3*d_model), it is split on dim=1."""
+ d_model = self.D_MODEL
+ c_attn = nn.Linear(d_model, 3 * d_model, bias=False)
+ # Overwrite with a (d_model, 3*d_model) weight: columns are the Q/K/V thirds.
+ thirds = [torch.full((d_model, d_model), float(c)) for c in (1, 2, 3)]
+ c_attn.weight = nn.Parameter(torch.cat(thirds, dim=1))
+
+ q, k, v = adapter._split_qkv_matrix(MockQwenAttention(c_attn))
+
+ # Transposing a constant block keeps it constant and shapes stay (d_model, d_model).
+ assert q.weight.shape == (d_model, d_model)
+ assert torch.all(q.weight == 1.0)
+ assert torch.all(k.weight == 2.0)
+ assert torch.all(v.weight == 3.0)
+
+ def test_unexpected_shape_raises(self, adapter: QwenArchitectureAdapter) -> None:
+ """A c_attn whose shape is neither layout should fail loudly, not quietly misbehave."""
+ mock = MockQwenAttention(nn.Linear(8, 8))
+ with pytest.raises(ValueError, match="Unexpected c_attn weight shape"):
+ adapter._split_qkv_matrix(mock)
+
+
+# Factory registration test
+
+
+class TestQwenFactoryRegistration:
+ """The factory must resolve Qwen's HF architecture string to this adapter."""
+
+ def test_factory_returns_qwen_adapter(self, cfg: TransformerBridgeConfig) -> None:
+ built = ArchitectureAdapterFactory.select_architecture_adapter(cfg)
+ assert isinstance(built, QwenArchitectureAdapter)
diff --git a/tests/unit/model_bridge/supported_architectures/test_smollm3_adapter.py b/tests/unit/model_bridge/supported_architectures/test_smollm3_adapter.py
index d29b5abc8..f60f3f977 100644
--- a/tests/unit/model_bridge/supported_architectures/test_smollm3_adapter.py
+++ b/tests/unit/model_bridge/supported_architectures/test_smollm3_adapter.py
@@ -383,28 +383,23 @@ def _hook(x: Any, hook: Any) -> Any:
# Identity RoPE inputs keep this test focused on hook reshaping, not rotation math.
cos = ones(1, self.SEQ, self.D_HEAD)
sin = zeros(1, self.SEQ, self.D_HEAD)
- out = attn_bridge(hidden, position_embeddings=(cos, sin))
- out_tensor = out[0] if isinstance(out, tuple) else out
+ attn_bridge(hidden, position_embeddings=(cos, sin))
- return captured["q"], captured["k"], captured["v"], out_tensor
+ return captured["q"], captured["k"], captured["v"]
def test_hook_q_uses_n_heads(
self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
) -> None:
- q, _, _, _ = self._run_and_capture(wired_attn_bridge)
+ q, _, _ = self._run_and_capture(wired_attn_bridge)
assert q.shape == (self.BATCH, self.SEQ, self.N_HEADS, self.D_HEAD)
def test_hook_kv_use_n_kv_heads(
self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
) -> None:
- _, k, v, _ = self._run_and_capture(wired_attn_bridge)
+ _, k, v = self._run_and_capture(wired_attn_bridge)
assert k.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
assert v.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
- def test_attn_output_shape(self, wired_attn_bridge: PositionEmbeddingsAttentionBridge) -> None:
- _, _, _, out = self._run_and_capture(wired_attn_bridge)
- assert out.shape == (self.BATCH, self.SEQ, self.D_MODEL)
-
class TestSmolLM3NoPE:
"""The one piece of real adapter logic: per-layer RoPE suppression.
diff --git a/tests/unit/model_bridge/supported_architectures/test_stablelm_adapter.py b/tests/unit/model_bridge/supported_architectures/test_stablelm_adapter.py
new file mode 100644
index 000000000..9b1144bc6
--- /dev/null
+++ b/tests/unit/model_bridge/supported_architectures/test_stablelm_adapter.py
@@ -0,0 +1,711 @@
+"""Unit tests for the StableLmArchitectureAdapter (download-free, tiny programmatic configs
+plus small synthetic tensors and a fake attention module, no real checkpoints).
+
+Covered:
+- Anti-drift config defaults: `normalization_type="LN"`, `uses_rms_norm=False`, and
+ `attn_implementation="eager"`. Each is read by a distinct consumer (bridge selection,
+ NormalizationBridge runtime behavior, model-boot attention class).
+- Weight conversions: QKVO weights via the base helper, Q/K/V biases inline with
+ GQA-aware head counts, plus the no-`n_key_value_heads` fallback.
+- Component-mapping structure, bridge types (NormalizationBridge for LN, not RMS), and HF paths.
+- Block submodules across both `parallel_attn_mlp` branches (with and without `ln2`).
+- GQA forward hook shapes: a fake attention module confirms Q uses `n_heads` while
+ K/V use `n_key_value_heads`.
+- `setup_hook_compatibility`: QK-LayerNorm hook injection on StableLM v2 models, including
+ the no-op-when-absent path and defensive guards.
+- `setup_component_testing`: rotary embedding wiring on the template attention bridge and
+ on each bridge-model block, plus forcing eager attention on the HF model and per-layer
+ self-attention configs.
+- Architecture guards against drift.
+"""
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+import torch.nn as nn
+from torch import ones, randn, zeros
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.conversion_utils.conversion_steps.rearrange_tensor_conversion import (
+ RearrangeTensorConversion,
+)
+from transformer_lens.conversion_utils.param_processing_conversion import (
+ ParamProcessingConversion,
+)
+from transformer_lens.hook_points import HookPoint
+from transformer_lens.model_bridge.generalized_components import (
+ BlockBridge,
+ EmbeddingBridge,
+ GatedMLPBridge,
+ LinearBridge,
+ NormalizationBridge,
+ ParallelBlockBridge,
+ PositionEmbeddingsAttentionBridge,
+ RotaryEmbeddingBridge,
+ UnembeddingBridge,
+)
+from transformer_lens.model_bridge.supported_architectures.stablelm import (
+ StableLmArchitectureAdapter,
+)
+
+
+@pytest.fixture(scope="class")
+def cfg() -> TransformerBridgeConfig:
+ return TransformerBridgeConfig(
+ d_model=64,
+ d_head=16,
+ n_layers=2,
+ n_ctx=128,
+ n_heads=4,
+ n_key_value_heads=2,
+ d_vocab=256,
+ architecture="StableLmForCausalLM",
+ )
+
+
+@pytest.fixture(scope="class")
+def adapter(cfg: TransformerBridgeConfig) -> StableLmArchitectureAdapter:
+ return StableLmArchitectureAdapter(cfg)
+
+
+def _cfg(
+ *,
+ n_key_value_heads: int | None = 2,
+ parallel_attn_mlp: bool = False,
+) -> TransformerBridgeConfig:
+ # Keep dimensions tiny so adapter tests do not need HF downloads or real checkpoints.
+ # n_key_value_heads=None exercises the GQA fallback to n_heads in the adapter.
+ return TransformerBridgeConfig(
+ d_model=64,
+ d_head=16,
+ n_layers=2,
+ n_ctx=128,
+ n_heads=4,
+ n_key_value_heads=n_key_value_heads,
+ d_vocab=256,
+ architecture="StableLmForCausalLM",
+ parallel_attn_mlp=parallel_attn_mlp,
+ )
+
+
+def _mapping(adapter: StableLmArchitectureAdapter) -> dict:
+ """Narrow component_mapping (Optional on the base class) to a non-None dict.
+
+ Factored into a helper so each test stays a one-liner instead of repeating the
+ `assert ... is not None` prelude in every method.
+ """
+ mapping = adapter.component_mapping
+ assert mapping is not None
+ return mapping
+
+
+def _conversions(adapter: StableLmArchitectureAdapter) -> dict:
+ """weight_processing_conversions is Optional on the base class, assert it is populated."""
+ conversions = adapter.weight_processing_conversions
+ assert conversions is not None
+ return conversions
+
+
+def _param_conversion(adapter: StableLmArchitectureAdapter, key: str) -> ParamProcessingConversion:
+ conv = _conversions(adapter)[key]
+ assert isinstance(conv, ParamProcessingConversion)
+ return conv
+
+
+def _rearrange(adapter: StableLmArchitectureAdapter, key: str) -> RearrangeTensorConversion:
+ tensor_conversion = _param_conversion(adapter, key).tensor_conversion
+ assert isinstance(tensor_conversion, RearrangeTensorConversion)
+ return tensor_conversion
+
+
+def _fake_hf_model(rotary_emb: object) -> SimpleNamespace:
+ """Minimal HF model exposing only model.rotary_emb (no config, no layers)."""
+ return SimpleNamespace(model=SimpleNamespace(rotary_emb=rotary_emb))
+
+
+def _fake_hf_model_with_eager_targets(rotary_emb: object) -> SimpleNamespace:
+ """HF model whose top-level and per-layer attention implementation start non-eager."""
+ layers = [
+ SimpleNamespace(
+ self_attn=SimpleNamespace(config=SimpleNamespace(_attn_implementation="sdpa"))
+ )
+ for _ in range(2)
+ ]
+ return SimpleNamespace(
+ config=SimpleNamespace(_attn_implementation="sdpa"),
+ model=SimpleNamespace(rotary_emb=rotary_emb, layers=layers),
+ )
+
+
+class DummyAttention:
+ def __init__(self) -> None:
+ self.rotary_emb = None
+
+ def set_rotary_emb(self, rotary_emb: object) -> None:
+ self.rotary_emb = rotary_emb
+
+
+class DummyBlock:
+ def __init__(self, has_attention: bool = True) -> None:
+ if has_attention:
+ self.attn = DummyAttention()
+
+
+class DummyBridgeModel:
+ def __init__(self, blocks: list[DummyBlock]) -> None:
+ self.blocks = blocks
+
+
+class FakeStableLMAttention(nn.Module):
+ """Minimal StableLM-style attention module for adapter hook-shape tests.
+
+ Q is n_heads-wide, K/V are n_key_value_heads-wide, with no biases by default
+ (use_qkv_bias is False on stock StableLM).
+ """
+
+ def __init__(self, cfg: TransformerBridgeConfig) -> None:
+ super().__init__()
+ # PositionEmbeddingsAttentionBridge reads these HF-style attributes during forward.
+ self.head_dim = cfg.d_head
+ self.num_key_value_groups = cfg.n_heads // (cfg.n_key_value_heads or cfg.n_heads)
+ self.scaling = cfg.d_head**-0.5
+ self.attention_dropout = 0.0
+
+ kv_width = (cfg.n_key_value_heads or cfg.n_heads) * cfg.d_head
+ self.q_proj = nn.Linear(cfg.d_model, cfg.n_heads * cfg.d_head, bias=False)
+ self.k_proj = nn.Linear(cfg.d_model, kv_width, bias=False)
+ self.v_proj = nn.Linear(cfg.d_model, kv_width, bias=False)
+ self.o_proj = nn.Linear(cfg.n_heads * cfg.d_head, cfg.d_model, bias=False)
+
+
+class _DummyHFAttn(nn.Module):
+ """HF attention stand-in for the QK-LayerNorm setup_hook_compatibility tests.
+
+ `qk_layernorm` is read as a bool flag. `q_layernorm` and `k_layernorm` are
+ nn.Modules whose `forward` the adapter wraps to fire the new hook points.
+ """
+
+ def __init__(self, qk_layernorm: bool = True) -> None:
+ super().__init__()
+ self.qk_layernorm = qk_layernorm
+ self.q_layernorm = nn.Identity()
+ self.k_layernorm = nn.Identity()
+
+
+class _DummyAttnBridge(nn.Module):
+ """nn.Module stand-in for the attention bridge.
+
+ The adapter calls `add_module` on the bridge to register HookPoints, which
+ only works on nn.Modules.
+ """
+
+ def __init__(self, original_component: Any | None = None) -> None:
+ super().__init__()
+ self.original_component = original_component
+
+
+class _DummyBlockWithAttn:
+ def __init__(self, attn: Any) -> None:
+ self.attn = attn
+
+
+class _DummyBlockNoAttn:
+ pass
+
+
+class _DummyHookBridge:
+ """Stand-in for a built TransformerBridge: exposes `blocks` and a `_hook_registry` dict."""
+
+ def __init__(self, blocks: list[Any]) -> None:
+ self.blocks = blocks
+ self._hook_registry: dict[str, HookPoint] = {}
+
+
+class TestStableLMAdapterConfig:
+ """Anti-drift config flags: StableLM deviates from the Llama-family default of RMSNorm
+ and deliberately forces eager attention for parity with the bridge's eager-mode
+ reimplementation of attention."""
+
+ def test_normalization_type_is_ln(self, adapter: StableLmArchitectureAdapter) -> None:
+ """StableLM uses standard LayerNorm, not RMSNorm like the rest of the Llama family."""
+ assert adapter.cfg.normalization_type == "LN"
+
+ def test_uses_rms_norm_is_false(self, adapter: StableLmArchitectureAdapter) -> None:
+ """Paired with normalization_type=LN, but consumed by a different code path:
+ NormalizationBridge.uses_rms_norm reads this flag at forward time to decide
+ whether the bridge behaves as RMSNorm or LayerNorm. The ComponentTypes tests
+ only verify class identity (NormalizationBridge), not runtime behavior, so a
+ silent flip here would slip past them. Sibling LN-anti-drift adapters
+ (cohere, gpt_bigcode) keep the same assertion for the same reason."""
+ assert adapter.cfg.uses_rms_norm is False
+
+ def test_attn_implementation_is_eager(self, adapter: StableLmArchitectureAdapter) -> None:
+ """cfg.attn_implementation is read at model boot in sources/transformers.py and
+ passed to HF's loader, which instantiates different attention modules per value.
+ TestStableLMSetupComponentTesting only covers the post-load fixup path. The
+ load-time selection path is only guarded by this assertion."""
+ assert adapter.cfg.attn_implementation == "eager"
+
+
+class TestStableLMWeightConversions:
+ """StableLM declares QKVO weight conversions via the base helper plus Q/K/V bias
+ conversions inline (for the optional `use_qkv_bias=True` variants like stable-code-3b).
+ The bias rearranges are the adapter's own choice, so their patterns and head-count
+ axes are worth asserting. The QKVO weight patterns come from the base helper and are
+ covered by the base class's own tests."""
+
+ def test_conversion_keys_are_exactly_qkvo_weights_plus_qkv_biases(
+ self, adapter: StableLmArchitectureAdapter
+ ) -> None:
+ """O has no bias (no `b_O` in the conversion map), and the MLP has no biases."""
+ assert set(_conversions(adapter).keys()) == {
+ "blocks.{i}.attn.q.weight",
+ "blocks.{i}.attn.k.weight",
+ "blocks.{i}.attn.v.weight",
+ "blocks.{i}.attn.o.weight",
+ "blocks.{i}.attn.q.bias",
+ "blocks.{i}.attn.k.bias",
+ "blocks.{i}.attn.v.bias",
+ }
+
+ def test_q_bias_rearrange_uses_n_heads(self, adapter: StableLmArchitectureAdapter) -> None:
+ rearrange = _rearrange(adapter, "blocks.{i}.attn.q.bias")
+ assert rearrange.pattern == "(n h) -> n h"
+ assert rearrange.axes_lengths.get("n") == 4
+
+ def test_kv_bias_rearrange_uses_n_kv_heads(self, adapter: StableLmArchitectureAdapter) -> None:
+ """GQA: K/V biases follow n_key_value_heads (2), not n_heads."""
+ for slot in ("k", "v"):
+ rearrange = _rearrange(adapter, f"blocks.{{i}}.attn.{slot}.bias")
+ assert rearrange.pattern == "(n h) -> n h"
+ assert rearrange.axes_lengths.get("n") == 2
+
+ def test_gqa_fallback_to_n_heads_without_kv_heads(self) -> None:
+ """Without n_key_value_heads, K/V biases fall back to n_heads (the `or self.cfg.n_heads`
+ clause in the adapter)."""
+ adapter = StableLmArchitectureAdapter(_cfg(n_key_value_heads=None))
+ for slot in ("k", "v"):
+ assert _rearrange(adapter, f"blocks.{{i}}.attn.{slot}.bias").axes_lengths["n"] == 4
+
+
+class TestStableLMComponentMapping:
+ """Structure of the component mapping: required keys and HF module paths."""
+
+ def test_has_required_top_level_keys(self, adapter: StableLmArchitectureAdapter) -> None:
+ mapping = _mapping(adapter)
+ for key in ("embed", "rotary_emb", "blocks", "ln_final", "unembed"):
+ assert key in mapping, f"Missing top-level key: {key!r}"
+
+ def test_top_level_hf_paths(self, adapter: StableLmArchitectureAdapter) -> None:
+ mapping = _mapping(adapter)
+ assert mapping["embed"].name == "model.embed_tokens"
+ assert mapping["rotary_emb"].name == "model.rotary_emb"
+ assert mapping["blocks"].name == "model.layers"
+ assert mapping["ln_final"].name == "model.norm"
+ assert mapping["unembed"].name == "lm_head"
+
+
+class TestStableLMComponentTypes:
+ """Bridge classes selected for each component slot. The norms are deliberately
+ NormalizationBridge (LayerNorm), not RMSNormalizationBridge: this is the structural
+ consequence of `normalization_type='LN'` in the AdapterConfig tests above."""
+
+ def test_rotary_emb_is_rotary_bridge(self, adapter: StableLmArchitectureAdapter) -> None:
+ assert isinstance(_mapping(adapter)["rotary_emb"], RotaryEmbeddingBridge)
+
+ def test_ln_final_is_layernorm_not_rms(self, adapter: StableLmArchitectureAdapter) -> None:
+ """LN, not RMSNorm. The anti-drift normalization_type flag has to wire here."""
+ ln_final = _mapping(adapter)["ln_final"]
+ assert isinstance(ln_final, NormalizationBridge)
+
+ def test_block_attn_is_position_embeddings_bridge(
+ self, adapter: StableLmArchitectureAdapter
+ ) -> None:
+ block = _mapping(adapter)["blocks"]
+ assert isinstance(block.submodules["attn"], PositionEmbeddingsAttentionBridge)
+
+ def test_block_mlp_is_gated_mlp_bridge(self, adapter: StableLmArchitectureAdapter) -> None:
+ block = _mapping(adapter)["blocks"]
+ assert isinstance(block.submodules["mlp"], GatedMLPBridge)
+
+ def test_block_norms_are_layernorm_not_rms(self, adapter: StableLmArchitectureAdapter) -> None:
+ block = _mapping(adapter)["blocks"]
+ assert isinstance(block.submodules["ln1"], NormalizationBridge)
+ assert isinstance(block.submodules["ln2"], NormalizationBridge)
+
+ def test_embed_and_unembed_are_correct_bridge_types(
+ self, adapter: StableLmArchitectureAdapter
+ ) -> None:
+ mapping = _mapping(adapter)
+ assert isinstance(mapping["embed"], EmbeddingBridge)
+ assert isinstance(mapping["unembed"], UnembeddingBridge)
+
+ def test_attn_q_k_v_o_are_linear_bridges(self, adapter: StableLmArchitectureAdapter) -> None:
+ attn = _mapping(adapter)["blocks"].submodules["attn"]
+ for slot in ("q", "k", "v", "o"):
+ assert isinstance(attn.submodules[slot], LinearBridge)
+
+ def test_mlp_gate_in_out_are_linear_bridges(self, adapter: StableLmArchitectureAdapter) -> None:
+ mlp = _mapping(adapter)["blocks"].submodules["mlp"]
+ for slot in ("gate", "in", "out"):
+ assert isinstance(mlp.submodules[slot], LinearBridge)
+
+
+class TestStableLMBlockSubmodulesDefault:
+ """Default branch (`parallel_attn_mlp=False`): sequential residual with separate ln1
+ and ln2 norms. The block carries four submodules: ln1, ln2, attn, mlp."""
+
+ def test_default_block_submodule_keys(self, adapter: StableLmArchitectureAdapter) -> None:
+ block = _mapping(adapter)["blocks"]
+ # Strict type identity: the sequential branch must NOT pick ParallelBlockBridge
+ # (which is a BlockBridge subclass), so a plain `isinstance` would not discriminate.
+ assert type(block) is BlockBridge
+ assert set(block.submodules.keys()) == {"ln1", "ln2", "attn", "mlp"}
+
+ def test_ln2_uses_post_attention_layernorm_hf_name(
+ self, adapter: StableLmArchitectureAdapter
+ ) -> None:
+ block = _mapping(adapter)["blocks"]
+ assert block.submodules["ln2"].name == "post_attention_layernorm"
+
+ def test_ln1_uses_input_layernorm_hf_name(self, adapter: StableLmArchitectureAdapter) -> None:
+ block = _mapping(adapter)["blocks"]
+ assert block.submodules["ln1"].name == "input_layernorm"
+
+
+class TestStableLMBlockSubmodulesParallelResidual:
+ """Parallel-residual branch (`parallel_attn_mlp=True`): both attn and MLP read from
+ ln1's output, so HF sets post_attention_layernorm=None. The block carries three
+ submodules (ln1, attn, mlp) inside a `ParallelBlockBridge` container."""
+
+ @pytest.fixture
+ def parallel_adapter(self) -> StableLmArchitectureAdapter:
+ return StableLmArchitectureAdapter(_cfg(parallel_attn_mlp=True))
+
+ def test_parallel_block_submodule_keys(
+ self, parallel_adapter: StableLmArchitectureAdapter
+ ) -> None:
+ """Container is ParallelBlockBridge so the no-ln2 layout is the supported shape.
+ BlockBridge's C15 guard rejects `attn + mlp` without `ln2`, which is exactly the
+ regression #1386 was filed for."""
+ block = _mapping(parallel_adapter)["blocks"]
+ assert isinstance(block, ParallelBlockBridge)
+ assert set(block.submodules.keys()) == {"ln1", "attn", "mlp"}
+
+
+class TestStableLMGQAHookShapes:
+ """Wire a fake attention module into the bridge and verify GQA hook shapes.
+
+ Spec assertions cannot prove the bridge reshapes activations correctly.
+ Here Q must surface n_heads while K/V surface n_key_value_heads,
+ which is the whole point of grouped-query attention.
+ """
+
+ N_HEADS = 4
+ N_KV_HEADS = 2
+ D_MODEL = 64
+ D_HEAD = D_MODEL // N_HEADS
+ BATCH = 2
+ SEQ = 8
+
+ @pytest.fixture
+ def wired_attn_bridge(self) -> PositionEmbeddingsAttentionBridge:
+ adapter = StableLmArchitectureAdapter(_cfg(n_key_value_heads=self.N_KV_HEADS))
+ fake_attn = FakeStableLMAttention(adapter.cfg)
+ attn_bridge = _mapping(adapter)["blocks"].submodules["attn"]
+ assert isinstance(attn_bridge, PositionEmbeddingsAttentionBridge)
+ attn_bridge.set_original_component(fake_attn)
+ # A full TransformerBridge build materializes these child bridge modules for us.
+ # This unit test wires them by hand so it can stay download-free.
+ for name, original in {
+ "q": fake_attn.q_proj,
+ "k": fake_attn.k_proj,
+ "v": fake_attn.v_proj,
+ "o": fake_attn.o_proj,
+ }.items():
+ submodule = attn_bridge.submodules[name]
+ submodule.set_original_component(original)
+ attn_bridge.add_module(name, submodule)
+ attn_bridge.setup_hook_compatibility()
+ return attn_bridge
+
+ def _run_and_capture(self, attn_bridge: PositionEmbeddingsAttentionBridge) -> tuple:
+ captured: dict = {}
+
+ def _capture(name: str) -> Any:
+ def _hook(x: Any, hook: Any) -> Any:
+ captured[name] = x.detach()
+ return x
+
+ return _hook
+
+ attn_bridge.q.hook_out.add_hook(_capture("q"))
+ attn_bridge.k.hook_out.add_hook(_capture("k"))
+ attn_bridge.v.hook_out.add_hook(_capture("v"))
+
+ hidden = randn(self.BATCH, self.SEQ, self.D_MODEL)
+ # Identity RoPE inputs keep this test focused on hook reshaping, not rotation math.
+ cos = ones(1, self.SEQ, self.D_HEAD)
+ sin = zeros(1, self.SEQ, self.D_HEAD)
+ attn_bridge(hidden, position_embeddings=(cos, sin))
+
+ return captured["q"], captured["k"], captured["v"]
+
+ def test_hook_q_uses_n_heads(
+ self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
+ ) -> None:
+ q, _, _ = self._run_and_capture(wired_attn_bridge)
+ assert q.shape == (self.BATCH, self.SEQ, self.N_HEADS, self.D_HEAD)
+
+ def test_hook_kv_use_n_kv_heads(
+ self, wired_attn_bridge: PositionEmbeddingsAttentionBridge
+ ) -> None:
+ _, k, v = self._run_and_capture(wired_attn_bridge)
+ assert k.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
+ assert v.shape == (self.BATCH, self.SEQ, self.N_KV_HEADS, self.D_HEAD)
+
+
+class TestStableLMSetupHookCompatibility:
+ """`setup_hook_compatibility` injects `hook_q_layernorm` and `hook_k_layernorm` HookPoints
+ onto every attention bridge whose underlying HF attention has `qk_layernorm=True`
+ (StableLM v2 models like stablelm-2-12b). The hooks are added as bridge submodules,
+ registered in `bridge._hook_registry` with canonical TL-style names, and the HF
+ q_layernorm/k_layernorm forward methods are wrapped to fire them.
+
+ Coverage: every branch of the override, including the no-op-when-disabled path and
+ each defensive guard.
+ """
+
+ @pytest.fixture
+ def adapter_only(self) -> StableLmArchitectureAdapter:
+ return StableLmArchitectureAdapter(_cfg())
+
+ def _hook_bridge(self, blocks: list[Any]) -> _DummyHookBridge:
+ return _DummyHookBridge(blocks)
+
+ def test_no_op_when_bridge_has_no_blocks(
+ self, adapter_only: StableLmArchitectureAdapter
+ ) -> None:
+ """Guard branch: a bridge without `.blocks` must not raise."""
+ bridge = SimpleNamespace() # no blocks attribute at all
+ adapter_only.setup_hook_compatibility(bridge)
+
+ def test_no_op_when_qk_layernorm_disabled(
+ self, adapter_only: StableLmArchitectureAdapter
+ ) -> None:
+ """Stock StableLM (v1) has qk_layernorm=False, so no hooks are injected."""
+ hf_attn = _DummyHFAttn(qk_layernorm=False)
+ attn_bridge = _DummyAttnBridge(original_component=hf_attn)
+ bridge = self._hook_bridge([_DummyBlockWithAttn(attn_bridge)])
+
+ adapter_only.setup_hook_compatibility(bridge)
+
+ assert not hasattr(attn_bridge, "hook_q_layernorm")
+ assert not hasattr(attn_bridge, "hook_k_layernorm")
+ assert bridge._hook_registry == {}
+
+ def test_skips_block_without_attn(self, adapter_only: StableLmArchitectureAdapter) -> None:
+ """A block with no `attn` attribute must be silently skipped."""
+ hf_attn = _DummyHFAttn(qk_layernorm=True)
+ attn_bridge_with = _DummyAttnBridge(original_component=hf_attn)
+ bridge = self._hook_bridge([_DummyBlockNoAttn(), _DummyBlockWithAttn(attn_bridge_with)])
+
+ adapter_only.setup_hook_compatibility(bridge)
+
+ # Block 1 still received hooks, block 0 had no attn and was skipped without error.
+ assert hasattr(attn_bridge_with, "hook_q_layernorm")
+ assert "blocks.1.attn.hook_q_layernorm" in bridge._hook_registry
+
+ def test_skips_block_with_none_original_component(
+ self, adapter_only: StableLmArchitectureAdapter
+ ) -> None:
+ """When the attention bridge has no original_component (not yet built), skip
+ rather than raise."""
+ attn_bridge = _DummyAttnBridge(original_component=None)
+ bridge = self._hook_bridge([_DummyBlockWithAttn(attn_bridge)])
+
+ adapter_only.setup_hook_compatibility(bridge)
+
+ assert not hasattr(attn_bridge, "hook_q_layernorm")
+ assert bridge._hook_registry == {}
+
+ def test_adds_hook_modules_to_attn_bridge(
+ self, adapter_only: StableLmArchitectureAdapter
+ ) -> None:
+ """Happy path: HookPoints are added as submodules of the attention bridge."""
+ hf_attn = _DummyHFAttn(qk_layernorm=True)
+ attn_bridge = _DummyAttnBridge(original_component=hf_attn)
+ bridge = self._hook_bridge([_DummyBlockWithAttn(attn_bridge)])
+
+ adapter_only.setup_hook_compatibility(bridge)
+
+ assert isinstance(attn_bridge.hook_q_layernorm, HookPoint)
+ assert isinstance(attn_bridge.hook_k_layernorm, HookPoint)
+
+ def test_registers_hooks_in_bridge_hook_registry_with_canonical_names(
+ self, adapter_only: StableLmArchitectureAdapter
+ ) -> None:
+ """Hooks are registered under TL-canonical names so the registry scanner can find them
+ (the scanner skips _original_component subtrees, so the adapter wires registry entries
+ directly)."""
+ blocks = []
+ for _ in range(3):
+ hf_attn = _DummyHFAttn(qk_layernorm=True)
+ attn_bridge = _DummyAttnBridge(original_component=hf_attn)
+ blocks.append(_DummyBlockWithAttn(attn_bridge))
+ bridge = self._hook_bridge(blocks)
+
+ adapter_only.setup_hook_compatibility(bridge)
+
+ expected = set()
+ for i in range(3):
+ expected.add(f"blocks.{i}.attn.hook_q_layernorm")
+ expected.add(f"blocks.{i}.attn.hook_k_layernorm")
+ assert set(bridge._hook_registry.keys()) == expected
+ # And every registry entry points at the same HookPoint object on the bridge.
+ for i, block in enumerate(blocks):
+ assert (
+ bridge._hook_registry[f"blocks.{i}.attn.hook_q_layernorm"]
+ is block.attn.hook_q_layernorm
+ )
+ assert (
+ bridge._hook_registry[f"blocks.{i}.attn.hook_k_layernorm"]
+ is block.attn.hook_k_layernorm
+ )
+
+ def test_q_layernorm_forward_wrap_fires_hook(
+ self, adapter_only: StableLmArchitectureAdapter
+ ) -> None:
+ """Behavioral assertion: calling the HF q_layernorm.forward fires the new hook
+ with the layernorm output. Without the wrap, the hook would never run."""
+ hf_attn = _DummyHFAttn(qk_layernorm=True)
+ attn_bridge = _DummyAttnBridge(original_component=hf_attn)
+ bridge = self._hook_bridge([_DummyBlockWithAttn(attn_bridge)])
+ adapter_only.setup_hook_compatibility(bridge)
+
+ captured: dict[str, torch.Tensor] = {}
+
+ def _capture(x: torch.Tensor, hook: HookPoint) -> torch.Tensor:
+ captured["q"] = x.detach()
+ return x
+
+ attn_bridge.hook_q_layernorm.add_hook(_capture)
+ x = randn(2, 4, 16)
+ out = hf_attn.q_layernorm.forward(x)
+ # nn.Identity passes through, so the hook saw exactly x.
+ assert "q" in captured
+ assert torch.equal(captured["q"], x)
+ assert torch.equal(out, x)
+
+ def test_k_layernorm_forward_wrap_fires_hook(
+ self, adapter_only: StableLmArchitectureAdapter
+ ) -> None:
+ """Same behavioral assertion for the K-side wrap."""
+ hf_attn = _DummyHFAttn(qk_layernorm=True)
+ attn_bridge = _DummyAttnBridge(original_component=hf_attn)
+ bridge = self._hook_bridge([_DummyBlockWithAttn(attn_bridge)])
+ adapter_only.setup_hook_compatibility(bridge)
+
+ captured: dict[str, torch.Tensor] = {}
+
+ def _capture(x: torch.Tensor, hook: HookPoint) -> torch.Tensor:
+ captured["k"] = x.detach()
+ return x
+
+ attn_bridge.hook_k_layernorm.add_hook(_capture)
+ x = randn(2, 4, 16)
+ out = hf_attn.k_layernorm.forward(x)
+ assert "k" in captured
+ assert torch.equal(captured["k"], x)
+ assert torch.equal(out, x)
+
+
+class TestStableLMSetupComponentTesting:
+ """`setup_component_testing` wires the shared rotary embedding onto the template
+ attention bridge and onto each bridge-model block's attention. It also forces eager
+ attention on the HF model (top-level config and per-layer self_attn.config) for
+ numerical parity with the bridge's eager-mode reimplementation."""
+
+ def test_sets_rotary_emb_on_template_attention(
+ self, adapter: StableLmArchitectureAdapter
+ ) -> None:
+ rotary_emb = object()
+ attn_template = adapter.get_generalized_component("blocks.0.attn")
+ assert isinstance(attn_template, PositionEmbeddingsAttentionBridge)
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb))
+
+ assert attn_template._rotary_emb is rotary_emb
+
+ def test_sets_rotary_emb_on_each_bridge_model_attention(
+ self, adapter: StableLmArchitectureAdapter
+ ) -> None:
+ rotary_emb = object()
+ bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(), DummyBlock()])
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model)
+
+ for block in bridge_model.blocks:
+ assert block.attn.rotary_emb is rotary_emb
+
+ def test_skips_bridge_blocks_without_attention(
+ self, adapter: StableLmArchitectureAdapter
+ ) -> None:
+ rotary_emb = object()
+ bridge_model = DummyBridgeModel([DummyBlock(), DummyBlock(has_attention=False)])
+
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb), bridge_model=bridge_model)
+
+ assert bridge_model.blocks[0].attn.rotary_emb is rotary_emb
+
+ def test_forces_eager_attention_implementation(
+ self, adapter: StableLmArchitectureAdapter
+ ) -> None:
+ """Bridge attention only matches HF under eager attention, so it is forced on
+ at both the top-level config and on each per-layer self_attn.config."""
+ hf_model = _fake_hf_model_with_eager_targets(object())
+
+ adapter.setup_component_testing(hf_model)
+
+ assert hf_model.config._attn_implementation == "eager"
+ for layer in hf_model.model.layers:
+ assert layer.self_attn.config._attn_implementation == "eager"
+
+ def test_tolerates_minimal_hf_model_without_config_or_layers(
+ self, adapter: StableLmArchitectureAdapter
+ ) -> None:
+ """The defensive hasattr branches must not raise when config/layers are absent."""
+ rotary_emb = object()
+ # _fake_hf_model exposes only model.rotary_emb (no config, no layers).
+ adapter.setup_component_testing(_fake_hf_model(rotary_emb))
+
+ attn_template = adapter.get_generalized_component("blocks.0.attn")
+ assert isinstance(attn_template, PositionEmbeddingsAttentionBridge)
+ assert attn_template._rotary_emb is rotary_emb
+
+
+class TestStableLMArchitectureGuards:
+ """Guards against drift from the StableLM conversion-map contract."""
+
+ def test_no_normalization_weights_in_conversions(
+ self, adapter: StableLmArchitectureAdapter
+ ) -> None:
+ """LayerNorm weights and biases are not folded by this adapter, so no norm keys
+ appear in the conversion map."""
+ for key in _conversions(adapter):
+ assert "ln1" not in key
+ assert "ln2" not in key
+ assert "ln_final" not in key
+
+ def test_no_output_or_mlp_bias_conversions(self, adapter: StableLmArchitectureAdapter) -> None:
+ """O has no bias (HF does not expose one) and the MLP projections are bias-free.
+ The only biases declared are Q/K/V on the attention input."""
+ for key in _conversions(adapter):
+ if key.endswith(".bias"):
+ assert key in {
+ "blocks.{i}.attn.q.bias",
+ "blocks.{i}.attn.k.bias",
+ "blocks.{i}.attn.v.bias",
+ }
diff --git a/tests/unit/tools/test_direct_logit_attribution.py b/tests/unit/tools/test_direct_logit_attribution.py
new file mode 100644
index 000000000..2b77805e7
--- /dev/null
+++ b/tests/unit/tools/test_direct_logit_attribution.py
@@ -0,0 +1,53 @@
+"""Unit tests for Direct Logit Attribution guards and argument validation.
+
+These exercise the fast-failing checks (argument validation, the Bridge
+compatibility-mode requirement, and the hybrid-architecture refusal) without
+loading a real model — using a ``spec``-ed mock TransformerBridge so the checks
+fire before any forward pass.
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.tools.analysis import direct_logit_attribution
+
+
+def _mock_bridge(compatibility_mode=True, layer_types=("attn+mlp",)):
+ """A mock TransformerBridge that satisfies isinstance checks."""
+ bridge = MagicMock(spec=TransformerBridge)
+ bridge.compatibility_mode = compatibility_mode
+ bridge.layer_types.return_value = list(layer_types)
+ return bridge
+
+
+def test_invalid_unit_raises():
+ bridge = _mock_bridge()
+ with pytest.raises(ValueError, match="unit must be one of"):
+ direct_logit_attribution(bridge, "hi", answer_tokens=" world", unit="neuron")
+
+
+def test_missing_answer_tokens_raises():
+ bridge = _mock_bridge()
+ with pytest.raises(ValueError, match="answer_tokens is required"):
+ direct_logit_attribution(bridge, "hi")
+
+
+def test_requires_compatibility_mode():
+ bridge = _mock_bridge(compatibility_mode=False)
+ with pytest.raises(ValueError, match="compatibility mode"):
+ direct_logit_attribution(bridge, "hi", answer_tokens=" world")
+
+
+def test_rejects_hybrid_architecture():
+ bridge = _mock_bridge(layer_types=("attn+mlp", "mamba+mlp"))
+ with pytest.raises(NotImplementedError, match="hybrid"):
+ direct_logit_attribution(bridge, "hi", answer_tokens=" world")
+
+
+def test_direct_logit_attribution_is_exported_from_analysis_package():
+ from transformer_lens.tools import analysis
+
+ assert analysis.direct_logit_attribution is direct_logit_attribution
+ assert hasattr(analysis, "DirectLogitAttribution")
diff --git a/transformer_lens/BertNextSentencePrediction.py b/transformer_lens/BertNextSentencePrediction.py
index eb38e6879..fb6f8e8fa 100644
--- a/transformer_lens/BertNextSentencePrediction.py
+++ b/transformer_lens/BertNextSentencePrediction.py
@@ -153,7 +153,7 @@ def forward(
"[CLS] Sentence A [SEP] Sentence B [SEP]", token_type_ids would be
[0, 0, ..., 0, 1, ..., 1, 1]. `0` represents tokens from Sentence A,
`1` from Sentence B. If not provided, BERT assumes a single sequence input.
- This parameter gets inferred from the the tokenizer if input is a string or list of strings.
+ This parameter gets inferred from the tokenizer if input is a string or list of strings.
Shape is (batch_size, sequence_length).
one_zero_attention_mask: Optional[torch.Tensor]: A binary mask which indicates
which tokens should be attended to (1) and which should be ignored (0).
diff --git a/transformer_lens/HookedEncoder.py b/transformer_lens/HookedEncoder.py
index f71a9c75a..9dffc07e4 100644
--- a/transformer_lens/HookedEncoder.py
+++ b/transformer_lens/HookedEncoder.py
@@ -251,7 +251,7 @@ def forward(
"[CLS] Sentence A [SEP] Sentence B [SEP]", token_type_ids would be
[0, 0, ..., 0, 1, ..., 1, 1]. `0` represents tokens from Sentence A,
`1` from Sentence B. If not provided, BERT assumes a single sequence input.
- This parameter gets inferred from the the tokenizer if input is a string or list of strings.
+ This parameter gets inferred from the tokenizer if input is a string or list of strings.
Shape is (batch_size, sequence_length).
one_zero_attention_mask: Optional[torch.Tensor]: A binary mask which indicates
which tokens should be attended to (1) and which should be ignored (0).
diff --git a/transformer_lens/HookedTransformer.py b/transformer_lens/HookedTransformer.py
index 9baf479e5..8370bbf8a 100644
--- a/transformer_lens/HookedTransformer.py
+++ b/transformer_lens/HookedTransformer.py
@@ -1509,7 +1509,7 @@ def init_weights(self):
The default PyTorch scheme is the following: all linear layers use uniform(-1/sqrt(fan_in),
1/sqrt(fan_in)) for weights, and uniform(-1/sqrt(fan_in), 1/sqrt(fan_in)) for biases. For
biases, fan_in is computed using the fan_in for the weight matrix of the linear layer. Note
- tha it *does not actually* use Kaiming initialization, despite the fact that it calls the
+ that it *does not actually* use Kaiming initialization, despite the fact that it calls the
function.
However, for Transformer blocks, it instead initializes biases to zero and weights using Xavier uniform, that
diff --git a/transformer_lens/benchmarks/backward_gradients.py b/transformer_lens/benchmarks/backward_gradients.py
index e612127f0..4e75009d7 100644
--- a/transformer_lens/benchmarks/backward_gradients.py
+++ b/transformer_lens/benchmarks/backward_gradients.py
@@ -11,6 +11,7 @@
make_grad_capture_hook,
safe_allclose,
)
+from transformer_lens.hook_points import HookPoint
from transformer_lens.model_bridge import TransformerBridge
@@ -44,12 +45,15 @@ def benchmark_backward_hooks(
hook_names = list(bridge._hook_registry.keys())
# Register backward hooks on bridge
- bridge_handles = []
+ bridge_hook_points: list[HookPoint] = []
for hook_name in hook_names:
if hook_name in bridge.hook_dict:
hook_point = bridge.hook_dict[hook_name]
- handle = hook_point.add_hook(make_grad_capture_hook(bridge_gradients, hook_name, return_none=True), dir="bwd") # type: ignore[func-returns-value]
- bridge_handles.append(handle)
+ hook_point.add_hook(
+ make_grad_capture_hook(bridge_gradients, hook_name, return_none=True),
+ dir="bwd",
+ )
+ bridge_hook_points.append(hook_point)
# Run bridge forward and backward
bridge_output = bridge(test_text)
@@ -57,9 +61,8 @@ def benchmark_backward_hooks(
bridge_loss.backward()
# Clean up hooks
- for handle in bridge_handles:
- if handle is not None:
- handle.remove()
+ for hook_point in bridge_hook_points:
+ hook_point.remove_hooks(dir="bwd")
if reference_model is None:
# No reference - just verify gradients were captured
@@ -77,12 +80,15 @@ def benchmark_backward_hooks(
return result
# Register backward hooks on reference model
- reference_handles = []
+ reference_hook_points: list[HookPoint] = []
for hook_name in hook_names:
if hook_name in reference_model.hook_dict:
hook_point = reference_model.hook_dict[hook_name]
- handle = hook_point.add_hook(make_grad_capture_hook(reference_gradients, hook_name, return_none=True), dir="bwd") # type: ignore[func-returns-value]
- reference_handles.append(handle)
+ hook_point.add_hook(
+ make_grad_capture_hook(reference_gradients, hook_name, return_none=True),
+ dir="bwd",
+ )
+ reference_hook_points.append(hook_point)
# Run reference forward and backward
reference_output = reference_model(test_text)
@@ -90,9 +96,8 @@ def benchmark_backward_hooks(
reference_loss.backward()
# Clean up hooks
- for handle in reference_handles:
- if handle is not None:
- handle.remove()
+ for hook_point in reference_hook_points:
+ hook_point.remove_hooks(dir="bwd")
# Compare gradients
common_hooks = set(bridge_gradients.keys()) & set(reference_gradients.keys())
@@ -295,12 +300,15 @@ def benchmark_critical_backward_hooks(
bridge_gradients: Dict[str, torch.Tensor] = {}
# Register backward hooks on bridge
- bridge_handles = []
+ bridge_hook_points: list[HookPoint] = []
for hook_name in critical_hooks:
if hook_name in bridge.hook_dict:
hook_point = bridge.hook_dict[hook_name]
- handle = hook_point.add_hook(make_grad_capture_hook(bridge_gradients, hook_name, return_none=True), dir="bwd") # type: ignore[func-returns-value]
- bridge_handles.append(handle)
+ hook_point.add_hook(
+ make_grad_capture_hook(bridge_gradients, hook_name, return_none=True),
+ dir="bwd",
+ )
+ bridge_hook_points.append(hook_point)
# Run bridge forward and backward
bridge_output = bridge(test_text)
@@ -308,9 +316,8 @@ def benchmark_critical_backward_hooks(
bridge_loss.backward()
# Clean up hooks
- for handle in bridge_handles:
- if handle is not None:
- handle.remove()
+ for hook_point in bridge_hook_points:
+ hook_point.remove_hooks(dir="bwd")
if reference_model is None:
# No reference - just verify gradients were captured
@@ -331,12 +338,15 @@ def benchmark_critical_backward_hooks(
# Register backward hooks on reference model
reference_gradients: Dict[str, torch.Tensor] = {}
- reference_handles = []
+ reference_hook_points: list[HookPoint] = []
for hook_name in critical_hooks:
if hook_name in reference_model.hook_dict:
hook_point = reference_model.hook_dict[hook_name]
- handle = hook_point.add_hook(make_grad_capture_hook(reference_gradients, hook_name, return_none=True), dir="bwd") # type: ignore[func-returns-value]
- reference_handles.append(handle)
+ hook_point.add_hook(
+ make_grad_capture_hook(reference_gradients, hook_name, return_none=True),
+ dir="bwd",
+ )
+ reference_hook_points.append(hook_point)
# Run reference forward and backward
reference_output = reference_model(test_text)
@@ -344,9 +354,8 @@ def benchmark_critical_backward_hooks(
reference_loss.backward()
# Clean up hooks
- for handle in reference_handles:
- if handle is not None:
- handle.remove()
+ for hook_point in reference_hook_points:
+ hook_point.remove_hooks(dir="bwd")
# Compare gradients
mismatches = []
diff --git a/transformer_lens/benchmarks/hook_registration.py b/transformer_lens/benchmarks/hook_registration.py
index 14b823e7f..d74932f08 100644
--- a/transformer_lens/benchmarks/hook_registration.py
+++ b/transformer_lens/benchmarks/hook_registration.py
@@ -13,6 +13,7 @@
filter_expected_missing_hooks,
make_capture_hook,
)
+from transformer_lens.hook_points import HookPoint
from transformer_lens.model_bridge import TransformerBridge
@@ -134,13 +135,13 @@ def benchmark_forward_hooks(
hook_names = list(bridge.hook_dict.keys())
# Register hooks on bridge and track missing hooks
- bridge_handles = []
+ bridge_hook_points: list[tuple[str, HookPoint]] = []
missing_from_bridge = []
for hook_name in hook_names:
if hook_name in bridge.hook_dict:
hook_point = bridge.hook_dict[hook_name]
- handle = hook_point.add_hook(make_capture_hook(bridge_activations, hook_name)) # type: ignore[func-returns-value]
- bridge_handles.append((hook_name, handle))
+ hook_point.add_hook(make_capture_hook(bridge_activations, hook_name))
+ bridge_hook_points.append((hook_name, hook_point))
else:
missing_from_bridge.append(hook_name)
@@ -152,12 +153,11 @@ def benchmark_forward_hooks(
_ = bridge(test_text)
# Clean up bridge hooks
- for hook_name, handle in bridge_handles:
- if handle is not None:
- handle.remove()
+ for _, hook_point in bridge_hook_points:
+ hook_point.remove_hooks()
# Check for hooks that didn't fire (registered but no activation captured)
- registered_hooks = {name for name, _ in bridge_handles}
+ registered_hooks = {name for name, _ in bridge_hook_points}
hooks_that_didnt_fire = registered_hooks - set(bridge_activations.keys())
if reference_model is None:
@@ -182,12 +182,12 @@ def benchmark_forward_hooks(
)
# Register hooks on reference model
- reference_handles = []
+ reference_hook_points: list[HookPoint] = []
for hook_name in hook_names:
if hook_name in reference_model.hook_dict:
hook_point = reference_model.hook_dict[hook_name]
- handle = hook_point.add_hook(make_capture_hook(reference_activations, hook_name)) # type: ignore[func-returns-value]
- reference_handles.append(handle)
+ hook_point.add_hook(make_capture_hook(reference_activations, hook_name))
+ reference_hook_points.append(hook_point)
# Run reference forward pass
with torch.no_grad():
@@ -197,9 +197,8 @@ def benchmark_forward_hooks(
_ = reference_model(test_text)
# Clean up reference hooks
- for handle in reference_handles:
- if handle is not None:
- handle.remove()
+ for hook_point in reference_hook_points:
+ hook_point.remove_hooks()
# CRITICAL CHECK: Bridge must have all hooks that reference has.
# Filter out hooks that bridge models inherently don't have.
@@ -363,7 +362,7 @@ def benchmark_gated_hooks_fire(
tested_flags.append(flag_name)
try:
activations: dict[str, torch.Tensor] = {}
- handles: list[tuple[str, object]] = []
+ bridge_hook_points: list[HookPoint] = []
target_hook_names = [
name
for name in bridge.hook_dict
@@ -375,8 +374,8 @@ def benchmark_gated_hooks_fire(
]
for hname in target_hook_names:
hp = bridge.hook_dict[hname]
- h = hp.add_hook(make_capture_hook(activations, hname)) # type: ignore[func-returns-value]
- handles.append((hname, h))
+ hp.add_hook(make_capture_hook(activations, hname))
+ bridge_hook_points.append(hp)
with torch.no_grad():
if prepend_bos is not None:
@@ -384,9 +383,8 @@ def benchmark_gated_hooks_fire(
else:
_ = bridge(test_text)
- for _, h in handles:
- if h is not None and hasattr(h, "remove"):
- h.remove()
+ for hp in bridge_hook_points:
+ hp.remove_hooks()
# Bucket fired counts per stem.
for stem in hook_stems:
@@ -497,21 +495,20 @@ def benchmark_critical_forward_hooks(
bridge_activations: Dict[str, torch.Tensor] = {}
# Register hooks on bridge
- bridge_handles = []
+ bridge_hook_points: list[HookPoint] = []
for hook_name in critical_hooks:
if hook_name in bridge.hook_dict:
hook_point = bridge.hook_dict[hook_name]
- handle = hook_point.add_hook(make_capture_hook(bridge_activations, hook_name)) # type: ignore[func-returns-value]
- bridge_handles.append(handle)
+ hook_point.add_hook(make_capture_hook(bridge_activations, hook_name))
+ bridge_hook_points.append(hook_point)
# Run bridge forward pass
with torch.no_grad():
_ = bridge(test_text)
# Clean up hooks
- for handle in bridge_handles:
- if handle is not None:
- handle.remove()
+ for hook_point in bridge_hook_points:
+ hook_point.remove_hooks()
if reference_model is None:
# No reference - just verify activations were captured
@@ -526,21 +523,20 @@ def benchmark_critical_forward_hooks(
# Compare with reference model
reference_activations: Dict[str, torch.Tensor] = {}
- reference_handles = []
+ reference_hook_points: list[HookPoint] = []
for hook_name in critical_hooks:
if hook_name in reference_model.hook_dict:
hook_point = reference_model.hook_dict[hook_name]
- handle = hook_point.add_hook(make_capture_hook(reference_activations, hook_name)) # type: ignore[func-returns-value]
- reference_handles.append(handle)
+ hook_point.add_hook(make_capture_hook(reference_activations, hook_name))
+ reference_hook_points.append(hook_point)
# Run reference forward pass
with torch.no_grad():
_ = reference_model(test_text)
# Clean up hooks
- for handle in reference_handles:
- if handle is not None:
- handle.remove()
+ for hook_point in reference_hook_points:
+ hook_point.remove_hooks()
# Compare activations — categorize by presence
bridge_missing = [] # Hooks in reference but not in bridge (BAD)
diff --git a/transformer_lens/benchmarks/hook_structure.py b/transformer_lens/benchmarks/hook_structure.py
index 324b8386f..b71900e07 100644
--- a/transformer_lens/benchmarks/hook_structure.py
+++ b/transformer_lens/benchmarks/hook_structure.py
@@ -15,6 +15,7 @@
make_capture_hook,
make_grad_capture_hook,
)
+from transformer_lens.hook_points import HookPoint
from transformer_lens.model_bridge import TransformerBridge
@@ -52,13 +53,13 @@ def benchmark_forward_hooks_structure(
hook_names = list(bridge.hook_dict.keys())
# Register hooks on bridge and track missing hooks
- bridge_handles = []
+ bridge_hook_points: list[tuple[str, HookPoint]] = []
missing_from_bridge = []
for hook_name in hook_names:
if hook_name in bridge.hook_dict:
hook_point = bridge.hook_dict[hook_name]
- handle = hook_point.add_hook(make_capture_hook(bridge_activations, hook_name)) # type: ignore[func-returns-value]
- bridge_handles.append((hook_name, handle))
+ hook_point.add_hook(make_capture_hook(bridge_activations, hook_name))
+ bridge_hook_points.append((hook_name, hook_point))
else:
missing_from_bridge.append(hook_name)
@@ -70,12 +71,11 @@ def benchmark_forward_hooks_structure(
_ = bridge(test_text)
# Clean up bridge hooks
- for hook_name, handle in bridge_handles:
- if handle is not None:
- handle.remove()
+ for _, hook_point in bridge_hook_points:
+ hook_point.remove_hooks()
# Check for hooks that didn't fire
- registered_hooks = {name for name, _ in bridge_handles}
+ registered_hooks = {name for name, _ in bridge_hook_points}
hooks_that_didnt_fire = registered_hooks - set(bridge_activations.keys())
if reference_model is None:
@@ -100,12 +100,12 @@ def benchmark_forward_hooks_structure(
)
# Register hooks on reference model
- reference_handles = []
+ reference_hook_points: list[HookPoint] = []
for hook_name in hook_names:
if hook_name in reference_model.hook_dict:
hook_point = reference_model.hook_dict[hook_name]
- handle = hook_point.add_hook(make_capture_hook(reference_activations, hook_name)) # type: ignore[func-returns-value]
- reference_handles.append(handle)
+ hook_point.add_hook(make_capture_hook(reference_activations, hook_name))
+ reference_hook_points.append(hook_point)
# Run reference forward pass
with torch.no_grad():
@@ -115,9 +115,8 @@ def benchmark_forward_hooks_structure(
_ = reference_model(test_text)
# Clean up reference hooks
- for handle in reference_handles:
- if handle is not None:
- handle.remove()
+ for hook_point in reference_hook_points:
+ hook_point.remove_hooks()
# CRITICAL CHECK: Bridge must have all hooks that reference has
if missing_from_bridge:
@@ -245,13 +244,13 @@ def benchmark_backward_hooks_structure(
]
# Register backward hooks on bridge
- bridge_handles = []
+ bridge_hook_points: list[tuple[str, HookPoint]] = []
missing_from_bridge = []
for hook_name in grad_hook_names:
if hook_name in bridge.hook_dict:
hook_point = bridge.hook_dict[hook_name]
- handle = hook_point.add_hook(make_grad_capture_hook(bridge_grads, hook_name), dir="bwd") # type: ignore[func-returns-value]
- bridge_handles.append((hook_name, handle))
+ hook_point.add_hook(make_grad_capture_hook(bridge_grads, hook_name), dir="bwd")
+ bridge_hook_points.append((hook_name, hook_point))
else:
missing_from_bridge.append(hook_name)
@@ -265,12 +264,11 @@ def benchmark_backward_hooks_structure(
loss.backward()
# Clean up bridge hooks
- for hook_name, handle in bridge_handles:
- if handle is not None:
- handle.remove()
+ for _, hook_point in bridge_hook_points:
+ hook_point.remove_hooks(dir="bwd")
# Check for hooks that didn't fire
- registered_hooks = {name for name, _ in bridge_handles}
+ registered_hooks = {name for name, _ in bridge_hook_points}
hooks_that_didnt_fire = registered_hooks - set(bridge_grads.keys())
if reference_model is None:
@@ -295,12 +293,12 @@ def benchmark_backward_hooks_structure(
)
# Register backward hooks on reference
- reference_handles = []
+ reference_hook_points: list[HookPoint] = []
for hook_name in grad_hook_names:
if hook_name in reference_model.hook_dict:
hook_point = reference_model.hook_dict[hook_name]
- handle = hook_point.add_hook(make_grad_capture_hook(reference_grads, hook_name), dir="bwd") # type: ignore[func-returns-value]
- reference_handles.append(handle)
+ hook_point.add_hook(make_grad_capture_hook(reference_grads, hook_name), dir="bwd")
+ reference_hook_points.append(hook_point)
# Run reference forward + backward pass
if prepend_bos is not None:
@@ -312,9 +310,8 @@ def benchmark_backward_hooks_structure(
ref_loss.backward()
# Clean up reference hooks
- for handle in reference_handles:
- if handle is not None:
- handle.remove()
+ for hook_point in reference_hook_points:
+ hook_point.remove_hooks(dir="bwd")
# CRITICAL CHECK: Bridge must have all backward hooks that reference has
if missing_from_bridge:
diff --git a/transformer_lens/components/abstract_attention.py b/transformer_lens/components/abstract_attention.py
index f2dd0338b..b7a8cabbd 100644
--- a/transformer_lens/components/abstract_attention.py
+++ b/transformer_lens/components/abstract_attention.py
@@ -175,7 +175,7 @@ def __init__(
self.register_buffer("rotary_sin", sin)
self.register_buffer("rotary_cos", cos)
elif self.cfg.positional_embedding_type == "alibi":
- # ALiBi bias wil be constructed on the first forward pass.
+ # ALiBi bias will be constructed on the first forward pass.
# Note: While computationally efficient, initializing an bias with max n_ctx (16, 1024, 1024) of float32 will occupy ~256MiB of contiguous GPU memory, which may not be optimal for memory usage.
self.alibi = None
diff --git a/transformer_lens/components/grouped_query_attention.py b/transformer_lens/components/grouped_query_attention.py
index d9fbbf8a8..550542960 100644
--- a/transformer_lens/components/grouped_query_attention.py
+++ b/transformer_lens/components/grouped_query_attention.py
@@ -153,7 +153,7 @@ def calculate_attention_scores(
k: Float[torch.Tensor, "batch key_pos kv_head_index d_head"],
) -> Float[torch.Tensor, "batch head_index query_pos key_pos"]:
"""Calculate attention scores from Q and the unexpanded K matrix.
- K will be expaned from [batch, pos, n_key_value_head, d_head] to [batch, pos, n_query_heads, d_head] using torch.repeat_interleave.
+ K will be expanded from [batch, pos, n_key_value_head, d_head] to [batch, pos, n_query_heads, d_head] using torch.repeat_interleave.
Args:
q (Float[torch.Tensor, "batch query_pos head_index d_head"]): The Q tensor.
@@ -172,7 +172,7 @@ def calculate_z_scores(
pattern: Float[torch.Tensor, "batch head_index query_pos key_pos"],
) -> Float[torch.Tensor, "batch query_pos head_index d_head"]:
"""Calculate z scores from the attention pattern and the unexpanded V matrix.
- V will be expaned from [batch, pos, n_key_value_head, d_head] to [batch, pos, n_query_heads, d_head] using torch.repeat_interleave.
+ V will be expanded from [batch, pos, n_key_value_head, d_head] to [batch, pos, n_query_heads, d_head] using torch.repeat_interleave.
Args:
v (Float[torch.Tensor, "batch query_pos head_index d_head"]): The V tensor.
diff --git a/transformer_lens/components/t5_block.py b/transformer_lens/components/t5_block.py
index 88d5467e0..e93b70e7f 100644
--- a/transformer_lens/components/t5_block.py
+++ b/transformer_lens/components/t5_block.py
@@ -16,7 +16,7 @@
class T5Block(nn.Module):
"""
- T5 decoder Block. Uses T5Layernorm, and T5attention insted of usual ones.
+ T5 decoder Block. Uses T5Layernorm, and T5attention instead of usual ones.
Also uses cross attention if is_decoder is True.
"""
diff --git a/transformer_lens/config/transformer_bridge_config.py b/transformer_lens/config/transformer_bridge_config.py
index 44ea7f055..fe8661162 100644
--- a/transformer_lens/config/transformer_bridge_config.py
+++ b/transformer_lens/config/transformer_bridge_config.py
@@ -83,7 +83,6 @@ def __init__(
NTK_by_parts_low_freq_factor: float = 1.0,
NTK_by_parts_high_freq_factor: float = 4.0,
NTK_by_parts_factor: float = 8.0,
- eps_attr: str = "eps",
rmsnorm_uses_offset: bool = False,
attn_implementation: Optional[str] = None,
# Audio model configuration
@@ -176,7 +175,6 @@ def __init__(
self.NTK_by_parts_low_freq_factor = NTK_by_parts_low_freq_factor
self.NTK_by_parts_high_freq_factor = NTK_by_parts_high_freq_factor
self.NTK_by_parts_factor = NTK_by_parts_factor
- self.eps_attr = eps_attr
self.rmsnorm_uses_offset = rmsnorm_uses_offset
self.attn_implementation = attn_implementation
# Audio model configuration
diff --git a/transformer_lens/head_detector.py b/transformer_lens/head_detector.py
index 9efd237ff..e05ad6412 100644
--- a/transformer_lens/head_detector.py
+++ b/transformer_lens/head_detector.py
@@ -28,7 +28,7 @@
SEQ_LEN_ERR = "The sequence must be non-empty and must fit within the model's context window."
-DET_PAT_NOT_SQUARE_ERR = "The detection pattern must be a lower triangular matrix of shape (sequence_length, sequence_length); sequence_length=%d; got detection patern of shape %s"
+DET_PAT_NOT_SQUARE_ERR = "The detection pattern must be a lower triangular matrix of shape (sequence_length, sequence_length); sequence_length=%d; got detection pattern of shape %s"
def detect_head(
@@ -87,7 +87,7 @@ def detect_head(
Currently available heads are: `["previous_token_head", "duplicate_token_head",
"induction_head"]`.
heads: If specific attention heads is given here, all other heads' score is set to -1.
- Useful for IOI-style circuit analysis. Heads can be spacified as a list tuples (layer,
+ Useful for IOI-style circuit analysis. Heads can be specified as a list of tuples (layer,
head) or a dictionary mapping a layer to heads within that layer that we want to
analyze. cache: Include the cache to save time if you want.
exclude_bos: Exclude attention paid to the beginning of sequence token.
diff --git a/transformer_lens/model_bridge/bridge.py b/transformer_lens/model_bridge/bridge.py
index 077007cd5..3badb26f0 100644
--- a/transformer_lens/model_bridge/bridge.py
+++ b/transformer_lens/model_bridge/bridge.py
@@ -2094,7 +2094,7 @@ def cache_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
try:
if hasattr(tensor, "detach"):
cache[name] = tensor.detach().to(cache_device)
- except:
+ except Exception:
pass
return tensor
@@ -2347,6 +2347,66 @@ def wrapped_hook_fn(tensor, hook, _orig_fn=original_hook_fn):
for hook_point, direction in added_hooks:
hook_point.remove_hooks(dir=direction)
+ def _resolve_stopping_criteria(
+ self,
+ stop_strings: Optional[Union[str, List[str]]],
+ stopping_criteria: Optional[Any],
+ ) -> Optional[Any]:
+ """Combine ``stop_strings`` and ``stopping_criteria`` into one StoppingCriteriaList.
+
+ Returns ``None`` when neither is supplied (or both reduce to no-ops),
+ so callers can cheaply check whether any extra stop signal is active.
+ ``stop_strings`` is turned into a HuggingFace ``StopStringCriteria`` (which reproduces
+ HF's exact partial-token-aware, end-anchored matching: it fires when the stop string
+ ends the generated text, even if the string straddles token boundaries) and therefore
+ requires a tokenizer.
+ A user-supplied ``stopping_criteria`` may be a single ``StoppingCriteria``,
+ a list of them, or a ``StoppingCriteriaList``.
+
+ Raises:
+ ValueError: if ``stop_strings`` is supplied without a tokenizer.
+ TypeError: if ``stopping_criteria`` is not a ``StoppingCriteria``, a
+ list/tuple of them, or a ``StoppingCriteriaList``.
+ """
+ if stop_strings is None and stopping_criteria is None:
+ return None
+
+ from transformers import ( # local import: matches the file's transformers usage
+ StoppingCriteria,
+ StoppingCriteriaList,
+ StopStringCriteria,
+ )
+
+ criteria = StoppingCriteriaList()
+
+ if stop_strings is not None:
+ strings = [stop_strings] if isinstance(stop_strings, str) else list(stop_strings)
+ strings = [s for s in strings if s] # drop empty strings (HF errors on them)
+ if strings:
+ if self.tokenizer is None:
+ raise ValueError(
+ "stop_strings requires a tokenizer (stop strings are detected by "
+ "matching against the tokenizer vocabulary), but this TransformerBridge "
+ "has no tokenizer. Pass a stopping_criteria callable that operates on "
+ "token ids instead, or use hf_generate()."
+ )
+ criteria.append(StopStringCriteria(tokenizer=self.tokenizer, stop_strings=strings))
+
+ if stopping_criteria is not None:
+ if isinstance(stopping_criteria, StoppingCriteriaList):
+ criteria.extend(stopping_criteria)
+ elif isinstance(stopping_criteria, (list, tuple)):
+ criteria.extend(stopping_criteria)
+ elif isinstance(stopping_criteria, StoppingCriteria):
+ criteria.append(stopping_criteria)
+ else:
+ raise TypeError(
+ "stopping_criteria must be a transformers.StoppingCriteria, a list of "
+ f"them, or a StoppingCriteriaList, but got {type(stopping_criteria).__name__}."
+ )
+
+ return criteria if len(criteria) > 0 else None
+
def _generate_tokens(
self,
current_tokens: torch.Tensor,
@@ -2377,14 +2437,21 @@ def _generate_tokens(
pixel_values: Optional[torch.Tensor],
multimodal_kwargs: Dict[str, Any],
verbose: bool,
+ stopping_criteria_list: Optional[Any] = None,
) -> Generator[Tuple[torch.Tensor, torch.Tensor, bool], None, None]:
"""Core generation loop. Yields (sampled_tokens, final_logits, all_finished) per step.
- Owns the forward pass, sampling, EOS handling, token accumulation, and
- KV cache management. Callers are responsible for try/finally cleanup of
- ``_capture_hf_cache``.
+ Owns the forward pass, sampling, stop handling (EOS and any
+ ``stopping_criteria_list``), token accumulation, and KV cache management. Callers
+ are responsible for try/finally cleanup of ``_capture_hf_cache``.
+
+ ``stopping_criteria_list`` (from ``_resolve_stopping_criteria``) is evaluated on
+ the running sequence each step and folded into the finished-sequence mask alongside
+ EOS, so when it is ``None`` the loop runs the EOS-only path unchanged.
"""
_hf_kv_cache = None
+ # A row may finish via EOS and/or any of the configured stopping criteria.
+ any_stop_active = stop_at_eos or stopping_criteria_list is not None
for gen_step_idx in tqdm.tqdm(range(max_new_tokens), disable=not verbose):
with torch.no_grad():
@@ -2520,9 +2587,13 @@ def _generate_tokens(
else (decoder_tokens if is_encoder_decoder else current_tokens),
).to(self.cfg.device)
- # Handle EOS
- if stop_at_eos:
+ # Freeze rows that finished on an earlier step so they stop emitting
+ # real tokens. Applies to every active stop mechanism, not just EOS.
+ if any_stop_active:
sampled_tokens[finished_sequences] = eos_token_for_padding
+
+ # Fold this step's EOS matches into the finished mask.
+ if stop_at_eos:
finished_sequences.logical_or_(
torch.isin(
sampled_tokens.to(self.cfg.device),
@@ -2544,7 +2615,26 @@ def _generate_tokens(
else:
current_tokens = torch.cat([current_tokens, sampled_tokens.unsqueeze(1)], dim=1)
- all_finished = bool(stop_at_eos and finished_sequences.all().item())
+ # Fold stop_strings / stopping_criteria into the finished mask. They are
+ # evaluated on the full running sequence (prompt + everything generated so
+ # far, including the token just appended) with this step's logits as the
+ # scores argument, matching transformers' StoppingCriteria contract. The
+ # combined list returns a per-row bool [batch] OR-ing every criterion.
+ # generate()/generate_stream() guarantee this is plain decoder-only token
+ # generation, so current_tokens is the running token sequence.
+ if stopping_criteria_list is not None:
+ criteria_finished = stopping_criteria_list(current_tokens, final_logits).to(
+ device=self.cfg.device, dtype=torch.bool
+ )
+ if criteria_finished.shape != finished_sequences.shape:
+ raise ValueError(
+ "A stopping criterion returned shape "
+ f"{tuple(criteria_finished.shape)}, expected a per-row bool of "
+ f"shape {tuple(finished_sequences.shape)} (one entry per sequence)."
+ )
+ finished_sequences.logical_or_(criteria_finished)
+
+ all_finished = bool(any_stop_active and finished_sequences.all().item())
yield sampled_tokens, final_logits, all_finished
@@ -2573,6 +2663,8 @@ def generate(
names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None,
device: Optional[Union[str, torch.device]] = None,
pixel_values: Optional[torch.Tensor] = None,
+ stop_strings: Optional[Union[str, List[str]]] = None,
+ stopping_criteria: Optional[Any] = None,
**multimodal_kwargs,
) -> (
str | list[str] | torch.Tensor | Any | tuple[Any, ActivationCache]
@@ -2621,6 +2713,22 @@ def generate(
pixel_values: Optional image tensor for multimodal models. Only passed on the
first generation step (the vision encoder processes the image once, then
embeddings are part of the token sequence for subsequent steps).
+ stop_strings: Optional string or list of strings. A sequence stops once its
+ generated text ends with one of these strings, using HuggingFace's
+ StopStringCriteria (partial-token-aware, end-anchored) matching.
+ Requires a tokenizer (raises ValueError otherwise).
+ Independent of stop_at_eos: either can stop a sequence.
+ stopping_criteria: Optional HuggingFace stopping criteria, a single
+ transformers.StoppingCriteria, a list of them, or a StoppingCriteriaList.
+ Each is called as criterion(input_ids, scores) after every step and ORed
+ with the other stop signals, where input_ids is the running sequence and
+ scores is this step's logits ([batch, d_vocab]). Each criterion must return
+ a per-row bool [batch] (or a scalar bool). stop_strings and stopping_criteria
+ are supported only for standard decoder-only text generation. Encoder-decoder,
+ inputs_embeds, and multimodal generation always raise NotImplementedError.
+ Stateful/SSM models raise only when run with use_past_kv_cache=False (the
+ default keeps them on the hooked loop). Each error names the supported
+ alternative.
Returns:
Generated sequence as string, list of strings, or tensor depending on input type and return_type.
@@ -2775,6 +2883,62 @@ def generate(
and pixel_values is None
and not multimodal_kwargs
)
+
+ # stop_strings / stopping_criteria are applied inside the hooked _generate_tokens
+ # loop, so they are supported only on the standard decoder-only text path. Reject
+ # the paths that route around that loop with a clear error rather than silently
+ # dropping the kwargs. This must run before the stateful delegation below.
+ stopping_criteria_list = self._resolve_stopping_criteria(stop_strings, stopping_criteria)
+ if stopping_criteria_list is not None:
+ if is_encoder_decoder:
+ _unsupported = "encoder-decoder models"
+ elif _generate_from_embeds:
+ _unsupported = "inputs_embeds generation"
+ elif pixel_values is not None or multimodal_kwargs:
+ _unsupported = "multimodal (pixel_values) generation"
+ else:
+ _unsupported = None
+ if _unsupported is not None:
+ raise NotImplementedError(
+ f"stop_strings/stopping_criteria are not supported for {_unsupported} in "
+ "TransformerBridge.generate(). Call hf_generate(...), which runs "
+ "HuggingFace's own generation loop and supports HF-native stopping on "
+ "those inputs."
+ )
+ if is_stateful_model and not use_stateful_cache:
+ # Reached only for a stateful/SSM model with use_past_kv_cache=False: the
+ # hooked loop needs the stateful cache, so generate() would otherwise fall
+ # back to hf_generate() and drop these kwargs. The default cache setting
+ # keeps generation on the hooked loop, where stopping is applied.
+ raise NotImplementedError(
+ "stop_strings/stopping_criteria on a stateful/SSM model require the "
+ "stateful cache path, which runs only with use_past_kv_cache=True (the "
+ "default). With use_past_kv_cache=False generate() falls back to "
+ "hf_generate(). Set use_past_kv_cache=True to keep stopping on the hooked "
+ "loop, or call hf_generate(...) directly for HF-native stopping."
+ )
+ # Finished rows are overwritten with this id so they stop emitting real tokens
+ # while the rest of a batch keeps going. stop_at_eos already set a sensible
+ # value, otherwise fall back to the tokenizer pad/eos id. (For a single
+ # sequence this id is never read: the loop exits when the row finishes.)
+ if not stop_at_eos:
+ _pad_id = None
+ if self.tokenizer is not None:
+ _pad_id = (
+ self.tokenizer.pad_token_id
+ if self.tokenizer.pad_token_id is not None
+ else self.tokenizer.eos_token_id
+ )
+ if _pad_id is not None:
+ eos_token_for_padding = _pad_id
+ elif batch_size > 1:
+ raise ValueError(
+ "Batched generation with stopping_criteria and stop_at_eos=False "
+ "needs a padding token to freeze finished rows, but no tokenizer "
+ "pad/eos id is available. Set stop_at_eos=True, use a tokenizer with "
+ "a pad or eos token, or generate one sequence at a time."
+ )
+
if is_stateful_model and not use_stateful_cache:
hf_kwargs: dict[str, Any] = {
"max_new_tokens": max_new_tokens,
@@ -2857,6 +3021,7 @@ def generate(
pixel_values=pixel_values,
multimodal_kwargs=multimodal_kwargs if multimodal_kwargs else {},
verbose=verbose,
+ stopping_criteria_list=stopping_criteria_list,
):
sampled_tokens_list.append(sampled_tokens.unsqueeze(1))
if logits_seq_list is not None:
@@ -2948,6 +3113,8 @@ def generate_stream(
padding_side: Optional[str] = None,
return_type: Optional[str] = "input",
verbose: bool = True,
+ stop_strings: Optional[Union[str, List[str]]] = None,
+ stopping_criteria: Optional[Any] = None,
) -> Generator[Union[torch.Tensor, str], None, None]:
"""Stream tokens from the model as they are generated.
@@ -2972,6 +3139,12 @@ def generate_stream(
is forced internally for batched generation.
return_type: 'input' (match input type), 'str', or 'tokens'.
verbose: Show progress bar.
+ stop_strings: Optional string or list of strings. A sequence stops once its
+ generated text ends with one of them (HF StopStringCriteria). Requires a
+ tokenizer. See generate() for details.
+ stopping_criteria: Optional transformers StoppingCriteria, list, or
+ StoppingCriteriaList, called as criterion(input_ids, scores) each step
+ (scores is the step's logits). See generate() for the full contract.
Yields:
Token tensors [batch, seq_len] or strings, accumulated up to
@@ -3036,6 +3209,28 @@ def generate_stream(
finished_sequences = torch.zeros(batch_size, dtype=torch.bool, device=self.cfg.device)
+ # stop_strings / stopping_criteria: build the combined criteria list (validates
+ # tokenizer for stop_strings). generate_stream only runs the decoder-only text
+ # path, so no path guards are needed here.
+ stopping_criteria_list = self._resolve_stopping_criteria(stop_strings, stopping_criteria)
+ if stopping_criteria_list is not None and not stop_at_eos:
+ _pad_id = None
+ if self.tokenizer is not None:
+ _pad_id = (
+ self.tokenizer.pad_token_id
+ if self.tokenizer.pad_token_id is not None
+ else self.tokenizer.eos_token_id
+ )
+ if _pad_id is not None:
+ eos_token_for_padding = _pad_id
+ elif batch_size > 1:
+ raise ValueError(
+ "Batched generate_stream with stopping_criteria and stop_at_eos=False "
+ "needs a padding token to freeze finished rows, but no tokenizer pad/eos "
+ "id is available. Set stop_at_eos=True or use a tokenizer with a pad/eos "
+ "token."
+ )
+
# --- Cache setup ---
if use_past_kv_cache:
self._capture_hf_cache = True
@@ -3087,6 +3282,7 @@ def _maybe_decode(
pixel_values=None,
multimodal_kwargs={},
verbose=verbose,
+ stopping_criteria_list=stopping_criteria_list,
)
):
new_tokens = sampled_tokens.unsqueeze(-1)
diff --git a/transformer_lens/model_bridge/supported_architectures/AGENTS.md b/transformer_lens/model_bridge/supported_architectures/AGENTS.md
index 8d1bf047d..628e8b7f3 100644
--- a/transformer_lens/model_bridge/supported_architectures/AGENTS.md
+++ b/transformer_lens/model_bridge/supported_architectures/AGENTS.md
@@ -98,7 +98,6 @@ HF raw config attributes are invisible to TL-side consumers unless propagated to
| `query_pre_attn_scalar` | `self.cfg.query_pre_attn_scalar` | Gemma2/3 — query scaling override |
| `sliding_window` | `self.cfg.sliding_window` | Mistral, Qwen2, Gemma2 — local-attention layers |
| `layer_types` | `self.cfg.layer_types` | Hybrid models with per-layer attention type lists |
-| Non-standard RMSNorm eps key | `self.cfg.eps_attr = ""` | Llama uses `"variance_epsilon"` instead of `"eps"` |
**Weight-fold attributes** (need BOTH surface-on-cfg AND fold-into-weight via `preprocess_weights` — see [the next section](#when-to-override-preprocess_weights)):
@@ -238,7 +237,6 @@ Failure message names the missing set. (`INTENTIONAL_EXCLUDES` in the test handl
| RoPE (rotary positional embeddings) | `llama.py`, `mistral.py`, `qwen2.py`+ | `RotaryEmbeddingBridge(name="model.rotary_emb")` + `cfg.positional_embedding_type = "rotary"` |
| GQA / MQA (`n_key_value_heads < n_heads`) | `llama.py`, `mistral.py`, `falcon.py`, `cohere.py` | Set `cfg.n_key_value_heads`; pass `n_kv_heads=` to `_qkvo_weight_conversions()` |
| RMSNorm with offset | `gemma1.py`, `gemma2.py`, `gemma3.py` | `cfg.rmsnorm_uses_offset = True` + `ArithmeticTensorConversion(ADDITION, 1.0)` |
-| Custom RMSNorm eps attribute | `llama.py` | `cfg.eps_attr = "variance_epsilon"` (Llama uses this instead of `eps`) |
| Standard LayerNorm | `gpt2.py`, `bloom.py` | `cfg.normalization_type = "LN"` |
| Gated MLP (`gate_proj`, `up_proj`, `down_proj`) | `llama.py`, `mistral.py`, `gemma1.py`, `qwen2.py`+ | `GatedMLPBridge` with submodules `{gate, in, out}` |
| Combined QKV (`c_attn`) | `gpt2.py`, `bloom.py` | `QKVSplitRearrangeConversion` to split + rearrange |
@@ -324,7 +322,7 @@ class TestMyArchHookCompatibility:
No weight load, no HF Hub access — synthetic cfg + structural assertions only. Runs in default `make unit-test`.
-Add one test per architecture quirk (softcaps, RMSNorm offsets, sliding window, custom `eps_attr`, MoE routing). Gemma1's "must NOT override `setup_hook_compatibility`" is a good one-quirk-one-test example.
+Add one test per architecture quirk (softcaps, RMSNorm offsets, sliding window, MoE routing). Gemma1's "must NOT override `setup_hook_compatibility`" is a good one-quirk-one-test example.
### 2. Integration parity test — `tests/integration/model_bridge/test__adapter.py`
diff --git a/transformer_lens/model_bridge/supported_architectures/baichuan.py b/transformer_lens/model_bridge/supported_architectures/baichuan.py
index a50fabc37..f78063a54 100644
--- a/transformer_lens/model_bridge/supported_architectures/baichuan.py
+++ b/transformer_lens/model_bridge/supported_architectures/baichuan.py
@@ -186,7 +186,6 @@ def __init__(self, cfg: Any) -> None:
self.cfg.gated_mlp = True
self.cfg.attn_only = False
self.cfg.uses_rms_norm = True
- self.cfg.eps_attr = "variance_epsilon"
# Fused W_pack prevents standard fold_ln from reaching Q/K/V separately.
# preprocess_weights() handles it instead.
diff --git a/transformer_lens/model_bridge/supported_architectures/cohere.py b/transformer_lens/model_bridge/supported_architectures/cohere.py
index 5dcb74149..97e8c3301 100644
--- a/transformer_lens/model_bridge/supported_architectures/cohere.py
+++ b/transformer_lens/model_bridge/supported_architectures/cohere.py
@@ -53,10 +53,8 @@ def __init__(self, cfg: Any) -> None:
# --- Normalization ---
# CohereLayerNorm is true LayerNorm (subtracts mean), NOT RMSNorm.
# uses_rms_norm=False tells NormalizationBridge to subtract the mean.
- # eps_attr="variance_epsilon": CohereLayerNorm stores eps as self.variance_epsilon.
self.cfg.normalization_type = "LN"
self.cfg.uses_rms_norm = False
- self.cfg.eps_attr = "variance_epsilon"
self.cfg.final_rms = False
# --- Position embeddings and MLP ---
diff --git a/transformer_lens/model_bridge/supported_architectures/gpt_bigcode.py b/transformer_lens/model_bridge/supported_architectures/gpt_bigcode.py
index 102c4f2c7..2f3d7aa2b 100644
--- a/transformer_lens/model_bridge/supported_architectures/gpt_bigcode.py
+++ b/transformer_lens/model_bridge/supported_architectures/gpt_bigcode.py
@@ -84,7 +84,6 @@ def __init__(self, cfg: Any) -> None:
self.cfg.gated_mlp = False
self.cfg.attn_only = False
self.cfg.uses_rms_norm = False
- self.cfg.eps_attr = "layer_norm_epsilon"
self.cfg.n_key_value_heads = 1 # MQA: always 1 KV head
# Mirror GPT-2 combined-QKV flags
diff --git a/transformer_lens/model_bridge/supported_architectures/gpt_oss.py b/transformer_lens/model_bridge/supported_architectures/gpt_oss.py
index 2e32d277c..2a808b6f7 100644
--- a/transformer_lens/model_bridge/supported_architectures/gpt_oss.py
+++ b/transformer_lens/model_bridge/supported_architectures/gpt_oss.py
@@ -30,8 +30,6 @@ def __init__(self, cfg: Any) -> None:
self.cfg.normalization_type = "RMS"
self.cfg.uses_rms_norm = True
- # GPT-OSS uses 'variance_epsilon' instead of 'eps' for RMSNorm
- self.cfg.eps_attr = "variance_epsilon"
# GPT-OSS uses rotary position embeddings, not learned embeddings
self.cfg.positional_embedding_type = "rotary"
# GPT-OSS attention returns (output, attn_weights), not a 3-tuple
diff --git a/transformer_lens/model_bridge/supported_architectures/granite.py b/transformer_lens/model_bridge/supported_architectures/granite.py
index c46081b0b..37e00598b 100644
--- a/transformer_lens/model_bridge/supported_architectures/granite.py
+++ b/transformer_lens/model_bridge/supported_architectures/granite.py
@@ -52,7 +52,6 @@ def _setup_common_config(self, cfg: Any) -> None:
self.cfg.attn_only = False
self.cfg.uses_rms_norm = True
self.cfg.default_prepend_bos = False
- self.cfg.eps_attr = "variance_epsilon"
self.default_config = {
"d_model": cfg.d_model,
diff --git a/transformer_lens/model_bridge/supported_architectures/internlm2.py b/transformer_lens/model_bridge/supported_architectures/internlm2.py
index a5405e807..80c056a23 100644
--- a/transformer_lens/model_bridge/supported_architectures/internlm2.py
+++ b/transformer_lens/model_bridge/supported_architectures/internlm2.py
@@ -93,7 +93,6 @@ def __init__(self, cfg: Any) -> None:
self.cfg.gated_mlp = True
self.cfg.attn_only = False
self.cfg.uses_rms_norm = True
- self.cfg.eps_attr = "variance_epsilon"
# Standard fold_ln silently skips attention when wqkv is fused (see class docstring).
# preprocess_weights() handles it instead — same approach as phi3.py.
diff --git a/transformer_lens/model_bridge/supported_architectures/llama.py b/transformer_lens/model_bridge/supported_architectures/llama.py
index 5e98ceb37..b7e731a39 100644
--- a/transformer_lens/model_bridge/supported_architectures/llama.py
+++ b/transformer_lens/model_bridge/supported_architectures/llama.py
@@ -63,8 +63,6 @@ def __init__(self, cfg: Any) -> None:
self.cfg.n_key_value_heads = cfg.n_key_value_heads
self.cfg.uses_rms_norm = True
- # Llama uses 'variance_epsilon' instead of 'eps' for RMSNorm
- self.cfg.eps_attr = "variance_epsilon"
self.weight_processing_conversions = {
**self._qkvo_weight_conversions(),
diff --git a/transformer_lens/model_bridge/supported_architectures/llava.py b/transformer_lens/model_bridge/supported_architectures/llava.py
index 407ae72e4..4993ffdb5 100644
--- a/transformer_lens/model_bridge/supported_architectures/llava.py
+++ b/transformer_lens/model_bridge/supported_architectures/llava.py
@@ -62,7 +62,6 @@ def __init__(self, cfg: Any) -> None:
self.cfg.attn_implementation = "eager"
self.cfg.final_rms = True
self.cfg.attn_only = False
- self.cfg.eps_attr = "variance_epsilon"
# GQA support
if hasattr(cfg, "n_key_value_heads") and cfg.n_key_value_heads is not None:
diff --git a/transformer_lens/model_bridge/supported_architectures/phi.py b/transformer_lens/model_bridge/supported_architectures/phi.py
index a9921ffee..851a0f650 100644
--- a/transformer_lens/model_bridge/supported_architectures/phi.py
+++ b/transformer_lens/model_bridge/supported_architectures/phi.py
@@ -119,23 +119,7 @@ def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> No
"""
# Get rotary embedding instance from the model
# Phi models have rotary_emb at model.model.rotary_emb
- if hasattr(hf_model, "model") and hasattr(hf_model.model, "rotary_emb"):
- rotary_emb = hf_model.model.rotary_emb
- else:
- # Fallback: try to get from first layer
- if hasattr(hf_model, "model") and hasattr(hf_model.model, "layers"):
- if len(hf_model.model.layers) > 0:
- first_layer = hf_model.model.layers[0]
- if hasattr(first_layer, "self_attn") and hasattr(
- first_layer.self_attn, "rotary_emb"
- ):
- rotary_emb = first_layer.self_attn.rotary_emb
- else:
- return # Can't find rotary_emb
- else:
- return
- else:
- return
+ rotary_emb = hf_model.model.rotary_emb
# Set rotary_emb on actual bridge instances in bridge_model if available
if bridge_model is not None and hasattr(bridge_model, "blocks"):
diff --git a/transformer_lens/model_bridge/supported_architectures/stablelm.py b/transformer_lens/model_bridge/supported_architectures/stablelm.py
index a09a8114f..4d42069ff 100644
--- a/transformer_lens/model_bridge/supported_architectures/stablelm.py
+++ b/transformer_lens/model_bridge/supported_architectures/stablelm.py
@@ -16,6 +16,7 @@
GatedMLPBridge,
LinearBridge,
NormalizationBridge,
+ ParallelBlockBridge,
PositionEmbeddingsAttentionBridge,
RotaryEmbeddingBridge,
UnembeddingBridge,
@@ -135,10 +136,13 @@ def __init__(self, cfg: Any) -> None:
},
)
+ # StableLM has both parallel (use_parallel_residual=True) and sequential variants.
+ block_cls = ParallelBlockBridge if use_parallel_residual else BlockBridge
+
self.component_mapping = {
"embed": EmbeddingBridge(name="model.embed_tokens"),
"rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"),
- "blocks": BlockBridge(
+ "blocks": block_cls(
name="model.layers",
submodules=block_submodules,
),
diff --git a/transformer_lens/pretrained/weight_conversions/neox.py b/transformer_lens/pretrained/weight_conversions/neox.py
index ff84b5b0d..70d8a6b48 100644
--- a/transformer_lens/pretrained/weight_conversions/neox.py
+++ b/transformer_lens/pretrained/weight_conversions/neox.py
@@ -14,7 +14,7 @@ def convert_neox_weights(neox, cfg: HookedTransformerConfig):
state_dict[f"blocks.{l}.ln1.b"] = neox.gpt_neox.layers[l].input_layernorm.bias
# For some inexplicable reason, NeoX both uses the concatenated QKV
- # matmul of GPT-2 (afaict this has a neglible performance impact) AND
+ # matmul of GPT-2 (afaict this has a negligible performance impact) AND
# has the flattened axis in the DIFFERENT order of (head_index qkv
# d_head) - this took me an hour to debug...
W = neox.gpt_neox.layers[l].attention.query_key_value.weight
diff --git a/transformer_lens/tools/__init__.py b/transformer_lens/tools/__init__.py
index beeef8216..2b78fa84e 100644
--- a/transformer_lens/tools/__init__.py
+++ b/transformer_lens/tools/__init__.py
@@ -4,9 +4,10 @@
including the model registry for discovering compatible HuggingFace models.
Subpackages:
+ - analysis: High-level interpretability analyses (e.g. Direct Logit Attribution)
- model_registry: Tools for discovering and documenting supported models
"""
-from . import model_registry
+from . import analysis, model_registry
-__all__ = ["model_registry"]
+__all__ = ["analysis", "model_registry"]
diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py
new file mode 100644
index 000000000..a9cc5d19e
--- /dev/null
+++ b/transformer_lens/tools/analysis/__init__.py
@@ -0,0 +1,17 @@
+"""Analysis tools for TransformerLens.
+
+This subpackage collects high-level, single-call interpretability analyses that
+sit on top of the hook/cache system. They work with both ``HookedTransformer``
+and the newer ``TransformerBridge`` (the two share the ``ActivationCache`` API).
+
+Tools:
+ - direct_logit_attribution: Direct Logit Attribution (DLA) over components,
+ layers, or attention heads.
+"""
+
+from transformer_lens.tools.analysis.direct_logit_attribution import (
+ DirectLogitAttribution,
+ direct_logit_attribution,
+)
+
+__all__ = ["DirectLogitAttribution", "direct_logit_attribution"]
diff --git a/transformer_lens/tools/analysis/direct_logit_attribution.py b/transformer_lens/tools/analysis/direct_logit_attribution.py
new file mode 100644
index 000000000..96c31d5e9
--- /dev/null
+++ b/transformer_lens/tools/analysis/direct_logit_attribution.py
@@ -0,0 +1,264 @@
+"""Direct Logit Attribution (DLA).
+
+Direct Logit Attribution decomposes a model's output logit (or a logit
+*difference* between a correct and an incorrect token) into the additive
+contributions of upstream components — the embedding, each attention and MLP
+sublayer, or each individual attention head. Because the unembedding is linear
+and the residual stream is a sum of component outputs, the final logit is
+(up to the final LayerNorm) a sum of per-component dot products with the
+unembedding direction of the token of interest. DLA reads off those dot
+products. See the `logit lens
+`_
+and `Interpretability in the Wild `_ for the
+canonical uses.
+
+This module exposes a single entry point, :func:`direct_logit_attribution`,
+that wraps the lower-level ``ActivationCache`` primitives
+(:meth:`~transformer_lens.ActivationCache.ActivationCache.decompose_resid`,
+:meth:`~transformer_lens.ActivationCache.ActivationCache.accumulated_resid`,
+:meth:`~transformer_lens.ActivationCache.ActivationCache.stack_head_results`
+and :meth:`~transformer_lens.ActivationCache.ActivationCache.logit_attrs`) into
+one call. It works unchanged with both ``HookedTransformer`` and
+``TransformerBridge`` because they share the cache API.
+
+Example::
+
+ from transformer_lens import HookedTransformer
+ from transformer_lens.tools.analysis import direct_logit_attribution
+
+ model = HookedTransformer.from_pretrained("gpt2", device="cpu")
+ result = direct_logit_attribution(
+ model,
+ "The Eiffel Tower is in the city of",
+ answer_tokens=" Paris",
+ incorrect_tokens=" London",
+ unit="component",
+ )
+ for label, value in zip(result.labels, result.attribution.squeeze()):
+ print(f"{label:>12}: {value.item():+.3f}")
+"""
+
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import torch
+from jaxtyping import Float
+
+from transformer_lens.ActivationCache import ActivationCache
+from transformer_lens.utilities import SliceInput
+
+# Token-like inputs accepted for the correct/incorrect answers, mirroring
+# ActivationCache.logit_attrs.
+TokenInput = Union[
+ str,
+ int,
+ torch.Tensor,
+]
+
+# Which structural unit the residual stream is decomposed into.
+Unit = str # one of: "component", "layer", "head"
+
+_VALID_UNITS = ("component", "layer", "head")
+
+# Block variants that lack the attn_out + mlp_out structure decompose_resid expects.
+# When TransformerBridge.layer_types() reports any of these we refuse early — the
+# downstream decompose_resid would otherwise raise a confusing KeyError.
+_HYBRID_VARIANT_NAMES = ("mamba", "ssm", "mixer", "linear_attn")
+
+
+@dataclass
+class DirectLogitAttribution:
+ """Result of a :func:`direct_logit_attribution` call.
+
+ Attributes:
+ attribution:
+ Tensor of logit (or logit-difference) attributions with shape
+ ``[component, *batch_and_pos]``. The leading axis is aligned with
+ ``labels``. When ``pos`` selects a single position (the default) the
+ position axis is dropped, leaving ``[component, batch]`` — or
+ ``[component]`` if the cache had its batch dimension removed.
+ labels:
+ Human-readable name for each component, aligned with the leading
+ axis of ``attribution`` (e.g. ``"embed"``, ``"0_attn_out"``,
+ ``"L3H7"``).
+ unit:
+ The decomposition unit used ("component", "layer", or "head").
+ """
+
+ attribution: Float[torch.Tensor, "component *batch_and_pos"]
+ labels: List[str]
+ unit: Unit
+
+ def top(self, k: int = 5) -> List[tuple]:
+ """Return the ``k`` highest-attribution ``(label, value)`` pairs.
+
+ Attribution is reduced to a scalar per component by meaning over any
+ remaining batch/position dimensions, so this is most meaningful when a
+ single position was selected.
+ """
+ flat = self.attribution
+ if flat.ndim > 1:
+ flat = flat.flatten(start_dim=1).mean(dim=-1)
+ values, indices = torch.topk(flat, min(k, flat.shape[0]))
+ return [(self.labels[i], values[j].item()) for j, i in enumerate(indices.tolist())]
+
+
+def _residual_stack_and_labels(
+ cache: ActivationCache,
+ unit: Unit,
+ pos_slice: SliceInput,
+):
+ """Decompose the residual stream into ``unit`` components plus labels.
+
+ LayerNorm is intentionally *not* applied here — ``logit_attrs`` applies the
+ final-layer scaling itself, so applying it twice would double-count.
+ """
+ if unit == "component":
+ # embed (+ pos_embed) and each layer's attn_out / mlp_out.
+ return cache.decompose_resid(apply_ln=False, pos_slice=pos_slice, return_labels=True)
+ if unit == "layer":
+ # Cumulative residual stream after each sublayer — logit-lens style.
+ return cache.accumulated_resid(
+ apply_ln=False, incl_mid=True, pos_slice=pos_slice, return_labels=True
+ )
+ if unit == "head":
+ # Each attention head's contribution, plus the MLP/embedding remainder.
+ return cache.stack_head_results(
+ apply_ln=False, pos_slice=pos_slice, incl_remainder=True, return_labels=True
+ )
+ raise ValueError(f"unit must be one of {_VALID_UNITS}, got {unit!r}")
+
+
+def _validate_bridge_compatibility(model) -> None:
+ """Reject Bridge inputs that DLA can't produce correct numbers for.
+
+ HookedTransformer always has LN folded into W_U, so these checks only fire
+ for TransformerBridge. The compatibility-mode check catches a silent-
+ correctness footgun: without folded LN, the projection direction in
+ ``logit_attrs`` is wrong on a Bridge. The hybrid-arch check catches Mamba/
+ SSM blocks early with a clear error rather than letting ``decompose_resid``
+ raise a confusing KeyError downstream.
+ """
+ # Lazy import — keeps the module importable without dragging in the bridge.
+ from transformer_lens.model_bridge import TransformerBridge
+
+ if not isinstance(model, TransformerBridge):
+ return
+
+ if not getattr(model, "compatibility_mode", False):
+ raise ValueError(
+ "DLA on a TransformerBridge requires compatibility mode so that LayerNorm "
+ "weights are folded into W_U. Call `model.enable_compatibility_mode()` "
+ "after loading the bridge, then re-run DLA."
+ )
+
+ layer_types = model.layer_types()
+ hybrid = [lt for lt in layer_types if any(p in _HYBRID_VARIANT_NAMES for p in lt.split("+"))]
+ if hybrid:
+ raise NotImplementedError(
+ f"DLA does not yet support hybrid architectures (found block types {hybrid}). "
+ f"Only standard attention + MLP transformers (e.g. GPT-2, LLaMA, Pythia) are "
+ f"supported; hybrid support requires extending ActivationCache.decompose_resid."
+ )
+
+
+def direct_logit_attribution(
+ model,
+ input: Union[str, List[str], torch.Tensor, None] = None,
+ answer_tokens: Optional[TokenInput] = None,
+ incorrect_tokens: Optional[TokenInput] = None,
+ *,
+ unit: Unit = "component",
+ pos: SliceInput = -1,
+ cache: Optional[ActivationCache] = None,
+) -> DirectLogitAttribution:
+ """Compute Direct Logit Attribution for a prompt.
+
+ Decomposes the contribution of model components to the logit of
+ ``answer_tokens`` (or, if ``incorrect_tokens`` is given, to the logit
+ *difference* ``answer - incorrect`` along the ``W_U`` direction, which is
+ usually what you want for circuit analysis).
+
+ The model is run once with caching unless a precomputed ``cache`` is passed.
+ Works with both ``HookedTransformer`` and ``TransformerBridge``.
+
+ Note that DLA attributes only the part of a logit that comes from the
+ residual stream through the unembedding direction; the unembedding bias
+ ``b_U`` is a per-token constant that no component produces. So a complete
+ decomposition reconstructs ``logit[token] - b_U[token]`` rather than the raw
+ logit.
+
+ On a ``TransformerBridge``, compatibility mode must be enabled (so the final
+ LayerNorm is folded into ``W_U``) — otherwise the projection direction is
+ wrong and DLA returns silently incorrect numbers. Hybrid architectures
+ (Mamba/SSM/Mixer/LinearAttention) are not yet supported because
+ ``decompose_resid`` only understands the ``attn_out + mlp_out`` block layout;
+ both conditions raise an explicit error at call time.
+
+ Args:
+ model:
+ A ``HookedTransformer`` or ``TransformerBridge`` (the latter with
+ ``enable_compatibility_mode()`` already called).
+ input:
+ Prompt to run — a string, list of strings, or token tensor. Optional
+ only when a precomputed ``cache`` is supplied.
+ answer_tokens:
+ The correct token(s) to attribute, as a string, id, or tensor. A
+ string is converted with ``model.to_single_token``.
+ incorrect_tokens:
+ Optional baseline token(s). When given, attribution is computed for
+ the ``answer - incorrect`` residual direction. Must broadcast to the
+ same shape as ``answer_tokens``.
+ unit:
+ Decomposition granularity:
+
+ - ``"component"`` (default): embedding + each layer's attention and
+ MLP output (via ``decompose_resid``).
+ - ``"layer"``: cumulative residual stream after each sublayer, i.e.
+ logit-lens trajectory (via ``accumulated_resid``).
+ - ``"head"``: each attention head individually, plus a remainder
+ term for everything else (via ``stack_head_results``).
+ pos:
+ Sequence position(s) to attribute. Defaults to ``-1`` (the final
+ token, the usual choice for next-token DLA). Pass ``None`` to keep
+ every position (the result then has a trailing position axis).
+ cache:
+ Optional precomputed ``ActivationCache`` to reuse instead of running
+ the model again.
+
+ Returns:
+ A :class:`DirectLogitAttribution` with ``attribution`` (shape
+ ``[component, *batch_and_pos]``) and aligned ``labels``.
+
+ Raises:
+ ValueError: If ``unit`` is invalid, ``answer_tokens`` is ``None``,
+ neither ``input`` nor ``cache`` is provided, or a
+ ``TransformerBridge`` is passed without compatibility mode enabled.
+ NotImplementedError: If a ``TransformerBridge`` reports a hybrid block
+ layout (Mamba/SSM/Mixer/LinearAttention).
+ """
+ if unit not in _VALID_UNITS:
+ raise ValueError(f"unit must be one of {_VALID_UNITS}, got {unit!r}")
+ if answer_tokens is None:
+ raise ValueError("answer_tokens is required")
+
+ _validate_bridge_compatibility(model)
+
+ if cache is None:
+ if input is None:
+ raise ValueError("provide either `input` to run the model, or a precomputed `cache`")
+ _, cache = model.run_with_cache(input)
+
+ residual_stack, labels = _residual_stack_and_labels(cache, unit, pos)
+
+ # logit_attrs applies the final LayerNorm scaling (with the same pos slice)
+ # and dots each component against the (correct - incorrect) unembed direction.
+ attribution = cache.logit_attrs(
+ residual_stack,
+ tokens=answer_tokens,
+ incorrect_tokens=incorrect_tokens,
+ pos_slice=pos,
+ has_batch_dim=cache.has_batch_dim,
+ )
+
+ return DirectLogitAttribution(attribution=attribution, labels=labels, unit=unit)
diff --git a/transformer_lens/utilities/defaults_utils.py b/transformer_lens/utilities/defaults_utils.py
index 84826e395..7745ba8ac 100644
--- a/transformer_lens/utilities/defaults_utils.py
+++ b/transformer_lens/utilities/defaults_utils.py
@@ -33,7 +33,7 @@ class LocallyOverridenDefaults:
WARNING: This context manager must be used for any function/method that directly accesses
default values which may be overridden by the user using the function/method's arguments,
e.g., `model.cfg.default_prepend_bos` and `model.tokenizer.padding_side` which can be
- overriden by `prepend_bos` and `padding_side` arguments, respectively, in the `to_tokens`.
+ overridden by `prepend_bos` and `padding_side` arguments, respectively, in the `to_tokens`.
"""
def __init__(self, model, **overrides):
diff --git a/transformer_lens/utilities/lm_utils.py b/transformer_lens/utilities/lm_utils.py
index 2c0ffed43..a3d4f7932 100644
--- a/transformer_lens/utilities/lm_utils.py
+++ b/transformer_lens/utilities/lm_utils.py
@@ -1,6 +1,6 @@
"""lm_utils.
-This module contains utility functions related to langauge models
+This module contains utility functions related to language models
"""
from __future__ import annotations
diff --git a/transformer_lens/utilities/logits_utils.py b/transformer_lens/utilities/logits_utils.py
index 2baacc22f..8d93331d4 100644
--- a/transformer_lens/utilities/logits_utils.py
+++ b/transformer_lens/utilities/logits_utils.py
@@ -117,7 +117,7 @@ def sample_logits(
len(tokens.shape) == 2
), "Frequency penalty do not support input in the form of embeddings"
for batch_index in range(final_logits.shape[0]):
- # torch.bincount returns a tensor of length d_vocab, with the number of occurences of each token in the tokens.
+ # torch.bincount returns a tensor of length d_vocab, with the number of occurrences of each token in the tokens.
final_logits[batch_index] = final_logits[
batch_index
] - freq_penalty * torch.bincount(
diff --git a/transformer_lens/utilities/slice.py b/transformer_lens/utilities/slice.py
index f1c4b0245..98cba6404 100644
--- a/transformer_lens/utilities/slice.py
+++ b/transformer_lens/utilities/slice.py
@@ -1,6 +1,6 @@
"""Slice.
-This module contains the functionailty for the Slice object
+This module contains the functionality for the Slice object
"""
from __future__ import annotations
diff --git a/transformer_lens/utilities/tensors.py b/transformer_lens/utilities/tensors.py
index 5e0971b13..f233fc97d 100644
--- a/transformer_lens/utilities/tensors.py
+++ b/transformer_lens/utilities/tensors.py
@@ -121,7 +121,7 @@ def get_offset_position_ids(
"""
Returns the indices of non-padded tokens, offset by the position of the first attended token.
"""
- # shift the position ids so that the id at the the first attended token position becomes zero.
+ # shift the position ids so that the id at the first attended token position becomes zero.
# The position ids of the prepending pad tokens are shifted to -1.
shifted_position_ids = attention_mask.cumsum(dim=1) - 1 # [batch, tokens_length]
diff --git a/uv.lock b/uv.lock
index 69338a25a..3f03b3e29 100644
--- a/uv.lock
+++ b/uv.lock
@@ -6458,7 +6458,6 @@ dependencies = [
{ name = "rich" },
{ name = "sentencepiece" },
{ name = "torch" },
- { name = "torchvision" },
{ name = "tqdm" },
{ name = "transformers" },
{ name = "transformers-stream-generator" },
@@ -6545,7 +6544,6 @@ requires-dist = [
{ name = "rich", specifier = ">=12.6.0" },
{ name = "sentencepiece" },
{ name = "torch", specifier = ">=2.6" },
- { name = "torchvision", specifier = ">=0.22,<0.23" },
{ name = "tqdm", specifier = ">=4.64.1" },
{ name = "transformers", specifier = ">=5.4.0" },
{ name = "transformers-stream-generator", specifier = ">=0.0.5,<0.1" },
@@ -6595,7 +6593,7 @@ jupyter = [
]
multimodal = [
{ name = "timm", specifier = ">=1.0.27" },
- { name = "torchvision", specifier = ">=0.22,<0.23" },
+ { name = "torchvision", specifier = ">=0.22" },
]
quantization = [
{ name = "bitsandbytes", specifier = ">=0.46.1" },