Coverage for transformer_lens/benchmarks/utils.py: 62%
186 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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. HookedTransformer creates 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 "mlp.hook_expert_indices",
72 "mlp.hook_expert_weights",
73 # Parallel attention+MLP architectures (GPT-J, GPT-NeoX): HF has a single
74 # shared layer norm (ln_1), while HT creates a virtual ln2 that shares weights
75 # with ln1. The Bridge only wraps the actual HF ln_1, so ln2 hooks don't exist.
76 # These patterns only match "missing" hooks when ln2 is absent from the Bridge;
77 # for non-parallel architectures, the Bridge HAS ln2 and these won't be missing.
78 "ln2.hook_scale",
79 "ln2.hook_normalized",
80]
83def filter_expected_missing_hooks(hook_names: Collection[str]) -> list[str]:
84 """Filter out hook names that bridge models are expected to be missing."""
85 return [
86 h
87 for h in hook_names
88 if not any(pattern in h for pattern in BRIDGE_EXPECTED_MISSING_PATTERNS)
89 ]
92_DEFAULT_WAVEFORM_SAMPLES = 16000
95def build_modality_input(
96 bridge: Any,
97 batch_size: int = 1,
98 device: Optional[Union[str, torch.device]] = None,
99 dtype: Optional[torch.dtype] = None,
100) -> Optional[torch.Tensor]:
101 """Build a synthetic model input for a non-text bridge, or None for text models.
103 Audio and vision models take a raw tensor where text models take token ids, and
104 the shape is architecture-specific: spectrogram encoders (AST) need
105 ``[batch, max_length, num_mel_bins]`` where waveform encoders (HuBERT, wav2vec2)
106 need ``[batch, samples]``, and vision encoders need
107 ``[batch, num_channels, image_size, image_size]``. Shapes come from the HF config
108 so a checkpoint with non-default dimensions is still handled correctly.
109 """
110 cfg = getattr(bridge, "cfg", None)
111 if cfg is None: 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true
112 return None
114 hf_config = getattr(getattr(bridge, "original_model", None), "config", None)
116 def _from_config(name: str, default: int) -> int:
117 value = getattr(hf_config, name, None) if hf_config is not None else None
118 if value is None: 118 ↛ 119line 118 didn't jump to line 119 because the condition on line 118 was never true
119 value = getattr(cfg, name, None)
120 return int(value) if isinstance(value, int) else default
122 if getattr(cfg, "is_visual_model", False):
123 image_size = _from_config("image_size", 224)
124 num_channels = _from_config("num_channels", 3)
125 return torch.randn(
126 batch_size, num_channels, image_size, image_size, device=device, dtype=dtype
127 )
129 if getattr(cfg, "is_audio_model", False):
130 # num_mel_bins is what distinguishes a spectrogram encoder from a waveform one.
131 num_mel_bins = getattr(hf_config, "num_mel_bins", None) if hf_config is not None else None
132 if num_mel_bins is None:
133 return torch.randn(batch_size, _DEFAULT_WAVEFORM_SAMPLES, device=device, dtype=dtype)
134 max_length = _from_config("max_length", 1024)
135 return torch.randn(batch_size, max_length, int(num_mel_bins), device=device, dtype=dtype)
137 return None
140def safe_allclose(
141 tensor1: torch.Tensor,
142 tensor2: torch.Tensor,
143 atol: float = 1e-5,
144 rtol: float = 1e-5,
145) -> bool:
146 """torch.allclose that handles dtype and device mismatches."""
147 if tensor1.device != tensor2.device: 147 ↛ 148line 147 didn't jump to line 148 because the condition on line 147 was never true
148 tensor1 = tensor1.cpu()
149 tensor2 = tensor2.cpu()
150 if tensor1.dtype != tensor2.dtype: 150 ↛ 151line 150 didn't jump to line 151 because the condition on line 150 was never true
151 tensor1 = tensor1.to(torch.float32)
152 tensor2 = tensor2.to(torch.float32)
153 return torch.allclose(tensor1, tensor2, atol=atol, rtol=rtol)
156class BenchmarkSeverity(Enum):
157 """Severity levels for benchmark results."""
159 INFO = "info" # ✅ PASS - Model working perfectly, all checks passed
160 WARNING = "warning" # ⚠️ PASS with notes - Acceptable differences worth noting
161 DANGER = "danger" # ❌ FAIL - Significant mismatches or failures
162 ERROR = "error" # ❌ ERROR - Test crashed or couldn't run
163 SKIPPED = "skipped" # ⏭️ SKIPPED - Test skipped (e.g., no reference model available)
166@dataclass
167class BenchmarkResult:
168 """Result of a benchmark test."""
170 name: str
171 severity: BenchmarkSeverity
172 message: str
173 details: Optional[Dict[str, Any]] = None
174 passed: bool = True
175 phase: Optional[int] = None # Phase number (1, 2, 3, etc.)
177 def __str__(self) -> str:
178 """Format result for console output."""
179 severity_icons = {
180 BenchmarkSeverity.INFO: "🟢",
181 BenchmarkSeverity.WARNING: "🟡",
182 BenchmarkSeverity.DANGER: "🔴",
183 BenchmarkSeverity.ERROR: "❌",
184 BenchmarkSeverity.SKIPPED: "⏭️",
185 }
186 icon = severity_icons[self.severity]
188 if self.severity == BenchmarkSeverity.SKIPPED:
189 status = "SKIPPED"
190 else:
191 status = "PASS" if self.passed else "FAIL"
193 result = f"{icon} [{status}] {self.name}: {self.message}"
195 if self.details:
196 detail_lines = []
197 for key, value in self.details.items():
198 detail_lines.append(f" {key}: {value}")
199 result += "\n" + "\n".join(detail_lines)
201 return result
203 def print_immediate(self) -> None:
204 """Print this result immediately to console."""
205 print(str(self))
208@dataclass
209class PhaseReferenceData:
210 """Float32 reference data from Phase 1 for Phase 3 equivalence comparison."""
212 hf_logits: Optional[torch.Tensor] = None
213 hf_loss: Optional[float] = None
214 test_text: Optional[str] = None
217def make_capture_hook(storage: dict, name: str):
218 """Create a forward hook that captures activations into a dict.
220 Handles both raw tensors and tuples (extracts first element).
221 """
223 def hook_fn(tensor, hook):
224 if isinstance(tensor, torch.Tensor):
225 storage[name] = tensor.detach().clone()
226 elif isinstance(tensor, tuple) and len(tensor) > 0:
227 if isinstance(tensor[0], torch.Tensor):
228 storage[name] = tensor[0].detach().clone()
229 return tensor
231 return hook_fn
234def make_grad_capture_hook(storage: dict, name: str, return_none: bool = False):
235 """Create a backward hook that captures gradients into a dict.
237 Args:
238 storage: Dict to store captured gradients
239 name: Key name for storage
240 return_none: If True, return None (for backward hooks that shouldn't modify grads)
241 """
243 def hook_fn(tensor, hook=None):
244 if isinstance(tensor, torch.Tensor):
245 storage[name] = tensor.detach().clone()
246 elif isinstance(tensor, tuple) and len(tensor) > 0: 246 ↛ 249line 246 didn't jump to line 249 because the condition on line 246 was always true
247 if tensor[0] is not None and isinstance(tensor[0], torch.Tensor): 247 ↛ 249line 247 didn't jump to line 249 because the condition on line 247 was always true
248 storage[name] = tensor[0].detach().clone()
249 return None if return_none else tensor
251 return hook_fn
254def _squeeze_batch_dim(t1: torch.Tensor, t2: torch.Tensor):
255 """Handle batch dimension differences (e.g., [seq, dim] vs [1, seq, dim]).
257 Returns (t1, t2) with matching shapes, or None if shapes are incompatible.
258 """
259 if t1.shape == t2.shape:
260 return t1, t2
261 if t1.ndim == t2.ndim - 1 and t2.shape[0] == 1 and t1.shape == t2.shape[1:]:
262 return t1.unsqueeze(0), t2
263 if t2.ndim == t1.ndim - 1 and t1.shape[0] == 1 and t2.shape == t1.shape[1:]: 263 ↛ 264line 263 didn't jump to line 264 because the condition on line 263 was never true
264 return t1, t2.unsqueeze(0)
265 return None
268def compare_activation_dicts(
269 dict1: Dict[str, torch.Tensor],
270 dict2: Dict[str, torch.Tensor],
271 atol: float = 1e-5,
272 rtol: float = 0.0,
273) -> List[str]:
274 """Compare two activation/gradient dicts, returning mismatch descriptions.
276 Handles batch-dim squeezing and dtype/device normalization.
277 """
278 mismatches = []
279 common_keys = sorted(set(dict1.keys()) & set(dict2.keys()))
280 for key in common_keys:
281 t1, t2 = dict1[key], dict2[key]
282 squeezed = _squeeze_batch_dim(t1, t2)
283 if squeezed is None:
284 mismatches.append(f"{key}: Shape mismatch - {t1.shape} vs {t2.shape}")
285 continue
286 t1, t2 = squeezed
287 if not safe_allclose(t1, t2, atol=atol, rtol=rtol):
288 b, r = t1.float(), t2.float()
289 max_diff = torch.max(torch.abs(b - r)).item()
290 mean_diff = torch.mean(torch.abs(b - r)).item()
291 mismatches.append(
292 f"{key}: Value mismatch - max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}"
293 )
294 return mismatches
297def compare_tensors(
298 tensor1: torch.Tensor,
299 tensor2: torch.Tensor,
300 atol: float = 1e-5,
301 rtol: float = 1e-5,
302 name: str = "tensors",
303) -> BenchmarkResult:
304 """Compare two tensors and return a benchmark result.
306 Args:
307 tensor1: First tensor
308 tensor2: Second tensor
309 atol: Absolute tolerance
310 rtol: Relative tolerance
311 name: Name of the comparison
313 Returns:
314 BenchmarkResult with comparison details
315 """
316 # Check shapes
317 if tensor1.shape != tensor2.shape: 317 ↛ 318line 317 didn't jump to line 318 because the condition on line 317 was never true
318 return BenchmarkResult(
319 name=name,
320 severity=BenchmarkSeverity.DANGER,
321 message=f"Shape mismatch: {tensor1.shape} vs {tensor2.shape}",
322 passed=False,
323 )
325 if tensor1.device != tensor2.device: 325 ↛ 326line 325 didn't jump to line 326 because the condition on line 325 was never true
326 tensor1 = tensor1.cpu()
327 tensor2 = tensor2.cpu()
329 if tensor1.dtype != tensor2.dtype: 329 ↛ 330line 329 didn't jump to line 330 because the condition on line 329 was never true
330 tensor1 = tensor1.to(torch.float32)
331 tensor2 = tensor2.to(torch.float32)
333 if torch.allclose(tensor1, tensor2, atol=atol, rtol=rtol): 333 ↛ 341line 333 didn't jump to line 341 because the condition on line 333 was always true
334 return BenchmarkResult(
335 name=name,
336 severity=BenchmarkSeverity.INFO,
337 message="Tensors match within tolerance",
338 details={"atol": atol, "rtol": rtol},
339 )
341 diff = torch.abs(tensor1 - tensor2)
342 max_diff = diff.max().item()
343 mean_diff = diff.mean().item()
344 rel_diff = diff / (torch.abs(tensor1) + 1e-10)
345 mean_rel = rel_diff.mean().item()
347 return BenchmarkResult(
348 name=name,
349 severity=BenchmarkSeverity.DANGER,
350 message=f"Tensors differ: max_diff={max_diff:.6f}, mean_rel={mean_rel:.6f}",
351 details={
352 "max_diff": max_diff,
353 "mean_diff": mean_diff,
354 "mean_rel": mean_rel,
355 "atol": atol,
356 "rtol": rtol,
357 },
358 passed=False,
359 )
362def compare_scalars(
363 scalar1: Union[float, int],
364 scalar2: Union[float, int],
365 atol: float = 1e-5,
366 name: str = "scalars",
367) -> BenchmarkResult:
368 """Compare two scalar values and return a benchmark result.
370 Args:
371 scalar1: First scalar
372 scalar2: Second scalar
373 atol: Absolute tolerance
374 name: Name of the comparison
376 Returns:
377 BenchmarkResult with comparison details
378 """
379 diff = abs(float(scalar1) - float(scalar2))
381 if diff < atol: 381 ↛ 389line 381 didn't jump to line 389 because the condition on line 381 was always true
382 return BenchmarkResult(
383 name=name,
384 severity=BenchmarkSeverity.INFO,
385 message=f"Scalars match: {scalar1:.6f} ≈ {scalar2:.6f}",
386 details={"diff": diff, "atol": atol},
387 )
388 else:
389 return BenchmarkResult(
390 name=name,
391 severity=BenchmarkSeverity.DANGER,
392 message=f"Scalars differ: {scalar1:.6f} vs {scalar2:.6f}",
393 details={"diff": diff, "atol": atol},
394 passed=False,
395 )
398def format_results(results: List[BenchmarkResult]) -> str:
399 """Format a list of benchmark results for console output.
401 Args:
402 results: List of benchmark results
404 Returns:
405 Formatted string for console output
406 """
407 output = []
408 output.append("=" * 80)
409 output.append("BENCHMARK RESULTS")
410 output.append("=" * 80)
412 # Count by severity
413 severity_counts = {
414 BenchmarkSeverity.INFO: 0,
415 BenchmarkSeverity.WARNING: 0,
416 BenchmarkSeverity.DANGER: 0,
417 BenchmarkSeverity.ERROR: 0,
418 BenchmarkSeverity.SKIPPED: 0,
419 }
421 passed = 0
422 failed = 0
423 skipped = 0
425 for result in results:
426 severity_counts[result.severity] += 1
427 if result.severity == BenchmarkSeverity.SKIPPED:
428 skipped += 1
429 elif result.passed:
430 passed += 1
431 else:
432 failed += 1
434 # Summary
435 total = len(results)
436 run_tests = total - skipped
437 output.append(f"\nTotal: {total} tests")
438 if skipped > 0:
439 output.append(f"Run: {run_tests} tests")
440 output.append(f"Skipped: {skipped} tests")
441 if run_tests > 0:
442 output.append(f"Passed: {passed} ({passed/run_tests*100:.1f}%)")
443 output.append(f"Failed: {failed} ({failed/run_tests*100:.1f}%)")
444 output.append("")
445 output.append(f"🟢 INFO: {severity_counts[BenchmarkSeverity.INFO]}")
446 output.append(f"🟡 WARNING: {severity_counts[BenchmarkSeverity.WARNING]}")
447 output.append(f"🔴 DANGER: {severity_counts[BenchmarkSeverity.DANGER]}")
448 output.append(f"❌ ERROR: {severity_counts[BenchmarkSeverity.ERROR]}")
449 if skipped > 0:
450 output.append(f"⏭️ SKIPPED: {severity_counts[BenchmarkSeverity.SKIPPED]}")
451 output.append("")
452 output.append("-" * 80)
454 # Individual results
455 for result in results:
456 output.append(str(result))
457 output.append("")
459 output.append("=" * 80)
461 return "\n".join(output)
464def bridge_self_target_loss(bridge, test_text: str):
465 """Loss with the tokenized input as explicit labels.
467 Seq2seq bridges refuse label-less return_type="loss" (encoder input_ids are
468 not decoder targets), so every benchmark loss call routes through here.
469 """
470 labels = bridge.to_tokens(test_text)
471 return bridge(test_text, labels=labels, return_type="loss")