Coverage for transformer_lens/benchmarks/encoder_common.py: 55%
87 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +0000
1"""Shared benchmark implementations for non-text encoder bridges.
3Audio and vision encoders take one raw tensor (waveform, spectrogram, or pixels)
4where text models take token ids; given that tensor, the forward, run_with_cache,
5and perturbation-stability checks are modality-independent. The modality modules
6(audio.py, vision.py) wrap these with their result names, HF reference-forward
7kwarg, and critical component lists.
8"""
10from typing import Optional, Sequence
12import torch
14from transformer_lens.benchmarks.utils import (
15 BenchmarkResult,
16 BenchmarkSeverity,
17 compare_tensors,
18 is_tiny_test_model,
19)
20from transformer_lens.model_bridge import TransformerBridge
23def extract_encoder_states(out) -> Optional[torch.Tensor]:
24 """Pull the hidden-state/logit tensor out of a tensor, BaseModelOutput, or head output."""
25 if isinstance(out, torch.Tensor):
26 return out
27 if hasattr(out, "last_hidden_state"):
28 return out.last_hidden_state
29 if hasattr(out, "logits") and out.logits is not None:
30 return out.logits
31 return None
34def benchmark_encoder_forward(
35 bridge: TransformerBridge,
36 test_input: torch.Tensor,
37 name: str,
38 ref_input_key: str,
39 reference_model: Optional[torch.nn.Module] = None,
40) -> BenchmarkResult:
41 """Forward-pass benchmark on a raw modality tensor.
43 Compares bridge output against the HF native model on the same input when a
44 reference is given. Bare encoders compare last_hidden_state; head models
45 (CTC, classification) compare logits.
47 Args:
48 bridge: TransformerBridge model to test
49 test_input: Raw modality tensor (waveform, spectrogram, or pixels)
50 name: Result name (e.g. "audio_forward", "vision_forward")
51 ref_input_key: Kwarg the HF reference takes the tensor under
52 ("input_values" for audio, "pixel_values" for vision)
53 reference_model: Optional HF reference model for comparison
54 """
55 try:
56 with torch.no_grad():
57 # Use return_type="logits" — for bare encoders without logits, this
58 # returns the BaseModelOutput object (bridge falls through to logits=output).
59 bridge_output_raw = bridge(test_input, return_type="logits")
61 # Extract the output tensor
62 if isinstance(bridge_output_raw, torch.Tensor): 62 ↛ 65line 62 didn't jump to line 65 because the condition on line 62 was always true
63 bridge_output = bridge_output_raw
64 output_key = "logits"
65 elif hasattr(bridge_output_raw, "logits") and bridge_output_raw.logits is not None:
66 bridge_output = bridge_output_raw.logits
67 output_key = "logits"
68 elif hasattr(bridge_output_raw, "last_hidden_state"):
69 bridge_output = bridge_output_raw.last_hidden_state
70 output_key = "last_hidden_state"
71 else:
72 return BenchmarkResult(
73 name=name,
74 severity=BenchmarkSeverity.DANGER,
75 message="Bridge produced no recognizable output (no logits or last_hidden_state)",
76 passed=False,
77 )
79 if bridge_output.numel() == 0: 79 ↛ 80line 79 didn't jump to line 80 because the condition on line 79 was never true
80 return BenchmarkResult(
81 name=name,
82 severity=BenchmarkSeverity.DANGER,
83 message="Bridge output is empty",
84 passed=False,
85 )
87 if torch.isnan(bridge_output).any() or torch.isinf(bridge_output).any():
88 return BenchmarkResult(
89 name=name,
90 severity=BenchmarkSeverity.DANGER,
91 message="Bridge output contains NaN or Inf values",
92 passed=False,
93 )
95 # Compare against HF reference if available
96 if reference_model is not None:
97 with torch.no_grad():
98 ref_output_raw = reference_model(**{ref_input_key: test_input})
99 if output_key == "logits": 99 ↛ 102line 99 didn't jump to line 102 because the condition on line 99 was always true
100 ref_output = ref_output_raw.logits
101 else:
102 ref_output = ref_output_raw.last_hidden_state
104 return compare_tensors(
105 bridge_output,
106 ref_output,
107 atol=1e-3,
108 rtol=3e-2,
109 name=name,
110 )
112 return BenchmarkResult(
113 name=name,
114 severity=BenchmarkSeverity.INFO,
115 message=f"Forward pass successful ({output_key} shape: {bridge_output.shape})",
116 details={"output_shape": str(bridge_output.shape), "output_key": output_key},
117 )
119 except Exception as e:
120 return BenchmarkResult(
121 name=name,
122 severity=BenchmarkSeverity.ERROR,
123 message=f"Forward pass failed: {str(e)}",
124 passed=False,
125 )
128def benchmark_encoder_cache(
129 bridge: TransformerBridge,
130 test_input: torch.Tensor,
131 name: str,
132 critical_components: Sequence[str],
133 min_found: int = 3,
134) -> BenchmarkResult:
135 """run_with_cache() benchmark on a raw modality tensor.
137 Verifies that critical hooks fire and produce valid tensors: the
138 ``critical_components`` this architecture actually declares in its
139 component mapping, plus the first and last block.
141 Args:
142 bridge: TransformerBridge model to test
143 test_input: Raw modality tensor (waveform, spectrogram, or pixels)
144 name: Result name (e.g. "audio_cache", "vision_cache")
145 critical_components: Component-mapping names whose hook_out must be
146 cached when the architecture declares them
147 min_found: Minimum critical hooks present to still pass with a warning
148 """
149 try:
150 with torch.no_grad():
151 _, cache = bridge.run_with_cache(test_input)
153 cache_keys = list(cache.keys())
154 if len(cache_keys) == 0: 154 ↛ 155line 154 didn't jump to line 155 because the condition on line 154 was never true
155 return BenchmarkResult(
156 name=name,
157 severity=BenchmarkSeverity.DANGER,
158 message="run_with_cache returned empty cache",
159 passed=False,
160 )
162 component_mapping = bridge.adapter.component_mapping or {}
163 critical_hooks = [
164 f"{comp}.hook_out" for comp in critical_components if comp in component_mapping
165 ]
166 # Also check at least the first and last block
167 n_layers = bridge.cfg.n_layers
168 critical_hooks.append("blocks.0.hook_out")
169 critical_hooks.append(f"blocks.{n_layers - 1}.hook_out")
171 missing = [h for h in critical_hooks if h not in cache_keys]
172 found = len(critical_hooks) - len(missing)
174 # Check for NaN/Inf in cached values
175 nan_hooks = []
176 for key in cache_keys[:20]: # Sample first 20 hooks
177 val = cache[key]
178 if isinstance(val, torch.Tensor) and (torch.isnan(val).any() or torch.isinf(val).any()): 178 ↛ 179line 178 didn't jump to line 179 because the condition on line 178 was never true
179 nan_hooks.append(key)
181 if missing:
182 return BenchmarkResult(
183 name=name,
184 severity=BenchmarkSeverity.WARNING,
185 message=f"Missing {len(missing)} critical hooks: {missing[:3]}",
186 passed=found >= min_found,
187 details={
188 "total_cached": len(cache_keys),
189 "critical_found": found,
190 "critical_expected": len(critical_hooks),
191 "missing": missing,
192 },
193 )
195 if nan_hooks: 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true
196 return BenchmarkResult(
197 name=name,
198 severity=BenchmarkSeverity.DANGER,
199 message=f"NaN/Inf found in {len(nan_hooks)} cached hooks",
200 passed=False,
201 details={"nan_hooks": nan_hooks[:5]},
202 )
204 return BenchmarkResult(
205 name=name,
206 severity=BenchmarkSeverity.INFO,
207 message=f"Cache successful: {len(cache_keys)} hooks captured, "
208 f"{found}/{len(critical_hooks)} critical hooks present",
209 details={
210 "total_cached": len(cache_keys),
211 "critical_found": found,
212 "critical_expected": len(critical_hooks),
213 },
214 )
216 except Exception as e:
217 return BenchmarkResult(
218 name=name,
219 severity=BenchmarkSeverity.ERROR,
220 message=f"Cache benchmark failed: {str(e)}",
221 passed=False,
222 )
225def benchmark_encoder_representation_stability(
226 bridge: TransformerBridge,
227 test_input: torch.Tensor,
228 name: str,
229) -> BenchmarkResult:
230 """Representation stability under small input perturbations.
232 Similar inputs should produce similar hidden states. Skipped for
233 tiny-random models (random weights won't produce stable representations).
235 Args:
236 bridge: TransformerBridge model to test
237 test_input: Raw modality tensor (waveform, spectrogram, or pixels)
238 name: Result name (e.g. "audio_representation_stability")
239 """
240 model_name = getattr(bridge.cfg, "model_name", "")
241 if is_tiny_test_model(model_name): 241 ↛ 248line 241 didn't jump to line 248 because the condition on line 241 was always true
242 return BenchmarkResult(
243 name=name,
244 severity=BenchmarkSeverity.SKIPPED,
245 message="Skipped for tiny-random model (random weights won't produce stable representations)",
246 )
248 try:
249 # Create a slightly perturbed version
250 noise = torch.randn_like(test_input) * 0.01
251 perturbed_input = test_input + noise
253 with torch.no_grad():
254 output_orig = bridge(test_input, return_type="logits")
255 output_pert = bridge(perturbed_input, return_type="logits")
257 orig_states = extract_encoder_states(output_orig)
258 pert_states = extract_encoder_states(output_pert)
260 if orig_states is None or pert_states is None:
261 return BenchmarkResult(
262 name=name,
263 severity=BenchmarkSeverity.WARNING,
264 message="Could not extract hidden states for stability check",
265 passed=False,
266 )
268 # Compute cosine similarity (flatten to 2D: [batch, features])
269 orig_flat = orig_states.reshape(orig_states.shape[0], -1)
270 pert_flat = pert_states.reshape(pert_states.shape[0], -1)
271 cosine_sim = (
272 torch.nn.functional.cosine_similarity(orig_flat, pert_flat, dim=-1).mean().item()
273 )
275 passed = cosine_sim > 0.95
276 return BenchmarkResult(
277 name=name,
278 severity=BenchmarkSeverity.INFO if passed else BenchmarkSeverity.WARNING,
279 message=f"Representation stability: cosine_similarity={cosine_sim:.4f} "
280 f"(threshold: 0.95)",
281 passed=passed,
282 details={"cosine_similarity": cosine_sim, "noise_std": 0.01},
283 )
285 except Exception as e:
286 return BenchmarkResult(
287 name=name,
288 severity=BenchmarkSeverity.ERROR,
289 message=f"Representation stability check failed: {str(e)}",
290 passed=False,
291 )