Coverage for transformer_lens/benchmarks/hook_registration.py: 53%
210 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Hook registration and behavior benchmarks for TransformerBridge."""
3from typing import Dict, Iterable, Optional
5import torch
7from transformer_lens.benchmarks.utils import (
8 BenchmarkResult,
9 BenchmarkSeverity,
10 bridge_self_target_loss,
11 compare_activation_dicts,
12 compare_scalars,
13 filter_expected_missing_hooks,
14 make_capture_hook,
15)
16from transformer_lens.hook_points import HookPoint
17from transformer_lens.model_bridge import TransformerBridge
20def benchmark_hook_registry(
21 bridge: TransformerBridge,
22 reference_hooks: Optional[Iterable[str]] = None,
23) -> BenchmarkResult:
24 """Benchmark hook registry completeness.
26 Args:
27 bridge: TransformerBridge model to test
28 reference_hooks: Optional reference hook-name collection (e.g. the keys of
29 a golden hook manifest). Structural self-check only if None.
31 Returns:
32 BenchmarkResult with registry comparison details
33 """
34 try:
35 if reference_hooks 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 the reference hook-name set
61 bridge_hooks = set(bridge.hook_dict.keys())
62 reference_hook_set = set(reference_hooks)
64 common_hooks = bridge_hooks & reference_hook_set
65 missing_hooks = reference_hook_set - bridge_hooks
66 extra_hooks = bridge_hooks - reference_hook_set
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 every reference hook is present in Bridge
88 return BenchmarkResult(
89 name="hook_registry",
90 severity=BenchmarkSeverity.INFO,
91 message=f"All {len(reference_hook_set)} 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_hook_set),
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_activations: Optional[Dict[str, torch.Tensor]] = 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 (must match the snapshot's prompt)
121 reference_activations: Optional reference activations keyed by hook name
122 (e.g. a golden fixture snapshot). Fire-only self-check if None.
123 tolerance: Tolerance for activation matching (fraction of mismatches allowed)
124 prepend_bos: Whether to prepend BOS token. If None, uses model default.
126 Returns:
127 BenchmarkResult with hook activation comparison details
128 """
129 try:
130 bridge_activations: Dict[str, torch.Tensor] = {}
132 # Reference hook names come from the snapshot when provided
133 if reference_activations 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_activations.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_activations is None: 164 ↛ 187line 164 didn't jump to line 187 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 # CRITICAL CHECK: Bridge must have all hooks that reference has.
186 # Filter out hooks that bridge models inherently don't have.
187 if missing_from_bridge:
188 missing_from_bridge = filter_expected_missing_hooks(missing_from_bridge)
190 if missing_from_bridge:
191 return BenchmarkResult(
192 name="forward_hooks",
193 severity=BenchmarkSeverity.DANGER,
194 message=f"Bridge is MISSING {len(missing_from_bridge)} hooks that exist in reference model",
195 details={
196 "missing_count": len(missing_from_bridge),
197 "missing_hooks": missing_from_bridge[:20], # Show first 20
198 "total_reference_hooks": len(hook_names),
199 },
200 passed=False,
201 )
203 # CRITICAL CHECK: All registered hooks must fire
204 # Filter out hooks expected to not fire due to architectural differences.
205 if hooks_that_didnt_fire:
206 hooks_that_didnt_fire = set(filter_expected_missing_hooks(hooks_that_didnt_fire))
208 if hooks_that_didnt_fire:
209 return BenchmarkResult(
210 name="forward_hooks",
211 severity=BenchmarkSeverity.DANGER,
212 message=f"{len(hooks_that_didnt_fire)} hooks exist but DIDN'T FIRE during forward pass",
213 details={
214 "didnt_fire_count": len(hooks_that_didnt_fire),
215 "didnt_fire_hooks": list(hooks_that_didnt_fire)[:20],
216 "total_registered": len(registered_hooks),
217 },
218 passed=False,
219 )
221 # Compare activations
222 common_hooks = set(bridge_activations.keys()) & set(reference_activations.keys())
223 mismatches = compare_activation_dicts(
224 bridge_activations, reference_activations, atol=tolerance
225 )
227 if mismatches:
228 # Detect Bloom-style residual-merged hooks: Bloom adds residual inside
229 # attn/MLP modules (dropout_add), so hook_attn_out and hook_mlp_out capture
230 # attn+residual instead of just attn. This is a known HF architectural difference.
231 has_bloom_blocks = any(type(m).__name__ == "BloomBlockBridge" for m in bridge.modules())
232 # Filter out known architectural differences
233 significant_mismatches = [
234 m
235 for m in mismatches
236 if "hook_attn_scores" not in m # Exclude attn_scores which have inf from masking
237 and not (has_bloom_blocks and ("hook_attn_out" in m or "hook_mlp_out" in m))
238 # QK norm hooks: Bridge preserves HF's 4D [batch, heads, seq, d_head]
239 # while HT flattens to [batch*seq*heads, d_head]. This is an intentional
240 # shape convention difference, not a computation error.
241 and "q_norm" not in m and "k_norm" not in m
242 ]
244 if significant_mismatches:
245 return BenchmarkResult(
246 name="forward_hooks",
247 severity=BenchmarkSeverity.DANGER,
248 message=f"Found {len(significant_mismatches)}/{len(common_hooks)} hooks with mismatches",
249 details={
250 "total_hooks": len(common_hooks),
251 "mismatches": len(significant_mismatches),
252 "sample_mismatches": significant_mismatches[:5],
253 },
254 passed=False,
255 )
256 else:
257 return BenchmarkResult(
258 name="forward_hooks",
259 severity=BenchmarkSeverity.WARNING,
260 message=f"All mismatches due to known architectural differences ({len(mismatches)} hooks)",
261 details={"total_hooks": len(common_hooks)},
262 )
264 return BenchmarkResult(
265 name="forward_hooks",
266 severity=BenchmarkSeverity.INFO,
267 message=f"All {len(common_hooks)} forward hooks match within tolerance",
268 details={"hook_count": len(common_hooks), "tolerance": tolerance},
269 )
271 except Exception as e:
272 return BenchmarkResult(
273 name="forward_hooks",
274 severity=BenchmarkSeverity.ERROR,
275 message=f"Forward hooks check failed: {str(e)}",
276 passed=False,
277 )
280# Configuration for cfg-gated attention hooks. Each entry names a config flag
281# and the hook-name stems that should fire on supporting layers when that flag
282# is on. Stems are matched against `hook_dict` keys via substring; both
283# block-level aliases (blocks.N.hook_X) and attn-level primaries
284# (blocks.N.attn.hook_X) are accepted.
285_GATED_HOOK_CONFIGS: list[tuple[str, tuple[str, ...]]] = [
286 ("use_attn_result", ("hook_result",)),
287 ("use_split_qkv_input", ("hook_q_input", "hook_k_input", "hook_v_input")),
288 ("use_attn_in", ("hook_attn_in",)),
289]
292def benchmark_gated_hooks_fire(
293 bridge: TransformerBridge,
294 test_text: str = "The quick brown fox",
295 prepend_bos: Optional[bool] = None,
296) -> BenchmarkResult:
297 """Verify each cfg-gated attention hook fires when its flag is enabled.
299 Hooks like `hook_result`, `hook_q_input`, `hook_attn_in` exist
300 unconditionally on the attention bridge but are only populated when the
301 corresponding config flag is set (keeping default-path cost at zero).
302 This benchmark toggles each flag in turn, runs a short forward, and asserts
303 at least one layer's matching hook actually captured an activation.
305 `use_attn_in` and `use_split_qkv_input` are mutually exclusive, so each
306 flag runs in its own forward pass. Plain `AttentionBridge` (non-PEA/JPEA)
307 adapters raise `NotImplementedError` from the setter — recorded as skipped
308 rather than failed, since the applicability gate is intentional.
309 """
310 try:
311 if not hasattr(bridge, "blocks") or not len(bridge.blocks): 311 ↛ 312line 311 didn't jump to line 312 because the condition on line 311 was never true
312 return BenchmarkResult(
313 name="gated_hooks_fire",
314 severity=BenchmarkSeverity.INFO,
315 message="Bridge has no blocks attribute; gated-hook check skipped",
316 )
318 fired: dict[str, int] = {}
319 skipped: list[tuple[str, str]] = []
320 failed: list[tuple[str, str]] = []
321 tested_flags: list[str] = []
323 for flag_name, hook_stems in _GATED_HOOK_CONFIGS:
324 # Force a clean baseline: all three flags off before toggling.
325 for reset_flag in ("use_attn_result", "use_split_qkv_input", "use_attn_in"):
326 setattr(bridge.cfg, reset_flag, False)
328 setter = getattr(bridge, f"set_{flag_name}", None)
329 if setter is None: 329 ↛ 330line 329 didn't jump to line 330 because the condition on line 329 was never true
330 skipped.append((flag_name, "setter missing on bridge"))
331 continue
332 try:
333 setter(True)
334 except NotImplementedError as e:
335 skipped.append((flag_name, str(e).split("\n", 1)[0][:120]))
336 continue
337 except ValueError as e:
338 # Defensive: mutual-exclusivity shouldn't trigger because we
339 # reset all flags first, but record if something upstream
340 # left the cfg dirty.
341 skipped.append((flag_name, f"setter refused: {e}"))
342 continue
344 tested_flags.append(flag_name)
345 try:
346 activations: dict[str, torch.Tensor] = {}
347 bridge_hook_points: list[HookPoint] = []
348 target_hook_names = [
349 name
350 for name in bridge.hook_dict
351 if any(f".{stem}" in name or name.endswith(stem) for stem in hook_stems)
352 # Exclude the cross-stem substring collisions, e.g. a hook
353 # named "...hook_q_input_foo" — not expected today but be
354 # defensive.
355 and any(name.rsplit(".", 1)[-1] == stem for stem in hook_stems)
356 ]
357 for hname in target_hook_names:
358 hp = bridge.hook_dict[hname]
359 hp.add_hook(make_capture_hook(activations, hname))
360 bridge_hook_points.append(hp)
362 with torch.no_grad():
363 if prepend_bos is not None: 363 ↛ 364line 363 didn't jump to line 364 because the condition on line 363 was never true
364 _ = bridge(test_text, prepend_bos=prepend_bos)
365 else:
366 _ = bridge(test_text)
368 for hp in bridge_hook_points:
369 hp.remove_hooks()
371 # Bucket fired counts per stem.
372 for stem in hook_stems:
373 fired_count = sum(1 for name in activations if name.rsplit(".", 1)[-1] == stem)
374 fired[stem] = fired_count
375 if fired_count == 0 and any(
376 name.rsplit(".", 1)[-1] == stem for name in target_hook_names
377 ):
378 failed.append((flag_name, stem))
379 except NotImplementedError as e:
380 # Some architectures reject the split path only at forward time (e.g.
381 # gated q_proj: Qwen3.5 / Qwen3-Next). Treat like the setter-time gate —
382 # skipped (applicability gate is intentional), not failed.
383 tested_flags.remove(flag_name)
384 skipped.append((flag_name, str(e).split("\n", 1)[0][:120]))
385 finally:
386 setter(False)
388 for reset_flag in ("use_attn_result", "use_split_qkv_input", "use_attn_in"):
389 setattr(bridge.cfg, reset_flag, False)
391 if failed:
392 return BenchmarkResult(
393 name="gated_hooks_fire",
394 severity=BenchmarkSeverity.DANGER,
395 message=f"{len(failed)} gated hooks did not fire when their flag was enabled",
396 details={
397 "failed": failed,
398 "fired_counts": fired,
399 "tested_flags": tested_flags,
400 "skipped": skipped,
401 },
402 passed=False,
403 )
405 if not tested_flags:
406 return BenchmarkResult(
407 name="gated_hooks_fire",
408 severity=BenchmarkSeverity.INFO,
409 message=(
410 "Architecture does not support any gated attention hooks "
411 f"({len(skipped)} flags skipped)"
412 ),
413 details={"skipped": skipped},
414 )
416 msg = (
417 f"All gated hooks fired on their supporting layers "
418 f"({sum(fired.values())} activations across {len(fired)} hook stems"
419 f", {len(tested_flags)} flags tested)"
420 )
421 if skipped: 421 ↛ 422line 421 didn't jump to line 422 because the condition on line 421 was never true
422 msg += f"; {len(skipped)} flags not applicable to this architecture"
423 return BenchmarkResult(
424 name="gated_hooks_fire",
425 severity=BenchmarkSeverity.INFO,
426 message=msg,
427 details={"fired_counts": fired, "tested_flags": tested_flags, "skipped": skipped},
428 )
430 except Exception as e:
431 return BenchmarkResult(
432 name="gated_hooks_fire",
433 severity=BenchmarkSeverity.ERROR,
434 message=f"Gated-hook check failed: {str(e)}",
435 passed=False,
436 )
439def benchmark_critical_forward_hooks(
440 bridge: TransformerBridge,
441 test_text: str,
442 reference_activations: Optional[Dict[str, torch.Tensor]] = None,
443 tolerance: float = 2e-2,
444) -> BenchmarkResult:
445 """Benchmark critical forward hooks commonly used in interpretability research.
447 Args:
448 bridge: TransformerBridge model to test
449 test_text: Input text for testing (must match the snapshot's prompt)
450 reference_activations: Optional reference activations keyed by hook name
451 (e.g. a golden fixture snapshot). Capture-only self-check if None.
452 tolerance: Tolerance for activation comparison
454 Returns:
455 BenchmarkResult with critical hook comparison details
456 """
457 # Scale tolerance for deep models — numerical precision differences
458 # accumulate through layers, especially for ln_final.hook_normalized
459 # which passes through the entire model. Cap at 3x base to avoid
460 # overly permissive tolerance for very deep models (70B+).
461 n_layers = getattr(bridge.cfg, "n_layers", 1)
462 if n_layers > 12: 462 ↛ 463line 462 didn't jump to line 463 because the condition on line 462 was never true
463 tolerance = min(tolerance * (1 + 0.05 * (n_layers - 12)), tolerance * 3.0)
465 # Critical hooks that are commonly used
466 critical_hooks = [
467 "hook_embed",
468 "hook_pos_embed",
469 "blocks.0.hook_resid_pre",
470 "blocks.0.hook_resid_mid",
471 "blocks.0.hook_resid_post",
472 "blocks.0.attn.hook_q",
473 "blocks.0.attn.hook_k",
474 "blocks.0.attn.hook_v",
475 "blocks.0.attn.hook_z",
476 "blocks.0.attn.hook_result",
477 "blocks.0.mlp.hook_pre",
478 "blocks.0.mlp.hook_post",
479 "blocks.0.hook_mlp_out",
480 "ln_final.hook_normalized",
481 ]
483 try:
484 bridge_activations: Dict[str, torch.Tensor] = {}
486 # Register hooks on bridge
487 bridge_hook_points: list[HookPoint] = []
488 for hook_name in critical_hooks:
489 if hook_name in bridge.hook_dict: 489 ↛ 488line 489 didn't jump to line 488 because the condition on line 489 was always true
490 hook_point = bridge.hook_dict[hook_name]
491 hook_point.add_hook(make_capture_hook(bridge_activations, hook_name))
492 bridge_hook_points.append(hook_point)
494 # Run bridge forward pass
495 with torch.no_grad():
496 _ = bridge(test_text)
498 # Clean up hooks
499 for hook_point in bridge_hook_points:
500 hook_point.remove_hooks()
502 if reference_activations is None: 502 ↛ 513line 502 didn't jump to line 513 because the condition on line 502 was always true
503 # No reference - just verify activations were captured
504 captured_count = len(bridge_activations)
505 return BenchmarkResult(
506 name="critical_forward_hooks",
507 severity=BenchmarkSeverity.INFO,
508 message=f"Bridge captured {captured_count}/{len(critical_hooks)} critical hooks",
509 details={"captured": captured_count, "expected": len(critical_hooks)},
510 )
512 # Compare activations — categorize by presence
513 bridge_missing = [] # Hooks in reference but not in bridge (BAD)
514 reference_missing = [] # Hooks in bridge but not in reference (OK)
516 for hook_name in critical_hooks:
517 if hook_name not in bridge_activations and hook_name not in reference_activations:
518 continue
519 if hook_name not in bridge_activations:
520 bridge_missing.append(f"{hook_name}: Not found in Bridge")
521 continue
522 if hook_name not in reference_activations:
523 reference_missing.append(
524 f"{hook_name}: Not in Reference (Bridge has additional hooks)"
525 )
527 mismatches = compare_activation_dicts(
528 bridge_activations, reference_activations, atol=tolerance
529 )
531 # Filter out hooks expected to be missing in bridge models.
532 if bridge_missing:
533 bridge_missing = filter_expected_missing_hooks(bridge_missing)
535 if bridge_missing:
536 return BenchmarkResult(
537 name="critical_forward_hooks",
538 severity=BenchmarkSeverity.DANGER,
539 message=f"Bridge is missing {len(bridge_missing)} critical hooks that exist in reference",
540 details={"missing_from_bridge": bridge_missing},
541 passed=False,
542 )
544 # Report if reference is missing hooks that bridge has (INFO - bridge has extras)
545 if reference_missing and not mismatches:
546 return BenchmarkResult(
547 name="critical_forward_hooks",
548 severity=BenchmarkSeverity.INFO,
549 message=f"All common hooks match. Bridge has {len(reference_missing)} additional hooks not in reference.",
550 details={
551 "bridge_extras": reference_missing,
552 "compared": len(critical_hooks) - len(reference_missing),
553 },
554 )
556 if mismatches:
557 # Detect Bloom-style residual-merged hooks
558 has_bloom_blocks = any(type(m).__name__ == "BloomBlockBridge" for m in bridge.modules())
559 # Filter out known architectural differences
560 significant_mismatches = [
561 m
562 for m in mismatches
563 if "hook_z" not in m
564 and not (has_bloom_blocks and ("hook_mlp_out" in m or "hook_attn_out" in m))
565 ]
567 if significant_mismatches:
568 return BenchmarkResult(
569 name="critical_forward_hooks",
570 severity=BenchmarkSeverity.DANGER,
571 message=f"Found {len(significant_mismatches)} significant mismatches in critical hooks",
572 details={
573 "mismatches": significant_mismatches[:5],
574 "bridge_extras": reference_missing,
575 },
576 passed=False,
577 )
578 else:
579 return BenchmarkResult(
580 name="critical_forward_hooks",
581 severity=BenchmarkSeverity.WARNING,
582 message="All mismatches due to known architectural differences (hook_z shape)",
583 details={
584 "total_hooks": len(critical_hooks),
585 "bridge_extras": reference_missing,
586 },
587 )
589 compared_count = len(critical_hooks) - len(reference_missing) - len(bridge_missing)
590 return BenchmarkResult(
591 name="critical_forward_hooks",
592 severity=BenchmarkSeverity.INFO,
593 message=f"All {compared_count} common critical hooks match",
594 details={
595 "matched": compared_count,
596 "bridge_extras": len(reference_missing),
597 "skipped": len(bridge_missing),
598 },
599 )
601 except Exception as e:
602 import traceback
604 return BenchmarkResult(
605 name="critical_forward_hooks",
606 severity=BenchmarkSeverity.ERROR,
607 message=f"Critical hooks check failed: {str(e)}",
608 details={
609 "error_type": type(e).__name__,
610 "error_message": str(e),
611 "traceback": traceback.format_exc(),
612 },
613 passed=False,
614 )
617def benchmark_hook_functionality(
618 bridge: TransformerBridge,
619 test_text: str,
620 reference_effect: Optional[float] = None,
621 atol: float = 2e-3,
622) -> BenchmarkResult:
623 """Benchmark hook system functionality through ablation effects.
625 Args:
626 bridge: TransformerBridge model to test
627 test_text: Input text for testing (must match the reference's prompt)
628 reference_effect: Optional reference ablation loss-delta (e.g. from a
629 golden fixture). Effect-only self-check if None.
630 atol: Absolute tolerance for effect comparison
632 Returns:
633 BenchmarkResult with hook functionality comparison details
634 """
635 try:
636 # For GQA models, V/K tensors have fewer heads than Q
637 # Use head 0 which always exists, or last head if we want to test a later one
638 # We need to dynamically determine the number of heads available
639 head_to_ablate = 0 # Use first head which always exists
641 def ablation_hook(activation, hook):
642 # Zero out an attention head in layer 0
643 # Clone to avoid in-place modification of autograd views
644 activation = activation.clone()
645 if activation.ndim == 4: 645 ↛ 651line 645 didn't jump to line 651 because the condition on line 645 was always true
646 # Standard: [batch, seq, n_heads, d_head]
647 # For GQA models, the head dimension may be smaller than n_heads
648 n_heads = activation.shape[2]
649 head_idx = min(head_to_ablate, n_heads - 1)
650 activation[:, :, head_idx, :] = 0
651 elif activation.ndim == 3:
652 # Bridge with joint QKV projection (e.g., Phi-3): [batch, seq, d_model]
653 # hook_conversion may not reshape when the underlying linear is a
654 # combined qkv_proj. Zero out a head-sized slice instead.
655 d_model = activation.shape[-1]
656 n_heads = bridge.cfg.n_heads
657 d_head = d_model // n_heads
658 head_idx = min(head_to_ablate, n_heads - 1)
659 start = head_idx * d_head
660 end = start + d_head
661 activation[:, :, start:end] = 0
662 return activation
664 # Test bridge
665 # Encoder-decoder bridges name their stacks; a bare blocks.* hook would
666 # silently no-op and make the ablation vacuous.
667 ablation_target = next(
668 (
669 name
670 for name in (
671 "blocks.0.attn.hook_v",
672 "encoder_blocks.0.attn.hook_v",
673 "decoder_blocks.0.attn.hook_v",
674 )
675 if name in bridge.hook_dict
676 ),
677 "blocks.0.attn.hook_v",
678 )
679 bridge_original = bridge_self_target_loss(bridge, test_text)
680 bridge_ablated = bridge.run_with_hooks(
681 test_text,
682 return_type="loss",
683 labels=bridge.to_tokens(test_text),
684 fwd_hooks=[(ablation_target, ablation_hook)],
685 )
686 bridge_effect = bridge_ablated - bridge_original
688 if reference_effect is None: 688 ↛ 706line 688 didn't jump to line 706 because the condition on line 688 was always true
689 # No reference - just verify ablation had an effect
690 effect_magnitude = abs(bridge_effect.item())
691 if effect_magnitude < 1e-6: 691 ↛ 692line 691 didn't jump to line 692 because the condition on line 691 was never true
692 return BenchmarkResult(
693 name="hook_functionality",
694 severity=BenchmarkSeverity.WARNING,
695 message=f"Ablation had minimal effect: {effect_magnitude:.6f}",
696 details={"effect": effect_magnitude},
697 )
699 return BenchmarkResult(
700 name="hook_functionality",
701 severity=BenchmarkSeverity.INFO,
702 message=f"Ablation hook functional with effect: {effect_magnitude:.6f}",
703 details={"effect": effect_magnitude},
704 )
706 return compare_scalars(
707 bridge_effect.item(),
708 reference_effect,
709 atol=atol,
710 name="hook_functionality",
711 )
713 except Exception as e:
714 import traceback
716 return BenchmarkResult(
717 name="hook_functionality",
718 severity=BenchmarkSeverity.ERROR,
719 message=f"Hook functionality check failed: {str(e)}",
720 details={
721 "error_type": type(e).__name__,
722 "error_message": str(e),
723 "traceback": traceback.format_exc(),
724 },
725 passed=False,
726 )