Coverage for transformer_lens/model_bridge/sources/inspect/transformers_provider.py: 88%

1"""Our HF-transformers ``inspect_ai`` model provider, registered as ``tl_bridge``. 

2 

3The model-runner side of the Inspect driver (this file uses torch; the consuming 

4``InspectDriver`` does not). On ``generate`` it reads the request from 

5``config.extra_body["extra_args"]`` (token ids, which ``<layer>:<kind>`` boundaries 

6to capture, and intervention specs), runs an HF causal LM with forward hooks that 

7capture residual/attn/mlp boundaries and apply affine interventions, and returns a 

8``ModelOutput`` whose ``metadata`` carries the activations (encoded by ``wire``) 

9plus the full-sequence logits. 

10""" 

11from __future__ import annotations 

12 

13import contextvars 

14import uuid 

15from collections import defaultdict 

16from contextlib import contextmanager 

17from typing import Any, Iterator, Mapping 

18 

19import numpy as np 

20import torch 

21from inspect_ai.model import ( 

22 ChatCompletionChoice, 

23 ChatMessageAssistant, 

24 GenerateConfig, 

25 Logprobs, 

26 ModelOutput, 

27 ModelUsage, 

28 StopReason, 

29 modelapi, 

30) 

31 

32from . import hooks, wire 

33from ._provider_base import ( 

34 _InspectModelAPIBase, 

35 _parse_tool_calls, 

36 _require_interveneable, 

37 _require_served, 

38 _warn_unsupported_config, 

39) 

40 

41# NOT "transformer_lens" — inspect_ai ships a built-in provider by that name (the 

42# reverse direction: serving a HookedTransformer as an Inspect model for generation). 

43PROVIDER_NAME = "tl_bridge" 

44 

45# Per-call hook isolation. Capture/intervene hooks consult this contextvar and only fire 

46# for their own call's id — so concurrent inspect_eval samples (each running with their 

47# own contextvars copy via asyncio.to_thread) don't cross-pollute each other's activations. 

48_current_call_id: contextvars.ContextVar[str] = contextvars.ContextVar( 

49 "tl_inspect_call_id", default="" 

50) 

51 

52# Decoder ModuleList by architecture family; each block's output is resid_post. 

53_LAYER_PATHS = ("model.layers", "transformer.h", "gpt_neox.layers", "model.decoder.layers") 

54# Attn/MLP submodule names within a block, by family. 

55_ATTN_ATTRS = ("self_attn", "attn", "attention", "self_attention") # self_attention: Falcon 

56_MLP_ATTRS = ("mlp", "feed_forward") 

57# fc-split blocks (OPT/XGLM) have no mlp container — fc1/fc2 sit on the block directly: 

58# mlp_in boundary = fc1's input, mlp_out boundary = fc2's output. 

59_MLP_SPLIT_ATTRS = ("fc1", "fc2") 

60# Separate q/k/v projections (Llama/Mistral/Qwen/OPT-family). Fused-qkv archs 

61# (GPT-2 c_attn, Falcon/GPTNeoX query_key_value) gate q/k/v — their packed layouts 

62# vary per family, so slicing them is per-arch work we don't hand-maintain here. 

63_Q_PROJ_ATTRS = ("q_proj", "query") 

64_K_PROJ_ATTRS = ("k_proj", "key") 

65_V_PROJ_ATTRS = ("v_proj", "value") 

66# Attention out-projection; its input is z (works for fused-qkv archs too). 

67_O_PROJ_ATTRS = ("o_proj", "out_proj", "dense", "c_proj") 

68 

69 

70@modelapi(name=PROVIDER_NAME) 

71def transformer_lens_provider(): 

72 """Lazy registration hook — returns the provider class on first use.""" 

73 return TransformerLensTransformersModelAPI 

74 

75 

76class TransformerLensTransformersModelAPI(_InspectModelAPIBase): 

77 """HF-backed provider: residual/attn/mlp capture + interventions + full logits. 

78 

79 Inherits generate() dispatch, _messages_to_ids, _logprob_entry, and the per-turn 

80 capture-config validation from :class:`_InspectModelAPIBase`; the HF-specific bits 

81 here are the model load, the forward-hook capture machinery, and the structural probe. 

82 """ 

83 

84 # Real per-position logits via direct HF forward — loss/both via RemoteBridge work. 

85 provides_sequence_logits = True 

86 

87 def __init__( 

88 self, 

89 model_name: str, 

90 base_url: str | None = None, 

91 api_key: str | None = None, 

92 config: GenerateConfig = GenerateConfig(), 

93 **model_args: Any, 

94 ) -> None: 

95 super().__init__(model_name, base_url, api_key, [], config) 

96 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer 

97 

98 from transformer_lens.model_bridge.sources._hf_format import ( 

99 determine_architecture_from_hf_config, 

100 ) 

