Coverage for transformer_lens/benchmarks/weight_processing.py: 5%
339 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Weight processing benchmarks for TransformerBridge."""
3from typing import Optional, cast
5import torch
7from transformer_lens import HookedTransformer
8from transformer_lens.benchmarks.utils import (
9 BenchmarkResult,
10 BenchmarkSeverity,
11 bridge_self_target_loss,
12 is_tiny_test_model,
13 safe_allclose,
14)
15from transformer_lens.model_bridge import TransformerBridge
16from transformer_lens.model_bridge.generalized_components.attention import (
17 PerLayerGeometryError,
18)
21def benchmark_weight_processing(
22 bridge: TransformerBridge,
23 test_text: str,
24 reference_model: Optional[HookedTransformer] = None,
25) -> BenchmarkResult:
26 """Benchmark weight processing (folding, centering) application.
28 Args:
29 bridge: TransformerBridge model to test
30 test_text: Input text for testing
31 reference_model: Optional HookedTransformer reference model
33 Returns:
34 BenchmarkResult with weight processing verification details
35 """
36 try:
37 from transformer_lens.components.layer_norm_pre import LayerNormPre
38 from transformer_lens.model_bridge.generalized_components.normalization import (
39 NormalizationBridge,
40 )
42 # Check layer norm folding
43 if not isinstance(bridge.ln_final, NormalizationBridge):
44 return BenchmarkResult(
45 name="weight_processing",
46 severity=BenchmarkSeverity.WARNING,
47 message=f"Bridge ln_final is {type(bridge.ln_final).__name__}, expected NormalizationBridge",
48 )
50 # Verify NormalizationBridge has LayerNormPre functionality
51 if not hasattr(bridge.ln_final, "_layernorm_pre_forward"):
52 return BenchmarkResult(
53 name="weight_processing",
54 severity=BenchmarkSeverity.WARNING,
55 message="Bridge ln_final missing LayerNormPre functionality",
56 )
58 if not hasattr(bridge.ln_final.config, "layer_norm_folding"):
59 return BenchmarkResult(
60 name="weight_processing",
61 severity=BenchmarkSeverity.WARNING,
62 message="Bridge ln_final missing layer_norm_folding config",
63 )
65 if reference_model is not None:
66 # Check that reference model has LayerNormPre
67 if not isinstance(reference_model.ln_final, LayerNormPre):
68 return BenchmarkResult(
69 name="weight_processing",
70 severity=BenchmarkSeverity.WARNING,
71 message=f"Reference ln_final is {type(reference_model.ln_final).__name__}, expected LayerNormPre",
72 )
74 # Check weight centering - writing weights should be approximately centered
75 mlp_blocks = bridge.blocks_with("mlp")
76 if not mlp_blocks:
77 return BenchmarkResult(
78 name="weight_processing",
79 severity=BenchmarkSeverity.WARNING,
80 message="No blocks have MLP submodule — cannot check centering",
81 )
82 _mlp_idx, mlp_block = mlp_blocks[0]
83 bridge_w_out = mlp_block.mlp.W_out
84 reference_w_out = reference_model.blocks[_mlp_idx].mlp.W_out
86 bridge_mean = torch.mean(torch.abs(torch.mean(bridge_w_out, dim=-1, keepdim=True)))
87 reference_mean = torch.mean(
88 torch.abs(torch.mean(reference_w_out, dim=-1, keepdim=True))
89 )
91 if bridge_mean.item() > 1e-3:
92 return BenchmarkResult(
93 name="weight_processing",
94 severity=BenchmarkSeverity.WARNING,
95 message=f"Bridge weights not well-centered: {bridge_mean.item():.6f}",
96 details={"bridge_mean": bridge_mean.item()},
97 )
99 if reference_mean.item() > 1e-3:
100 return BenchmarkResult(
101 name="weight_processing",
102 severity=BenchmarkSeverity.WARNING,
103 message=f"Reference weights not well-centered: {reference_mean.item():.6f}",
104 details={"reference_mean": reference_mean.item()},
105 )
107 return BenchmarkResult(
108 name="weight_processing",
109 severity=BenchmarkSeverity.INFO,
110 message="Weight processing verified (folding and centering applied)",
111 details={
112 "bridge_mean": bridge_mean.item(),
113 "reference_mean": reference_mean.item(),
114 },
115 )
117 return BenchmarkResult(
118 name="weight_processing",
119 severity=BenchmarkSeverity.INFO,
120 message="Weight processing structure verified",
121 )
123 except Exception as e:
124 return BenchmarkResult(
125 name="weight_processing",
126 severity=BenchmarkSeverity.ERROR,
127 message=f"Weight processing check failed: {str(e)}",
128 passed=False,
129 )
132def benchmark_weight_sharing(
133 bridge: TransformerBridge,
134 test_text: str,
135 reference_model: Optional[HookedTransformer] = None,
136 atol: float = 1e-3,
137) -> BenchmarkResult:
138 """Benchmark weight sharing and modification effects.
140 Args:
141 bridge: TransformerBridge model to test
142 test_text: Input text for testing
143 reference_model: Optional HookedTransformer reference model
144 atol: Absolute tolerance for effect comparison
146 Returns:
147 BenchmarkResult with weight sharing verification details
148 """
149 try:
150 # Get baseline loss
151 bridge_original = bridge_self_target_loss(bridge, test_text)
153 if reference_model is not None:
154 reference_original = reference_model(test_text, return_type="loss")
156 bridge_attn_blocks = bridge.blocks_with("attn")
157 if not bridge_attn_blocks:
158 return BenchmarkResult(
159 name="weight_sharing",
160 severity=BenchmarkSeverity.INFO,
161 message="No blocks have attention submodule — skipping weight sharing check",
162 )
163 bridge_attn_idx, bridge_attn_block = bridge_attn_blocks[0]
165 # Verify weights are identical before modification
166 bridge_W_V = torch.clone(cast(torch.Tensor, bridge_attn_block.attn.W_V))
167 reference_W_V = torch.clone(
168 cast(torch.Tensor, reference_model.blocks[bridge_attn_idx].attn.W_V)
169 )
171 # Check if models have GQA (different head counts for K/V vs Q)
172 has_gqa = (
173 hasattr(bridge.cfg, "n_key_value_heads")
174 and bridge.cfg.n_key_value_heads != bridge.cfg.n_heads
175 )
177 # For GQA models, HookedTransformer may not support GQA correctly yet
178 # Skip the weight comparison if shapes don't match
179 if bridge_W_V.shape != reference_W_V.shape: # type: ignore[union-attr]
180 if has_gqa:
181 # This is expected - HookedTransformer doesn't support GQA yet
182 # Skip this benchmark for GQA models
183 return BenchmarkResult(
184 name="weight_sharing",
185 severity=BenchmarkSeverity.INFO,
186 message=f"GQA model detected - skipping HT comparison (Bridge W_V: {bridge_W_V.shape}, HT W_V: {reference_W_V.shape})", # type: ignore[union-attr]
187 details={
188 "bridge_shape": str(bridge_W_V.shape), # type: ignore[union-attr]
189 "reference_shape": str(reference_W_V.shape), # type: ignore[union-attr]
190 },
191 )
192 else:
193 return BenchmarkResult(
194 name="weight_sharing",
195 severity=BenchmarkSeverity.WARNING,
196 message=f"Weight shapes differ: Bridge {bridge_W_V.shape} vs Reference {reference_W_V.shape}", # type: ignore[union-attr]
197 details={
198 "bridge_shape": str(bridge_W_V.shape), # type: ignore[union-attr]
199 "reference_shape": str(reference_W_V.shape), # type: ignore[union-attr]
200 },
201 )
203 if not safe_allclose(bridge_W_V, reference_W_V): # type: ignore[arg-type]
204 return BenchmarkResult(
205 name="weight_sharing",
206 severity=BenchmarkSeverity.WARNING,
207 message="Weights differ before modification",
208 )
210 # Modify weights in both models
211 with torch.no_grad():
212 bridge_attn_block.attn.W_V[0, :, :] = 0 # type: ignore[union-attr,operator]
213 reference_model.blocks[bridge_attn_idx].attn.W_V[0, :, :] = 0
215 # Test modified losses
216 bridge_modified = bridge_self_target_loss(bridge, test_text)
217 reference_modified = reference_model(test_text, return_type="loss")
219 bridge_change = bridge_modified - bridge_original
220 reference_change = reference_modified - reference_original
222 # Restore weights
223 with torch.no_grad():
224 bridge_attn_block.attn.W_V.copy_(bridge_W_V) # type: ignore[union-attr,operator,arg-type]
225 reference_model.blocks[bridge_attn_idx].attn.W_V.copy_(reference_W_V)
227 diff = abs(bridge_change - reference_change)
228 if diff < atol:
229 return BenchmarkResult(
230 name="weight_sharing",
231 severity=BenchmarkSeverity.INFO,
232 message=f"Weight modifications have similar effects: {bridge_change:.6f} ≈ {reference_change:.6f}",
233 details={"diff": diff.item(), "atol": atol},
234 )
235 else:
236 return BenchmarkResult(
237 name="weight_sharing",
238 severity=BenchmarkSeverity.WARNING,
239 message=f"Weight modification effects differ: {bridge_change:.6f} vs {reference_change:.6f}",
240 details={"diff": diff.item(), "atol": atol},
241 )
243 # No reference model - just verify modification has an effect
244 # Find first block with attention (hybrid models may not have attn on block 0)
245 bridge_attn_blocks = bridge.blocks_with("attn")
246 if not bridge_attn_blocks:
247 return BenchmarkResult(
248 name="weight_sharing",
249 severity=BenchmarkSeverity.INFO,
250 message="No blocks have attention submodule — skipping weight sharing check",
251 )
252 _ws_idx, ws_attn_block = bridge_attn_blocks[0]
254 original_W_V = ws_attn_block.attn.W_V.clone()
255 with torch.no_grad():
256 ws_attn_block.attn.W_V[0, :, :] = 0
258 bridge_modified = bridge_self_target_loss(bridge, test_text)
259 change = abs(bridge_modified - bridge_original)
261 # Restore weights
262 with torch.no_grad():
263 ws_attn_block.attn.W_V.copy_(original_W_V)
265 if change < 1e-6:
266 return BenchmarkResult(
267 name="weight_sharing",
268 severity=BenchmarkSeverity.WARNING,
269 message=f"Weight modification had minimal effect: {change:.6f}",
270 details={"change": change.item()},
271 )
273 return BenchmarkResult(
274 name="weight_sharing",
275 severity=BenchmarkSeverity.INFO,
276 message=f"Weight modification affects forward pass: change={change:.6f}",
277 details={"change": change.item()},
278 )
280 except Exception as e:
281 return BenchmarkResult(
282 name="weight_sharing",
283 severity=BenchmarkSeverity.ERROR,
284 message=f"Weight sharing check failed: {str(e)}",
285 passed=False,
286 )
289def benchmark_weight_modification(
290 bridge: TransformerBridge,
291 test_text: str,
292 reference_model: Optional[HookedTransformer] = None,
293) -> BenchmarkResult:
294 """Benchmark that weight modifications propagate correctly.
296 Args:
297 bridge: TransformerBridge model to test
298 test_text: Input text for testing
299 reference_model: Optional HookedTransformer reference model (not used)
301 Returns:
302 BenchmarkResult with weight modification verification details
303 """
304 try:
305 # Get original loss
306 original_loss = bridge_self_target_loss(bridge, test_text)
308 # Find first block with attention (hybrid models may not have attn on block 0)
309 wm_attn_blocks = bridge.blocks_with("attn")
310 if not wm_attn_blocks:
311 return BenchmarkResult(
312 name="weight_modification",
313 severity=BenchmarkSeverity.INFO,
314 message="No blocks have attention submodule — skipping weight modification check",
315 )
316 _wm_idx, wm_attn_block = wm_attn_blocks[0]
318 # Modify W_V weights
319 with torch.no_grad():
320 original_w_v = wm_attn_block.attn.W_V.clone()
321 # Check dimensionality - GQA models may have 2D tensors instead of 3D
322 if original_w_v.ndim == 3:
323 # Standard 3D tensor: [n_heads, d_model, d_head]
324 wm_attn_block.attn.W_V[0, :, :] = 0
325 elif original_w_v.ndim == 2:
326 # 2D tensor (e.g., GQA models): [n_heads * d_head, d_model] or similar
327 wm_attn_block.attn.W_V[0, :] = 0
328 else:
329 return BenchmarkResult(
330 name="weight_modification",
331 severity=BenchmarkSeverity.WARNING,
332 message=f"Unexpected W_V shape: {original_w_v.shape} (ndim={original_w_v.ndim})",
333 passed=False,
334 )
336 # Get modified loss (with error handling to restore weights)
337 try:
338 modified_loss = bridge_self_target_loss(bridge, test_text)
339 except Exception as forward_error:
340 # Restore weights before reporting error
341 with torch.no_grad():
342 wm_attn_block.attn.W_V.copy_(original_w_v)
344 # Some models (e.g., models with complex attention mechanisms) may have
345 # forward pass issues after weight modification. Report as skipped.
346 return BenchmarkResult(
347 name="weight_modification",
348 severity=BenchmarkSeverity.SKIPPED,
349 message=f"Weight modification not testable for this architecture: {str(forward_error)}",
350 details={"error": str(forward_error), "architecture_limitation": True},
351 )
353 # Restore weights
354 with torch.no_grad():
355 wm_attn_block.attn.W_V.copy_(original_w_v)
357 # Loss should change
358 change = abs(modified_loss - original_loss)
359 if change < 1e-6:
360 # W_V modification didn't propagate. This can happen in models with
361 # combined QKV projections (e.g., Bloom) where the split V weight
362 # is separate from the combined QKV weight used in forward.
363 # Try MLP weight modification as fallback.
364 mlp_fallback_error = None
365 mlp_blocks = bridge.blocks_with("mlp")
366 mlp_block = mlp_blocks[0][1] if mlp_blocks else None
367 try:
368 if mlp_block is None:
369 raise AttributeError("No blocks have mlp submodule")
370 with torch.no_grad():
371 original_mlp_w = mlp_block.mlp.out.weight.clone()
372 mlp_block.mlp.out.weight[0, :] = 0
373 mlp_modified_loss = bridge_self_target_loss(bridge, test_text)
374 with torch.no_grad():
375 mlp_block.mlp.out.weight.copy_(original_mlp_w)
376 mlp_change = abs(mlp_modified_loss - original_loss)
377 if mlp_change > 1e-6:
378 return BenchmarkResult(
379 name="weight_modification",
380 severity=BenchmarkSeverity.INFO,
381 message=f"Weight modification propagates via MLP (change: {mlp_change:.6f}). "
382 f"W_V not propagated (combined QKV architecture).",
383 details={"change": mlp_change.item(), "fallback": "mlp"},
384 )
385 except Exception as mlp_err:
386 mlp_fallback_error = str(mlp_err)
388 details = {"change": change.item()}
389 if mlp_fallback_error is not None:
390 details["mlp_fallback_error"] = mlp_fallback_error
391 return BenchmarkResult(
392 name="weight_modification",
393 severity=BenchmarkSeverity.DANGER,
394 message=f"Weight modification did not affect loss (change: {change:.6f})",
395 details=details,
396 passed=False,
397 )
399 return BenchmarkResult(
400 name="weight_modification",
401 severity=BenchmarkSeverity.INFO,
402 message=f"Weight modification propagates correctly (change: {change:.6f})",
403 details={"change": change.item()},
404 )
406 except Exception as e:
407 # Some architectures (e.g., Gemma 3 with complex attention, OpenELM with
408 # combined QKV) don't expose W_V. Report as skipped, not passed.
409 if (
410 "cannot be multiplied" in str(e)
411 or "shape" in str(e).lower()
412 or "has no attribute" in str(e)
413 ):
414 return BenchmarkResult(
415 name="weight_modification",
416 severity=BenchmarkSeverity.SKIPPED,
417 message=f"Weight modification not testable for this architecture: {str(e)}",
418 details={"error": str(e), "architecture_limitation": True},
419 )
420 return BenchmarkResult(
421 name="weight_modification",
422 severity=BenchmarkSeverity.ERROR,
423 message=f"Weight modification check failed: {str(e)}",
424 passed=False,
425 )
428def benchmark_layer_norm_folding(
429 bridge: TransformerBridge,
430 test_text: str,
431 reference_model: Optional[HookedTransformer] = None,
432) -> BenchmarkResult:
433 """Benchmark layer norm folding - norm weights should be identity after folding.
435 Args:
436 bridge: TransformerBridge model to test
437 test_text: Input text for testing
438 reference_model: Optional HookedTransformer reference model (not used)
440 Returns:
441 BenchmarkResult with layer norm folding verification details
442 """
443 try:
444 # Skip for architectures that don't support fold_ln (e.g., post-LN like BERT)
445 adapter = getattr(bridge, "adapter", None)
446 if adapter and not getattr(adapter, "supports_fold_ln", True):
447 return BenchmarkResult(
448 name="layer_norm_folding",
449 severity=BenchmarkSeverity.SKIPPED,
450 message="Skipped (post-LN architecture does not support fold_ln)",
451 passed=True,
452 )
454 # Get state dict from bridge (should return TransformerLens format keys)
455 state_dict = bridge.state_dict()
457 # Check both ln1 (attention LN) and ln2 (MLP LN) in TransformerLens format.
458 # Models with combined QKV projections (e.g., OpenELM's qkv_proj) cannot
459 # fold ln1 into attention weights, but ln2 should always be foldable.
460 tolerance = 0.01
461 # For rmsnorm_uses_offset models (Gemma/Gemma2), HF computes x*(1+weight),
462 # so the identity weight after folding is 0.0 (gives 1+0=1). For standard
463 # models, identity is 1.0.
464 cfg = getattr(getattr(bridge, "adapter", None), "cfg", None)
465 rmsnorm_uses_offset = getattr(cfg, "rmsnorm_uses_offset", False)
466 expected_val = 0.0 if rmsnorm_uses_offset else 1.0
467 folded = []
468 not_folded = []
470 for ln_name in ["ln1", "ln2"]:
471 ln_key = f"blocks.0.{ln_name}.weight"
472 if ln_key not in state_dict:
473 continue
474 ln_weight = state_dict[ln_key]
475 mean_val = torch.mean(ln_weight).item()
476 if abs(mean_val - expected_val) < tolerance:
477 folded.append((ln_name, ln_key, mean_val))
478 else:
479 not_folded.append((ln_name, ln_key, mean_val))
481 if not folded and not not_folded:
482 # No LN weights found — model uses non-parametric LayerNorm
483 # (e.g., OLMo v1 has fixed weight=1, bias=0 with no learnable params).
484 # Nothing to fold, so this is a pass.
485 return BenchmarkResult(
486 name="layer_norm_folding",
487 severity=BenchmarkSeverity.INFO,
488 message="No learnable layer norm weights (non-parametric LayerNorm)",
489 passed=True,
490 )
492 if folded and not not_folded:
493 # All LN weights are folded
494 names = ", ".join(f"{n} (mean={m:.6f})" for n, _, m in folded)
495 return BenchmarkResult(
496 name="layer_norm_folding",
497 severity=BenchmarkSeverity.INFO,
498 message=f"Layer norm folding verified: {names}",
499 details={"folded": [n for n, _, _ in folded]},
500 )
501 elif folded and not_folded:
502 # Partial folding — some LN weights folded, some not.
503 # This is expected for models with combined QKV (ln1 can't fold).
504 folded_names = ", ".join(f"{n} (mean={m:.6f})" for n, _, m in folded)
505 unfolded_names = ", ".join(f"{n} (mean={m:.6f})" for n, _, m in not_folded)
506 return BenchmarkResult(
507 name="layer_norm_folding",
508 severity=BenchmarkSeverity.WARNING,
509 message=(
510 f"Partial LN folding: {folded_names} folded; "
511 f"{unfolded_names} preserved (expected for combined QKV models)"
512 ),
513 details={
514 "folded": [n for n, _, _ in folded],
515 "not_folded": [n for n, _, _ in not_folded],
516 },
517 passed=True,
518 )
519 else:
520 # No LN weights folded
521 names = ", ".join(f"{n} (mean={m:.6f})" for n, _, m in not_folded)
522 return BenchmarkResult(
523 name="layer_norm_folding",
524 severity=BenchmarkSeverity.WARNING,
525 message=f"Layer norm weights not identity after folding: {names}",
526 details={"not_folded": [n for n, _, _ in not_folded]},
527 passed=False,
528 )
530 except Exception as e:
531 return BenchmarkResult(
532 name="layer_norm_folding",
533 severity=BenchmarkSeverity.ERROR,
534 message=f"Layer norm folding check failed: {str(e)}",
535 passed=False,
536 )
539def benchmark_attention_output_centering(
540 bridge: TransformerBridge,
541 test_text: str,
542 reference_model: Optional[HookedTransformer] = None,
543) -> BenchmarkResult:
544 """Benchmark attention output centering - W_O should have mean ≈ 0.
546 Args:
547 bridge: TransformerBridge model to test
548 test_text: Input text for testing
549 reference_model: Optional HookedTransformer reference model (not used)
551 Returns:
552 BenchmarkResult with attention output centering verification details
553 """
554 try:
555 # Skip centering check for tiny/test models — random weights don't
556 # center meaningfully and produce false failures.
557 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""):
558 return BenchmarkResult(
559 name="attention_output_centering",
560 severity=BenchmarkSeverity.INFO,
561 message="Skipped for tiny/test model (random weights don't center meaningfully)",
562 )
564 attn_blocks = bridge.blocks_with("attn")
565 if not attn_blocks:
566 # Attention-less (pure SSM) or hybrids whose attention lives inside a
567 # passthrough mixer (NemotronH): nothing to center — skip, don't fail.
568 return BenchmarkResult(
569 name="attention_output_centering",
570 severity=BenchmarkSeverity.SKIPPED,
571 message="No blocks expose an attention submodule (SSM / passthrough-mixer hybrid)",
572 )
574 # Check W_O accessibility on first attention block. Explicit probe,
575 # not hasattr: hasattr invokes the property and only swallows
576 # AttributeError, so the per-layer-geometry ValueError would escape to
577 # the generic handler before the loop's raw-weight fallback runs.
578 first_idx, first_attn_block = attn_blocks[0]
579 w_o_missing = False
580 try:
581 _ = first_attn_block.attn.W_O
582 except AttributeError:
583 w_o_missing = True
584 except PerLayerGeometryError:
585 pass # per-layer geometry; the loop below reads the raw projection
586 if w_o_missing:
587 # No mapped output projection (JetMoe's MoA keeps per-expert W_O
588 # inside the delegated module): structurally nothing to center —
589 # skip like the SSM case above rather than fail.
590 if getattr(first_attn_block.attn, "o", None) is None:
591 return BenchmarkResult(
592 name="attention_output_centering",
593 severity=BenchmarkSeverity.SKIPPED,
594 message="No mapped output projection (delegated per-expert W_O)",
595 )
596 return BenchmarkResult(
597 name="attention_output_centering",
598 severity=BenchmarkSeverity.WARNING,
599 message="W_O not accessible on bridge model",
600 passed=False,
601 )
603 # Compute mean across all attention blocks
604 tolerance = 0.01 # 1% tolerance
605 worst_mean = 0.0
606 for idx, block in attn_blocks:
607 try:
608 column_means = torch.mean(block.attn.W_O, dim=-1)
609 except PerLayerGeometryError:
610 # Per-layer attention geometry (OpenELM varies head counts per
611 # layer): the factorized accessor refuses, but centering is
612 # head-agnostic — the d_model mean reads straight off the 2D
613 # projection.
614 raw = block.attn.o.weight
615 in_out = block.attn._weight_layout_in_out(block.attn.o)
616 column_means = raw.mean(dim=-1) if in_out else raw.mean(dim=0)
617 mean_abs = torch.mean(torch.abs(column_means)).item()
618 worst_mean = max(worst_mean, mean_abs)
620 n_attn = len(attn_blocks)
621 n_total = len(bridge.blocks)
622 block_info = f" ({n_attn}/{n_total} blocks have attention)" if n_attn < n_total else ""
624 if worst_mean < tolerance:
625 return BenchmarkResult(
626 name="attention_output_centering",
627 severity=BenchmarkSeverity.INFO,
628 message=f"Attention output centering verified (worst_mean={worst_mean:.6f}){block_info}",
629 details={"mean": worst_mean, "tolerance": tolerance, "n_attn_blocks": n_attn},
630 )
631 else:
632 return BenchmarkResult(
633 name="attention_output_centering",
634 severity=BenchmarkSeverity.WARNING,
635 message=f"Attention output weights not well-centered (worst_mean={worst_mean:.6f}){block_info}",
636 details={"mean": worst_mean, "tolerance": tolerance, "n_attn_blocks": n_attn},
637 passed=False,
638 )
640 except Exception as e:
641 return BenchmarkResult(
642 name="attention_output_centering",
643 severity=BenchmarkSeverity.ERROR,
644 message=f"Attention output centering check failed: {str(e)}",
645 passed=False,
646 )
649def benchmark_mlp_output_centering(
650 bridge: TransformerBridge,
651 test_text: str,
652 reference_model: Optional[HookedTransformer] = None,
653) -> BenchmarkResult:
654 """Benchmark MLP output centering - MLP output weights should have mean ≈ 0.
656 Args:
657 bridge: TransformerBridge model to test
658 test_text: Input text for testing
659 reference_model: Optional HookedTransformer reference model (not used)
661 Returns:
662 BenchmarkResult with MLP output centering verification details
663 """
664 try:
665 # Skip centering check for tiny/test models — random weights don't
666 # center meaningfully and produce false failures.
667 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""):
668 return BenchmarkResult(
669 name="mlp_output_centering",
670 severity=BenchmarkSeverity.INFO,
671 message="Skipped for tiny/test model (random weights don't center meaningfully)",
672 )
674 # Find an MLP-like submodule (may be "mlp", "shared_mlp", etc.)
675 from transformer_lens.model_bridge.generalized_components.moe import MoEBridge
677 mlp_module = None
678 for block in bridge.blocks:
679 for name in ("mlp", "shared_mlp"):
680 if name in block._modules:
681 mlp_module = block._modules[name]
682 break
683 if mlp_module is not None:
684 break
685 if mlp_module is None:
686 # Pure SSM, or a hybrid whose MLP lives inside a passthrough mixer:
687 # no standalone MLP to center — skip, don't fail.
688 return BenchmarkResult(
689 name="mlp_output_centering",
690 severity=BenchmarkSeverity.SKIPPED,
691 message="No block exposes an MLP submodule (SSM / passthrough-mixer hybrid)",
692 )
694 if isinstance(mlp_module, MoEBridge):
695 return BenchmarkResult(
696 name="mlp_output_centering",
697 severity=BenchmarkSeverity.INFO,
698 message="Skipped for MoE models (no single W_out weight)",
699 details={"is_moe": True},
700 )
702 # Check if W_out exists and is accessible (HT format or bridge format)
703 w_out = None
704 if hasattr(mlp_module, "W_out"):
705 w_out = mlp_module.W_out
706 elif hasattr(mlp_module, "out"):
707 out_module = mlp_module.out
708 if hasattr(out_module, "original_component") and hasattr(
709 out_module.original_component, "weight"
710 ):
711 w_out = out_module.original_component.weight
712 elif hasattr(out_module, "weight"):
713 w_out = out_module.weight
714 if w_out is None:
715 return BenchmarkResult(
716 name="mlp_output_centering",
717 severity=BenchmarkSeverity.WARNING,
718 message="W_out not accessible on bridge model",
719 passed=False,
720 )
722 # Compute mean along output dimension
723 mean_abs = torch.mean(torch.abs(torch.mean(w_out, dim=-1))).item()
725 tolerance = 0.01 # 1% tolerance
727 if mean_abs < tolerance:
728 return BenchmarkResult(
729 name="mlp_output_centering",
730 severity=BenchmarkSeverity.INFO,
731 message=f"MLP output centering verified (mean={mean_abs:.6f})",
732 details={"mean": mean_abs, "tolerance": tolerance},
733 )
734 else:
735 return BenchmarkResult(
736 name="mlp_output_centering",
737 severity=BenchmarkSeverity.WARNING,
738 message=f"MLP output weights not well-centered (mean={mean_abs:.6f})",
739 details={"mean": mean_abs, "tolerance": tolerance},
740 passed=False,
741 )
743 except Exception as e:
744 return BenchmarkResult(
745 name="mlp_output_centering",
746 severity=BenchmarkSeverity.ERROR,
747 message=f"MLP output centering check failed: {str(e)}",
748 passed=False,
749 )
752def benchmark_unembed_centering(
753 bridge: TransformerBridge,
754 test_text: str,
755 reference_model: Optional[HookedTransformer] = None,
756) -> BenchmarkResult:
757 """Benchmark unembed centering - unembed matrix should have mean ≈ 0.
759 Args:
760 bridge: TransformerBridge model to test
761 test_text: Input text for testing
762 reference_model: Optional HookedTransformer reference model (not used)
764 Returns:
765 BenchmarkResult with unembed centering verification details
766 """
767 try:
768 # Get state dict from bridge (should return TransformerLens format keys)
769 state_dict = bridge.state_dict()
771 # Check for unembed weight in TransformerLens format
772 unembed_key = "unembed.weight"
774 # Fallback: if TL format key doesn't exist, try common HF format patterns
775 if unembed_key not in state_dict:
776 # Try standard HF format
777 if "lm_head.weight" in state_dict:
778 unembed_key = "lm_head.weight"
779 else:
780 return BenchmarkResult(
781 name="unembed_centering",
782 severity=BenchmarkSeverity.WARNING,
783 message="Could not find unembed weights in state dict",
784 passed=False,
785 )
787 # Get the unembed weight tensor
788 w_u = state_dict[unembed_key]
790 # Compute mean along vocabulary dimension (dim 0)
791 mean_abs = torch.mean(torch.abs(torch.mean(w_u, dim=0))).item()
793 tolerance = 0.01 # 1% tolerance (consistent with attn/mlp centering)
795 if mean_abs < tolerance:
796 return BenchmarkResult(
797 name="unembed_centering",
798 severity=BenchmarkSeverity.INFO,
799 message=f"Unembed centering verified (mean={mean_abs:.6f})",
800 details={"mean": mean_abs, "tolerance": tolerance, "key": unembed_key},
801 )
802 else:
803 return BenchmarkResult(
804 name="unembed_centering",
805 severity=BenchmarkSeverity.WARNING,
806 message=f"Unembed matrix not well-centered (mean={mean_abs:.6f})",
807 details={"mean": mean_abs, "tolerance": tolerance, "key": unembed_key},
808 passed=False,
809 )
811 except Exception as e:
812 return BenchmarkResult(
813 name="unembed_centering",
814 severity=BenchmarkSeverity.ERROR,
815 message=f"Unembed centering check failed: {str(e)}",
816 passed=False,
817 )
820def benchmark_value_bias_folding(
821 bridge: TransformerBridge,
822 test_text: str,
823 reference_model: Optional[HookedTransformer] = None,
824) -> BenchmarkResult:
825 """Benchmark value bias folding - b_V should be zero after folding.
827 Args:
828 bridge: TransformerBridge model to test
829 test_text: Input text for testing
830 reference_model: Optional HookedTransformer reference model (not used)
832 Returns:
833 BenchmarkResult with value bias folding verification details
834 """
835 try:
836 # Skip for GQA models (where n_key_value_heads != n_heads)
837 # Value bias folding doesn't work the same way because V outputs are repeated
838 if hasattr(bridge.cfg, "n_key_value_heads") and bridge.cfg.n_key_value_heads is not None:
839 if bridge.cfg.n_key_value_heads != bridge.cfg.n_heads:
840 return BenchmarkResult(
841 name="value_bias_folding",
842 severity=BenchmarkSeverity.INFO,
843 message="Skipped for GQA models (n_key_value_heads != n_heads)",
844 details={
845 "is_gqa": True,
846 "n_heads": bridge.cfg.n_heads,
847 "n_kv_heads": bridge.cfg.n_key_value_heads,
848 },
849 )
851 attn_blocks = bridge.blocks_with("attn")
852 if not attn_blocks:
853 return BenchmarkResult(
854 name="value_bias_folding",
855 severity=BenchmarkSeverity.INFO,
856 message="No blocks have attention submodule (expected for hybrid models without mapped attn)",
857 details={"has_bias": False},
858 )
860 first_idx, first_attn_block = attn_blocks[0]
862 # Check if b_V exists
863 if not hasattr(first_attn_block.attn, "b_V"):
864 return BenchmarkResult(
865 name="value_bias_folding",
866 severity=BenchmarkSeverity.INFO,
867 message="No value bias found (expected for models without biases)",
868 details={"has_bias": False},
869 )
871 b_v = first_attn_block.attn.b_V
873 if b_v is None:
874 return BenchmarkResult(
875 name="value_bias_folding",
876 severity=BenchmarkSeverity.INFO,
877 message="Value bias is None (expected for models without biases)",
878 details={"has_bias": False},
879 )
881 # Check if b_V is approximately zero
882 max_abs = torch.max(torch.abs(b_v)).item()
883 tolerance = 1e-6
885 if max_abs < tolerance:
886 return BenchmarkResult(
887 name="value_bias_folding",
888 severity=BenchmarkSeverity.INFO,
889 message=f"Value bias folding verified (max_abs={max_abs:.6e})",
890 details={"max_abs": max_abs, "tolerance": tolerance},
891 )
892 else:
893 return BenchmarkResult(
894 name="value_bias_folding",
895 severity=BenchmarkSeverity.WARNING,
896 message=f"Value bias not zero after folding (max_abs={max_abs:.6e})",
897 details={"max_abs": max_abs, "tolerance": tolerance},
898 passed=False,
899 )
901 except Exception as e:
902 return BenchmarkResult(
903 name="value_bias_folding",
904 severity=BenchmarkSeverity.ERROR,
905 message=f"Value bias folding check failed: {str(e)}",
906 passed=False,
907 )
910def benchmark_no_nan_inf(
911 bridge: TransformerBridge,
912 test_text: str,
913 reference_model: Optional[HookedTransformer] = None,
914) -> BenchmarkResult:
915 """Benchmark that weights contain no NaN or Inf values.
917 Args:
918 bridge: TransformerBridge model to test
919 test_text: Input text for testing
920 reference_model: Optional HookedTransformer reference model (not used)
922 Returns:
923 BenchmarkResult with NaN/Inf verification details
924 """
925 try:
926 # Get state dict from original model
927 state_dict = bridge.state_dict()
929 # Check for NaN/Inf in all tensors
930 nan_keys = []
931 inf_keys = []
933 for key, value in state_dict.items():
934 if torch.isnan(value).any():
935 nan_keys.append(key)
936 if torch.isinf(value).any():
937 inf_keys.append(key)
939 if nan_keys or inf_keys:
940 message_parts = []
941 if nan_keys:
942 message_parts.append(f"NaN in {len(nan_keys)} tensors")
943 if inf_keys:
944 message_parts.append(f"Inf in {len(inf_keys)} tensors")
946 return BenchmarkResult(
947 name="no_nan_inf",
948 severity=BenchmarkSeverity.DANGER,
949 message=f"Invalid values found: {', '.join(message_parts)}",
950 details={"nan_keys": nan_keys, "inf_keys": inf_keys},
951 passed=False,
952 )
954 return BenchmarkResult(
955 name="no_nan_inf",
956 severity=BenchmarkSeverity.INFO,
957 message="No NaN or Inf values found in weights",
958 details={"num_tensors_checked": len(state_dict)},
959 )
961 except Exception as e:
962 return BenchmarkResult(
963 name="no_nan_inf",
964 severity=BenchmarkSeverity.ERROR,
965 message=f"NaN/Inf check failed: {str(e)}",
966 passed=False,
967 )
970def benchmark_weight_magnitudes(
971 bridge: TransformerBridge,
972 test_text: str,
973 reference_model: Optional[HookedTransformer] = None,
974) -> BenchmarkResult:
975 """Benchmark that weight magnitudes are in reasonable ranges.
977 Args:
978 bridge: TransformerBridge model to test
979 test_text: Input text for testing
980 reference_model: Optional HookedTransformer reference model (not used)
982 Returns:
983 BenchmarkResult with weight magnitude verification details
984 """
985 try:
986 # Get state dict from original model
987 state_dict = bridge.state_dict()
989 # Check magnitude ranges
990 too_small_keys = []
991 too_large_keys = []
993 min_threshold = 1e-6
994 max_threshold = 1000.0
996 # For rmsnorm_uses_offset models (Gemma/Gemma2), fold_ln sets LN weights
997 # to 0.0 (identity for (1+w) normalization). Skip LN weights for these models.
998 cfg = getattr(getattr(bridge, "adapter", None), "cfg", None)
999 rmsnorm_uses_offset = getattr(cfg, "rmsnorm_uses_offset", False)
1001 for key, value in state_dict.items():
1002 # Skip non-weight tensors (buffers, etc.)
1003 if "weight" not in key and "bias" not in key:
1004 continue
1006 # Skip internal _original_component keys - these are implementation details
1007 if "_original_component" in key:
1008 continue
1010 # Skip value biases - they are expected to be zero after folding
1011 if ".v.bias" in key:
1012 continue
1014 # Skip attention projection biases - they can be zero in some models
1015 if (
1016 ".k_proj.bias" in key
1017 or ".q_proj.bias" in key
1018 or ".v_proj.bias" in key
1019 or ".o_proj.bias" in key
1020 or ".k.bias" in key
1021 or ".q.bias" in key
1022 or ".v.bias" in key
1023 or ".o.bias" in key
1024 ):
1025 continue
1027 # Skip layer norm biases - they are expected to be zero after folding
1028 if (
1029 "ln1.bias" in key
1030 or "ln2.bias" in key
1031 or "ln_1.bias" in key
1032 or "ln_2.bias" in key
1033 or "ln_final.bias" in key
1034 or "input_layernorm.bias" in key
1035 or "post_attention_layernorm.bias" in key
1036 ):
1037 continue
1039 # For rmsnorm_uses_offset models, fold_ln sets LN weights to 0.0
1040 # (identity for (1+w) normalization). Skip all LN weight keys —
1041 # including post-norms (ln1_post, ln2_post) which aren't folded but
1042 # use the same (1+w) convention — to avoid false magnitude warnings.
1043 if rmsnorm_uses_offset and (
1044 "ln1.weight" in key
1045 or "ln2.weight" in key
1046 or "ln1_post.weight" in key
1047 or "ln2_post.weight" in key
1048 or "ln_1.weight" in key
1049 or "ln_2.weight" in key
1050 or "ln_final.weight" in key
1051 or "input_layernorm.weight" in key
1052 or "post_attention_layernorm.weight" in key
1053 ):
1054 continue
1056 # Skip unembed bias - it may be zero after processing
1057 if "unembed.bias" in key:
1058 continue
1060 # Skip zero biases - many models initialize biases to zero which is
1061 # mathematically equivalent to having no bias. This is a valid state.
1062 if "bias" in key and torch.all(value == 0).item():
1063 continue
1065 mean_abs = torch.mean(torch.abs(value)).item()
1066 max_abs = torch.max(torch.abs(value)).item()
1068 if mean_abs > 0.0 and mean_abs < min_threshold:
1069 # For non-zero weights, check if they're suspiciously small
1070 too_small_keys.append((key, mean_abs))
1072 if max_abs > max_threshold:
1073 too_large_keys.append((key, max_abs))
1075 if too_small_keys or too_large_keys:
1076 message_parts = []
1077 if too_small_keys:
1078 message_parts.append(f"{len(too_small_keys)} too small")
1079 if too_large_keys:
1080 message_parts.append(f"{len(too_large_keys)} too large")
1082 return BenchmarkResult(
1083 name="weight_magnitudes",
1084 severity=BenchmarkSeverity.WARNING,
1085 message=f"Weight magnitude issues: {', '.join(message_parts)}",
1086 details={
1087 "too_small": too_small_keys[:5], # Limit to first 5
1088 "too_large": too_large_keys[:5], # Limit to first 5
1089 },
1090 passed=False,
1091 )
1093 return BenchmarkResult(
1094 name="weight_magnitudes",
1095 severity=BenchmarkSeverity.INFO,
1096 message="All weight magnitudes in reasonable ranges",
1097 details={"min_threshold": min_threshold, "max_threshold": max_threshold},
1098 )
1100 except Exception as e:
1101 return BenchmarkResult(
1102 name="weight_magnitudes",
1103 severity=BenchmarkSeverity.ERROR,
1104 message=f"Weight magnitude check failed: {str(e)}",
1105 passed=False,
1106 )