Coverage for transformer_lens/benchmarks/component_outputs.py: 66%
467 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Comprehensive component benchmarking utility for TransformerBridge.
3This module provides utilities to benchmark all standard components in a TransformerBridge
4model against their HuggingFace equivalents, ensuring output parity.
5"""
7from __future__ import annotations
9from dataclasses import dataclass, field
10from typing import Any, Callable, Dict, List, Optional, Tuple, cast
12import torch
13from torch import nn
15from transformer_lens.benchmarks.utils import build_modality_input
16from transformer_lens.config import TransformerBridgeConfig
17from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
18from transformer_lens.model_bridge.generalized_components.base import (
19 GeneralizedComponent,
20)
23@dataclass
24class ComponentTestResult:
25 """Result of testing a single component."""
27 component_path: str
28 component_type: str
29 passed: bool
30 max_diff: float
31 mean_diff: float
32 output_shape: Tuple[int, ...]
33 error_message: Optional[str] = None
34 percentile_diffs: Optional[Dict[str, float]] = None # 50th, 90th, 99th percentile diffs
36 def get_failure_severity(self) -> str:
37 """Categorize the severity of a failure.
39 Returns:
40 Severity level: "critical", "high", "medium", "low", or "pass"
41 """
42 if self.passed:
43 return "pass"
44 if self.error_message:
45 return "critical"
46 if self.max_diff > 1e-1:
47 return "critical"
48 elif self.max_diff > 1e-3:
49 return "high"
50 elif self.max_diff > 1e-4:
51 return "medium"
52 else:
53 return "low"
56@dataclass
57class BenchmarkReport:
58 """Complete benchmark report for all components."""
60 model_name: str
61 total_components: int
62 passed_components: int
63 failed_components: int
64 component_results: List[ComponentTestResult] = field(default_factory=list)
66 @property
67 def pass_rate(self) -> float:
68 """Calculate the pass rate as a percentage."""
69 if self.total_components == 0: 69 ↛ 70line 69 didn't jump to line 70 because the condition on line 69 was never true
70 return 0.0
71 return (self.passed_components / self.total_components) * 100
73 def print_summary(self, verbose: bool = False) -> None:
74 """Print a summary of the benchmark results.
76 Args:
77 verbose: If True, print details for all components. If False, only print failures.
78 """
79 print("\n" + "=" * 80)
80 print(f"Component Benchmark Report: {self.model_name}")
81 print("=" * 80)
82 print(f"Total components tested: {self.total_components}")
83 print(f"Passed: {self.passed_components} ({self.pass_rate:.1f}%)")
84 print(f"Failed: {self.failed_components}")
85 print("=" * 80)
87 if verbose: 87 ↛ 88line 87 didn't jump to line 88 because the condition on line 87 was never true
88 print("\nAll Component Results:")
89 print("-" * 80)
90 for result in self.component_results:
91 self._print_component_result(result)
92 elif self.failed_components > 0: 92 ↛ 93line 92 didn't jump to line 93 because the condition on line 92 was never true
93 print("\nFailed Components:")
94 print("-" * 80)
95 for result in self.component_results:
96 if not result.passed:
97 self._print_component_result(result)
99 print("=" * 80 + "\n")
101 def _print_component_result(self, result: ComponentTestResult) -> None:
102 """Print details of a single component result."""
103 status = "✓ PASS" if result.passed else "✗ FAIL"
104 severity = result.get_failure_severity()
106 # Add severity indicator for failures
107 if not result.passed and severity != "critical":
108 status = f"{status} [{severity.upper()}]"
110 print(f"{status} | {result.component_path}")
111 print(f" Type: {result.component_type}")
112 print(f" Shape: {result.output_shape}")
113 print(f" Max diff: {result.max_diff:.6e}")
114 print(f" Mean diff: {result.mean_diff:.6e}")
116 if result.percentile_diffs:
117 print(f" Percentile diffs:")
118 for percentile, diff in sorted(result.percentile_diffs.items()):
119 print(f" {percentile}: {diff:.6e}")
121 if result.error_message:
122 print(f" Error: {result.error_message}")
123 print()
125 def get_component_type_summary(self) -> Dict[str, Dict[str, int]]:
126 """Get a summary of results grouped by component type.
128 Returns:
129 Dictionary mapping component types to their pass/fail counts
130 """
131 summary: Dict[str, Dict[str, int]] = {}
133 for result in self.component_results:
134 comp_type = result.component_type
135 if comp_type not in summary:
136 summary[comp_type] = {"passed": 0, "failed": 0, "total": 0}
138 summary[comp_type]["total"] += 1
139 if result.passed: 139 ↛ 142line 139 didn't jump to line 142 because the condition on line 139 was always true
140 summary[comp_type]["passed"] += 1
141 else:
142 summary[comp_type]["failed"] += 1
144 return summary
146 def get_failure_by_severity(self) -> Dict[str, List[ComponentTestResult]]:
147 """Group failures by severity level.
149 Returns:
150 Dictionary mapping severity levels to lists of failed components
151 """
152 failures: Dict[str, List[ComponentTestResult]] = {
153 "critical": [],
154 "high": [],
155 "medium": [],
156 "low": [],
157 }
159 for result in self.component_results:
160 if not result.passed:
161 severity = result.get_failure_severity()
162 if severity in failures:
163 failures[severity].append(result)
165 return failures
167 def print_detailed_analysis(self) -> None:
168 """Print detailed analysis of benchmark results."""
169 print("\n" + "=" * 80)
170 print("Detailed Benchmark Analysis")
171 print("=" * 80)
173 # Component type summary
174 print("\nResults by Component Type:")
175 print("-" * 80)
176 type_summary = self.get_component_type_summary()
177 for comp_type, stats in sorted(type_summary.items()):
178 pass_rate = (stats["passed"] / stats["total"]) * 100 if stats["total"] > 0 else 0
179 print(
180 f"{comp_type:30s}: {stats['passed']:3d}/{stats['total']:3d} passed ({pass_rate:5.1f}%)"
181 )
183 # Failure severity analysis
184 if self.failed_components > 0:
185 print("\nFailures by Severity:")
186 print("-" * 80)
187 failures_by_severity = self.get_failure_by_severity()
188 for severity in ["critical", "high", "medium", "low"]:
189 count = len(failures_by_severity[severity])
190 if count > 0:
191 print(f"{severity.upper():10s}: {count} component(s)")
192 for result in failures_by_severity[severity][:3]: # Show first 3
193 print(f" - {result.component_path} (max_diff: {result.max_diff:.2e})")
194 if count > 3:
195 print(f" ... and {count - 3} more")
197 print("=" * 80 + "\n")
200def _is_ssm_mixer_internal(component_path: str) -> bool:
201 """True for a submodule *inside* an SSM/recurrent mixer slot (``.mixer``/``.linear_attn``).
203 Such submodules wrap the identical HF module (parity is covered by forward_pass_logits)
204 and take SSM-internal shapes, not the ``[b, seq, d_model]`` residual the isolated harness
205 feeds — so they are skipped. The mixer node itself (path ending in the slot) is not.
206 """
207 return any(slot in component_path.split(".")[:-1] for slot in ("mixer", "linear_attn"))
210class ComponentBenchmarker:
211 """Benchmarking utility for testing TransformerBridge components against HuggingFace."""
213 def _is_delegated_block(self) -> bool:
214 """Return True if the blocks component has maintain_native_attention set."""
215 blocks = (
216 getattr(self.adapter, "component_mapping", {}).get("blocks")
217 if self.adapter is not None
218 else None
219 )
220 return getattr(blocks, "maintain_native_attention", False)
222 def __init__(
223 self,
224 bridge_model: nn.Module,
225 hf_model: nn.Module,
226 adapter: ArchitectureAdapter,
227 cfg: TransformerBridgeConfig,
228 atol: float = 1e-4,
229 rtol: float = 1e-4,
230 ):
231 """Initialize the component benchmarker.
233 Args:
234 bridge_model: The TransformerBridge model
235 hf_model: The HuggingFace model
236 adapter: The architecture adapter for mapping components
237 cfg: The model configuration
238 atol: Absolute tolerance for comparing outputs
239 rtol: Relative tolerance for comparing outputs
240 """
241 self.bridge_model = bridge_model
242 self.hf_model = hf_model
243 self.adapter = adapter
244 self.cfg = cfg
246 # Reconcile dtypes: upcast both models to the higher-precision dtype.
247 self._bridge_was_upcast = False
248 self._bridge_original_dtype: Optional[torch.dtype] = None
249 try:
250 hf_dtype = next(hf_model.parameters()).dtype
251 except StopIteration:
252 hf_dtype = torch.float32
253 try:
254 bridge_dtype = next(bridge_model.parameters()).dtype
255 except StopIteration:
256 bridge_dtype = torch.float32
257 if hf_dtype != bridge_dtype:
258 # Upcast to the higher-precision dtype
259 target = hf_dtype if hf_dtype.itemsize >= bridge_dtype.itemsize else bridge_dtype
260 if bridge_dtype != target: 260 ↛ 261line 260 didn't jump to line 261 because the condition on line 260 was never true
261 self._bridge_original_dtype = bridge_dtype
262 bridge_model.to(target)
263 self._bridge_was_upcast = True
264 if hf_dtype != target: 264 ↛ 266line 264 didn't jump to line 266 because the condition on line 264 was always true
265 hf_model.to(target)
266 self.test_dtype = hf_dtype if hf_dtype.itemsize >= bridge_dtype.itemsize else bridge_dtype
268 # Adjust tolerances based on dtype for reduced precision formats
269 model_dtype = getattr(cfg, "dtype", torch.float32)
270 if model_dtype == torch.bfloat16: 270 ↛ 275line 270 didn't jump to line 275 because the condition on line 270 was never true
271 # bfloat16 has ~7 bits of precision (3 decimal digits)
272 # Use more lenient tolerance
273 # Normalization layers (RMSNorm/LayerNorm) can have larger errors due to
274 # square roots and divisions, so use 0.3 tolerance
275 self.atol = max(atol, 0.3)
276 self.rtol = max(rtol, 0.3)
277 elif model_dtype == torch.float16: 277 ↛ 279line 277 didn't jump to line 279 because the condition on line 277 was never true
278 # float16 has ~10 bits of precision (3-4 decimal digits)
279 self.atol = max(atol, 5e-3)
280 self.rtol = max(rtol, 5e-3)
281 else:
282 # float32 or float64 - use provided tolerances
283 self.atol = atol
284 self.rtol = rtol
286 def benchmark_all_components(
287 self,
288 test_inputs: Optional[Dict[str, torch.Tensor]] = None,
289 skip_components: Optional[List[str]] = None,
290 ) -> BenchmarkReport:
291 """Benchmark all components in the model.
293 Args:
294 test_inputs: Optional dictionary of pre-generated test inputs.
295 If None, will generate default inputs.
296 skip_components: Optional list of component paths to skip
298 Returns:
299 BenchmarkReport with results for all tested components
300 """
301 skip_components = skip_components or []
302 # Adapters can exclude subcomponents whose isolated forward cannot run
303 # on synthesized probes (e.g. fused top-k routers) by path suffix.
304 skip_suffixes = tuple(getattr(self.adapter, "component_test_skip_suffixes", ()))
305 component_mapping = self.adapter.get_component_mapping()
307 # Generate test inputs if not provided
308 if test_inputs is None: 308 ↛ 311line 308 didn't jump to line 311 because the condition on line 308 was always true
309 test_inputs = self._generate_test_inputs()
311 results: List[ComponentTestResult] = []
313 # Block-type components that need to be tested recursively by layer
314 # (they are ModuleLists that don't have direct forward methods)
315 block_components = {"blocks", "encoder_blocks", "decoder_blocks", "L_blocks", "H_blocks"}
317 # Test top-level components (embed, pos_embed, ln_final, unembed)
318 for comp_name, component in component_mapping.items():
319 if comp_name in skip_components: 319 ↛ 320line 319 didn't jump to line 320 because the condition on line 319 was never true
320 continue
322 if comp_name in block_components:
323 # Handle blocks separately - test their subcomponents by layer
324 continue
326 result = self._test_component(comp_name, component, test_inputs)
327 if result is not None: 327 ↛ 318line 327 didn't jump to line 318 because the condition on line 327 was always true
328 results.append(result)
330 # Test block components recursively
331 for block_type in block_components:
332 if block_type in component_mapping and block_type not in skip_components:
333 blocks_component = component_mapping[block_type]
334 # Derive the length from the actual stack: asymmetric
335 # encoder-decoder models (e.g. Blenderbot's 2/12) have
336 # per-stack counts that cfg.n_layers cannot represent.
337 bound_blocks = getattr(self.bridge_model, block_type)
338 n_layers = len(bound_blocks)
340 for layer_idx in range(n_layers):
341 # Get the actual block to check which submodules were bound
342 actual_block = getattr(self.bridge_model, block_type)[layer_idx]
343 for subcomp_name, subcomponent in blocks_component.submodules.items():
344 # Skip optional submodules absent on this layer (hybrid architectures)
345 if subcomp_name not in actual_block._modules: 345 ↛ 346line 345 didn't jump to line 346 because the condition on line 345 was never true
346 continue
347 comp_path = f"{block_type}.{layer_idx}.{subcomp_name}"
348 self._test_component_recursive(
349 comp_path, subcomponent, test_inputs, results, skip_components
350 )
352 # Clean up test inputs to free memory
353 if test_inputs is not None: 353 ↛ 361line 353 didn't jump to line 361 because the condition on line 353 was always true
354 for key in list(test_inputs.keys()):
355 tensor = test_inputs[key]
356 if tensor is not None and isinstance(tensor, torch.Tensor): 356 ↛ 354line 356 didn't jump to line 354 because the condition on line 356 was always true
357 del tensor
358 test_inputs.clear()
360 # Create report
361 passed = sum(1 for r in results if r.passed)
362 failed = sum(1 for r in results if not r.passed)
364 report = BenchmarkReport(
365 model_name=getattr(self.cfg, "model_name", "unknown"),
366 total_components=len(results),
367 passed_components=passed,
368 failed_components=failed,
369 component_results=results,
370 )
372 # Restore bridge to its original dtype if we upcast it
373 if self._bridge_was_upcast and self._bridge_original_dtype is not None: 373 ↛ 374line 373 didn't jump to line 374 because the condition on line 373 was never true
374 self.bridge_model.to(self._bridge_original_dtype)
376 return report
378 def _test_component_recursive(
379 self,
380 component_path: str,
381 component: GeneralizedComponent,
382 test_inputs: Dict[str, torch.Tensor],
383 results: List[ComponentTestResult],
384 skip_components: Optional[List[str]] = None,
385 ) -> None:
386 """Recursively test a component and all its subcomponents.
388 This method tests the given component and then recursively tests all its
389 nested subcomponents (e.g., attn.q, attn.k, mlp.gate, etc.).
391 Note: We skip testing q/k/v subcomponents when the parent attention module
392 uses joint QKV projection (JointQKVAttentionBridge), as these are virtual
393 components that don't exist as separate modules in HuggingFace.
395 Args:
396 component_path: Path to the component (e.g., "blocks.0.attn")
397 component: The generalized component bridge
398 test_inputs: Dictionary of test inputs
399 results: List to append results to
400 skip_components: Optional list of component paths to skip
401 """
402 skip_components = skip_components or []
404 # Skip if in skip list
405 if component_path in skip_components: 405 ↛ 406line 405 didn't jump to line 406 because the condition on line 405 was never true
406 return
408 # Adapter-declared suffix exclusions (isolated forward can't run on
409 # synthesized probes, e.g. fused top-k routers).
410 skip_suffixes = tuple(getattr(self.adapter, "component_test_skip_suffixes", ()))
411 if any(component_path.endswith(suffix) for suffix in skip_suffixes): 411 ↛ 412line 411 didn't jump to line 412 because the condition on line 411 was never true
412 return
414 # SSM/recurrent mixer internal submodules can't be tested in isolation (see
415 # _is_ssm_mixer_internal); the mixer node itself is still tested against HF.
416 if _is_ssm_mixer_internal(component_path): 416 ↛ 417line 416 didn't jump to line 417 because the condition on line 416 was never true
417 return
419 # Skip MLP components that don't exist as separate modules in HF (name=None)
420 # These are virtual components where fc1/fc2 are directly on the layer
421 # Component testing doesn't work for these because get_component returns the parent layer
422 if "mlp" in component_path and hasattr(component, "name") and component.name is None:
423 return
425 # Skip MLPs whose forward needs more than a hidden_states probe (BLOOM
426 # residual; transformers 5.x fused MoE experts need router indices/weights);
427 # full-model forward_pass_logits already covers their parity. Fetch the
428 # layer-bound component since the per-block template's original_component is None.
429 if "mlp" in component_path:
430 import inspect
432 try:
433 bound = self.adapter.get_component(self.bridge_model, component_path)
434 inner = getattr(bound, "original_component", None)
435 if inner is not None: 435 ↛ 453line 435 didn't jump to line 453 because the condition on line 435 was always true
436 required = [
437 p
438 for p in inspect.signature(inner.forward).parameters.values()
439 if p.default is inspect.Parameter.empty
440 and p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)
441 and p.name != "self"
442 ]
443 # An isolatable MLP needs only hidden_states; more required
444 # positional args means it can't run standalone.
445 if len(required) > 1:
446 return
447 except (AttributeError, ValueError, TypeError):
448 # Can't resolve/inspect — fall through and let the test run.
449 pass
451 # Skip attention components that require position embeddings in Phase 3
452 # These can't be tested in isolation without full model context for position embeddings
453 if (
454 "attn" in component_path
455 and hasattr(component, "requires_position_embeddings")
456 and component.requires_position_embeddings
457 ):
458 return
460 # Skip attention components that use native HF attention (maintain_native_attention=True)
461 # These have custom forward signatures (e.g., BLOOM requires residual, alibi, attention_mask)
462 # and can't be tested in isolation without full model context
463 if ( 463 ↛ 468line 463 didn't jump to line 468 because the condition on line 463 was never true
464 "attn" in component_path
465 and hasattr(component, "maintain_native_attention")
466 and component.maintain_native_attention
467 ):
468 return
470 # Skip attention and PLE submodules when using DelegatedAttentionBlockBridge.
471 # These architectures delegate all math to HF; the benchmark can't call the HF
472 # attention in isolation (missing position_embeddings, attention_mask, etc.) and
473 # PLE submodules receive per-layer inputs at a different dimension than hidden_states.
474 _is_delegated = self._is_delegated_block()
475 if _is_delegated and "attn" in component_path: 475 ↛ 476line 475 didn't jump to line 476 because the condition on line 475 was never true
476 return
477 if _is_delegated and any( 477 ↛ 485line 477 didn't jump to line 485 because the condition on line 477 was never true
478 name in component_path
479 for name in (
480 "per_layer_input_gate",
481 "per_layer_projection",
482 "post_per_layer_input_norm",
483 )
484 ):
485 return
487 # Skip models whose MLP/attn forward signatures require extra context from the block:
488 # - BLOOM: MLP requires residual and alibi bias
489 # - T5: requires cache_position for relative position embeddings
490 # - MPT: MLP.forward(hidden_states, residual) performs the residual addition internally
491 if "attn" in component_path or "mlp" in component_path:
492 hf_model_config = getattr(self.hf_model, "config", None)
493 if hf_model_config and hasattr(hf_model_config, "model_type"): 493 ↛ 500line 493 didn't jump to line 500 because the condition on line 493 was always true
494 if hf_model_config.model_type in ["bloom", "t5", "mpt"]: 494 ↛ 495line 494 didn't jump to line 495 because the condition on line 494 was never true
495 return
497 # Skip components that require specific shaped inputs from their parent modules
498 # These components expect intermediate outputs from their parent attention/MLP
499 # modules and can't be tested with generic hidden state inputs
500 path_parts = component_path.split(".")
501 if len(path_parts) >= 3: # e.g., "blocks.0.attn.o" or "blocks.0.mlp.out" 501 ↛ 544line 501 didn't jump to line 544 because the condition on line 501 was always true
502 last_part = path_parts[-1]
504 # Skip attention output projection (expects concatenated attn output)
505 # Skip MLP output projection (expects MLP intermediate activations)
506 # — including *_out variants (dense_out on MoE shared/dense paths).
507 # Note: q_norm/k_norm are handled specially in _run_component
508 if last_part in ["o", "out"] or last_part.endswith("_out"):
509 return
511 # Skip MLA intermediates (expect compressed-dim inputs, not hidden_states)
512 if last_part in [ 512 ↛ 520line 512 didn't jump to line 520 because the condition on line 512 was never true
513 "q_a_proj",
514 "q_a_layernorm",
515 "q_b_proj",
516 "kv_a_proj_with_mqa",
517 "kv_a_layernorm",
518 "kv_b_proj",
519 ]:
520 return
522 # Skip virtual splits from fused projections (no standalone HF equivalent)
523 if last_part in ["q", "k", "v", "gate", "in"]:
524 parent_path = ".".join(path_parts[:-1])
525 try:
526 parent_component = self.adapter.get_component(self.bridge_model, parent_path)
527 if hasattr(parent_component, "submodules"): 527 ↛ 544line 527 didn't jump to line 544 because the condition on line 527 was always true
528 parent_bridge = cast(GeneralizedComponent, parent_component)
529 subs = parent_bridge.submodules
530 # Joint QKV: q/k/v are splits from fused qkv_proj/c_attn
531 if last_part in ["q", "k", "v"] and ("qkv" in subs or "c_attn" in subs):
532 return
533 # Joint gate+up: gate/in are splits from fused gate_up_proj
534 if last_part in ["gate", "in"] and ( 534 ↛ 538line 534 didn't jump to line 538 because the condition on line 534 was never true
535 "gate_up" in subs
536 or type(parent_bridge).__name__ == "JointGateUpMLPBridge"
537 ):
538 return
539 except Exception:
540 pass
542 # Skip components not wired on this layer (per-layer or per-config variation).
543 # Only report as failure if the HF model has it but the bridge doesn't.
544 try:
545 self.adapter.get_component(self.bridge_model, component_path)
546 except (AttributeError, ValueError):
547 parts = component_path.split(".")
548 if len(parts) >= 3 and parts[1].isdigit(): 548 ↛ 566line 548 didn't jump to line 566 because the condition on line 548 was always true
549 subpath = ".".join([parts[0]] + ["{layer}"] + parts[2:])
550 # Per-layer variation: exists on some other layer (e.g., MoE vs dense)
551 for probe_layer in range(self.cfg.n_layers):
552 probe_path = subpath.replace("{layer}", str(probe_layer))
553 try:
554 self.adapter.get_component(self.bridge_model, probe_path)
555 return # Found on another layer — skip this one
556 except (AttributeError, ValueError):
557 continue
558 # Per-config absence: HF model also lacks it (e.g., q_lora_rank=None)
559 try:
560 self.adapter.get_component(self.hf_model, component_path)
561 except (AttributeError, ValueError):
562 return
563 # Bridge is missing a component that HF has — likely misconfiguration
565 # Test this component
566 result = self._test_component(component_path, component, test_inputs)
567 if result is not None: 567 ↛ 571line 567 didn't jump to line 571 because the condition on line 567 was always true
568 results.append(result)
570 # Recursively test subcomponents
571 if hasattr(component, "submodules") and component.submodules:
572 for subcomp_name, subcomponent in component.submodules.items():
573 sub_path = f"{component_path}.{subcomp_name}"
574 self._test_component_recursive(
575 sub_path, subcomponent, test_inputs, results, skip_components
576 )
578 def _test_component(
579 self,
580 component_path: str,
581 component: GeneralizedComponent,
582 test_inputs: Dict[str, torch.Tensor],
583 ) -> Optional[ComponentTestResult]:
584 """Test a single component.
586 Args:
587 component_path: Path to the component (e.g., "embed", "blocks.0.attn")
588 component: The generalized component bridge
589 test_inputs: Dictionary of test inputs
591 Returns:
592 ComponentTestResult or None if the component cannot be tested
593 """
594 try:
595 # Skip rotary_emb for DelegatedAttentionBlockBridge architectures.
596 # Gemma4's RotaryEmbeddingBridge wraps a rotary that returns a set-like
597 # structure which the benchmark comparison can't subscript.
598 if self._is_delegated_block() and component_path == "rotary_emb": 598 ↛ 599line 598 didn't jump to line 599 because the condition on line 598 was never true
599 return None
601 # Get bridge component
602 # The adapter returns nn.Module, but for bridge models it's actually GeneralizedComponent
603 bridge_component = cast(
604 GeneralizedComponent, self.adapter.get_component(self.bridge_model, component_path)
605 )
607 # Get HuggingFace component
608 hf_component = self.adapter.get_component(self.hf_model, component_path)
610 # Determine appropriate test input based on component type
611 test_input = self._get_test_input_for_component(component_path, test_inputs)
612 if test_input is None: 612 ↛ 613line 612 didn't jump to line 613 because the condition on line 612 was never true
613 return None
615 # Get input args/kwargs from the Bridge component
616 # All bridge components inherit from GeneralizedComponent and have get_dummy_inputs()
617 batch, seq_len, _ = test_input.shape
618 pos_indices = (
619 torch.arange(seq_len, device=test_input.device).unsqueeze(0).expand(batch, -1)
620 )
622 # For embedding components, generate the embedding input once
623 shared_embed_input = None
624 if component_path in ("embed", "encoder_embed", "decoder_embed"):
625 shared_embed_input = test_inputs.get("modality_input")
626 if shared_embed_input is None: 626 ↛ 634line 626 didn't jump to line 634 because the condition on line 626 was always true
627 batch, seq_len, _ = test_input.shape
628 shared_embed_input = torch.randint(
629 0, self.cfg.d_vocab, (batch, seq_len), device=test_input.device
630 )
632 # Generate shared inputs for attention/MLP/rotary components that have get_random_inputs()
633 # This is needed for model-specific inputs like position_embeddings or attention_mask
634 shared_inputs = None
635 if (
636 (
637 "attn" in component_path
638 or "mlp" in component_path
639 or "rotary" in component_path
640 or "conv" in component_path
641 )
642 and hasattr(bridge_component, "get_random_inputs")
643 and callable(getattr(bridge_component, "get_random_inputs"))
644 ):
645 batch_size, seq_len = test_input.shape[:2]
646 # Cast to callable to satisfy mypy - we've already verified it exists and is callable
647 get_random_inputs_fn = cast(
648 Callable[..., Dict[str, Any]], bridge_component.get_random_inputs
649 )
650 shared_inputs = get_random_inputs_fn(
651 batch_size=batch_size,
652 seq_len=seq_len,
653 device=test_input.device,
654 dtype=test_input.dtype,
655 )
656 if "attn" in component_path:
657 self._add_direct_attention_mask_if_needed(
658 shared_inputs, hf_component, batch_size, seq_len
659 )
661 # Override position_embeddings with correct values from HF model's rotary_emb
662 # This is needed for models with partial RoPE or non-standard rotary dims
663 if ( 663 ↛ 668line 663 didn't jump to line 668 because the condition on line 663 was never true
664 "attn" in component_path
665 and "position_embeddings" in shared_inputs
666 and hasattr(self.hf_model, "model")
667 ):
668 rotary_attr = getattr(self.hf_model.model, "rotary_emb", None)
669 if callable(rotary_attr):
670 try:
671 position_ids = (
672 torch.arange(seq_len, device=test_input.device)
673 .unsqueeze(0)
674 .expand(batch_size, -1)
675 )
676 position_embeddings = rotary_attr(test_input, position_ids)
677 shared_inputs["position_embeddings"] = position_embeddings
678 except Exception:
679 # If rotary_emb fails, keep the fallback position_embeddings from get_random_inputs()
680 pass
682 # Run through both components with shared inputs (for attention) or standard inputs (for others)
683 bridge_output = self._run_component(
684 bridge_component, test_input, component_path, shared_embed_input, shared_inputs
685 )
686 hf_output = self._run_component(
687 hf_component, test_input, component_path, shared_embed_input, shared_inputs
688 )
690 # Extract tensors if outputs are tuples
691 # Legacy modules (e.g. GPT-1's Attention) return lists, not tuples.
692 bridge_tensor = (
693 bridge_output[0] if isinstance(bridge_output, (tuple, list)) else bridge_output
694 )
695 hf_tensor = hf_output[0] if isinstance(hf_output, (tuple, list)) else hf_output
697 # Marian/Bart apply a trained final_logits_bias AFTER lm_head inside the
698 # model forward. The bridge folds it into the unembed bias (b_U), so the
699 # isolated HF lm_head must add it too for a fair comparison — otherwise
700 # the two diverge by exactly the bias (a false failure; the assembled
701 # unembed+bias path is already covered by forward_pass_logits).
702 if "unembed" in component_path or "lm_head" in component_path:
703 flb = getattr(self.hf_model, "final_logits_bias", None)
704 if (
705 flb is not None
706 and isinstance(hf_tensor, torch.Tensor)
707 and hf_tensor.shape[-1] == flb.shape[-1]
708 ):
709 hf_tensor = hf_tensor + flb.to(hf_tensor.dtype)
711 # Ensure both are tensors
712 if not isinstance(bridge_tensor, torch.Tensor) or not isinstance( 712 ↛ 715line 712 didn't jump to line 715 because the condition on line 712 was never true
713 hf_tensor, torch.Tensor
714 ):
715 return ComponentTestResult(
716 component_path=component_path,
717 component_type=type(component).__name__,
718 passed=False,
719 max_diff=float("inf"),
720 mean_diff=float("inf"),
721 output_shape=(),
722 error_message=f"Outputs are not tensors: bridge={type(bridge_tensor)}, hf={type(hf_tensor)}",
723 )
725 # Compare outputs
726 passed, max_diff, mean_diff, percentile_diffs = self._compare_outputs(
727 bridge_tensor, hf_tensor
728 )
730 # Get output shape before deleting tensors
731 output_shape = tuple(bridge_tensor.shape)
733 # Clean up output tensors immediately to free memory
734 del bridge_output, hf_output, bridge_tensor, hf_tensor
735 if shared_inputs is not None:
736 # Clean up shared inputs
737 for key in list(shared_inputs.keys()):
738 val = shared_inputs[key]
739 if val is not None and isinstance(val, torch.Tensor):
740 del val
741 shared_inputs[key] = None
742 if shared_embed_input is not None:
743 del shared_embed_input
745 return ComponentTestResult(
746 component_path=component_path,
747 component_type=type(component).__name__,
748 passed=passed,
749 max_diff=max_diff,
750 mean_diff=mean_diff,
751 output_shape=output_shape,
752 percentile_diffs=percentile_diffs,
753 )
755 except Exception as e:
756 return ComponentTestResult(
757 component_path=component_path,
758 component_type=type(component).__name__,
759 passed=False,
760 max_diff=float("inf"),
761 mean_diff=float("inf"),
762 output_shape=(),
763 error_message=str(e),
764 )
766 @staticmethod
767 def _add_direct_attention_mask_if_needed(
768 shared_inputs: Dict[str, Any],
769 hf_component: Any,
770 batch_size: int,
771 seq_len: int,
772 ) -> None:
773 """Add a causal mask for direct HF attention calls that need parent context."""
774 if "attention_mask" in shared_inputs: 774 ↛ 775line 774 didn't jump to line 775 because the condition on line 774 was never true
775 return
776 hidden_states = shared_inputs.get("hidden_states")
777 if not isinstance(hidden_states, torch.Tensor): 777 ↛ 778line 777 didn't jump to line 778 because the condition on line 777 was never true
778 return
779 if not getattr(hf_component, "is_causal", False):
780 return
781 if getattr(hf_component, "is_cross_attention", False): 781 ↛ 782line 781 didn't jump to line 782 because the condition on line 781 was never true
782 return
784 min_dtype = torch.finfo(hidden_states.dtype).min
785 causal_mask = torch.ones(seq_len, seq_len, device=hidden_states.device, dtype=torch.bool)
786 causal_mask = torch.tril(causal_mask).view(1, 1, seq_len, seq_len)
787 attention_mask = torch.zeros(
788 batch_size,
789 1,
790 seq_len,
791 seq_len,
792 device=hidden_states.device,
793 dtype=hidden_states.dtype,
794 )
795 shared_inputs["attention_mask"] = attention_mask.masked_fill(~causal_mask, min_dtype)
797 def _run_component(
798 self,
799 component: nn.Module,
800 test_input: torch.Tensor,
801 component_path: str,
802 shared_embed_input: Optional[torch.Tensor] = None,
803 shared_inputs: Optional[dict] = None,
804 ) -> Any:
805 """Run a component with appropriate arguments.
807 Args:
808 component: The component to run
809 test_input: The test input tensor
810 component_path: Path to the component for debugging
811 shared_embed_input: Pre-generated input for embedding components (token ids,
812 or a spectrogram/pixel tensor for audio and vision models)
813 shared_inputs: Pre-generated inputs from get_random_inputs() to use for both bridge and HF components
815 Returns:
816 The component output
817 """
818 # q_norm/k_norm expect d_head, not d_model
819 if component_path.endswith(".q_norm") or component_path.endswith(".k_norm"): 819 ↛ 821line 819 didn't jump to line 821 because the condition on line 819 was never true
820 # Reshape test_input from (batch, seq, d_model) to (batch, seq, d_head)
821 batch, seq, d_model = test_input.shape
822 d_head = self.cfg.d_head
823 # Use just d_head dimensions as test input
824 test_input_reshaped = test_input[..., :d_head]
825 return component(test_input_reshaped)
827 # Use shared inputs if provided (generated from bridge component's get_random_inputs())
828 if shared_inputs is not None:
829 # Check if shared_inputs contains positional args
830 if "args" in shared_inputs:
831 # Call with positional args (e.g., for rotary embeddings)
832 return component(*shared_inputs["args"])
833 else:
834 # Call with keyword args (e.g., for attention)
835 try:
836 return component(**shared_inputs)
837 except TypeError:
838 # Pre-kwarg-era modules (e.g. GPT-1's Attention) take the
839 # input positionally and reject the hidden_states keyword.
840 hidden = shared_inputs.get("hidden_states")
841 if hidden is None:
842 raise
843 mask = shared_inputs.get("attention_mask")
844 if mask is not None:
845 return component(hidden, attention_mask=mask)
846 return component(hidden)
848 # Fallback: Use legacy calling conventions for components without get_random_inputs()
849 if "attn" in component_path and "attn" == component_path.split(".")[-1]: 849 ↛ 851line 849 didn't jump to line 851 because the condition on line 849 was never true
850 # Attention components (legacy fallback)
851 try:
852 # Try TransformerLens-style attention
853 return component(
854 query_input=test_input,
855 key_input=test_input,
856 value_input=test_input,
857 past_kv_cache_entry=None,
858 attention_mask=None,
859 )
860 except TypeError:
861 try:
862 # Try HuggingFace-style attention
863 return component(hidden_states=test_input)
864 except TypeError:
865 # Try simple call
866 return component(test_input)
867 elif component_path in ("embed", "encoder_embed", "decoder_embed"):
868 # Token ids for text models; a spectrogram or pixel tensor for audio/vision.
869 if shared_embed_input is not None: 869 ↛ 872line 869 didn't jump to line 872 because the condition on line 869 was always true
870 embed_input = shared_embed_input
871 else:
872 batch, seq_len, _ = test_input.shape
873 embed_input = torch.randint(
874 0, self.cfg.d_vocab, (batch, seq_len), device=test_input.device
875 )
876 return component(embed_input)
877 elif component_path == "pos_embed" or "pos_embed" in component_path:
878 # Position embedding expects integer position indices
879 batch, seq_len, _ = test_input.shape
880 # For positional embeddings, we need position indices
881 pos_indices = (
882 torch.arange(seq_len, device=test_input.device).unsqueeze(0).expand(batch, -1)
883 )
884 try:
885 return component(pos_indices)
886 except (TypeError, IndexError):
887 # Some pos embeds just return their embeddings directly
888 # or may not take inputs
889 try:
890 if hasattr(component, "weight") and isinstance(component.weight, torch.Tensor): 890 ↛ 893line 890 didn't jump to line 893 because the condition on line 890 was always true
891 return component.weight[:seq_len]
892 else:
893 raise AttributeError("Component has no weight attribute")
894 except AttributeError:
895 # Skip this component
896 raise ValueError("Cannot test pos_embed - unclear interface")
897 elif isinstance(getattr(component, "original_component", component), torch.nn.Embedding):
898 # Any other embedding table (BERT's token_type_embed) rejects the
899 # float default: embeddings index with integer ids. The HF side is
900 # the bare nn.Embedding, the bridge side wraps one. Ids are derived
901 # from test_input so both sides index identically.
902 embedding_table = getattr(component, "original_component", component)
903 assert isinstance(embedding_table, torch.nn.Embedding) # narrowed by the elif
904 num_ids = int(embedding_table.num_embeddings)
905 id_input = (test_input.abs().sum(dim=-1) * 1e3).long() % num_ids
906 return component(id_input)
907 elif component_path == "project_in": 907 ↛ 909line 907 didn't jump to line 909 because the condition on line 907 was never true
908 # project_in expects word_embed_proj_dim, not d_model.
909 word_embed_proj_dim = getattr(self.cfg, "word_embed_proj_dim", None)
910 if word_embed_proj_dim is not None and word_embed_proj_dim != self.cfg.d_model:
911 test_input = test_input[..., :word_embed_proj_dim]
912 return component(test_input)
913 elif (
914 component_path == "unembed"
915 or "unembed" in component_path
916 or "lm_head" in component_path
917 ):
918 # Unembed may expect word_embed_proj_dim (e.g., OPT-350m project_out).
919 word_embed_proj_dim = getattr(self.cfg, "word_embed_proj_dim", None)
920 if ( 920 ↛ 925line 920 didn't jump to line 925 because the condition on line 920 was never true
921 word_embed_proj_dim is not None
922 and word_embed_proj_dim != self.cfg.d_model
923 and test_input.shape[-1] != word_embed_proj_dim
924 ):
925 test_input = test_input[..., :word_embed_proj_dim]
926 return component(test_input)
927 else:
928 # Standard components (MLP, LayerNorm, etc.)
929 try:
930 return component(test_input)
931 except TypeError:
932 # Try with hidden_states kwarg
933 return component(hidden_states=test_input)
935 def _get_test_input_for_component(
936 self, component_path: str, test_inputs: Dict[str, torch.Tensor]
937 ) -> Optional[torch.Tensor]:
938 """Get the appropriate test input for a component.
940 Args:
941 component_path: Path to the component
942 test_inputs: Dictionary of available test inputs
944 Returns:
945 The appropriate test input tensor, or None if not applicable
946 """
947 # Use standard hidden state input for most components
948 return test_inputs.get("hidden_states")
950 def _generate_test_inputs(self) -> Dict[str, torch.Tensor]:
951 """Generate default test inputs for benchmarking.
953 Returns:
954 Dictionary of test input tensors
955 """
956 batch_size = 2
957 seq_len = 8
958 d_model = self.cfg.d_model
960 # Use the reconciled dtype from __init__.
961 dtype = self.test_dtype
962 try:
963 device = next(self.hf_model.parameters()).device
964 except StopIteration:
965 device = torch.device("cpu")
967 inputs = {
968 "hidden_states": torch.randn(batch_size, seq_len, d_model, dtype=dtype, device=device),
969 }
971 # Vision/audio models have no token vocabulary — d_vocab stays at its -1 sentinel.
972 if self.cfg.d_vocab > 0: 972 ↛ 978line 972 didn't jump to line 978 because the condition on line 972 was always true
973 inputs["token_ids"] = torch.randint(
974 0, self.cfg.d_vocab, (batch_size, seq_len), device=device
975 )
977 # Audio/vision embeddings consume a raw spectrogram or pixel tensor, not token ids.
978 modality_input = build_modality_input(
979 self.bridge_model, batch_size=batch_size, device=device, dtype=dtype
980 )
981 if modality_input is not None: 981 ↛ 982line 981 didn't jump to line 982 because the condition on line 981 was never true
982 inputs["modality_input"] = modality_input
984 return inputs
986 def _compare_outputs(
987 self, bridge_output: torch.Tensor, hf_output: torch.Tensor
988 ) -> Tuple[bool, float, float, Dict[str, float]]:
989 """Compare two output tensors.
991 Args:
992 bridge_output: Output from TransformerBridge component
993 hf_output: Output from HuggingFace component
995 Returns:
996 Tuple of (passed, max_diff, mean_diff, percentile_diffs)
997 """
998 # Check shapes match
999 if bridge_output.shape != hf_output.shape: 999 ↛ 1000line 999 didn't jump to line 1000 because the condition on line 999 was never true
1000 return False, float("inf"), float("inf"), {}
1002 # Compute differences (upcast to float32 for safety)
1003 bo = bridge_output.float()
1004 ho = hf_output.float()
1005 diff = torch.abs(bo - ho)
1006 max_diff = diff.max().item()
1007 mean_diff = diff.mean().item()
1009 # Compute percentile differences
1010 flat_diff = diff.flatten()
1011 percentile_diffs = {
1012 "50th": torch.quantile(flat_diff, 0.5).item(),
1013 "90th": torch.quantile(flat_diff, 0.9).item(),
1014 "99th": torch.quantile(flat_diff, 0.99).item(),
1015 }
1017 # Check if within tolerance
1018 passed = torch.allclose(bo, ho, atol=self.atol, rtol=self.rtol)
1020 return passed, max_diff, mean_diff, percentile_diffs
1023def benchmark_model(
1024 model_name: str,
1025 device: str = "cpu",
1026 atol: float = 1e-4,
1027 rtol: float = 1e-4,
1028 skip_components: Optional[List[str]] = None,
1029 verbose: bool = False,
1030) -> BenchmarkReport:
1031 """Benchmark all components in a model.
1033 Args:
1034 model_name: Name of the HuggingFace model to benchmark
1035 device: Device to run on
1036 atol: Absolute tolerance for comparisons
1037 rtol: Relative tolerance for comparisons
1038 skip_components: Optional list of component paths to skip
1039 verbose: If True, print detailed results for all components
1041 Returns:
1042 BenchmarkReport with results for all components
1043 """
1044 from transformers import AutoModelForCausalLM
1046 from transformer_lens.model_bridge import TransformerBridge
1048 # Load models
1049 print(f"Loading models: {model_name}")
1050 bridge_model = TransformerBridge.boot_transformers(model_name, device=device) # type: ignore[attr-defined]
1052 # Load HF model with same attn_implementation as bridge model (if specified)
1053 # This ensures numerical consistency between bridge and HF models
1054 hf_kwargs = {"device_map": device}
1055 if (
1056 hasattr(bridge_model.adapter.cfg, "attn_implementation")
1057 and bridge_model.adapter.cfg.attn_implementation is not None
1058 ):
1059 hf_kwargs["attn_implementation"] = bridge_model.adapter.cfg.attn_implementation
1061 hf_model = AutoModelForCausalLM.from_pretrained(model_name, **hf_kwargs)
1063 # Set models to eval mode (disable dropout, etc.)
1064 bridge_model.eval()
1065 hf_model.eval()
1067 # Get adapter
1068 adapter = bridge_model.adapter
1070 # Set up component testing (e.g., sync rotary_emb references for Gemma-3)
1071 # Pass bridge_model so adapter can set up actual bridge instances, not just templates
1072 adapter.setup_component_testing(hf_model, bridge_model=bridge_model)
1074 # Create benchmarker
1075 benchmarker = ComponentBenchmarker(
1076 bridge_model=bridge_model,
1077 hf_model=hf_model,
1078 adapter=adapter,
1079 cfg=bridge_model.cfg,
1080 atol=atol,
1081 rtol=rtol,
1082 )
1084 # Run benchmark
1085 print("Running component benchmark...")
1086 report = benchmarker.benchmark_all_components(skip_components=skip_components)
1088 # Print report
1089 report.print_summary(verbose=verbose)
1091 return report