101 from transformer_lens.model_bridge.sources.transformers.helpers import ( 

102 get_hf_model_class_for_architecture, 

103 ) 

104 

105 self._device = model_args.pop("device", "cpu") 

106 hf_kwargs = model_args.pop("model_kwargs", {}) 

107 # Config-first class selection: vision/seq2seq/masked-LM archs need a different 

108 # AutoModel entry point than AutoModelForCausalLM. 

109 config_kwargs = { 

110 k: hf_kwargs[k] for k in ("token", "trust_remote_code", "revision") if k in hf_kwargs 

111 } 

112 hf_config = AutoConfig.from_pretrained(model_name, **config_kwargs) 

113 try: 

114 model_class = get_hf_model_class_for_architecture( 

115 determine_architecture_from_hf_config(hf_config) 

116 ) 

117 except ValueError: 

118 # Arch unknown to TL's registry — keep loading it as a plain HF causal LM 

119 # (standalone Inspect usage; the structural probe handles the layout). 

120 model_class = AutoModelForCausalLM 

121 self._hf = model_class.from_pretrained(model_name, **hf_kwargs).to(self._device).eval() 

122 self._tokenizer = AutoTokenizer.from_pretrained(model_name) 

123 self._layers = _locate_layers(self._hf) 

124 # Head-split reshape geometry; None per-field when the config lacks it (head 

125 # kinds are then gated by _detect_capabilities' width checks). 

126 self._geometry = _attn_geometry(self._hf.config) 

127 self._kinds, self._capability_note = _detect_capabilities( 

128 self._hf, self._layers, self._geometry 

129 ) 

130 # Per-turn capture during plain generation (agent rollouts): every _generate_eval 

131 # stashes these hooks in ModelOutput.metadata. Gated by the structural self-check. 

132 self._eval_capture = self._parse_eval_capture(model_args) 

133 for key in self._eval_capture: 

134 if key.endswith(":pattern"): 134 ↛ 137line 134 didn't jump to line 137 because the condition on line 134 was never true

135 # pattern rides the forward's output_attentions, which the eval path's 

136 # hf.generate doesn't thread — reject rather than silently omit. 

137 raise ValueError( 

138 "Per-turn eval capture of attn.hook_pattern is not supported (it needs " 

139 "output_attentions on the forward). Use the driver path " 

140 "(bridge.run_with_cache) for pattern capture." 

141 ) 

142 

143 def _generate_capture(self, input: Any, extra_args: Mapping[str, Any], config: GenerateConfig): 

144 """TL-driven single forward: capture residual/attn/mlp boundaries + full logits.""" 

145 input_ids = extra_args.get("input_ids") 

146 if input_ids is None: 

147 input_ids = self._messages_to_ids(input)[0].tolist() 

148 # capture/interventions are keyed by "<layer>:<kind>" (hooks.wire_key). 

149 capture_keys = list(extra_args.get("capture", [])) 

150 for key in capture_keys: 

151 _, _, kind = key.partition(":") 

152 _require_served(kind, self._kinds, self._capability_note, f"capture {key!r}") 

153 interventions: Mapping[str, Any] = extra_args.get("interventions", {}) 

154 # extra_args is a documented surface: gated/capture-only kinds must fail here, 

155 # not silently no-op when no hook installs for them. 

156 for key in interventions: 

157 _, _, kind = key.partition(":") 

158 _require_interveneable( 

159 kind, self._kinds, self._capability_note, f"intervention {key!r}" 

160 ) 

161 want_logits = bool(extra_args.get("return_logits", True)) 

162 

163 capture, intervene = _plan(capture_keys, interventions) 

164 # pattern comes from the forward's output_attentions, not a module hook. 

165 pattern_layers = [layer for layer, kinds in capture.items() if "pattern" in kinds] 

166 raw: dict[tuple[int, str], np.ndarray] = {} 

167 call_id = uuid.uuid4().hex 

168 token = _current_call_id.set(call_id) 

169 handles = self._install_hooks(capture, intervene, raw, call_id) 

170 try: 

171 with torch.no_grad(): 

172 ids = torch.tensor([list(input_ids)], device=self._device) 

173 outputs = self._hf(ids, output_attentions=bool(pattern_layers)) 

174 logits = outputs.logits # (1, seq, vocab) 

175 finally: 

176 for handle in handles: 

177 handle.remove() 

178 _current_call_id.reset(token) 

179 

180 attentions = getattr(outputs, "attentions", None) 

181 for layer in pattern_layers: 

182 attn_l = attentions[layer] if attentions is not None else None 

183 if attn_l is not None: # None → driver's missing-hook warning handles it 183 ↛ 181line 183 didn't jump to line 181 because the condition on line 183 was always true

