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 transformer_lens.utilities.position_ids import accepts_mask_derived_position_ids 

33 

34from . import hooks, wire 

35from ._provider_base import ( 

36 _InspectModelAPIBase, 

37 _parse_tool_calls, 

38 _require_interveneable, 

39 _require_served, 

40 _warn_unsupported_config, 

41) 

42from .masks import normalize_attention_mask 

43 

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

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

46PROVIDER_NAME = "tl_bridge" 

47 

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

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

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

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

52 "tl_inspect_call_id", default="" 

53) 

54 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

71 

72 

73@modelapi(name=PROVIDER_NAME) 

74def transformer_lens_provider(): 

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

76 return TransformerLensTransformersModelAPI 

77 

78 

79class TransformerLensTransformersModelAPI(_InspectModelAPIBase): 

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

81 

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

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

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

85 """ 

86 

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

88 provides_sequence_logits = True 

89 supports_attention_mask = True 

90 

91 def __init__( 

92 self, 

93 model_name: str, 

94 base_url: str | None = None, 

95 api_key: str | None = None, 

96 config: GenerateConfig = GenerateConfig(), 

97 **model_args: Any, 

98 ) -> None: 

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

100 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer 

101 

102 from transformer_lens.model_bridge.sources._hf_format import ( 

103 determine_architecture_from_hf_config, 

104 ) 

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

106 get_hf_model_class_for_architecture, 

107 ) 

108 

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

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

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

112 # AutoModel entry point than AutoModelForCausalLM. 

113 config_kwargs = { 

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

115 } 

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

117 try: 

118 model_class = get_hf_model_class_for_architecture( 

119 determine_architecture_from_hf_config(hf_config) 

120 ) 

121 except ValueError: 

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

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

124 model_class = AutoModelForCausalLM 

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

126 self._tokenizer = AutoTokenizer.from_pretrained(model_name) 

127 self._layers = _locate_layers(self._hf) 

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

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

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

131 self._kinds, self._capability_note = _detect_capabilities( 

132 self._hf, self._layers, self._geometry 

133 ) 

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

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

136 self._eval_capture = self._parse_eval_capture(model_args) 

137 for key in self._eval_capture: 

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

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

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

141 raise ValueError( 

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

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

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

145 ) 

146 

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

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

149 input_ids = extra_args.get("input_ids") 

150 if input_ids is None: 

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

152 mask_value = extra_args.get("attention_mask") 

153 mask = ( 

154 normalize_attention_mask(mask_value, len(input_ids)) if mask_value is not None else None 

155 ) 

156 forward_kwargs: dict[str, Any] = {} 

157 if mask is not None: 

158 attention_mask = torch.tensor([mask], device=self._device) 

159 forward_kwargs["attention_mask"] = attention_mask 

160 positions = (attention_mask.cumsum(-1) - 1).clamp_min(0) 

161 default_positions = torch.arange(len(input_ids), device=self._device) 

162 needs_positions = bool(((positions != default_positions) & attention_mask.bool()).any()) 

163 if needs_positions and accepts_mask_derived_position_ids(self._hf): 

164 forward_kwargs["position_ids"] = positions 

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

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

167 for key in capture_keys: 

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

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

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

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

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

173 for key in interventions: 

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

175 _require_interveneable( 

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

177 ) 

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

179 

180 capture, intervene = _plan(capture_keys, interventions) 

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

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

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

184 call_id = uuid.uuid4().hex 

185 token = _current_call_id.set(call_id) 

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

187 try: 

188 with torch.no_grad(): 

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

190 outputs = self._hf(ids, output_attentions=bool(pattern_layers), **forward_kwargs) 

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

192 finally: 

193 for handle in handles: 

194 handle.remove() 

195 _current_call_id.reset(token) 

196 

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

198 for layer in pattern_layers: 

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

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

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

202 

203 captured = _assemble(raw, capture_keys) 

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

205 if want_logits: 

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

207 

208 last_token = max(i for i, keep in enumerate(mask) if keep) if mask is not None else -1 

209 next_logits = logits[0, last_token] 

210 next_id = int(next_logits.argmax()) 

211 logprobs = ( 

212 Logprobs(content=[self._logprob_entry(next_id, next_logits, config.top_logprobs)]) 

213 if config.logprobs 

214 else None 

215 ) 

216 return ModelOutput( 

217 model=self.model_name, 

218 choices=[ 

219 ChatCompletionChoice( 

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

221 stop_reason="stop", 

222 logprobs=logprobs, 

223 ) 

224 ], 

225 metadata=metadata, 

226 ) 

227 

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

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

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

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

232 _warn_unsupported_config(config, PROVIDER_NAME) 

233 ids = self._messages_to_ids(input, tools) 

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

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

236 temperature = config.temperature 

237 do_sample = temperature is not None and temperature > 0 

238 gen: dict[str, Any] = { 

239 "max_new_tokens": max_new, 

240 "do_sample": do_sample, 

241 "return_dict_in_generate": True, 

242 "output_scores": True, 

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

244 } 

245 if config.stop_seqs: 

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

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

248 gen["tokenizer"] = self._tokenizer 

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

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

251 if config.top_p is not None: 

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

253 if config.top_k is not None: 

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

255 

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

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

258 rng_state = None 

259 cuda_rng_state = None 

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

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

262 rng_state = torch.get_rng_state() 

263 if on_cuda: 

264 cuda_rng_state = torch.cuda.get_rng_state_all() 

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

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

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

268 with self._eval_capture_scope() as metadata: 

269 try: 

270 with torch.no_grad(): 

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

272 finally: 

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

274 torch.set_rng_state(rng_state) 

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

276 torch.cuda.set_rng_state_all(cuda_rng_state) 

277 

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

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

280 logprobs = None 

281 if config.logprobs: 

282 content = [ 

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

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

285 ] 

286 logprobs = Logprobs(content=content) 

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

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

289 eos = self._tokenizer.eos_token_id 

290 stop_reason: StopReason 

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

292 stop_reason = "tool_calls" 

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

294 stop_reason = "stop" 

295 else: 

296 stop_reason = "max_tokens" 

297 return ModelOutput( 

298 model=self.model_name, 

299 choices=[ 

300 ChatCompletionChoice( 

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

302 stop_reason=stop_reason, 

303 logprobs=logprobs, 

304 ) 

305 ], 

306 usage=ModelUsage( 

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

308 ), 

309 metadata=metadata or None, 

310 ) 

311 

312 @contextmanager 

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

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

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

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

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

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

319 if not self._eval_capture: 

320 yield {} 

321 return 

322 wire_keys = list(self._eval_capture) 

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

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

325 call_id = uuid.uuid4().hex 

326 token = _current_call_id.set(call_id) 

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

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

329 try: 

330 yield metadata 

331 finally: 

332 for handle in handles: 

333 handle.remove() 

334 _current_call_id.reset(token) 

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

336 

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

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

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

340 handles = [] 

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

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

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

344 if not cap_kinds and not iv_kinds: 

345 continue 

346 attn = _first_attr(block, _ATTN_ATTRS) 

347 _, mlp_out_mod = _locate_mlp(block) 

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

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

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

351 handles.append( 

352 block.register_forward_pre_hook( 

353 _pre_hook( 

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

355 ), 

356 with_kwargs=True, 

357 ) 

358 ) 

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

360 handles.append( 

361 attn.register_forward_hook( 

362 _out_hook( 

363 layer, 

364 "attn_out", 

365 "attn_out" in cap_kinds, 

366 iv_kinds.get("attn_out"), 

367 raw, 

368 call_id, 

369 ) 

370 ) 

371 ) 

372 if attn is not None: 

373 handles.extend( 

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

375 ) 

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

377 handles.append( 

378 mlp_out_mod.register_forward_hook( 

379 _out_hook( 

380 layer, 

381 "mlp_out", 

382 "mlp_out" in cap_kinds, 

383 iv_kinds.get("mlp_out"), 

384 raw, 

385 call_id, 

386 ) 

387 ) 

388 ) 

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

390 handles.append( 

391 block.register_forward_hook( 

392 _out_hook( 

393 layer, 

394 "resid_post", 

395 "resid_post" in cap_kinds, 

396 iv_kinds.get("resid_post"), 

397 raw, 

398 call_id, 

399 ) 

400 ) 

401 ) 

402 return handles 

403 

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

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

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

407 

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

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

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

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

412 ``hook_q/k/v/z``. 

413 """ 

