Coverage for transformer_lens/benchmarks/utils.py: 59%
186 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Utility types and functions for benchmarking."""
3from contextlib import contextmanager
4from dataclasses import dataclass
5from enum import Enum
6from typing import Any, Collection, Dict, Iterator, List, Optional, Union
8import torch
10# Prefixes used by tiny/random test models that produce degenerate weights and
11# should be skipped for certain benchmarks (centering, generation, etc.).
12TINY_TEST_MODEL_PATTERNS = (
13 "tiny-random",
14 "trl-internal-testing/tiny",
15 "peft-internal-testing/tiny",
16)
19@contextmanager
20def deterministic_rng(seed: int = 42) -> Iterator[None]:
21 """Seed sampling locally, leaving the caller's RNG stream untouched.
23 Benchmarks that sample need reproducible verdicts, but a bare manual_seed
24 rewrites the process-global stream for every benchmark that follows.
25 """
26 with torch.random.fork_rng(devices=[]):
27 torch.manual_seed(seed)
28 yield
31def is_tiny_test_model(model_name: str) -> bool:
32 """Check if a model name belongs to a tiny/random test model."""
33 return any(pattern in model_name for pattern in TINY_TEST_MODEL_PATTERNS)
36# Hook patterns that bridge models inherently don't have because they use HF's
37# native implementation rather than reimplementing attention/MLP internals.
38BRIDGE_EXPECTED_MISSING_PATTERNS = [
39 "mlp.hook_pre",
40 "mlp.hook_post",
41 "hook_mlp_in",
42 "hook_mlp_out",
43 "attn.hook_rot_q",
44 "attn.hook_rot_k",
45 "hook_pos_embed",
46 "embed.ln.hook_scale",
47 "embed.ln.hook_normalized",
48 "attn.hook_q",
49 "attn.hook_k",
50 "attn.hook_v",
51 # cfg-gated attention hooks. These exist unconditionally on the attention
52 # bridge (so `run_with_cache` key lookups never KeyError) but only fire
53 # when their config flag is on. `benchmark_forward_hooks` runs with
54 # defaults (flags=False) so these correctly don't fire during that
55 # benchmark — suppressing them here prevents false "didn't fire"
56 # failures. The affirmative verification that they DO fire when flags
57 # are on lives in `benchmark_gated_hooks_fire`, which toggles each flag
58 # and asserts the relevant hooks capture activations.
59 "hook_result",
60 "hook_attn_in",
61 "hook_q_input",
62 "hook_k_input",
63 "hook_v_input",
64 "attn.hook_attn_scores",
65 "attn.hook_pattern",
66 # MoE per-expert hooks: Bridge uses HF's batched MoE forward pass via MoEBridge,
67 # which wraps the entire MoE module. The legacy reference created individual expert
68 # modules with per-expert hooks (e.g., blocks.0.mlp.experts.3.hook_pre).
69 "mlp.experts.",
70 "mlp.hook_experts",
71 # Parallel attention+MLP architectures (GPT-J, GPT-NeoX): HF has a single
72 # shared layer norm (ln_1), while HT creates a virtual ln2 that shares weights
73 # with ln1. The Bridge only wraps the actual HF ln_1, so ln2 hooks don't exist.
74 # These patterns only match "missing" hooks when ln2 is absent from the Bridge;
75 # for non-parallel architectures, the Bridge HAS ln2 and these won't be missing.
76 "ln2.hook_scale",
77 "ln2.hook_normalized",
78]
81def filter_expected_missing_hooks(hook_names: Collection[str]) -> list[str]:
82 """Filter out hook names that bridge models are expected to be missing."""
83 return [
84 h
85 for h in hook_names
86 if not any(pattern in h for pattern in BRIDGE_EXPECTED_MISSING_PATTERNS)
87 ]
90_DEFAULT_WAVEFORM_SAMPLES = 16000
93def build_modality_input(
94 bridge: Any,
95 batch_size: int = 1,
96 device: Optional[Union[str, torch.device]] = None,
97 dtype: Optional[torch.dtype] = None,
98) -> Optional[torch.Tensor]:
99 """Build a synthetic model input for a non-text bridge, or None for text models.
101 Audio and vision models take a raw tensor where text models take token ids, and
102 the shape is architecture-specific: spectrogram encoders (AST) need
103 ``[batch, max_length, num_mel_bins]`` where waveform encoders (HuBERT, wav2vec2)
104 need ``[batch, samples]``, and vision encoders need
105 ``[batch, num_channels, image_size, image_size]``. Shapes come from the HF config
106 so a checkpoint with non-default dimensions is still handled correctly.
107 """
108 cfg = getattr(bridge, "cfg", None)
109 if cfg is None: 109 ↛ 110line 109 didn't jump to line 110 because the condition on line 109 was never true
110 return None
112 hf_config = getattr(getattr(bridge, "original_model", None), "config", None)
114 def _from_config(name: str, default: int) -> int:
115 value = getattr(hf_config, name, None) if hf_config is not None else None
116 if value is None: 116 ↛ 117line 116 didn't jump to line 117 because the condition on line 116 was never true
117 value = getattr(cfg, name, None)
118 return int(value) if isinstance(value, int) else default
120 if getattr(cfg, "is_visual_model", False):
121 image_size = _from_config("image_size", 224)
122 num_channels = _from_config("num_channels", 3)
123 return torch.randn(
124 batch_size, num_channels, image_size, image_size, device=device, dtype=dtype
125 )
127 if getattr(cfg, "is_audio_model", False):
128 # num_mel_bins is what distinguishes a spectrogram encoder from a waveform one.
129 num_mel_bins = getattr(hf_config, "num_mel_bins", None) if hf_config is not None else None
130 if num_mel_bins is None:
131 return torch.randn(batch_size, _DEFAULT_WAVEFORM_SAMPLES, device=device, dtype=dtype)
132 max_length = _from_config("max_length", 1024)
133 return torch.randn(batch_size, max_length, int(num_mel_bins), device=device, dtype=dtype)
135 return None
138def safe_allclose(
139 tensor1: torch.Tensor,
140 tensor2: torch.Tensor,
141 atol: float = 1e-5,
142 rtol: float = 1e-5,
143) -> bool:
144 """torch.allclose that handles dtype and device mismatches."""
145 if tensor1.device != tensor2.device: 145 ↛ 146line 145 didn't jump to line 146 because the condition on line 145 was never true
146 tensor1 = tensor1.cpu()
147 tensor2 = tensor2.cpu()
148 if tensor1.dtype != tensor2.dtype: 148 ↛ 149line 148 didn't jump to line 149 because the condition on line 148 was never true
149 tensor1 = tensor1.to(torch.float32)
150 tensor2 = tensor2.to(torch.float32)
151 return torch.allclose(tensor1, tensor2, atol=atol, rtol=rtol)
154class BenchmarkSeverity(Enum):
155 """Severity levels for benchmark results."""
157 INFO = "info" # ✅ PASS - Model working perfectly, all checks passed
158 WARNING = "warning" # ⚠️ PASS with notes - Acceptable differences worth noting
159 DANGER = "danger" # ❌ FAIL - Significant mismatches or failures
160 ERROR = "error" # ❌ ERROR - Test crashed or couldn't run
161 SKIPPED = "skipped" # ⏭️ SKIPPED - Test skipped (e.g., no reference model available)
164@dataclass
165class BenchmarkResult:
166 """Result of a benchmark test."""
168 name: str
169 severity: BenchmarkSeverity
170 message: str
171 details: Optional[Dict[str, Any]] = None
172 passed: bool = True
173 phase: Optional[int] = None # Phase number (1, 2, 3, etc.)
175 def __str__(self) -> str:
176 """Format result for console output."""
177 severity_icons = {
178 BenchmarkSeverity.INFO: "🟢",
179 BenchmarkSeverity.WARNING: "🟡",
180 BenchmarkSeverity.DANGER: "🔴",
181 BenchmarkSeverity.ERROR: "❌",
182 BenchmarkSeverity.SKIPPED: "⏭️",
183 }
184 icon = severity_icons[self.severity]
186 if self.severity == BenchmarkSeverity.SKIPPED:
187 status = "SKIPPED"
188 else:
189 status = "PASS" if self.passed else "FAIL"
191 result = f"{icon} [{status}] {self.name}: {self.message}"
193 if self.details:
194 detail_lines = []
195 for key, value in self.details.items():
196 detail_lines.append(f" {key}: {value}")
197 result += "\n" + "\n".join(detail_lines)
199 return result
201 def print_immediate(self) -> None:
202 """Print this result immediately to console."""
203 print(str(self))
206@dataclass
207class PhaseReferenceData:
208 """Float32 reference data from Phase 1 for Phase 3 equivalence comparison."""
210 hf_logits: Optional[torch.Tensor] = None
211 hf_loss: Optional[float] = None
212 test_text: Optional[str] = None
215def make_capture_hook(storage: dict, name: str):
216 """Create a forward hook that captures activations into a dict.
218 Handles both raw tensors and tuples (extracts first element).
219 """
221 def hook_fn(tensor, hook):
222 if isinstance(tensor, torch.Tensor):
223 storage[name] = tensor.detach().clone()
224 elif isinstance(tensor, tuple) and len(tensor) > 0:
225 if isinstance(tensor[0], torch.Tensor):
226 storage[name] = tensor[0].detach().clone()
227 return tensor
229 return hook_fn
232def make_grad_capture_hook(storage: dict, name: str, return_none: bool = False):
233 """Create a backward hook that captures gradients into a dict.
235 Args:
236 storage: Dict to store captured gradients
237 name: Key name for storage
238 return_none: If True, return None (for backward hooks that shouldn't modify grads)
239 """
241 def hook_fn(tensor, hook=None):
242 if isinstance(tensor, torch.Tensor):
243 storage[name] = tensor.detach().clone()
244 elif isinstance(tensor, tuple) and len(tensor) > 0: 244 ↛ 247line 244 didn't jump to line 247 because the condition on line 244 was always true
245 if tensor[0] is not None and isinstance(tensor[0], torch.Tensor): 245 ↛ 247line 245 didn't jump to line 247 because the condition on line 245 was always true
246 storage[name] = tensor[0].detach().clone()
247 return None if return_none else tensor
249 return hook_fn
252def _squeeze_batch_dim(t1: torch.Tensor, t2: torch.Tensor):
253 """Handle batch dimension differences (e.g., [seq, dim] vs [1, seq, dim]).
255 Returns (t1, t2) with matching shapes, or None if shapes are incompatible.
256 """
257 if t1.shape == t2.shape:
258 return t1, t2
259 if t1.ndim == t2.ndim - 1 and t2.shape[0] == 1 and t1.shape == t2.shape[1:]:
260 return t1.unsqueeze(0), t2
261 if t2.ndim == t1.ndim - 1 and t1.shape[0] == 1 and t2.shape == t1.shape[1:]: 261 ↛ 262line 261 didn't jump to line 262 because the condition on line 261 was never true
262 return t1, t2.unsqueeze(0)
263 return None
266def compare_activation_dicts(
267 dict1: Dict[str, torch.Tensor],
268 dict2: Dict[str, torch.Tensor],
269 atol: float = 1e-5,
270 rtol: float = 0.0,
271) -> List[str]:
272 """Compare two activation/gradient dicts, returning mismatch descriptions.
274 Handles batch-dim squeezing and dtype/device normalization.
275 """
276 mismatches = []
277 common_keys = sorted(set(dict1.keys()) & set(dict2.keys()))
278 for key in common_keys:
279 t1, t2 = dict1[key], dict2[key]
280 squeezed = _squeeze_batch_dim(t1, t2)
281 if squeezed is None:
282 mismatches.append(f"{key}: Shape mismatch - {t1.shape} vs {t2.shape}")
283 continue
284 t1, t2 = squeezed
285 if not safe_allclose(t1, t2, atol=atol, rtol=rtol):
286 b, r = t1.float(), t2.float()
287 max_diff = torch.max(torch.abs(b - r)).item()
288 mean_diff = torch.mean(torch.abs(b - r)).item()
289 mismatches.append(
290 f"{key}: Value mismatch - max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}"
291 )
292 return mismatches
295def compare_tensors(
296 tensor1: torch.Tensor,
297 tensor2: torch.Tensor,
298 atol: float = 1e-5,
299 rtol: float = 1e-5,
300 name: str = "tensors",
301) -> BenchmarkResult:
302 """Compare two tensors and return a benchmark result.
304 Args:
305 tensor1: First tensor
306 tensor2: Second tensor
307 atol: Absolute tolerance
308 rtol: Relative tolerance
309 name: Name of the comparison
311 Returns:
312 BenchmarkResult with comparison details
313 """
314 # Check shapes
315 if tensor1.shape != tensor2.shape: 315 ↛ 316line 315 didn't jump to line 316 because the condition on line 315 was never true
316 return BenchmarkResult(
317 name=name,
318 severity=BenchmarkSeverity.DANGER,
319 message=f"Shape mismatch: {tensor1.shape} vs {tensor2.shape}",
320 passed=False,
321 )
323 if tensor1.device != tensor2.device: 323 ↛ 324line 323 didn't jump to line 324 because the condition on line 323 was never true
324 tensor1 = tensor1.cpu()
325 tensor2 = tensor2.cpu()
327 if tensor1.dtype != tensor2.dtype: 327 ↛ 328line 327 didn't jump to line 328 because the condition on line 327 was never true
328 tensor1 = tensor1.to(torch.float32)
329 tensor2 = tensor2.to(torch.float32)
331 if torch.allclose(tensor1, tensor2, atol=atol, rtol=rtol): 331 ↛ 339line 331 didn't jump to line 339 because the condition on line 331 was always true
332 return BenchmarkResult(
333 name=name,
334 severity=BenchmarkSeverity.INFO,
335 message="Tensors match within tolerance",
336 details={"atol": atol, "rtol": rtol},
337 )
339 diff = torch.abs(tensor1 - tensor2)
340 max_diff = diff.max().item()
341 mean_diff = diff.mean().item()
342 rel_diff = diff / (torch.abs(tensor1) + 1e-10)
343 mean_rel = rel_diff.mean().item()
345 return BenchmarkResult(
346 name=name,
347 severity=BenchmarkSeverity.DANGER,
348 message=f"Tensors differ: max_diff={max_diff:.6f}, mean_rel={mean_rel:.6f}",
349 details={
350 "max_diff": max_diff,
351 "mean_diff": mean_diff,
352 "mean_rel": mean_rel,
353 "atol": atol,
354 "rtol": rtol,
355 },
356 passed=False,
357 )
360def compare_scalars(
361 scalar1: Union[float, int],
362 scalar2: Union[float, int],
363 atol: float = 1e-5,
364 name: str = "scalars",
365) -> BenchmarkResult:
366 """Compare two scalar values and return a benchmark result.
368 Args:
369 scalar1: First scalar
370 scalar2: Second scalar
371 atol: Absolute tolerance
372 name: Name of the comparison
374 Returns:
375 BenchmarkResult with comparison details
376 """
377 diff = abs(float(scalar1) - float(scalar2))
379 if diff < atol: 379 ↛ 387line 379 didn't jump to line 387 because the condition on line 379 was always true
380 return BenchmarkResult(
381 name=name,
382 severity=BenchmarkSeverity.INFO,
383 message=f"Scalars match: {scalar1:.6f} ≈ {scalar2:.6f}",
384 details={"diff": diff, "atol": atol},
385 )
386 else:
387 return BenchmarkResult(
388 name=name,
389 severity=BenchmarkSeverity.DANGER,
390 message=f"Scalars differ: {scalar1:.6f} vs {scalar2:.6f}",
391 details={"diff": diff, "atol": atol},
392 passed=False,
393 )
396def format_results(results: List[BenchmarkResult]) -> str:
397 """Format a list of benchmark results for console output.
399 Args:
400 results: List of benchmark results
402 Returns:
403 Formatted string for console output
404 """
405 output = []
406 output.append("=" * 80)
407 output.append("BENCHMARK RESULTS")
408 output.append("=" * 80)
410 # Count by severity
411 severity_counts = {
412 BenchmarkSeverity.INFO: 0,
413 BenchmarkSeverity.WARNING: 0,
414 BenchmarkSeverity.DANGER: 0,
415 BenchmarkSeverity.ERROR: 0,
416 BenchmarkSeverity.SKIPPED: 0,
417 }
419 passed = 0
420 failed = 0
421 skipped = 0
423 for result in results:
424 severity_counts[result.severity] += 1
425 if result.severity == BenchmarkSeverity.SKIPPED:
426 skipped += 1
427 elif result.passed:
428 passed += 1
429 else:
430 failed += 1
432 # Summary
433 total = len(results)
434 run_tests = total - skipped
435 output.append(f"\nTotal: {total} tests")
436 if skipped > 0:
437 output.append(f"Run: {run_tests} tests")
438 output.append(f"Skipped: {skipped} tests")
439 if run_tests > 0:
440 output.append(f"Passed: {passed} ({passed/run_tests*100:.1f}%)")
441 output.append(f"Failed: {failed} ({failed/run_tests*100:.1f}%)")
442 output.append("")
443 output.append(f"🟢 INFO: {severity_counts[BenchmarkSeverity.INFO]}")
444 output.append(f"🟡 WARNING: {severity_counts[BenchmarkSeverity.WARNING]}")
445 output.append(f"🔴 DANGER: {severity_counts[BenchmarkSeverity.DANGER]}")
446 output.append(f"❌ ERROR: {severity_counts[BenchmarkSeverity.ERROR]}")
447 if skipped > 0:
448 output.append(f"⏭️ SKIPPED: {severity_counts[BenchmarkSeverity.SKIPPED]}")
449 output.append("")
450 output.append("-" * 80)
452 # Individual results
453 for result in results:
454 output.append(str(result))
455 output.append("")
457 output.append("=" * 80)
459 return "\n".join(output)
462def bridge_self_target_loss(bridge, test_text: str):
463 """Loss with the tokenized input as explicit labels.
465 Seq2seq bridges refuse label-less return_type="loss" (encoder input_ids are
466 not decoder targets), so every benchmark loss call routes through here.
467 """
468 labels = bridge.to_tokens(test_text)
469 return bridge(test_text, labels=labels, return_type="loss")