184 raw[(layer, "pattern")] = attn_l[0].detach().float().cpu().numpy() 

185 

186 captured = _assemble(raw, capture_keys) 

187 metadata: dict[str, Any] = {"activations": wire.encode_activations(captured)} 

188 if want_logits: 

189 metadata["tl_logits"] = wire.encode_array(logits[0].float().cpu().numpy()) 

190 

191 next_id = int(logits[0, -1].argmax()) 

192 logprobs = ( 

193 Logprobs(content=[self._logprob_entry(next_id, logits[0, -1], config.top_logprobs)]) 

194 if config.logprobs 

195 else None 

196 ) 

197 return ModelOutput( 

198 model=self.model_name, 

199 choices=[ 

200 ChatCompletionChoice( 

201 message=ChatMessageAssistant(content=str(self._tokenizer.decode([next_id]))), 

202 stop_reason="stop", 

203 logprobs=logprobs, 

204 ) 

205 ], 

206 metadata=metadata, 

207 ) 

208 

209 def _generate_eval(self, input: Any, config: GenerateConfig, tools: Any): 

210 """Plain Inspect generation: HF generate from the chat input (rendering ``tools`` 

211 into the template), honoring max_tokens/sampling, with optional per-token logprobs, 

212 token usage, parsed tool calls, and per-turn activation capture (agent rollouts).""" 

213 _warn_unsupported_config(config, PROVIDER_NAME) 

214 ids = self._messages_to_ids(input, tools) 

215 prompt_len = int(ids.shape[1]) 

216 max_new = int(config.max_tokens) if config.max_tokens else 16 

217 temperature = config.temperature 

218 do_sample = temperature is not None and temperature > 0 

219 gen: dict[str, Any] = { 

220 "max_new_tokens": max_new, 

221 "do_sample": do_sample, 

222 "return_dict_in_generate": True, 

223 "output_scores": True, 

224 "pad_token_id": self._tokenizer.pad_token_id or self._tokenizer.eos_token_id, 

225 } 

226 if config.stop_seqs: 

227 # transformers turns stop_strings into a StopStringCriteria; needs the tokenizer. 

228 gen["stop_strings"] = list(config.stop_seqs) 

229 gen["tokenizer"] = self._tokenizer 

230 if temperature is not None and temperature > 0: 230 ↛ 231line 230 didn't jump to line 231 because the condition on line 230 was never true

231 gen["temperature"] = float(temperature) 

232 if config.top_p is not None: 

233 gen["top_p"] = float(config.top_p) 

234 if config.top_k is not None: 

235 gen["top_k"] = int(config.top_k) 

236 

237 # Save BOTH CPU and CUDA RNG state — get_rng_state() is CPU-only, so seeding on a 

238 # CUDA model would otherwise leak its CUDA seed past this generate. 

239 rng_state = None 

240 cuda_rng_state = None 

241 on_cuda = "cuda" in str(self._device) 

242 if do_sample and config.seed is not None: 242 ↛ 243line 242 didn't jump to line 243 because the condition on line 242 was never true

243 rng_state = torch.get_rng_state() 

244 if on_cuda: 

245 cuda_rng_state = torch.cuda.get_rng_state_all() 

246 torch.manual_seed(int(config.seed)) 

247 # Install per-turn capture hooks AROUND generate (not pre-): first-write-wins lets 

248 # the prompt forward populate them and decode forwards skip — no extra forward. 

249 with self._eval_capture_scope() as metadata: 

250 try: 

251 with torch.no_grad(): 

252 out = self._hf.generate(ids, **gen) 

253 finally: 

254 if rng_state is not None: 254 ↛ 255line 254 didn't jump to line 255 because the condition on line 254 was never true

255 torch.set_rng_state(rng_state) 

256 if cuda_rng_state is not None: 256 ↛ 257line 256 didn't jump to line 257 because the condition on line 256 was never true

257 torch.cuda.set_rng_state_all(cuda_rng_state) 

258 

259 new_ids = out.sequences[0, prompt_len:] 

260 completion = str(self._tokenizer.decode(new_ids, skip_special_tokens=True)) 

261 logprobs = None 

262 if config.logprobs: 

263 content = [ 

264 self._logprob_entry(int(tok), step[0], config.top_logprobs) 

265 for tok, step in zip(new_ids.tolist(), out.scores) 

266 ] 

267 logprobs = Logprobs(content=content) 

268 n_new = int(new_ids.shape[0]) 

269 tool_calls = _parse_tool_calls(completion) if len(tools) else None 

270 eos = self._tokenizer.eos_token_id 

271 stop_reason: StopReason 

272 if tool_calls: 272 ↛ 273line 272 didn't jump to line 273 because the condition on line 272 was never true

273 stop_reason = "tool_calls" 

