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