diff --git a/demos/vLLM_Bridge_Integration_Test.ipynb b/demos/vLLM_Bridge_Integration_Test.ipynb index 031b5c494..ade7aa1e2 100644 --- a/demos/vLLM_Bridge_Integration_Test.ipynb +++ b/demos/vLLM_Bridge_Integration_Test.ipynb @@ -42,7 +42,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": "# Install vllm and TransformerLens @ feature/vllm-batched. ~3-5 minutes.\n# Branch must match this notebook: feature/vllm-batched has the hook-fire counter\n# (Step 7) and the batched capture surface (Step 10). feature/driver-system lacks\n# both, so collective_rpc raises NotImplementedError on tl_reset_counter / the\n# batched RPCs.\n# vllm pinned to 0.20.2 — the version the internal-API walks in\n# transformer_lens/model_bridge/sources/vllm/{plugin,internals}.py were validated against.\n# vLLM rearranges its internal class paths every 4-6 weeks; re-validate before bumping.\n%pip install -q \"vllm==0.20.2\"\n%pip install -q git+https://github.com/TransformerLensOrg/TransformerLens.git@feature/vllm-batched", + "source": "# Install vllm and TransformerLens @ feature/vllm-batched. ~3-5 minutes.\n# Branch must match this notebook: feature/vllm-batched has the hook-fire counter\n# (Step 7) and the batched capture surface (Step 10). feature/driver-system lacks\n# both, so collective_rpc raises NotImplementedError on tl_reset_counter / the\n# batched RPCs.\n#\n# vLLM pinned to 0.20.2 — the version the internal-API walks in\n# transformer_lens/model_bridge/sources/vllm/{plugin,internals}.py were validated\n# against. Installed from the CUDA-12.9 *release wheel*, NOT PyPI's default\n# `vllm==0.20.2`: the PyPI wheel is a CUDA-13 build that fails on Colab's CUDA-12.8\n# image with `libcudart.so.13: cannot open shared object file` and drags in\n# cuda-toolkit 13, which conflicts with Colab's cu12 RAPIDS stack. The cu129 wheel\n# uses libcudart.so.12 (present) and stays in the cu12 family. vLLM has no cu128\n# wheel for 0.20.2; cu129 is CUDA-12.x-compatible with the cu128 runtime.\n# Re-validate the internals walks before bumping vLLM (it moves them every 4-6 weeks).\n%pip install -q \"https://github.com/vllm-project/vllm/releases/download/v0.20.2/vllm-0.20.2%2Bcu129-cp38-abi3-manylinux_2_31_x86_64.whl\"\n%pip install -q git+https://github.com/TransformerLensOrg/TransformerLens.git@feature/vllm-batched", "id": "4ab1eb60e6b1" }, { @@ -50,7 +50,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": "import gc\nimport os\nimport sys\n\nimport torch\n\n# HF_TOKEN comes from Colab secrets. Falls back to env var for non-Colab runs.\ntry:\n from google.colab import userdata\n os.environ.setdefault(\"HF_TOKEN\", userdata.get(\"HF_TOKEN\"))\nexcept (ImportError, Exception):\n pass\nassert os.environ.get(\"HF_TOKEN\"), \"Set HF_TOKEN in Colab Secrets (gear icon, left sidebar).\"\n\n# Colab/Jupyter compatibility: ipykernel's stdout doesn't expose a fileno();\n# vLLM's worker init calls sys.stdout.fileno() during parallel-state setup\n# and crashes with UnsupportedOperation: fileno. Patch fileno to return the\n# underlying process FDs (1, 2) — Colab writes back to those anyway.\nif \"ipykernel\" in sys.modules:\n sys.stdout.fileno = lambda: 1 # type: ignore[method-assign]\n sys.stderr.fileno = lambda: 2 # type: ignore[method-assign]\n\n# Read the installed vllm version from package metadata, NOT `import vllm` —\n# importing vllm loads its CUDA C extension (vllm._C), which is exactly what\n# fails with `libcudart.so.NN not found` when a newer wheel built for a CUDA\n# version Colab doesn't ship gets installed. Metadata read works regardless.\nfrom importlib.metadata import PackageNotFoundError\nfrom importlib.metadata import version as _pkg_version\n\n_PINNED_VLLM = \"0.20.2\"\ntry:\n _vllm_ver = _pkg_version(\"vllm\")\n print(f\"vllm version: {_vllm_ver}\")\n if _vllm_ver != _PINNED_VLLM:\n print(\n f\"⚠ expected vllm=={_PINNED_VLLM} (the version the capture plugin is validated \"\n f\"against); got {_vllm_ver}. Newer wheels target a CUDA the Colab image may not \"\n \"ship (→ libcudart.so error at boot) and may move vLLM internals the plugin walks. \"\n f\"Re-pin with %pip install -q 'vllm=={_PINNED_VLLM}' and restart the runtime.\"\n )\nexcept PackageNotFoundError:\n print(\"⚠ vllm is not installed — run the install cell above.\")\n\nMODEL = \"meta-llama/Llama-3.2-1B\"\nPROMPT = \"The quick brown fox jumps over the\"\nDTYPE = torch.float16\ntorch.manual_seed(0)\n\nprint(f\"CUDA available: {torch.cuda.is_available()}\")\nprint(f\"Device: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU only — abort'}\")\nassert torch.cuda.is_available(), \"GPU runtime required.\"", + "source": "import gc\nimport os\nimport sys\n\nimport torch\n\n# HF_TOKEN comes from Colab secrets. Falls back to env var for non-Colab runs.\ntry:\n from google.colab import userdata\n os.environ.setdefault(\"HF_TOKEN\", userdata.get(\"HF_TOKEN\"))\nexcept (ImportError, Exception):\n pass\nassert os.environ.get(\"HF_TOKEN\"), \"Set HF_TOKEN in Colab Secrets (gear icon, left sidebar).\"\n\n# Colab/Jupyter compatibility: ipykernel's stdout doesn't expose a fileno();\n# vLLM's worker init calls sys.stdout.fileno() during parallel-state setup\n# and crashes with UnsupportedOperation: fileno. Patch fileno to return the\n# underlying process FDs (1, 2) — Colab writes back to those anyway.\nif \"ipykernel\" in sys.modules:\n sys.stdout.fileno = lambda: 1 # type: ignore[method-assign]\n sys.stderr.fileno = lambda: 2 # type: ignore[method-assign]\n\n# Read the installed vllm version from package metadata, NOT `import vllm` —\n# importing vllm loads its CUDA C extension (vllm._C), which is exactly what\n# fails with `libcudart.so.NN not found` when a wheel built for a CUDA version\n# Colab doesn't ship gets installed. Metadata read works regardless.\nfrom importlib.metadata import PackageNotFoundError\nfrom importlib.metadata import version as _pkg_version\n\n_PINNED_VLLM = \"0.20.2\"\ntry:\n _vllm_ver = _pkg_version(\"vllm\")\n print(f\"vllm version: {_vllm_ver}\")\n # The install cell pins the cu129 release wheel, so the local tag is\n # \"0.20.2+cu129\" — compare the base version only.\n if _vllm_ver.split(\"+\")[0] != _PINNED_VLLM:\n print(\n f\"⚠ expected vllm base {_PINNED_VLLM} (the version the capture plugin is validated \"\n f\"against); got {_vllm_ver}. A mismatched wheel may target a CUDA the Colab image \"\n \"doesn't ship (→ libcudart.so error at boot) and may move vLLM internals the plugin \"\n \"walks. Re-run the install cell above (it fetches the cu129 release wheel) and \"\n \"restart the runtime. Do NOT `pip install vllm==0.20.2` — PyPI's default is the \"\n \"CUDA-13 build that fails on Colab's CUDA-12.8 image.\"\n )\n elif \"+cu\" in _vllm_ver and \"cu129\" not in _vllm_ver:\n print(\n f\"⚠ got {_vllm_ver}: not the cu129 wheel. PyPI's default vllm==0.20.2 is a CUDA-13 \"\n \"build (libcudart.so.13 missing on Colab's CUDA-12.8). Re-run the install cell above.\"\n )\nexcept PackageNotFoundError:\n print(\"⚠ vllm is not installed — run the install cell above.\")\n\nMODEL = \"meta-llama/Llama-3.2-1B\"\nPROMPT = \"The quick brown fox jumps over the\"\nDTYPE = torch.float16\ntorch.manual_seed(0)\n\nprint(f\"CUDA available: {torch.cuda.is_available()}\")\nprint(f\"Device: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU only — abort'}\")\nassert torch.cuda.is_available(), \"GPU runtime required.\"", "id": "1748b73ca8b0" }, { @@ -121,23 +121,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": [ - "from transformer_lens.model_bridge.remote_bridge import RemoteBridge\n", - "\n", - "# max_model_len=2048 caps the KV cache reservation. Llama-3.2-1B's native\n", - "# context is 131072 (128k) — the default reservation is ~4 GiB and overshoots\n", - "# the free T4 budget. The test prompt is ~10 tokens, so 2048 is plenty.\n", - "bridge = RemoteBridge.boot_vllm(\n", - " MODEL,\n", - " dtype=DTYPE,\n", - " gpu_memory_utilization=0.5,\n", - " max_model_len=2048,\n", - ")\n", - "assert isinstance(bridge, RemoteBridge), f\"Expected RemoteBridge, got {type(bridge).__name__}\"\n", - "print(f\"Architecture: {bridge.cfg.architecture}\")\n", - "print(f\"Fireable hooks: {len(bridge._driver.supported_hook_points)}\")\n", - "print(f\"Non-fireable hooks (fused kernels): {len(bridge._driver.non_fireable_hook_points)}\")" - ], + "source": "from transformer_lens.model_bridge.remote_bridge import RemoteBridge\n\n# max_model_len=2048 caps the KV cache reservation. Llama-3.2-1B's native\n# context is 131072 (128k) — the default reservation is ~4 GiB and overshoots\n# the free T4 budget. The test prompt is ~10 tokens, so 2048 is plenty.\n#\n# enable_position_interventions=True widens each hook's affine scale/bias buffers\n# from (width,) to (max_num_batched_tokens, width) so Step 6½ can patch specific\n# sequence positions via a spec 'pos' field. Costs ~2× the affine buffers' resident\n# memory (~0.8 GB here); the buffers stay identity until an intervention writes them,\n# so every other step is unchanged. Drop the flag (and skip Step 6½) to run leaner.\nbridge = RemoteBridge.boot_vllm(\n MODEL,\n dtype=DTYPE,\n gpu_memory_utilization=0.5,\n max_model_len=2048,\n enable_position_interventions=True,\n)\nassert isinstance(bridge, RemoteBridge), f\"Expected RemoteBridge, got {type(bridge).__name__}\"\nprint(f\"Architecture: {bridge.cfg.architecture}\")\nprint(f\"Fireable hooks: {len(bridge._driver.supported_hook_points)}\")\nprint(f\"Non-fireable hooks (fused kernels): {len(bridge._driver.non_fireable_hook_points)}\")", "id": "49d50957a4c7" }, { @@ -208,18 +192,18 @@ "execution_count": null, "outputs": [] }, + { + "cell_type": "code", + "id": "85900fee", + "source": "# ============================================================================\n# VERIFY — the driver returns real full-sequence logits (branch: review/dev-4.x)\n# ============================================================================\n# With VLLMDriver._reconstruct_logits + provides_sequence_logits=True, run_with_cache's\n# logits should be (1, seq, d_vocab) real logits at EVERY position (old: final-only\n# log-probs, -inf elsewhere). Reuses v_lnf / hf_lnf / W / norm_w from the Step 4½ cell.\nlg = logits_vllm.float()\nseq = tokens.shape[1]\nprint(f\"driver logits shape: {tuple(lg.shape)} (expect (1, {seq}, {bridge.cfg.d_vocab}))\")\nfinite_pos = torch.isfinite(lg[0]).all(-1)\nprint(f\"all-finite positions: {int(finite_pos.sum())}/{lg.shape[1]} (old behavior: 1, final only)\")\nassert lg.shape == (1, seq, bridge.cfg.d_vocab), \"driver logits are not full-sequence\"\nassert finite_pos.all(), \"some positions non-finite — reconstruction didn't fill them\"\n\ndrv = lg[0, :, : W.shape[0]] # trim any vocab padding to compare against W's real vocab\n\n# (a) driver output == the inline formula (Step 4½) → the driver wired the reconstruction right.\nd_inline = (drv - (v_lnf @ W.T)).abs().max().item()\nprint(f\"driver vs inline (ln_final @ W_U): max|Δ|={d_inline:.2e} (~0 => correct formula in driver)\")\n\n# (b) driver logits vs HF's real logits (HF post-weight = HF_ln_final * norm_w) @ W_U, all positions.\nhf_real = (hf_lnf * norm_w) @ W.T\nper_pos = (drv - hf_real).abs().amax(-1)\nagree = (drv.argmax(-1) == hf_real.argmax(-1)).float().mean().item() * 100\nprint(f\"driver vs HF real logits: max|Δ|={per_pos.max():.2e} (fp16) argmax agreement={agree:.1f}%\")\nprint(\"\\n>>> full-sequence + inline~0 + argmax 100% => VLLMDriver._reconstruct_logits is correct end-to-end.\")\n", + "metadata": {}, + "execution_count": null, + "outputs": [] + }, { "cell_type": "markdown", "metadata": {}, - "source": [ - "## Step 5 — Per-hook L2 (acceptance gate)\n", - "\n", - "Target is relative L2 < 5e-3 in fp16 for every fireable hook. One hook remains exempted from the strict gate:\n", - "\n", - "- **`ln_final.hook_normalized`** — vLLM's `model.norm` is invoked as part of the fused-residual norm kernel and the captured value scales ~2× HF's. Open investigation (likely a residual-fusion semantic discrepancy at the model boundary rather than a real divergence; argmax + downstream parity work correctly).\n", - "\n", - "All other fireable hooks (including `blocks.{i}.hook_out`, which is now materialized from vLLM's `(mlp_delta, residual)` tuple) must be within target." - ], + "source": "## Step 5 — Per-hook L2 (acceptance gate)\n\nTarget is relative L2 < 5e-3 in fp16 for every fireable hook. One hook remains exempted from the strict gate:\n\n- **`ln_final.hook_normalized`** — vLLM exposes the **post-weight** RMSNorm value (`normed × norm_weight`), while HF/HookedTransformer exposes the **pre-weight** value, so vLLM's capture reads larger by exactly the final-norm weight. **Resolved** by Step 4½: the vLLM/HF ratio equals `norm_w` (median ≈ 2.45 for Llama-3.2-1B — the earlier \"~2×\" was just this weight's magnitude, not a fixed factor). This is the documented post/pre-weight convention gap (see `sources.vllm.overlays.decoder_only`), **not** a divergence — and it's exactly why the Step 4½ logit reconstruction (`ln_final @ lm_head.weightᵀ`) works directly off the raw capture. Still exempted from the strict L2 gate below because the *values* legitimately differ from HF's pre-weight hook.\n\nAll other fireable hooks (including `blocks.{i}.hook_out`, which is now materialized from vLLM's `(mlp_delta, residual)` tuple) must be within target.", "id": "3e66a40c8c0a" }, { @@ -321,6 +305,20 @@ ], "id": "8c94e6f538f4" }, + { + "cell_type": "markdown", + "id": "57e058a4", + "source": "## Step 6½ — Per-position interventions (`pos`)\n\nWhole-sequence interventions (Step 6) edit every position. With the driver booted\n`enable_position_interventions=True`, a spec can carry a **`pos`** field (int or\n`list[int]`) that scopes the affine to specific sequence rows — position-scoped\nactivation patching / tensor injection, the same surface the Inspect/HF backend\nalready offers. This exercises the `(max_num_batched_tokens, width)` affine buffers\nand the per-row `narrow()` path in the compiled hook.\n\nThe test edits `embed.hook_out` at chosen positions and verifies, from the captured\n(post-affine) cache, that **only** those rows changed and every other row stayed\nidentical to the clean run.", + "metadata": {} + }, + { + "cell_type": "code", + "id": "91576f5a", + "source": "# Per-position interventions: a spec's 'pos' (int or list[int]) scopes the affine to\n# specific sequence rows. Verified against the captured (post-affine) embed.hook_out:\n# only the targeted rows change; every other row matches the clean baseline (fp16).\nseq = tokens.shape[1]\nd_model = bridge.cfg.d_model\nHOOK = \"embed.hook_out\"\n\n_, cache_base = bridge.run_with_cache(tokens)\nbase = cache_base[HOOK][0].float().clone()\n\n# (a) add a constant vector at ONE position; all other rows must be untouched.\nPOS = 2\n_, cache_add = bridge.run_with_cache(\n tokens, intervene={HOOK: {\"op\": \"add\", \"value\": [3.0] * d_model, \"pos\": POS}}\n)\nadd = cache_add[HOOK][0].float()\noff = [p for p in range(seq) if p != POS]\nprint(f\"add@{POS}: Δ@POS mean={(add[POS] - base[POS]).mean():.3f} (→3.0) \"\n f\"max|Δ| off-POS={(add[off] - base[off]).abs().max():.2e} (→0)\")\nassert torch.allclose(add[POS], base[POS] + 3.0, atol=1e-2), \"add@pos wrong at the target row\"\nassert torch.allclose(add[off], base[off], atol=1e-2), \"add@pos leaked to other rows\"\n\n# (b) suppress a LIST of positions; only those rows go to zero, the rest pass through.\nPOSL = [0, seq - 1]\n_, cache_sup = bridge.run_with_cache(tokens, intervene={HOOK: {\"op\": \"suppress\", \"pos\": POSL}})\nsup = cache_sup[HOOK][0].float()\nkeep = [p for p in range(seq) if p not in POSL]\nprint(f\"suppress@{POSL}: |max|@POS={sup[POSL].abs().max():.2e} (→0) \"\n f\"max|Δ| kept-rows={(sup[keep] - base[keep]).abs().max():.2e} (→0)\")\nassert sup[POSL].abs().max() < 1e-4, \"suppress@pos did not zero the target rows\"\nassert torch.allclose(sup[keep], base[keep], atol=1e-2), \"suppress@pos leaked to kept rows\"\n\n# (c) reset: a clean forward reverts — pos interventions are per-forward, not sticky.\n_, cache_reset = bridge.run_with_cache(tokens)\nassert torch.allclose(cache_reset[HOOK][0].float(), base, atol=1e-2), \"pos intervention persisted\"\nprint(\"✅ per-position add / suppress / reset all correct\")", + "metadata": {}, + "execution_count": null, + "outputs": [] + }, { "cell_type": "markdown", "id": "a21f3155a525", @@ -555,7 +553,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": "## Summary\n\nIf all asserts above passed, the v4 Driver-protocol vLLM backend is sound on this architecture:\n\n- `boot_vllm` returns `RemoteBridge` end-to-end.\n- `collective_rpc → tl_read_captures` populates the cache.\n- Greedy argmax matches HF (greedy parity).\n- Per-hook L2 < 5e-3 vs HF across every fireable hook (hook-fidelity gate; `ln_final.hook_normalized` exempted — see overlay docstring).\n- Affine interventions (suppress / scale / add / set) apply per-forward and reset cleanly.\n- Each hook fires exactly once per forward under compile (no double-fire / silent overwrite).\n- Back-to-back captures are deterministic (stream-safe).\n- GPU lifetime behaves as expected for vLLM 0.20.x (full release needs a runtime restart).\n- **Batched mode** (`enable_batching=True`, Step 10): `batch_size > 1` parity vs single-prompt runs, chunked-prefill accumulation reconstructs full sequences, and global interventions apply + reset across the batch.\n\nKnown divergences (documented, not bugs): `ln_final.hook_normalized` post-weight vs pre-weight, Gemma `embed.hook_out` scaling. Both reconciled by conversions noted in `sources.vllm.overlays.decoder_only`.\n\nOut of scope: multi-token generation (`max_new_tokens > 1`), per-request (non-global) batched interventions, tensor/pipeline parallelism.", + "source": "## Summary\n\nIf all asserts above passed, the v4 Driver-protocol vLLM backend is sound on this architecture:\n\n- `boot_vllm` returns `RemoteBridge` end-to-end.\n- `collective_rpc → tl_read_captures` populates the cache.\n- Greedy argmax matches HF (greedy parity).\n- Per-hook L2 < 5e-3 vs HF across every fireable hook (hook-fidelity gate; `ln_final.hook_normalized` exempted — see overlay docstring).\n- Affine interventions (suppress / scale / add / set) apply per-forward and reset cleanly — both whole-sequence (Step 6) and scoped to specific positions via a spec `pos` field (Step 6½, `enable_position_interventions=True`).\n- Each hook fires exactly once per forward under compile (no double-fire / silent overwrite).\n- Back-to-back captures are deterministic (stream-safe).\n- GPU lifetime behaves as expected for vLLM 0.20.x (full release needs a runtime restart).\n- **Batched mode** (`enable_batching=True`, Step 10): `batch_size > 1` parity vs single-prompt runs, chunked-prefill accumulation reconstructs full sequences, and global interventions apply + reset across the batch.\n\nKnown divergences (documented, not bugs): `ln_final.hook_normalized` post-weight vs pre-weight, Gemma `embed.hook_out` scaling. Both reconciled by conversions noted in `sources.vllm.overlays.decoder_only`.\n\nOut of scope: multi-token generation (`max_new_tokens > 1`), per-request (non-global) batched interventions, tensor/pipeline parallelism.", "id": "fedb17c9a539" } ], diff --git a/scripts/inspect_parity_report.py b/scripts/inspect_parity_report.py index 4d643f0c4..9a6083ca8 100644 --- a/scripts/inspect_parity_report.py +++ b/scripts/inspect_parity_report.py @@ -60,13 +60,21 @@ ATOL, RTOL = 1e-3, 1e-3 -# Boundary kind -> TransformerBridge-native hook suffix. +# Kind -> TransformerBridge-native hook suffix: the five d_model boundaries plus the +# head-split kinds (served where the structural probe finds the projections — q/k/v are +# gated on fused-qkv archs, pattern under non-eager attention; gated kinds show up in +# the report's `gated=` field, not as failures). KIND_SUFFIX = { "resid_pre": "hook_in", "resid_mid": "ln2.hook_in", "resid_post": "hook_out", "attn_out": "attn.hook_out", "mlp_out": "mlp.hook_out", + "q": "attn.hook_q", + "k": "attn.hook_k", + "v": "attn.hook_v", + "z": "attn.hook_z", + "pattern": "attn.hook_pattern", } @@ -103,6 +111,11 @@ def verify(model_id: str) -> dict: mism.append(f"{hk} missing") continue a, b = hf_cache[hk].float(), i_cache[hk].float() + if a.shape != b.shape and a.numel() == b.numel(): + # OPT-style blocks flatten the FFN to (batch·seq, d) and the bridge + # caches that raw 2-D layout; the driver emits the conventional + # (1, seq, d). Same values, pure reshape — normalize to compare. + a = a.reshape(b.shape) if a.shape != b.shape: mism.append(f"{hk} shape {tuple(a.shape)}!={tuple(b.shape)}") continue diff --git a/tests/acceptance/model_bridge/test_inspect_provider.py b/tests/acceptance/model_bridge/test_inspect_provider.py index 6bc51e725..46bbac2a7 100644 --- a/tests/acceptance/model_bridge/test_inspect_provider.py +++ b/tests/acceptance/model_bridge/test_inspect_provider.py @@ -35,6 +35,10 @@ "blocks.0.hook_out", # resid_post "blocks.6.attn.hook_out", "blocks.11.hook_out", + # Head-split kinds gpt2 serves despite its fused c_attn (q/k/v are gated): z reads + # the c_proj (Conv1D) input, pattern rides output_attentions under eager. + "blocks.0.attn.hook_z", + "blocks.0.attn.hook_pattern", ] # Boundary kind -> TransformerBridge-native hook suffix. @@ -48,12 +52,15 @@ # Token-free tiny-random checkpoints, one per structural-check code path: standard # sequential (nothing gated), parallel-residual (resid_mid gated by the causal probe), -# post-norm (resid_mid gated by the linear-identity probe). Pins cross-family behavior -# so a detector/load-path regression on non-gpt2 archs fails CI instead of shipping silently. +# post-norm (resid_mid gated by the linear-identity probe), and fc-split FFN (OPT/XGLM: +# no mlp container — fc1/fc2 on the block; mlp_out must still be located and match). Pins +# cross-family behavior so a detector/load-path regression on non-gpt2 archs fails CI. STRUCTURAL_FAMILIES = [ ("hf-internal-testing/tiny-random-LlamaForCausalLM", frozenset()), ("hf-internal-testing/tiny-random-GPTJForCausalLM", frozenset({"resid_mid"})), ("hf-internal-testing/tiny-random-Gemma2ForCausalLM", frozenset({"resid_mid"})), + ("hf-internal-testing/tiny-random-OPTForCausalLM", frozenset()), + ("hf-internal-testing/tiny-random-XGLMForCausalLM", frozenset()), ] @@ -245,9 +252,11 @@ def test_sequential_offers_resid_mid(self): ) m = self._toy_model() # standard sequential, linear residual - kinds, note = _detect_capabilities(m, m.layers) + kinds, note = _detect_capabilities(m, m.layers, (None, None, None)) assert "resid_mid" in kinds - assert note == "" # nothing gated + # No boundary kind gated. (Head kinds ARE gated — the toy has no head geometry — + # so the note mentions only those.) + assert "resid_mid" not in note and "attn_out" not in note and "mlp_out" not in note def test_parallel_gates_resid_mid(self): from transformer_lens.model_bridge.sources.inspect.transformers_provider import ( @@ -255,7 +264,7 @@ def test_parallel_gates_resid_mid(self): ) m = self._toy_model(parallel=True) # mlp reads block input, not attn output - kinds, note = _detect_capabilities(m, m.layers) + kinds, note = _detect_capabilities(m, m.layers, (None, None, None)) assert "resid_mid" not in kinds assert {"resid_pre", "resid_post", "attn_out", "mlp_out"} <= kinds # rest still served assert "resid_mid" in note # note explains the gate @@ -268,7 +277,7 @@ def test_nonlinear_residual_gates_resid_mid(self): # Sequential (attn feeds mlp) but outputs are scaled before the residual add, so # resid_post != resid_pre + attn_out + mlp_out — resid_mid must still be gated. m = self._toy_model(resid_scale=2.0) - kinds, note = _detect_capabilities(m, m.layers) + kinds, note = _detect_capabilities(m, m.layers, (None, None, None)) assert "resid_mid" not in kinds assert {"resid_pre", "resid_post", "attn_out", "mlp_out"} <= kinds assert "resid_mid" in note @@ -282,7 +291,7 @@ def test_probe_leaves_global_rng_untouched(self): m = self._toy_model() before = torch.get_rng_state() - _detect_capabilities(m, m.layers) + _detect_capabilities(m, m.layers, (None, None, None)) assert torch.equal(torch.get_rng_state(), before) @@ -327,6 +336,137 @@ def test_gating_and_parity(self, model_id, expected_gated): inspect.close() +HEAD_MODEL = "hf-internal-testing/tiny-random-LlamaForCausalLM" +HEAD_HOOKS = [ + "blocks.0.attn.hook_q", + "blocks.0.attn.hook_k", + "blocks.0.attn.hook_v", + "blocks.0.attn.hook_z", + "blocks.0.attn.hook_pattern", + "blocks.1.attn.hook_z", +] + + +class TestHeadSplitHooks: + """Head-split q/k/v/z/pattern on a separate-projection arch (tiny-random Llama): + capture parity vs boot_transformers, interventions, and the fused-qkv gating path + (via the module-scoped gpt2 fixtures).""" + + @pytest.fixture(scope="class") + def head_pair(self): + from transformer_lens.model_bridge.remote_bridge import RemoteBridge + from transformer_lens.model_bridge.transformer_bridge import TransformerBridge + + hf = TransformerBridge.boot_transformers(HEAD_MODEL, device="cpu", dtype=torch.float32) + inspect = RemoteBridge.boot_inspect(HEAD_MODEL, dtype=torch.float32) + yield hf, inspect + inspect.close() + + def test_all_head_hooks_served(self, head_pair): + _, inspect = head_pair + for hook in HEAD_HOOKS: + assert hook in inspect._driver.supported_hook_points, f"{hook} not served" + + def test_head_capture_parity(self, head_pair): + hf, inspect = head_pair + toks = hf.to_tokens(PROMPT) + _, hf_cache = hf.run_with_cache(toks) + _, i_cache = inspect.run_with_cache(toks) + for hook in HEAD_HOOKS: + a, b = hf_cache[hook].float(), i_cache[hook].float() + assert a.shape == b.shape, f"{hook}: {tuple(b.shape)} vs bridge {tuple(a.shape)}" + assert torch.allclose( + a, b, atol=1e-4, rtol=1e-4 + ), f"{hook} diverges: max {(a - b).abs().max().item():.2e}" + + def test_suppress_v_zeroes_capture_and_moves_logits(self, head_pair): + _, inspect = head_pair + toks = inspect.to_tokens(PROMPT) + base = inspect.forward(toks) + logits, cache = inspect.run_with_cache( + toks, intervene={"blocks.0.attn.hook_v": {"op": "suppress"}} + ) + assert cache["blocks.0.attn.hook_v"].abs().max().item() == 0.0 + assert not torch.allclose(base, logits) + + def test_per_position_q_patch_is_position_scoped(self, head_pair): + _, inspect = head_pair + toks = inspect.to_tokens(PROMPT) + _, base_cache = inspect.run_with_cache(toks) + _, cache = inspect.run_with_cache( + toks, intervene={"blocks.0.attn.hook_q": {"op": "add", "value": 5.0, "pos": 1}} + ) + q_base, q_new = base_cache["blocks.0.attn.hook_q"][0], cache["blocks.0.attn.hook_q"][0] + others = [p for p in range(q_base.shape[0]) if p != 1] + assert torch.allclose(q_new[1], q_base[1] + 5.0, atol=1e-5) + assert torch.equal(q_new[others], q_base[others]) + + def test_pattern_intervention_rejected(self, head_pair): + _, inspect = head_pair + with pytest.raises(ValueError, match="capture-only"): + inspect.forward( + inspect.to_tokens(PROMPT), + intervene={"blocks.0.attn.hook_pattern": {"op": "suppress"}}, + ) + + def test_gpt2_fused_qkv_gated_but_z_pattern_served(self, inspect_bridge): + supported = inspect_bridge._driver.supported_hook_points + nonfireable = inspect_bridge._driver.non_fireable_hook_points + for hook in ("blocks.0.attn.hook_q", "blocks.0.attn.hook_k", "blocks.0.attn.hook_v"): + assert hook not in supported and hook in nonfireable + assert "blocks.0.attn.hook_z" in supported + assert "blocks.0.attn.hook_pattern" in supported + + def test_gptneo_wrapper_attention_serves_head_hooks(self): + """GPT-Neo wraps the real attention (standard q_proj/out_proj) at attn.attention — + the projection-host descent must find it, and captures must match the bridge.""" + from transformer_lens.model_bridge.remote_bridge import RemoteBridge + from transformer_lens.model_bridge.transformer_bridge import TransformerBridge + + mid = "hf-internal-testing/tiny-random-GPTNeoForCausalLM" + hf = TransformerBridge.boot_transformers(mid, device="cpu", dtype=torch.float32) + inspect = RemoteBridge.boot_inspect(mid, dtype=torch.float32) + try: + for kind in ("q", "k", "v", "z", "pattern"): + assert f"blocks.0.attn.hook_{kind}" in inspect._driver.supported_hook_points + toks = hf.to_tokens(PROMPT) + _, hf_cache = hf.run_with_cache(toks) + _, i_cache = inspect.run_with_cache(toks) + for hook in ("blocks.0.attn.hook_q", "blocks.0.attn.hook_z"): + a, b = hf_cache[hook].float(), i_cache[hook].float() + assert a.shape == b.shape + assert torch.allclose(a, b, atol=1e-4, rtol=1e-4) + finally: + inspect.close() + + def test_provider_direct_interventions_validated(self, inspect_bridge, tokens): + """The provider's documented extra_args interface must reject gated and + capture-only intervention kinds instead of silently no-op'ing (the driver path + already rejects; this covers callers that speak to the provider directly).""" + import asyncio + + from inspect_ai.model import GenerateConfig + + api = inspect_bridge._driver._model.api + + def call(interventions): + cfg = GenerateConfig( + extra_body={ + "extra_args": { + "input_ids": tokens[0].tolist(), + "capture": [], + "interventions": interventions, + } + } + ) + return asyncio.run(api.generate("", [], None, cfg)) + + with pytest.raises(ValueError, match="gated"): # q gated on gpt2's fused c_attn + call({"0:q": {"op": "suppress"}}) + with pytest.raises(ValueError, match="capture-only"): + call({"0:resid_mid": {"op": "suppress"}}) + + class TestRebootSemantics: """inspect_ai memoizes get_model by name; boot_inspect passes memoize=False so each boot honors its own args and independently owns (and frees) its model.""" diff --git a/tests/acceptance/model_bridge/test_symbolic_bridge_hooks.py b/tests/acceptance/model_bridge/test_symbolic_bridge_hooks.py new file mode 100644 index 000000000..8a28edf74 --- /dev/null +++ b/tests/acceptance/model_bridge/test_symbolic_bridge_hooks.py @@ -0,0 +1,56 @@ +"""SymbolicBridge hook mirroring: fc-split archs (OPT/XGLM) must fire the placeholder's +own ``mlp.hook_in``/``mlp.hook_out`` — wired at setup from the ``in``/``out`` subcomponents +(see ``component_setup._wire_symbolic_hooks``). Without the mirror those HookPoints exist +in the registry but never fire, so caching misses them and interventions silently no-op. + +Runs on CPU with token-free tiny-random checkpoints. +""" +from __future__ import annotations + +import pytest +import torch + +OPT = "hf-internal-testing/tiny-random-OPTForCausalLM" +XGLM = "hf-internal-testing/tiny-random-XGLMForCausalLM" + + +@pytest.fixture(scope="module") +def opt_bridge(): + from transformer_lens.model_bridge.transformer_bridge import TransformerBridge + + return TransformerBridge.boot_transformers(OPT, device="cpu", dtype=torch.float32) + + +class TestSymbolicMirror: + def test_mlp_hooks_fire_and_match_subcomponents(self, opt_bridge): + tokens = opt_bridge.to_tokens("Hello world") + _, cache = opt_bridge.run_with_cache(tokens) + assert "blocks.0.mlp.hook_out" in cache.cache_dict + assert "blocks.0.mlp.hook_in" in cache.cache_dict + # The mirror re-fires the subcomponent tensors: out = fc2's output, in = fc1's input. + assert torch.equal(cache["blocks.0.mlp.hook_out"], cache["blocks.0.mlp.out.hook_out"]) + assert torch.equal(cache["blocks.0.mlp.hook_in"], cache["blocks.0.mlp.in.hook_in"]) + + def test_intervention_via_symbolic_hook_propagates(self, opt_bridge): + tokens = opt_bridge.to_tokens("Hello world") + base = opt_bridge.forward(tokens) + patched = opt_bridge.run_with_hooks( + tokens, fwd_hooks=[("blocks.0.mlp.hook_out", lambda t, hook: torch.zeros_like(t))] + ) + assert not torch.allclose(base, patched), "symbolic-hook edit did not reach the logits" + + def test_mirror_survives_reset_hooks(self, opt_bridge): + """The wiring is permanent — a reset_hooks (any run_with_* teardown) must not + sever it, else the second run_with_cache silently loses the mlp entries.""" + tokens = opt_bridge.to_tokens("Hello world") + opt_bridge.reset_hooks() + _, cache = opt_bridge.run_with_cache(tokens) + assert "blocks.0.mlp.hook_out" in cache.cache_dict + + def test_xglm_mirror_fires(self): + from transformer_lens.model_bridge.transformer_bridge import TransformerBridge + + bridge = TransformerBridge.boot_transformers(XGLM, device="cpu", dtype=torch.float32) + _, cache = bridge.run_with_cache(bridge.to_tokens("Hi")) + assert "blocks.0.mlp.hook_out" in cache.cache_dict + assert torch.equal(cache["blocks.0.mlp.hook_out"], cache["blocks.0.mlp.out.hook_out"]) diff --git a/tests/unit/model_bridge/test_inspect_driver.py b/tests/unit/model_bridge/test_inspect_driver.py index b4b789445..1ee943f37 100644 --- a/tests/unit/model_bridge/test_inspect_driver.py +++ b/tests/unit/model_bridge/test_inspect_driver.py @@ -423,6 +423,124 @@ def test_context_manager_closes(self): assert driver._model is None +class TestHeadSplitKinds: + """Head-split q/k/v/z/pattern: registry entries, driver assembly, intervention gating. + + The provider serves these only when its structural probe finds the projections; here + a profile with an explicit kind set stands in for that detection. + """ + + def _head_profile(self): + return profiles.TLBridgeProfile(supported_kinds=hooks.ALL_KINDS | hooks.HEAD_KINDS) + + def test_resolve_head_names(self): + assert hooks.resolve("blocks.3.attn.hook_q") == (3, "q") + assert hooks.resolve("blocks.0.attn.hook_z") == (0, "z") + assert hooks.resolve("blocks.1.attn.hook_pattern") == (1, "pattern") + assert hooks.resolve("blocks.0.attn.hook_attn_scores") is None # never served + + def test_default_supported_excludes_head_kinds(self): + """kinds=None stays boundary-only — head kinds are opt-in via detection, so + profiles that pass None (vllm-lens, defaults) don't silently claim them.""" + default = hooks.supported_hook_points(N_LAYERS) + assert "blocks.0.attn.hook_q" not in default + assert "blocks.0.attn.hook_pattern" not in default + + def test_all_hook_points_is_boundary_plus_head(self): + universe = hooks.all_hook_points(N_LAYERS) + assert hooks.supported_hook_points(N_LAYERS) < universe + assert "blocks.0.attn.hook_q" in universe + assert len(universe) == (len(hooks.ALL_KINDS) + len(hooks.HEAD_KINDS)) * N_LAYERS + + def test_interveneable_includes_qkvz_not_pattern(self): + assert {"q", "k", "v", "z"} <= hooks.INTERVENEABLE_KINDS + assert "pattern" not in hooks.INTERVENEABLE_KINDS + + def test_attn_scores_always_nonfireable(self): + driver = InspectDriver( + _fake_model(), _adapter(), tokenizer=None, profile=self._head_profile() + ) + assert "blocks.0.attn.hook_attn_scores" in driver.non_fireable_hook_points + + def test_driver_serves_head_kinds_via_profile(self): + driver = InspectDriver( + _fake_model(), _adapter(), tokenizer=None, profile=self._head_profile() + ) + for name in ("blocks.0.attn.hook_q", "blocks.1.attn.hook_z", "blocks.0.attn.hook_pattern"): + assert name in driver.supported_hook_points + assert driver.supported_hook_points.isdisjoint(driver.non_fireable_hook_points) + + def test_default_driver_moves_head_kinds_to_nonfireable(self): + driver = _driver() # default profile: boundary kinds only + assert "blocks.0.attn.hook_q" in driver.non_fireable_hook_points + assert "blocks.0.attn.hook_pattern" in driver.non_fireable_hook_points + + def test_head_captures_get_batch_dim(self): + """3-D wire arrays (seq,heads,d_head) / (heads,q,k) unsqueeze to exactly one batch dim.""" + heads, d_head, seq = 2, 2, 3 + caps = { + "0:q": np.ones((seq, heads, d_head), np.float32), + "0:pattern": np.ones((heads, seq, seq), np.float32), + "0:resid_post": np.ones((seq, D_MODEL), np.float32), + } + driver = InspectDriver( + _fake_model(captures=caps, logits=np.zeros((seq, D_VOCAB), np.float32)), + _adapter(), + tokenizer=None, + profile=self._head_profile(), + ) + result = driver.forward( + torch.tensor([[1, 2, 3]]), + capture=("blocks.0.attn.hook_q", "blocks.0.attn.hook_pattern", "blocks.0.hook_out"), + ) + assert tuple(result.captured["blocks.0.attn.hook_q"].shape) == (1, seq, heads, d_head) + assert tuple(result.captured["blocks.0.attn.hook_pattern"].shape) == (1, heads, seq, seq) + assert tuple(result.captured["blocks.0.hook_out"].shape) == (1, seq, D_MODEL) + + def test_qkvz_intervention_translates_to_wire_key(self): + supported = hooks.supported_hook_points(N_LAYERS, hooks.ALL_KINDS | hooks.HEAD_KINDS) + out = intervention.build_interventions( + {"blocks.0.attn.hook_v": {"op": "suppress"}}, supported + ) + assert out == {"0:v": {"op": "suppress"}} + + def test_pattern_intervention_rejected(self): + supported = hooks.supported_hook_points(N_LAYERS, hooks.ALL_KINDS | hooks.HEAD_KINDS) + with pytest.raises(ValueError, match="capture-only"): + intervention.build_interventions( + {"blocks.0.attn.hook_pattern": {"op": "suppress"}}, supported + ) + + def test_head_hook_rejected_when_gated(self): + """A gated head hook (default boundary-only profile) can't be intervened on.""" + supported = hooks.supported_hook_points(N_LAYERS) # no head kinds + with pytest.raises(ValueError, match="not in supported_hook_points"): + intervention.build_interventions( + {"blocks.0.attn.hook_q": {"op": "suppress"}}, supported + ) + + def test_turn_activations_batch_dim_is_rank_aware(self): + """Mixed 2-D boundary and 3-D head-split arrays in one turn all get exactly one + batch dim — a rank-2-only rule would leave head-split arrays batchless, and a + batchless (seq, heads, d_head) is shape-indistinguishable from a batched 2-D.""" + from transformer_lens.model_bridge.sources.inspect.eval import turn_activations + + payload = wire.encode_activations( + { + "0:resid_post": np.ones((3, D_MODEL), np.float32), + "0:q": np.ones((3, 2, 2), np.float32), + "0:pattern": np.ones((2, 3, 3), np.float32), + } + ) + sample = SimpleNamespace( + events=[SimpleNamespace(output=SimpleNamespace(metadata={"activations": payload}))] + ) + (turn,) = turn_activations(sample) + assert turn["blocks.0.hook_out"].shape == (1, 3, D_MODEL) + assert turn["blocks.0.attn.hook_q"].shape == (1, 3, 2, 2) + assert turn["blocks.0.attn.hook_pattern"].shape == (1, 2, 3, 3) + + def test_driver_imports_no_torch(): """The acceptance gate: the driver file must not import torch (data-only boundary).""" import transformer_lens.model_bridge.sources.inspect.driver as drv diff --git a/tests/unit/model_bridge/test_vllm_boot.py b/tests/unit/model_bridge/test_vllm_boot.py index cd0edaff0..7d79bbc63 100644 --- a/tests/unit/model_bridge/test_vllm_boot.py +++ b/tests/unit/model_bridge/test_vllm_boot.py @@ -107,6 +107,20 @@ def test_rejects_locked_kwarg_override(): boot_vllm("any-model", tensor_parallel_size=2) +def test_rejects_position_interventions_with_batching(): + """Position interventions need the compiled path's affine buffers — fails fast.""" + with pytest.raises(ValueError, match="incompatible with enable_batching"): + boot_vllm("any-model", enable_position_interventions=True, enable_batching=True) + + +def test_position_interventions_flag_reaches_driver(mocked_boot): + """boot_vllm threads enable_position_interventions through to the driver.""" + bridge = boot_vllm("any-model", enable_position_interventions=True) + assert bridge._driver._enable_position_interventions is True + # Default stays off. + assert boot_vllm("any-model")._driver._enable_position_interventions is False + + @pytest.mark.parametrize( "raw, expected", [ diff --git a/tests/unit/model_bridge/test_vllm_driver.py b/tests/unit/model_bridge/test_vllm_driver.py index f676a005d..80140367c 100644 --- a/tests/unit/model_bridge/test_vllm_driver.py +++ b/tests/unit/model_bridge/test_vllm_driver.py @@ -89,6 +89,7 @@ def _driver( max_num_batched_tokens=2048, generated_token=None, top_logprobs=None, + enable_position_interventions=False, ) -> VLLMDriver: """Build a VLLMDriver. ``captures`` populates llm.collective_rpc; ``generated_token`` and ``top_logprobs`` populate llm.generate's RequestOutput so _synthesize_logits @@ -104,6 +105,7 @@ def _driver( overlay=_overlay(), hf_config=hf_config or _hf_config(), max_num_batched_tokens=max_num_batched_tokens, + enable_position_interventions=enable_position_interventions, ) @@ -115,11 +117,19 @@ def test_passes_validate_driver(self): assert isinstance(driver, Driver) validate_driver(driver) - def test_no_torch_capability_flags(self): - """vLLM owns the model in a worker — no torch-specific capability surface.""" + def test_capability_flags(self): + """vLLM owns the model in a worker — no torch module surface (parameters/ + state_dict/grads), but named-weight reads ARE served via the get_param RPC.""" driver = _driver() - for feature in ("parameters", "state_dict", "gradients", "weight_access"): + for feature in ("parameters", "state_dict", "gradients"): assert driver.supports(feature) is False, f"vLLM shouldn't support {feature!r}" + assert driver.supports("weight_access") is True + + def test_zero_layers_config_fails_loud(self): + """A missing/zero num_hidden_layers must raise at boot, not leave a raw '{i}' + template in non_fireable_hook_points.""" + with pytest.raises(ValueError, match="num_hidden_layers"): + _driver(hf_config=_hf_config(num_hidden_layers=0)) def test_non_fireable_expanded_per_layer(self): """``blocks.{i}.attn.hook_pattern`` template expands to one entry per layer.""" @@ -168,7 +178,7 @@ def test_forward_logits_from_sampler_logprobs(self): result = _driver( captures={"embed.hook_out": torch.randn(3, 4)}, top_logprobs={7: 2.5, 3: 1.0, 11: -0.5}, - ).forward(torch.tensor([[1, 2, 3]])) + ).forward(torch.tensor([[1, 2, 3]]), capture=("embed.hook_out",)) assert isinstance(result, ForwardResult) assert tuple(result.captured["embed.hook_out"].shape) == (1, 3, 4) assert result.logits is not None and tuple(result.logits.shape) == (1, 3, 16) @@ -234,6 +244,67 @@ def test_forward_rejects_set_missing_value(self): torch.tensor([[1]]), intervene={"embed.hook_out": {"op": "set"}} ) + +class TestVLLMDriverPositionInterventions: + """`pos`-scoped interventions require the opt-in (max_n, width) affine buffers.""" + + def test_pos_rejected_without_flag(self): + with pytest.raises(NotImplementedError, match="enable_position_interventions=True"): + _driver()._validate_interventions({"embed.hook_out": {"op": "suppress", "pos": 0}}) + + def test_pos_accepted_and_preserved_with_flag(self): + driver = _driver(enable_position_interventions=True) + out = driver._validate_interventions( + {"embed.hook_out": {"op": "add", "value": 1.0, "pos": [0, 2]}} + ) + assert out["embed.hook_out"]["pos"] == [0, 2] + + def test_pos_int_accepted_with_flag(self): + driver = _driver(enable_position_interventions=True) + out = driver._validate_interventions({"embed.hook_out": {"op": "suppress", "pos": 3}}) + assert out["embed.hook_out"]["pos"] == 3 + + def test_pos_wrong_type_raises(self): + driver = _driver(enable_position_interventions=True) + with pytest.raises(ValueError, match="must be an int or list of ints"): + driver._validate_interventions({"embed.hook_out": {"op": "suppress", "pos": "last"}}) + + def test_pos_negative_raises(self): + driver = _driver(enable_position_interventions=True) + with pytest.raises(ValueError, match="must be non-negative"): + driver._validate_interventions({"embed.hook_out": {"op": "suppress", "pos": [-1]}}) + + def test_pos_rejected_on_batched_path(self): + """The batched/eager path has no affine buffers, so 'pos' is unsupported there.""" + driver = _driver() + driver._enable_batching = True + with pytest.raises(NotImplementedError, match="batched/eager path"): + driver._validate_interventions({"embed.hook_out": {"op": "suppress", "pos": 0}}) + + def test_no_pos_still_works_with_flag(self): + """A whole-sequence spec (no pos) is unaffected by the flag.""" + driver = _driver(enable_position_interventions=True) + out = driver._validate_interventions({"embed.hook_out": {"op": "suppress"}}) + assert "pos" not in out["embed.hook_out"] + + def test_pos_beyond_prompt_length_fails_loud(self): + """pos within buffer capacity but past the prompt length must raise, not silently no-op. + + The rejection fires before driver.forward reaches the vllm import, so no install needed. + """ + driver = _driver(captures={}, enable_position_interventions=True) + with pytest.raises(ValueError, match="beyond the prompt length"): + driver.forward( # 3-token prompt, pos=50 is unreadable by the hook + torch.tensor([[1, 2, 3]]), + intervene={"embed.hook_out": {"op": "suppress", "pos": 50}}, + ) + + def test_reject_pos_beyond_seq_bounds_against_actual_length(self): + driver = _driver(enable_position_interventions=True) + driver._reject_pos_beyond_seq({"embed.hook_out": {"op": "suppress", "pos": [0, 2]}}, 3) + with pytest.raises(ValueError, match="beyond the prompt length"): + driver._reject_pos_beyond_seq({"embed.hook_out": {"op": "suppress", "pos": 3}}, 3) + def test_forward_pushes_interventions_before_generate(self): """The driver pushes spec dicts via collective_rpc('tl_set_interventions', ...).""" pytest.importorskip("vllm") @@ -267,6 +338,29 @@ def test_forward_rejects_prompt_exceeding_buffer(self): with pytest.raises(ValueError, match="exceeds max_num_batched_tokens"): _driver(captures={}, max_num_batched_tokens=4).forward(torch.tensor([[1, 2, 3, 4, 5]])) + def test_zero_capture_skips_worker_read(self): + """capture=() must not trigger a GPU→CPU copy: with logits off there is no + tl_read_captures RPC at all (an empty tuple used to collapse to None = read + every buffer); with logits on, only the forced ln_final is read.""" + pytest.importorskip("vllm") + driver = _driver(captures={"embed.hook_out": torch.zeros(3, 4)}) + driver.forward(torch.tensor([[1, 2, 3]]), return_logits=False) + reads = [ + c.args + for c in driver._llm.collective_rpc.call_args_list + if c.args[0] == "tl_read_captures" + ] + assert reads == [], f"hookless forward still read the buffers: {reads}" + + driver2 = _driver(captures={"ln_final.hook_normalized": torch.zeros(3, 4)}) + driver2.forward(torch.tensor([[1, 2, 3]])) # return_logits=True default + reads2 = [ + c.args + for c in driver2._llm.collective_rpc.call_args_list + if c.args[0] == "tl_read_captures" + ] + assert len(reads2) == 1 and reads2[0][1][1] == ["ln_final.hook_normalized"] + def _batched_request_output(request_id, generated_token=None, top_logprobs=None): """RequestOutput-shaped mock carrying a request_id for the accumulator join.""" @@ -338,7 +432,7 @@ def test_req_id_join_not_positional(self): "req-A": {"embed.hook_out": torch.full((3, 4), 1.0)}, } result = _batched_driver(outputs=outputs, captures_by_req=captures_by_req).forward( - [[1, 2, 3], [4, 5]] + [[1, 2, 3], [4, 5]], capture=("embed.hook_out",) ) emb = result.captured["embed.hook_out"] assert tuple(emb.shape) == (2, 3, 4) # (batch, max_seq, width) @@ -360,7 +454,7 @@ def test_join_strips_internal_req_id_suffix(self): "11-a1b2c3d4": {"embed.hook_out": torch.full((2, 4), 2.0)}, } result = _batched_driver(outputs=outputs, captures_by_req=captures_by_req).forward( - [[1, 2, 3], [4, 5]] + [[1, 2, 3], [4, 5]], capture=("embed.hook_out",) ) emb = result.captured["embed.hook_out"] assert tuple(emb.shape) == (2, 3, 4) @@ -379,7 +473,7 @@ def test_join_prefix_does_not_collide_on_numeric_ids(self): "10-bbbbbbbb": {"embed.hook_out": torch.full((3, 4), 9.0)}, } result = _batched_driver(outputs=outputs, captures_by_req=captures_by_req).forward( - [[1, 2], [1, 2, 3]] + [[1, 2], [1, 2, 3]], capture=("embed.hook_out",) ) emb = result.captured["embed.hook_out"] assert torch.equal(emb[0, :2], torch.full((2, 4), 1.0)) # req "1" @@ -396,12 +490,31 @@ def test_shorter_rows_right_padded_with_zeros(self): "r1": {"embed.hook_out": torch.ones(1, 4)}, } result = _batched_driver(outputs=outputs, captures_by_req=captures_by_req).forward( - [[1, 2, 3], [9]] + [[1, 2, 3], [9]], capture=("embed.hook_out",) ) emb = result.captured["embed.hook_out"] # Row 1 has 1 real token; positions 1,2 are zero pad. assert torch.equal(emb[1, 1:], torch.zeros(2, 4)) + def test_zero_capture_skips_batched_join(self): + """capture=() with logits off must return empty captures, not crash: _assemble_padded + needs one worker key per request, so it can't run on the empty read — mirror the + single path. (Regression: the batched join was called unconditionally.)""" + pytest.importorskip("vllm") + outputs = [ + _batched_request_output("r0", top_logprobs={1: 1.0}), + _batched_request_output("r1", top_logprobs={1: 1.0}), + ] + driver = _batched_driver(outputs=outputs, captures_by_req={}) + result = driver.forward([[1, 2, 3], [9]], return_logits=False) # capture=() default + assert result.captured == {} + reads = [ + c.args + for c in driver._llm.collective_rpc.call_args_list + if c.args[0] == "tl_read_batched_captures" + ] + assert reads == [], f"batched hookless forward still read captures: {reads}" + def test_logits_at_per_row_last_token(self): """Next-token logits land at prompt_len-1 per row, never -1 (a pad row).""" pytest.importorskip("vllm") @@ -414,7 +527,7 @@ def test_logits_at_per_row_last_token(self): "r1": {"embed.hook_out": torch.ones(1, 4)}, } result = _batched_driver(outputs=outputs, captures_by_req=captures_by_req).forward( - [[1, 2, 3], [9]] + [[1, 2, 3], [9]], capture=("embed.hook_out",) ) logits = result.logits assert logits is not None and tuple(logits.shape) == (2, 3, 16) @@ -481,12 +594,18 @@ def test_bridge_replays_vllm_captures(self): assert len(fired) == 1 assert tuple(fired[0].shape) == (1, 3, 4) - def test_forward_rejects_loss_return_types(self): - """Synthesized logits are last-position-only; loss/both would be nan, so refuse. + def test_loss_return_type_gated_on_sequence_logits(self): + """The bridge's loss/both guard fires only when the driver lacks full-sequence + logits. The vLLM driver now reconstructs them (provides_sequence_logits=True), + so loss/both are permitted; a driver without that capability is still refused. - The guard is the first thing forward() does, so this needs no vllm install. + The rejection path fires before driver.forward, so it needs no vllm install. """ - bridge = RemoteBridge(adapter=_adapter(), tokenizer=None, driver=_driver(captures={})) + driver = _driver(captures={}) + assert driver.provides_sequence_logits is True # reconstruction path is live + bridge = RemoteBridge(adapter=_adapter(), tokenizer=None, driver=driver) + # Drop the capability → the guard must refuse loss/both (nan over -inf positions). + driver.provides_sequence_logits = False for rt in ("loss", "both"): with pytest.raises(NotImplementedError, match="return_type"): bridge.forward(torch.tensor([[1, 2, 3]]), return_type=rt) diff --git a/tests/unit/model_bridge/test_vllm_worker_extension.py b/tests/unit/model_bridge/test_vllm_worker_extension.py index 5d06ca40a..3d254347c 100644 --- a/tests/unit/model_bridge/test_vllm_worker_extension.py +++ b/tests/unit/model_bridge/test_vllm_worker_extension.py @@ -135,6 +135,71 @@ def test_value_shape_mismatch_raises(self): _apply_intervention(scale, bias, {"op": "add", "value": [1.0, 2.0]}) +class TestApplyInterventionPerPosition: + """2-D (max_n, width) affine buffers let a spec's 'pos' scope the edit to rows. + + Buffers start at identity (scale=1, bias=0), mirroring tl_set_interventions' + reset before each apply, so a pos-scoped edit must leave the other rows untouched. + """ + + def _buffers(self, max_n=5, width=4): + return torch.ones(max_n, width), torch.zeros(max_n, width) + + def test_set_at_single_pos_leaves_others_identity(self): + scale, bias = self._buffers() + _apply_intervention(scale, bias, {"op": "set", "value": 2.0, "pos": 2}) + assert torch.equal(scale[2], torch.zeros(4)) + assert torch.equal(bias[2], torch.full((4,), 2.0)) + for r in (0, 1, 3, 4): + assert torch.equal(scale[r], torch.ones(4)) + assert torch.equal(bias[r], torch.zeros(4)) + + def test_add_vector_at_pos_list(self): + scale, bias = self._buffers() + vec = [1.0, 2.0, 3.0, 4.0] + _apply_intervention(scale, bias, {"op": "add", "value": vec, "pos": [0, 3]}) + for r in (0, 3): + assert torch.equal(scale[r], torch.ones(4)) + assert torch.equal(bias[r], torch.tensor(vec)) + for r in (1, 2, 4): + assert torch.equal(bias[r], torch.zeros(4)) + + def test_suppress_at_pos(self): + scale, bias = self._buffers() + _apply_intervention(scale, bias, {"op": "suppress", "pos": 1}) + assert torch.equal(scale[1], torch.zeros(4)) + assert torch.equal(bias[1], torch.zeros(4)) + assert torch.equal(scale[0], torch.ones(4)) # untouched + + def test_scale_at_pos(self): + scale, bias = self._buffers() + _apply_intervention(scale, bias, {"op": "scale", "factor": 0.5, "pos": 4}) + assert torch.equal(scale[4], torch.full((4,), 0.5)) + assert torch.equal(scale[0], torch.ones(4)) + + def test_no_pos_writes_all_rows(self): + """A whole-sequence spec on 2-D buffers broadcasts across every row.""" + scale, bias = self._buffers() + _apply_intervention(scale, bias, {"op": "add", "value": [1.0, 2.0, 3.0, 4.0]}) + for r in range(5): + assert torch.equal(bias[r], torch.tensor([1.0, 2.0, 3.0, 4.0])) + + def test_pos_out_of_range_raises(self): + import pytest + + scale, bias = self._buffers(max_n=3) + with pytest.raises(ValueError, match="out of range"): + _apply_intervention(scale, bias, {"op": "suppress", "pos": 5}) + + def test_pos_on_1d_buffer_raises(self): + """A 'pos' spec against 1-D buffers (flag off) is a misconfiguration.""" + import pytest + + scale, bias = torch.ones(4), torch.zeros(4) + with pytest.raises(ValueError, match="requires 2-D affine buffers"): + _apply_intervention(scale, bias, {"op": "suppress", "pos": 0}) + + class TestApplyOp: """_apply_op is the eager batched path's tensor-level intervention.""" diff --git a/transformer_lens/model_bridge/component_setup.py b/transformer_lens/model_bridge/component_setup.py index 7821d0354..cb7b026a5 100644 --- a/transformer_lens/model_bridge/component_setup.py +++ b/transformer_lens/model_bridge/component_setup.py @@ -58,6 +58,22 @@ def set_original_components( setup_components(component_mapping, bridge_module, architecture_adapter, original_model) +def _wire_symbolic_hooks(symbolic: SymbolicBridge) -> None: + """Fire a SymbolicBridge's own hook_in/hook_out from its ``in``/``out`` subcomponents. + + Permanent hooks (survive reset_hooks) re-fire the subcomponent activation through the + placeholder's HookPoint, so ``blocks.{i}.mlp.hook_in/hook_out`` (and their compat + aliases) behave like every non-symbolic arch's. A hook that returns a modified tensor + on the symbolic point propagates: the mirror returns it into the subcomponent's chain. + """ + sub_in = symbolic.submodules.get("in") + sub_out = symbolic.submodules.get("out") + if sub_in is not None: + sub_in.hook_in.add_hook(lambda tensor, hook: symbolic.hook_in(tensor), is_permanent=True) + if sub_out is not None: + sub_out.hook_out.add_hook(lambda tensor, hook: symbolic.hook_out(tensor), is_permanent=True) + + def setup_submodules( component: GeneralizedComponent, architecture_adapter: ArchitectureAdapter, @@ -92,6 +108,13 @@ def setup_submodules( for sub_name, (sub_path, sub_comp) in submodule.real_components.items(): prefixed_key = f"{module_name}.{sub_name}" component.real_components[prefixed_key] = (sub_path, sub_comp) + + # The placeholder has no forward, so its own hook_in/hook_out would never + # fire — mirror them from the designated subcomponents (fc-split archs: + # mlp.hook_in = fc1's input, mlp.hook_out = fc2's output). Firing through + # the parent HookPoint keeps caching AND interventions working: the return + # value feeds back into the subcomponent's hook chain. + _wire_symbolic_hooks(submodule) else: # Set up original_component if not already set if submodule.original_component is None: diff --git a/transformer_lens/model_bridge/sources/inspect/driver.py b/transformer_lens/model_bridge/sources/inspect/driver.py index b7dc954ed..518036b27 100644 --- a/transformer_lens/model_bridge/sources/inspect/driver.py +++ b/transformer_lens/model_bridge/sources/inspect/driver.py @@ -48,9 +48,11 @@ def __init__(self, model: Any, adapter: Any, tokenizer: Any, profile: Any = None # Provider-specific: loss/both allowed only if the provider returns full logits. self.provides_sequence_logits = self._profile.provides_sequence_logits self.supported_hook_points = self._profile.supported_hooks(self._n_layers) - full = hooks.supported_hook_points(self._n_layers) + # Everything the registry could serve (boundaries + head-split) that this + # provider/model doesn't, plus the never-fireable set (embed, ln_final, scores). + universe = hooks.all_hook_points(self._n_layers) self.non_fireable_hook_points = hooks.nonfireable_hook_points(self._n_layers) | ( - full - self.supported_hook_points + universe - self.supported_hook_points ) # Background event loop, created lazily on first forward. self._loop: asyncio.AbstractEventLoop | None = None @@ -106,8 +108,9 @@ async def _generate(self, prompt: Any, extra_args: dict[str, Any]) -> Any: return await self._model.generate(prompt, config=config) def _assemble_captures(self, output: Any, names: list[str]) -> dict[str, np.ndarray]: - """Decode the requested boundaries → ``{hook_name: (1, seq, d_model)}``; names the - provider didn't return are skipped (and warned once).""" + """Decode the requested hooks → ``{hook_name: (1, ...)}`` (batch dim added onto the + provider's batchless array — rank 2 for boundaries, 3 for head-split/pattern); + names the provider didn't return are skipped (and warned once).""" metadata = getattr(output, "metadata", None) or {} decoded = wire.decode_activations(metadata, self._wire_keys(names)) captured: dict[str, np.ndarray] = {} @@ -120,7 +123,8 @@ def _assemble_captures(self, output: Any, names: list[str]) -> dict[str, np.ndar if arr is None: missing.append(name) continue - captured[name] = arr[np.newaxis, ...] if arr.ndim == 2 else arr + batchless = hooks.WIRE_BATCHLESS_NDIM.get(resolved[1], 2) + captured[name] = arr[np.newaxis, ...] if arr.ndim == batchless else arr self._warn_missing(missing) return captured diff --git a/transformer_lens/model_bridge/sources/inspect/eval.py b/transformer_lens/model_bridge/sources/inspect/eval.py index eebad3593..176d6036a 100644 --- a/transformer_lens/model_bridge/sources/inspect/eval.py +++ b/transformer_lens/model_bridge/sources/inspect/eval.py @@ -84,9 +84,10 @@ async def solve(state: TaskState, generate: Generate) -> TaskState: def turn_activations(sample: Any) -> list[dict[str, np.ndarray]]: """Per-turn activations from an eval sample's model events, for a provider booted with - ``capture=[...]`` (e.g. ``model_args={"capture": [...]}``). Returns one - ``{hook_name: (1, seq, d)}`` dict per model generation, in turn order — the activations - of an agentic/multi-turn rollout. + ``capture=[...]`` (e.g. ``model_args={"capture": [...]}``). Returns one dict per model + generation, in turn order — the activations of an agentic/multi-turn rollout. Every + array gets a leading batch dim: boundaries are ``(1, seq, d_model)``, head-split + q/k/v/z ``(1, seq, heads, d_head)``. """ turns = [] for event in getattr(sample, "events", []) or []: @@ -94,11 +95,15 @@ def turn_activations(sample: Any) -> list[dict[str, np.ndarray]]: if not metadata or "activations" not in metadata: continue decoded = wire.decode_activations(metadata, list(metadata["activations"])) - named = { - name: arr[np.newaxis, ...] if arr.ndim == 2 else arr - for wk, arr in decoded.items() - if (name := hooks.name_from_wire_key(wk)) is not None - } + named = {} + for wk, arr in decoded.items(): + name = hooks.name_from_wire_key(wk) + if name is None: + continue + # Rank-aware batch dim (mirrors driver._assemble_captures): boundary kinds + # arrive rank-2, head-split kinds rank-3 — unsqueeze exactly once either way. + batchless = hooks.WIRE_BATCHLESS_NDIM.get(wk.partition(":")[2], 2) + named[name] = arr[np.newaxis, ...] if arr.ndim == batchless else arr if named: turns.append(named) return turns diff --git a/transformer_lens/model_bridge/sources/inspect/hooks.py b/transformer_lens/model_bridge/sources/inspect/hooks.py index ed287bca0..228247221 100644 --- a/transformer_lens/model_bridge/sources/inspect/hooks.py +++ b/transformer_lens/model_bridge/sources/inspect/hooks.py @@ -1,8 +1,10 @@ """Canonical hook names ↔ (layer, kind) for the Inspect HF provider. Torch-free; shared by the provider (capture/intervene) and the driver (supported set, -decode). Covers the ``d_model``-shaped decoder-layer boundaries; head-split hooks -(q/k/v/z, pattern), ``embed``, and ``ln_final`` (fold-LN convention) are non-fireable. +decode). Covers the ``d_model``-shaped decoder-layer boundaries plus the head-split +attention hooks (q/k/v/z, pattern) where the provider's structural probe finds the +projections; ``attn_scores``, ``embed``, and ``ln_final`` (fold-LN convention) stay +non-fireable. Names are TransformerBridge-native (``blocks.{i}.hook_out``, ``.attn.hook_out``, ...), not the HookedTransformer aliases. A bridge cache carries both with identical values, @@ -10,8 +12,9 @@ Which boundaries are actually fireable is decided per-model by the provider's structural self-check, not a hand-kept architecture list: it locates attn/mlp and probes whether the -``resid_pre + attn_out`` derivation holds, gating ``resid_mid`` otherwise. -``supported_hook_points(n_layers, kinds=...)`` filters to that detected set. +``resid_pre + attn_out`` derivation holds, gating ``resid_mid`` otherwise; head-split +kinds need separate q/k/v projections (q/k/v), a locatable out-projection (z), or eager +attention (pattern). ``supported_hook_points(n_layers, kinds=...)`` filters to that set. """ from __future__ import annotations @@ -19,6 +22,10 @@ from typing import Iterable, Optional ALL_KINDS = frozenset({"resid_pre", "resid_mid", "resid_post", "attn_out", "mlp_out"}) +# Head-split attention kinds, served only when structurally detected (never by default): +# q/k/v are the pre-RoPE projection outputs, z is the out-projection input (all +# ``(seq, heads, d_head)``); pattern is post-softmax attention ``(heads, q_pos, k_pos)``. +HEAD_KINDS = frozenset({"q", "k", "v", "z", "pattern"}) # One canonical TransformerBridge name per boundary (no aliases — avoids duplicate # HookPoints/cache entries). resid_mid (ln2.hook_in) is derived (resid_pre + @@ -29,8 +36,27 @@ "resid_post": "blocks.{i}.hook_out", "attn_out": "blocks.{i}.attn.hook_out", "mlp_out": "blocks.{i}.mlp.hook_out", + "q": "blocks.{i}.attn.hook_q", + "k": "blocks.{i}.attn.hook_k", + "v": "blocks.{i}.attn.hook_v", + "z": "blocks.{i}.attn.hook_z", + "pattern": "blocks.{i}.attn.hook_pattern", +} +# pattern is read from the forward's output_attentions — nothing to write back, so it's +# capture-only (like the derived resid_mid). +INTERVENEABLE_KINDS = frozenset( + {"resid_pre", "attn_out", "mlp_out", "resid_post", "q", "k", "v", "z"} +) + +# Rank of each kind's batchless wire array; the driver unsqueezes exactly one batch dim. +WIRE_BATCHLESS_NDIM = { + **{kind: 2 for kind in ALL_KINDS}, # (seq, d_model) + "q": 3, + "k": 3, + "v": 3, + "z": 3, # (seq, heads, d_head) + "pattern": 3, # (heads, q_pos, k_pos) } -INTERVENEABLE_KINDS = frozenset({"resid_pre", "attn_out", "mlp_out", "resid_post"}) _SUFFIX_TO_KIND = { "hook_in": "resid_pre", @@ -38,29 +64,35 @@ "hook_out": "resid_post", "attn.hook_out": "attn_out", "mlp.hook_out": "mlp_out", + "attn.hook_q": "q", + "attn.hook_k": "k", + "attn.hook_v": "v", + "attn.hook_z": "z", + "attn.hook_pattern": "pattern", } _BLOCK = re.compile(r"^blocks\.(\d+)\.(.+)$") def supported_hook_points(n_layers: int, kinds: Optional[Iterable[str]] = None) -> frozenset[str]: - """Fireable hook names across all layers. ``kinds=None`` means all boundaries; + """Fireable hook names across all layers. ``kinds=None`` means all *boundary* kinds + (head-split kinds are opt-in — a provider must detect and list them explicitly); pass the provider's detected kinds to gate (e.g. drop ``resid_mid`` for parallel).""" - selected = _KIND_NAMES if kinds is None else {k: _KIND_NAMES[k] for k in kinds} - return frozenset(name.format(i=i) for i in range(n_layers) for name in selected.values()) + selected = ALL_KINDS if kinds is None else kinds + return frozenset(_KIND_NAMES[k].format(i=i) for i in range(n_layers) for k in selected) + + +def all_hook_points(n_layers: int) -> frozenset[str]: + """Every hook name the registry can serve (boundaries + head-split) — the universe a + driver subtracts its supported set from to build ``non_fireable_hook_points``.""" + return supported_hook_points(n_layers, ALL_KINDS | HEAD_KINDS) def nonfireable_hook_points(n_layers: int) -> frozenset[str]: - """Hooks the residual/attn/mlp provider can't fire (head-split, embed, ln_final).""" + """Hooks no provider configuration can fire (embed, ln_final, pre-softmax scores). + Head-split q/k/v/z/pattern are *conditionally* fireable and belong here only when a + model's structural probe gates them — the driver handles that subtraction.""" names = ["embed.hook_out", "ln_final.hook_normalized", "unembed.hook_out"] - for i in range(n_layers): - names += [ - f"blocks.{i}.attn.hook_pattern", - f"blocks.{i}.attn.hook_attn_scores", - f"blocks.{i}.attn.hook_q", - f"blocks.{i}.attn.hook_k", - f"blocks.{i}.attn.hook_v", - f"blocks.{i}.attn.hook_z", - ] + names += [f"blocks.{i}.attn.hook_attn_scores" for i in range(n_layers)] return frozenset(names) diff --git a/transformer_lens/model_bridge/sources/inspect/intervention.py b/transformer_lens/model_bridge/sources/inspect/intervention.py index fcda9bc80..0cd1ee9dc 100644 --- a/transformer_lens/model_bridge/sources/inspect/intervention.py +++ b/transformer_lens/model_bridge/sources/inspect/intervention.py @@ -47,7 +47,7 @@ def build_interventions( if resolved is None or resolved[1] not in hooks.INTERVENEABLE_KINDS: raise ValueError( f"Cannot intervene on {hook_name!r}: capture-only " - f"(intervene on resid_pre/attn_out/mlp_out/resid_post instead)." + f"(intervene on resid_pre/attn_out/mlp_out/resid_post or attn q/k/v/z instead)." ) if op == "scale" and "factor" not in spec: raise ValueError(f"Intervention {hook_name!r}: op='scale' requires 'factor' (float).") @@ -55,6 +55,15 @@ def build_interventions( raise ValueError( f"Intervention {hook_name!r}: op={op!r} requires 'value' (scalar or width-shaped)." ) + pos = spec.get("pos") + if pos is not None and not ( + isinstance(pos, int) + or (isinstance(pos, (list, tuple)) and all(isinstance(p, int) for p in pos)) + ): + raise ValueError( + f"Intervention {hook_name!r}: 'pos' must be an int or list of ints " + f"(sequence positions to patch); got {pos!r}." + ) layer, kind = resolved out[hooks.wire_key(layer, kind)] = dict(spec) return out diff --git a/transformer_lens/model_bridge/sources/inspect/source.py b/transformer_lens/model_bridge/sources/inspect/source.py index b23f565c8..656d31558 100644 --- a/transformer_lens/model_bridge/sources/inspect/source.py +++ b/transformer_lens/model_bridge/sources/inspect/source.py @@ -43,12 +43,16 @@ def boot_inspect( Fireable hooks (``tl_bridge``, TransformerBridge-native names): ``blocks.{i}.hook_in`` (resid_pre) / ``ln2.hook_in`` (resid_mid) / - ``hook_out`` (resid_post) / ``attn.hook_out`` / ``mlp.hook_out``. The provider runs - a structural self-check per model and gates any boundary it can't serve faithfully: - ``resid_mid`` for parallel-residual or norm-variant blocks, and ``attn_out``/ - ``mlp_out`` when their submodule isn't locatable (it warns when it gates one). - Head-split hooks (q/k/v/z, pattern), ``embed``, and ``ln_final`` are always - non-fireable — use ``boot_transformers()`` for those. + ``hook_out`` (resid_post) / ``attn.hook_out`` / ``mlp.hook_out``, plus the head-split + attention hooks where the structural probe finds them: ``attn.hook_q/k/v`` (pre-RoPE + projection outputs; separate-projection archs only — fused qkv gates them), + ``attn.hook_z`` (out-projection input), and ``attn.hook_pattern`` (post-softmax, + capture-only, eager attention required). The provider runs a structural self-check per + model and gates any boundary it can't serve faithfully: ``resid_mid`` for + parallel-residual or norm-variant blocks, ``attn_out``/``mlp_out`` when their submodule + isn't locatable (it warns when it gates one). ``embed``, ``ln_final``, and + ``attn.hook_attn_scores`` are always non-fireable — use ``boot_transformers()`` for + those. For parity with ``boot_transformers`` the provider loads with the same dtype (fp32 by default) and eager attention. Full-sequence logits ride on ``return_logits=True`` (the diff --git a/transformer_lens/model_bridge/sources/inspect/transformers_provider.py b/transformer_lens/model_bridge/sources/inspect/transformers_provider.py index be8a26cf8..74ca10b51 100644 --- a/transformer_lens/model_bridge/sources/inspect/transformers_provider.py +++ b/transformer_lens/model_bridge/sources/inspect/transformers_provider.py @@ -29,7 +29,7 @@ modelapi, ) -from . import wire +from . import hooks, wire from ._provider_base import _InspectModelAPIBase, _parse_tool_calls, _require_served # NOT "transformer_lens" — inspect_ai ships a built-in provider by that name (the @@ -48,6 +48,17 @@ # Attn/MLP submodule names within a block, by family. _ATTN_ATTRS = ("self_attn", "attn", "attention", "self_attention") # self_attention: Falcon _MLP_ATTRS = ("mlp", "feed_forward") +# fc-split blocks (OPT/XGLM) have no mlp container — fc1/fc2 sit on the block directly: +# mlp_in boundary = fc1's input, mlp_out boundary = fc2's output. +_MLP_SPLIT_ATTRS = ("fc1", "fc2") +# Separate q/k/v projections (Llama/Mistral/Qwen/OPT-family). Fused-qkv archs +# (GPT-2 c_attn, Falcon/GPTNeoX query_key_value) gate q/k/v — their packed layouts +# vary per family, so slicing them is per-arch work we don't hand-maintain here. +_Q_PROJ_ATTRS = ("q_proj", "query") +_K_PROJ_ATTRS = ("k_proj", "key") +_V_PROJ_ATTRS = ("v_proj", "value") +# Attention out-projection; its input is z (works for fused-qkv archs too). +_O_PROJ_ATTRS = ("o_proj", "out_proj", "dense", "c_proj") @modelapi(name=PROVIDER_NAME) @@ -85,10 +96,24 @@ def __init__( ) self._tokenizer = AutoTokenizer.from_pretrained(model_name) self._layers = _locate_layers(self._hf) - self._kinds, self._capability_note = _detect_capabilities(self._hf, self._layers) + # Head-split reshape geometry; None per-field when the config lacks it (head + # kinds are then gated by _detect_capabilities' width checks). + self._geometry = _attn_geometry(self._hf.config) + self._kinds, self._capability_note = _detect_capabilities( + self._hf, self._layers, self._geometry + ) # Per-turn capture during plain generation (agent rollouts): every _generate_eval # stashes these hooks in ModelOutput.metadata. Gated by the structural self-check. self._eval_capture = self._parse_eval_capture(model_args) + for key in self._eval_capture: + if key.endswith(":pattern"): + # pattern rides the forward's output_attentions, which the eval path's + # hf.generate doesn't thread — reject rather than silently omit. + raise ValueError( + "Per-turn eval capture of attn.hook_pattern is not supported (it needs " + "output_attentions on the forward). Use the driver path " + "(bridge.run_with_cache) for pattern capture." + ) def _generate_capture(self, input: Any, extra_args: Mapping[str, Any], config: GenerateConfig): """TL-driven single forward: capture residual/attn/mlp boundaries + full logits.""" @@ -101,9 +126,22 @@ def _generate_capture(self, input: Any, extra_args: Mapping[str, Any], config: G _, _, kind = key.partition(":") _require_served(kind, self._kinds, self._capability_note, f"capture {key!r}") interventions: Mapping[str, Any] = extra_args.get("interventions", {}) + # The driver validates before sending, but this direct interface (extra_args) is + # documented — validate here too so a gated/capture-only kind fails loud instead + # of silently no-op'ing when no hook installs for it. + for key in interventions: + _, _, kind = key.partition(":") + _require_served(kind, self._kinds, self._capability_note, f"intervention {key!r}") + if kind not in hooks.INTERVENEABLE_KINDS: + raise ValueError( + f"intervention {key!r}: kind {kind!r} is capture-only " + f"(interveneable: {sorted(hooks.INTERVENEABLE_KINDS)})." + ) want_logits = bool(extra_args.get("return_logits", True)) capture, intervene = _plan(capture_keys, interventions) + # pattern comes from the forward's output_attentions, not a module hook. + pattern_layers = [layer for layer, kinds in capture.items() if "pattern" in kinds] raw: dict[tuple[int, str], np.ndarray] = {} call_id = uuid.uuid4().hex token = _current_call_id.set(call_id) @@ -111,12 +149,19 @@ def _generate_capture(self, input: Any, extra_args: Mapping[str, Any], config: G try: with torch.no_grad(): ids = torch.tensor([list(input_ids)], device=self._device) - logits = self._hf(ids).logits # (1, seq, vocab) + outputs = self._hf(ids, output_attentions=bool(pattern_layers)) + logits = outputs.logits # (1, seq, vocab) finally: for handle in handles: handle.remove() _current_call_id.reset(token) + attentions = getattr(outputs, "attentions", None) + for layer in pattern_layers: + attn_l = attentions[layer] if attentions is not None else None + if attn_l is not None: # None → driver's missing-hook warning handles it + raw[(layer, "pattern")] = attn_l[0].detach().float().cpu().numpy() + captured = _assemble(raw, capture_keys) metadata: dict[str, Any] = {"activations": wire.encode_activations(captured)} if want_logits: @@ -254,7 +299,7 @@ def _install_hooks(self, capture, intervene, raw, call_id: str) -> list: if not cap_kinds and not iv_kinds: continue attn = _first_attr(block, _ATTN_ATTRS) - mlp = _first_attr(block, _MLP_ATTRS) + _, mlp_out_mod = _locate_mlp(block) if "resid_pre" in cap_kinds or "resid_pre" in iv_kinds: handles.append( block.register_forward_pre_hook( @@ -277,9 +322,13 @@ def _install_hooks(self, capture, intervene, raw, call_id: str) -> list: ) ) ) - if mlp is not None and ("mlp_out" in cap_kinds or "mlp_out" in iv_kinds): + if attn is not None: + handles.extend( + self._install_head_hooks(layer, attn, cap_kinds, iv_kinds, raw, call_id) + ) + if mlp_out_mod is not None and ("mlp_out" in cap_kinds or "mlp_out" in iv_kinds): handles.append( - mlp.register_forward_hook( + mlp_out_mod.register_forward_hook( _out_hook( layer, "mlp_out", @@ -305,6 +354,58 @@ def _install_hooks(self, capture, intervene, raw, call_id: str) -> list: ) return handles + def _install_head_hooks(self, layer, attn, cap_kinds, iv_kinds, raw, call_id: str) -> list: + """Hooks for the head-split kinds: q/k/v on their projection outputs, z on the + out-projection's input. pattern isn't hooked (it rides output_attentions). + + Interventions apply to the module's natural *flat* tensor ``(..., seq, + heads·d_head)`` — a spec ``value`` is scalar, ``(heads·d_head,)``, or per-position + ``(len(pos), heads·d_head)`` (flatten a captured head-split tensor to build one). + Captures are emitted head-split ``(seq, heads, d_head)`` to match the bridge's + ``hook_q/k/v/z``. + """ + handles: list = [] + d_head = self._geometry[2] + if d_head is None: # geometry undetectable → head kinds were gated at detection + return handles + host = _projection_host(attn) + for kind, attrs in (("q", _Q_PROJ_ATTRS), ("k", _K_PROJ_ATTRS), ("v", _V_PROJ_ATTRS)): + if kind not in cap_kinds and kind not in iv_kinds: + continue + proj = _first_attr(host, attrs) + if proj is None: + # Detection ran on layers[0]; a later layer missing the projection would + # silently skip a validated intervention — fail loud (capture-only misses + # surface through the driver's missing-hook warning instead). + if kind in iv_kinds: + raise RuntimeError( + f"Intervention on blocks.{layer}.attn.hook_{kind} cannot be applied: " + f"layer {layer} has no {kind} projection (heterogeneous layers)." + ) + continue + handles.append( + proj.register_forward_hook( + _proj_hook( + layer, kind, d_head, kind in cap_kinds, iv_kinds.get(kind), raw, call_id + ) + ) + ) + if "z" in cap_kinds or "z" in iv_kinds: + o_proj = _first_attr(host, _O_PROJ_ATTRS) + if o_proj is None and "z" in iv_kinds: + raise RuntimeError( + f"Intervention on blocks.{layer}.attn.hook_z cannot be applied: layer " + f"{layer} has no out-projection (heterogeneous layers)." + ) + if o_proj is not None: + handles.append( + o_proj.register_forward_pre_hook( + _zin_hook(layer, d_head, "z" in cap_kinds, iv_kinds.get("z"), raw, call_id), + with_kwargs=True, + ) + ) + return handles + def _plan(capture_keys, interventions): """Resolve wire keys → per-layer kinds to capture (resid_mid needs pre+attn) and intervene.""" @@ -370,7 +471,10 @@ def hook(_module, _inputs, output): if spec is not None: hidden = _apply_affine(hidden, spec) if want_capture and (layer, kind) not in raw: - raw[(layer, kind)] = hidden[0].detach().float().cpu().numpy() + # OPT-style blocks flatten the FFN to (batch·seq, d) — with batch_size=1 + # that IS (seq, d) already; only 3-D (batch, seq, d) needs the batch strip. + flat = hidden if hidden.ndim == 2 else hidden[0] + raw[(layer, kind)] = flat.detach().float().cpu().numpy() if spec is None: return None return (hidden, *output[1:]) if is_tuple else hidden @@ -378,45 +482,114 @@ def hook(_module, _inputs, output): return hook -def _apply_affine(t: torch.Tensor, spec: Mapping[str, Any]) -> torch.Tensor: - """suppress→0, scale→·factor, add→+value, set→value (value scalar or width-shaped).""" +def _proj_hook(layer, kind, d_head, want_capture, spec, raw, call_id): + """q/k/v projection output: affine on the flat ``(..., seq, heads·d_head)`` tensor, + captured head-split ``(seq, heads, d_head)`` (pre-RoPE — matches the bridge's + ``hook_q``/``hook_k``/``hook_v``). Mutations feed the downstream attention math.""" + + def hook(_module, _inputs, output): + if _current_call_id.get() != call_id: + return None # different concurrent call's hook + hidden = output + if spec is not None: + hidden = _apply_affine(hidden, spec) + if want_capture and (layer, kind) not in raw: + flat = hidden[0].detach().float().cpu() + raw[(layer, kind)] = flat.reshape(flat.shape[0], -1, d_head).numpy() + return hidden if spec is not None else None + + return hook + + +def _zin_hook(layer, d_head, want_capture, spec, raw, call_id): + """z — the out-projection's *input* (attention-weighted values, all heads): pre-hook, + affine on the flat tensor, captured head-split to match the bridge's ``hook_z``.""" + + def hook(_module, args, kwargs): + if _current_call_id.get() != call_id: + return None # different concurrent call's hook + z = args[0] + if spec is not None: + z = _apply_affine(z, spec) + if want_capture and (layer, "z") not in raw: + flat = z[0].detach().float().cpu() + raw[(layer, "z")] = flat.reshape(flat.shape[0], -1, d_head).numpy() + if spec is None: + return None + return (z, *args[1:]), kwargs + + return hook + + +def _affine_op(sub: torch.Tensor, spec: Mapping[str, Any]) -> torch.Tensor: + """One affine op: suppress→0, scale→·factor, add→+value, set→value. ``value`` broadcasts + (scalar, width-shaped, or per-position ``(n_pos, width)``).""" op = spec["op"] if op == "suppress": - return torch.zeros_like(t) + return torch.zeros_like(sub) if op == "scale": - return t * float(spec["factor"]) - value = torch.as_tensor(spec["value"], dtype=t.dtype, device=t.device) + return sub * float(spec["factor"]) + value = torch.as_tensor(spec["value"], dtype=sub.dtype, device=sub.device) if op == "add": - return t + value - return torch.zeros_like(t) + value # set + return sub + value + return torch.zeros_like(sub) + value # set + + +def _apply_affine(t: torch.Tensor, spec: Mapping[str, Any]) -> torch.Tensor: + """Affine intervention on a captured tensor ``(..., seq, width)``. + Without ``pos`` the op spans every position (the original width-broadcast form). With + ``pos`` (an int or list of sequence indices) it touches only those positions — the + activation-patching primitive — and ``value`` may be per-position ``(len(pos), width)`` + to transplant a captured activation (path/causal tracing) rather than a single vector. + """ + pos = spec.get("pos") + if pos is None: + return _affine_op(t, spec) + idx = [pos] if isinstance(pos, int) else list(pos) + out = t.clone() + out[..., idx, :] = _affine_op(t[..., idx, :], spec) + return out -def _detect_capabilities(model: Any, layers: Any) -> tuple[frozenset, str]: - """Structural self-check: which boundary kinds this model can serve faithfully. + +def _detect_capabilities( + model: Any, layers: Any, geometry: tuple[Any, Any, Any] +) -> tuple[frozenset, str]: + """Structural self-check: which kinds this model can serve faithfully. resid_pre/resid_post are the block in/out (always); attn_out/mlp_out need their submodules locatable; resid_mid is gated unless its derivation holds (see - :func:`_resid_mid_derivable`). Returns (kinds, note); note explains any gating, '' if none. + :func:`_resid_mid_derivable`). Head-split kinds: q/k/v need separate projections whose + widths match ``heads·d_head``; z needs an out-projection of in-width ``n_heads·d_head``; + pattern needs eager attention (output_attentions is a no-op under sdpa/flash). + Returns (kinds, note); note explains any gating, '' if none. """ block = layers[0] attn = _first_attr(block, _ATTN_ATTRS) - mlp = _first_attr(block, _MLP_ATTRS) + mlp_in_mod, mlp_out_mod = _locate_mlp(block) kinds = {"resid_pre", "resid_post"} gated = [] if attn is not None: kinds.add("attn_out") else: gated.append("attn_out (no attention submodule found)") - if mlp is not None: + if mlp_out_mod is not None: kinds.add("mlp_out") else: gated.append("mlp_out (no MLP submodule found)") - if attn is not None and mlp is not None and _resid_mid_derivable(model, block, attn, mlp): + if ( + attn is not None + and mlp_out_mod is not None + and _resid_mid_derivable(model, block, attn, mlp_in_mod, mlp_out_mod) + ): kinds.add("resid_mid") else: gated.append( "resid_mid (resid_pre + attn_out doesn't hold — parallel or norm-variant block)" ) + head_kinds, head_gated = _detect_head_capabilities(model, attn, geometry) + kinds |= head_kinds + gated += head_gated note = ( "" if not gated @@ -427,11 +600,85 @@ def _detect_capabilities(model: Any, layers: Any) -> tuple[frozenset, str]: return frozenset(kinds), note -def _resid_mid_derivable(model: Any, block: Any, attn: Any, mlp: Any) -> bool: +def _detect_head_capabilities( + model: Any, attn: Any, geometry: tuple[Any, Any, Any] +) -> tuple[set, list]: + """Head-split kinds this model serves: q/k/v iff separate projections with the + expected widths, z iff the out-projection's in-width is ``n_heads·d_head``, pattern + iff attention runs eager (otherwise ``output_attentions`` returns None/garbage).""" + n_heads, n_kv_heads, d_head = geometry + kinds: set = set() + gated: list = [] + if attn is None or d_head is None: + gated.append("q/k/v/z/pattern (no attention submodule or head geometry in config)") + return kinds, gated + + host = _projection_host(attn) + q = _first_attr(host, _Q_PROJ_ATTRS) + k = _first_attr(host, _K_PROJ_ATTRS) + v = _first_attr(host, _V_PROJ_ATTRS) + expected = {"q": n_heads * d_head, "k": n_kv_heads * d_head, "v": n_kv_heads * d_head} + if all( + proj is not None and _out_width(proj) == expected[kind] + for kind, proj in (("q", q), ("k", k), ("v", v)) + ): + kinds |= {"q", "k", "v"} + else: + gated.append("q/k/v (fused or nonstandard qkv projections)") + + o_proj = _first_attr(host, _O_PROJ_ATTRS) + if o_proj is not None and _in_width(o_proj) == n_heads * d_head: + kinds.add("z") + else: + gated.append("z (out-projection missing or nonstandard width)") + + if getattr(model.config, "_attn_implementation", "eager") == "eager": + kinds.add("pattern") + else: + gated.append("pattern (attention implementation is not eager)") + return kinds, gated + + +def _out_width(module: Any) -> Any: + """Output width of a projection: ``nn.Linear.out_features`` or GPT-2 ``Conv1D.nf``.""" + out = getattr(module, "out_features", None) + if out is not None: + return int(out) + nf = getattr(module, "nf", None) # transformers Conv1D + return int(nf) if nf is not None else None + + +def _in_width(module: Any) -> Any: + """Input width of a projection: ``nn.Linear.in_features`` or Conv1D ``weight.shape[0]``.""" + in_f = getattr(module, "in_features", None) + if in_f is not None: + return int(in_f) + if getattr(module, "nf", None) is not None and hasattr(module, "weight"): + return int(module.weight.shape[0]) # Conv1D stores weight (in, out) + return None + + +def _attn_geometry(config: Any) -> tuple[Any, Any, Any]: + """(n_heads, n_kv_heads, d_head) from an HF config; (None, None, None) if underivable.""" + n_heads = getattr(config, "num_attention_heads", None) or getattr(config, "n_head", None) + hidden = getattr(config, "hidden_size", None) or getattr(config, "n_embd", None) + if n_heads is None: + return None, None, None + n_kv = getattr(config, "num_key_value_heads", None) or n_heads + d_head = getattr(config, "head_dim", None) + if d_head is None and hidden is not None: + d_head = hidden // n_heads + return (int(n_heads), int(n_kv), int(d_head) if d_head is not None else None) + + +def _resid_mid_derivable( + model: Any, block: Any, attn: Any, mlp_in_mod: Any, mlp_out_mod: Any +) -> bool: """True iff ``resid_mid = resid_pre + attn_out`` holds, via two tiny probe forwards. Requires both the linear identity ``resid_post = resid_pre + attn_out + mlp_out`` (broken by post-norm/multiplier blocks — Gemma2/OLMo2/Granite) and attn feeding mlp (broken by - parallel blocks — GPTNeoX/GPT-J, where mlp reads resid_pre directly).""" + parallel blocks — GPTNeoX/GPT-J, where mlp reads resid_pre directly). ``mlp_in_mod``/ + ``mlp_out_mod`` are the same module for container archs, (fc1, fc2) for fc-split.""" cap: dict[str, Any] = {} def grab(key: str): # type: ignore[no-untyped-def] @@ -463,8 +710,8 @@ def perturb_attn(_m: Any, _i: Any, out: Any): # type: ignore[no-untyped-def] handles = [ block.register_forward_pre_hook(grab_in("resid_pre"), with_kwargs=True), attn.register_forward_hook(grab("attn_out")), - mlp.register_forward_hook(grab("mlp_out")), - mlp.register_forward_pre_hook(grab_in("mlp_in"), with_kwargs=True), + mlp_out_mod.register_forward_hook(grab("mlp_out")), + mlp_in_mod.register_forward_pre_hook(grab_in("mlp_in"), with_kwargs=True), block.register_forward_hook(grab("resid_post")), ] model(ids) @@ -472,7 +719,7 @@ def perturb_attn(_m: Any, _i: Any, out: Any): # type: ignore[no-untyped-def] h.remove() mlp_in_clean = cap.pop("mlp_in", None) handles = [ - mlp.register_forward_pre_hook(grab_in("mlp_in"), with_kwargs=True), + mlp_in_mod.register_forward_pre_hook(grab_in("mlp_in"), with_kwargs=True), attn.register_forward_hook(perturb_attn), ] model(ids) @@ -490,7 +737,13 @@ def perturb_attn(_m: Any, _i: Any, out: Any): # type: ignore[no-untyped-def] return False if mlp_in_clean is None or mlp_in_perturbed is None: return False - if not (rp.shape == ao.shape == mo.shape == rpost.shape == mlp_in_clean.shape): + # OPT-style blocks flatten the FFN to (batch·seq, d); with the probe's batch=1 that + # is a pure reshape of the block-level (1, seq, d) — normalize before comparing. + if mo.shape != rpost.shape and mo.numel() == rpost.numel(): + mo = mo.reshape(rpost.shape) + if not (rp.shape == ao.shape == mo.shape == rpost.shape): + return False + if mlp_in_clean.shape != mlp_in_perturbed.shape or mlp_in_clean.numel() != rp.numel(): return False # (1) sub-block outputs add to the residual without intervening norm/scale. identity = (rpost - rp - ao - mo).abs().max().item() @@ -514,6 +767,33 @@ def _locate_layers(model: Any) -> Any: ) +def _locate_mlp(block: Any) -> tuple[Any, Any]: + """The modules bounding the MLP: ``(in_module, out_module)`` — the mlp_in boundary is + in_module's input, mlp_out is out_module's output. Container archs return the mlp + module twice; fc-split blocks (OPT/XGLM: fc1/fc2 directly on the block) return + ``(fc1, fc2)``. ``(None, None)`` when neither layout is found (mlp_out gated).""" + mlp = _first_attr(block, _MLP_ATTRS) + if mlp is not None: + return mlp, mlp + fc1, fc2 = (getattr(block, name, None) for name in _MLP_SPLIT_ATTRS) + if fc1 is not None and fc2 is not None: + return fc1, fc2 + return None, None + + +def _projection_host(attn: Any) -> Any: + """The module whose direct attrs are the q/k/v/out projections. Usually ``attn`` + itself; GPT-Neo-style blocks wrap the real attention (with its standard q_proj/ + out_proj) one level down at ``attn.attention``. Descend only when the located module + has neither a q- nor an out-projection — GPTNeoX's ``block.attention`` (fused + query_key_value + dense) has ``dense`` directly, so it never descends.""" + if _first_attr(attn, _Q_PROJ_ATTRS) is None and _first_attr(attn, _O_PROJ_ATTRS) is None: + inner = getattr(attn, "attention", None) + if inner is not None: + return inner + return attn + + def _first_attr(obj: Any, names: tuple[str, ...]) -> Any: for name in names: found = getattr(obj, name, None) diff --git a/transformer_lens/model_bridge/sources/vllm/driver.py b/transformer_lens/model_bridge/sources/vllm/driver.py index e78aa47ff..91e0270dc 100644 --- a/transformer_lens/model_bridge/sources/vllm/driver.py +++ b/transformer_lens/model_bridge/sources/vllm/driver.py @@ -20,10 +20,16 @@ class VLLMDriver(DriverBase): """Driver wrapping a vLLM ``LLM``; captures via ``collective_rpc``.""" - # vLLM owns the model in a worker — no torch surface (parameters/state_dict/grads). - _supported_features = frozenset() - # Logits synthesized for the final position only (sampler bypass). - provides_sequence_logits = False + # vLLM owns the model in a worker — no torch module surface (parameters/state_dict/ + # grads). Named-weight reads ARE served: get_param() returns CPU clones via the + # tl_get_param RPC (what logit reconstruction and direct-logit-attribution use). + _supported_features = frozenset({"weight_access"}) + # Full-sequence logits reconstructed host-side (ln_final @ lm_head.weight.T); vLLM's + # sampler only hands back the final position, so the driver rebuilds the rest. + provides_sequence_logits = True + + # Post-weight final-norm capture that lm_head consumes — reconstruction reads it. + _LN_FINAL = "ln_final.hook_normalized" def __init__( self, @@ -34,11 +40,15 @@ def __init__( hf_config: Any, max_num_batched_tokens: int, enable_batching: bool = False, + enable_position_interventions: bool = False, ) -> None: super().__init__(adapter.cfg, tokenizer) self._llm = llm self._max_num_batched_tokens = max_num_batched_tokens self._enable_batching = enable_batching + # Position-scoped 'pos' interventions need (max_n, width) affine buffers, + # allocated at boot only when this is set (see plugin.patched_load_model). + self._enable_position_interventions = enable_position_interventions # Logprobs per forward = real vocab (boot's max_logprobs). d_vocab can be # padded larger, which vLLM would reject; the logits tensor stays d_vocab. self._n_logprobs = int(getattr(hf_config, "vocab_size", self.bridge_config.d_vocab)) @@ -46,9 +56,16 @@ def __init__( self.supported_hook_points = frozenset(overlay.capture_specs(hf_config).keys()) n_layers = getattr(hf_config, "num_hidden_layers", 0) + if not isinstance(n_layers, int) or n_layers <= 0: + # A raw "{i}" template would land unexpanded in non_fireable_hook_points — + # a broken config should fail at boot, not surface as a garbled hook name. + raise ValueError( + f"VLLMDriver: hf_config.num_hidden_layers={n_layers!r} — expected a " + "positive int; the config is missing or malformed." + ) nonfiring: list[str] = [] for tmpl in overlay.nonfiring_hooks(): - if "{i}" in tmpl and isinstance(n_layers, int) and n_layers > 0: + if "{i}" in tmpl: nonfiring.extend(tmpl.replace("{i}", str(i)) for i in range(n_layers)) else: nonfiring.append(tmpl) @@ -73,9 +90,11 @@ def forward( ) intervene_specs = self._validate_interventions(intervene or {}) - # Restrict the GPU→CPU read to these hooks (None = all). run_with_cache - # doesn't derive this from names_filter yet — only explicit forward(capture=). - names = list(capture) or None + # capture is authoritative — the bridge sends exactly the hooked names, so () + # means "capture nothing" and a plain forward(tokens) skips the GPU→CPU copy + # entirely. (The worker's None-means-all convention is never triggered from here; + # an empty tuple used to collapse to None and silently copy every buffer.) + names = list(capture) if self._enable_batching: return self._forward_batched(input_ids, intervene_specs, return_logits, names) @@ -88,6 +107,9 @@ def forward( f"{self._max_num_batched_tokens}; raise the boot_vllm kwarg or " "shorten the prompt." ) + # 'pos'-scoped edits target affine-buffer rows, but the compiled hook only reads + # rows [0, len(ids_list)); a pos past the prompt length would be a silent no-op. + self._reject_pos_beyond_seq(intervene_specs, len(ids_list)) from vllm import SamplingParams from vllm.inputs import TokensPrompt @@ -110,15 +132,31 @@ def forward( ) n_tokens = len(ids_list) - # collective_rpc returns one result per worker; single-rank, so [0]. - worker_captures = self._llm.collective_rpc("tl_read_captures", args=([n_tokens], names))[0] - # Add batch dim: vLLM hands back (n_tokens, width); bridge expects (1, n_tokens, width). - captured = {name: t.unsqueeze(0) for name, t in worker_captures.items()} + # Reconstructing logits needs ln_final; force it into the read even if the caller + # didn't request it (dropped from `captured` below so the surface stays as asked). + read_names = names + if return_logits and self._LN_FINAL not in names: + read_names = names + [self._LN_FINAL] + # collective_rpc returns one result per worker; single-rank, so [0]. Nothing to + # read (no captures, logits off) → skip the crossing altogether. + worker_captures = ( + self._llm.collective_rpc("tl_read_captures", args=([n_tokens], read_names))[0] + if read_names + else {} + ) logits: torch.Tensor | None = None if return_logits: - logits = self._synthesize_logits(outputs[0], n_tokens, self.bridge_config.d_vocab) + recon = self._reconstruct_logits(worker_captures.get(self._LN_FINAL)) + # Fall back to final-position log-probs if the unembedding isn't fetchable. + logits = ( + recon.unsqueeze(0) + if recon is not None + else self._synthesize_logits(outputs[0], n_tokens, self.bridge_config.d_vocab) + ) + # Add batch dim; expose only the caller's requested hooks (drop the forced ln_final). + captured = {name: t.unsqueeze(0) for name, t in worker_captures.items() if name in names} return ForwardResult(logits=logits, captured=captured, raw_output=outputs[0]) def _forward_batched( @@ -126,13 +164,13 @@ def _forward_batched( input_ids: TensorLike, intervene_specs: dict, return_logits: bool, - names: list[str] | None = None, + names: list[str], ) -> ForwardResult: """Eager batched path: per-request capture, right-padded to (B, S, W). No per-prompt length gate — chunked prefill accumulates long prompts - across forwards. Interventions are global across the batch. ``names`` - restricts the returned hooks (``None`` = all). + across forwards. Interventions are global across the batch. ``names`` is + authoritative: exactly the hooks to return (empty = none). """ from vllm import SamplingParams from vllm.inputs import TokensPrompt @@ -155,14 +193,35 @@ def _forward_batched( ), ) + # Force ln_final into the read so logits can be reconstructed (dropped below if + # the caller didn't ask for it). + read_names = names + if return_logits and self._LN_FINAL not in names: + read_names = names + [self._LN_FINAL] # Keyed by req_id (no guaranteed order) — _assemble_padded joins to slot - # k via outputs[k].request_id, not by position. - worker_captures = self._llm.collective_rpc("tl_read_batched_captures", args=(names,))[0] - captured = self._assemble_padded(outputs, worker_captures, prompt_lens) + # k via outputs[k].request_id, not by position. Empty read → skip both the + # crossing AND the join (_assemble_padded requires one worker key per request, + # so it can't run on an empty dict — mirror the single path's empty captured). + if read_names: + worker_captures = self._llm.collective_rpc( + "tl_read_batched_captures", args=(read_names,) + )[0] + captured = self._assemble_padded(outputs, worker_captures, prompt_lens) + else: + captured = {} logits: torch.Tensor | None = None if return_logits: - logits = self._synthesize_logits_batched(outputs, prompt_lens, d_vocab) + recon = self._reconstruct_logits( + captured.get(self._LN_FINAL) + ) # (batch, max_seq, d_vocab) + logits = ( + recon + if recon is not None + else self._synthesize_logits_batched(outputs, prompt_lens, d_vocab) + ) + if self._LN_FINAL not in names: + captured.pop(self._LN_FINAL, None) return ForwardResult(logits=logits, captured=captured, raw_output=outputs) @@ -245,6 +304,37 @@ def get_param(self, dotted_name: str) -> torch.Tensor | None: return None return self._llm.collective_rpc("tl_get_param", args=(dotted_name,))[0] + def _reconstruct_logits(self, ln_final: Any) -> torch.Tensor | None: + """Rebuild real logits from the captured post-weight ln_final: + ``ln_final @ lm_head.weight.T`` (+ bias, + Gemma-family tanh soft-cap). + + vLLM's ``ln_final.hook_normalized`` is the POST-weight RMSNorm value lm_head + consumes (verified empirically: it equals HF's pre-weight value times the norm + weight), so no un-fold is needed. Accepts any ``(..., d_model)`` tensor and returns + ``(..., d_vocab)`` on CPU. ``None`` if ln_final wasn't captured or no unembedding + weight is fetchable — the caller then falls back to the sampler's log-probs. + """ + if ln_final is None: + return None + weight = self.get_param("lm_head.weight") + if weight is None: # tied embeddings expose no separate lm_head + weight = self.get_param("model.embed_tokens.weight") + if weight is None: + return None + lf = ln_final.to(device=weight.device, dtype=torch.float32) + logits = lf @ weight.to(torch.float32).T + bias = self.get_param("lm_head.bias") + if bias is not None: + logits = logits + bias.to(device=logits.device, dtype=torch.float32) + cap = getattr(self.bridge_config, "output_logits_soft_cap", None) + if cap is not None and cap > 0: # Gemma-family cap; -1.0 is the "disabled" sentinel + logits = float(cap) * torch.tanh(logits / float(cap)) + d_vocab = int(self.bridge_config.d_vocab) + if logits.shape[-1] < d_vocab: # pad the padded-vocab tail (never predicted) + pad = logits.new_full((*logits.shape[:-1], d_vocab - logits.shape[-1]), float("-inf")) + logits = torch.cat([logits, pad], dim=-1) + return logits.cpu() + def close(self) -> None: # Detach hooks before dropping the LLM so they don't stay registered on # worker modules for the life of the process (long-running notebooks). @@ -334,9 +424,62 @@ def _validate_interventions(self, intervene: Mapping[str, Any]) -> dict: raise ValueError( f"Cannot intervene on {hook_name!r}: not in supported_hook_points." ) + pos = spec.get("pos") + if pos is not None: + if self._enable_batching: + # The batched/eager path applies ops to the raw tensor, not the + # (max_n, width) affine buffers, so it has no position surface. + raise NotImplementedError( + f"Intervention {hook_name!r}: per-position 'pos' is not supported on the " + "batched/eager path. Boot the compiled path (enable_batching=False) with " + "enable_position_interventions=True." + ) + if not self._enable_position_interventions: + # Default affine buffers are (width,) and broadcast across every position; + # honoring 'pos' needs the (max_n, width) buffers allocated at boot. + raise NotImplementedError( + f"Intervention {hook_name!r}: per-position 'pos' requires " + "boot_vllm(enable_position_interventions=True) (its default affine " + "buffers broadcast across all positions). Use the Inspect/HF backend for " + "position-scoped patching, or drop 'pos' for a whole-sequence edit." + ) + if not ( + isinstance(pos, int) + or (isinstance(pos, (list, tuple)) and all(isinstance(p, int) for p in pos)) + ): + raise ValueError( + f"Intervention {hook_name!r}: 'pos' must be an int or list of ints " + f"(sequence positions to patch); got {pos!r}." + ) + bad = [p for p in ([pos] if isinstance(pos, int) else pos) if p < 0] + if bad: + raise ValueError( + f"Intervention {hook_name!r}: 'pos' must be non-negative; got {bad}." + ) out[hook_name] = dict(spec) return out + @staticmethod + def _reject_pos_beyond_seq(specs: Mapping[str, Any], seq_len: int) -> None: + """Fail loud if a spec's 'pos' targets a row past the actual prompt length. + + The compiled hook applies the affine over rows ``[0, seq_len)`` only, so a ``pos`` + in ``[seq_len, max_num_batched_tokens)`` clears the driver's non-negativity check + and the worker's buffer-capacity check yet is never read — a silent no-op. Bound it + against the real sequence length (known once ``ids_list`` exists) and raise instead. + """ + for hook_name, spec in specs.items(): + pos = spec.get("pos") + if pos is None: + continue + idx = [pos] if isinstance(pos, int) else list(pos) + bad = [p for p in idx if p >= seq_len] + if bad: + raise ValueError( + f"Intervention {hook_name!r}: 'pos' {bad} is beyond the prompt length " + f"{seq_len} (positions are 0-indexed); the edit would be silently ignored." + ) + @staticmethod def _normalize_input_ids(input_ids: Any) -> list: """Coerce input_ids to a flat list[int] for ``TokensPrompt``; batch_size=1 only.""" diff --git a/transformer_lens/model_bridge/sources/vllm/plugin.py b/transformer_lens/model_bridge/sources/vllm/plugin.py index 6b54d4649..51e5d70fb 100644 --- a/transformer_lens/model_bridge/sources/vllm/plugin.py +++ b/transformer_lens/model_bridge/sources/vllm/plugin.py @@ -49,12 +49,14 @@ def configure( max_num_batched_tokens: int, dtype: torch.dtype, enable_batching: bool = False, + enable_position_interventions: bool = False, ) -> None: """Set capture specs, buffer length, dtype, and hook flavor before ``LLM(...)``.""" _config["capture_specs"] = capture_specs _config["max_num_batched_tokens"] = max_num_batched_tokens _config["dtype"] = dtype _config["enable_batching"] = enable_batching + _config["enable_position_interventions"] = enable_position_interventions def register() -> None: @@ -84,6 +86,9 @@ def patched_load_model(self): max_n = _config["max_num_batched_tokens"] dtype = _config["dtype"] enable_batching = _config.get("enable_batching", False) + # Per-position affine buffers are (max_n, width) instead of (width,), so each + # sequence row can carry a distinct scale/bias (position-scoped patching). + per_position = _config.get("enable_position_interventions", False) device = next(self.model_runner.model.parameters()).device # Detach prior handles before reassigning — vLLM doesn't double-load @@ -124,9 +129,12 @@ def patched_load_model(self): else: capture_buf = torch.zeros(max_n, width, device=device, dtype=dtype) # Affine identity at install. Driver swaps via tl_set_interventions - # to enable suppress/scale/add/set ops between forwards. - scale_buf = torch.ones(width, device=device, dtype=dtype) - bias_buf = torch.zeros(width, device=device, dtype=dtype) + # to enable suppress/scale/add/set ops between forwards. Shape is + # (max_n, width) when position interventions are enabled so each row + # can differ; (width,) otherwise (broadcast across all positions). + affine_shape = (max_n, width) if per_position else (width,) + scale_buf = torch.ones(affine_shape, device=device, dtype=dtype) + bias_buf = torch.zeros(affine_shape, device=device, dtype=dtype) # Default closed — opened explicitly by tl_reset_capture_flags for the # next forward(s) that need to capture. capture_flag = torch.ones(1, device=device, dtype=torch.int64) @@ -142,6 +150,7 @@ def patched_load_model(self): self._tl_fire_counter, capture_flag, materialize=materialize, + per_position=per_position, ) ) self._tl_hook_handles.append(handle) @@ -158,6 +167,7 @@ def _make_capture_hook( capture_flag: torch.Tensor, *, materialize: bool = False, + per_position: bool = False, ): """GPU-only, dynamic-shape-safe affine + first-write-wins capture into pre-allocated buffers. @@ -174,18 +184,30 @@ def _make_capture_hook( steps self-copy (no overwrite). Interventions still apply on every forward regardless of the flag — the gate only affects the capture write. + ``per_position`` (compile-time constant): when set, ``scale_buf``/``bias_buf`` are + ``(max_n, width)`` and the affine is row-scoped (``buf.narrow(0, 0, n)``) so a + ``pos``-scoped intervention edits only its rows; otherwise they are ``(width,)`` + and broadcast across every position. + ``fire_counter`` is incremented per call for the fire-once check. """ + def _affine(t: torch.Tensor, n: Any) -> torch.Tensor: + # per_position is a closure constant → torch.compile specializes this branch. + # narrow(0, 0, n) keeps the SymInt dynamic shape (same trick as _gated_capture); + # the (width,) path broadcasts across all rows. + if per_position: + return t * scale_buf.narrow(0, 0, n) + bias_buf.narrow(0, 0, n) + return t * scale_buf + bias_buf + @torch.no_grad() def hook(_module, _inputs, output): fire_counter.add_(1) if materialize and isinstance(output, tuple) and len(output) == 2: hidden, residual = output if isinstance(hidden, torch.Tensor) and isinstance(residual, torch.Tensor): - t = hidden + residual - modified = t * scale_buf + bias_buf - n = t.shape[0] + n = hidden.shape[0] + modified = _affine(hidden + residual, n) _gated_capture(capture_buf, n, modified, capture_flag) # Reconstructs ``modified`` in the next layer's fused norm: exact at # identity, bounded fp16 error under intervention. @@ -201,10 +223,8 @@ def hook(_module, _inputs, output): return None # Affine transform; default scale=1 / bias=0 means identity. Driver # swaps buffer contents to enable interventions. - modified = t * scale_buf + bias_buf - # narrow() keeps the dynamic shape; [:n] gets erased under fake-tensor - # tracing and copy_ then sees the full buffer ("expand s72 -> max_n"). n = t.shape[0] + modified = _affine(t, n) _gated_capture(capture_buf, n, modified, capture_flag) # Gate on isinstance, not truthy tuple_tail — a 1-tuple has an empty tail. if isinstance(output, tuple): diff --git a/transformer_lens/model_bridge/sources/vllm/source.py b/transformer_lens/model_bridge/sources/vllm/source.py index bb6d03d50..6bb635110 100644 --- a/transformer_lens/model_bridge/sources/vllm/source.py +++ b/transformer_lens/model_bridge/sources/vllm/source.py @@ -39,6 +39,7 @@ def boot_vllm( max_model_len: Optional[int] = None, max_num_batched_tokens: int = 2048, enable_batching: bool = False, + enable_position_interventions: bool = False, **vllm_kwargs: Any, ) -> RemoteBridge: """Boot a model via vLLM and wrap it in a :class:`RemoteBridge` via :class:`VLLMDriver`. @@ -86,7 +87,20 @@ def boot_vllm( ``batch_size > 1``) — the throughput path for SAE/probe data collection. Default ``False`` keeps the compile-validated single-prompt path. Batched caches are right-padded with zeros to the longest sequence. + + ``enable_position_interventions`` widens each hook's affine scale/bias buffers + from ``(width,)`` to ``(max_num_batched_tokens, width)`` so an intervention spec + can carry a ``pos`` field (int or list[int]) that scopes the edit to specific + sequence positions — position-scoped activation patching / tensor injection. + Costs ~2× extra resident GPU memory across all hooks (the scale and bias buffers + join the already-``(max_n, width)`` capture buffer), so it is opt-in and defaults + ``False``. Compiled-path only — incompatible with ``enable_batching``. """ + if enable_position_interventions and enable_batching: + raise ValueError( + "enable_position_interventions requires the compiled path and is incompatible " + "with enable_batching=True (the batched/eager path has no affine buffers)." + ) _reject_locked_overrides(vllm_kwargs) from transformers import AutoConfig, AutoTokenizer @@ -104,6 +118,7 @@ def boot_vllm( max_num_batched_tokens=max_num_batched_tokens, dtype=resolved_dtype, enable_batching=enable_batching, + enable_position_interventions=enable_position_interventions, ) plugin.register() @@ -177,6 +192,7 @@ def boot_vllm( hf_config=hf_config, max_num_batched_tokens=max_num_batched_tokens, enable_batching=enable_batching, + enable_position_interventions=enable_position_interventions, ) bridge = RemoteBridge(adapter=adapter, tokenizer=tokenizer, driver=driver) _log_hook_summary(model_name, architecture, driver) diff --git a/transformer_lens/model_bridge/sources/vllm/worker_extension.py b/transformer_lens/model_bridge/sources/vllm/worker_extension.py index 5873be05c..b8cbcca04 100644 --- a/transformer_lens/model_bridge/sources/vllm/worker_extension.py +++ b/transformer_lens/model_bridge/sources/vllm/worker_extension.py @@ -170,37 +170,72 @@ def _apply_op(t: torch.Tensor, spec: Dict[str, Any]) -> torch.Tensor: return torch.zeros_like(t) + value # set +def _write_rows(buf: torch.Tensor, idx: Optional[List[int]], value: Any) -> None: + """In-place write of ``value`` (Python scalar or ``(width,)`` tensor) into ``buf``. + + ``idx is None`` writes the whole buffer — for a 2-D ``(max_n, width)`` buffer that + broadcasts a ``(width,)`` value across every row. ``idx`` (a list of row indices) + writes only those rows of a 2-D buffer. + """ + if idx is None: + if isinstance(value, torch.Tensor): + buf[...] = value # (width,) broadcasts across rows for a 2-D buffer + else: + buf.fill_(value) + else: + buf[idx] = value # advanced-index rows; scalar or (width,) broadcasts + + def _apply_intervention( scale_buf: torch.Tensor, bias_buf: torch.Tensor, spec: Dict[str, Any] ) -> None: - """Translate a spec dict to in-place buffer writes.""" + """Translate a spec dict to in-place buffer writes. + + ``spec['pos']`` (int or list[int]) scopes the edit to those sequence rows and + requires 2-D ``(max_n, width)`` affine buffers (position interventions enabled at + boot). Absent ``pos`` writes the whole buffer, applying to every position. The + caller resets both buffers to identity first, so a ``pos`` edit leaves other rows + untouched. + """ op = spec.get("op") if op not in SUPPORTED_OPS: raise ValueError(f"Unsupported intervention op: {op!r}. Supported: {sorted(SUPPORTED_OPS)}") + pos = spec.get("pos") + idx: Optional[List[int]] = None + if pos is not None: + if scale_buf.ndim != 2: + raise ValueError( + "Per-position 'pos' requires 2-D affine buffers; boot with " + "enable_position_interventions=True." + ) + idx = [pos] if isinstance(pos, int) else list(pos) + max_n = scale_buf.shape[0] + bad = [p for p in idx if p < 0 or p >= max_n] + if bad: + raise ValueError(f"Intervention 'pos' {bad} out of range [0, {max_n}).") + if op == "suppress": - scale_buf.zero_() - bias_buf.zero_() + _write_rows(scale_buf, idx, 0.0) + _write_rows(bias_buf, idx, 0.0) return if op == "scale": - scale_buf.fill_(float(spec["factor"])) - bias_buf.zero_() + _write_rows(scale_buf, idx, float(spec["factor"])) + _write_rows(bias_buf, idx, 0.0) return value = torch.as_tensor(spec["value"], device=bias_buf.device, dtype=bias_buf.dtype) + width = bias_buf.shape[-1] # 0-d broadcasts across width (e.g. "shift all dims by 0.5"); width-shaped # writes element-wise (e.g. SAE steering vector). Anything else is an error. - if value.ndim == 0: - bias_buf.fill_(value.item()) - elif value.shape == bias_buf.shape: - bias_buf.copy_(value) - else: + if value.ndim != 0 and value.shape != (width,): raise ValueError( - f"Intervention 'value' must be a scalar or shape {tuple(bias_buf.shape)}; " + f"Intervention 'value' must be a scalar or shape {(width,)}; " f"got shape {tuple(value.shape)}" ) + _write_rows(bias_buf, idx, value.item() if value.ndim == 0 else value) if op == "add": - scale_buf.fill_(1.0) + _write_rows(scale_buf, idx, 1.0) elif op == "set": - scale_buf.zero_() + _write_rows(scale_buf, idx, 0.0) else: raise RuntimeError( f"op {op!r} is in SUPPORTED_OPS but _apply_intervention has no branch for it."