274 elif n_new and eos is not None and int(new_ids[-1]) == eos: 274 ↛ 275line 274 didn't jump to line 275 because the condition on line 274 was never true

275 stop_reason = "stop" 

276 else: 

277 stop_reason = "max_tokens" 

278 return ModelOutput( 

279 model=self.model_name, 

280 choices=[ 

281 ChatCompletionChoice( 

282 message=ChatMessageAssistant(content=completion, tool_calls=tool_calls), 

283 stop_reason=stop_reason, 

284 logprobs=logprobs, 

285 ) 

286 ], 

287 usage=ModelUsage( 

288 input_tokens=prompt_len, output_tokens=n_new, total_tokens=prompt_len + n_new 

289 ), 

290 metadata=metadata or None, 

291 ) 

292 

293 @contextmanager 

294 def _eval_capture_scope(self) -> Iterator[dict[str, Any]]: 

295 """Install per-turn ``capture=[...]`` hooks for the duration of a generate. First- 

296 write-wins lets the prompt forward populate the raw dict (decode forwards find it 

297 populated and skip), so this adds no extra forward. Yields a metadata dict (empty 

298 when capture isn't configured) that the caller folds into ``ModelOutput.metadata``. 

299 Contextvar-isolated so concurrent inspect_eval samples don't cross-pollute.""" 

300 if not self._eval_capture: 

301 yield {} 

302 return 

303 wire_keys = list(self._eval_capture) 

304 capture, _ = _plan(wire_keys, {}) 

305 raw: dict[tuple[int, str], np.ndarray] = {} 

306 call_id = uuid.uuid4().hex 

307 token = _current_call_id.set(call_id) 

308 handles = self._install_hooks(capture, {}, raw, call_id) 

309 metadata: dict[str, Any] = {} 

310 try: 

311 yield metadata 

312 finally: 

313 for handle in handles: 

314 handle.remove() 

315 _current_call_id.reset(token) 

316 metadata["activations"] = wire.encode_activations(_assemble(raw, wire_keys)) 

317 

318 def _install_hooks(self, capture, intervene, raw, call_id: str) -> list: 

319 """Hook each block's pre/attn/mlp/post boundaries that need capture or intervention. 

320 Hooks consult ``_current_call_id`` and only fire for ``call_id`` (concurrent calls).""" 

321 handles = [] 

322 for layer, block in enumerate(self._layers): 

323 cap_kinds = capture.get(layer, set()) 

324 iv_kinds = intervene.get(layer, {}) 

325 if not cap_kinds and not iv_kinds: 

326 continue 

327 attn = _first_attr(block, _ATTN_ATTRS) 

328 _, mlp_out_mod = _locate_mlp(block) 

329 # Capabilities were probed on layer 0; hybrid archs can lack the module here. 

330 _require_layer_modules(layer, attn, mlp_out_mod, cap_kinds, iv_kinds) 

331 if "resid_pre" in cap_kinds or "resid_pre" in iv_kinds: 

332 handles.append( 

333 block.register_forward_pre_hook( 

334 _pre_hook( 

335 layer, "resid_pre" in cap_kinds, iv_kinds.get("resid_pre"), raw, call_id 

336 ), 

337 with_kwargs=True, 

338 ) 

339 ) 

340 if attn is not None and ("attn_out" in cap_kinds or "attn_out" in iv_kinds): 

341 handles.append( 

342 attn.register_forward_hook( 

343 _out_hook( 

344 layer, 

345 "attn_out", 

346 "attn_out" in cap_kinds, 

347 iv_kinds.get("attn_out"), 

348 raw, 

349 call_id, 

350 ) 

351 ) 

352 ) 

353 if attn is not None: 

354 handles.extend( 

355 self._install_head_hooks(layer, attn, cap_kinds, iv_kinds, raw, call_id) 

356 ) 

357 if mlp_out_mod is not None and ("mlp_out" in cap_kinds or "mlp_out" in iv_kinds): 

358 handles.append( 

359 mlp_out_mod.register_forward_hook( 

360 _out_hook( 

361 layer, 

362 "mlp_out", 

363 "mlp_out" in cap_kinds, 

364 iv_kinds.get("mlp_out"), 

365 raw, 

366 call_id, 

367 ) 

368 ) 

369 ) 

370 if "resid_post" in cap_kinds or "resid_post" in iv_kinds: 

371 handles.append( 

372 block.register_forward_hook( 

373 _out_hook( 

374 layer, 

375 "resid_post", 

376 "resid_post" in cap_kinds, 

377 iv_kinds.get("resid_post"), 

378 raw, 

379 call_id, 

380 ) 

381 ) 

382 ) 

383 return handles 

384 

385 def _install_head_hooks(self, layer, attn, cap_kinds, iv_kinds, raw, call_id: str) -> list: 

