Coverage for transformer_lens/benchmarks/weight_processing.py: 3%
260 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Weight processing benchmarks for TransformerBridge."""
4import torch
6from transformer_lens.benchmarks.utils import (
7 BenchmarkResult,
8 BenchmarkSeverity,
9 bridge_self_target_loss,
10 is_tiny_test_model,
11)
12from transformer_lens.model_bridge import TransformerBridge
13from transformer_lens.model_bridge.generalized_components.attention import (
14 PerLayerGeometryError,
15)
18def benchmark_weight_modification(
19 bridge: TransformerBridge,
20 test_text: str,
21) -> BenchmarkResult:
22 """Benchmark that weight modifications propagate correctly.
24 Args:
25 bridge: TransformerBridge model to test
26 test_text: Input text for testing
28 Returns:
29 BenchmarkResult with weight modification verification details
30 """
31 try:
32 # Get original loss
33 original_loss = bridge_self_target_loss(bridge, test_text)
35 # Find first block with attention (hybrid models may not have attn on block 0)
36 wm_attn_blocks = bridge.blocks_with("attn")
37 if not wm_attn_blocks:
38 return BenchmarkResult(
39 name="weight_modification",
40 severity=BenchmarkSeverity.INFO,
41 message="No blocks have attention submodule — skipping weight modification check",
42 )
43 _wm_idx, wm_attn_block = wm_attn_blocks[0]
45 # Modify W_V weights
46 with torch.no_grad():
47 original_w_v = wm_attn_block.attn.W_V.clone()
48 # Check dimensionality - GQA models may have 2D tensors instead of 3D
49 if original_w_v.ndim == 3:
50 # Standard 3D tensor: [n_heads, d_model, d_head]
51 wm_attn_block.attn.W_V[0, :, :] = 0
52 elif original_w_v.ndim == 2:
53 # 2D tensor (e.g., GQA models): [n_heads * d_head, d_model] or similar
54 wm_attn_block.attn.W_V[0, :] = 0
55 else:
56 return BenchmarkResult(
57 name="weight_modification",
58 severity=BenchmarkSeverity.WARNING,
59 message=f"Unexpected W_V shape: {original_w_v.shape} (ndim={original_w_v.ndim})",
60 passed=False,
61 )
63 # Get modified loss (with error handling to restore weights)
64 try:
65 modified_loss = bridge_self_target_loss(bridge, test_text)
66 except Exception as forward_error:
67 # Restore weights before reporting error
68 with torch.no_grad():
69 wm_attn_block.attn.W_V.copy_(original_w_v)
71 # Some models (e.g., models with complex attention mechanisms) may have
72 # forward pass issues after weight modification. Report as skipped.
73 return BenchmarkResult(
74 name="weight_modification",
75 severity=BenchmarkSeverity.SKIPPED,
76 message=f"Weight modification not testable for this architecture: {str(forward_error)}",
77 details={"error": str(forward_error), "architecture_limitation": True},
78 )
80 # Restore weights
81 with torch.no_grad():
82 wm_attn_block.attn.W_V.copy_(original_w_v)
84 # Loss should change
85 change = abs(modified_loss - original_loss)
86 if change < 1e-6:
87 # W_V modification didn't propagate. This can happen in models with
88 # combined QKV projections (e.g., Bloom) where the split V weight
89 # is separate from the combined QKV weight used in forward.
90 # Try MLP weight modification as fallback.
91 mlp_fallback_error = None
92 mlp_blocks = bridge.blocks_with("mlp")
93 mlp_block = mlp_blocks[0][1] if mlp_blocks else None
94 try:
95 if mlp_block is None:
96 raise AttributeError("No blocks have mlp submodule")
97 with torch.no_grad():
98 original_mlp_w = mlp_block.mlp.out.weight.clone()
99 mlp_block.mlp.out.weight[0, :] = 0
100 mlp_modified_loss = bridge_self_target_loss(bridge, test_text)
101 with torch.no_grad():
102 mlp_block.mlp.out.weight.copy_(original_mlp_w)
103 mlp_change = abs(mlp_modified_loss - original_loss)
104 if mlp_change > 1e-6:
105 return BenchmarkResult(
106 name="weight_modification",
107 severity=BenchmarkSeverity.INFO,
108 message=f"Weight modification propagates via MLP (change: {mlp_change:.6f}). "
109 f"W_V not propagated (combined QKV architecture).",
110 details={"change": mlp_change.item(), "fallback": "mlp"},
111 )
112 except Exception as mlp_err:
113 mlp_fallback_error = str(mlp_err)
115 details = {"change": change.item()}
116 if mlp_fallback_error is not None:
117 details["mlp_fallback_error"] = mlp_fallback_error
118 return BenchmarkResult(
119 name="weight_modification",
120 severity=BenchmarkSeverity.DANGER,
121 message=f"Weight modification did not affect loss (change: {change:.6f})",
122 details=details,
123 passed=False,
124 )
126 return BenchmarkResult(
127 name="weight_modification",
128 severity=BenchmarkSeverity.INFO,
129 message=f"Weight modification propagates correctly (change: {change:.6f})",
130 details={"change": change.item()},
131 )
133 except Exception as e:
134 # Some architectures (e.g., Gemma 3 with complex attention, OpenELM with
135 # combined QKV) don't expose W_V. Report as skipped, not passed.
136 if (
137 "cannot be multiplied" in str(e)
138 or "shape" in str(e).lower()
139 or "has no attribute" in str(e)
140 ):
141 return BenchmarkResult(
142 name="weight_modification",
143 severity=BenchmarkSeverity.SKIPPED,
144 message=f"Weight modification not testable for this architecture: {str(e)}",
145 details={"error": str(e), "architecture_limitation": True},
146 )
147 return BenchmarkResult(
148 name="weight_modification",
149 severity=BenchmarkSeverity.ERROR,
150 message=f"Weight modification check failed: {str(e)}",
151 passed=False,
152 )
155def benchmark_layer_norm_folding(
156 bridge: TransformerBridge,
157 test_text: str,
158) -> BenchmarkResult:
159 """Benchmark layer norm folding - norm weights should be identity after folding.
161 Args:
162 bridge: TransformerBridge model to test
163 test_text: Input text for testing
165 Returns:
166 BenchmarkResult with layer norm folding verification details
167 """
168 try:
169 # Skip for architectures that don't support fold_ln (e.g., post-LN like BERT)
170 adapter = getattr(bridge, "adapter", None)
171 if adapter and not getattr(adapter, "supports_fold_ln", True):
172 return BenchmarkResult(
173 name="layer_norm_folding",
174 severity=BenchmarkSeverity.SKIPPED,
175 message="Skipped (post-LN architecture does not support fold_ln)",
176 passed=True,
177 )
179 # Get state dict from bridge (should return TransformerLens format keys)
180 state_dict = bridge.state_dict()
182 # Check both ln1 (attention LN) and ln2 (MLP LN) in TransformerLens format.
183 # Models with combined QKV projections (e.g., OpenELM's qkv_proj) cannot
184 # fold ln1 into attention weights, but ln2 should always be foldable.
185 tolerance = 0.01
186 # For rmsnorm_uses_offset models (Gemma/Gemma2), HF computes x*(1+weight),
187 # so the identity weight after folding is 0.0 (gives 1+0=1). For standard
188 # models, identity is 1.0.
189 cfg = getattr(getattr(bridge, "adapter", None), "cfg", None)
190 rmsnorm_uses_offset = getattr(cfg, "rmsnorm_uses_offset", False)
191 expected_val = 0.0 if rmsnorm_uses_offset else 1.0
192 folded = []
193 not_folded = []
195 for ln_name in ["ln1", "ln2"]:
196 ln_key = f"blocks.0.{ln_name}.weight"
197 if ln_key not in state_dict:
198 continue
199 ln_weight = state_dict[ln_key]
200 mean_val = torch.mean(ln_weight).item()
201 if abs(mean_val - expected_val) < tolerance:
202 folded.append((ln_name, ln_key, mean_val))
203 else:
204 not_folded.append((ln_name, ln_key, mean_val))
206 if not folded and not not_folded:
207 # No LN weights found — model uses non-parametric LayerNorm
208 # (e.g., OLMo v1 has fixed weight=1, bias=0 with no learnable params).
209 # Nothing to fold, so this is a pass.
210 return BenchmarkResult(
211 name="layer_norm_folding",
212 severity=BenchmarkSeverity.INFO,
213 message="No learnable layer norm weights (non-parametric LayerNorm)",
214 passed=True,
215 )
217 if folded and not not_folded:
218 # All LN weights are folded
219 names = ", ".join(f"{n} (mean={m:.6f})" for n, _, m in folded)
220 return BenchmarkResult(
221 name="layer_norm_folding",
222 severity=BenchmarkSeverity.INFO,
223 message=f"Layer norm folding verified: {names}",
224 details={"folded": [n for n, _, _ in folded]},
225 )
226 elif folded and not_folded:
227 # Partial folding — some LN weights folded, some not.
228 # This is expected for models with combined QKV (ln1 can't fold).
229 folded_names = ", ".join(f"{n} (mean={m:.6f})" for n, _, m in folded)
230 unfolded_names = ", ".join(f"{n} (mean={m:.6f})" for n, _, m in not_folded)
231 return BenchmarkResult(
232 name="layer_norm_folding",
233 severity=BenchmarkSeverity.WARNING,
234 message=(
235 f"Partial LN folding: {folded_names} folded; "
236 f"{unfolded_names} preserved (expected for combined QKV models)"
237 ),
238 details={
239 "folded": [n for n, _, _ in folded],
240 "not_folded": [n for n, _, _ in not_folded],
241 },
242 passed=True,
243 )
244 else:
245 # No LN weights folded
246 names = ", ".join(f"{n} (mean={m:.6f})" for n, _, m in not_folded)
247 return BenchmarkResult(
248 name="layer_norm_folding",
249 severity=BenchmarkSeverity.WARNING,
250 message=f"Layer norm weights not identity after folding: {names}",
251 details={"not_folded": [n for n, _, _ in not_folded]},
252 passed=False,
253 )
255 except Exception as e:
256 return BenchmarkResult(
257 name="layer_norm_folding",
258 severity=BenchmarkSeverity.ERROR,
259 message=f"Layer norm folding check failed: {str(e)}",
260 passed=False,
261 )
264def benchmark_attention_output_centering(
265 bridge: TransformerBridge,
266 test_text: str,
267) -> BenchmarkResult:
268 """Benchmark attention output centering - W_O should have mean ≈ 0.
270 Args:
271 bridge: TransformerBridge model to test
272 test_text: Input text for testing
274 Returns:
275 BenchmarkResult with attention output centering verification details
276 """
277 try:
278 # Skip centering check for tiny/test models — random weights don't
279 # center meaningfully and produce false failures.
280 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""):
281 return BenchmarkResult(
282 name="attention_output_centering",
283 severity=BenchmarkSeverity.INFO,
284 message="Skipped for tiny/test model (random weights don't center meaningfully)",
285 )
287 attn_blocks = bridge.blocks_with("attn")
288 if not attn_blocks:
289 # Attention-less (pure SSM) or hybrids whose attention lives inside a
290 # passthrough mixer (NemotronH): nothing to center — skip, don't fail.
291 return BenchmarkResult(
292 name="attention_output_centering",
293 severity=BenchmarkSeverity.SKIPPED,
294 message="No blocks expose an attention submodule (SSM / passthrough-mixer hybrid)",
295 )
297 # Check W_O accessibility on first attention block. Explicit probe,
298 # not hasattr: hasattr invokes the property and only swallows
299 # AttributeError, so the per-layer-geometry ValueError would escape to
300 # the generic handler before the loop's raw-weight fallback runs.
301 first_idx, first_attn_block = attn_blocks[0]
302 w_o_missing = False
303 try:
304 _ = first_attn_block.attn.W_O
305 except AttributeError:
306 w_o_missing = True
307 except PerLayerGeometryError:
308 pass # per-layer geometry; the loop below reads the raw projection
309 if w_o_missing:
310 # No mapped output projection (JetMoe's MoA keeps per-expert W_O
311 # inside the delegated module): structurally nothing to center —
312 # skip like the SSM case above rather than fail.
313 if getattr(first_attn_block.attn, "o", None) is None:
314 return BenchmarkResult(
315 name="attention_output_centering",
316 severity=BenchmarkSeverity.SKIPPED,
317 message="No mapped output projection (delegated per-expert W_O)",
318 )
319 return BenchmarkResult(
320 name="attention_output_centering",
321 severity=BenchmarkSeverity.WARNING,
322 message="W_O not accessible on bridge model",
323 passed=False,
324 )
326 # Compute mean across all attention blocks
327 tolerance = 0.01 # 1% tolerance
328 worst_mean = 0.0
329 for idx, block in attn_blocks:
330 try:
331 column_means = torch.mean(block.attn.W_O, dim=-1)
332 except PerLayerGeometryError:
333 # Per-layer attention geometry (OpenELM varies head counts per
334 # layer): the factorized accessor refuses, but centering is
335 # head-agnostic — the d_model mean reads straight off the 2D
336 # projection.
337 raw = block.attn.o.weight
338 in_out = block.attn._weight_layout_in_out(block.attn.o)
339 column_means = raw.mean(dim=-1) if in_out else raw.mean(dim=0)
340 mean_abs = torch.mean(torch.abs(column_means)).item()
341 worst_mean = max(worst_mean, mean_abs)
343 n_attn = len(attn_blocks)
344 n_total = len(bridge.blocks)
345 block_info = f" ({n_attn}/{n_total} blocks have attention)" if n_attn < n_total else ""
347 if worst_mean < tolerance:
348 return BenchmarkResult(
349 name="attention_output_centering",
350 severity=BenchmarkSeverity.INFO,
351 message=f"Attention output centering verified (worst_mean={worst_mean:.6f}){block_info}",
352 details={"mean": worst_mean, "tolerance": tolerance, "n_attn_blocks": n_attn},
353 )
354 else:
355 return BenchmarkResult(
356 name="attention_output_centering",
357 severity=BenchmarkSeverity.WARNING,
358 message=f"Attention output weights not well-centered (worst_mean={worst_mean:.6f}){block_info}",
359 details={"mean": worst_mean, "tolerance": tolerance, "n_attn_blocks": n_attn},
360 passed=False,
361 )
363 except Exception as e:
364 return BenchmarkResult(
365 name="attention_output_centering",
366 severity=BenchmarkSeverity.ERROR,
367 message=f"Attention output centering check failed: {str(e)}",
368 passed=False,
369 )
372def benchmark_mlp_output_centering(
373 bridge: TransformerBridge,
374 test_text: str,
375) -> BenchmarkResult:
376 """Benchmark MLP output centering - MLP output weights should have mean ≈ 0.
378 Args:
379 bridge: TransformerBridge model to test
380 test_text: Input text for testing
382 Returns:
383 BenchmarkResult with MLP output centering verification details
384 """
385 try:
386 # Skip centering check for tiny/test models — random weights don't
387 # center meaningfully and produce false failures.
388 if is_tiny_test_model(getattr(bridge.cfg, "model_name", "") or ""):
389 return BenchmarkResult(
390 name="mlp_output_centering",
391 severity=BenchmarkSeverity.INFO,
392 message="Skipped for tiny/test model (random weights don't center meaningfully)",
393 )
395 # Find an MLP-like submodule (may be "mlp", "shared_mlp", etc.)
396 from transformer_lens.model_bridge.generalized_components.moe import MoEBridge
398 mlp_module = None
399 for block in bridge.blocks:
400 for name in ("mlp", "shared_mlp"):
401 if name in block._modules:
402 mlp_module = block._modules[name]
403 break
404 if mlp_module is not None:
405 break
406 if mlp_module is None:
407 # Pure SSM, or a hybrid whose MLP lives inside a passthrough mixer:
408 # no standalone MLP to center — skip, don't fail.
409 return BenchmarkResult(
410 name="mlp_output_centering",
411 severity=BenchmarkSeverity.SKIPPED,
412 message="No block exposes an MLP submodule (SSM / passthrough-mixer hybrid)",
413 )
415 if isinstance(mlp_module, MoEBridge):
416 return BenchmarkResult(
417 name="mlp_output_centering",
418 severity=BenchmarkSeverity.INFO,
419 message="Skipped for MoE models (no single W_out weight)",
420 details={"is_moe": True},
421 )
423 # Check if W_out exists and is accessible (HT format or bridge format)
424 w_out = None
425 if hasattr(mlp_module, "W_out"):
426 w_out = mlp_module.W_out
427 elif hasattr(mlp_module, "out"):
428 out_module = mlp_module.out
429 if hasattr(out_module, "original_component") and hasattr(
430 out_module.original_component, "weight"
431 ):
432 w_out = out_module.original_component.weight
433 elif hasattr(out_module, "weight"):
434 w_out = out_module.weight
435 if w_out is None:
436 return BenchmarkResult(
437 name="mlp_output_centering",
438 severity=BenchmarkSeverity.WARNING,
439 message="W_out not accessible on bridge model",
440 passed=False,
441 )
443 # Compute mean along output dimension
444 mean_abs = torch.mean(torch.abs(torch.mean(w_out, dim=-1))).item()
446 tolerance = 0.01 # 1% tolerance
448 if mean_abs < tolerance:
449 return BenchmarkResult(
450 name="mlp_output_centering",
451 severity=BenchmarkSeverity.INFO,
452 message=f"MLP output centering verified (mean={mean_abs:.6f})",
453 details={"mean": mean_abs, "tolerance": tolerance},
454 )
455 else:
456 return BenchmarkResult(
457 name="mlp_output_centering",
458 severity=BenchmarkSeverity.WARNING,
459 message=f"MLP output weights not well-centered (mean={mean_abs:.6f})",
460 details={"mean": mean_abs, "tolerance": tolerance},
461 passed=False,
462 )
464 except Exception as e:
465 return BenchmarkResult(
466 name="mlp_output_centering",
467 severity=BenchmarkSeverity.ERROR,
468 message=f"MLP output centering check failed: {str(e)}",
469 passed=False,
470 )
473def benchmark_unembed_centering(
474 bridge: TransformerBridge,
475 test_text: str,
476) -> BenchmarkResult:
477 """Benchmark unembed centering - unembed matrix should have mean ≈ 0.
479 Args:
480 bridge: TransformerBridge model to test
481 test_text: Input text for testing
483 Returns:
484 BenchmarkResult with unembed centering verification details
485 """
486 try:
487 # Get state dict from bridge (should return TransformerLens format keys)
488 state_dict = bridge.state_dict()
490 # Check for unembed weight in TransformerLens format
491 unembed_key = "unembed.weight"
493 # Fallback: if TL format key doesn't exist, try common HF format patterns
494 if unembed_key not in state_dict:
495 # Try standard HF format
496 if "lm_head.weight" in state_dict:
497 unembed_key = "lm_head.weight"
498 else:
499 return BenchmarkResult(
500 name="unembed_centering",
501 severity=BenchmarkSeverity.WARNING,
502 message="Could not find unembed weights in state dict",
503 passed=False,
504 )
506 # Get the unembed weight tensor
507 w_u = state_dict[unembed_key]
509 # Compute mean along vocabulary dimension (dim 0)
510 mean_abs = torch.mean(torch.abs(torch.mean(w_u, dim=0))).item()
512 tolerance = 0.01 # 1% tolerance (consistent with attn/mlp centering)
514 if mean_abs < tolerance:
515 return BenchmarkResult(
516 name="unembed_centering",
517 severity=BenchmarkSeverity.INFO,
518 message=f"Unembed centering verified (mean={mean_abs:.6f})",
519 details={"mean": mean_abs, "tolerance": tolerance, "key": unembed_key},
520 )
521 else:
522 return BenchmarkResult(
523 name="unembed_centering",
524 severity=BenchmarkSeverity.WARNING,
525 message=f"Unembed matrix not well-centered (mean={mean_abs:.6f})",
526 details={"mean": mean_abs, "tolerance": tolerance, "key": unembed_key},
527 passed=False,
528 )
530 except Exception as e:
531 return BenchmarkResult(
532 name="unembed_centering",
533 severity=BenchmarkSeverity.ERROR,
534 message=f"Unembed centering check failed: {str(e)}",
535 passed=False,
536 )
539def benchmark_value_bias_folding(
540 bridge: TransformerBridge,
541 test_text: str,
542) -> BenchmarkResult:
543 """Benchmark value bias folding - b_V should be zero after folding.
545 Args:
546 bridge: TransformerBridge model to test
547 test_text: Input text for testing
549 Returns:
550 BenchmarkResult with value bias folding verification details
551 """
552 try:
553 # Skip for GQA models (where n_key_value_heads != n_heads)
554 # Value bias folding doesn't work the same way because V outputs are repeated
555 if hasattr(bridge.cfg, "n_key_value_heads") and bridge.cfg.n_key_value_heads is not None:
556 if bridge.cfg.n_key_value_heads != bridge.cfg.n_heads:
557 return BenchmarkResult(
558 name="value_bias_folding",
559 severity=BenchmarkSeverity.INFO,
560 message="Skipped for GQA models (n_key_value_heads != n_heads)",
561 details={
562 "is_gqa": True,
563 "n_heads": bridge.cfg.n_heads,
564 "n_kv_heads": bridge.cfg.n_key_value_heads,
565 },
566 )
568 attn_blocks = bridge.blocks_with("attn")
569 if not attn_blocks:
570 return BenchmarkResult(
571 name="value_bias_folding",
572 severity=BenchmarkSeverity.INFO,
573 message="No blocks have attention submodule (expected for hybrid models without mapped attn)",
574 details={"has_bias": False},
575 )
577 first_idx, first_attn_block = attn_blocks[0]
579 # Check if b_V exists
580 if not hasattr(first_attn_block.attn, "b_V"):
581 return BenchmarkResult(
582 name="value_bias_folding",
583 severity=BenchmarkSeverity.INFO,
584 message="No value bias found (expected for models without biases)",
585 details={"has_bias": False},
586 )
588 b_v = first_attn_block.attn.b_V
590 if b_v is None:
591 return BenchmarkResult(
592 name="value_bias_folding",
593 severity=BenchmarkSeverity.INFO,
594 message="Value bias is None (expected for models without biases)",
595 details={"has_bias": False},
596 )
598 # Check if b_V is approximately zero
599 max_abs = torch.max(torch.abs(b_v)).item()
600 tolerance = 1e-6
602 if max_abs < tolerance:
603 return BenchmarkResult(
604 name="value_bias_folding",
605 severity=BenchmarkSeverity.INFO,
606 message=f"Value bias folding verified (max_abs={max_abs:.6e})",
607 details={"max_abs": max_abs, "tolerance": tolerance},
608 )
609 else:
610 return BenchmarkResult(
611 name="value_bias_folding",
612 severity=BenchmarkSeverity.WARNING,
613 message=f"Value bias not zero after folding (max_abs={max_abs:.6e})",
614 details={"max_abs": max_abs, "tolerance": tolerance},
615 passed=False,
616 )
618 except Exception as e:
619 return BenchmarkResult(
620 name="value_bias_folding",
621 severity=BenchmarkSeverity.ERROR,
622 message=f"Value bias folding check failed: {str(e)}",
623 passed=False,
624 )
627def benchmark_no_nan_inf(
628 bridge: TransformerBridge,
629 test_text: str,
630) -> BenchmarkResult:
631 """Benchmark that weights contain no NaN or Inf values.
633 Args:
634 bridge: TransformerBridge model to test
635 test_text: Input text for testing
637 Returns:
638 BenchmarkResult with NaN/Inf verification details
639 """
640 try:
641 # Get state dict from original model
642 state_dict = bridge.state_dict()
644 # Check for NaN/Inf in all tensors
645 nan_keys = []
646 inf_keys = []
648 for key, value in state_dict.items():
649 if torch.isnan(value).any():
650 nan_keys.append(key)
651 if torch.isinf(value).any():
652 inf_keys.append(key)
654 if nan_keys or inf_keys:
655 message_parts = []
656 if nan_keys:
657 message_parts.append(f"NaN in {len(nan_keys)} tensors")
658 if inf_keys:
659 message_parts.append(f"Inf in {len(inf_keys)} tensors")
661 return BenchmarkResult(
662 name="no_nan_inf",
663 severity=BenchmarkSeverity.DANGER,
664 message=f"Invalid values found: {', '.join(message_parts)}",
665 details={"nan_keys": nan_keys, "inf_keys": inf_keys},
666 passed=False,
667 )
669 return BenchmarkResult(
670 name="no_nan_inf",
671 severity=BenchmarkSeverity.INFO,
672 message="No NaN or Inf values found in weights",
673 details={"num_tensors_checked": len(state_dict)},
674 )
676 except Exception as e:
677 return BenchmarkResult(
678 name="no_nan_inf",
679 severity=BenchmarkSeverity.ERROR,
680 message=f"NaN/Inf check failed: {str(e)}",
681 passed=False,
682 )
685def benchmark_weight_magnitudes(
686 bridge: TransformerBridge,
687 test_text: str,
688) -> BenchmarkResult:
689 """Benchmark that weight magnitudes are in reasonable ranges.
691 Args:
692 bridge: TransformerBridge model to test
693 test_text: Input text for testing
695 Returns:
696 BenchmarkResult with weight magnitude verification details
697 """
698 try:
699 # Get state dict from original model
700 state_dict = bridge.state_dict()
702 # Check magnitude ranges
703 too_small_keys = []
704 too_large_keys = []
706 min_threshold = 1e-6
707 max_threshold = 1000.0
709 # For rmsnorm_uses_offset models (Gemma/Gemma2), fold_ln sets LN weights
710 # to 0.0 (identity for (1+w) normalization). Skip LN weights for these models.
711 cfg = getattr(getattr(bridge, "adapter", None), "cfg", None)
712 rmsnorm_uses_offset = getattr(cfg, "rmsnorm_uses_offset", False)
714 for key, value in state_dict.items():
715 # Skip non-weight tensors (buffers, etc.)
716 if "weight" not in key and "bias" not in key:
717 continue
719 # Skip internal _original_component keys - these are implementation details
720 if "_original_component" in key:
721 continue
723 # Skip value biases - they are expected to be zero after folding
724 if ".v.bias" in key:
725 continue
727 # Skip attention projection biases - they can be zero in some models
728 if (
729 ".k_proj.bias" in key
730 or ".q_proj.bias" in key
731 or ".v_proj.bias" in key
732 or ".o_proj.bias" in key
733 or ".k.bias" in key
734 or ".q.bias" in key
735 or ".v.bias" in key
736 or ".o.bias" in key
737 ):
738 continue
740 # Skip layer norm biases - they are expected to be zero after folding
741 if (
742 "ln1.bias" in key
743 or "ln2.bias" in key
744 or "ln_1.bias" in key
745 or "ln_2.bias" in key
746 or "ln_final.bias" in key
747 or "input_layernorm.bias" in key
748 or "post_attention_layernorm.bias" in key
749 ):
750 continue
752 # For rmsnorm_uses_offset models, fold_ln sets LN weights to 0.0
753 # (identity for (1+w) normalization). Skip all LN weight keys —
754 # including post-norms (ln1_post, ln2_post) which aren't folded but
755 # use the same (1+w) convention — to avoid false magnitude warnings.
756 if rmsnorm_uses_offset and (
757 "ln1.weight" in key
758 or "ln2.weight" in key
759 or "ln1_post.weight" in key
760 or "ln2_post.weight" in key
761 or "ln_1.weight" in key
762 or "ln_2.weight" in key
763 or "ln_final.weight" in key
764 or "input_layernorm.weight" in key
765 or "post_attention_layernorm.weight" in key
766 ):
767 continue
769 # Skip unembed bias - it may be zero after processing
770 if "unembed.bias" in key:
771 continue
773 # Skip zero biases - many models initialize biases to zero which is
774 # mathematically equivalent to having no bias. This is a valid state.
775 if "bias" in key and torch.all(value == 0).item():
776 continue
778 mean_abs = torch.mean(torch.abs(value)).item()
779 max_abs = torch.max(torch.abs(value)).item()
781 if mean_abs > 0.0 and mean_abs < min_threshold:
782 # For non-zero weights, check if they're suspiciously small
783 too_small_keys.append((key, mean_abs))
785 if max_abs > max_threshold:
786 too_large_keys.append((key, max_abs))
788 if too_small_keys or too_large_keys:
789 message_parts = []
790 if too_small_keys:
791 message_parts.append(f"{len(too_small_keys)} too small")
792 if too_large_keys:
793 message_parts.append(f"{len(too_large_keys)} too large")
795 return BenchmarkResult(
796 name="weight_magnitudes",
797 severity=BenchmarkSeverity.WARNING,
798 message=f"Weight magnitude issues: {', '.join(message_parts)}",
799 details={
800 "too_small": too_small_keys[:5], # Limit to first 5
801 "too_large": too_large_keys[:5], # Limit to first 5
802 },
803 passed=False,
804 )
806 return BenchmarkResult(
807 name="weight_magnitudes",
808 severity=BenchmarkSeverity.INFO,
809 message="All weight magnitudes in reasonable ranges",
810 details={"min_threshold": min_threshold, "max_threshold": max_threshold},
811 )
813 except Exception as e:
814 return BenchmarkResult(
815 name="weight_magnitudes",
816 severity=BenchmarkSeverity.ERROR,
817 message=f"Weight magnitude check failed: {str(e)}",
818 passed=False,
819 )