Docs/additional hooked transformer cleanup (#1818)

* Removed remaining references to HookedTransformer

* Fragile MPS test fix

Author: jlarson4

.claude/commands/add-model-support.md

@@ -46,7 +46,7 @@ Each step names the doc to read **when you reach that step** — don't load all
 
 6. **Add the HF repo entry** to [`data/supported_models.json`](../../transformer_lens/tools/model_registry/data/supported_models.json) per [§Adding the HF repo to the registry](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#adding-the-hf-repo-to-the-registry). Ask the user about adding canonical sibling variants from `CANONICAL_AUTHORS_BY_ARCH[<HFArchClass>]`.
 
-7. **Verify** end-to-end: `/verify-model $ARGUMENTS`. Read both `status` AND per-phase scores. `STATUS_VERIFIED` means hard gates passed (see [§Phase-score thresholds](../../transformer_lens/tools/model_registry/AGENTS.md#phase-score-thresholds)) — but P4's 50% bar is intentionally lenient. P4 well below 100% on a small parity-test model + `status==1` → suspect missing `preprocess_weights` fold or wrong `default_prepend_bos`; investigate before step 8.
+7. **Verify** end-to-end: `/verify-model $ARGUMENTS`. Read both `status` AND per-phase scores. `STATUS_VERIFIED` means hard gates passed (see [§Phase-score thresholds](../../transformer_lens/tools/model_registry/AGENTS.md#phase-score-thresholds)) — but P4's ~54.5% bar (`p4_pass_threshold()`, not a fixed number) is intentionally lenient. P4 well below 100% on a small parity-test model + `status==1` → suspect missing `preprocess_weights` fold or wrong `default_prepend_bos`; investigate before step 8.
 
 8. **Write tests** per [§Required tests](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#required-tests) (unit + integration). Copy the closest sibling.
 

.claude/commands/verify-model.md

@@ -37,7 +37,7 @@ Full reference: [tools/model_registry/AGENTS.md §Flag reference](../../transfor
 - `--max-memory <gb>` — skip if param estimate exceeds; e.g. `16` on a 24 GB GPU leaves headroom for activations
 - `--phases 1 2 3` — restrict (P4 is slowest; restrict when debugging P1 forward parity)
 - `--dry-run` — see above; always first
-- `--no-hf-reference` / `--no-ht-reference` — skip HF / HT comparison (faster, lower confidence)
+- `--no-hf-reference` — skip the HF reference comparison (faster, lower confidence; the run can only reach `STATUS_PROVISIONAL`, never verified)
 - `--reverify` — re-test `status==1`
 - `--retry-failed` — re-test `status==3` (read existing `note` first)
 
@@ -52,9 +52,10 @@ Hard thresholds (`_MIN_PHASE_SCORES` in `verify_models.py`):
 | 1 | 100% | — | `STATUS_FAILED` |
 | 2 | 75% | `logits_equivalence`, `loss_equivalence` | `STATUS_FAILED` |
 | 3 | 75% | `logits_equivalence`, `loss_equivalence` | `STATUS_FAILED` |
-| 4 | 50% | — | **Non-gating** — adds `"low text quality"` to `note`; never fails. |
+| 4 | ~54.5% (`p4_pass_threshold()`, derived from the judge bake-off noise floor — not a fixed number) | — | **Non-gating** — adds `"low text quality"` to `note`; never fails. |
 | 7 | 75% | `multimodal_forward` | `STATUS_FAILED`. NULL = fail. |
-| 8 | 75% | `audio_forward` | `STATUS_FAILED`. NULL = fail. |
+| 8 | 75% | `audio_forward`, `audio_text_forward` | `STATUS_FAILED`. NULL = fail. |
+| 9 | 75% | `vision_forward`, `vision_cache` | `STATUS_FAILED`. NULL = fail. |
 
 `STATUS_VERIFIED` means hard gates passed. `note` carries quality flags or failure details.
 

AGENTS.md

@@ -58,7 +58,7 @@ make unit-test          # fast, no model loads
 make integration-test   # cross-component
 make acceptance-test    # end-to-end
 make docstring-test     # doctest + doctest-plus
-make notebook-test      # slow; subset run in CI
+make notebook-test      # slow, local-only; CI runs its own per-notebook matrix (notebook-checks in checks.yml), not this target
 make test-pr            # unit + docstring + acceptance + integration (PR-review surface)
 make test               # everything (long; includes benchmarks + notebooks)
 
@@ -213,9 +213,9 @@ Load-bearing pins live in [pyproject.toml](pyproject.toml):
 
 | Pin | Where | Why it matters |
 |---|---|---|
-| `transformers>=5.4.0` | `[project] dependencies` | The Bridge adapter contract is written against HF module layouts; every minor HF release can break adapter component-mappings. Bumping is a real test pass. |
+| `transformers>=5.9.0` | `[project] dependencies` | The Bridge adapter contract is written against HF module layouts; every minor HF release can break adapter component-mappings. Bumping is a real test pass. |
 | `torch>=2.6` | `[project] dependencies` | Hook system relies on PyTorch's forward / backward hook semantics; major torch bumps occasionally change ordering. |
-| `accelerate>=0.23.0` | `[project] dependencies` | Required for Llama-family loading. |
+| `accelerate>=1.1.0` | `[project] dependencies` | Required for Llama-family loading and `device_map` disk offload (`align_module_device`). |
 | `numpy>=1.24` / `>=1.26` | `[project] dependencies` (python-version-conditional) | Doctest float formatting can drift across NumPy versions. |
 | `isort==5.8.0` | `[dependency-groups] dev` (exact) | Format check pins to exactly this version; a bump flips the formatting of every file. |
 

CLAUDE.md

@@ -23,7 +23,7 @@
 
 - [AGENTS.md §10](AGENTS.md#10-hard-rules) — hard rules; load-bearing.
 - [AGENTS.md §2](AGENTS.md#2-two-systems-live-in-this-repo) — the single model system (the `Hooked*` model classes were removed in 4.0).
-- [tests/QUARANTINES.md](tests/QUARANTINES.md) — check before debugging any failing test. The macOS-arm64 KV-cache skip is the most common time-sink.
+- [tests/QUARANTINES.md](tests/QUARANTINES.md) — check before debugging any failing test. A skip with a matching reason is not your bug; a failure not on that list is real.
 - [debugging_numerical_divergence.md](docs/source/content/debugging_numerical_divergence.md) — Bridge-vs-HF logit drift bisection.
 - [compatibility_mode.md](docs/source/content/compatibility_mode.md) — `bridge.enable_compatibility_mode()` contract; read before adding tests that use it.
 - [devtools/adapter_builder/](devtools/adapter_builder/README.md) — autonomous adapter builder (contributor tooling; agent-teams mode needs Max, solo mode runs on any tier); manual path is `/add-model-support`.

README.md

@@ -177,14 +177,14 @@ cover:
   `dt_proj`, `out_proj` for Mamba-1; `in_proj`, `conv1d`, `inner_norm`, `out_proj` for
   Mamba-2)
 * Stateful generation with cache-aware decode steps
-* The `compute_effective_attention` utility (in
-  `transformer_lens.model_bridge.supported_architectures.mamba2`) that materializes
-  Mamba-2's SSD-derived attention matrix for comparison with transformer attention
-  patterns
+* `ActivationCache.compute_ssm_effective_attention(layer=...)`, which materializes the
+  SSD-derived attention matrix of a Mamba-2 (or any other SSM / hybrid) layer for
+  comparison with transformer attention patterns; the older
+  `mamba2.compute_effective_attention` remains only as a deprecated alias
 
 Verification lives in the integration tests at
 `tests/integration/model_bridge/test_mamba_adapter.py` and
-`tests/integration/model_bridge/test_mamba2_adapter.py` (81 tests total), and the
+`tests/integration/model_bridge/test_mamba2_adapter.py`, and the
 `verify_models` benchmark suite now covers the SSM and hybrid families. Mamba-1,
 Mamba-2, gated-delta-net (Qwen3.5 / Qwen3-Next), NemotronH, and GraniteMoeHybrid
 all declare `applicable_phases = [1, 2, 3, 4]`, so their forward parity (P1, vs raw

demos/Exploratory_Analysis_Demo.ipynb

@@ -1936,7 +1936,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The above suggests that it would be a useful bit of infrastructure to have a \"wiki\" for the heads of a model, giving their scores according to some metrics re head functions, like the ones we've seen here. TransformerLens makes this easy to make, as just changing the name input to `HookedTransformer.from_pretrained` gives a different model but in the same architecture, so the same code should work. If you want to make this, I'd love to see it! \n",
+    "The above suggests that it would be a useful bit of infrastructure to have a \"wiki\" for the heads of a model, giving their scores according to some metrics re head functions, like the ones we've seen here. TransformerLens makes this easy to make, as just changing the name input to `TransformerBridge.boot_transformers` gives a different model but in the same architecture, so the same code should work. If you want to make this, I'd love to see it! \n",
     "\n",
     "As a proof of concept, [I made a mosaic of all induction heads across the 40 models then in TransformerLens](https://www.neelnanda.io/mosaic).\n",
     "\n",
@@ -2094,7 +2094,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "transformer-lens",
+   "display_name": "transformer-lens (3.12.12)",
    "language": "python",
    "name": "python3"
   },

demos/Inspect_Bridge_Demo.ipynb

@@ -8,7 +8,7 @@
     "# Inspect Driver — Demo\n",
     "\n",
     "Turn a model served through [`inspect_ai`](https://inspect.aisi.org.uk) into a\n",
-    "TransformerLens `HookedTransformer`: `run_with_cache`, named hook points, and\n",
+    "TransformerLens `TransformerBridge`: `run_with_cache`, named hook points, and\n",
     "interventions all work over the Inspect boundary.\n",
     "\n",
     "`boot_inspect` boots the model behind an `inspect_ai` provider and wraps it in a\n",
@@ -679,7 +679,15 @@
     }
    },
    "outputs": [],
-   "source": "# (9) Tool-aware generation: for a tool-capable instruct model the provider renders tool\n# schemas into the chat template and parses tool calls from the output (so Inspect agent\n# loops run). gpt2 has no tool template, so here we just show the tool-call parser:\nfrom transformer_lens.model_bridge.sources.inspect.transformers_provider import _parse_tool_calls\n\ncalls = _parse_tool_calls('<tool_call>{\"name\": \"add\", \"arguments\": {\"a\": 2, \"b\": 2}}</tool_call>')\nprint(\"parsed tool calls:\", [(c.function, c.arguments) for c in calls])"
+   "source": [
+    "# (9) Tool-aware generation: for a tool-capable instruct model the provider renders tool\n",
+    "# schemas into the chat template and parses tool calls from the output (so Inspect agent\n",
+    "# loops run). gpt2 has no tool template, so here we just show the tool-call parser:\n",
+    "from transformer_lens.model_bridge.sources.inspect.transformers_provider import _parse_tool_calls\n",
+    "\n",
+    "calls = _parse_tool_calls('<tool_call>{\"name\": \"add\", \"arguments\": {\"a\": 2, \"b\": 2}}</tool_call>')\n",
+    "print(\"parsed tool calls:\", [(c.function, c.arguments) for c in calls])"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -706,7 +714,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "transformer-lens",
+   "display_name": "transformer-lens (3.12.12)",
    "language": "python",
    "name": "python3"
   },
@@ -3112,4 +3120,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

demos/LIT_Integration_Demo.ipynb

@@ -122,8 +122,8 @@
     "\n",
     "# LIT integration imports\n",
     "from transformer_lens.lit import (\n",
-    "    HookedTransformerLIT,\n",
-    "    HookedTransformerLITConfig,\n",
+    "    TransformerLensLIT,\n",
+    "    TransformerLensLITConfig,\n",
     "    SimpleTextDataset,\n",
     "    PromptCompletionDataset,\n",
     "    IOIDataset,\n",
@@ -155,36 +155,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "18cbabbd",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Loading gpt2-small...\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "`torch_dtype` is deprecated! Use `dtype` instead!\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Loaded pretrained model gpt2-small into HookedTransformer\n",
-      "Loaded model: gpt2\n",
-      "  Layers: 12\n",
-      "  Heads: 12\n",
-      "  d_model: 768\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# Load GPT-2 (124M parameters)\n",
     "# Other options: \"gpt2-medium\", \"gpt2-large\", \"gpt2-xl\", \"EleutherAI/pythia-70m\", etc.\n",
@@ -231,7 +205,7 @@
    ],
    "source": [
     "# Configure the wrapper\n",
-    "config = HookedTransformerLITConfig(\n",
+    "config = TransformerLensLITConfig(\n",
     "    max_seq_length=256,         # Maximum input length\n",
     "    batch_size=4,               # Batch size for inference\n",
     "    top_k=10,                   # Number of top predictions to show\n",
@@ -243,7 +217,7 @@
     ")\n",
     "\n",
     "# Create the wrapper\n",
-    "lit_model = HookedTransformerLIT(model, config=config)\n",
+    "lit_model = TransformerLensLIT(model, config=config)\n",
     "\n",
     "print(f\"Created LIT wrapper: {lit_model.description()}\")\n",
     "print(f\"\\nInput spec keys: {list(lit_model.input_spec().keys())}\")\n",

demos/Main_Demo.ipynb

@@ -995,7 +995,7 @@
     "\n",
     "LayerNorm is a normalization technique used by transformers, analogous to BatchNorm but more friendly to massive parallelisation. No one *really* knows why it works, but it seems to improve model numerical stability. Unlike BatchNorm, LayerNorm actually changes the functional form of the model, which makes it a massive pain for interpretability! \n",
     "\n",
-    "Folding LayerNorm is a technique to make it lower overhead to deal with, and the flags `center_writing_weights` and `fold_ln` in `HookedTransformer.from_pretrained` apply this automatically (they default to True). These simplify the internal structure without changing the weights.\n",
+    "Folding LayerNorm is a technique to make it lower overhead to deal with, and the flags `fold_ln` and `center_writing_weights` in `model.enable_compatibility_mode()` apply this automatically (they default to True). These simplify the internal structure without changing the weights.\n",
     "\n",
     "Intuitively, LayerNorm acts on each residual stream vector (ie for each batch element and token position) independently, sets their mean to 0 (centering) and standard deviation to 1 (normalizing) (*across* the residual stream dimension - very weird!), and then applies a learned elementwise scaling and translation to each vector.\n",
     "\n",
@@ -1807,7 +1807,7 @@
    "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",
+    "**Note:** `TransformerBridge.boot_transformers` mainly works with HuggingFace revisions, and is only compatible with checkpoints for a few model families. Legacy TransformerLens-format repos (`NeelNanda/*`, `ArthurConmy/*`, `Baidicoot/*`) have no HF `model_type`, so `boot_transformers` cannot load them; use `TransformerBridge.boot_tl_legacy(model_name, checkpoint_index=...)` (or `checkpoint_value=...`) for those, as the code below does.\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",

devtools/adapter_builder/CLAUDE.md

@@ -107,8 +107,8 @@ Nothing is copied into the TransformerLens worktree except `.claude/agents/` (re
 - `.env` — see `.env.example` for all variables; gitignored, never commit it
 - `HF_TOKEN` — HuggingFace API token (falls back to the repo root `.env` if unset here)
 - `DEFAULT_TARGET_REPO` — optional override; defaults to the containing repo
-- `DEFAULT_BASE_BRANCH` — default branch (dev-4.x)
-- `DEFAULT_MAX_MEMORY_GB` — memory limit for verification (96)
+- `DEFAULT_BASE_BRANCH` — base branch for worktrees (`dev` in `.env.example`; the launch scripts fall back to the stale, fully-merged `dev-4.x` if it is unset, so set it)
+- `DEFAULT_MAX_MEMORY_GB` — memory limit for verification (`48` in `.env.example`; script fallback `96`)
 - `WORKTREE_BASE` — where agent pair worktrees live; optional, defaults to `<parent-of-TransformerLens>/worktrees`
 - `NOTIFICATION_WEBHOOK_URL` — Slack webhook for notifications
 - `NOTIFICATION_NUMBER` — iMessage fallback