386 """Hooks for the head-split kinds: q/k/v on their projection outputs, z on the 

387 out-projection's input. pattern isn't hooked (it rides output_attentions). 

388 

389 Interventions apply to the module's natural *flat* tensor ``(..., seq, 

390 heads·d_head)`` — a spec ``value`` is scalar, ``(heads·d_head,)``, or per-position 

391 ``(len(pos), heads·d_head)`` (flatten a captured head-split tensor to build one). 

392 Captures are emitted head-split ``(seq, heads, d_head)`` to match the bridge's 

393 ``hook_q/k/v/z``. 

394 """ 

395 handles: list = [] 

396 d_head = self._geometry[2] 

397 if d_head is None: # geometry undetectable → head kinds were gated at detection 397 ↛ 398line 397 didn't jump to line 398 because the condition on line 397 was never true

398 return handles 

399 host = _projection_host(attn) 

400 for kind, attrs in (("q", _Q_PROJ_ATTRS), ("k", _K_PROJ_ATTRS), ("v", _V_PROJ_ATTRS)): 

401 if kind not in cap_kinds and kind not in iv_kinds: 

402 continue 

403 proj = _first_attr(host, attrs) 

404 if proj is None: 404 ↛ 408line 404 didn't jump to line 408 because the condition on line 404 was never true

405 # Detection ran on layers[0]; a later layer missing the projection would 

406 # silently skip a validated intervention — fail loud (capture-only misses 

407 # surface through the driver's missing-hook warning instead). 

408 if kind in iv_kinds: 

409 raise RuntimeError( 

410 f"Intervention on blocks.{layer}.attn.hook_{kind} cannot be applied: " 

411 f"layer {layer} has no {kind} projection (heterogeneous layers)." 

412 ) 

413 continue 

414 handles.append( 

415 proj.register_forward_hook( 

416 _proj_hook( 

417 layer, kind, d_head, kind in cap_kinds, iv_kinds.get(kind), raw, call_id 

418 ) 

419 ) 

420 ) 

421 if "z" in cap_kinds or "z" in iv_kinds: 

422 o_proj = _first_attr(host, _O_PROJ_ATTRS) 

423 if o_proj is None and "z" in iv_kinds: 423 ↛ 424line 423 didn't jump to line 424 because the condition on line 423 was never true

424 raise RuntimeError( 

425 f"Intervention on blocks.{layer}.attn.hook_z cannot be applied: layer " 

426 f"{layer} has no out-projection (heterogeneous layers)." 

427 ) 

428 if o_proj is not None: 428 ↛ 435line 428 didn't jump to line 435 because the condition on line 428 was always true

429 handles.append( 

430 o_proj.register_forward_pre_hook( 

431 _zin_hook(layer, d_head, "z" in cap_kinds, iv_kinds.get("z"), raw, call_id), 

432 with_kwargs=True, 

433 ) 

434 ) 

435 return handles 

436 

437 

438def _require_layer_modules( 

439 layer: int, attn: Any, mlp_out_mod: Any, cap_kinds: Any, iv_kinds: Any 

440) -> None: 

441 """Fail loud when this layer lacks a targeted submodule — detection ran on layer 0, 

442 so hybrid attn/SSM stacks would otherwise install nothing and silently no-op.""" 

443 wanted = set(cap_kinds) | set(iv_kinds) 

444 if attn is None: 

445 needs_attn = wanted & ({"attn_out"} | (hooks.HEAD_KINDS - {"pattern"})) 

446 if needs_attn: 

447 raise RuntimeError( 

448 f"blocks.{layer} has no attention submodule: cannot serve " 

449 f"{sorted(needs_attn)} at layer {layer} (heterogeneous layers)." 

450 ) 

451 if mlp_out_mod is None and "mlp_out" in wanted: 

452 raise RuntimeError( 

453 f"blocks.{layer} has no MLP submodule: cannot serve mlp_out at layer " 

454 f"{layer} (heterogeneous layers)." 

455 ) 

456 

457 

458def _plan(capture_keys, interventions): 

459 """Resolve wire keys → per-layer kinds to capture (resid_mid needs pre+attn) and intervene.""" 

460 capture: dict[int, set[str]] = defaultdict(set) 

461 for key in capture_keys: 

462 layer, _, kind = key.partition(":") 

463 layer = int(layer) 

464 if kind == "resid_mid": 

465 capture[layer] |= {"resid_pre", "attn_out"} # derived = resid_pre + attn_out 

466 else: 

467 capture[layer].add(kind) 

468 intervene: dict[int, dict[str, Any]] = defaultdict(dict) 

469 for key, spec in interventions.items(): 

470 layer, _, kind = key.partition(":") 

471 intervene[int(layer)][kind] = spec 

