Coverage for transformer_lens/benchmarks/main_benchmark.py: 32%
1040 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Main benchmark runner for TransformerBridge.
3This module provides the main benchmark suite that compares TransformerBridge
4against reference implementations in an optimized multi-phase approach:
5Phase 1: HF + Bridge (unprocessed) - Compare against raw HuggingFace model
6Phase 2: Bridge (unprocessed) + HT (unprocessed) - Compare unprocessed models
7Phase 3: Bridge (processed) + HT (processed) - Full compatibility mode testing
8Phase 4: Text Quality - profile prompts scored by a pinned judge's perplexity ratio
9Phase 5: Granular Weight Processing Tests (optional, individual flags)
10Phase 6: Granular Weight Processing Tests (optional, combined flags)
11Phase 7: Multimodal Tests (only for multimodal models with pixel_values support)
12Phase 8: Audio Tests (only for audio encoder models / audio-conditioned decoders)
13Phase 9: Vision Tests (only for vision-only encoder models, e.g. ViT/DeiT)
14"""
16import gc
17from typing import Dict, List, Optional, Union
19import torch
20from transformers import (
21 AutoConfig,
22 AutoModelForCausalLM,
23 PreTrainedModel,
24 PreTrainedTokenizerBase,
25)
27from transformer_lens import HookedTransformer
28from transformer_lens.benchmarks.activation_cache import (
29 benchmark_activation_cache,
30 benchmark_run_with_cache,
31)
32from transformer_lens.benchmarks.backward_gradients import (
33 benchmark_backward_hooks,
34 benchmark_critical_backward_hooks,
35 benchmark_gradient_computation,
36 needs_fp32_gradients,
37)
38from transformer_lens.benchmarks.component_benchmark import benchmark_all_components
39from transformer_lens.benchmarks.forward_pass import (
40 _compute_self_target_loss,
41 benchmark_forward_pass,
42 benchmark_logits_equivalence,
43 benchmark_loss_equivalence,
44)
45from transformer_lens.benchmarks.generation import (
46 benchmark_generation,
47 benchmark_generation_with_kv_cache,
48 benchmark_multiple_generation_calls,
49)
50from transformer_lens.benchmarks.hook_registration import (
51 benchmark_critical_forward_hooks,
52 benchmark_forward_hooks,
53 benchmark_gated_hooks_fire,
54 benchmark_hook_functionality,
55 benchmark_hook_registry,
56)
57from transformer_lens.benchmarks.text_quality import benchmark_text_quality
58from transformer_lens.benchmarks.utils import (
59 BenchmarkResult,
60 BenchmarkSeverity,
61 PhaseReferenceData,
62 build_modality_input,
63 compare_tensors,
64 format_results,
65)
66from transformer_lens.benchmarks.weight_processing import (
67 benchmark_attention_output_centering,
68 benchmark_layer_norm_folding,
69 benchmark_mlp_output_centering,
70 benchmark_no_nan_inf,
71 benchmark_unembed_centering,
72 benchmark_value_bias_folding,
73 benchmark_weight_magnitudes,
74 benchmark_weight_modification,
75 benchmark_weight_processing,
76 benchmark_weight_sharing,
77)
78from transformer_lens.config import TransformerBridgeConfig
79from transformer_lens.factories.architecture_adapter_factory import (
80 ArchitectureAdapterFactory,
81)
82from transformer_lens.model_bridge import TransformerBridge
84# Architecture classification — single source of truth in utilities.architectures
85from transformer_lens.utilities.architectures import (
86 NO_HT_COMPARISON_ARCHITECTURES,
87 get_architectures_for_config,
88 is_audio_model,
89 is_encoder_decoder_model,
90 is_masked_lm_model,
91)
92from transformer_lens.utilities.hf_utils import get_hf_token as _hf_token
95def should_skip_ht_comparison(model_name: str, trust_remote_code: bool = False) -> bool:
96 """Benchmark-specific: skip Phase 2/3 for architectures with different hook shapes."""
97 try:
98 config = AutoConfig.from_pretrained(
99 model_name, trust_remote_code=trust_remote_code, token=_hf_token()
100 )
101 architectures = get_architectures_for_config(config)
102 return any(arch in NO_HT_COMPARISON_ARCHITECTURES for arch in architectures)
103 except Exception:
104 return False
107def get_auto_model_class(model_name: str, trust_remote_code: bool = False):
108 """Delegates to the bridge's architecture detection for consistency."""
109 from transformer_lens.model_bridge.sources.transformers import (
110 determine_architecture_from_hf_config,
111 get_hf_model_class_for_architecture,
112 )
114 try:
115 config = AutoConfig.from_pretrained(
116 model_name, trust_remote_code=trust_remote_code, token=_hf_token()
117 )
118 architecture = determine_architecture_from_hf_config(config)
119 return get_hf_model_class_for_architecture(architecture)
120 except Exception:
121 return AutoModelForCausalLM
124def _fixup_custom_model(hf_model) -> None:
125 """Apply post-load fixups for models with custom code (e.g., OpenELM).
127 Recomputes non-persistent buffers (inv_freq, causal_mask) that may be
128 zeroed during HuggingFace's meta-device loading.
129 """
130 # OpenELM fixups
131 if hasattr(hf_model, "transformer") and hasattr(hf_model.transformer, "layers"): 131 ↛ 133line 131 didn't jump to line 133 because the condition on line 131 was never true
132 # Ensure use_cache is set (OpenELM custom config omits it)
133 if not hasattr(hf_model.config, "use_cache") or "use_cache" not in hf_model.config.__dict__:
134 hf_model.config.use_cache = False
136 # Fix 1: Always recompute causal_mask (non-persistent buffer).
137 # After meta→real materialization, the buffer may contain garbage values
138 # rather than clean zeros, so we always recompute.
139 if hasattr(hf_model.transformer, "causal_mask"):
140 cm = hf_model.transformer.causal_mask
141 if cm is not None and cm.numel() > 0:
142 seq_len = cm.shape[-1]
143 correct_mask = torch.triu(
144 torch.ones(seq_len, seq_len, dtype=cm.dtype, device=cm.device),
145 diagonal=1,
146 )
147 hf_model.transformer.causal_mask = correct_mask
149 # Fix 2: Always recompute RoPE inv_freq and sin/cos (non-persistent buffers).
150 rope_max = getattr(hf_model.config, "rope_max_length", None)
151 if rope_max is not None:
152 for layer in hf_model.transformer.layers:
153 if hasattr(layer, "attn") and hasattr(layer.attn, "pos_embedding"):
154 rope = layer.attn.pos_embedding
155 if hasattr(rope, "inv_freq"):
156 correct_inv_freq = 1.0 / (
157 rope.freq_constant
158 ** (
159 torch.arange(0, rope.model_dim, 2, dtype=torch.float32)
160 / rope.model_dim
161 )
162 )
163 rope.inv_freq = correct_inv_freq.to(rope.inv_freq.device)
164 # Force-recompute sin/cos
165 rope._cached_cos = None
166 rope._cached_sin = None
167 rope._compute_sin_cos_embeddings(rope_max)
169 # Create synthetic lm_head for weight-tied models (share_input_output_layers)
170 if getattr(hf_model, "lm_head", None) is None:
171 embed = hf_model.transformer.token_embeddings
172 lm_head = torch.nn.Linear(embed.embedding_dim, embed.num_embeddings, bias=False)
173 lm_head.weight = embed.weight
174 hf_model.lm_head = lm_head
176 # Rotary tables destroyed by meta-device loading, for ANY architecture: the
177 # reference needs the same repair the adapter applies to the bridge, or
178 # Phase 1 compares a correct model against a corrupt one. Only tables that
179 # fail a validity check are touched, so scaled RoPE is never clobbered.
180 from transformer_lens.model_bridge.buffer_restore import restore_rotary_inv_freq
182 restore_rotary_inv_freq(hf_model)
184 if type(hf_model).__name__ == "GiddForDiffusionLM": 184 ↛ 185line 184 didn't jump to line 185 because the condition on line 184 was never true
185 from transformer_lens.model_bridge.supported_architectures.gidd import (
186 restore_frequencies,
187 )
189 restore_frequencies(hf_model)
192def _hf_forward_with_mask_fallback(hf_model, tokens):
193 """Run an HF decoder forward, retrying with a 2D then 4D mask for models that
194 dereference ``attention_mask`` unconditionally (e.g. LLaDA2) -- else the Phase-1
195 capture raises, gets swallowed, and silently degrades to a shape-only check."""
196 try:
197 return hf_model(tokens)
198 except (AttributeError, ValueError):
199 b, s = tokens.shape[0], tokens.shape[-1]
200 for mask in (
201 torch.ones(b, s, dtype=torch.long, device=tokens.device),
202 torch.ones(b, 1, s, s, dtype=torch.long, device=tokens.device),
203 ):
204 try:
205 return hf_model(tokens, attention_mask=mask)
206 except (AttributeError, ValueError):
207 continue
208 raise
211def run_comparison_benchmarks(
212 bridge_model: TransformerBridge,
213 reference_model: Optional[HookedTransformer],
214 test_text: str,
215 phase_name: str,
216 is_processed: bool,
217 verbose: bool = True,
218 phase1_reference: Optional[PhaseReferenceData] = None,
219 restore_dtype_after_equivalence: Optional[torch.dtype] = None,
220) -> List[BenchmarkResult]:
221 """Run standardized comparison benchmarks between Bridge and reference model.
223 This function runs the same comprehensive test suite for both unprocessed (Phase 2)
224 and processed (Phase 3) modes to ensure parity in testing coverage.
226 Args:
227 bridge_model: TransformerBridge model to test
228 reference_model: HookedTransformer reference (same architecture) or None
229 test_text: Input text for testing
230 phase_name: Name of the phase ("Phase 2" or "Phase 3") for logging
231 is_processed: Whether models have processed weights (for weight-specific tests)
232 verbose: Whether to print detailed results
233 phase1_reference: Optional saved Phase 1 HF reference data for equivalence testing
234 restore_dtype_after_equivalence: If set, downcast bridge_model to this dtype after
235 the equivalence comparison but before hook/cache/gradient tests. Used when the
236 bridge was upcast to float32 for precise equivalence testing.
238 Returns:
239 List of BenchmarkResult objects
240 """
241 results: List[BenchmarkResult] = []
243 def add_result(result: BenchmarkResult) -> None:
244 """Add a result and optionally print it immediately."""
245 results.append(result)
246 if verbose: 246 ↛ 247line 246 didn't jump to line 247 because the condition on line 246 was never true
247 result.print_immediate()
249 # Check if we have a same-architecture reference
250 ht_available = reference_model is not None
252 # ========================================================================
253 # 1. Weight Processing Benchmarks (only for processed mode)
254 # MOST BASIC: Check weights are valid before testing anything else
255 # ========================================================================
256 if is_processed: 256 ↛ 257line 256 didn't jump to line 257 because the condition on line 256 was never true
257 if verbose:
258 print("1. Weight Processing Benchmarks (Foundation)")
259 try:
260 # Critical weight validation tests (run first - most basic)
261 add_result(benchmark_no_nan_inf(bridge_model, test_text))
262 add_result(benchmark_weight_magnitudes(bridge_model, test_text))
264 # Detailed weight processing validation benchmarks (don't need reference model)
265 add_result(benchmark_layer_norm_folding(bridge_model, test_text))
266 add_result(benchmark_attention_output_centering(bridge_model, test_text))
267 add_result(benchmark_mlp_output_centering(bridge_model, test_text))
268 add_result(benchmark_unembed_centering(bridge_model, test_text))
269 add_result(benchmark_value_bias_folding(bridge_model, test_text))
271 # Weight comparison tests (require reference model)
272 if ht_available:
273 add_result(
274 benchmark_weight_processing(
275 bridge_model, test_text, reference_model=reference_model
276 )
277 )
278 add_result(
279 benchmark_weight_sharing(
280 bridge_model, test_text, reference_model=reference_model
281 )
282 )
283 else:
284 if verbose:
285 print("⏭️ weight_processing and weight_sharing skipped (no HT reference)")
286 for benchmark_name in ["weight_processing", "weight_sharing"]:
287 add_result(
288 BenchmarkResult(
289 name=benchmark_name,
290 severity=BenchmarkSeverity.SKIPPED,
291 message="Skipped (HookedTransformer not available for this model)",
292 passed=True,
293 )
294 )
296 # weight_modification doesn't need reference model
297 add_result(benchmark_weight_modification(bridge_model, test_text))
298 gc.collect()
299 except Exception as e:
300 if verbose:
301 print(f"✗ Weight processing benchmark failed: {e}\n")
303 # ========================================================================
304 # 2. Model Equivalence Benchmarks (Forward Pass)
305 # Tests basic forward computation - depends on weights being correct
306 # ========================================================================
307 if verbose: 307 ↛ 308line 307 didn't jump to line 308 because the condition on line 307 was never true
308 print("2. Model Equivalence Benchmarks (Forward Pass)")
310 has_phase1_ref = phase1_reference is not None and phase1_reference.hf_logits is not None
312 if ht_available: 312 ↛ 313line 312 didn't jump to line 313 because the condition on line 312 was never true
313 try:
314 add_result(
315 benchmark_logits_equivalence(
316 bridge_model, test_text, reference_model=reference_model
317 )
318 )
319 add_result(
320 benchmark_loss_equivalence(bridge_model, test_text, reference_model=reference_model)
321 )
322 gc.collect()
323 except Exception as e:
324 if verbose:
325 print(f"✗ Equivalence benchmark failed: {e}\n")
326 elif has_phase1_ref: 326 ↛ 330line 326 didn't jump to line 330 because the condition on line 326 was never true
327 # Compare processed bridge against unprocessed Phase 1 reference.
328 # We use log_softmax because center_unembed shifts raw logits by a
329 # softmax-invariant constant. Both passes run in float32 (no bf16 round-trip).
330 try:
331 if verbose:
332 print("Using saved Phase 1 bridge reference for equivalence comparison")
334 assert phase1_reference is not None
335 assert phase1_reference.hf_logits is not None
337 # Compare log_softmax (centering-invariant) instead of raw logits.
338 bridge_logits = bridge_model(test_text, return_type="logits")
339 ref_logits = phase1_reference.hf_logits.to(bridge_logits.device)
340 bridge_log_probs = torch.nn.functional.log_softmax(bridge_logits, dim=-1)
341 ref_log_probs = torch.nn.functional.log_softmax(ref_logits, dim=-1)
343 # Both passes in float32 — remaining error is float32 non-associativity
344 # in weight processing (~0.006 max_diff on 24-layer Qwen2).
345 logits_atol = 0.01
346 logits_rtol = 1e-4
347 loss_atol = 1e-3
349 add_result(
350 compare_tensors(
351 bridge_log_probs,
352 ref_log_probs,
353 atol=logits_atol,
354 rtol=logits_rtol,
355 name="logits_equivalence",
356 )
357 )
358 if phase1_reference.hf_loss is not None:
359 add_result(
360 benchmark_loss_equivalence(
361 bridge_model,
362 test_text,
363 reference_loss=phase1_reference.hf_loss,
364 atol=loss_atol,
365 )
366 )
367 else:
368 add_result(
369 BenchmarkResult(
370 name="loss_equivalence",
371 severity=BenchmarkSeverity.SKIPPED,
372 message="Skipped (no Phase 1 loss reference available)",
373 passed=True,
374 )
375 )
376 gc.collect()
377 except Exception as e:
378 if verbose:
379 print(f"✗ Phase 1 reference comparison failed: {e}\n")
380 else:
381 if verbose: 381 ↛ 382line 381 didn't jump to line 382 because the condition on line 381 was never true
382 print("⏭️ Skipped (no HookedTransformer reference)\n")
383 for benchmark_name in ["logits_equivalence", "loss_equivalence"]:
384 add_result(
385 BenchmarkResult(
386 name=benchmark_name,
387 severity=BenchmarkSeverity.SKIPPED,
388 message="Skipped (HookedTransformer not available for this model)",
389 passed=True,
390 )
391 )
393 # Restore native dtype so remaining tests run in the model's real dtype.
394 # Both bridge and reference must be downcast so hook comparisons use the
395 # same precision — otherwise bridge activations (bfloat16) are compared
396 # against reference activations (float32), producing spurious mismatches.
397 if restore_dtype_after_equivalence is not None: 397 ↛ 398line 397 didn't jump to line 398 because the condition on line 397 was never true
398 try:
399 bridge_model.to(restore_dtype_after_equivalence)
400 if reference_model is not None:
401 reference_model.to(restore_dtype_after_equivalence)
402 if verbose:
403 print(f" (restored to {restore_dtype_after_equivalence} for remaining tests)\n")
404 except Exception as e:
405 if verbose:
406 print(f"⚠ Could not restore dtype: {e}\n")
408 # ========================================================================
409 # 3. Hook Registration Benchmarks
410 # Tests hooks exist and are registered - depends on model structure
411 # ========================================================================
412 if verbose: 412 ↛ 413line 412 didn't jump to line 413 because the condition on line 412 was never true
413 print("3. Hook Registration Benchmarks")
415 if ht_available: 415 ↛ 416line 415 didn't jump to line 416 because the condition on line 415 was never true
416 try:
417 add_result(benchmark_hook_registry(bridge_model, reference_model=reference_model))
418 gc.collect()
419 except Exception as e:
420 if verbose:
421 print(f"✗ Hook registry benchmark failed: {e}\n")
422 else:
423 try:
424 add_result(benchmark_hook_registry(bridge_model))
425 gc.collect()
426 except Exception as e:
427 if verbose:
428 print(f"✗ Hook registry benchmark failed: {e}\n")
430 # ========================================================================
431 # 4. Forward Hook Functionality Benchmarks
432 # Tests hooks fire and produce correct values - depends on forward pass + hooks
433 # ========================================================================
434 if verbose: 434 ↛ 435line 434 didn't jump to line 435 because the condition on line 434 was never true
435 print("4. Forward Hook Functionality Benchmarks")
437 if ht_available: 437 ↛ 438line 437 didn't jump to line 438 because the condition on line 437 was never true
438 try:
439 add_result(
440 benchmark_hook_functionality(
441 bridge_model, test_text, reference_model=reference_model
442 )
443 )
444 add_result(
445 benchmark_critical_forward_hooks(
446 bridge_model, test_text, reference_model=reference_model
447 )
448 )
449 add_result(
450 benchmark_forward_hooks(bridge_model, test_text, reference_model=reference_model)
451 )
452 add_result(benchmark_gated_hooks_fire(bridge_model, test_text))
453 # Reset hooks to prevent handle leaks
454 if hasattr(bridge_model, "reset_hooks"):
455 bridge_model.reset_hooks()
456 if reference_model is not None and hasattr(reference_model, "reset_hooks"):
457 reference_model.reset_hooks()
458 gc.collect()
459 except Exception as e:
460 if verbose:
461 print(f"✗ Forward hook benchmark failed: {e}\n")
462 else:
463 try:
464 add_result(benchmark_hook_functionality(bridge_model, test_text))
465 add_result(benchmark_critical_forward_hooks(bridge_model, test_text))
466 add_result(benchmark_forward_hooks(bridge_model, test_text))
467 add_result(benchmark_gated_hooks_fire(bridge_model, test_text))
468 # Reset hooks to prevent handle leaks
469 if hasattr(bridge_model, "reset_hooks"): 469 ↛ 471line 469 didn't jump to line 471 because the condition on line 469 was always true
470 bridge_model.reset_hooks()
471 gc.collect()
472 except Exception as e:
473 if verbose:
474 print(f"✗ Forward hook benchmark failed: {e}\n")
476 # ========================================================================
477 # 5. Activation Cache Benchmarks
478 # Tests caching mechanism - depends on forward pass + hooks working
479 # ========================================================================
480 if verbose: 480 ↛ 481line 480 didn't jump to line 481 because the condition on line 480 was never true
481 print("5. Activation Cache Benchmarks")
483 if ht_available: 483 ↛ 484line 483 didn't jump to line 484 because the condition on line 483 was never true
484 try:
485 add_result(
486 benchmark_run_with_cache(bridge_model, test_text, reference_model=reference_model)
487 )
488 add_result(
489 benchmark_activation_cache(bridge_model, test_text, reference_model=reference_model)
490 )
491 # Reset hooks to prevent handle leaks
492 if hasattr(bridge_model, "reset_hooks"):
493 bridge_model.reset_hooks()
494 if reference_model is not None and hasattr(reference_model, "reset_hooks"):
495 reference_model.reset_hooks()
496 gc.collect()
497 except Exception as e:
498 if verbose:
499 print(f"✗ Activation cache benchmark failed: {e}\n")
500 else:
501 try:
502 add_result(benchmark_run_with_cache(bridge_model, test_text))
503 add_result(benchmark_activation_cache(bridge_model, test_text))
504 # Reset hooks to prevent handle leaks
505 if hasattr(bridge_model, "reset_hooks"): 505 ↛ 507line 505 didn't jump to line 507 because the condition on line 505 was always true
506 bridge_model.reset_hooks()
507 gc.collect()
508 except Exception as e:
509 if verbose:
510 print(f"✗ Activation cache benchmark failed: {e}\n")
512 # ========================================================================
513 # 6. Backward Gradient Benchmarks
514 # MOST COMPLEX: Tests gradients and backward hooks - depends on everything above
515 # ========================================================================
516 if verbose: 516 ↛ 517line 516 didn't jump to line 517 because the condition on line 516 was never true
517 print("6. Backward Gradient Benchmarks")
519 # Gradient comparisons are graded against fp32-calibrated thresholds
520 # (REL_L2_TOLERANCE): bf16's rounding floor alone is ~2e-3 rel_l2, inside the
521 # measured bug band, so reduced-precision gradients cannot be graded at all.
522 # Upcast for the gradient section on every device (MPS additionally lacks
523 # bf16 autograd), then restore below.
524 bridge_grad_dtype = bridge_model.cfg.dtype if hasattr(bridge_model, "cfg") else None
525 grad_fp32_upcast = needs_fp32_gradients(bridge_grad_dtype)
526 if grad_fp32_upcast: 526 ↛ 527line 526 didn't jump to line 527 because the condition on line 526 was never true
527 try:
528 bridge_model.to(torch.float32)
529 if reference_model is not None:
530 reference_model.to(torch.float32)
531 except Exception:
532 grad_fp32_upcast = False # Upcast failed; proceed as-is
534 if ht_available: 534 ↛ 535line 534 didn't jump to line 535 because the condition on line 534 was never true
535 try:
536 add_result(
537 benchmark_gradient_computation(
538 bridge_model, test_text, reference_model=reference_model
539 )
540 )
541 add_result(
542 benchmark_critical_backward_hooks(
543 bridge_model, test_text, reference_model=reference_model
544 )
545 )
546 add_result(
547 benchmark_backward_hooks(bridge_model, test_text, reference_model=reference_model)
548 )
549 # Reset hooks to prevent handle leaks
550 if hasattr(bridge_model, "reset_hooks"):
551 bridge_model.reset_hooks()
552 if reference_model is not None and hasattr(reference_model, "reset_hooks"):
553 reference_model.reset_hooks()
554 gc.collect()
555 except Exception as e:
556 if verbose:
557 print(f"✗ Gradient benchmark failed: {e}\n")
558 else:
559 try:
560 add_result(benchmark_gradient_computation(bridge_model, test_text))
561 add_result(benchmark_critical_backward_hooks(bridge_model, test_text))
562 add_result(benchmark_backward_hooks(bridge_model, test_text))
563 # Reset hooks to prevent handle leaks
564 if hasattr(bridge_model, "reset_hooks"): 564 ↛ 566line 564 didn't jump to line 566 because the condition on line 564 was always true
565 bridge_model.reset_hooks()
566 gc.collect()
567 except Exception as e:
568 if verbose:
569 print(f"✗ Gradient benchmark failed: {e}\n")
571 if grad_fp32_upcast and bridge_grad_dtype is not None: 571 ↛ 572line 571 didn't jump to line 572 because the condition on line 571 was never true
572 try:
573 bridge_model.to(bridge_grad_dtype)
574 if reference_model is not None:
575 reference_model.to(bridge_grad_dtype)
576 except Exception:
577 pass
579 return results
582def run_benchmark_suite(
583 model_name: str,
584 device: str = "cpu",
585 dtype: torch.dtype = torch.float32,
586 test_text: Optional[str] = None,
587 use_hf_reference: bool = True,
588 use_ht_reference: bool = True,
589 enable_compatibility_mode: bool = True,
590 verbose: bool = True,
591 track_memory: bool = False,
592 test_weight_processing_individually: bool = False,
593 phases: list[int] | None = None,
594 trust_remote_code: bool = False,
595 judge_model: PreTrainedModel | None = None,
596 judge_tokenizer: PreTrainedTokenizerBase | None = None,
597 prompt_profile: str | None = None,
598) -> List[BenchmarkResult]:
599 """Run comprehensive benchmark suite for TransformerBridge.
601 This function implements an optimized multi-phase approach to minimize model reloading:
602 Phase 1: HF + Bridge (unprocessed) - Compare against raw HuggingFace model
603 Phase 2: Bridge (unprocessed) + HT (unprocessed) - Compare unprocessed models
604 Phase 3: Bridge (processed) + HT (processed) - Full compatibility mode testing
605 Phase 4: Text Quality - profile prompts scored by a pinned judge's perplexity ratio
606 Phase 5: Individual Weight Processing Flags (optional)
607 Phase 6: Combined Weight Processing Flags (optional)
609 When test_weight_processing_individually=True, Phases 5 & 6 run after
610 Phase 3, testing each weight processing flag individually and in combinations.
612 Args:
613 model_name: Name of the model to benchmark (e.g., "gpt2")
614 device: Device to run on ("cpu" or "cuda")
615 dtype: Precision for model loading (default: torch.float32). Use
616 torch.bfloat16 to halve memory for larger models. Phase 2/3
617 comparisons automatically upcast to float32 for precision.
618 test_text: Optional test text (default: standard test prompt)
619 use_hf_reference: Whether to compare against HuggingFace model
620 use_ht_reference: Whether to compare against HookedTransformer
621 enable_compatibility_mode: Whether to enable compatibility mode on bridge
622 verbose: Whether to print results to console
623 track_memory: Whether to track and report memory usage (requires psutil)
624 test_weight_processing_individually: Whether to run granular weight processing
625 tests that check each processing flag individually (default: False)
626 phases: Optional list of phase numbers to run (e.g., [1, 2, 3]). If None, runs all phases.
627 trust_remote_code: Whether to trust remote code for custom architectures.
628 judge_model: Optional pre-loaded Phase-4 judge. When provided with
629 judge_tokenizer, avoids reloading for each model in batch.
630 judge_tokenizer: Optional pre-loaded tokenizer for the Phase-4 judge.
631 prompt_profile: Optional Phase-4 prompt profile (e.g. "chat",
632 "task:translation@en-de"). Resolved from curation + the registry
633 when None.
635 Returns:
636 List of BenchmarkResult objects
637 """
638 if test_text is None: 638 ↛ 645line 638 didn't jump to line 645 because the condition on line 638 was always true
639 test_text = (
640 "Natural language processing tasks, such as question answering, "
641 "machine translation, reading comprehension, and summarization, "
642 "are typically approached with supervised learning."
643 )
645 results: List[BenchmarkResult] = []
647 # Memory tracking setup
648 memory_tracker = None
649 if track_memory: 649 ↛ 650line 649 didn't jump to line 650 because the condition on line 649 was never true
650 try:
651 import psutil
653 process = psutil.Process()
654 initial_memory = process.memory_info().rss / 1024 / 1024 # MB
656 def get_memory_mb():
657 return process.memory_info().rss / 1024 / 1024
659 memory_tracker = {"initial": initial_memory, "checkpoints": []}
660 if verbose:
661 print(f"Memory tracking enabled (initial: {initial_memory:.1f} MB)")
662 except ImportError:
663 if verbose:
664 print("⚠ psutil not available - memory tracking disabled")
665 track_memory = False
667 if verbose: 667 ↛ 668line 667 didn't jump to line 668 because the condition on line 667 was never true
668 print(f"\n{'='*80}")
669 print(f"Running TransformerBridge Benchmark Suite")
670 print(f"Model: {model_name}")
671 print(f"Device: {device}")
672 print(f"{'='*80}\n")
674 # Auto-skip HT comparison for architectures with intentionally different hook shapes
675 if use_ht_reference and should_skip_ht_comparison(model_name, trust_remote_code): 675 ↛ 676line 675 didn't jump to line 676 because the condition on line 675 was never true
676 use_ht_reference = False
677 if verbose:
678 print(
679 "Note: Skipping HookedTransformer comparison (architecture uses "
680 "different hook shapes by design). Phase 1 is the gold standard.\n"
681 )
683 # Early exit if only running Phase 5/6 (they load their own models independently)
684 if phases is not None and all(p in [5, 6] for p in phases): 684 ↛ 685line 684 didn't jump to line 685 because the condition on line 684 was never true
685 if verbose:
686 print(f"Skipping Phase 1-4 (only running Phase {', '.join(map(str, sorted(phases)))})")
687 print("Phase 5/6 load their own models independently\n")
689 from transformer_lens.benchmarks.granular_weight_processing import (
690 run_granular_weight_processing_benchmarks,
691 )
693 if 5 in phases and test_weight_processing_individually and enable_compatibility_mode:
694 phase5_results = run_granular_weight_processing_benchmarks(
695 model_name=model_name,
696 device=device,
697 test_text=test_text,
698 verbose=verbose,
699 phase=5,
700 )
701 for config_name, config_results in phase5_results.items():
702 for result in config_results:
703 result.phase = 5
704 results.append(result)
705 if verbose:
706 result.print_immediate()
708 if 6 in phases and test_weight_processing_individually and enable_compatibility_mode:
709 phase6_results = run_granular_weight_processing_benchmarks(
710 model_name=model_name,
711 device=device,
712 test_text=test_text,
713 verbose=verbose,
714 phase=6,
715 )
716 for config_name, config_results in phase6_results.items():
717 for result in config_results:
718 result.phase = 6
719 results.append(result)
720 if verbose:
721 result.print_immediate()
723 return results
725 # Track current phase for result tagging
726 current_phase: List[Optional[int]] = [None] # Use list to allow modification in nested function
728 def should_run_phase(phase_num: int) -> bool:
729 """Check if a phase should run based on the phases filter."""
730 return phases is None or phase_num in phases
732 def add_result(result: BenchmarkResult) -> None:
733 """Add a result and optionally print it immediately."""
734 # Tag result with current phase
735 if current_phase[0] is not None and result.phase is None: 735 ↛ 737line 735 didn't jump to line 737 because the condition on line 735 was always true
736 result.phase = current_phase[0]
737 results.append(result)
738 if verbose: 738 ↛ 739line 738 didn't jump to line 739 because the condition on line 738 was never true
739 result.print_immediate()
741 def cleanup_tensors(*tensors) -> None:
742 """Free memory from tensors and caches."""
743 for tensor in tensors:
744 if tensor is not None:
745 # If it's an ActivationCache, clear all tensors
746 if hasattr(tensor, "cache_dict"):
747 for key in list(tensor.cache_dict.keys()):
748 val = tensor.cache_dict[key]
749 if val is not None and isinstance(val, torch.Tensor):
750 del val
751 tensor.cache_dict[key] = None
752 tensor.cache_dict.clear()
753 # If it's a regular tensor, just delete it
754 elif isinstance(tensor, torch.Tensor):
755 del tensor
756 # Force cleanup
757 gc.collect()
758 if device != "cpu" and torch.cuda.is_available():
759 torch.cuda.empty_cache()
760 if device == "mps" and hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
761 torch.mps.synchronize()
762 torch.mps.empty_cache()
764 def cleanup_model(model, model_name_str: str):
765 """Free up memory by deleting a model and forcing garbage collection."""
766 import gc
768 if verbose: 768 ↛ 769line 768 didn't jump to line 769 because the condition on line 768 was never true
769 print(f"Cleaning up {model_name_str}...")
771 # Track memory before cleanup
772 if track_memory and memory_tracker is not None: 772 ↛ 773line 772 didn't jump to line 773 because the condition on line 772 was never true
773 memory_before = get_memory_mb()
775 # Move model to CPU first to free GPU memory immediately
776 if device != "cpu" and hasattr(model, "cpu"): 776 ↛ 777line 776 didn't jump to line 777 because the condition on line 776 was never true
777 try:
778 model.cpu()
779 if torch.cuda.is_available():
780 torch.cuda.empty_cache()
781 if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
782 torch.mps.synchronize()
783 torch.mps.empty_cache()
784 except Exception:
785 pass
787 # Explicitly remove all hooks to prevent memory leaks
788 if hasattr(model, "modules"): 788 ↛ 824line 788 didn't jump to line 824 because the condition on line 788 was always true
789 try:
790 for module in model.modules():
791 # Clear PyTorch hooks
792 if hasattr(module, "_forward_hooks"): 792 ↛ 794line 792 didn't jump to line 794 because the condition on line 792 was always true
793 module._forward_hooks.clear()
794 if hasattr(module, "_backward_hooks"): 794 ↛ 796line 794 didn't jump to line 796 because the condition on line 794 was always true
795 module._backward_hooks.clear()
796 if hasattr(module, "_forward_pre_hooks"): 796 ↛ 798line 796 didn't jump to line 798 because the condition on line 796 was always true
797 module._forward_pre_hooks.clear()
798 if hasattr(module, "_backward_pre_hooks"): 798 ↛ 800line 798 didn't jump to line 800 because the condition on line 798 was always true
799 module._backward_pre_hooks.clear()
800 if hasattr(module, "_state_dict_hooks"): 800 ↛ 802line 800 didn't jump to line 802 because the condition on line 800 was always true
801 module._state_dict_hooks.clear()
802 if hasattr(module, "_state_dict_pre_hooks"): 802 ↛ 804line 802 didn't jump to line 804 because the condition on line 802 was always true
803 module._state_dict_pre_hooks.clear()
804 if hasattr(module, "_load_state_dict_pre_hooks"): 804 ↛ 806line 804 didn't jump to line 806 because the condition on line 804 was always true
805 module._load_state_dict_pre_hooks.clear()
806 if hasattr(module, "_load_state_dict_post_hooks"): 806 ↛ 810line 806 didn't jump to line 810 because the condition on line 806 was always true
807 module._load_state_dict_post_hooks.clear()
809 # Clear TransformerLens-specific hooks
810 if hasattr(module, "remove_all_hooks"): 810 ↛ 811line 810 didn't jump to line 811 because the condition on line 810 was never true
811 module.remove_all_hooks()
813 # Clear gradients
814 if hasattr(module, "zero_grad"): 814 ↛ 790line 814 didn't jump to line 790 because the condition on line 814 was always true
815 try:
816 module.zero_grad(set_to_none=True)
817 except Exception:
818 pass
819 except Exception:
820 # If hook cleanup fails, continue anyway
821 pass
823 # Clear top-level hooks
824 if hasattr(model, "_forward_hooks"): 824 ↛ 826line 824 didn't jump to line 826 because the condition on line 824 was always true
825 model._forward_hooks.clear()
826 if hasattr(model, "_backward_hooks"): 826 ↛ 828line 826 didn't jump to line 828 because the condition on line 826 was always true
827 model._backward_hooks.clear()
828 if hasattr(model, "_forward_pre_hooks"): 828 ↛ 832line 828 didn't jump to line 832 because the condition on line 828 was always true
829 model._forward_pre_hooks.clear()
831 # Clear top-level gradients
832 if hasattr(model, "zero_grad"): 832 ↛ 839line 832 didn't jump to line 839 because the condition on line 832 was always true
833 try:
834 model.zero_grad(set_to_none=True)
835 except Exception:
836 pass
838 # Break circular references to help GC
839 if hasattr(model, "_modules"): 839 ↛ 853line 839 didn't jump to line 853 because the condition on line 839 was always true
840 # Clear each submodule's __dict__ to break circular references
841 for name, submodule in list(model._modules.items()):
842 if submodule is not None: 842 ↛ 841line 842 didn't jump to line 841 because the condition on line 842 was always true
843 # Clear submodule hooks
844 if hasattr(submodule, "_forward_hooks"): 844 ↛ 846line 844 didn't jump to line 846 because the condition on line 844 was always true
845 submodule._forward_hooks.clear()
846 if hasattr(submodule, "_backward_hooks"): 846 ↛ 849line 846 didn't jump to line 849 because the condition on line 846 was always true
847 submodule._backward_hooks.clear()
848 # Break reference
849 model._modules[name] = None
850 model._modules.clear()
852 # Clear parameters dict
853 if hasattr(model, "_parameters"): 853 ↛ 862line 853 didn't jump to line 862 because the condition on line 853 was always true
854 for param_name in list(model._parameters.keys()): 854 ↛ 855line 854 didn't jump to line 855 because the loop on line 854 never started
855 param = model._parameters[param_name]
856 if param is not None:
857 del param
858 model._parameters[param_name] = None
859 model._parameters.clear()
861 # Clear buffers dict
862 if hasattr(model, "_buffers"): 862 ↛ 870line 862 didn't jump to line 870 because the condition on line 862 was always true
863 for buffer_name in list(model._buffers.keys()): 863 ↛ 864line 863 didn't jump to line 864 because the loop on line 863 never started
864 buffer = model._buffers[buffer_name]
865 if buffer is not None:
866 del buffer
867 model._buffers[buffer_name] = None
868 model._buffers.clear()
870 del model
872 # Aggressive garbage collection (multiple passes to break circular references)
873 for _ in range(3):
874 gc.collect()
876 # Clear GPU cache
877 if device != "cpu" and torch.cuda.is_available(): 877 ↛ 878line 877 didn't jump to line 878 because the condition on line 877 was never true
878 torch.cuda.empty_cache()
879 torch.cuda.synchronize()
880 if device == "mps" and hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"): 880 ↛ 881line 880 didn't jump to line 881 because the condition on line 880 was never true
881 torch.mps.synchronize()
882 torch.mps.empty_cache()
884 # Track memory after cleanup
885 if track_memory and memory_tracker is not None: 885 ↛ 886line 885 didn't jump to line 886 because the condition on line 885 was never true
886 memory_after = get_memory_mb()
887 freed_mb = memory_before - memory_after
888 memory_tracker["checkpoints"].append(
889 {
890 "label": f"Cleanup: {model_name_str}",
891 "memory_mb": memory_after,
892 "freed_mb": freed_mb,
893 }
894 )
895 if verbose and freed_mb > 0:
896 print(f" Freed {freed_mb:.1f} MB")
898 # ========================================================================
899 # PHASE 1: HuggingFace + Bridge (unprocessed)
900 # ========================================================================
901 current_phase[0] = 1
902 if verbose: 902 ↛ 903line 902 didn't jump to line 903 because the condition on line 902 was never true
903 print(f"\n{'='*80}")
904 print("PHASE 1: HuggingFace + TransformerBridge (unprocessed)")
905 print(f"{'='*80}\n")
907 bridge_unprocessed = None
908 hf_model = None
909 phase1_reference = PhaseReferenceData()
911 # Load bridge without weights first to detect attn_implementation and dtype
912 if verbose: 912 ↛ 913line 912 didn't jump to line 913 because the condition on line 912 was never true
913 print("Detecting model configuration...")
914 bridge_dtype = dtype
915 attn_implementation = None
916 try:
917 # Load a lightweight version without weights to get config
918 bridge_config_only = TransformerBridge.boot_transformers(model_name, device=device, dtype=bridge_dtype, load_weights=False, trust_remote_code=trust_remote_code) # type: ignore[attr-defined]
919 # Match bridge's attn_implementation: check adapter config first, then
920 # default to "eager" (bridge uses output_attentions=True which forces eager).
921 if hasattr(bridge_config_only.adapter.cfg, "attn_implementation"): 921 ↛ 923line 921 didn't jump to line 923 because the condition on line 921 was always true
922 attn_implementation = bridge_config_only.adapter.cfg.attn_implementation
923 if attn_implementation is None: 923 ↛ 925line 923 didn't jump to line 925 because the condition on line 923 was always true
924 attn_implementation = "eager"
925 if verbose: 925 ↛ 926line 925 didn't jump to line 926 because the condition on line 925 was never true
926 print(f"✓ Detected attn_implementation={attn_implementation}")
927 # Clean up config-only bridge immediately to free memory
928 del bridge_config_only
929 gc.collect()
930 except Exception as e:
931 if verbose:
932 print(f"⚠ Could not detect config (will use defaults): {str(e)}")
933 # Config-only bridge failed; apply architecture patches directly to prevent
934 # _init_weights from re-randomizing loaded weights.
935 if trust_remote_code:
936 try:
937 from transformer_lens.model_bridge.sources.transformers import (
938 determine_architecture_from_hf_config,
939 map_default_transformer_lens_config,
940 )
942 hf_cfg = AutoConfig.from_pretrained(
943 model_name, trust_remote_code=True, token=_hf_token()
944 )
945 tl_cfg = map_default_transformer_lens_config(hf_cfg)
946 arch = determine_architecture_from_hf_config(hf_cfg)
947 bridge_cfg = TransformerBridgeConfig.from_dict(tl_cfg.__dict__)
948 bridge_cfg.architecture = arch
949 bridge_cfg.model_name = model_name
950 adapter = ArchitectureAdapterFactory.select_architecture_adapter(bridge_cfg)
951 adapter.prepare_loading(model_name, {})
952 if verbose:
953 print("✓ Applied architecture patches for custom code model")
954 del adapter, bridge_cfg, tl_cfg, hf_cfg
955 except Exception as patch_err:
956 if verbose:
957 print(f"⚠ Could not apply architecture patches: {patch_err}")
959 hf_saved_logits = None
960 hf_saved_loss = None
962 if use_hf_reference and should_run_phase(1): 962 ↛ 1037line 962 didn't jump to line 1037 because the condition on line 962 was always true
963 try:
964 if verbose: 964 ↛ 965line 964 didn't jump to line 965 because the condition on line 964 was never true
965 print("Loading HuggingFace reference model...")
966 # Match bridge loading path: no device_map, explicit .to(device),
967 # and matching torch_dtype. When dtype=float32, loading in float32
968 # ensures non-persistent buffers (e.g., Gemma3's embed_scale) are
969 # computed at full precision. When dtype=bfloat16, both HF and
970 # Bridge load in bfloat16 so comparisons are apples-to-apples.
971 hf_kwargs: dict[str, object] = {
972 "low_cpu_mem_usage": True, # Reduce memory spikes during loading
973 "torch_dtype": dtype,
974 }
975 if _hf_token(): 975 ↛ 977line 975 didn't jump to line 977 because the condition on line 975 was always true
976 hf_kwargs["token"] = _hf_token()
977 if attn_implementation is not None: 977 ↛ 982line 977 didn't jump to line 982 because the condition on line 977 was always true
978 hf_kwargs["attn_implementation"] = attn_implementation
979 if verbose: 979 ↛ 980line 979 didn't jump to line 980 because the condition on line 979 was never true
980 print(f"Using attn_implementation={attn_implementation}")
981 # Use appropriate AutoModel class (e.g., AutoModelForSeq2SeqLM for T5)
982 auto_model_class = get_auto_model_class(model_name, trust_remote_code=trust_remote_code)
983 if verbose and auto_model_class != AutoModelForCausalLM: 983 ↛ 984line 983 didn't jump to line 984 because the condition on line 983 was never true
984 print(f"Using {auto_model_class.__name__}")
985 # Ensure pad_token_id exists (some models crash without it during init).
986 hf_config = AutoConfig.from_pretrained(
987 model_name, trust_remote_code=trust_remote_code, token=_hf_token()
988 )
989 if not hasattr(hf_config, "pad_token_id") or "pad_token_id" not in hf_config.__dict__: 989 ↛ 990line 989 didn't jump to line 990 because the condition on line 989 was never true
990 eos = getattr(hf_config, "eos_token_id", None)
991 hf_config.pad_token_id = eos[0] if isinstance(eos, (list, tuple)) else eos
992 hf_kwargs["config"] = hf_config
993 if trust_remote_code: 993 ↛ 994line 993 didn't jump to line 994 because the condition on line 993 was never true
994 hf_kwargs["trust_remote_code"] = True
995 hf_model = auto_model_class.from_pretrained(model_name, **hf_kwargs) # type: ignore[arg-type]
996 hf_model = hf_model.to(device)
997 # Post-load fixup for custom code models (e.g., OpenELM).
998 # Must run AFTER .to(device) so non-persistent buffers (RoPE sin/cos,
999 # causal_mask) are recomputed on the target device, matching the bridge
1000 # which also recomputes after .to(device).
1001 _fixup_custom_model(hf_model)
1002 hf_model.eval()
1003 # Detect dtype from HF model
1004 try:
1005 bridge_dtype = next(hf_model.parameters()).dtype
1006 if verbose: 1006 ↛ 1007line 1006 didn't jump to line 1007 because the condition on line 1006 was never true
1007 print(f"Detected dtype={bridge_dtype}")
1008 except StopIteration:
1009 pass
1010 # When float32 was requested but the model natively uses reduced
1011 # precision, upcast for maximum benchmark accuracy. When dtype was
1012 # explicitly set to bfloat16/float16 (e.g., to fit larger models in
1013 # memory), respect it — both HF and Bridge will run in that precision.
1014 if dtype == torch.float32 and bridge_dtype in (torch.float16, torch.bfloat16): 1014 ↛ 1015line 1014 didn't jump to line 1015 because the condition on line 1014 was never true
1015 if verbose:
1016 print(f"⚠ {bridge_dtype} detected, upcasting to float32 for benchmarking...")
1017 hf_model.to(torch.float32)
1018 bridge_dtype = torch.float32
1019 if verbose:
1020 print("✓ Upcast to float32 in-place")
1021 elif bridge_dtype != dtype: 1021 ↛ 1022line 1021 didn't jump to line 1022 because the condition on line 1021 was never true
1022 bridge_dtype = dtype # Trust the requested dtype
1023 if verbose: 1023 ↛ 1024line 1023 didn't jump to line 1024 because the condition on line 1023 was never true
1024 print("✓ HuggingFace model loaded")
1026 # HF reference logits will be captured AFTER the bridge is
1027 # loaded so we can use bridge.to_tokens() for consistent
1028 # tokenization (e.g. BOS prepending). This happens right
1029 # after the component benchmark, while both models are still
1030 # in memory, before the HF model is deleted.
1032 except Exception as e:
1033 if verbose:
1034 print(f"✗ Could not load HuggingFace model: {str(e)}\n")
1036 # Now load the full bridge with correct dtype (GPU is mostly free)
1037 if verbose: 1037 ↛ 1038line 1037 didn't jump to line 1038 because the condition on line 1037 was never true
1038 print("Loading TransformerBridge (unprocessed)...")
1039 try:
1040 bridge_unprocessed = TransformerBridge.boot_transformers(model_name, device=device, dtype=bridge_dtype, trust_remote_code=trust_remote_code) # type: ignore[attr-defined]
1041 if verbose: 1041 ↛ 1042line 1041 didn't jump to line 1042 because the condition on line 1041 was never true
1042 print("✓ TransformerBridge loaded (unprocessed)\n")
1043 # Apply the adapter's prepare_model() to the HF reference model so
1044 # both bridge and reference have the same fixups (e.g., weight tying).
1045 # This keeps model-specific logic in the adapter, not the benchmark.
1046 if hf_model is not None and hasattr(bridge_unprocessed, "adapter"): 1046 ↛ 1066line 1046 didn't jump to line 1066 because the condition on line 1046 was always true
1047 bridge_unprocessed.adapter.prepare_model(hf_model)
1048 except Exception as e:
1049 import traceback
1051 error_trace = traceback.format_exc()
1052 add_result(
1053 BenchmarkResult(
1054 name="load_bridge_unprocessed",
1055 severity=BenchmarkSeverity.ERROR,
1056 message=f"Failed to load unprocessed TransformerBridge: {str(e)}",
1057 passed=False,
1058 )
1059 )
1060 if verbose:
1061 print(f"✗ Failed to load TransformerBridge: {str(e)}")
1062 print(f"\nStack trace:\n{error_trace}")
1063 return results
1065 # Detect audio/vision models once for use across all phases
1066 _is_audio = bridge_unprocessed is not None and getattr(
1067 bridge_unprocessed.cfg, "is_audio_model", False
1068 )
1069 _is_visual = bridge_unprocessed is not None and getattr(
1070 bridge_unprocessed.cfg, "is_visual_model", False
1071 )
1072 # Shared non-text input (spectrogram, waveform, or pixels) — the same tensor is used
1073 # for the HF reference capture and the bridge forward so they stay comparable.
1074 _test_modality_input = (
1075 build_modality_input(bridge_unprocessed, device=device, dtype=dtype)
1076 if (_is_audio or _is_visual)
1077 else None
1078 )
1080 # Run Phase 1 benchmarks
1081 if should_run_phase(1) and bridge_unprocessed: 1081 ↛ 1259line 1081 didn't jump to line 1259 because the condition on line 1081 was always true
1082 if verbose: 1082 ↛ 1083line 1082 didn't jump to line 1083 because the condition on line 1082 was never true
1083 print("Running Phase 1 benchmarks...\n")
1085 # Component-level benchmarks
1086 if verbose: 1086 ↛ 1087line 1086 didn't jump to line 1087 because the condition on line 1086 was never true
1087 print("1. Component-Level Benchmarks")
1088 if hf_model is not None: 1088 ↛ 1172line 1088 didn't jump to line 1172 because the condition on line 1088 was always true
1089 # Full mode: component benchmark with independent HF model (brief 2.0x)
1090 try:
1091 component_result = benchmark_all_components(bridge_unprocessed, hf_model)
1092 add_result(component_result)
1093 if verbose: 1093 ↛ 1094line 1093 didn't jump to line 1094 because the condition on line 1093 was never true
1094 status = "✓" if component_result.passed else "✗"
1095 print(f"{status} {component_result.message}\n")
1096 gc.collect()
1097 if device != "cpu" and torch.cuda.is_available(): 1097 ↛ 1098line 1097 didn't jump to line 1098 because the condition on line 1097 was never true
1098 torch.cuda.empty_cache()
1099 if device == "mps" and hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"): 1099 ↛ 1100line 1099 didn't jump to line 1100 because the condition on line 1099 was never true
1100 torch.mps.synchronize()
1101 torch.mps.empty_cache()
1102 except Exception as e:
1103 if verbose:
1104 print(f"✗ Component benchmark failed: {e}\n")
1106 # Capture HF reference outputs. Both models are still in memory (2.0x window).
1107 if verbose: 1107 ↛ 1108line 1107 didn't jump to line 1108 because the condition on line 1107 was never true
1108 print("Capturing HF reference outputs to CPU...")
1109 try:
1110 if _test_modality_input is not None: 1110 ↛ 1112line 1110 didn't jump to line 1112 because the condition on line 1110 was never true
1111 # Audio/vision models: use the shared non-text input for HF vs bridge
1112 with torch.no_grad():
1113 if _is_visual:
1114 hf_out = hf_model(pixel_values=_test_modality_input)
1115 else:
1116 hf_out = hf_model(input_values=_test_modality_input)
1117 # Bare encoders output last_hidden_state, not logits
1118 if hasattr(hf_out, "logits") and hf_out.logits is not None:
1119 hf_saved_logits = hf_out.logits.detach().cpu().clone()
1120 else:
1121 hf_saved_logits = hf_out.last_hidden_state.detach().cpu().clone()
1122 # No loss computation — there are no next-token labels here
1123 if verbose:
1124 kind = "vision" if _is_visual else "audio"
1125 print(
1126 f"✓ Captured HF {kind} output {hf_saved_logits.shape}, "
1127 f"loss=N/A (no token labels)\n"
1128 )
1129 else:
1130 hf_tokens = bridge_unprocessed.to_tokens(test_text)
1131 is_enc_dec = is_encoder_decoder_model(
1132 model_name, trust_remote_code=trust_remote_code
1133 )
1134 with torch.no_grad():
1135 if is_enc_dec: 1135 ↛ 1136line 1135 didn't jump to line 1136 because the condition on line 1135 was never true
1136 decoder_start_id = getattr(
1137 getattr(hf_model, "config", None),
1138 "decoder_start_token_id",
1139 0,
1140 )
1141 dec_ids = torch.tensor([[decoder_start_id]]).to(hf_tokens.device)
1142 hf_out = hf_model(hf_tokens, decoder_input_ids=dec_ids)
1143 else:
1144 hf_out = _hf_forward_with_mask_fallback(hf_model, hf_tokens)
1145 hf_saved_logits = hf_out.logits.detach().cpu().clone()
1147 # Compute causal LM loss (shift logits and labels)
1148 if not is_enc_dec and hf_saved_logits.shape[1] > 1: 1148 ↛ 1157line 1148 didn't jump to line 1157
1149 shift_logits = hf_out.logits[..., :-1, :].contiguous()
1150 shift_labels = hf_tokens[..., 1:].contiguous()
1151 loss_fn = torch.nn.CrossEntropyLoss()
1152 hf_saved_loss = loss_fn(
1153 shift_logits.view(-1, shift_logits.size(-1)),
1154 shift_labels.view(-1),
1155 ).item()
1157 if verbose: 1157 ↛ 1158line 1157 didn't jump to line 1158 because the condition on line 1157 was never true
1158 loss_str = f"{hf_saved_loss:.4f}" if hf_saved_loss is not None else "N/A"
1159 print(
1160 f"✓ Captured HF logits {hf_saved_logits.shape}, " f"loss={loss_str}\n"
1161 )
1162 del hf_tokens
1163 except Exception as e:
1164 if verbose:
1165 print(f"⚠ Could not capture HF reference outputs: {e}\n")
1167 # Delete HF model immediately after component benchmark + logit capture.
1168 # From here on, Phase 1 runs at 1.0x using saved HF tensors.
1169 cleanup_model(hf_model, "HuggingFace model")
1170 hf_model = None
1171 else:
1172 if verbose:
1173 print("⏭️ Skipped (no HF reference model available)\n")
1175 # Forward pass benchmarks
1176 if verbose: 1176 ↛ 1177line 1176 didn't jump to line 1177 because the condition on line 1176 was never true
1177 print("2. Forward Pass Benchmarks")
1179 # Widen tolerance for reduced-precision benchmarking — MPS bfloat16
1180 # matmul non-determinism can exceed the float32 default of 1e-3
1181 p1_atol = 1e-3 if dtype == torch.float32 else 5e-3
1183 # For audio/vision models, reuse the input from HF reference capture
1184 _p1_input: Union[str, torch.Tensor] = test_text
1185 if _test_modality_input is not None: 1185 ↛ 1186line 1185 didn't jump to line 1186 because the condition on line 1185 was never true
1186 _p1_input = _test_modality_input
1188 if hf_saved_logits is not None: 1188 ↛ 1203line 1188 didn't jump to line 1203 because the condition on line 1188 was always true
1189 # Full mode: use pre-captured HF logits (bridge only, 1.0x)
1190 try:
1191 add_result(
1192 benchmark_forward_pass(
1193 bridge_unprocessed,
1194 _p1_input,
1195 reference_logits=hf_saved_logits.to(device),
1196 atol=p1_atol,
1197 )
1198 )
1199 except Exception as e:
1200 if verbose:
1201 print(f"✗ Forward pass benchmark failed: {e}\n")
1202 else:
1203 try:
1204 add_result(benchmark_forward_pass(bridge_unprocessed, _p1_input, atol=p1_atol))
1205 except Exception as e:
1206 if verbose:
1207 print(f"✗ Forward pass benchmark failed: {e}\n")
1209 # Capture Phase 1 reference for Phase 3 equivalence comparison.
1210 # Skip for audio/vision models (Phase 3 won't run — no HookedTransformer
1211 # support — and the capture below feeds text, which they cannot accept).
1212 # When dtype==float32 (default) and the model natively uses reduced
1213 # precision, upcast for maximum accuracy. When the user explicitly
1214 # requested a non-float32 dtype, run the reference pass in that dtype
1215 # so the entire pipeline honours the requested precision.
1216 if bridge_unprocessed is not None and not _is_audio and not _is_visual: 1216 ↛ 1259line 1216 didn't jump to line 1259 because the condition on line 1216 was always true
1217 try:
1218 original_dtype = bridge_unprocessed.cfg.dtype
1219 needs_upcast = dtype == torch.float32 and original_dtype not in (
1220 torch.float32,
1221 torch.float64,
1222 )
1223 # Snapshot registered buffers before the round-trip. HF's
1224 # RotaryEmbedding recomputes inv_freq during the float32 forward
1225 # pass, and the downcast back to bfloat16 would produce different
1226 # values than the original, corrupting the model for Phase 2.
1227 saved_buffers = {}
1228 if needs_upcast: 1228 ↛ 1229line 1228 didn't jump to line 1229 because the condition on line 1228 was never true
1229 for bname, buf in bridge_unprocessed.named_buffers():
1230 saved_buffers[bname] = buf.data.clone()
1231 bridge_unprocessed.to(torch.float32)
1232 with torch.no_grad():
1233 bridge_logits = bridge_unprocessed(test_text, return_type="logits")
1234 phase1_reference.hf_logits = bridge_logits.detach().cpu().clone()
1235 bridge_loss = _compute_self_target_loss(bridge_unprocessed, test_text)
1236 phase1_reference.hf_loss = bridge_loss.item()
1237 phase1_reference.test_text = test_text
1238 if needs_upcast: 1238 ↛ 1239line 1238 didn't jump to line 1239 because the condition on line 1238 was never true
1239 bridge_unprocessed.to(original_dtype)
1240 # Restore buffers that were corrupted by the round-trip.
1241 # Use direct assignment (not copy_) to preserve original dtype.
1242 # HF's RotaryEmbedding keeps inv_freq in float32 even when the
1243 # model is bfloat16. After to(bfloat16), the buffer becomes
1244 # bfloat16, and copy_() would truncate the float32 saved values.
1245 for bname, buf in bridge_unprocessed.named_buffers():
1246 if bname in saved_buffers:
1247 buf.data = saved_buffers[bname]
1248 if verbose: 1248 ↛ 1249line 1248 didn't jump to line 1249 because the condition on line 1248 was never true
1249 dtype_note = " (upcast to float32)" if needs_upcast else ""
1250 print(
1251 f"✓ Saved Phase 1 reference data "
1252 f"(logits: {phase1_reference.hf_logits.shape}){dtype_note}"
1253 )
1254 except Exception as e:
1255 if verbose:
1256 print(f"⚠ Could not save Phase 1 reference data: {e}")
1258 # Free saved HF tensors now that Phase 1 is done
1259 del hf_saved_logits, hf_saved_loss
1261 # Save bridge_dtype before potential cleanup (needed for Phase 3)
1262 saved_bridge_dtype = bridge_dtype
1264 # Clean up HF model if still alive (e.g., Phase 1 was skipped)
1265 if hf_model is not None: 1265 ↛ 1266line 1265 didn't jump to line 1266 because the condition on line 1265 was never true
1266 cleanup_model(hf_model, "HuggingFace model")
1267 hf_model = None
1269 # ========================================================================
1270 # PHASE 2: Bridge (unprocessed) + HookedTransformer (unprocessed)
1271 # ========================================================================
1272 current_phase[0] = 2
1274 # OPTIMIZATION: Run generation benchmarks first (only bridge in memory)
1275 # Then cleanup bridge before loading HT to reduce peak memory
1276 if should_run_phase(2) and bridge_unprocessed: 1276 ↛ 1364line 1276 didn't jump to line 1364 because the condition on line 1276 was always true
1277 if verbose: 1277 ↛ 1278line 1277 didn't jump to line 1278 because the condition on line 1277 was never true
1278 print(f"\n{'='*80}")
1279 print("PHASE 2: TransformerBridge (unprocessed) + HookedTransformer (unprocessed)")
1280 print(f"{'='*80}\n")
1281 if verbose: 1281 ↛ 1282line 1281 didn't jump to line 1282 because the condition on line 1281 was never true
1282 print("Running Phase 2 benchmarks...\n")
1284 # Generation benchmarks (unprocessed only) - RUN FIRST
1285 # Skip for encoder-decoder and audio models (no text generation capability)
1286 # Diffusion LMs generate through a native sampler; the benchmarks route
1287 # to it, so only architectures with neither path are skipped.
1288 _adapter = getattr(bridge_unprocessed, "adapter", None)
1289 _no_generate = not getattr(_adapter, "supports_generation", True) and not getattr(
1290 _adapter, "native_sampler", None
1291 )
1292 _skip_generation = (
1293 is_encoder_decoder_model(model_name)
1294 or getattr(bridge_unprocessed.cfg, "is_audio_model", False)
1295 or _no_generate
1296 )
1297 _skip_reason = (
1298 "Skipped (model does not support generation)"
1299 if _no_generate
1300 else "Skipped (encoder-decoder model)"
1301 )
1302 if verbose: 1302 ↛ 1303line 1302 didn't jump to line 1303 because the condition on line 1302 was never true
1303 print("1. Generation Benchmarks (unprocessed)")
1304 if _skip_generation: 1304 ↛ 1340line 1304 didn't jump to line 1340 because the condition on line 1304 was always true
1305 if verbose: 1305 ↛ 1306line 1305 didn't jump to line 1306 because the condition on line 1305 was never true
1306 print(f"⏭️ {_skip_reason}\n")
1307 add_result(
1308 BenchmarkResult(
1309 name="generation",
1310 severity=BenchmarkSeverity.INFO,
1311 passed=True,
1312 message=_skip_reason,
1313 )
1314 )
1315 add_result(
1316 BenchmarkResult(
1317 name="generation_with_kv_cache",
1318 severity=BenchmarkSeverity.INFO,
1319 passed=True,
1320 message=_skip_reason,
1321 )
1322 )
1323 add_result(
1324 BenchmarkResult(
1325 name="multiple_generation_calls",
1326 severity=BenchmarkSeverity.INFO,
1327 passed=True,
1328 message=_skip_reason,
1329 )
1330 )
1331 add_result(
1332 BenchmarkResult(
1333 name="text_quality",
1334 severity=BenchmarkSeverity.INFO,
1335 passed=True,
1336 message=_skip_reason,
1337 )
1338 )
1339 else:
1340 try:
1341 add_result(benchmark_generation(bridge_unprocessed, test_text, max_new_tokens=10))
1342 add_result(
1343 benchmark_generation_with_kv_cache(
1344 bridge_unprocessed, test_text, max_new_tokens=10
1345 )
1346 )
1347 add_result(
1348 benchmark_multiple_generation_calls(
1349 bridge_unprocessed,
1350 test_prompts=[
1351 "The quick brown fox",
1352 "Hello world",
1353 "Machine learning is",
1354 ],
1355 max_new_tokens=5,
1356 )
1357 )
1358 gc.collect() # Force cleanup after generation benchmarks
1359 except Exception as e:
1360 if verbose:
1361 print(f"✗ Generation benchmark failed: {e}\n")
1363 # Match bridge's default_prepend_bos setting in HookedTransformer.
1364 ht_prepend_bos = None
1365 if bridge_unprocessed is not None and hasattr(bridge_unprocessed, "cfg"): 1365 ↛ 1373line 1365 didn't jump to line 1373 because the condition on line 1365 was always true
1366 bridge_bos = getattr(bridge_unprocessed.cfg, "default_prepend_bos", None)
1367 if bridge_bos is not None: 1367 ↛ 1373line 1367 didn't jump to line 1373 because the condition on line 1367 was always true
1368 ht_prepend_bos = bridge_bos
1370 # HookedTransformer is a causal decoder: loading a masked LM into it runs a
1371 # bidirectional model under a causal mask, so it can never be a valid
1372 # reference — numerical comparisons fall back to the Phase 1 HF logits.
1373 if use_ht_reference and is_masked_lm_model(model_name, trust_remote_code=trust_remote_code): 1373 ↛ 1379line 1373 didn't jump to line 1379 because the condition on line 1373 was always true
1374 if verbose: 1374 ↛ 1375line 1374 didn't jump to line 1375 because the condition on line 1374 was never true
1375 print("Skipping HookedTransformer reference: masked-LM is not representable causally.")
1376 use_ht_reference = False
1378 # Load HookedTransformer for comparison (after generation benchmarks)
1379 ht_model_unprocessed = None
1380 if should_run_phase(2) and use_ht_reference: 1380 ↛ 1381line 1380 didn't jump to line 1381 because the condition on line 1380 was never true
1381 try:
1382 if verbose:
1383 print("Loading HookedTransformer (unprocessed) for comparison...")
1384 ht_model_unprocessed = HookedTransformer.from_pretrained(
1385 model_name,
1386 device=device,
1387 dtype=bridge_dtype,
1388 fold_ln=False,
1389 center_writing_weights=False,
1390 center_unembed=False,
1391 fold_value_biases=False,
1392 refactor_factored_attn_matrices=False,
1393 default_prepend_bos=ht_prepend_bos,
1394 )
1395 if verbose:
1396 print("✓ HookedTransformer loaded (unprocessed)\n")
1397 except Exception as e:
1398 if verbose:
1399 print(f"✗ Could not load unprocessed HookedTransformer: {str(e)}\n")
1401 # Run Phase 2 comparison benchmarks using unified function
1402 if should_run_phase(2) and bridge_unprocessed: 1402 ↛ 1440line 1402 didn't jump to line 1440 because the condition on line 1402 was always true
1403 if verbose: 1403 ↛ 1404line 1403 didn't jump to line 1404 because the condition on line 1403 was never true
1404 print("2. Running Unprocessed Model Comparison Benchmarks\n")
1406 # When dtype==float32 (default) but the model natively loaded in
1407 # reduced precision, upcast for maximum benchmark accuracy. When the
1408 # user explicitly requested bfloat16/float16, honour that — run the
1409 # entire comparison in the requested precision.
1410 phase2_restore_dtype = None
1411 if dtype == torch.float32 and bridge_dtype in (torch.bfloat16, torch.float16): 1411 ↛ 1412line 1411 didn't jump to line 1412 because the condition on line 1411 was never true
1412 try:
1413 bridge_unprocessed.to(torch.float32)
1414 if ht_model_unprocessed is not None:
1415 ht_model_unprocessed.to(torch.float32)
1416 phase2_restore_dtype = bridge_dtype
1417 if verbose:
1418 print(f" (upcast from {bridge_dtype} to float32 for comparison)\n")
1419 except Exception:
1420 phase2_restore_dtype = None # Upcast failed; proceed as-is
1422 phase2_results = run_comparison_benchmarks(
1423 bridge_model=bridge_unprocessed,
1424 reference_model=ht_model_unprocessed,
1425 test_text=test_text,
1426 phase_name="Phase 2",
1427 is_processed=False, # Unprocessed mode - skip weight processing tests
1428 verbose=verbose,
1429 restore_dtype_after_equivalence=phase2_restore_dtype,
1430 )
1431 # Tag all phase 2 results with phase number
1432 for result in phase2_results:
1433 if result.phase is None: 1433 ↛ 1432line 1433 didn't jump to line 1432 because the condition on line 1433 was always true
1434 result.phase = 2
1435 results.extend(phase2_results)
1437 # Generation benchmarks already run above (before loading HT)
1439 # Clean up unprocessed HT model - no longer needed
1440 if ht_model_unprocessed is not None: 1440 ↛ 1441line 1440 didn't jump to line 1441 because the condition on line 1440 was never true
1441 cleanup_model(ht_model_unprocessed, "HookedTransformer (unprocessed)")
1442 ht_model_unprocessed = None
1443 # bridge_unprocessed is kept alive for Phase 3 and Phase 4 — reusing the
1444 # same instance avoids non-deterministic loading in some architectures
1445 # (e.g., OpenELM).
1447 # ========================================================================
1448 # PHASE 4: Text Quality (profile prompts, judge perplexity-ratio scoring)
1449 # Runs before Phase 3 so it can reuse bridge_unprocessed (Phase 3
1450 # destructively processes the weights, consuming the bridge).
1451 # ========================================================================
1452 current_phase[0] = 4
1454 if ( 1454 ↛ 1467line 1454 didn't jump to line 1467 because the condition on line 1454 was never true
1455 should_run_phase(4)
1456 and bridge_unprocessed is not None
1457 # applicable_phases and supports_generation are independent switches;
1458 # without this check a disagreement surfaces as an ERROR, not a skip.
1459 # Native-sampler architectures generate too, just not autoregressively.
1460 and (
1461 getattr(bridge_unprocessed.adapter, "supports_generation", True)
1462 or getattr(bridge_unprocessed.adapter, "native_sampler", None) is not None
1463 )
1464 and not is_masked_lm_model(model_name, trust_remote_code=trust_remote_code)
1465 and not is_audio_model(model_name, trust_remote_code=trust_remote_code)
1466 ):
1467 if prompt_profile is None:
1468 from transformer_lens.benchmarks.text_quality_profiles import (
1469 resolve_profile,
1470 )
1471 from transformer_lens.tools.model_registry.registry_io import (
1472 registry_prompt_profile,
1473 )
1475 config = getattr(bridge_unprocessed, "original_model", None)
1476 archs = getattr(getattr(config, "config", None), "architectures", None) or []
1477 prompt_profile = str(
1478 resolve_profile(
1479 model_name, archs[0] if archs else None, registry_prompt_profile(model_name)
1480 )
1481 )
1483 if verbose:
1484 print(f"\n{'='*80}")
1485 print(f"PHASE 2.5: Text Quality (profile {prompt_profile}, judge ratio scoring)")
1486 print(f"{'='*80}\n")
1488 try:
1489 text_quality_result = benchmark_text_quality(
1490 bridge_unprocessed,
1491 prompt_profile,
1492 judge_model=judge_model,
1493 judge_tokenizer=judge_tokenizer,
1494 model_name=model_name,
1495 )
1496 text_quality_result.phase = 4
1497 add_result(text_quality_result)
1498 except Exception as e:
1499 if verbose:
1500 print(f"✗ Text quality benchmark failed: {e}\n")
1502 # ========================================================================
1503 # Phase 7: Multimodal Tests (only for multimodal models)
1504 # Runs before Phase 3 so we can reuse bridge_unprocessed before cleanup.
1505 # ========================================================================
1506 if ( 1506 ↛ 1511line 1506 didn't jump to line 1511 because the condition on line 1506 was never true
1507 bridge_unprocessed is not None
1508 and getattr(bridge_unprocessed.cfg, "is_multimodal", False)
1509 and should_run_phase(7)
1510 ):
1511 current_phase[0] = 7
1512 if verbose:
1513 print("\n" + "=" * 80)
1514 print("PHASE 7: MULTIMODAL TESTS")
1515 print("=" * 80)
1516 print("Testing multimodal forward pass, generation, and caching with images.")
1517 print("=" * 80 + "\n")
1519 try:
1520 from transformer_lens.benchmarks.multimodal import (
1521 benchmark_multimodal_cache,
1522 benchmark_multimodal_forward,
1523 benchmark_multimodal_generation,
1524 )
1526 mm_results = [
1527 benchmark_multimodal_forward(bridge_unprocessed, test_text=test_text),
1528 benchmark_multimodal_generation(bridge_unprocessed, test_text=test_text),
1529 benchmark_multimodal_cache(bridge_unprocessed, test_text=test_text),
1530 ]
1531 for result in mm_results:
1532 result.phase = 7
1533 results.append(result)
1534 if verbose:
1535 print(result)
1537 if verbose:
1538 print("\n" + "=" * 80)
1539 print("PHASE 7 COMPLETE")
1540 print("=" * 80)
1542 except Exception as e:
1543 if verbose:
1544 print(f"\n⚠ Multimodal tests failed: {e}\n")
1545 results.append(
1546 BenchmarkResult(
1547 name="multimodal_suite",
1548 passed=False,
1549 severity=BenchmarkSeverity.ERROR,
1550 message=f"Failed to run multimodal tests: {str(e)}",
1551 details={"error": str(e)},
1552 phase=7,
1553 )
1554 )
1556 # ========================================================================
1557 # Phase 8: Audio Tests (only for audio encoder models)
1558 # Runs before Phase 3 so we can reuse bridge_unprocessed before cleanup.
1559 # ========================================================================
1560 if ( 1560 ↛ 1565line 1560 didn't jump to line 1565 because the condition on line 1560 was never true
1561 bridge_unprocessed is not None
1562 and getattr(bridge_unprocessed.cfg, "is_audio_model", False)
1563 and should_run_phase(8)
1564 ):
1565 current_phase[0] = 8
1566 if verbose:
1567 print("\n" + "=" * 80)
1568 print("PHASE 8: AUDIO TESTS")
1569 print("=" * 80)
1570 print("Testing audio forward pass, caching, representation stability, and features.")
1571 print("=" * 80 + "\n")
1573 try:
1574 from transformer_lens.benchmarks.audio import run_audio_benchmarks
1576 audio_results = run_audio_benchmarks(
1577 bridge_unprocessed,
1578 test_audio=_test_modality_input,
1579 verbose=verbose,
1580 )
1581 for result in audio_results:
1582 result.phase = 8
1583 results.append(result)
1584 if verbose:
1585 print(result)
1587 if verbose:
1588 print("\n" + "=" * 80)
1589 print("PHASE 8 COMPLETE")
1590 print("=" * 80)
1592 except Exception as e:
1593 if verbose:
1594 print(f"\n⚠ Audio tests failed: {e}\n")
1595 results.append(
1596 BenchmarkResult(
1597 name="audio_suite",
1598 passed=False,
1599 severity=BenchmarkSeverity.ERROR,
1600 message=f"Failed to run audio tests: {str(e)}",
1601 details={"error": str(e)},
1602 phase=8,
1603 )
1604 )
1606 # ========================================================================
1607 # PHASE 8 (audio-text): audio-conditioned forward for audio decoders
1608 # (Qwen2Audio etc.) — is_multimodal with an audio processor, not an encoder.
1609 # Image Phase 7 feeds pixel_values and encoder Phase 8 feeds a raw waveform;
1610 # neither exercises these models' processed-feature audio path.
1611 # ========================================================================
1612 _audio_text = (
1613 bridge_unprocessed is not None
1614 and getattr(bridge_unprocessed.cfg, "is_multimodal", False)
1615 and not getattr(bridge_unprocessed.cfg, "is_audio_model", False)
1616 and getattr(getattr(bridge_unprocessed, "processor", None), "audio_token", None) is not None
1617 )
1618 if _audio_text and should_run_phase(8): 1618 ↛ 1619line 1618 didn't jump to line 1619 because the condition on line 1618 was never true
1619 current_phase[0] = 8
1620 if verbose:
1621 print("\n" + "=" * 80 + "\nPHASE 8: AUDIO-TEXT FORWARD\n" + "=" * 80 + "\n")
1622 from transformer_lens.benchmarks.audio import benchmark_audio_text_forward
1624 result = benchmark_audio_text_forward(bridge_unprocessed)
1625 result.phase = 8
1626 add_result(result)
1628 # ========================================================================
1629 # Phase 9: Vision Tests (only for vision encoder models — ViT/DeiT, not
1630 # vision+text multimodal models, which Phase 7 covers)
1631 # Runs before Phase 3 so we can reuse bridge_unprocessed before cleanup.
1632 # ========================================================================
1633 if ( 1633 ↛ 1639line 1633 didn't jump to line 1639 because the condition on line 1633 was never true
1634 bridge_unprocessed is not None
1635 and getattr(bridge_unprocessed.cfg, "is_visual_model", False)
1636 and not getattr(bridge_unprocessed.cfg, "is_multimodal", False)
1637 and should_run_phase(9)
1638 ):
1639 current_phase[0] = 9
1640 if verbose:
1641 print("\n" + "=" * 80)
1642 print("PHASE 9: VISION TESTS")
1643 print("=" * 80)
1644 print("Testing pixel forward pass, caching, representation stability, and decoding.")
1645 print("=" * 80 + "\n")
1647 try:
1648 from transformer_lens.benchmarks.vision import run_vision_benchmarks
1650 vision_results = run_vision_benchmarks(
1651 bridge_unprocessed,
1652 test_pixels=_test_modality_input,
1653 verbose=verbose,
1654 )
1655 for result in vision_results:
1656 result.phase = 9
1657 results.append(result)
1658 if verbose:
1659 print(result)
1661 if verbose:
1662 print("\n" + "=" * 80)
1663 print("PHASE 9 COMPLETE")
1664 print("=" * 80)
1666 except Exception as e:
1667 if verbose:
1668 print(f"\n⚠ Vision tests failed: {e}\n")
1669 results.append(
1670 BenchmarkResult(
1671 name="vision_suite",
1672 passed=False,
1673 severity=BenchmarkSeverity.ERROR,
1674 message=f"Failed to run vision tests: {str(e)}",
1675 details={"error": str(e)},
1676 phase=9,
1677 )
1678 )
1680 # ========================================================================
1681 # PHASE 3: Bridge (processed) + HookedTransformer (processed)
1682 # ========================================================================
1683 current_phase[0] = 3
1685 def _cleanup_bridge_unprocessed():
1686 """Clean up the kept-alive bridge_unprocessed if Phase 3 is skipped."""
1687 nonlocal bridge_unprocessed
1688 if bridge_unprocessed is not None: 1688 ↛ exitline 1688 didn't return from function '_cleanup_bridge_unprocessed' because the condition on line 1688 was always true
1689 cleanup_model(bridge_unprocessed, "TransformerBridge (unprocessed)")
1690 bridge_unprocessed = None
1692 _skip_phase3 = False
1693 if not enable_compatibility_mode: 1693 ↛ 1698line 1693 didn't jump to line 1698 because the condition on line 1693 was always true
1694 _cleanup_bridge_unprocessed()
1695 _skip_phase3 = True
1696 if verbose: 1696 ↛ 1697line 1696 didn't jump to line 1697 because the condition on line 1696 was never true
1697 print("\n⚠ Compatibility mode disabled - skipping Phase 3\n")
1698 elif not should_run_phase(3):
1699 _cleanup_bridge_unprocessed()
1700 _skip_phase3 = True
1701 if verbose:
1702 print("\n⚠ Phase 3 skipped (excluded by phases filter or adapter applicable_phases)\n")
1703 elif is_encoder_decoder_model(model_name):
1704 _cleanup_bridge_unprocessed()
1705 _skip_phase3 = True
1706 if verbose:
1707 print("\n⚠ Phase 3 skipped (encoder-decoder model - weight processing not supported)\n")
1709 bridge_processed = None
1710 ht_model_processed = None
1712 if not _skip_phase3: 1712 ↛ 1713line 1712 didn't jump to line 1713 because the condition on line 1712 was never true
1713 if verbose:
1714 print(f"\n{'='*80}")
1715 print("PHASE 3: TransformerBridge (processed) + HookedTransformer (processed)")
1716 print(f"{'='*80}\n")
1718 if not _skip_phase3: 1718 ↛ 1724line 1718 didn't jump to line 1724 because the condition on line 1718 was never true
1719 # Reuse the Phase 1 bridge instance and process weights in-place.
1720 # When dtype==float32 (default) and the model natively uses reduced
1721 # precision, upcast before processing to avoid bf16 quantization
1722 # round-trips. When the user explicitly requested bfloat16/float16,
1723 # process weights in the requested precision — no upcast.
1724 phase3_native_dtype = None # Set if we upcast; used to restore later
1725 if bridge_unprocessed is not None:
1726 try:
1727 if verbose:
1728 print("Processing weights on existing bridge (reusing Phase 1 instance)...")
1729 bridge_processed = bridge_unprocessed
1730 bridge_unprocessed = None # Transfer ownership
1731 phase3_native_dtype = bridge_processed.cfg.dtype
1732 if dtype == torch.float32 and phase3_native_dtype not in (
1733 torch.float32,
1734 torch.float64,
1735 ):
1736 bridge_processed.to(torch.float32)
1737 if verbose:
1738 print(f" (upcast from {phase3_native_dtype} to float32 before processing)")
1739 else:
1740 phase3_native_dtype = None # No restore needed
1741 bridge_processed.enable_compatibility_mode(disable_warnings=True)
1742 if verbose:
1743 print("✓ TransformerBridge compatibility mode enabled (processed)\n")
1744 except Exception as e:
1745 import traceback
1747 error_trace = traceback.format_exc()
1748 add_result(
1749 BenchmarkResult(
1750 name="process_bridge_weights",
1751 severity=BenchmarkSeverity.ERROR,
1752 message=f"Failed to process bridge weights: {str(e)}",
1753 passed=False,
1754 details={"error": str(e), "traceback": error_trace},
1755 )
1756 )
1757 if verbose:
1758 print(f"✗ Failed to process bridge weights: {str(e)}")
1759 print(f"\nStack trace:\n{error_trace}")
1760 else:
1761 # Fallback: load a fresh bridge if Phase 1 bridge was not available
1762 try:
1763 if verbose:
1764 print("Loading TransformerBridge (processed)...")
1765 bridge_dtype = saved_bridge_dtype
1766 if verbose:
1767 print(f"Using dtype={bridge_dtype} from Phase 1")
1768 bridge_processed = TransformerBridge.boot_transformers(model_name, device=device, dtype=bridge_dtype, trust_remote_code=trust_remote_code) # type: ignore[attr-defined]
1769 bridge_processed.enable_compatibility_mode(disable_warnings=True)
1770 if verbose:
1771 print("✓ TransformerBridge compatibility mode enabled (processed)\n")
1772 except Exception as e:
1773 import traceback
1775 error_trace = traceback.format_exc()
1776 add_result(
1777 BenchmarkResult(
1778 name="load_bridge_processed",
1779 severity=BenchmarkSeverity.ERROR,
1780 message=f"Failed to load processed TransformerBridge: {str(e)}",
1781 passed=False,
1782 details={"error": str(e), "traceback": error_trace},
1783 )
1784 )
1785 if verbose:
1786 print(f"✗ Failed to load processed TransformerBridge: {str(e)}")
1787 print(f"\nStack trace:\n{error_trace}")
1789 if bridge_processed is None:
1790 # Add failure results for all Phase 3 tests
1791 phase3_tests = [
1792 "no_nan_inf",
1793 "weight_magnitudes",
1794 "layer_norm_folding",
1795 "attention_output_centering",
1796 "mlp_output_centering",
1797 "unembed_centering",
1798 "value_bias_folding",
1799 "weight_processing",
1800 "weight_sharing",
1801 "weight_modification",
1802 "logits_equivalence",
1803 "loss_equivalence",
1804 "hook_registry",
1805 "hook_functionality",
1806 "critical_forward_hooks",
1807 "forward_hooks",
1808 "run_with_cache",
1809 "activation_cache",
1810 "gradient_computation",
1811 "critical_backward_hooks",
1812 "backward_hooks",
1813 ]
1815 for test_name in phase3_tests:
1816 add_result(
1817 BenchmarkResult(
1818 name=test_name,
1819 severity=BenchmarkSeverity.ERROR,
1820 message=f"Skipped due to weight processing failure",
1821 passed=False,
1822 details={"reason": "bridge_processing_failed"},
1823 )
1824 )
1826 if verbose:
1827 print("\n" + format_results(results))
1829 # Load HT in the same dtype that was requested for the benchmark.
1830 # This ensures a fair comparison — both bridge and HT operate in
1831 # the same precision throughout.
1832 phase3_ht_dtype = dtype
1834 if use_ht_reference:
1835 try:
1836 if verbose:
1837 print("Loading HookedTransformer (processed)...")
1838 ht_model_processed = HookedTransformer.from_pretrained(
1839 model_name,
1840 device=device,
1841 dtype=phase3_ht_dtype,
1842 fold_ln=True,
1843 center_writing_weights=True,
1844 center_unembed=True,
1845 fold_value_biases=True,
1846 refactor_factored_attn_matrices=False,
1847 default_prepend_bos=ht_prepend_bos,
1848 )
1849 if verbose:
1850 print("✓ HookedTransformer loaded (processed)\n")
1851 except Exception as e:
1852 if verbose:
1853 print(f"✗ Could not load processed HookedTransformer: {str(e)}\n")
1855 # Run Phase 3 benchmarks using unified function
1856 if bridge_processed:
1857 if verbose:
1858 print("Running Phase 3 benchmarks...\n")
1860 # Phase 3 runs in the requested dtype end-to-end. Both bridge and HT
1861 # operate in the same precision — no dtype restoration needed.
1862 phase3_results = run_comparison_benchmarks(
1863 bridge_model=bridge_processed,
1864 reference_model=ht_model_processed,
1865 test_text=test_text,
1866 phase_name="Phase 3",
1867 is_processed=True, # Processed mode - include weight processing tests
1868 verbose=verbose,
1869 phase1_reference=phase1_reference, # Saved HF logits/loss for equivalence testing
1870 )
1871 # Tag all phase 3 results with phase number
1872 for result in phase3_results:
1873 if result.phase is None:
1874 result.phase = 3
1875 results.extend(phase3_results)
1877 # Clean up Phase 3 models
1878 if bridge_processed is not None:
1879 cleanup_model(bridge_processed, "TransformerBridge (processed)")
1880 bridge_processed = None
1881 if ht_model_processed is not None:
1882 cleanup_model(ht_model_processed, "HookedTransformer (processed)")
1883 ht_model_processed = None
1885 # ========================================================================
1886 # Phase 5/6: Granular Weight Processing Tests (Optional)
1887 # ========================================================================
1888 if test_weight_processing_individually and enable_compatibility_mode: 1888 ↛ 1889line 1888 didn't jump to line 1889 because the condition on line 1888 was never true
1889 if verbose:
1890 print("\n" + "=" * 80)
1891 print("PHASE 5/6: GRANULAR WEIGHT PROCESSING TESTS")
1892 print("=" * 80)
1893 print("Testing each weight processing flag individually and in combinations")
1894 print("to isolate which specific processing steps cause issues.")
1895 print("=" * 80 + "\n")
1897 try:
1898 from transformer_lens.benchmarks.granular_weight_processing import (
1899 run_granular_weight_processing_benchmarks,
1900 )
1902 granular_results = run_granular_weight_processing_benchmarks(
1903 model_name=model_name,
1904 device=device,
1905 test_text=test_text,
1906 verbose=verbose,
1907 )
1909 # Convert granular results to BenchmarkResult format and add to main results
1910 for config_name, config_results in granular_results.items():
1911 for result in config_results:
1912 # Prefix the name with the config for clarity
1913 result.name = f"granular_{config_name}_{result.name}"
1914 results.append(result)
1916 if verbose:
1917 print("\n" + "=" * 80)
1918 print("PHASE 5/6 COMPLETE")
1919 print("=" * 80)
1921 except Exception as e:
1922 if verbose:
1923 print(f"\n⚠ Granular weight processing tests failed: {e}\n")
1924 results.append(
1925 BenchmarkResult(
1926 name="granular_weight_processing_suite",
1927 passed=False,
1928 severity=BenchmarkSeverity.ERROR,
1929 message=f"Failed to run granular weight processing tests: {str(e)}",
1930 details={"error": str(e)},
1931 )
1932 )
1934 # Print summary (individual results already printed immediately)
1935 if verbose: 1935 ↛ 1936line 1935 didn't jump to line 1936 because the condition on line 1935 was never true
1936 print("\n" + "=" * 80)
1937 print("BENCHMARK SUMMARY")
1938 print("=" * 80)
1940 # Group results by phase
1941 results_by_phase: Dict[Union[int, str], List[BenchmarkResult]] = {}
1942 for r in results:
1943 phase = r.phase if r.phase is not None else "Other"
1944 if phase not in results_by_phase:
1945 results_by_phase[phase] = []
1946 results_by_phase[phase].append(r)
1948 # Print phase-by-phase summary
1949 for phase in sorted(
1950 results_by_phase.keys(), key=lambda x: x if isinstance(x, int) else 999
1951 ):
1952 phase_results = results_by_phase[phase]
1953 phase_name = f"Phase {phase}" if isinstance(phase, int) else phase
1955 phase_passed = sum(
1956 1 for r in phase_results if r.passed and r.severity != BenchmarkSeverity.SKIPPED
1957 )
1958 phase_failed = sum(
1959 1 for r in phase_results if not r.passed and r.severity != BenchmarkSeverity.SKIPPED
1960 )
1961 phase_skipped = sum(1 for r in phase_results if r.severity == BenchmarkSeverity.SKIPPED)
1962 phase_total = len(phase_results)
1963 phase_run = phase_total - phase_skipped
1965 print(f"\n{phase_name}: {phase_run} tests run")
1966 if phase_run > 0:
1967 print(f" Passed: {phase_passed}/{phase_run} ({phase_passed/phase_run*100:.1f}%)")
1968 print(f" Failed: {phase_failed}/{phase_run} ({phase_failed/phase_run*100:.1f}%)")
1969 if phase_skipped > 0:
1970 print(f" Skipped: {phase_skipped}")
1972 # Overall summary
1973 passed = sum(1 for r in results if r.passed and r.severity != BenchmarkSeverity.SKIPPED)
1974 failed = sum(1 for r in results if not r.passed and r.severity != BenchmarkSeverity.SKIPPED)
1975 skipped = sum(1 for r in results if r.severity == BenchmarkSeverity.SKIPPED)
1976 total = len(results)
1977 run_tests = total - skipped
1979 print(f"\nOverall:")
1980 print(f"Total: {total} tests")
1981 if skipped > 0:
1982 print(f"Run: {run_tests} tests")
1983 print(f"Skipped: {skipped} tests")
1984 if run_tests > 0:
1985 print(f"Passed: {passed}/{run_tests} ({passed/run_tests*100:.1f}%)")
1986 print(f"Failed: {failed}/{run_tests} ({failed/run_tests*100:.1f}%)")
1987 print("=" * 80)
1989 # Print memory summary
1990 if track_memory and memory_tracker is not None: 1990 ↛ 1991line 1990 didn't jump to line 1991 because the condition on line 1990 was never true
1991 final_memory = get_memory_mb()
1992 total_increase = final_memory - memory_tracker["initial"]
1994 if verbose:
1995 print("\n" + "=" * 80)
1996 print("MEMORY USAGE SUMMARY")
1997 print("=" * 80)
1998 print(f"Initial memory: {memory_tracker['initial']:>8.1f} MB")
1999 print(f"Final memory: {final_memory:>8.1f} MB")
2000 print(f"Net increase: {total_increase:>+8.1f} MB")
2002 if memory_tracker["checkpoints"]:
2003 print("\nCleanup operations:")
2004 for cp in memory_tracker["checkpoints"]:
2005 if cp.get("freed_mb", 0) > 0:
2006 print(
2007 f" {cp['label']:<40} freed {cp['freed_mb']:>7.1f} MB "
2008 f"(after: {cp['memory_mb']:.1f} MB)"
2009 )
2010 print("=" * 80)
2012 return results
2015def update_model_registry(
2016 model_name: str, results: List[BenchmarkResult], use_hf_reference: bool = False
2017) -> bool:
2018 """Update the model registry with benchmark results.
2020 Args:
2021 model_name: The model that was benchmarked
2022 results: List of benchmark results
2023 use_hf_reference: Whether the run numerically compared against an HF
2024 reference. Defaults to False so an unstated reference state records
2025 a passing run as PROVISIONAL, never VERIFIED.
2027 Returns:
2028 True if registry was updated successfully
2029 """
2030 from transformer_lens.tools.model_registry.registry_io import (
2031 STATUS_FAILED,
2032 STATUS_PROVISIONAL,
2033 add_verification_record,
2034 update_model_status,
2035 )
2037 # Threshold/note logic shared with verify_models so the two paths can't drift.
2038 from transformer_lens.tools.model_registry.verify_models import (
2039 _build_verified_note,
2040 _check_phase_scores,
2041 _extract_phase_scores,
2042 _extract_prompt_profile,
2043 _pass_status,
2044 _sanitize_note,
2045 )
2047 phase_scores = _extract_phase_scores(results)
2049 score_error = _check_phase_scores(phase_scores, results)
2050 if score_error:
2051 status = STATUS_FAILED
2052 note = score_error
2053 else:
2054 status = _pass_status(use_hf_reference)
2055 note = _build_verified_note(phase_scores, results)
2056 if status == STATUS_PROVISIONAL:
2057 note = f"Structural only (no HF reference): {note}"
2059 # Try to determine architecture
2060 architecture_id = "Unknown"
2061 try:
2062 from transformers import AutoConfig
2064 config = AutoConfig.from_pretrained(model_name, token=_hf_token())
2065 archs = getattr(config, "architectures", []) or []
2066 if archs: 2066 ↛ 2071line 2066 didn't jump to line 2071 because the condition on line 2066 was always true
2067 architecture_id = archs[0]
2068 except Exception:
2069 pass
2071 updated = update_model_status(
2072 model_id=model_name,
2073 arch_id=architecture_id,
2074 status=status,
2075 phase_scores=phase_scores,
2076 note=note,
2077 sanitize_fn=_sanitize_note,
2078 prompt_profile=_extract_prompt_profile(results),
2079 )
2081 # No history record for provisional runs — VerificationHistory.is_verified()
2082 # treats any record as verified, which would bypass the provisional gate.
2083 if status != STATUS_PROVISIONAL:
2084 add_verification_record(
2085 model_id=model_name,
2086 arch_id=architecture_id,
2087 notes=note,
2088 verified_by="main_benchmark",
2089 sanitize_fn=_sanitize_note,
2090 )
2092 label = {STATUS_FAILED: "FAILED", STATUS_PROVISIONAL: "PROVISIONAL"}.get(status, "VERIFIED")
2093 score_parts = ", ".join(f"P{p}={s}%" for p, s in sorted(phase_scores.items()))
2094 print(f"Updated registry for {model_name} ({label}): {score_parts or 'no phase results'}")
2095 return updated
2098def main():
2099 """Run benchmarks from command line."""
2100 import argparse
2102 parser = argparse.ArgumentParser(description="Run TransformerBridge benchmarks")
2103 parser.add_argument(
2104 "--model",
2105 type=str,
2106 default="gpt2",
2107 help="Model name to benchmark (default: gpt2)",
2108 )
2109 parser.add_argument(
2110 "--device",
2111 type=str,
2112 default="cpu",
2113 help="Device to run on (default: cpu)",
2114 )
2115 parser.add_argument(
2116 "--no-hf-reference",
2117 action="store_true",
2118 help="Disable HuggingFace reference comparison",
2119 )
2120 parser.add_argument(
2121 "--no-ht-reference",
2122 action="store_true",
2123 help="Disable HookedTransformer reference comparison",
2124 )
2125 parser.add_argument(
2126 "--no-compat",
2127 action="store_true",
2128 help="Disable compatibility mode",
2129 )
2130 parser.add_argument(
2131 "--quiet",
2132 action="store_true",
2133 help="Suppress verbose output",
2134 )
2135 parser.add_argument(
2136 "--update-registry",
2137 action="store_true",
2138 help="Update model registry with benchmark results (default: false)",
2139 )
2140 parser.add_argument(
2141 "--trust-remote-code",
2142 action="store_true",
2143 help="Trust remote code for custom architectures (e.g., OpenELM)",
2144 )
2145 args = parser.parse_args()
2147 results = run_benchmark_suite(
2148 model_name=args.model,
2149 device=args.device,
2150 use_hf_reference=not args.no_hf_reference,
2151 use_ht_reference=not args.no_ht_reference,
2152 enable_compatibility_mode=not args.no_compat,
2153 verbose=not args.quiet,
2154 trust_remote_code=args.trust_remote_code,
2155 )
2157 if args.update_registry:
2158 # Same requested-reference state verify_models feeds pass_status(): a
2159 # --no-hf-reference run can only mint PROVISIONAL, never VERIFIED.
2160 update_model_registry(args.model, results, use_hf_reference=not args.no_hf_reference)
2163if __name__ == "__main__": 2163 ↛ 2164line 2163 didn't jump to line 2164 because the condition on line 2163 was never true
2164 main()