Coverage for transformer_lens/benchmarks/hook_registration.py: 48%
238 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Hook registration and behavior benchmarks for TransformerBridge."""
3from typing import Dict, Optional
5import torch
7from transformer_lens import HookedTransformer
8from transformer_lens.benchmarks.utils import (
9 BenchmarkResult,
10 BenchmarkSeverity,
11 bridge_self_target_loss,
12 compare_activation_dicts,
13 compare_scalars,
14 filter_expected_missing_hooks,
15 make_capture_hook,
16)
17from transformer_lens.hook_points import HookPoint
18from transformer_lens.model_bridge import TransformerBridge
21def benchmark_hook_registry(
22 bridge: TransformerBridge,
23 reference_model: Optional[HookedTransformer] = None,
24) -> BenchmarkResult:
25 """Benchmark hook registry completeness.
27 Args:
28 bridge: TransformerBridge model to test
29 reference_model: Optional HookedTransformer reference model
31 Returns:
32 BenchmarkResult with registry comparison details
33 """
34 try:
35 if reference_model is None: 35 ↛ 61line 35 didn't jump to line 61 because the condition on line 35 was always true
36 # No reference - just verify hooks exist
37 if not hasattr(bridge, "_hook_registry"): 37 ↛ 38line 37 didn't jump to line 38 because the condition on line 37 was never true
38 return BenchmarkResult(
39 name="hook_registry",
40 severity=BenchmarkSeverity.DANGER,
41 message="Bridge does not have _hook_registry attribute",
42 passed=False,
43 )
45 hook_count = len(bridge._hook_registry)
46 if hook_count == 0: 46 ↛ 47line 46 didn't jump to line 47 because the condition on line 46 was never true
47 return BenchmarkResult(
48 name="hook_registry",
49 severity=BenchmarkSeverity.WARNING,
50 message="Bridge hook registry is empty",
51 )
53 return BenchmarkResult(
54 name="hook_registry",
55 severity=BenchmarkSeverity.INFO,
56 message=f"Bridge has {hook_count} registered hooks",
57 details={"hook_count": hook_count},
58 )
60 # Compare with reference model
61 bridge_hooks = set(bridge.hook_dict.keys())
62 reference_hooks = set(reference_model.hook_dict.keys())
64 common_hooks = bridge_hooks & reference_hooks
65 missing_hooks = reference_hooks - bridge_hooks
66 extra_hooks = bridge_hooks - reference_hooks
68 # Filter out hooks that are expected to differ due to architectural differences.
69 if missing_hooks:
70 missing_hooks = set(filter_expected_missing_hooks(missing_hooks))
72 if missing_hooks:
73 return BenchmarkResult(
74 name="hook_registry",
75 severity=BenchmarkSeverity.DANGER,
76 message=f"Bridge is missing {len(missing_hooks)} hooks from reference model",
77 details={
78 "missing_hooks": len(missing_hooks),
79 "extra_hooks": len(extra_hooks),
80 "common_hooks": len(common_hooks),
81 "sample_missing": list(missing_hooks)[:5],
82 },
83 passed=False,
84 )
86 # Bridge having extra hooks is fine - it just means Bridge has more granular hooks
87 # What matters is that all HookedTransformer hooks are present in Bridge
88 return BenchmarkResult(
89 name="hook_registry",
90 severity=BenchmarkSeverity.INFO,
91 message=f"All {len(reference_hooks)} reference hooks present in Bridge"
92 + (f" (Bridge has {len(extra_hooks)} additional hooks)" if extra_hooks else ""),
93 details={
94 "reference_hooks": len(reference_hooks),
95 "bridge_hooks": len(bridge_hooks),
96 "extra_hooks": len(extra_hooks) if extra_hooks else 0,
97 },
98 )
100 except Exception as e:
101 return BenchmarkResult(
102 name="hook_registry",
103 severity=BenchmarkSeverity.ERROR,
104 message=f"Hook registry check failed: {str(e)}",
105 passed=False,
106 )
109def benchmark_forward_hooks(
110 bridge: TransformerBridge,
111 test_text: str,
112 reference_model: Optional[HookedTransformer] = None,
113 tolerance: float = 0.5,
114 prepend_bos: Optional[bool] = None,
115) -> BenchmarkResult:
116 """Benchmark all forward hooks for activation matching.
118 Args:
119 bridge: TransformerBridge model to test
120 test_text: Input text for testing
121 reference_model: Optional HookedTransformer for comparison
122 tolerance: Tolerance for activation matching (fraction of mismatches allowed)
123 prepend_bos: Whether to prepend BOS token. If None, uses model default.
125 Returns:
126 BenchmarkResult with hook activation comparison details
127 """
128 try:
129 bridge_activations: Dict[str, torch.Tensor] = {}
130 reference_activations: Dict[str, torch.Tensor] = {}
132 # Get all hook names
133 if reference_model is not None: 133 ↛ 134line 133 didn't jump to line 134 because the condition on line 133 was never true
134 hook_names = list(reference_model.hook_dict.keys())
135 else:
136 hook_names = list(bridge.hook_dict.keys())
138 # Register hooks on bridge and track missing hooks
139 bridge_hook_points: list[tuple[str, HookPoint]] = []
140 missing_from_bridge = []
141 for hook_name in hook_names:
142 if hook_name in bridge.hook_dict: 142 ↛ 147line 142 didn't jump to line 147 because the condition on line 142 was always true
143 hook_point = bridge.hook_dict[hook_name]
144 hook_point.add_hook(make_capture_hook(bridge_activations, hook_name))
145 bridge_hook_points.append((hook_name, hook_point))
146 else:
147 missing_from_bridge.append(hook_name)
149 # Run bridge forward pass
150 with torch.no_grad():
151 if prepend_bos is not None: 151 ↛ 152line 151 didn't jump to line 152 because the condition on line 151 was never true
152 _ = bridge(test_text, prepend_bos=prepend_bos)
153 else:
154 _ = bridge(test_text)
156 # Clean up bridge hooks
157 for _, hook_point in bridge_hook_points:
158 hook_point.remove_hooks()
160 # Check for hooks that didn't fire (registered but no activation captured)
161 registered_hooks = {name for name, _ in bridge_hook_points}
162 hooks_that_didnt_fire = registered_hooks - set(bridge_activations.keys())
164 if reference_model is None: 164 ↛ 186line 164 didn't jump to line 186 because the condition on line 164 was always true
165 # No reference - just verify activations were captured
166 if hooks_that_didnt_fire: 166 ↛ 178line 166 didn't jump to line 178 because the condition on line 166 was always true
167 return BenchmarkResult(
168 name="forward_hooks",
169 severity=BenchmarkSeverity.WARNING,
170 message=f"{len(hooks_that_didnt_fire)}/{len(registered_hooks)} hooks didn't fire during forward pass",
171 details={
172 "captured": len(bridge_activations),
173 "registered": len(registered_hooks),
174 "didnt_fire": list(hooks_that_didnt_fire)[:10],
175 },
176 )
178 return BenchmarkResult(
179 name="forward_hooks",
180 severity=BenchmarkSeverity.INFO,
181 message=f"Bridge captured {len(bridge_activations)} forward hook activations",
182 details={"activation_count": len(bridge_activations)},
183 )
185 # Register hooks on reference model
186 reference_hook_points: list[HookPoint] = []
187 for hook_name in hook_names:
188 if hook_name in reference_model.hook_dict:
189 hook_point = reference_model.hook_dict[hook_name]
190 hook_point.add_hook(make_capture_hook(reference_activations, hook_name))
191 reference_hook_points.append(hook_point)
193 # Run reference forward pass
194 with torch.no_grad():
195 if prepend_bos is not None:
196 _ = reference_model(test_text, prepend_bos=prepend_bos)
197 else:
198 _ = reference_model(test_text)
200 # Clean up reference hooks
201 for hook_point in reference_hook_points:
202 hook_point.remove_hooks()
204 # CRITICAL CHECK: Bridge must have all hooks that reference has.
205 # Filter out hooks that bridge models inherently don't have.
206 if missing_from_bridge:
207 missing_from_bridge = filter_expected_missing_hooks(missing_from_bridge)
209 if missing_from_bridge:
210 return BenchmarkResult(
211 name="forward_hooks",
212 severity=BenchmarkSeverity.DANGER,
213 message=f"Bridge is MISSING {len(missing_from_bridge)} hooks that exist in reference model",
214 details={
215 "missing_count": len(missing_from_bridge),
216 "missing_hooks": missing_from_bridge[:20], # Show first 20
217 "total_reference_hooks": len(hook_names),
218 },
219 passed=False,
220 )
222 # CRITICAL CHECK: All registered hooks must fire
223 # Filter out hooks expected to not fire due to architectural differences.
224 if hooks_that_didnt_fire:
225 hooks_that_didnt_fire = set(filter_expected_missing_hooks(hooks_that_didnt_fire))
227 if hooks_that_didnt_fire:
228 return BenchmarkResult(
229 name="forward_hooks",
230 severity=BenchmarkSeverity.DANGER,
231 message=f"{len(hooks_that_didnt_fire)} hooks exist but DIDN'T FIRE during forward pass",
232 details={
233 "didnt_fire_count": len(hooks_that_didnt_fire),
234 "didnt_fire_hooks": list(hooks_that_didnt_fire)[:20],
235 "total_registered": len(registered_hooks),
236 },
237 passed=False,
238 )
240 # Compare activations
241 common_hooks = set(bridge_activations.keys()) & set(reference_activations.keys())
242 mismatches = compare_activation_dicts(
243 bridge_activations, reference_activations, atol=tolerance
244 )
246 if mismatches:
247 # Detect Bloom-style residual-merged hooks: Bloom adds residual inside
248 # attn/MLP modules (dropout_add), so hook_attn_out and hook_mlp_out capture
249 # attn+residual instead of just attn. This is a known HF architectural difference.
250 has_bloom_blocks = any(type(m).__name__ == "BloomBlockBridge" for m in bridge.modules())
251 # Filter out known architectural differences
252 significant_mismatches = [
253 m
254 for m in mismatches
255 if "hook_attn_scores" not in m # Exclude attn_scores which have inf from masking
256 and not (has_bloom_blocks and ("hook_attn_out" in m or "hook_mlp_out" in m))
257 # QK norm hooks: Bridge preserves HF's 4D [batch, heads, seq, d_head]
258 # while HT flattens to [batch*seq*heads, d_head]. This is an intentional
259 # shape convention difference, not a computation error.
260 and "q_norm" not in m and "k_norm" not in m
261 ]
263 if significant_mismatches:
264 return BenchmarkResult(
265 name="forward_hooks",
266 severity=BenchmarkSeverity.DANGER,
267 message=f"Found {len(significant_mismatches)}/{len(common_hooks)} hooks with mismatches",
268 details={
269 "total_hooks": len(common_hooks),
270 "mismatches": len(significant_mismatches),
271 "sample_mismatches": significant_mismatches[:5],
272 },
273 passed=False,
274 )
275 else:
276 return BenchmarkResult(
277 name="forward_hooks",
278 severity=BenchmarkSeverity.WARNING,
279 message=f"All mismatches due to known architectural differences ({len(mismatches)} hooks)",
280 details={"total_hooks": len(common_hooks)},
281 )
283 return BenchmarkResult(
284 name="forward_hooks",
285 severity=BenchmarkSeverity.INFO,
286 message=f"All {len(common_hooks)} forward hooks match within tolerance",
287 details={"hook_count": len(common_hooks), "tolerance": tolerance},
288 )
290 except Exception as e:
291 return BenchmarkResult(
292 name="forward_hooks",
293 severity=BenchmarkSeverity.ERROR,
294 message=f"Forward hooks check failed: {str(e)}",
295 passed=False,
296 )
299# Configuration for cfg-gated attention hooks. Each entry names a config flag
300# and the hook-name stems that should fire on supporting layers when that flag
301# is on. Stems are matched against `hook_dict` keys via substring; both
302# block-level aliases (blocks.N.hook_X) and attn-level primaries
303# (blocks.N.attn.hook_X) are accepted.
304_GATED_HOOK_CONFIGS: list[tuple[str, tuple[str, ...]]] = [
305 ("use_attn_result", ("hook_result",)),
306 ("use_split_qkv_input", ("hook_q_input", "hook_k_input", "hook_v_input")),
307 ("use_attn_in", ("hook_attn_in",)),
308]
311def benchmark_gated_hooks_fire(
312 bridge: TransformerBridge,
313 test_text: str = "The quick brown fox",
314 prepend_bos: Optional[bool] = None,
315) -> BenchmarkResult:
316 """Verify each cfg-gated attention hook fires when its flag is enabled.
318 Hooks like `hook_result`, `hook_q_input`, `hook_attn_in` exist
319 unconditionally on the attention bridge but are only populated when the
320 corresponding config flag is set (keeping default-path cost at zero).
321 This benchmark toggles each flag in turn, runs a short forward, and asserts
322 at least one layer's matching hook actually captured an activation.
324 `use_attn_in` and `use_split_qkv_input` are mutually exclusive, so each
325 flag runs in its own forward pass. Plain `AttentionBridge` (non-PEA/JPEA)
326 adapters raise `NotImplementedError` from the setter — recorded as skipped
327 rather than failed, since the applicability gate is intentional.
328 """
329 try:
330 if not hasattr(bridge, "blocks") or not len(bridge.blocks): 330 ↛ 331line 330 didn't jump to line 331 because the condition on line 330 was never true
331 return BenchmarkResult(
332 name="gated_hooks_fire",
333 severity=BenchmarkSeverity.INFO,
334 message="Bridge has no blocks attribute; gated-hook check skipped",
335 )
337 fired: dict[str, int] = {}
338 skipped: list[tuple[str, str]] = []
339 failed: list[tuple[str, str]] = []
340 tested_flags: list[str] = []
342 for flag_name, hook_stems in _GATED_HOOK_CONFIGS:
343 # Force a clean baseline: all three flags off before toggling.
344 for reset_flag in ("use_attn_result", "use_split_qkv_input", "use_attn_in"):
345 setattr(bridge.cfg, reset_flag, False)
347 setter = getattr(bridge, f"set_{flag_name}", None)
348 if setter is None: 348 ↛ 349line 348 didn't jump to line 349 because the condition on line 348 was never true
349 skipped.append((flag_name, "setter missing on bridge"))
350 continue
351 try:
352 setter(True)
353 except NotImplementedError as e:
354 skipped.append((flag_name, str(e).split("\n", 1)[0][:120]))
355 continue
356 except ValueError as e:
357 # Defensive: mutual-exclusivity shouldn't trigger because we
358 # reset all flags first, but record if something upstream
359 # left the cfg dirty.
360 skipped.append((flag_name, f"setter refused: {e}"))
361 continue
363 tested_flags.append(flag_name)
364 try:
365 activations: dict[str, torch.Tensor] = {}
366 bridge_hook_points: list[HookPoint] = []
367 target_hook_names = [
368 name
369 for name in bridge.hook_dict
370 if any(f".{stem}" in name or name.endswith(stem) for stem in hook_stems)
371 # Exclude the cross-stem substring collisions, e.g. a hook
372 # named "...hook_q_input_foo" — not expected today but be
373 # defensive.
374 and any(name.rsplit(".", 1)[-1] == stem for stem in hook_stems)
375 ]
376 for hname in target_hook_names:
377 hp = bridge.hook_dict[hname]
378 hp.add_hook(make_capture_hook(activations, hname))
379 bridge_hook_points.append(hp)
381 with torch.no_grad():
382 if prepend_bos is not None: 382 ↛ 383line 382 didn't jump to line 383 because the condition on line 382 was never true
383 _ = bridge(test_text, prepend_bos=prepend_bos)
384 else:
385 _ = bridge(test_text)
387 for hp in bridge_hook_points:
388 hp.remove_hooks()
390 # Bucket fired counts per stem.
391 for stem in hook_stems:
392 fired_count = sum(1 for name in activations if name.rsplit(".", 1)[-1] == stem)
393 fired[stem] = fired_count
394 if fired_count == 0 and any(
395 name.rsplit(".", 1)[-1] == stem for name in target_hook_names
396 ):
397 failed.append((flag_name, stem))
398 except NotImplementedError as e:
399 # Some architectures reject the split path only at forward time (e.g.
400 # gated q_proj: Qwen3.5 / Qwen3-Next). Treat like the setter-time gate —
401 # skipped (applicability gate is intentional), not failed.
402 tested_flags.remove(flag_name)
403 skipped.append((flag_name, str(e).split("\n", 1)[0][:120]))
404 finally:
405 setter(False)
407 for reset_flag in ("use_attn_result", "use_split_qkv_input", "use_attn_in"):
408 setattr(bridge.cfg, reset_flag, False)
410 if failed:
411 return BenchmarkResult(
412 name="gated_hooks_fire",
413 severity=BenchmarkSeverity.DANGER,
414 message=f"{len(failed)} gated hooks did not fire when their flag was enabled",
415 details={
416 "failed": failed,
417 "fired_counts": fired,
418 "tested_flags": tested_flags,
419 "skipped": skipped,
420 },
421 passed=False,
422 )
424 if not tested_flags:
425 return BenchmarkResult(
426 name="gated_hooks_fire",
427 severity=BenchmarkSeverity.INFO,
428 message=(
429 "Architecture does not support any gated attention hooks "
430 f"({len(skipped)} flags skipped)"
431 ),
432 details={"skipped": skipped},
433 )
435 msg = (
436 f"All gated hooks fired on their supporting layers "
437 f"({sum(fired.values())} activations across {len(fired)} hook stems"
438 f", {len(tested_flags)} flags tested)"
439 )
440 if skipped: 440 ↛ 441line 440 didn't jump to line 441 because the condition on line 440 was never true
441 msg += f"; {len(skipped)} flags not applicable to this architecture"
442 return BenchmarkResult(
443 name="gated_hooks_fire",
444 severity=BenchmarkSeverity.INFO,
445 message=msg,
446 details={"fired_counts": fired, "tested_flags": tested_flags, "skipped": skipped},
447 )
449 except Exception as e:
450 return BenchmarkResult(
451 name="gated_hooks_fire",
452 severity=BenchmarkSeverity.ERROR,
453 message=f"Gated-hook check failed: {str(e)}",
454 passed=False,
455 )
458def benchmark_critical_forward_hooks(
459 bridge: TransformerBridge,
460 test_text: str,
461 reference_model: Optional[HookedTransformer] = None,
462 tolerance: float = 2e-2,
463) -> BenchmarkResult:
464 """Benchmark critical forward hooks commonly used in interpretability research.
466 Args:
467 bridge: TransformerBridge model to test
468 test_text: Input text for testing
469 reference_model: Optional HookedTransformer reference model
470 tolerance: Tolerance for activation comparison
472 Returns:
473 BenchmarkResult with critical hook comparison details
474 """
475 # Scale tolerance for deep models — numerical precision differences
476 # accumulate through layers, especially for ln_final.hook_normalized
477 # which passes through the entire model. Cap at 3x base to avoid
478 # overly permissive tolerance for very deep models (70B+).
479 n_layers = getattr(bridge.cfg, "n_layers", 1)
480 if n_layers > 12: 480 ↛ 481line 480 didn't jump to line 481 because the condition on line 480 was never true
481 tolerance = min(tolerance * (1 + 0.05 * (n_layers - 12)), tolerance * 3.0)
483 # Critical hooks that are commonly used
484 critical_hooks = [
485 "hook_embed",
486 "hook_pos_embed",
487 "blocks.0.hook_resid_pre",
488 "blocks.0.hook_resid_mid",
489 "blocks.0.hook_resid_post",
490 "blocks.0.attn.hook_q",
491 "blocks.0.attn.hook_k",
492 "blocks.0.attn.hook_v",
493 "blocks.0.attn.hook_z",
494 "blocks.0.attn.hook_result",
495 "blocks.0.mlp.hook_pre",
496 "blocks.0.mlp.hook_post",
497 "blocks.0.hook_mlp_out",
498 "ln_final.hook_normalized",
499 ]
501 try:
502 bridge_activations: Dict[str, torch.Tensor] = {}
504 # Register hooks on bridge
505 bridge_hook_points: list[HookPoint] = []
506 for hook_name in critical_hooks:
507 if hook_name in bridge.hook_dict: 507 ↛ 506line 507 didn't jump to line 506 because the condition on line 507 was always true
508 hook_point = bridge.hook_dict[hook_name]
509 hook_point.add_hook(make_capture_hook(bridge_activations, hook_name))
510 bridge_hook_points.append(hook_point)
512 # Run bridge forward pass
513 with torch.no_grad():
514 _ = bridge(test_text)
516 # Clean up hooks
517 for hook_point in bridge_hook_points:
518 hook_point.remove_hooks()
520 if reference_model is None: 520 ↛ 531line 520 didn't jump to line 531 because the condition on line 520 was always true
521 # No reference - just verify activations were captured
522 captured_count = len(bridge_activations)
523 return BenchmarkResult(
524 name="critical_forward_hooks",
525 severity=BenchmarkSeverity.INFO,
526 message=f"Bridge captured {captured_count}/{len(critical_hooks)} critical hooks",
527 details={"captured": captured_count, "expected": len(critical_hooks)},
528 )
530 # Compare with reference model
531 reference_activations: Dict[str, torch.Tensor] = {}
533 reference_hook_points: list[HookPoint] = []
534 for hook_name in critical_hooks:
535 if hook_name in reference_model.hook_dict:
536 hook_point = reference_model.hook_dict[hook_name]
537 hook_point.add_hook(make_capture_hook(reference_activations, hook_name))
538 reference_hook_points.append(hook_point)
540 # Run reference forward pass
541 with torch.no_grad():
542 _ = reference_model(test_text)
544 # Clean up hooks
545 for hook_point in reference_hook_points:
546 hook_point.remove_hooks()
548 # Compare activations — categorize by presence
549 bridge_missing = [] # Hooks in reference but not in bridge (BAD)
550 reference_missing = [] # Hooks in bridge but not in reference (OK)
552 for hook_name in critical_hooks:
553 if hook_name not in bridge_activations and hook_name not in reference_activations:
554 continue
555 if hook_name not in bridge_activations:
556 bridge_missing.append(f"{hook_name}: Not found in Bridge")
557 continue
558 if hook_name not in reference_activations:
559 reference_missing.append(
560 f"{hook_name}: Not in Reference (Bridge has additional hooks)"
561 )
563 mismatches = compare_activation_dicts(
564 bridge_activations, reference_activations, atol=tolerance
565 )
567 # Filter out hooks expected to be missing in bridge models.
568 if bridge_missing:
569 bridge_missing = filter_expected_missing_hooks(bridge_missing)
571 if bridge_missing:
572 return BenchmarkResult(
573 name="critical_forward_hooks",
574 severity=BenchmarkSeverity.DANGER,
575 message=f"Bridge is missing {len(bridge_missing)} critical hooks that exist in reference",
576 details={"missing_from_bridge": bridge_missing},
577 passed=False,
578 )
580 # Report if reference is missing hooks that bridge has (INFO - bridge has extras)
581 if reference_missing and not mismatches:
582 return BenchmarkResult(
583 name="critical_forward_hooks",
584 severity=BenchmarkSeverity.INFO,
585 message=f"All common hooks match. Bridge has {len(reference_missing)} additional hooks not in reference.",
586 details={
587 "bridge_extras": reference_missing,
588 "compared": len(critical_hooks) - len(reference_missing),
589 },
590 )
592 if mismatches:
593 # Detect Bloom-style residual-merged hooks
594 has_bloom_blocks = any(type(m).__name__ == "BloomBlockBridge" for m in bridge.modules())
595 # Filter out known architectural differences
596 significant_mismatches = [
597 m
598 for m in mismatches
599 if "hook_z" not in m
600 and not (has_bloom_blocks and ("hook_mlp_out" in m or "hook_attn_out" in m))
601 ]
603 if significant_mismatches:
604 return BenchmarkResult(
605 name="critical_forward_hooks",
606 severity=BenchmarkSeverity.DANGER,
607 message=f"Found {len(significant_mismatches)} significant mismatches in critical hooks",
608 details={
609 "mismatches": significant_mismatches[:5],
610 "bridge_extras": reference_missing,
611 },
612 passed=False,
613 )
614 else:
615 return BenchmarkResult(
616 name="critical_forward_hooks",
617 severity=BenchmarkSeverity.WARNING,
618 message="All mismatches due to known architectural differences (hook_z shape)",
619 details={
620 "total_hooks": len(critical_hooks),
621 "bridge_extras": reference_missing,
622 },
623 )
625 compared_count = len(critical_hooks) - len(reference_missing) - len(bridge_missing)
626 return BenchmarkResult(
627 name="critical_forward_hooks",
628 severity=BenchmarkSeverity.INFO,
629 message=f"All {compared_count} common critical hooks match",
630 details={
631 "matched": compared_count,
632 "bridge_extras": len(reference_missing),
633 "skipped": len(bridge_missing),
634 },
635 )
637 except Exception as e:
638 import traceback
640 return BenchmarkResult(
641 name="critical_forward_hooks",
642 severity=BenchmarkSeverity.ERROR,
643 message=f"Critical hooks check failed: {str(e)}",
644 details={
645 "error_type": type(e).__name__,
646 "error_message": str(e),
647 "traceback": traceback.format_exc(),
648 },
649 passed=False,
650 )
653def benchmark_hook_functionality(
654 bridge: TransformerBridge,
655 test_text: str,
656 reference_model: Optional[HookedTransformer] = None,
657 atol: float = 2e-3,
658) -> BenchmarkResult:
659 """Benchmark hook system functionality through ablation effects.
661 Args:
662 bridge: TransformerBridge model to test
663 test_text: Input text for testing
664 reference_model: Optional HookedTransformer reference model
665 atol: Absolute tolerance for effect comparison
667 Returns:
668 BenchmarkResult with hook functionality comparison details
669 """
670 try:
671 # For GQA models, V/K tensors have fewer heads than Q
672 # Use head 0 which always exists, or last head if we want to test a later one
673 # We need to dynamically determine the number of heads available
674 head_to_ablate = 0 # Use first head which always exists
676 def ablation_hook(activation, hook):
677 # Zero out an attention head in layer 0
678 # Clone to avoid in-place modification of autograd views
679 activation = activation.clone()
680 if activation.ndim == 4: 680 ↛ 686line 680 didn't jump to line 686 because the condition on line 680 was always true
681 # Standard: [batch, seq, n_heads, d_head]
682 # For GQA models, the head dimension may be smaller than n_heads
683 n_heads = activation.shape[2]
684 head_idx = min(head_to_ablate, n_heads - 1)
685 activation[:, :, head_idx, :] = 0
686 elif activation.ndim == 3:
687 # Bridge with joint QKV projection (e.g., Phi-3): [batch, seq, d_model]
688 # hook_conversion may not reshape when the underlying linear is a
689 # combined qkv_proj. Zero out a head-sized slice instead.
690 d_model = activation.shape[-1]
691 n_heads = bridge.cfg.n_heads
692 d_head = d_model // n_heads
693 head_idx = min(head_to_ablate, n_heads - 1)
694 start = head_idx * d_head
695 end = start + d_head
696 activation[:, :, start:end] = 0
697 return activation
699 # Test bridge
700 # Encoder-decoder bridges name their stacks; a bare blocks.* hook would
701 # silently no-op and make the ablation vacuous.
702 ablation_target = next(
703 (
704 name
705 for name in (
706 "blocks.0.attn.hook_v",
707 "encoder_blocks.0.attn.hook_v",
708 "decoder_blocks.0.attn.hook_v",
709 )
710 if name in bridge.hook_dict
711 ),
712 "blocks.0.attn.hook_v",
713 )
714 bridge_original = bridge_self_target_loss(bridge, test_text)
715 bridge_ablated = bridge.run_with_hooks(
716 test_text,
717 return_type="loss",
718 labels=bridge.to_tokens(test_text),
719 fwd_hooks=[(ablation_target, ablation_hook)],
720 )
721 bridge_effect = bridge_ablated - bridge_original
723 if reference_model is None: 723 ↛ 742line 723 didn't jump to line 742 because the condition on line 723 was always true
724 # No reference - just verify ablation had an effect
725 effect_magnitude = abs(bridge_effect.item())
726 if effect_magnitude < 1e-6: 726 ↛ 727line 726 didn't jump to line 727 because the condition on line 726 was never true
727 return BenchmarkResult(
728 name="hook_functionality",
729 severity=BenchmarkSeverity.WARNING,
730 message=f"Ablation had minimal effect: {effect_magnitude:.6f}",
731 details={"effect": effect_magnitude},
732 )
734 return BenchmarkResult(
735 name="hook_functionality",
736 severity=BenchmarkSeverity.INFO,
737 message=f"Ablation hook functional with effect: {effect_magnitude:.6f}",
738 details={"effect": effect_magnitude},
739 )
741 # Test reference model
742 reference_original = reference_model(test_text, return_type="loss")
743 reference_ablated = reference_model.run_with_hooks(
744 test_text, return_type="loss", fwd_hooks=[("blocks.0.attn.hook_v", ablation_hook)]
745 )
746 reference_effect = reference_ablated - reference_original
748 return compare_scalars(
749 bridge_effect.item(),
750 reference_effect.item(),
751 atol=atol,
752 name="hook_functionality",
753 )
755 except Exception as e:
756 import traceback
758 return BenchmarkResult(
759 name="hook_functionality",
760 severity=BenchmarkSeverity.ERROR,
761 message=f"Hook functionality check failed: {str(e)}",
762 details={
763 "error_type": type(e).__name__,
764 "error_message": str(e),
765 "traceback": traceback.format_exc(),
766 },
767 passed=False,
768 )