472 return capture, intervene 

473 

474 

475def _assemble(raw, capture_keys) -> dict[str, np.ndarray]: 

476 """Build the emitted ``{wire_key: (seq, d)}`` map, deriving resid_mid as needed.""" 

477 out: dict[str, np.ndarray] = {} 

478 for key in capture_keys: 

479 layer, _, kind = key.partition(":") 

480 layer = int(layer) 

481 if kind == "resid_mid": 

482 pre, attn = raw.get((layer, "resid_pre")), raw.get((layer, "attn_out")) 

483 if pre is not None and attn is not None: 483 ↛ 478line 483 didn't jump to line 478 because the condition on line 483 was always true

484 out[key] = pre + attn 

485 elif (layer, kind) in raw: 485 ↛ 478line 485 didn't jump to line 478 because the condition on line 485 was always true

486 out[key] = raw[(layer, kind)] 

487 return out 

488 

489 

490def _pre_hook(layer, want_capture, spec, raw, call_id): 

491 # with_kwargs=True: hidden_states is args[0] for most decoders, but some pass it as 

492 # the hidden_states kwarg — handle both so the right tensor is read/modified. 

493 # First-write-wins on raw so install-around-generate captures the prompt forward 

494 # (subsequent decode forwards find raw populated and skip — no extra prompt forward). 

495 def hook(_module, args, kwargs): 

496 if _current_call_id.get() != call_id: 496 ↛ 497line 496 didn't jump to line 497 because the condition on line 496 was never true

497 return None # different concurrent call's hook 

498 kw_key = None if args else "hidden_states" 

499 hidden = args[0] if args else kwargs["hidden_states"] 

500 if spec is not None: 

501 hidden = _apply_affine(hidden, spec) 

502 if want_capture and (layer, "resid_pre") not in raw: 502 ↛ 504line 502 didn't jump to line 504 because the condition on line 502 was always true

503 raw[(layer, "resid_pre")] = hidden[0].detach().float().cpu().numpy() 

504 if spec is None: 

505 return None 

506 if kw_key is None: 

507 return (hidden, *args[1:]), kwargs 

508 return args, {**kwargs, kw_key: hidden} 

509 

510 return hook 

511 

512 

513def _out_hook(layer, kind, want_capture, spec, raw, call_id): 

514 def hook(_module, _inputs, output): 

515 if _current_call_id.get() != call_id: 

516 return None # different concurrent call's hook 

517 is_tuple = isinstance(output, tuple) 

518 hidden = output[0] if is_tuple else output 

519 if spec is not None: 

520 hidden = _apply_affine(hidden, spec) 

521 if want_capture and (layer, kind) not in raw: 

522 # OPT-style blocks flatten the FFN to (batch·seq, d) — with batch_size=1 

523 # that IS (seq, d) already; only 3-D (batch, seq, d) needs the batch strip. 

524 flat = hidden if hidden.ndim == 2 else hidden[0] 

525 raw[(layer, kind)] = flat.detach().float().cpu().numpy() 

526 if spec is None: 

527 return None 

528 return (hidden, *output[1:]) if is_tuple else hidden 

529 

530 return hook 

531 

532 

533def _proj_hook(layer, kind, d_head, want_capture, spec, raw, call_id): 

534 """q/k/v projection output: affine on the flat ``(..., seq, heads·d_head)`` tensor, 

535 captured head-split ``(seq, heads, d_head)`` (pre-RoPE — matches the bridge's 

536 ``hook_q``/``hook_k``/``hook_v``). Mutations feed the downstream attention math.""" 

537 

538 def hook(_module, _inputs, output): 

539 if _current_call_id.get() != call_id: 539 ↛ 540line 539 didn't jump to line 540 because the condition on line 539 was never true

540 return None # different concurrent call's hook 

541 hidden = output 

542 if spec is not None: 

543 hidden = _apply_affine(hidden, spec) 

544 if want_capture and (layer, kind) not in raw: 544 ↛ 547line 544 didn't jump to line 547 because the condition on line 544 was always true

545 flat = hidden[0].detach().float().cpu() 

546 raw[(layer, kind)] = flat.reshape(flat.shape[0], -1, d_head).numpy() 

547 return hidden if spec is not None else None 

548 

549 return hook 

550 

551 

552def _zin_hook(layer, d_head, want_capture, spec, raw, call_id): 

553 """z — the out-projection's *input* (attention-weighted values, all heads): pre-hook, 

554 affine on the flat tensor, captured head-split to match the bridge's ``hook_z``.""" 

555 

556 def hook(_module, args, kwargs): 

557 if _current_call_id.get() != call_id: 557 ↛ 558line 557 didn't jump to line 558 because the condition on line 557 was never true

558 return None # different concurrent call's hook 

559 z = args[0] 