414 handles: list = [] 

415 d_head = self._geometry[2] 

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

417 return handles 

418 host = _projection_host(attn) 

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

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

421 continue 

422 proj = _first_attr(host, attrs) 

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

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

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

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

427 if kind in iv_kinds: 

428 raise RuntimeError( 

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

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

431 ) 

432 continue 

433 handles.append( 

434 proj.register_forward_hook( 

435 _proj_hook( 

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

437 ) 

438 ) 

439 ) 

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

441 o_proj = _first_attr(host, _O_PROJ_ATTRS) 

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

443 raise RuntimeError( 

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

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

446 ) 

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

448 handles.append( 

449 o_proj.register_forward_pre_hook( 

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

451 with_kwargs=True, 

452 ) 

453 ) 

454 return handles 

455 

456 

457def _require_layer_modules( 

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

459) -> None: 

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

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

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

463 if attn is None: 

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

465 if needs_attn: 

466 raise RuntimeError( 

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

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

469 ) 

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

471 raise RuntimeError( 

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

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

474 ) 

475 

476 

477def _plan(capture_keys, interventions): 

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

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

480 for key in capture_keys: 

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

482 layer = int(layer) 

483 if kind == "resid_mid": 

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

485 else: 

486 capture[layer].add(kind) 

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

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

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

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

