Coverage for transformer_lens/benchmarks/utils.py: 62%
195 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +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 # Sampling draws from the logits' device, so seed CUDA and MPS too; only the
27 # generators snapshotted here are seeded, so each one is restored afterwards.
28 cuda_devices = list(range(torch.cuda.device_count())) if torch.cuda.is_available() else []
29 mps_state = torch.mps.get_rng_state() if torch.backends.mps.is_available() else None
30 with torch.random.fork_rng(devices=cuda_devices):
31 try:
32 torch.default_generator.manual_seed(seed)
33 if cuda_devices:
34 torch.cuda.manual_seed_all(seed)
35 if mps_state is not None:
36 torch.mps.manual_seed(seed)
37 yield
38 finally:
39 if mps_state is not None:
40 torch.mps.set_rng_state(mps_state)
43def is_tiny_test_model(model_name: str) -> bool:
44 """Check if a model name belongs to a tiny/random test model."""
45 return any(pattern in model_name for pattern in TINY_TEST_MODEL_PATTERNS)
48# Hook patterns that bridge models inherently don't have because they use HF's
49# native implementation rather than reimplementing attention/MLP internals.
50BRIDGE_EXPECTED_MISSING_PATTERNS = [
51 "mlp.hook_pre",
52 "mlp.hook_post",
53 "hook_mlp_in",
54 "hook_mlp_out",
55 "attn.hook_rot_q",
56 "attn.hook_rot_k",
57 "hook_pos_embed",
58 "embed.ln.hook_scale",
59 "embed.ln.hook_normalized",
60 "attn.hook_q",
61 "attn.hook_k",
62 "attn.hook_v",
63 # cfg-gated attention hooks. These exist unconditionally on the attention
64 # bridge (so `run_with_cache` key lookups never KeyError) but only fire
65 # when their config flag is on. `benchmark_forward_hooks` runs with
66 # defaults (flags=False) so these correctly don't fire during that
67 # benchmark — suppressing them here prevents false "didn't fire"
68 # failures. The affirmative verification that they DO fire when flags
69 # are on lives in `benchmark_gated_hooks_fire`, which toggles each flag
70 # and asserts the relevant hooks capture activations.
71 "hook_result",
72 "hook_attn_in",
73 "hook_q_input",
74 "hook_k_input",
75 "hook_v_input",
76 "attn.hook_attn_scores",
77 "attn.hook_pattern",
78 # MoE per-expert hooks: Bridge uses HF's batched MoE forward pass via MoEBridge,
79 # which wraps the entire MoE module. The legacy reference created individual expert
80 # modules with per-expert hooks (e.g., blocks.0.mlp.experts.3.hook_pre).
81 "mlp.experts.",
82 "mlp.hook_experts",
83 # Parallel attention+MLP architectures (GPT-J, GPT-NeoX): HF has a single
84 # shared layer norm (ln_1), while HT creates a virtual ln2 that shares weights
85 # with ln1. The Bridge only wraps the actual HF ln_1, so ln2 hooks don't exist.
86 # These patterns only match "missing" hooks when ln2 is absent from the Bridge;
87 # for non-parallel architectures, the Bridge HAS ln2 and these won't be missing.
88 "ln2.hook_scale",
89 "ln2.hook_normalized",
90]
93def filter_expected_missing_hooks(hook_names: Collection[str]) -> list[str]:
94 """Filter out hook names that bridge models are expected to be missing."""
95 return [
96 h
97 for h in hook_names
98 if not any(pattern in h for pattern in BRIDGE_EXPECTED_MISSING_PATTERNS)
99 ]
102_DEFAULT_WAVEFORM_SAMPLES = 16000
105def build_modality_input(
106 bridge: Any,
107 batch_size: int = 1,
108 device: Optional[Union[str, torch.device]] = None,
109 dtype: Optional[torch.dtype] = None,
110) -> Optional[torch.Tensor]:
111 """Build a synthetic model input for a non-text bridge, or None for text models.
113 Audio and vision models take a raw tensor where text models take token ids, and
114 the shape is architecture-specific: spectrogram encoders (AST) need
115 ``[batch, max_length, num_mel_bins]`` where waveform encoders (HuBERT, wav2vec2)
116 need ``[batch, samples]``, and vision encoders need
117 ``[batch, num_channels, image_size, image_size]``. Shapes come from the HF config
118 so a checkpoint with non-default dimensions is still handled correctly.
119 """
120 cfg = getattr(bridge, "cfg", None)
121 if cfg is None: 121 ↛ 122line 121 didn't jump to line 122 because the condition on line 121 was never true
122 return None
124 hf_config = getattr(getattr(bridge, "original_model", None), "config", None)
126 def _from_config(name: str, default: int) -> int:
127 value = getattr(hf_config, name, None) if hf_config is not None else None
128 if value is None: 128 ↛ 129line 128 didn't jump to line 129 because the condition on line 128 was never true
129 value = getattr(cfg, name, None)
130 return int(value) if isinstance(value, int) else default
132 if getattr(cfg, "is_visual_model", False):
133 image_size = _from_config("image_size", 224)
134 num_channels = _from_config("num_channels", 3)
135 return torch.randn(
136 batch_size, num_channels, image_size, image_size, device=device, dtype=dtype
137 )
139 if getattr(cfg, "is_audio_model", False):
140 # num_mel_bins is what distinguishes a spectrogram encoder from a waveform one.
141 num_mel_bins = getattr(hf_config, "num_mel_bins", None) if hf_config is not None else None
142 if num_mel_bins is None:
143 return torch.randn(batch_size, _DEFAULT_WAVEFORM_SAMPLES, device=device, dtype=dtype)
144 max_length = _from_config("max_length", 1024)
145 return torch.randn(batch_size, max_length, int(num_mel_bins), device=device, dtype=dtype)
147 return None
150def safe_allclose(
151 tensor1: torch.Tensor,
152 tensor2: torch.Tensor,
153 atol: float = 1e-5,
154 rtol: float = 1e-5,
155) -> bool:
156 """torch.allclose that handles dtype and device mismatches."""
157 if tensor1.device != tensor2.device: 157 ↛ 158line 157 didn't jump to line 158 because the condition on line 157 was never true
158 tensor1 = tensor1.cpu()
159 tensor2 = tensor2.cpu()
160 if tensor1.dtype != tensor2.dtype: 160 ↛ 161line 160 didn't jump to line 161 because the condition on line 160 was never true
161 tensor1 = tensor1.to(torch.float32)
162 tensor2 = tensor2.to(torch.float32)
163 return torch.allclose(tensor1, tensor2, atol=atol, rtol=rtol)
166class BenchmarkSeverity(Enum):
167 """Severity levels for benchmark results."""
169 INFO = "info" # ✅ PASS - Model working perfectly, all checks passed
170 WARNING = "warning" # ⚠️ PASS with notes - Acceptable differences worth noting
171 DANGER = "danger" # ❌ FAIL - Significant mismatches or failures
172 ERROR = "error" # ❌ ERROR - Test crashed or couldn't run
173 SKIPPED = "skipped" # ⏭️ SKIPPED - Test skipped (e.g., no reference model available)
176@dataclass
177class BenchmarkResult:
178 """Result of a benchmark test."""
180 name: str
181 severity: BenchmarkSeverity
182 message: str
183 details: Optional[Dict[str, Any]] = None
184 passed: bool = True
185 phase: Optional[int] = None # Phase number (1, 2, 3, etc.)
187 def __str__(self) -> str:
188 """Format result for console output."""
189 severity_icons = {
190 BenchmarkSeverity.INFO: "🟢",
191 BenchmarkSeverity.WARNING: "🟡",
192 BenchmarkSeverity.DANGER: "🔴",
193 BenchmarkSeverity.ERROR: "❌",
194 BenchmarkSeverity.SKIPPED: "⏭️",
195 }
196 icon = severity_icons[self.severity]
198 if self.severity == BenchmarkSeverity.SKIPPED:
199 status = "SKIPPED"
200 else:
201 status = "PASS" if self.passed else "FAIL"
203 result = f"{icon} [{status}] {self.name}: {self.message}"
205 if self.details:
206 detail_lines = []
207 for key, value in self.details.items():
208 detail_lines.append(f" {key}: {value}")
209 result += "\n" + "\n".join(detail_lines)
211 return result
213 def print_immediate(self) -> None:
214 """Print this result immediately to console."""
215 print(str(self))
218@dataclass
219class PhaseReferenceData:
220 """Float32 reference data from Phase 1 for Phase 3 equivalence comparison."""
222 hf_logits: Optional[torch.Tensor] = None
223 hf_loss: Optional[float] = None
224 test_text: Optional[str] = None
227def make_capture_hook(storage: dict, name: str):
228 """Create a forward hook that captures activations into a dict.
230 Handles both raw tensors and tuples (extracts first element).
231 """
233 def hook_fn(tensor, hook):
234 if isinstance(tensor, torch.Tensor):
235 storage[name] = tensor.detach().clone()
236 elif isinstance(tensor, tuple) and len(tensor) > 0:
237 if isinstance(tensor[0], torch.Tensor):
238 storage[name] = tensor[0].detach().clone()
239 return tensor
241 return hook_fn
244def make_grad_capture_hook(storage: dict, name: str, return_none: bool = False):
245 """Create a backward hook that captures gradients into a dict.
247 Args:
248 storage: Dict to store captured gradients
249 name: Key name for storage
250 return_none: If True, return None (for backward hooks that shouldn't modify grads)
251 """
253 def hook_fn(tensor, hook=None):
254 if isinstance(tensor, torch.Tensor):
255 storage[name] = tensor.detach().clone()
256 elif isinstance(tensor, tuple) and len(tensor) > 0: 256 ↛ 259line 256 didn't jump to line 259 because the condition on line 256 was always true
257 if tensor[0] is not None and isinstance(tensor[0], torch.Tensor): 257 ↛ 259line 257 didn't jump to line 259 because the condition on line 257 was always true
258 storage[name] = tensor[0].detach().clone()
259 return None if return_none else tensor
261 return hook_fn
264def _squeeze_batch_dim(t1: torch.Tensor, t2: torch.Tensor):
265 """Handle batch dimension differences (e.g., [seq, dim] vs [1, seq, dim]).
267 Returns (t1, t2) with matching shapes, or None if shapes are incompatible.
268 """
269 if t1.shape == t2.shape:
270 return t1, t2
271 if t1.ndim == t2.ndim - 1 and t2.shape[0] == 1 and t1.shape == t2.shape[1:]:
272 return t1.unsqueeze(0), t2
273 if t2.ndim == t1.ndim - 1 and t1.shape[0] == 1 and t2.shape == t1.shape[1:]: 273 ↛ 274line 273 didn't jump to line 274 because the condition on line 273 was never true
274 return t1, t2.unsqueeze(0)
275 return None
278def compare_activation_dicts(
279 dict1: Dict[str, torch.Tensor],
280 dict2: Dict[str, torch.Tensor],
281 atol: float = 1e-5,
282 rtol: float = 0.0,
283) -> List[str]:
284 """Compare two activation/gradient dicts, returning mismatch descriptions.
286 Handles batch-dim squeezing and dtype/device normalization.
287 """
288 mismatches = []
289 common_keys = sorted(set(dict1.keys()) & set(dict2.keys()))
290 for key in common_keys:
291 t1, t2 = dict1[key], dict2[key]
292 squeezed = _squeeze_batch_dim(t1, t2)
293 if squeezed is None:
294 mismatches.append(f"{key}: Shape mismatch - {t1.shape} vs {t2.shape}")
295 continue
296 t1, t2 = squeezed
297 if not safe_allclose(t1, t2, atol=atol, rtol=rtol):
298 b, r = t1.float(), t2.float()
299 max_diff = torch.max(torch.abs(b - r)).item()
300 mean_diff = torch.mean(torch.abs(b - r)).item()
301 mismatches.append(
302 f"{key}: Value mismatch - max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}"
303 )
304 return mismatches
307def compare_tensors(
308 tensor1: torch.Tensor,
309 tensor2: torch.Tensor,
310 atol: float = 1e-5,
311 rtol: float = 1e-5,
312 name: str = "tensors",
313) -> BenchmarkResult:
314 """Compare two tensors and return a benchmark result.
316 Args:
317 tensor1: First tensor
318 tensor2: Second tensor
319 atol: Absolute tolerance
320 rtol: Relative tolerance
321 name: Name of the comparison
323 Returns:
324 BenchmarkResult with comparison details
325 """
326 # Check shapes
327 if tensor1.shape != tensor2.shape: 327 ↛ 328line 327 didn't jump to line 328 because the condition on line 327 was never true
328 return BenchmarkResult(
329 name=name,
330 severity=BenchmarkSeverity.DANGER,
331 message=f"Shape mismatch: {tensor1.shape} vs {tensor2.shape}",
332 passed=False,
333 )
335 if tensor1.device != tensor2.device: 335 ↛ 336line 335 didn't jump to line 336 because the condition on line 335 was never true
336 tensor1 = tensor1.cpu()
337 tensor2 = tensor2.cpu()
339 if tensor1.dtype != tensor2.dtype: 339 ↛ 340line 339 didn't jump to line 340 because the condition on line 339 was never true
340 tensor1 = tensor1.to(torch.float32)
341 tensor2 = tensor2.to(torch.float32)
343 if torch.allclose(tensor1, tensor2, atol=atol, rtol=rtol): 343 ↛ 351line 343 didn't jump to line 351 because the condition on line 343 was always true
344 return BenchmarkResult(
345 name=name,
346 severity=BenchmarkSeverity.INFO,
347 message="Tensors match within tolerance",
348 details={"atol": atol, "rtol": rtol},
349 )
351 diff = torch.abs(tensor1 - tensor2)
352 max_diff = diff.max().item()
353 mean_diff = diff.mean().item()
354 rel_diff = diff / (torch.abs(tensor1) + 1e-10)
355 mean_rel = rel_diff.mean().item()
357 return BenchmarkResult(
358 name=name,
359 severity=BenchmarkSeverity.DANGER,
360 message=f"Tensors differ: max_diff={max_diff:.6f}, mean_rel={mean_rel:.6f}",
361 details={
362 "max_diff": max_diff,
363 "mean_diff": mean_diff,
364 "mean_rel": mean_rel,
365 "atol": atol,
366 "rtol": rtol,
367 },
368 passed=False,
369 )
372def compare_scalars(
373 scalar1: Union[float, int],
374 scalar2: Union[float, int],
375 atol: float = 1e-5,
376 name: str = "scalars",
377) -> BenchmarkResult:
378 """Compare two scalar values and return a benchmark result.
380 Args:
381 scalar1: First scalar
382 scalar2: Second scalar
383 atol: Absolute tolerance
384 name: Name of the comparison
386 Returns:
387 BenchmarkResult with comparison details
388 """
389 diff = abs(float(scalar1) - float(scalar2))
391 if diff < atol: 391 ↛ 399line 391 didn't jump to line 399 because the condition on line 391 was always true
392 return BenchmarkResult(
393 name=name,
394 severity=BenchmarkSeverity.INFO,
395 message=f"Scalars match: {scalar1:.6f} ≈ {scalar2:.6f}",
396 details={"diff": diff, "atol": atol},
397 )
398 else:
399 return BenchmarkResult(
400 name=name,
401 severity=BenchmarkSeverity.DANGER,
402 message=f"Scalars differ: {scalar1:.6f} vs {scalar2:.6f}",
403 details={"diff": diff, "atol": atol},
404 passed=False,
405 )
408def format_results(results: List[BenchmarkResult]) -> str:
409 """Format a list of benchmark results for console output.
411 Args:
412 results: List of benchmark results
414 Returns:
415 Formatted string for console output
416 """
417 output = []
418 output.append("=" * 80)
419 output.append("BENCHMARK RESULTS")
420 output.append("=" * 80)
422 # Count by severity
423 severity_counts = {
424 BenchmarkSeverity.INFO: 0,
425 BenchmarkSeverity.WARNING: 0,
426 BenchmarkSeverity.DANGER: 0,
427 BenchmarkSeverity.ERROR: 0,
428 BenchmarkSeverity.SKIPPED: 0,
429 }
431 passed = 0
432 failed = 0
433 skipped = 0
435 for result in results:
436 severity_counts[result.severity] += 1
437 if result.severity == BenchmarkSeverity.SKIPPED:
438 skipped += 1
439 elif result.passed:
440 passed += 1
441 else:
442 failed += 1
444 # Summary
445 total = len(results)
446 run_tests = total - skipped
447 output.append(f"\nTotal: {total} tests")
448 if skipped > 0:
449 output.append(f"Run: {run_tests} tests")
450 output.append(f"Skipped: {skipped} tests")
451 if run_tests > 0:
452 output.append(f"Passed: {passed} ({passed/run_tests*100:.1f}%)")
453 output.append(f"Failed: {failed} ({failed/run_tests*100:.1f}%)")
454 output.append("")
455 output.append(f"🟢 INFO: {severity_counts[BenchmarkSeverity.INFO]}")
456 output.append(f"🟡 WARNING: {severity_counts[BenchmarkSeverity.WARNING]}")
457 output.append(f"🔴 DANGER: {severity_counts[BenchmarkSeverity.DANGER]}")
458 output.append(f"❌ ERROR: {severity_counts[BenchmarkSeverity.ERROR]}")
459 if skipped > 0:
460 output.append(f"⏭️ SKIPPED: {severity_counts[BenchmarkSeverity.SKIPPED]}")
461 output.append("")
462 output.append("-" * 80)
464 # Individual results
465 for result in results:
466 output.append(str(result))
467 output.append("")
469 output.append("=" * 80)
471 return "\n".join(output)
474def bridge_self_target_loss(bridge, test_text: str):
475 """Loss with the tokenized input as explicit labels.
477 Seq2seq bridges refuse label-less return_type="loss" (encoder input_ids are
478 not decoder targets), so every benchmark loss call routes through here.
479 """
480 labels = bridge.to_tokens(test_text)
481 return bridge(test_text, labels=labels, return_type="loss")