560 if spec is not None: 560 ↛ 561line 560 didn't jump to line 561 because the condition on line 560 was never true

561 z = _apply_affine(z, spec) 

562 if want_capture and (layer, "z") not in raw: 562 ↛ 565line 562 didn't jump to line 565 because the condition on line 562 was always true

563 flat = z[0].detach().float().cpu() 

564 raw[(layer, "z")] = flat.reshape(flat.shape[0], -1, d_head).numpy() 

565 if spec is None: 565 ↛ 567line 565 didn't jump to line 567 because the condition on line 565 was always true

566 return None 

567 return (z, *args[1:]), kwargs 

568 

569 return hook 

570 

571 

572def _affine_op(sub: torch.Tensor, spec: Mapping[str, Any]) -> torch.Tensor: 

573 """One affine op: suppress→0, scale→·factor, add→+value, set→value. ``value`` broadcasts 

574 (scalar, width-shaped, or per-position ``(n_pos, width)``).""" 

575 op = spec["op"] 

576 if op == "suppress": 

577 return torch.zeros_like(sub) 

578 if op == "scale": 578 ↛ 579line 578 didn't jump to line 579 because the condition on line 578 was never true

579 return sub * float(spec["factor"]) 

580 value = torch.as_tensor(spec["value"], dtype=sub.dtype, device=sub.device) 

581 if op == "add": 581 ↛ 583line 581 didn't jump to line 583 because the condition on line 581 was always true

582 return sub + value 

583 return torch.zeros_like(sub) + value # set 

584 

585 

586def _apply_affine(t: torch.Tensor, spec: Mapping[str, Any]) -> torch.Tensor: 

587 """Affine intervention on a captured tensor ``(..., seq, width)``. 

588 

589 Without ``pos`` the op spans every position (the original width-broadcast form). With 

590 ``pos`` (an int or list of sequence indices) it touches only those positions — the 

591 activation-patching primitive — and ``value`` may be per-position ``(len(pos), width)`` 

592 to transplant a captured activation (path/causal tracing) rather than a single vector. 

593 """ 

594 pos = spec.get("pos") 

595 if pos is None: 

596 return _affine_op(t, spec) 

597 idx = [pos] if isinstance(pos, int) else list(pos) 

598 out = t.clone() 

599 out[..., idx, :] = _affine_op(t[..., idx, :], spec) 

600 return out 

601 

602 

603def _detect_capabilities( 

604 model: Any, layers: Any, geometry: tuple[Any, Any, Any] 

605) -> tuple[frozenset, str]: 

606 """Structural self-check: which kinds this model can serve faithfully. 

607 

608 resid_pre/resid_post are the block in/out (always); attn_out/mlp_out need their 

609 submodules locatable; resid_mid is gated unless its derivation holds (see 

610 :func:`_resid_mid_derivable`). Head-split kinds: q/k/v need separate projections whose 

611 widths match ``heads·d_head``; z needs an out-projection of in-width ``n_heads·d_head``; 

612 pattern needs eager attention (output_attentions is a no-op under sdpa/flash). 

613 Returns (kinds, note); note explains any gating, '' if none. 

614 """ 

615 block = layers[0] 

616 attn = _first_attr(block, _ATTN_ATTRS) 

617 mlp_in_mod, mlp_out_mod = _locate_mlp(block) 

618 kinds = {"resid_pre", "resid_post"} 

619 gated = [] 

620 if attn is not None: 

621 kinds.add("attn_out") 

622 else: 

623 gated.append("attn_out (no attention submodule found)") 

624 if mlp_out_mod is not None: 

625 kinds.add("mlp_out") 

626 else: 

627 gated.append("mlp_out (no MLP submodule found)") 

628 if ( 

629 attn is not None 

630 and mlp_out_mod is not None 

631 and _resid_mid_derivable(model, block, attn, mlp_in_mod, mlp_out_mod) 

632 ): 

633 kinds.add("resid_mid") 

634 else: 

635 gated.append( 

636 "resid_mid (resid_pre + attn_out doesn't hold — parallel or norm-variant block)" 

637 ) 

638 head_kinds, head_gated = _detect_head_capabilities(model, attn, geometry) 

639 kinds |= head_kinds 

640 gated += head_gated 

641 note = ( 

642 "" 

643 if not gated 

644 else "InspectDriver: this architecture's block layout gates " 

645 + ", ".join(gated) 

646 + ". Remaining boundaries are served; use boot_transformers() for the gated ones." 

647 ) 

648 return frozenset(kinds), note 

649 

650 

651def _detect_head_capabilities( 

652 model: Any, attn: Any, geometry: tuple[Any, Any, Any] 

653) -> tuple[set, list]: 

