Coverage for transformer_lens/benchmarks/backward_gradients.py: 32%
220 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Backward gradient 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 make_grad_capture_hook,
12 safe_allclose,
13)
14from transformer_lens.hook_points import HookPoint
15from transformer_lens.model_bridge import TransformerBridge
17# Grading band for numerical (non-convention) gradient mismatches. Registering
18# backward hooks forces normalization off HF's native autograd onto the python
19# norm, which shifts results at float-rounding scale; measured noise is ~1e-5
20# rel_l2 with a single over-tolerance element, while injected bugs start at
21# ~1e-3 rel_l2 with 60+ elements over. Valid for fp32 gradients only — the
22# gradient section upcasts reduced-precision models before comparing.
23REL_L2_TOLERANCE = 1e-4
24OVER_TOLERANCE_MAX_ELEMENTS = 3
27def needs_fp32_gradients(dtype: Optional[torch.dtype]) -> bool:
28 """Reduced-precision gradients cannot be graded against the fp32-calibrated
29 band — bf16's rounding floor alone is ~2e-3 rel_l2, inside the bug band."""
30 return dtype is not None and dtype not in (torch.float32, torch.float64)
33def gradient_mismatch_stats(
34 bridge_finite: torch.Tensor,
35 reference_finite: torch.Tensor,
36 abs_tolerance: float,
37 rel_tolerance: float,
38) -> dict:
39 """Scale-aware statistics for grading one recorded gradient mismatch.
41 A zero reference with a nonzero bridge gradient is the maximally divergent
42 case, not perfect agreement, so rel_l2 is inf there rather than 0.
43 """
44 bf, rf = bridge_finite.float(), reference_finite.float()
45 ref_norm = torch.norm(rf)
46 diff_norm = torch.norm(bf - rf)
47 if ref_norm > 0:
48 rel_l2 = (diff_norm / ref_norm).item()
49 else:
50 rel_l2 = 0.0 if diff_norm == 0 else float("inf")
51 over_count = int(
52 (torch.abs(bf - rf) > abs_tolerance + rel_tolerance * torch.abs(rf)).sum().item()
53 )
54 return {"rel_l2": rel_l2, "over_count": over_count}
57def gradient_mismatch_is_numerical_noise(rel_l2: float, over_count: int) -> bool:
58 """True when a gradient mismatch is diffuse and tiny rather than a divergence.
60 Elementwise worst-case cannot separate the two: one element of 55k crossing
61 the tolerance scores the same as a head scaled by 1%. rel_l2 separates them
62 by 58x or more, and the element COUNT guards the localized case rel_l2 would
63 dilute. A count (not a fraction) keeps the band reachable on small tensors:
64 detection guarantees count >= 1, so a fractional guard of 1e-4 was
65 arithmetically unsatisfiable below 10,000 elements (gemma-3-270m's MQA
66 hook_rot_k is 6,912).
67 """
68 return rel_l2 <= REL_L2_TOLERANCE and over_count <= OVER_TOLERANCE_MAX_ELEMENTS
71def benchmark_backward_hooks(
72 bridge: TransformerBridge,
73 test_text: str,
74 reference_model: Optional[HookedTransformer] = None,
75 abs_tolerance: float = 0.2,
76 rel_tolerance: float = 3e-4,
77) -> BenchmarkResult:
78 """Benchmark all backward hooks for gradient matching.
80 Args:
81 bridge: TransformerBridge model to test
82 test_text: Input text for testing
83 reference_model: Optional HookedTransformer reference model
84 abs_tolerance: Absolute tolerance for gradient comparison
85 rel_tolerance: Relative tolerance for gradient comparison
87 Returns:
88 BenchmarkResult with backward hook comparison details
89 """
90 try:
91 bridge_gradients: Dict[str, torch.Tensor] = {}
92 reference_gradients: Dict[str, torch.Tensor] = {}
94 # Get all hook names
95 if reference_model is not None: 95 ↛ 96line 95 didn't jump to line 96 because the condition on line 95 was never true
96 hook_names = list(reference_model.hook_dict.keys())
97 else:
98 hook_names = list(bridge._hook_registry.keys())
100 # Register backward hooks on bridge
101 bridge_hook_points: list[HookPoint] = []
102 for hook_name in hook_names:
103 if hook_name in bridge.hook_dict: 103 ↛ 102line 103 didn't jump to line 102 because the condition on line 103 was always true
104 hook_point = bridge.hook_dict[hook_name]
105 hook_point.add_hook(
106 make_grad_capture_hook(bridge_gradients, hook_name, return_none=True),
107 dir="bwd",
108 )
109 bridge_hook_points.append(hook_point)
111 # Run bridge forward and backward
112 bridge_output = bridge(test_text)
113 bridge_loss = bridge_output[:, -1, :].sum()
114 bridge_loss.backward()
116 # Clean up hooks
117 for hook_point in bridge_hook_points:
118 hook_point.remove_hooks(dir="bwd")
120 if reference_model is None: 120 ↛ 136line 120 didn't jump to line 136 because the condition on line 120 was always true
121 # No reference - just verify gradients were captured
122 result = BenchmarkResult(
123 name="backward_hooks",
124 severity=BenchmarkSeverity.INFO,
125 message=f"Bridge captured {len(bridge_gradients)} backward hook gradients",
126 details={"gradient_count": len(bridge_gradients)},
127 )
129 # Clear model gradients (variables will be GC'd when function returns)
130 if hasattr(bridge, "zero_grad"): 130 ↛ 133line 130 didn't jump to line 133 because the condition on line 130 was always true
131 bridge.zero_grad()
133 return result
135 # Register backward hooks on reference model
136 reference_hook_points: list[HookPoint] = []
137 for hook_name in hook_names:
138 if hook_name in reference_model.hook_dict:
139 hook_point = reference_model.hook_dict[hook_name]
140 hook_point.add_hook(
141 make_grad_capture_hook(reference_gradients, hook_name, return_none=True),
142 dir="bwd",
143 )
144 reference_hook_points.append(hook_point)
146 # Run reference forward and backward
147 reference_output = reference_model(test_text)
148 reference_loss = reference_output[:, -1, :].sum()
149 reference_loss.backward()
151 # Clean up hooks
152 for hook_point in reference_hook_points:
153 hook_point.remove_hooks(dir="bwd")
155 # Compare gradients
156 common_hooks = set(bridge_gradients.keys()) & set(reference_gradients.keys())
158 # Hooks with known numerical differences due to architectural bridging
159 excluded_hooks = [
160 "blocks.0.attn.hook_pattern",
161 "blocks.0.attn.hook_z",
162 "blocks.0.hook_resid_pre",
163 "blocks.0.ln1.hook_scale",
164 "blocks.0.ln2.hook_normalized",
165 "blocks.3.mlp.hook_post",
166 "blocks.4.attn.hook_pattern",
167 "blocks.6.attn.hook_pattern",
168 "blocks.7.ln2.hook_scale",
169 "hook_embed",
170 "hook_pos_embed",
171 "blocks.1.attn.hook_pattern",
172 ]
174 mismatches = []
175 mismatch_stats: dict = {}
176 for hook_name in sorted(common_hooks):
177 if hook_name in excluded_hooks:
178 continue
180 bridge_grad = bridge_gradients[hook_name]
181 reference_grad = reference_gradients[hook_name]
183 # Check shapes
184 if bridge_grad.shape != reference_grad.shape:
185 mismatches.append(
186 f"{hook_name}: Shape mismatch - Bridge{bridge_grad.shape} vs Ref{reference_grad.shape}"
187 )
188 continue
190 # Handle special cases with inf or nan
191 bridge_finite = bridge_grad[torch.isfinite(bridge_grad)]
192 reference_finite = reference_grad[torch.isfinite(reference_grad)]
194 if bridge_finite.numel() > 0 and reference_finite.numel() > 0:
195 # Compare finite values
196 if not safe_allclose(
197 bridge_finite, reference_finite, atol=abs_tolerance, rtol=rel_tolerance
198 ):
199 bf = bridge_finite.float()
200 rf = reference_finite.float()
201 max_diff = torch.max(torch.abs(bf - rf)).item()
202 mean_diff = torch.mean(torch.abs(bf - rf)).item()
203 rel_diff = torch.abs(bf - rf) / (torch.abs(bf) + 1e-8)
204 mean_rel = rel_diff.mean().item()
205 # Scale-aware stats for grading. Elementwise worst-case alone
206 # cannot separate a real divergence from the float-rounding
207 # shift the python-norm fallback introduces when backward
208 # hooks force normalization off HF's native autograd path.
209 stats = gradient_mismatch_stats(bf, rf, abs_tolerance, rel_tolerance)
210 mismatch_stats[hook_name] = stats
211 mismatches.append(
212 f"{hook_name}: Value mismatch - max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}, "
213 f"mean_rel={mean_rel:.6f}, rel_l2={stats['rel_l2']:.3e}, "
214 f"over_count={stats['over_count']}"
215 )
217 tested_hooks = len(common_hooks) - len(excluded_hooks)
218 matching_hooks = tested_hooks - len(mismatches)
220 if mismatches:
221 # Check if mismatches are acceptable patterns
222 acceptable_patterns = [
223 "hook_attn_scores",
224 "hook_z",
225 "hook_pattern",
226 "hook_attn_out",
227 "hook_v",
228 "hook_q",
229 "hook_k",
230 "q_norm", # QK norm: Bridge uses 4D, HT uses 2D (shape convention)
231 "k_norm", # QK norm: Bridge uses 4D, HT uses 2D (shape convention)
232 "ln1.hook_",
233 "ln2.hook_",
234 # Sandwich norms (gemma-2/3): same class as ln1/ln2 above, which
235 # predate them.
236 "ln1_post.hook_",
237 "ln2_post.hook_",
238 "ln_final.hook_",
239 "hook_resid_mid",
240 "hook_resid_pre",
241 "hook_resid_post",
242 "hook_embed",
243 "hook_pos_embed",
244 "unembed.hook_",
245 "mlp.hook_post",
246 "mlp.hook_pre",
247 "hook_mlp_out",
248 ]
250 def within_noise_band(entry: str) -> bool:
251 """Diffuse, tiny deviation — the fallback's rounding, not a divergence.
253 Measured noise across architectures is rel_l2 ~1e-5 with a single
254 over-tolerance element on the rotary hooks (the only ones outside
255 the pattern list); injected bugs of a 1% head scale or a 0.1%
256 uniform scale land at rel_l2 1e-3+ with 60+ elements over.
257 """
258 name = entry.split(":")[0]
259 stats = mismatch_stats.get(name)
260 if stats is None:
261 return False
262 return gradient_mismatch_is_numerical_noise(stats["rel_l2"], stats["over_count"])
264 acceptable_mismatches = [
265 m
266 for m in mismatches
267 if any(pattern in m for pattern in acceptable_patterns) or within_noise_band(m)
268 ]
270 if len(acceptable_mismatches) == len(mismatches):
271 result = BenchmarkResult(
272 name="backward_hooks",
273 severity=BenchmarkSeverity.WARNING,
274 message=f"All mismatches due to known architectural differences ({len(mismatches)} hooks)",
275 details={
276 "total_hooks": tested_hooks,
277 "matching": matching_hooks,
278 "excluded": len(excluded_hooks),
279 },
280 )
282 # Clear model gradients (variables will be GC'd when function returns)
283 if hasattr(bridge, "zero_grad"):
284 bridge.zero_grad()
285 if hasattr(reference_model, "zero_grad"):
286 reference_model.zero_grad()
288 return result
289 else:
290 significant_mismatches = [m for m in mismatches if m not in acceptable_mismatches]
291 result = BenchmarkResult(
292 name="backward_hooks",
293 severity=BenchmarkSeverity.DANGER,
294 message=f"Found {len(significant_mismatches)} significant numerical mismatches",
295 details={
296 "total_hooks": tested_hooks,
297 "mismatches": len(significant_mismatches),
298 "sample_mismatches": significant_mismatches[:5],
299 },
300 passed=False,
301 )
303 # Clear model gradients (variables will be GC'd when function returns)
304 if hasattr(bridge, "zero_grad"):
305 bridge.zero_grad()
306 if hasattr(reference_model, "zero_grad"):
307 reference_model.zero_grad()
309 return result
311 result = BenchmarkResult(
312 name="backward_hooks",
313 severity=BenchmarkSeverity.INFO,
314 message=f"All {matching_hooks}/{tested_hooks} hooks match within tolerance",
315 details={
316 "matching_hooks": matching_hooks,
317 "tested_hooks": tested_hooks,
318 "excluded": len(excluded_hooks),
319 "abs_tolerance": abs_tolerance,
320 "rel_tolerance": rel_tolerance,
321 },
322 )
324 # Clear model gradients (variables will be GC'd when function returns)
325 if hasattr(bridge, "zero_grad"):
326 bridge.zero_grad()
327 if reference_model is not None and hasattr(reference_model, "zero_grad"):
328 reference_model.zero_grad()
330 return result
332 except Exception as e:
333 import traceback
335 return BenchmarkResult(
336 name="backward_hooks",
337 severity=BenchmarkSeverity.ERROR,
338 message=f"Backward hooks check failed: {str(e)}",
339 details={
340 "error_type": type(e).__name__,
341 "error_message": str(e),
342 "traceback": traceback.format_exc(),
343 },
344 passed=False,
345 )
348def benchmark_critical_backward_hooks(
349 bridge: TransformerBridge,
350 test_text: str,
351 reference_model: Optional[HookedTransformer] = None,
352 abs_tolerance: float = 0.2,
353 rel_tolerance: float = 3e-4,
354) -> BenchmarkResult:
355 """Benchmark critical backward hooks for gradient matching.
357 Args:
358 bridge: TransformerBridge model to test
359 test_text: Input text for testing
360 reference_model: Optional HookedTransformer reference model
361 abs_tolerance: Absolute tolerance for gradient comparison
362 rel_tolerance: Relative tolerance for gradient comparison
364 Returns:
365 BenchmarkResult with critical backward hook comparison details
366 """
367 critical_hooks = [
368 "hook_embed",
369 "blocks.0.hook_resid_pre",
370 "blocks.0.hook_resid_mid",
371 "blocks.0.hook_resid_post",
372 "blocks.0.attn.hook_q",
373 "blocks.0.attn.hook_k",
374 "blocks.0.attn.hook_v",
375 "blocks.0.attn.hook_z",
376 "blocks.0.attn.hook_result",
377 "blocks.0.mlp.hook_pre",
378 "blocks.0.mlp.hook_post",
379 "blocks.0.hook_mlp_out",
380 ]
382 try:
383 bridge_gradients: Dict[str, torch.Tensor] = {}
385 # Register backward hooks on bridge
386 bridge_hook_points: list[HookPoint] = []
387 for hook_name in critical_hooks:
388 if hook_name in bridge.hook_dict: 388 ↛ 387line 388 didn't jump to line 387 because the condition on line 388 was always true
389 hook_point = bridge.hook_dict[hook_name]
390 hook_point.add_hook(
391 make_grad_capture_hook(bridge_gradients, hook_name, return_none=True),
392 dir="bwd",
393 )
394 bridge_hook_points.append(hook_point)
396 # Run bridge forward and backward
397 bridge_output = bridge(test_text)
398 bridge_loss = bridge_output[:, -1, :].sum()
399 bridge_loss.backward()
401 # Clean up hooks
402 for hook_point in bridge_hook_points:
403 hook_point.remove_hooks(dir="bwd")
405 if reference_model is None: 405 ↛ 422line 405 didn't jump to line 422 because the condition on line 405 was always true
406 # No reference - just verify gradients were captured
407 captured_count = len(bridge_gradients)
408 result = BenchmarkResult(
409 name="critical_backward_hooks",
410 severity=BenchmarkSeverity.INFO,
411 message=f"Bridge captured {captured_count}/{len(critical_hooks)} critical backward gradients",
412 details={"captured": captured_count, "expected": len(critical_hooks)},
413 )
415 # Clear model gradients (variables will be GC'd when function returns)
416 if hasattr(bridge, "zero_grad"): 416 ↛ 419line 416 didn't jump to line 419 because the condition on line 416 was always true
417 bridge.zero_grad()
419 return result
421 # Register backward hooks on reference model
422 reference_gradients: Dict[str, torch.Tensor] = {}
424 reference_hook_points: list[HookPoint] = []
425 for hook_name in critical_hooks:
426 if hook_name in reference_model.hook_dict:
427 hook_point = reference_model.hook_dict[hook_name]
428 hook_point.add_hook(
429 make_grad_capture_hook(reference_gradients, hook_name, return_none=True),
430 dir="bwd",
431 )
432 reference_hook_points.append(hook_point)
434 # Run reference forward and backward
435 reference_output = reference_model(test_text)
436 reference_loss = reference_output[:, -1, :].sum()
437 reference_loss.backward()
439 # Clean up hooks
440 for hook_point in reference_hook_points:
441 hook_point.remove_hooks(dir="bwd")
443 # Compare gradients
444 mismatches = []
445 for hook_name in critical_hooks:
446 if hook_name not in bridge_gradients:
447 continue
448 if hook_name not in reference_gradients:
449 continue
451 bridge_grad = bridge_gradients[hook_name]
452 reference_grad = reference_gradients[hook_name]
454 # Check shapes
455 if bridge_grad.shape != reference_grad.shape:
456 mismatches.append(
457 f"{hook_name}: Shape mismatch - Bridge{bridge_grad.shape} vs Ref{reference_grad.shape}"
458 )
459 continue
461 # Compare only finite values
462 bridge_finite = bridge_grad[torch.isfinite(bridge_grad)]
463 reference_finite = reference_grad[torch.isfinite(reference_grad)]
465 if bridge_finite.numel() > 0 and reference_finite.numel() > 0:
466 if not safe_allclose(
467 bridge_finite, reference_finite, atol=abs_tolerance, rtol=rel_tolerance
468 ):
469 max_diff = torch.max(
470 torch.abs(bridge_finite.float() - reference_finite.float())
471 ).item()
472 mismatches.append(f"{hook_name}: max_diff={max_diff:.6f}")
474 if mismatches:
475 # Filter out known architectural differences
476 acceptable_patterns = [
477 "hook_z",
478 "hook_attn_scores",
479 "hook_pattern",
480 "hook_result",
481 "hook_v",
482 "hook_q",
483 "hook_k",
484 "q_norm", # QK norm: Bridge uses 4D, HT uses 2D (shape convention)
485 "k_norm", # QK norm: Bridge uses 4D, HT uses 2D (shape convention)
486 "ln1.hook_",
487 "ln2.hook_",
488 # Sandwich norms (gemma-2/3): same class as ln1/ln2 above.
489 "ln1_post.hook_",
490 "ln2_post.hook_",
491 "hook_resid_pre",
492 "hook_resid_mid",
493 "hook_resid_post",
494 "hook_embed",
495 "mlp.hook_post",
496 "mlp.hook_pre",
497 "hook_mlp_out",
498 ]
499 significant_mismatches = [
500 m for m in mismatches if not any(pattern in m for pattern in acceptable_patterns)
501 ]
503 if significant_mismatches:
504 result = BenchmarkResult(
505 name="critical_backward_hooks",
506 severity=BenchmarkSeverity.DANGER,
507 message=f"Found {len(significant_mismatches)} significant mismatches in critical hooks",
508 details={"mismatches": significant_mismatches[:5]},
509 passed=False,
510 )
511 else:
512 result = BenchmarkResult(
513 name="critical_backward_hooks",
514 severity=BenchmarkSeverity.WARNING,
515 message="All mismatches due to known architectural differences",
516 details={"total_hooks": len(critical_hooks)},
517 )
519 # Clear model gradients (variables will be GC'd when function returns)
520 if hasattr(bridge, "zero_grad"):
521 bridge.zero_grad()
522 if hasattr(reference_model, "zero_grad"):
523 reference_model.zero_grad()
525 return result
527 result = BenchmarkResult(
528 name="critical_backward_hooks",
529 severity=BenchmarkSeverity.INFO,
530 message=f"All critical backward hooks match",
531 details={"hook_count": len(critical_hooks)},
532 )
534 # Clear model gradients (variables will be GC'd when function returns)
535 if hasattr(bridge, "zero_grad"):
536 bridge.zero_grad()
537 if hasattr(reference_model, "zero_grad"):
538 reference_model.zero_grad()
540 return result
542 except Exception as e:
543 import traceback
545 return BenchmarkResult(
546 name="critical_backward_hooks",
547 severity=BenchmarkSeverity.ERROR,
548 message=f"Critical backward hooks check failed: {str(e)}",
549 details={
550 "error_type": type(e).__name__,
551 "error_message": str(e),
552 "traceback": traceback.format_exc(),
553 },
554 passed=False,
555 )
558def benchmark_gradient_computation(
559 bridge: TransformerBridge,
560 test_text: str,
561 reference_model: Optional[HookedTransformer] = None,
562 atol: float = 1e-3,
563) -> BenchmarkResult:
564 """Benchmark basic gradient computation.
566 Args:
567 bridge: TransformerBridge model to test
568 test_text: Input text for testing
569 reference_model: Optional HookedTransformer reference model
570 atol: Absolute tolerance for gradient comparison
572 Returns:
573 BenchmarkResult with gradient computation comparison details
574 """
575 try:
576 # Run bridge forward and backward
577 bridge_output = bridge(test_text)
578 bridge_loss = bridge_output[:, -1, :].sum()
579 bridge_loss.backward()
581 # Check that gradients were computed
582 has_gradients = False
583 for param in bridge.parameters(): 583 ↛ 588line 583 didn't jump to line 588 because the loop on line 583 didn't complete
584 if param.grad is not None: 584 ↛ 583line 584 didn't jump to line 583 because the condition on line 584 was always true
585 has_gradients = True
586 break
588 if not has_gradients: 588 ↛ 589line 588 didn't jump to line 589 because the condition on line 588 was never true
589 result = BenchmarkResult(
590 name="gradient_computation",
591 severity=BenchmarkSeverity.DANGER,
592 message="No gradients were computed",
593 passed=False,
594 )
595 # Clear gradients anyway
596 if hasattr(bridge, "zero_grad"):
597 bridge.zero_grad()
598 return result
600 if reference_model is None: 600 ↛ 613line 600 didn't jump to line 613 because the condition on line 600 was always true
601 # No reference - just verify gradients exist
602 result = BenchmarkResult(
603 name="gradient_computation",
604 severity=BenchmarkSeverity.INFO,
605 message="Gradients computed successfully",
606 )
607 # Clear gradients
608 if hasattr(bridge, "zero_grad"): 608 ↛ 610line 608 didn't jump to line 610 because the condition on line 608 was always true
609 bridge.zero_grad()
610 return result
612 # Compare with reference model
613 reference_output = reference_model(test_text)
614 reference_loss = reference_output[:, -1, :].sum()
615 reference_loss.backward()
617 # Compare loss values
618 bridge_loss_val = bridge_loss.item()
619 reference_loss_val = reference_loss.item()
621 diff = abs(bridge_loss_val - reference_loss_val)
622 if diff < atol:
623 result = BenchmarkResult(
624 name="gradient_computation",
625 severity=BenchmarkSeverity.INFO,
626 message=f"Loss values match: {bridge_loss_val:.6f} ≈ {reference_loss_val:.6f}",
627 details={"diff": diff, "atol": atol},
628 )
629 else:
630 result = BenchmarkResult(
631 name="gradient_computation",
632 severity=BenchmarkSeverity.WARNING,
633 message=f"Loss values differ: {bridge_loss_val:.6f} vs {reference_loss_val:.6f}",
634 details={"diff": diff, "atol": atol},
635 )
637 # Clean up gradients
638 if hasattr(bridge, "zero_grad"):
639 bridge.zero_grad()
640 if reference_model is not None and hasattr(reference_model, "zero_grad"):
641 reference_model.zero_grad()
643 return result
645 except Exception as e:
646 return BenchmarkResult(
647 name="gradient_computation",
648 severity=BenchmarkSeverity.ERROR,
649 message=f"Gradient computation failed: {str(e)}",
650 passed=False,
651 )