491 return capture, intervene 

492 

493 

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

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

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

497 for key in capture_keys: 

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

499 layer = int(layer) 

500 if kind == "resid_mid": 

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

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

503 out[key] = pre + attn 

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

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

506 return out 

507 

508 

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

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

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

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

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

514 def hook(_module, args, kwargs): 

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

516 return None # different concurrent call's hook 

517 kw_key = None if args else "hidden_states" 

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

519 if spec is not None: 

520 hidden = _apply_affine(hidden, spec) 

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

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

523 if spec is None: 

524 return None 

525 if kw_key is None: 

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

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

528 

529 return hook 

530 

531 

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

533 def hook(_module, _inputs, output): 

534 if _current_call_id.get() != call_id: 

535 return None # different concurrent call's hook 

536 is_tuple = isinstance(output, tuple) 

537 hidden = output[0] if is_tuple else output 

538 if spec is not None: 

539 hidden = _apply_affine(hidden, spec) 

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

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

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

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

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

545 if spec is None: 

546 return None 

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

548 

549 return hook 

550 

551 

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

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

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

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

556 

557 def hook(_module, _inputs, output): 

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

559 return None # different concurrent call's hook 

560 hidden = output 

561 if spec is not None: 

562 hidden = _apply_affine(hidden, spec) 

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

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

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

566 return hidden if spec is not None else None 

567 

568 return hook 

569 

570 

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

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

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

574 

575 def hook(_module, args, kwargs): 

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

577 return None # different concurrent call's hook 

578 z = args[0] 

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

580 z = _apply_affine(z, spec) 

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

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

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

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

585 return None 

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

587 

588 return hook 

589 

590 

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

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

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

594 op = spec["op"] 

595 if op == "suppress": 

596 return torch.zeros_like(sub) 

597 if op == "scale": 

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

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

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

601 return sub + value 

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

603 

604 

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

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

607 

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

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

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

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

612 """ 

613 pos = spec.get("pos") 

614 if pos is None: 

615 return _affine_op(t, spec) 

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

617 out = t.clone() 

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

619 return out 

620 

621 

622def _detect_capabilities( 

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

624) -> tuple[frozenset, str]: 

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

626 

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

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

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

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

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

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

633 """ 

634 block = layers[0] 

635 attn = _first_attr(block, _ATTN_ATTRS) 

636 mlp_in_mod, mlp_out_mod = _locate_mlp(block) 

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

638 gated = [] 

639 if attn is not None: 

640 kinds.add("attn_out") 

641 else: 

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

643 if mlp_out_mod is not None: 

644 kinds.add("mlp_out") 

645 else: 

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

647 if ( 

648 attn is not None 

649 and mlp_out_mod is not None 

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

651 ): 

652 kinds.add("resid_mid") 

653 else: 

654 gated.append( 

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

656 ) 

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

658 kinds |= head_kinds 

659 gated += head_gated 

660 note = ( 

661 "" 

662 if not gated 

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

664 + ", ".join(gated) 

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

666 ) 

667 return frozenset(kinds), note 

668 

669 

670def _detect_head_capabilities( 

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

672) -> tuple[set, list]: 

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

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

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

676 n_heads, n_kv_heads, d_head = geometry 

677 kinds: set = set() 

678 gated: list = [] 

679 if attn is None or d_head is None: 

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

681 return kinds, gated 

682 

683 host = _projection_host(attn) 

684 q = _first_attr(host, _Q_PROJ_ATTRS) 

685 k = _first_attr(host, _K_PROJ_ATTRS) 

686 v = _first_attr(host, _V_PROJ_ATTRS) 

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

688 if all( 

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

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

691 ): 

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

693 else: 

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

695 

696 o_proj = _first_attr(host, _O_PROJ_ATTRS) 

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

698 kinds.add("z") 

699 else: 

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

701 

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

703 kinds.add("pattern") 

704 else: 

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

706 return kinds, gated 

707 

708 

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

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

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

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

713 return int(out) 

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

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

716 

717 

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

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

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

721 if in_f is not None: 

722 return int(in_f) 

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

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

725 return None 

726 

727 

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

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

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

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

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

733 return None, None, None 

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