654 """Head-split kinds this model serves: q/k/v iff separate projections with the 

655 expected widths, z iff the out-projection's in-width is ``n_heads·d_head``, pattern 

656 iff attention runs eager (otherwise ``output_attentions`` returns None/garbage).""" 

657 n_heads, n_kv_heads, d_head = geometry 

658 kinds: set = set() 

659 gated: list = [] 

660 if attn is None or d_head is None: 

661 gated.append("q/k/v/z/pattern (no attention submodule or head geometry in config)") 

662 return kinds, gated 

663 

664 host = _projection_host(attn) 

665 q = _first_attr(host, _Q_PROJ_ATTRS) 

666 k = _first_attr(host, _K_PROJ_ATTRS) 

667 v = _first_attr(host, _V_PROJ_ATTRS) 

668 expected = {"q": n_heads * d_head, "k": n_kv_heads * d_head, "v": n_kv_heads * d_head} 

669 if all( 

670 proj is not None and _out_width(proj) == expected[kind] 

671 for kind, proj in (("q", q), ("k", k), ("v", v)) 

672 ): 

673 kinds |= {"q", "k", "v"} 

674 else: 

675 gated.append("q/k/v (fused or nonstandard qkv projections)") 

676 

677 o_proj = _first_attr(host, _O_PROJ_ATTRS) 

678 if o_proj is not None and _in_width(o_proj) == n_heads * d_head: 678 ↛ 681line 678 didn't jump to line 681 because the condition on line 678 was always true

679 kinds.add("z") 

680 else: 

681 gated.append("z (out-projection missing or nonstandard width)") 

682 

683 if getattr(model.config, "_attn_implementation", "eager") == "eager": 

684 kinds.add("pattern") 

685 else: 

686 gated.append("pattern (attention implementation is not eager)") 

687 return kinds, gated 

688 

689 

690def _out_width(module: Any) -> Any: 

691 """Output width of a projection: ``nn.Linear.out_features`` or GPT-2 ``Conv1D.nf``.""" 

692 out = getattr(module, "out_features", None) 

693 if out is not None: 693 ↛ 695line 693 didn't jump to line 695 because the condition on line 693 was always true

694 return int(out) 

695 nf = getattr(module, "nf", None) # transformers Conv1D 

696 return int(nf) if nf is not None else None 

697 

698 

699def _in_width(module: Any) -> Any: 

700 """Input width of a projection: ``nn.Linear.in_features`` or Conv1D ``weight.shape[0]``.""" 

701 in_f = getattr(module, "in_features", None) 

702 if in_f is not None: 

703 return int(in_f) 

704 if getattr(module, "nf", None) is not None and hasattr(module, "weight"): 704 ↛ 706line 704 didn't jump to line 706 because the condition on line 704 was always true

705 return int(module.weight.shape[0]) # Conv1D stores weight (in, out) 

706 return None 

707 

708 

709def _attn_geometry(config: Any) -> tuple[Any, Any, Any]: 

710 """(n_heads, n_kv_heads, d_head) from an HF config; (None, None, None) if underivable.""" 

711 n_heads = getattr(config, "num_attention_heads", None) or getattr(config, "n_head", None) 

712 hidden = getattr(config, "hidden_size", None) or getattr(config, "n_embd", None) 

713 if n_heads is None: 713 ↛ 714line 713 didn't jump to line 714 because the condition on line 713 was never true

714 return None, None, None 

715 n_kv = getattr(config, "num_key_value_heads", None) or n_heads 

716 d_head = getattr(config, "head_dim", None) 

717 if d_head is None and hidden is not None: 

718 d_head = hidden // n_heads 

719 return (int(n_heads), int(n_kv), int(d_head) if d_head is not None else None) 

720 

721 

722def _resid_mid_derivable( 

723 model: Any, block: Any, attn: Any, mlp_in_mod: Any, mlp_out_mod: Any 

724) -> bool: 

725 """True iff ``resid_mid = resid_pre + attn_out`` holds, via two tiny probe forwards. 

726 Requires both the linear identity ``resid_post = resid_pre + attn_out + mlp_out`` (broken 

727 by post-norm/multiplier blocks — Gemma2/OLMo2/Granite) and attn feeding mlp (broken by 

728 parallel blocks — GPTNeoX/GPT-J, where mlp reads resid_pre directly). ``mlp_in_mod``/ 

729 ``mlp_out_mod`` are the same module for container archs, (fc1, fc2) for fc-split.""" 

730 cap: dict[str, Any] = {} 

731 

732 def grab(key: str): # type: ignore[no-untyped-def] 

733 def hook(_m: Any, _i: Any, out: Any) -> None: 

734 t = out[0] if isinstance(out, tuple) else out 

735 cap[key] = t.detach().float() 

736 

737 return hook 

738 

739 def grab_in(key: str): # type: ignore[no-untyped-def]