Coverage for transformer_lens/benchmarks/main_benchmark.py: 37%
891 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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) - Runtime self-checks + HF logits/loss equivalence
7Phase 3: Bridge (processed) - Compatibility mode + HF logits/loss equivalence
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.benchmarks.activation_cache import (
28 benchmark_activation_cache,
29 benchmark_run_with_cache,
30)
31from transformer_lens.benchmarks.backward_gradients import (
32 benchmark_backward_hooks,
33 benchmark_critical_backward_hooks,
34 benchmark_gradient_computation,
35 needs_fp32_gradients,
36)
37from transformer_lens.benchmarks.component_benchmark import benchmark_all_components
38from transformer_lens.benchmarks.forward_pass import (
39 _compute_self_target_loss,
40 benchmark_forward_pass,
41 benchmark_loss_equivalence,
42)
43from transformer_lens.benchmarks.generation import (
44 benchmark_generation,
45 benchmark_generation_with_kv_cache,
46 benchmark_multiple_generation_calls,
47)
48from transformer_lens.benchmarks.hook_registration import (
49 benchmark_critical_forward_hooks,
50 benchmark_forward_hooks,
51 benchmark_gated_hooks_fire,
52 benchmark_hook_functionality,
53 benchmark_hook_registry,
54)
55from transformer_lens.benchmarks.text_quality import benchmark_text_quality
56from transformer_lens.benchmarks.utils import (
57 BenchmarkResult,
58 BenchmarkSeverity,
59 PhaseReferenceData,
60 build_modality_input,
61 compare_tensors,
62 format_results,
63)
64from transformer_lens.benchmarks.weight_processing import (
65 benchmark_attention_output_centering,
66 benchmark_layer_norm_folding,
67 benchmark_mlp_output_centering,
68 benchmark_no_nan_inf,
69 benchmark_unembed_centering,
70 benchmark_value_bias_folding,
71 benchmark_weight_magnitudes,
72 benchmark_weight_modification,
73)
74from transformer_lens.config import TransformerBridgeConfig
75from transformer_lens.factories.architecture_adapter_factory import (
76 ArchitectureAdapterFactory,
77)
78from transformer_lens.model_bridge import TransformerBridge
79from transformer_lens.tools.model_registry.registry_io import TEXT_PHASES
81# Architecture classification — single source of truth in utilities.architectures
82from transformer_lens.utilities.architectures import (
83 get_architectures_for_config,
84 is_audio_model,
85 is_encoder_decoder_model,
86 is_masked_lm_model,
87)
88from transformer_lens.utilities.hf_utils import get_hf_token as _hf_token
91def _adapter_applicable_phases(model_name: str, trust_remote_code: bool = False) -> list[int]:
92 """Text phases (1-4) the model's adapter declares applicable (default all)."""
93 from transformer_lens.factories.architecture_adapter_factory import (
94 SUPPORTED_ARCHITECTURES,
95 )
97 try:
98 config = AutoConfig.from_pretrained(
99 model_name, trust_remote_code=trust_remote_code, token=_hf_token()
100 )
101 for arch in get_architectures_for_config(config):
102 adapter_cls = SUPPORTED_ARCHITECTURES.get(arch)
103 if adapter_cls is not None:
104 return getattr(adapter_cls, "applicable_phases", list(TEXT_PHASES))
105 except Exception:
106 pass
107 return list(TEXT_PHASES)
110def _phase_enabled(
111 phase_num: int, phases: Optional[List[int]], applicable_phases: List[int]
112) -> bool:
113 """Phase gating shared by run_benchmark_suite's should_run_phase.
115 An adapter's ``applicable_phases`` declares which text phases (1-4) it covers.
116 Phases 7/8/9 are gated separately by ``is_multimodal``/``is_audio_model``/
117 ``is_visual_model`` at their call sites, so they are never filtered out here
118 (mirrors verify_models._phases_to_run).
119 """
120 if phases is not None and phase_num not in phases:
121 return False
122 return phase_num not in TEXT_PHASES or phase_num in applicable_phases
125def get_auto_model_class(model_name: str, trust_remote_code: bool = False):
126 """Delegates to the bridge's architecture detection for consistency."""
127 from transformer_lens.model_bridge.sources.transformers import (
128 determine_architecture_from_hf_config,
129 get_hf_model_class_for_architecture,
130 )
132 try:
133 config = AutoConfig.from_pretrained(
134 model_name, trust_remote_code=trust_remote_code, token=_hf_token()
135 )
136 architecture = determine_architecture_from_hf_config(config)
137 return get_hf_model_class_for_architecture(architecture)
138 except Exception:
139 return AutoModelForCausalLM
142def _fixup_custom_model(hf_model) -> None:
143 """Apply post-load fixups for models with custom code (e.g., OpenELM).
145 Recomputes non-persistent buffers (inv_freq, causal_mask) that may be
146 zeroed during HuggingFace's meta-device loading.
147 """
148 # OpenELM fixups
149 if hasattr(hf_model, "transformer") and hasattr(hf_model.transformer, "layers"): 149 ↛ 151line 149 didn't jump to line 151 because the condition on line 149 was never true
150 # Ensure use_cache is set (OpenELM custom config omits it)
151 if not hasattr(hf_model.config, "use_cache") or "use_cache" not in hf_model.config.__dict__:
152 hf_model.config.use_cache = False
154 # Fix 1: Always recompute causal_mask (non-persistent buffer).
155 # After meta→real materialization, the buffer may contain garbage values
156 # rather than clean zeros, so we always recompute.
157 if hasattr(hf_model.transformer, "causal_mask"):
158 cm = hf_model.transformer.causal_mask
159 if cm is not None and cm.numel() > 0:
160 seq_len = cm.shape[-1]
161 correct_mask = torch.triu(
162 torch.ones(seq_len, seq_len, dtype=cm.dtype, device=cm.device),
163 diagonal=1,
164 )
165 hf_model.transformer.causal_mask = correct_mask
167 # Fix 2: Always recompute RoPE inv_freq and sin/cos (non-persistent buffers).
168 rope_max = getattr(hf_model.config, "rope_max_length", None)
169 if rope_max is not None:
170 for layer in hf_model.transformer.layers:
171 if hasattr(layer, "attn") and hasattr(layer.attn, "pos_embedding"):
172 rope = layer.attn.pos_embedding
173 if hasattr(rope, "inv_freq"):
174 correct_inv_freq = 1.0 / (
175 rope.freq_constant
176 ** (
177 torch.arange(0, rope.model_dim, 2, dtype=torch.float32)
178 / rope.model_dim
179 )
180 )
181 rope.inv_freq = correct_inv_freq.to(rope.inv_freq.device)
182 # Force-recompute sin/cos
183 rope._cached_cos = None
184 rope._cached_sin = None
185 rope._compute_sin_cos_embeddings(rope_max)
187 # Create synthetic lm_head for weight-tied models (share_input_output_layers)
188 if getattr(hf_model, "lm_head", None) is None:
189 embed = hf_model.transformer.token_embeddings
190 lm_head = torch.nn.Linear(embed.embedding_dim, embed.num_embeddings, bias=False)
191 lm_head.weight = embed.weight
192 hf_model.lm_head = lm_head
194 # Rotary tables destroyed by meta-device loading, for ANY architecture: the
195 # reference needs the same repair the adapter applies to the bridge, or
196 # Phase 1 compares a correct model against a corrupt one. Only tables that
197 # fail a validity check are touched, so scaled RoPE is never clobbered.
198 from transformer_lens.model_bridge.buffer_restore import restore_rotary_inv_freq
200 restore_rotary_inv_freq(hf_model)
202 if type(hf_model).__name__ == "GiddForDiffusionLM": 202 ↛ 203line 202 didn't jump to line 203 because the condition on line 202 was never true
203 from transformer_lens.model_bridge.supported_architectures.gidd import (
204 restore_frequencies,
205 )
207 restore_frequencies(hf_model)
210def _hf_forward_with_mask_fallback(hf_model, tokens):
211 """Run an HF decoder forward, retrying with a 2D then 4D mask for models that
212 dereference ``attention_mask`` unconditionally (e.g. LLaDA2) -- else the Phase-1
213 capture raises, gets swallowed, and silently degrades to a shape-only check."""
214 try:
215 return hf_model(tokens)
216 except (AttributeError, ValueError):
217 b, s = tokens.shape[0], tokens.shape[-1]
218 for mask in (
219 torch.ones(b, s, dtype=torch.long, device=tokens.device),
220 torch.ones(b, 1, s, s, dtype=torch.long, device=tokens.device),
221 ):
222 try:
223 return hf_model(tokens, attention_mask=mask)
224 except (AttributeError, ValueError):
225 continue
226 raise
229def run_comparison_benchmarks(
230 bridge_model: TransformerBridge,
231 test_text: str,
232 phase_name: str,
233 is_processed: bool,
234 verbose: bool = True,
235 phase1_reference: Optional[PhaseReferenceData] = None,
236 restore_dtype_after_equivalence: Optional[torch.dtype] = None,
237) -> List[BenchmarkResult]:
238 """Run standardized runtime benchmarks on the bridge.
240 This function runs the same comprehensive test suite for both unprocessed (Phase 2)
241 and processed (Phase 3) modes: HF-anchored logits/loss equivalence (via the saved
242 Phase 1 reference) plus reference-free structural self-checks for hooks, cache,
243 and gradients.
245 Args:
246 bridge_model: TransformerBridge model to test
247 test_text: Input text for testing
248 phase_name: Name of the phase ("Phase 2" or "Phase 3") for logging
249 is_processed: Whether models have processed weights (for weight-specific tests)
250 verbose: Whether to print detailed results
251 phase1_reference: Optional saved Phase 1 HF reference data for equivalence testing
252 restore_dtype_after_equivalence: If set, downcast bridge_model to this dtype after
253 the equivalence comparison but before hook/cache/gradient tests. Used when the
254 bridge was upcast to float32 for precise equivalence testing.
256 Returns:
257 List of BenchmarkResult objects
258 """
259 results: List[BenchmarkResult] = []
261 def add_result(result: BenchmarkResult) -> None:
262 """Add a result and optionally print it immediately."""
263 results.append(result)
264 if verbose: 264 ↛ 265line 264 didn't jump to line 265 because the condition on line 264 was never true
265 result.print_immediate()
267 # ========================================================================
268 # 1. Weight Processing Benchmarks (only for processed mode)
269 # MOST BASIC: Check weights are valid before testing anything else
270 # ========================================================================
271 if is_processed: 271 ↛ 272line 271 didn't jump to line 272 because the condition on line 271 was never true
272 if verbose:
273 print("1. Weight Processing Benchmarks (Foundation)")
274 try:
275 # Critical weight validation tests (run first - most basic)
276 add_result(benchmark_no_nan_inf(bridge_model, test_text))
277 add_result(benchmark_weight_magnitudes(bridge_model, test_text))
279 # Detailed weight processing validation benchmarks (don't need reference model)
280 add_result(benchmark_layer_norm_folding(bridge_model, test_text))
281 add_result(benchmark_attention_output_centering(bridge_model, test_text))
282 add_result(benchmark_mlp_output_centering(bridge_model, test_text))
283 add_result(benchmark_unembed_centering(bridge_model, test_text))
284 add_result(benchmark_value_bias_folding(bridge_model, test_text))
286 # weight_modification doesn't need a reference model
287 add_result(benchmark_weight_modification(bridge_model, test_text))
288 gc.collect()
289 except Exception as e:
290 if verbose:
291 print(f"✗ Weight processing benchmark failed: {e}\n")
293 # ========================================================================
294 # 2. Model Equivalence Benchmarks (Forward Pass)
295 # Tests basic forward computation - depends on weights being correct
296 # ========================================================================
297 if verbose: 297 ↛ 298line 297 didn't jump to line 298 because the condition on line 297 was never true
298 print("2. Model Equivalence Benchmarks (Forward Pass)")
300 has_phase1_ref = phase1_reference is not None and phase1_reference.hf_logits is not None
302 if has_phase1_ref: 302 ↛ 357line 302 didn't jump to line 357 because the condition on line 302 was always true
303 # Compare the bridge against the saved Phase 1 HF reference.
304 # We use log_softmax because center_unembed shifts raw logits by a
305 # softmax-invariant constant. Both passes run in float32 (no bf16 round-trip).
306 try:
307 if verbose: 307 ↛ 308line 307 didn't jump to line 308 because the condition on line 307 was never true
308 print("Using saved Phase 1 bridge reference for equivalence comparison")
310 assert phase1_reference is not None
311 assert phase1_reference.hf_logits is not None
313 # Compare log_softmax (centering-invariant) instead of raw logits.
314 bridge_logits = bridge_model(test_text, return_type="logits")
315 ref_logits = phase1_reference.hf_logits.to(bridge_logits.device)
316 bridge_log_probs = torch.nn.functional.log_softmax(bridge_logits, dim=-1)
317 ref_log_probs = torch.nn.functional.log_softmax(ref_logits, dim=-1)
319 # Both passes in float32 — remaining error is float32 non-associativity
320 # in weight processing (~0.006 max_diff on 24-layer Qwen2).
321 logits_atol = 0.01
322 logits_rtol = 1e-4
323 loss_atol = 1e-3
325 add_result(
326 compare_tensors(
327 bridge_log_probs,
328 ref_log_probs,
329 atol=logits_atol,
330 rtol=logits_rtol,
331 name="logits_equivalence",
332 )
333 )
334 if phase1_reference.hf_loss is not None: 334 ↛ 344line 334 didn't jump to line 344 because the condition on line 334 was always true
335 add_result(
336 benchmark_loss_equivalence(
337 bridge_model,
338 test_text,
339 reference_loss=phase1_reference.hf_loss,
340 atol=loss_atol,
341 )
342 )
343 else:
344 add_result(
345 BenchmarkResult(
346 name="loss_equivalence",
347 severity=BenchmarkSeverity.SKIPPED,
348 message="Skipped (no Phase 1 loss reference available)",
349 passed=True,
350 )
351 )
352 gc.collect()
353 except Exception as e:
354 if verbose:
355 print(f"✗ Phase 1 reference comparison failed: {e}\n")
356 else:
357 if verbose:
358 print("⏭️ Skipped (no Phase 1 HF reference)\n")
359 for benchmark_name in ["logits_equivalence", "loss_equivalence"]:
360 add_result(
361 BenchmarkResult(
362 name=benchmark_name,
363 severity=BenchmarkSeverity.SKIPPED,
364 message="Skipped (no Phase 1 HF reference available)",
365 passed=True,
366 )
367 )
369 # Restore native dtype so remaining tests run in the model's real dtype.
370 if restore_dtype_after_equivalence is not None: 370 ↛ 371line 370 didn't jump to line 371 because the condition on line 370 was never true
371 try:
372 bridge_model.to(restore_dtype_after_equivalence)
373 if verbose:
374 print(f" (restored to {restore_dtype_after_equivalence} for remaining tests)\n")
375 except Exception as e:
376 if verbose:
377 print(f"⚠ Could not restore dtype: {e}\n")
379 # ========================================================================
380 # 3. Hook Registration Benchmarks
381 # Tests hooks exist and are registered - depends on model structure
382 # ========================================================================
383 if verbose: 383 ↛ 384line 383 didn't jump to line 384 because the condition on line 383 was never true
384 print("3. Hook Registration Benchmarks")
386 try:
387 add_result(benchmark_hook_registry(bridge_model))
388 gc.collect()
389 except Exception as e:
390 if verbose:
391 print(f"✗ Hook registry benchmark failed: {e}\n")
393 # ========================================================================
394 # 4. Forward Hook Functionality Benchmarks
395 # Tests hooks fire and produce correct values - depends on forward pass + hooks
396 # ========================================================================
397 if verbose: 397 ↛ 398line 397 didn't jump to line 398 because the condition on line 397 was never true
398 print("4. Forward Hook Functionality Benchmarks")
400 try:
401 add_result(benchmark_hook_functionality(bridge_model, test_text))
402 add_result(benchmark_critical_forward_hooks(bridge_model, test_text))
403 add_result(benchmark_forward_hooks(bridge_model, test_text))
404 add_result(benchmark_gated_hooks_fire(bridge_model, test_text))
405 # Reset hooks to prevent handle leaks
406 if hasattr(bridge_model, "reset_hooks"): 406 ↛ 408line 406 didn't jump to line 408 because the condition on line 406 was always true
407 bridge_model.reset_hooks()
408 gc.collect()
409 except Exception as e:
410 if verbose:
411 print(f"✗ Forward hook benchmark failed: {e}\n")
413 # ========================================================================
414 # 5. Activation Cache Benchmarks
415 # Tests caching mechanism - depends on forward pass + hooks working
416 # ========================================================================
417 if verbose: 417 ↛ 418line 417 didn't jump to line 418 because the condition on line 417 was never true
418 print("5. Activation Cache Benchmarks")
420 try:
421 add_result(benchmark_run_with_cache(bridge_model, test_text))
422 add_result(benchmark_activation_cache(bridge_model, test_text))
423 # Reset hooks to prevent handle leaks
424 if hasattr(bridge_model, "reset_hooks"): 424 ↛ 426line 424 didn't jump to line 426 because the condition on line 424 was always true
425 bridge_model.reset_hooks()
426 gc.collect()
427 except Exception as e:
428 if verbose:
429 print(f"✗ Activation cache benchmark failed: {e}\n")
431 # ========================================================================
432 # 6. Backward Gradient Benchmarks
433 # MOST COMPLEX: Tests gradients and backward hooks - depends on everything above
434 # ========================================================================
435 if verbose: 435 ↛ 436line 435 didn't jump to line 436 because the condition on line 435 was never true
436 print("6. Backward Gradient Benchmarks")
438 # Gradient comparisons are graded against fp32-calibrated thresholds
439 # (REL_L2_TOLERANCE): bf16's rounding floor alone is ~2e-3 rel_l2, inside the
440 # measured bug band, so reduced-precision gradients cannot be graded at all.
441 # Upcast for the gradient section on every device (MPS additionally lacks
442 # bf16 autograd), then restore below.
443 bridge_grad_dtype = bridge_model.cfg.dtype if hasattr(bridge_model, "cfg") else None
444 grad_fp32_upcast = needs_fp32_gradients(bridge_grad_dtype)
445 if grad_fp32_upcast: 445 ↛ 446line 445 didn't jump to line 446 because the condition on line 445 was never true
446 try:
447 bridge_model.to(torch.float32)
448 except Exception:
449 grad_fp32_upcast = False # Upcast failed; proceed as-is
451 try:
452 add_result(benchmark_gradient_computation(bridge_model, test_text))
453 add_result(benchmark_critical_backward_hooks(bridge_model, test_text))
454 add_result(benchmark_backward_hooks(bridge_model, test_text))
455 # Reset hooks to prevent handle leaks
456 if hasattr(bridge_model, "reset_hooks"): 456 ↛ 458line 456 didn't jump to line 458 because the condition on line 456 was always true
457 bridge_model.reset_hooks()
458 gc.collect()
459 except Exception as e:
460 if verbose:
461 print(f"✗ Gradient benchmark failed: {e}\n")
463 if grad_fp32_upcast and bridge_grad_dtype is not None: 463 ↛ 464line 463 didn't jump to line 464 because the condition on line 463 was never true
464 try:
465 bridge_model.to(bridge_grad_dtype)
466 except Exception:
467 pass
469 return results
472def run_benchmark_suite(
473 model_name: str,
474 device: str = "cpu",
475 dtype: torch.dtype = torch.float32,
476 test_text: Optional[str] = None,
477 use_hf_reference: bool = True,
478 enable_compatibility_mode: bool = True,
479 verbose: bool = True,
480 track_memory: bool = False,
481 phases: list[int] | None = None,
482 trust_remote_code: bool = False,
483 judge_model: PreTrainedModel | None = None,
484 judge_tokenizer: PreTrainedTokenizerBase | None = None,
485 prompt_profile: str | None = None,
486) -> List[BenchmarkResult]:
487 """Run comprehensive benchmark suite for TransformerBridge.
489 This function implements an optimized multi-phase approach to minimize model reloading:
490 Phase 1: HF + Bridge (unprocessed) - Compare against raw HuggingFace model
491 Phase 2: Bridge (unprocessed) - Runtime self-checks + HF logits/loss equivalence
492 Phase 3: Bridge (processed) - Compatibility mode + HF logits/loss equivalence
493 Phase 4: Text Quality - profile prompts scored by a pinned judge's perplexity ratio
495 Args:
496 model_name: Name of the model to benchmark (e.g., "gpt2")
497 device: Device to run on ("cpu" or "cuda")
498 dtype: Precision for model loading (default: torch.float32). Use
499 torch.bfloat16 to halve memory for larger models. Phase 2/3
500 comparisons automatically upcast to float32 for precision.
501 test_text: Optional test text (default: standard test prompt)
502 use_hf_reference: Whether to compare against HuggingFace model
503 enable_compatibility_mode: Whether to enable compatibility mode on bridge
504 verbose: Whether to print results to console
505 track_memory: Whether to track and report memory usage (requires psutil)
506 phases: Optional list of phase numbers to run (e.g., [1, 2, 3]). If None, runs all phases.
507 trust_remote_code: Whether to trust remote code for custom architectures.
508 judge_model: Optional pre-loaded Phase-4 judge. When provided with
509 judge_tokenizer, avoids reloading for each model in batch.
510 judge_tokenizer: Optional pre-loaded tokenizer for the Phase-4 judge.
511 prompt_profile: Optional Phase-4 prompt profile (e.g. "chat",
512 "task:translation@en-de"). Resolved from curation + the registry
513 when None.
515 Returns:
516 List of BenchmarkResult objects
517 """
518 if test_text is None: 518 ↛ 525line 518 didn't jump to line 525 because the condition on line 518 was always true
519 test_text = (
520 "Natural language processing tasks, such as question answering, "
521 "machine translation, reading comprehension, and summarization, "
522 "are typically approached with supervised learning."
523 )
525 results: List[BenchmarkResult] = []
527 # Memory tracking setup
528 memory_tracker = None
529 if track_memory: 529 ↛ 530line 529 didn't jump to line 530 because the condition on line 529 was never true
530 try:
531 import psutil
533 process = psutil.Process()
534 initial_memory = process.memory_info().rss / 1024 / 1024 # MB
536 def get_memory_mb():
537 return process.memory_info().rss / 1024 / 1024
539 memory_tracker = {"initial": initial_memory, "checkpoints": []}
540 if verbose:
541 print(f"Memory tracking enabled (initial: {initial_memory:.1f} MB)")
542 except ImportError:
543 if verbose:
544 print("⚠ psutil not available - memory tracking disabled")
545 track_memory = False
547 if verbose: 547 ↛ 548line 547 didn't jump to line 548 because the condition on line 547 was never true
548 print(f"\n{'='*80}")
549 print(f"Running TransformerBridge Benchmark Suite")
550 print(f"Model: {model_name}")
551 print(f"Device: {device}")
552 print(f"{'='*80}\n")
554 # Track current phase for result tagging
555 current_phase: List[Optional[int]] = [None] # Use list to allow modification in nested function
557 adapter_applicable = _adapter_applicable_phases(model_name, trust_remote_code)
559 def should_run_phase(phase_num: int) -> bool:
560 """Check if a phase should run based on the phases filter and adapter applicability."""
561 return _phase_enabled(phase_num, phases, adapter_applicable)
563 def add_result(result: BenchmarkResult) -> None:
564 """Add a result and optionally print it immediately."""
565 # Tag result with current phase
566 if current_phase[0] is not None and result.phase is None: 566 ↛ 568line 566 didn't jump to line 568 because the condition on line 566 was always true
567 result.phase = current_phase[0]
568 results.append(result)
569 if verbose: 569 ↛ 570line 569 didn't jump to line 570 because the condition on line 569 was never true
570 result.print_immediate()
572 def cleanup_tensors(*tensors) -> None:
573 """Free memory from tensors and caches."""
574 for tensor in tensors:
575 if tensor is not None:
576 # If it's an ActivationCache, clear all tensors
577 if hasattr(tensor, "cache_dict"):
578 for key in list(tensor.cache_dict.keys()):
579 val = tensor.cache_dict[key]
580 if val is not None and isinstance(val, torch.Tensor):
581 del val
582 tensor.cache_dict[key] = None
583 tensor.cache_dict.clear()
584 # If it's a regular tensor, just delete it
585 elif isinstance(tensor, torch.Tensor):
586 del tensor
587 # Force cleanup
588 gc.collect()
589 if device != "cpu" and torch.cuda.is_available():
590 torch.cuda.empty_cache()
591 if device == "mps" and hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
592 torch.mps.synchronize()
593 torch.mps.empty_cache()
595 def cleanup_model(model, model_name_str: str):
596 """Free up memory by deleting a model and forcing garbage collection."""
597 import gc
599 if verbose: 599 ↛ 600line 599 didn't jump to line 600 because the condition on line 599 was never true
600 print(f"Cleaning up {model_name_str}...")
602 # Track memory before cleanup
603 if track_memory and memory_tracker is not None: 603 ↛ 604line 603 didn't jump to line 604 because the condition on line 603 was never true
604 memory_before = get_memory_mb()
606 # Move model to CPU first to free GPU memory immediately
607 if device != "cpu" and hasattr(model, "cpu"): 607 ↛ 608line 607 didn't jump to line 608 because the condition on line 607 was never true
608 try:
609 model.cpu()
610 if torch.cuda.is_available():
611 torch.cuda.empty_cache()
612 if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"):
613 torch.mps.synchronize()
614 torch.mps.empty_cache()
615 except Exception:
616 pass
618 # Explicitly remove all hooks to prevent memory leaks
619 if hasattr(model, "modules"): 619 ↛ 655line 619 didn't jump to line 655 because the condition on line 619 was always true
620 try:
621 for module in model.modules():
622 # Clear PyTorch hooks
623 if hasattr(module, "_forward_hooks"): 623 ↛ 625line 623 didn't jump to line 625 because the condition on line 623 was always true
624 module._forward_hooks.clear()
625 if hasattr(module, "_backward_hooks"): 625 ↛ 627line 625 didn't jump to line 627 because the condition on line 625 was always true
626 module._backward_hooks.clear()
627 if hasattr(module, "_forward_pre_hooks"): 627 ↛ 629line 627 didn't jump to line 629 because the condition on line 627 was always true
628 module._forward_pre_hooks.clear()
629 if hasattr(module, "_backward_pre_hooks"): 629 ↛ 631line 629 didn't jump to line 631 because the condition on line 629 was always true
630 module._backward_pre_hooks.clear()
631 if hasattr(module, "_state_dict_hooks"): 631 ↛ 633line 631 didn't jump to line 633 because the condition on line 631 was always true
632 module._state_dict_hooks.clear()
633 if hasattr(module, "_state_dict_pre_hooks"): 633 ↛ 635line 633 didn't jump to line 635 because the condition on line 633 was always true
634 module._state_dict_pre_hooks.clear()
635 if hasattr(module, "_load_state_dict_pre_hooks"): 635 ↛ 637line 635 didn't jump to line 637 because the condition on line 635 was always true
636 module._load_state_dict_pre_hooks.clear()
637 if hasattr(module, "_load_state_dict_post_hooks"): 637 ↛ 641line 637 didn't jump to line 641 because the condition on line 637 was always true
638 module._load_state_dict_post_hooks.clear()
640 # Clear TransformerLens-specific hooks
641 if hasattr(module, "remove_all_hooks"): 641 ↛ 642line 641 didn't jump to line 642 because the condition on line 641 was never true
642 module.remove_all_hooks()
644 # Clear gradients
645 if hasattr(module, "zero_grad"): 645 ↛ 621line 645 didn't jump to line 621 because the condition on line 645 was always true
646 try:
647 module.zero_grad(set_to_none=True)
648 except Exception:
649 pass
650 except Exception:
651 # If hook cleanup fails, continue anyway
652 pass
654 # Clear top-level hooks
655 if hasattr(model, "_forward_hooks"): 655 ↛ 657line 655 didn't jump to line 657 because the condition on line 655 was always true
656 model._forward_hooks.clear()
657 if hasattr(model, "_backward_hooks"): 657 ↛ 659line 657 didn't jump to line 659 because the condition on line 657 was always true
658 model._backward_hooks.clear()
659 if hasattr(model, "_forward_pre_hooks"): 659 ↛ 663line 659 didn't jump to line 663 because the condition on line 659 was always true
660 model._forward_pre_hooks.clear()
662 # Clear top-level gradients
663 if hasattr(model, "zero_grad"): 663 ↛ 670line 663 didn't jump to line 670 because the condition on line 663 was always true
664 try:
665 model.zero_grad(set_to_none=True)
666 except Exception:
667 pass
669 # Break circular references to help GC
670 if hasattr(model, "_modules"): 670 ↛ 684line 670 didn't jump to line 684 because the condition on line 670 was always true
671 # Clear each submodule's __dict__ to break circular references
672 for name, submodule in list(model._modules.items()):
673 if submodule is not None: 673 ↛ 672line 673 didn't jump to line 672 because the condition on line 673 was always true
674 # Clear submodule hooks
675 if hasattr(submodule, "_forward_hooks"): 675 ↛ 677line 675 didn't jump to line 677 because the condition on line 675 was always true
676 submodule._forward_hooks.clear()
677 if hasattr(submodule, "_backward_hooks"): 677 ↛ 680line 677 didn't jump to line 680 because the condition on line 677 was always true
678 submodule._backward_hooks.clear()
679 # Break reference
680 model._modules[name] = None
681 model._modules.clear()
683 # Clear parameters dict
684 if hasattr(model, "_parameters"): 684 ↛ 693line 684 didn't jump to line 693 because the condition on line 684 was always true
685 for param_name in list(model._parameters.keys()): 685 ↛ 686line 685 didn't jump to line 686 because the loop on line 685 never started
686 param = model._parameters[param_name]
687 if param is not None:
688 del param
689 model._parameters[param_name] = None
690 model._parameters.clear()
692 # Clear buffers dict
693 if hasattr(model, "_buffers"): 693 ↛ 701line 693 didn't jump to line 701 because the condition on line 693 was always true
694 for buffer_name in list(model._buffers.keys()): 694 ↛ 695line 694 didn't jump to line 695 because the loop on line 694 never started
695 buffer = model._buffers[buffer_name]
696 if buffer is not None:
697 del buffer
698 model._buffers[buffer_name] = None
699 model._buffers.clear()
701 del model
703 # Aggressive garbage collection (multiple passes to break circular references)
704 for _ in range(3):
705 gc.collect()
707 # Clear GPU cache
708 if device != "cpu" and torch.cuda.is_available(): 708 ↛ 709line 708 didn't jump to line 709 because the condition on line 708 was never true
709 torch.cuda.empty_cache()
710 torch.cuda.synchronize()
711 if device == "mps" and hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"): 711 ↛ 712line 711 didn't jump to line 712 because the condition on line 711 was never true
712 torch.mps.synchronize()
713 torch.mps.empty_cache()
715 # Track memory after cleanup
716 if track_memory and memory_tracker is not None: 716 ↛ 717line 716 didn't jump to line 717 because the condition on line 716 was never true
717 memory_after = get_memory_mb()
718 freed_mb = memory_before - memory_after
719 memory_tracker["checkpoints"].append(
720 {
721 "label": f"Cleanup: {model_name_str}",
722 "memory_mb": memory_after,
723 "freed_mb": freed_mb,
724 }
725 )
726 if verbose and freed_mb > 0:
727 print(f" Freed {freed_mb:.1f} MB")
729 # ========================================================================
730 # PHASE 1: HuggingFace + Bridge (unprocessed)
731 # ========================================================================
732 current_phase[0] = 1
733 if verbose: 733 ↛ 734line 733 didn't jump to line 734 because the condition on line 733 was never true
734 print(f"\n{'='*80}")
735 print("PHASE 1: HuggingFace + TransformerBridge (unprocessed)")
736 print(f"{'='*80}\n")
738 bridge_unprocessed = None
739 hf_model = None
740 phase1_reference = PhaseReferenceData()
742 # Load bridge without weights first to detect attn_implementation and dtype
743 if verbose: 743 ↛ 744line 743 didn't jump to line 744 because the condition on line 743 was never true
744 print("Detecting model configuration...")
745 bridge_dtype = dtype
746 attn_implementation = None
747 try:
748 # Load a lightweight version without weights to get config
749 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]
750 # Match bridge's attn_implementation: check adapter config first, then
751 # default to "eager" (bridge uses output_attentions=True which forces eager).
752 if hasattr(bridge_config_only.adapter.cfg, "attn_implementation"): 752 ↛ 754line 752 didn't jump to line 754 because the condition on line 752 was always true
753 attn_implementation = bridge_config_only.adapter.cfg.attn_implementation
754 if attn_implementation is None: 754 ↛ 756line 754 didn't jump to line 756 because the condition on line 754 was always true
755 attn_implementation = "eager"
756 if verbose: 756 ↛ 757line 756 didn't jump to line 757 because the condition on line 756 was never true
757 print(f"✓ Detected attn_implementation={attn_implementation}")
758 # Clean up config-only bridge immediately to free memory
759 del bridge_config_only
760 gc.collect()
761 except Exception as e:
762 if verbose:
763 print(f"⚠ Could not detect config (will use defaults): {str(e)}")
764 # Config-only bridge failed; apply architecture patches directly to prevent
765 # _init_weights from re-randomizing loaded weights.
766 if trust_remote_code:
767 try:
768 from transformer_lens.model_bridge.sources.transformers import (
769 determine_architecture_from_hf_config,
770 map_default_transformer_lens_config,
771 )
773 hf_cfg = AutoConfig.from_pretrained(
774 model_name, trust_remote_code=True, token=_hf_token()
775 )
776 tl_cfg = map_default_transformer_lens_config(hf_cfg)
777 arch = determine_architecture_from_hf_config(hf_cfg)
778 bridge_cfg = TransformerBridgeConfig.from_dict(tl_cfg.__dict__)
779 bridge_cfg.architecture = arch
780 bridge_cfg.model_name = model_name
781 adapter = ArchitectureAdapterFactory.select_architecture_adapter(bridge_cfg)
782 adapter.prepare_loading(model_name, {})
783 if verbose:
784 print("✓ Applied architecture patches for custom code model")
785 del adapter, bridge_cfg, tl_cfg, hf_cfg
786 except Exception as patch_err:
787 if verbose:
788 print(f"⚠ Could not apply architecture patches: {patch_err}")
790 hf_saved_logits = None
791 hf_saved_loss = None
793 if use_hf_reference and should_run_phase(1): 793 ↛ 868line 793 didn't jump to line 868 because the condition on line 793 was always true
794 try:
795 if verbose: 795 ↛ 796line 795 didn't jump to line 796 because the condition on line 795 was never true
796 print("Loading HuggingFace reference model...")
797 # Match bridge loading path: no device_map, explicit .to(device),
798 # and matching torch_dtype. When dtype=float32, loading in float32
799 # ensures non-persistent buffers (e.g., Gemma3's embed_scale) are
800 # computed at full precision. When dtype=bfloat16, both HF and
801 # Bridge load in bfloat16 so comparisons are apples-to-apples.
802 hf_kwargs: dict[str, object] = {
803 "low_cpu_mem_usage": True, # Reduce memory spikes during loading
804 "torch_dtype": dtype,
805 }
806 if _hf_token(): 806 ↛ 808line 806 didn't jump to line 808 because the condition on line 806 was always true
807 hf_kwargs["token"] = _hf_token()
808 if attn_implementation is not None: 808 ↛ 813line 808 didn't jump to line 813 because the condition on line 808 was always true
809 hf_kwargs["attn_implementation"] = attn_implementation
810 if verbose: 810 ↛ 811line 810 didn't jump to line 811 because the condition on line 810 was never true
811 print(f"Using attn_implementation={attn_implementation}")
812 # Use appropriate AutoModel class (e.g., AutoModelForSeq2SeqLM for T5)
813 auto_model_class = get_auto_model_class(model_name, trust_remote_code=trust_remote_code)
814 if verbose and auto_model_class != AutoModelForCausalLM: 814 ↛ 815line 814 didn't jump to line 815 because the condition on line 814 was never true
815 print(f"Using {auto_model_class.__name__}")
816 # Ensure pad_token_id exists (some models crash without it during init).
817 hf_config = AutoConfig.from_pretrained(
818 model_name, trust_remote_code=trust_remote_code, token=_hf_token()
819 )
820 if not hasattr(hf_config, "pad_token_id") or "pad_token_id" not in hf_config.__dict__: 820 ↛ 821line 820 didn't jump to line 821 because the condition on line 820 was never true
821 eos = getattr(hf_config, "eos_token_id", None)
822 hf_config.pad_token_id = eos[0] if isinstance(eos, (list, tuple)) else eos
823 hf_kwargs["config"] = hf_config
824 if trust_remote_code: 824 ↛ 825line 824 didn't jump to line 825 because the condition on line 824 was never true
825 hf_kwargs["trust_remote_code"] = True
826 hf_model = auto_model_class.from_pretrained(model_name, **hf_kwargs) # type: ignore[arg-type]
827 hf_model = hf_model.to(device)
828 # Post-load fixup for custom code models (e.g., OpenELM).
829 # Must run AFTER .to(device) so non-persistent buffers (RoPE sin/cos,
830 # causal_mask) are recomputed on the target device, matching the bridge
831 # which also recomputes after .to(device).
832 _fixup_custom_model(hf_model)
833 hf_model.eval()
834 # Detect dtype from HF model
835 try:
836 bridge_dtype = next(hf_model.parameters()).dtype
837 if verbose: 837 ↛ 838line 837 didn't jump to line 838 because the condition on line 837 was never true
838 print(f"Detected dtype={bridge_dtype}")
839 except StopIteration:
840 pass
841 # When float32 was requested but the model natively uses reduced
842 # precision, upcast for maximum benchmark accuracy. When dtype was
843 # explicitly set to bfloat16/float16 (e.g., to fit larger models in
844 # memory), respect it — both HF and Bridge will run in that precision.
845 if dtype == torch.float32 and bridge_dtype in (torch.float16, torch.bfloat16): 845 ↛ 846line 845 didn't jump to line 846 because the condition on line 845 was never true
846 if verbose:
847 print(f"⚠ {bridge_dtype} detected, upcasting to float32 for benchmarking...")
848 hf_model.to(torch.float32)
849 bridge_dtype = torch.float32
850 if verbose:
851 print("✓ Upcast to float32 in-place")
852 elif bridge_dtype != dtype: 852 ↛ 853line 852 didn't jump to line 853 because the condition on line 852 was never true
853 bridge_dtype = dtype # Trust the requested dtype
854 if verbose: 854 ↛ 855line 854 didn't jump to line 855 because the condition on line 854 was never true
855 print("✓ HuggingFace model loaded")
857 # HF reference logits will be captured AFTER the bridge is
858 # loaded so we can use bridge.to_tokens() for consistent
859 # tokenization (e.g. BOS prepending). This happens right
860 # after the component benchmark, while both models are still
861 # in memory, before the HF model is deleted.
863 except Exception as e:
864 if verbose:
865 print(f"✗ Could not load HuggingFace model: {str(e)}\n")
867 # Now load the full bridge with correct dtype (GPU is mostly free)
868 if verbose: 868 ↛ 869line 868 didn't jump to line 869 because the condition on line 868 was never true
869 print("Loading TransformerBridge (unprocessed)...")
870 try:
871 bridge_unprocessed = TransformerBridge.boot_transformers(model_name, device=device, dtype=bridge_dtype, trust_remote_code=trust_remote_code) # type: ignore[attr-defined]
872 if verbose: 872 ↛ 873line 872 didn't jump to line 873 because the condition on line 872 was never true
873 print("✓ TransformerBridge loaded (unprocessed)\n")
874 # Apply the adapter's prepare_model() to the HF reference model so
875 # both bridge and reference have the same fixups (e.g., weight tying).
876 # This keeps model-specific logic in the adapter, not the benchmark.
877 if hf_model is not None and hasattr(bridge_unprocessed, "adapter"): 877 ↛ 897line 877 didn't jump to line 897 because the condition on line 877 was always true
878 bridge_unprocessed.adapter.prepare_model(hf_model)
879 except Exception as e:
880 import traceback
882 error_trace = traceback.format_exc()
883 add_result(
884 BenchmarkResult(
885 name="load_bridge_unprocessed",
886 severity=BenchmarkSeverity.ERROR,
887 message=f"Failed to load unprocessed TransformerBridge: {str(e)}",
888 passed=False,
889 )
890 )
891 if verbose:
892 print(f"✗ Failed to load TransformerBridge: {str(e)}")
893 print(f"\nStack trace:\n{error_trace}")
894 return results
896 # Detect audio/vision models once for use across all phases
897 _is_audio = bridge_unprocessed is not None and getattr(
898 bridge_unprocessed.cfg, "is_audio_model", False
899 )
900 _is_visual = bridge_unprocessed is not None and getattr(
901 bridge_unprocessed.cfg, "is_visual_model", False
902 )
903 # Shared non-text input (spectrogram, waveform, or pixels) — the same tensor is used
904 # for the HF reference capture and the bridge forward so they stay comparable.
905 _test_modality_input = (
906 build_modality_input(bridge_unprocessed, device=device, dtype=dtype)
907 if (_is_audio or _is_visual)
908 else None
909 )
911 # Run Phase 1 benchmarks
912 if should_run_phase(1) and bridge_unprocessed: 912 ↛ 1090line 912 didn't jump to line 1090 because the condition on line 912 was always true
913 if verbose: 913 ↛ 914line 913 didn't jump to line 914 because the condition on line 913 was never true
914 print("Running Phase 1 benchmarks...\n")
916 # Component-level benchmarks
917 if verbose: 917 ↛ 918line 917 didn't jump to line 918 because the condition on line 917 was never true
918 print("1. Component-Level Benchmarks")
919 if hf_model is not None: 919 ↛ 1003line 919 didn't jump to line 1003 because the condition on line 919 was always true
920 # Full mode: component benchmark with independent HF model (brief 2.0x)
921 try:
922 component_result = benchmark_all_components(bridge_unprocessed, hf_model)
923 add_result(component_result)
924 if verbose: 924 ↛ 925line 924 didn't jump to line 925 because the condition on line 924 was never true
925 status = "✓" if component_result.passed else "✗"
926 print(f"{status} {component_result.message}\n")
927 gc.collect()
928 if device != "cpu" and torch.cuda.is_available(): 928 ↛ 929line 928 didn't jump to line 929 because the condition on line 928 was never true
929 torch.cuda.empty_cache()
930 if device == "mps" and hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"): 930 ↛ 931line 930 didn't jump to line 931 because the condition on line 930 was never true
931 torch.mps.synchronize()
932 torch.mps.empty_cache()
933 except Exception as e:
934 if verbose:
935 print(f"✗ Component benchmark failed: {e}\n")
937 # Capture HF reference outputs. Both models are still in memory (2.0x window).
938 if verbose: 938 ↛ 939line 938 didn't jump to line 939 because the condition on line 938 was never true
939 print("Capturing HF reference outputs to CPU...")
940 try:
941 if _test_modality_input is not None: 941 ↛ 943line 941 didn't jump to line 943 because the condition on line 941 was never true
942 # Audio/vision models: use the shared non-text input for HF vs bridge
943 with torch.no_grad():
944 if _is_visual:
945 hf_out = hf_model(pixel_values=_test_modality_input)
946 else:
947 hf_out = hf_model(input_values=_test_modality_input)
948 # Bare encoders output last_hidden_state, not logits
949 if hasattr(hf_out, "logits") and hf_out.logits is not None:
950 hf_saved_logits = hf_out.logits.detach().cpu().clone()
951 else:
952 hf_saved_logits = hf_out.last_hidden_state.detach().cpu().clone()
953 # No loss computation — there are no next-token labels here
954 if verbose:
955 kind = "vision" if _is_visual else "audio"
956 print(
957 f"✓ Captured HF {kind} output {hf_saved_logits.shape}, "
958 f"loss=N/A (no token labels)\n"
959 )
960 else:
961 hf_tokens = bridge_unprocessed.to_tokens(test_text)
962 is_enc_dec = is_encoder_decoder_model(
963 model_name, trust_remote_code=trust_remote_code
964 )
965 with torch.no_grad():
966 if is_enc_dec: 966 ↛ 967line 966 didn't jump to line 967 because the condition on line 966 was never true
967 decoder_start_id = getattr(
968 getattr(hf_model, "config", None),
969 "decoder_start_token_id",
970 0,
971 )
972 dec_ids = torch.tensor([[decoder_start_id]]).to(hf_tokens.device)
973 hf_out = hf_model(hf_tokens, decoder_input_ids=dec_ids)
974 else:
975 hf_out = _hf_forward_with_mask_fallback(hf_model, hf_tokens)
976 hf_saved_logits = hf_out.logits.detach().cpu().clone()
978 # Compute causal LM loss (shift logits and labels)
979 if not is_enc_dec and hf_saved_logits.shape[1] > 1: 979 ↛ 988line 979 didn't jump to line 988
980 shift_logits = hf_out.logits[..., :-1, :].contiguous()
981 shift_labels = hf_tokens[..., 1:].contiguous()
982 loss_fn = torch.nn.CrossEntropyLoss()
983 hf_saved_loss = loss_fn(
984 shift_logits.view(-1, shift_logits.size(-1)),
985 shift_labels.view(-1),
986 ).item()
988 if verbose: 988 ↛ 989line 988 didn't jump to line 989 because the condition on line 988 was never true
989 loss_str = f"{hf_saved_loss:.4f}" if hf_saved_loss is not None else "N/A"
990 print(
991 f"✓ Captured HF logits {hf_saved_logits.shape}, " f"loss={loss_str}\n"
992 )
993 del hf_tokens
994 except Exception as e:
995 if verbose:
996 print(f"⚠ Could not capture HF reference outputs: {e}\n")
998 # Delete HF model immediately after component benchmark + logit capture.
999 # From here on, Phase 1 runs at 1.0x using saved HF tensors.
1000 cleanup_model(hf_model, "HuggingFace model")
1001 hf_model = None
1002 else:
1003 if verbose:
1004 print("⏭️ Skipped (no HF reference model available)\n")
1006 # Forward pass benchmarks
1007 if verbose: 1007 ↛ 1008line 1007 didn't jump to line 1008 because the condition on line 1007 was never true
1008 print("2. Forward Pass Benchmarks")
1010 # Widen tolerance for reduced-precision benchmarking — MPS bfloat16
1011 # matmul non-determinism can exceed the float32 default of 1e-3
1012 p1_atol = 1e-3 if dtype == torch.float32 else 5e-3
1014 # For audio/vision models, reuse the input from HF reference capture
1015 _p1_input: Union[str, torch.Tensor] = test_text
1016 if _test_modality_input is not None: 1016 ↛ 1017line 1016 didn't jump to line 1017 because the condition on line 1016 was never true
1017 _p1_input = _test_modality_input
1019 if hf_saved_logits is not None: 1019 ↛ 1034line 1019 didn't jump to line 1034 because the condition on line 1019 was always true
1020 # Full mode: use pre-captured HF logits (bridge only, 1.0x)
1021 try:
1022 add_result(
1023 benchmark_forward_pass(
1024 bridge_unprocessed,
1025 _p1_input,
1026 reference_logits=hf_saved_logits.to(device),
1027 atol=p1_atol,
1028 )
1029 )
1030 except Exception as e:
1031 if verbose:
1032 print(f"✗ Forward pass benchmark failed: {e}\n")
1033 else:
1034 try:
1035 add_result(benchmark_forward_pass(bridge_unprocessed, _p1_input, atol=p1_atol))
1036 except Exception as e:
1037 if verbose:
1038 print(f"✗ Forward pass benchmark failed: {e}\n")
1040 # Capture Phase 1 reference for Phase 3 equivalence comparison.
1041 # Skip for audio/vision models (Phase 3 won't run — weight processing
1042 # unsupported — and the capture below feeds text, which they cannot accept).
1043 # When dtype==float32 (default) and the model natively uses reduced
1044 # precision, upcast for maximum accuracy. When the user explicitly
1045 # requested a non-float32 dtype, run the reference pass in that dtype
1046 # so the entire pipeline honours the requested precision.
1047 if bridge_unprocessed is not None and not _is_audio and not _is_visual: 1047 ↛ 1090line 1047 didn't jump to line 1090 because the condition on line 1047 was always true
1048 try:
1049 original_dtype = bridge_unprocessed.cfg.dtype
1050 needs_upcast = dtype == torch.float32 and original_dtype not in (
1051 torch.float32,
1052 torch.float64,
1053 )
1054 # Snapshot registered buffers before the round-trip. HF's
1055 # RotaryEmbedding recomputes inv_freq during the float32 forward
1056 # pass, and the downcast back to bfloat16 would produce different
1057 # values than the original, corrupting the model for Phase 2.
1058 saved_buffers = {}
1059 if needs_upcast: 1059 ↛ 1060line 1059 didn't jump to line 1060 because the condition on line 1059 was never true
1060 for bname, buf in bridge_unprocessed.named_buffers():
1061 saved_buffers[bname] = buf.data.clone()
1062 bridge_unprocessed.to(torch.float32)
1063 with torch.no_grad():
1064 bridge_logits = bridge_unprocessed(test_text, return_type="logits")
1065 phase1_reference.hf_logits = bridge_logits.detach().cpu().clone()
1066 bridge_loss = _compute_self_target_loss(bridge_unprocessed, test_text)
1067 phase1_reference.hf_loss = bridge_loss.item()
1068 phase1_reference.test_text = test_text
1069 if needs_upcast: 1069 ↛ 1070line 1069 didn't jump to line 1070 because the condition on line 1069 was never true
1070 bridge_unprocessed.to(original_dtype)
1071 # Restore buffers that were corrupted by the round-trip.
1072 # Use direct assignment (not copy_) to preserve original dtype.
1073 # HF's RotaryEmbedding keeps inv_freq in float32 even when the
1074 # model is bfloat16. After to(bfloat16), the buffer becomes
1075 # bfloat16, and copy_() would truncate the float32 saved values.
1076 for bname, buf in bridge_unprocessed.named_buffers():
1077 if bname in saved_buffers:
1078 buf.data = saved_buffers[bname]
1079 if verbose: 1079 ↛ 1080line 1079 didn't jump to line 1080 because the condition on line 1079 was never true
1080 dtype_note = " (upcast to float32)" if needs_upcast else ""
1081 print(
1082 f"✓ Saved Phase 1 reference data "
1083 f"(logits: {phase1_reference.hf_logits.shape}){dtype_note}"
1084 )
1085 except Exception as e:
1086 if verbose:
1087 print(f"⚠ Could not save Phase 1 reference data: {e}")
1089 # Free saved HF tensors now that Phase 1 is done
1090 del hf_saved_logits, hf_saved_loss
1092 # Save bridge_dtype before potential cleanup (needed for Phase 3)
1093 saved_bridge_dtype = bridge_dtype
1095 # Clean up HF model if still alive (e.g., Phase 1 was skipped)
1096 if hf_model is not None: 1096 ↛ 1097line 1096 didn't jump to line 1097 because the condition on line 1096 was never true
1097 cleanup_model(hf_model, "HuggingFace model")
1098 hf_model = None
1100 # ========================================================================
1101 # PHASE 2: Bridge (unprocessed) — runtime self-checks + HF equivalence
1102 # ========================================================================
1103 current_phase[0] = 2
1105 # OPTIMIZATION: Run generation benchmarks first (only bridge in memory).
1106 if should_run_phase(2) and bridge_unprocessed: 1106 ↛ 1194line 1106 didn't jump to line 1194 because the condition on line 1106 was always true
1107 if verbose: 1107 ↛ 1108line 1107 didn't jump to line 1108 because the condition on line 1107 was never true
1108 print(f"\n{'='*80}")
1109 print("PHASE 2: TransformerBridge (unprocessed) — runtime self-checks + HF equivalence")
1110 print(f"{'='*80}\n")
1111 if verbose: 1111 ↛ 1112line 1111 didn't jump to line 1112 because the condition on line 1111 was never true
1112 print("Running Phase 2 benchmarks...\n")
1114 # Generation benchmarks (unprocessed only) - RUN FIRST
1115 # Skip for encoder-decoder and audio models (no text generation capability)
1116 # Diffusion LMs generate through a native sampler; the benchmarks route
1117 # to it, so only architectures with neither path are skipped.
1118 _adapter = getattr(bridge_unprocessed, "adapter", None)
1119 _no_generate = not getattr(_adapter, "supports_generation", True) and not getattr(
1120 _adapter, "native_sampler", None
1121 )
1122 _skip_generation = (
1123 is_encoder_decoder_model(model_name)
1124 or getattr(bridge_unprocessed.cfg, "is_audio_model", False)
1125 or _no_generate
1126 )
1127 _skip_reason = (
1128 "Skipped (model does not support generation)"
1129 if _no_generate
1130 else "Skipped (encoder-decoder model)"
1131 )
1132 if verbose: 1132 ↛ 1133line 1132 didn't jump to line 1133 because the condition on line 1132 was never true
1133 print("1. Generation Benchmarks (unprocessed)")
1134 if _skip_generation: 1134 ↛ 1170line 1134 didn't jump to line 1170 because the condition on line 1134 was always true
1135 if verbose: 1135 ↛ 1136line 1135 didn't jump to line 1136 because the condition on line 1135 was never true
1136 print(f"⏭️ {_skip_reason}\n")
1137 add_result(
1138 BenchmarkResult(
1139 name="generation",
1140 severity=BenchmarkSeverity.INFO,
1141 passed=True,
1142 message=_skip_reason,
1143 )
1144 )
1145 add_result(
1146 BenchmarkResult(
1147 name="generation_with_kv_cache",
1148 severity=BenchmarkSeverity.INFO,
1149 passed=True,
1150 message=_skip_reason,
1151 )
1152 )
1153 add_result(
1154 BenchmarkResult(
1155 name="multiple_generation_calls",
1156 severity=BenchmarkSeverity.INFO,
1157 passed=True,
1158 message=_skip_reason,
1159 )
1160 )
1161 add_result(
1162 BenchmarkResult(
1163 name="text_quality",
1164 severity=BenchmarkSeverity.INFO,
1165 passed=True,
1166 message=_skip_reason,
1167 )
1168 )
1169 else:
1170 try:
1171 add_result(benchmark_generation(bridge_unprocessed, test_text, max_new_tokens=10))
1172 add_result(
1173 benchmark_generation_with_kv_cache(
1174 bridge_unprocessed, test_text, max_new_tokens=10
1175 )
1176 )
1177 add_result(
1178 benchmark_multiple_generation_calls(
1179 bridge_unprocessed,
1180 test_prompts=[
1181 "The quick brown fox",
1182 "Hello world",
1183 "Machine learning is",
1184 ],
1185 max_new_tokens=5,
1186 )
1187 )
1188 gc.collect() # Force cleanup after generation benchmarks
1189 except Exception as e:
1190 if verbose:
1191 print(f"✗ Generation benchmark failed: {e}\n")
1193 # Run Phase 2 runtime benchmarks using unified function
1194 if should_run_phase(2) and bridge_unprocessed: 1194 ↛ 1236line 1194 didn't jump to line 1236 because the condition on line 1194 was always true
1195 if verbose: 1195 ↛ 1196line 1195 didn't jump to line 1196 because the condition on line 1195 was never true
1196 print("2. Running Unprocessed Model Runtime Benchmarks\n")
1198 # When dtype==float32 (default) but the model natively loaded in
1199 # reduced precision, upcast for maximum benchmark accuracy. When the
1200 # user explicitly requested bfloat16/float16, honour that — run the
1201 # entire comparison in the requested precision.
1202 phase2_restore_dtype = None
1203 if dtype == torch.float32 and bridge_dtype in (torch.bfloat16, torch.float16): 1203 ↛ 1204line 1203 didn't jump to line 1204 because the condition on line 1203 was never true
1204 try:
1205 bridge_unprocessed.to(torch.float32)
1206 phase2_restore_dtype = bridge_dtype
1207 if verbose:
1208 print(f" (upcast from {bridge_dtype} to float32 for comparison)\n")
1209 except Exception:
1210 phase2_restore_dtype = None # Upcast failed; proceed as-is
1212 phase2_results = run_comparison_benchmarks(
1213 bridge_model=bridge_unprocessed,
1214 test_text=test_text,
1215 phase_name="Phase 2",
1216 is_processed=False, # Unprocessed mode - skip weight processing tests
1217 verbose=verbose,
1218 phase1_reference=phase1_reference, # Saved HF logits/loss for equivalence testing
1219 restore_dtype_after_equivalence=phase2_restore_dtype,
1220 )
1221 # Tag all phase 2 results with phase number
1222 for result in phase2_results:
1223 if result.phase is None: 1223 ↛ 1222line 1223 didn't jump to line 1222 because the condition on line 1223 was always true
1224 result.phase = 2
1225 results.extend(phase2_results)
1227 # bridge_unprocessed is kept alive for Phase 3 and Phase 4 — reusing the
1228 # same instance avoids non-deterministic loading in some architectures
1229 # (e.g., OpenELM).
1231 # ========================================================================
1232 # PHASE 4: Text Quality (profile prompts, judge perplexity-ratio scoring)
1233 # Runs before Phase 3 so it can reuse bridge_unprocessed (Phase 3
1234 # destructively processes the weights, consuming the bridge).
1235 # ========================================================================
1236 current_phase[0] = 4
1238 if ( 1238 ↛ 1251line 1238 didn't jump to line 1251 because the condition on line 1238 was never true
1239 should_run_phase(4)
1240 and bridge_unprocessed is not None
1241 # applicable_phases and supports_generation are independent switches;
1242 # without this check a disagreement surfaces as an ERROR, not a skip.
1243 # Native-sampler architectures generate too, just not autoregressively.
1244 and (
1245 getattr(bridge_unprocessed.adapter, "supports_generation", True)
1246 or getattr(bridge_unprocessed.adapter, "native_sampler", None) is not None
1247 )
1248 and not is_masked_lm_model(model_name, trust_remote_code=trust_remote_code)
1249 and not is_audio_model(model_name, trust_remote_code=trust_remote_code)
1250 ):
1251 if prompt_profile is None:
1252 from transformer_lens.benchmarks.text_quality_profiles import (
1253 resolve_profile,
1254 )
1255 from transformer_lens.tools.model_registry.registry_io import (
1256 registry_prompt_profile,
1257 )
1259 config = getattr(bridge_unprocessed, "original_model", None)
1260 archs = getattr(getattr(config, "config", None), "architectures", None) or []
1261 prompt_profile = str(
1262 resolve_profile(
1263 model_name, archs[0] if archs else None, registry_prompt_profile(model_name)
1264 )
1265 )
1267 if verbose:
1268 print(f"\n{'='*80}")
1269 print(f"PHASE 2.5: Text Quality (profile {prompt_profile}, judge ratio scoring)")
1270 print(f"{'='*80}\n")
1272 try:
1273 text_quality_result = benchmark_text_quality(
1274 bridge_unprocessed,
1275 prompt_profile,
1276 judge_model=judge_model,
1277 judge_tokenizer=judge_tokenizer,
1278 model_name=model_name,
1279 )
1280 text_quality_result.phase = 4
1281 add_result(text_quality_result)
1282 except Exception as e:
1283 if verbose:
1284 print(f"✗ Text quality benchmark failed: {e}\n")
1286 # ========================================================================
1287 # Phase 7: Multimodal Tests (only for multimodal models)
1288 # Runs before Phase 3 so we can reuse bridge_unprocessed before cleanup.
1289 # ========================================================================
1290 if ( 1290 ↛ 1295line 1290 didn't jump to line 1295 because the condition on line 1290 was never true
1291 bridge_unprocessed is not None
1292 and getattr(bridge_unprocessed.cfg, "is_multimodal", False)
1293 and should_run_phase(7)
1294 ):
1295 current_phase[0] = 7
1296 if verbose:
1297 print("\n" + "=" * 80)
1298 print("PHASE 7: MULTIMODAL TESTS")
1299 print("=" * 80)
1300 print("Testing multimodal forward pass, generation, and caching with images.")
1301 print("=" * 80 + "\n")
1303 try:
1304 from transformer_lens.benchmarks.multimodal import (
1305 benchmark_multimodal_cache,
1306 benchmark_multimodal_forward,
1307 benchmark_multimodal_generation,
1308 )
1310 mm_results = [
1311 benchmark_multimodal_forward(bridge_unprocessed, test_text=test_text),
1312 benchmark_multimodal_generation(bridge_unprocessed, test_text=test_text),
1313 benchmark_multimodal_cache(bridge_unprocessed, test_text=test_text),
1314 ]
1315 for result in mm_results:
1316 result.phase = 7
1317 results.append(result)
1318 if verbose:
1319 print(result)
1321 if verbose:
1322 print("\n" + "=" * 80)
1323 print("PHASE 7 COMPLETE")
1324 print("=" * 80)
1326 except Exception as e:
1327 if verbose:
1328 print(f"\n⚠ Multimodal tests failed: {e}\n")
1329 results.append(
1330 BenchmarkResult(
1331 name="multimodal_suite",
1332 passed=False,
1333 severity=BenchmarkSeverity.ERROR,
1334 message=f"Failed to run multimodal tests: {str(e)}",
1335 details={"error": str(e)},
1336 phase=7,
1337 )
1338 )
1340 # ========================================================================
1341 # Phase 8: Audio Tests (only for audio encoder models)
1342 # Runs before Phase 3 so we can reuse bridge_unprocessed before cleanup.
1343 # ========================================================================
1344 if ( 1344 ↛ 1349line 1344 didn't jump to line 1349 because the condition on line 1344 was never true
1345 bridge_unprocessed is not None
1346 and getattr(bridge_unprocessed.cfg, "is_audio_model", False)
1347 and should_run_phase(8)
1348 ):
1349 current_phase[0] = 8
1350 if verbose:
1351 print("\n" + "=" * 80)
1352 print("PHASE 8: AUDIO TESTS")
1353 print("=" * 80)
1354 print("Testing audio forward pass, caching, representation stability, and features.")
1355 print("=" * 80 + "\n")
1357 try:
1358 from transformer_lens.benchmarks.audio import run_audio_benchmarks
1360 audio_results = run_audio_benchmarks(
1361 bridge_unprocessed,
1362 test_audio=_test_modality_input,
1363 verbose=verbose,
1364 )
1365 for result in audio_results:
1366 result.phase = 8
1367 results.append(result)
1368 if verbose:
1369 print(result)
1371 if verbose:
1372 print("\n" + "=" * 80)
1373 print("PHASE 8 COMPLETE")
1374 print("=" * 80)
1376 except Exception as e:
1377 if verbose:
1378 print(f"\n⚠ Audio tests failed: {e}\n")
1379 results.append(
1380 BenchmarkResult(
1381 name="audio_suite",
1382 passed=False,
1383 severity=BenchmarkSeverity.ERROR,
1384 message=f"Failed to run audio tests: {str(e)}",
1385 details={"error": str(e)},
1386 phase=8,
1387 )
1388 )
1390 # ========================================================================
1391 # PHASE 8 (audio-text): audio-conditioned forward for audio decoders
1392 # (Qwen2Audio etc.) — is_multimodal with an audio processor, not an encoder.
1393 # Image Phase 7 feeds pixel_values and encoder Phase 8 feeds a raw waveform;
1394 # neither exercises these models' processed-feature audio path.
1395 # ========================================================================
1396 _audio_text = (
1397 bridge_unprocessed is not None
1398 and getattr(bridge_unprocessed.cfg, "is_multimodal", False)
1399 and not getattr(bridge_unprocessed.cfg, "is_audio_model", False)
1400 and getattr(getattr(bridge_unprocessed, "processor", None), "audio_token", None) is not None
1401 )
1402 if _audio_text and should_run_phase(8): 1402 ↛ 1403line 1402 didn't jump to line 1403 because the condition on line 1402 was never true
1403 current_phase[0] = 8
1404 if verbose:
1405 print("\n" + "=" * 80 + "\nPHASE 8: AUDIO-TEXT FORWARD\n" + "=" * 80 + "\n")
1406 from transformer_lens.benchmarks.audio import benchmark_audio_text_forward
1408 result = benchmark_audio_text_forward(bridge_unprocessed)
1409 result.phase = 8
1410 add_result(result)
1412 # ========================================================================
1413 # Phase 9: Vision Tests (only for vision encoder models — ViT/DeiT, not
1414 # vision+text multimodal models, which Phase 7 covers)
1415 # Runs before Phase 3 so we can reuse bridge_unprocessed before cleanup.
1416 # ========================================================================
1417 if ( 1417 ↛ 1423line 1417 didn't jump to line 1423 because the condition on line 1417 was never true
1418 bridge_unprocessed is not None
1419 and getattr(bridge_unprocessed.cfg, "is_visual_model", False)
1420 and not getattr(bridge_unprocessed.cfg, "is_multimodal", False)
1421 and should_run_phase(9)
1422 ):
1423 current_phase[0] = 9
1424 if verbose:
1425 print("\n" + "=" * 80)
1426 print("PHASE 9: VISION TESTS")
1427 print("=" * 80)
1428 print("Testing pixel forward pass, caching, representation stability, and decoding.")
1429 print("=" * 80 + "\n")
1431 try:
1432 from transformer_lens.benchmarks.vision import run_vision_benchmarks
1434 vision_results = run_vision_benchmarks(
1435 bridge_unprocessed,
1436 test_pixels=_test_modality_input,
1437 verbose=verbose,
1438 )
1439 for result in vision_results:
1440 result.phase = 9
1441 results.append(result)
1442 if verbose:
1443 print(result)
1445 if verbose:
1446 print("\n" + "=" * 80)
1447 print("PHASE 9 COMPLETE")
1448 print("=" * 80)
1450 except Exception as e:
1451 if verbose:
1452 print(f"\n⚠ Vision tests failed: {e}\n")
1453 results.append(
1454 BenchmarkResult(
1455 name="vision_suite",
1456 passed=False,
1457 severity=BenchmarkSeverity.ERROR,
1458 message=f"Failed to run vision tests: {str(e)}",
1459 details={"error": str(e)},
1460 phase=9,
1461 )
1462 )
1464 # ========================================================================
1465 # PHASE 3: Bridge (processed/compatibility mode) — HF equivalence
1466 # ========================================================================
1467 current_phase[0] = 3
1469 def _cleanup_bridge_unprocessed():
1470 """Clean up the kept-alive bridge_unprocessed if Phase 3 is skipped."""
1471 nonlocal bridge_unprocessed
1472 if bridge_unprocessed is not None: 1472 ↛ exitline 1472 didn't return from function '_cleanup_bridge_unprocessed' because the condition on line 1472 was always true
1473 cleanup_model(bridge_unprocessed, "TransformerBridge (unprocessed)")
1474 bridge_unprocessed = None
1476 _skip_phase3 = False
1477 if not enable_compatibility_mode: 1477 ↛ 1482line 1477 didn't jump to line 1482 because the condition on line 1477 was always true
1478 _cleanup_bridge_unprocessed()
1479 _skip_phase3 = True
1480 if verbose: 1480 ↛ 1481line 1480 didn't jump to line 1481 because the condition on line 1480 was never true
1481 print("\n⚠ Compatibility mode disabled - skipping Phase 3\n")
1482 elif not should_run_phase(3):
1483 _cleanup_bridge_unprocessed()
1484 _skip_phase3 = True
1485 if verbose:
1486 print("\n⚠ Phase 3 skipped (excluded by phases filter or adapter applicable_phases)\n")
1487 elif is_encoder_decoder_model(model_name):
1488 _cleanup_bridge_unprocessed()
1489 _skip_phase3 = True
1490 if verbose:
1491 print("\n⚠ Phase 3 skipped (encoder-decoder model - weight processing not supported)\n")
1493 bridge_processed = None
1495 if not _skip_phase3: 1495 ↛ 1496line 1495 didn't jump to line 1496 because the condition on line 1495 was never true
1496 if verbose:
1497 print(f"\n{'='*80}")
1498 print("PHASE 3: TransformerBridge (processed/compatibility mode) — HF equivalence")
1499 print(f"{'='*80}\n")
1501 if not _skip_phase3: 1501 ↛ 1507line 1501 didn't jump to line 1507 because the condition on line 1501 was never true
1502 # Reuse the Phase 1 bridge instance and process weights in-place.
1503 # When dtype==float32 (default) and the model natively uses reduced
1504 # precision, upcast before processing to avoid bf16 quantization
1505 # round-trips. When the user explicitly requested bfloat16/float16,
1506 # process weights in the requested precision — no upcast.
1507 phase3_native_dtype = None # Set if we upcast; used to restore later
1508 if bridge_unprocessed is not None:
1509 try:
1510 if verbose:
1511 print("Processing weights on existing bridge (reusing Phase 1 instance)...")
1512 bridge_processed = bridge_unprocessed
1513 bridge_unprocessed = None # Transfer ownership
1514 phase3_native_dtype = bridge_processed.cfg.dtype
1515 if dtype == torch.float32 and phase3_native_dtype not in (
1516 torch.float32,
1517 torch.float64,
1518 ):
1519 bridge_processed.to(torch.float32)
1520 if verbose:
1521 print(f" (upcast from {phase3_native_dtype} to float32 before processing)")
1522 else:
1523 phase3_native_dtype = None # No restore needed
1524 bridge_processed.enable_compatibility_mode(disable_warnings=True)
1525 if verbose:
1526 print("✓ TransformerBridge compatibility mode enabled (processed)\n")
1527 except Exception as e:
1528 import traceback
1530 error_trace = traceback.format_exc()
1531 add_result(
1532 BenchmarkResult(
1533 name="process_bridge_weights",
1534 severity=BenchmarkSeverity.ERROR,
1535 message=f"Failed to process bridge weights: {str(e)}",
1536 passed=False,
1537 details={"error": str(e), "traceback": error_trace},
1538 )
1539 )
1540 if verbose:
1541 print(f"✗ Failed to process bridge weights: {str(e)}")
1542 print(f"\nStack trace:\n{error_trace}")
1543 else:
1544 # Fallback: load a fresh bridge if Phase 1 bridge was not available
1545 try:
1546 if verbose:
1547 print("Loading TransformerBridge (processed)...")
1548 bridge_dtype = saved_bridge_dtype
1549 if verbose:
1550 print(f"Using dtype={bridge_dtype} from Phase 1")
1551 bridge_processed = TransformerBridge.boot_transformers(model_name, device=device, dtype=bridge_dtype, trust_remote_code=trust_remote_code) # type: ignore[attr-defined]
1552 bridge_processed.enable_compatibility_mode(disable_warnings=True)
1553 if verbose:
1554 print("✓ TransformerBridge compatibility mode enabled (processed)\n")
1555 except Exception as e:
1556 import traceback
1558 error_trace = traceback.format_exc()
1559 add_result(
1560 BenchmarkResult(
1561 name="load_bridge_processed",
1562 severity=BenchmarkSeverity.ERROR,
1563 message=f"Failed to load processed TransformerBridge: {str(e)}",
1564 passed=False,
1565 details={"error": str(e), "traceback": error_trace},
1566 )
1567 )
1568 if verbose:
1569 print(f"✗ Failed to load processed TransformerBridge: {str(e)}")
1570 print(f"\nStack trace:\n{error_trace}")
1572 if bridge_processed is None:
1573 # Add failure results for all Phase 3 tests
1574 phase3_tests = [
1575 "no_nan_inf",
1576 "weight_magnitudes",
1577 "layer_norm_folding",
1578 "attention_output_centering",
1579 "mlp_output_centering",
1580 "unembed_centering",
1581 "value_bias_folding",
1582 "weight_modification",
1583 "logits_equivalence",
1584 "loss_equivalence",
1585 "hook_registry",
1586 "hook_functionality",
1587 "critical_forward_hooks",
1588 "forward_hooks",
1589 "run_with_cache",
1590 "activation_cache",
1591 "gradient_computation",
1592 "critical_backward_hooks",
1593 "backward_hooks",
1594 ]
1596 for test_name in phase3_tests:
1597 add_result(
1598 BenchmarkResult(
1599 name=test_name,
1600 severity=BenchmarkSeverity.ERROR,
1601 message=f"Skipped due to weight processing failure",
1602 passed=False,
1603 details={"reason": "bridge_processing_failed"},
1604 )
1605 )
1607 if verbose:
1608 print("\n" + format_results(results))
1610 # Run Phase 3 benchmarks using unified function
1611 if bridge_processed:
1612 if verbose:
1613 print("Running Phase 3 benchmarks...\n")
1615 # Phase 3 runs in the requested dtype end-to-end, so no dtype
1616 # restoration is needed.
1617 phase3_results = run_comparison_benchmarks(
1618 bridge_model=bridge_processed,
1619 test_text=test_text,
1620 phase_name="Phase 3",
1621 is_processed=True, # Processed mode - include weight processing tests
1622 verbose=verbose,
1623 phase1_reference=phase1_reference, # Saved HF logits/loss for equivalence testing
1624 )
1625 # Tag all phase 3 results with phase number
1626 for result in phase3_results:
1627 if result.phase is None:
1628 result.phase = 3
1629 results.extend(phase3_results)
1631 # Clean up Phase 3 models
1632 if bridge_processed is not None:
1633 cleanup_model(bridge_processed, "TransformerBridge (processed)")
1634 bridge_processed = None
1636 # Print summary (individual results already printed immediately)
1637 if verbose: 1637 ↛ 1638line 1637 didn't jump to line 1638 because the condition on line 1637 was never true
1638 print("\n" + "=" * 80)
1639 print("BENCHMARK SUMMARY")
1640 print("=" * 80)
1642 # Group results by phase
1643 results_by_phase: Dict[Union[int, str], List[BenchmarkResult]] = {}
1644 for r in results:
1645 phase = r.phase if r.phase is not None else "Other"
1646 if phase not in results_by_phase:
1647 results_by_phase[phase] = []
1648 results_by_phase[phase].append(r)
1650 # Print phase-by-phase summary
1651 for phase in sorted(
1652 results_by_phase.keys(), key=lambda x: x if isinstance(x, int) else 999
1653 ):
1654 phase_results = results_by_phase[phase]
1655 phase_name = f"Phase {phase}" if isinstance(phase, int) else phase
1657 phase_passed = sum(
1658 1 for r in phase_results if r.passed and r.severity != BenchmarkSeverity.SKIPPED
1659 )
1660 phase_failed = sum(
1661 1 for r in phase_results if not r.passed and r.severity != BenchmarkSeverity.SKIPPED
1662 )
1663 phase_skipped = sum(1 for r in phase_results if r.severity == BenchmarkSeverity.SKIPPED)
1664 phase_total = len(phase_results)
1665 phase_run = phase_total - phase_skipped
1667 print(f"\n{phase_name}: {phase_run} tests run")
1668 if phase_run > 0:
1669 print(f" Passed: {phase_passed}/{phase_run} ({phase_passed/phase_run*100:.1f}%)")
1670 print(f" Failed: {phase_failed}/{phase_run} ({phase_failed/phase_run*100:.1f}%)")
1671 if phase_skipped > 0:
1672 print(f" Skipped: {phase_skipped}")
1674 # Overall summary
1675 passed = sum(1 for r in results if r.passed and r.severity != BenchmarkSeverity.SKIPPED)
1676 failed = sum(1 for r in results if not r.passed and r.severity != BenchmarkSeverity.SKIPPED)
1677 skipped = sum(1 for r in results if r.severity == BenchmarkSeverity.SKIPPED)
1678 total = len(results)
1679 run_tests = total - skipped
1681 print(f"\nOverall:")
1682 print(f"Total: {total} tests")
1683 if skipped > 0:
1684 print(f"Run: {run_tests} tests")
1685 print(f"Skipped: {skipped} tests")
1686 if run_tests > 0:
1687 print(f"Passed: {passed}/{run_tests} ({passed/run_tests*100:.1f}%)")
1688 print(f"Failed: {failed}/{run_tests} ({failed/run_tests*100:.1f}%)")
1689 print("=" * 80)
1691 # Print memory summary
1692 if track_memory and memory_tracker is not None: 1692 ↛ 1693line 1692 didn't jump to line 1693 because the condition on line 1692 was never true
1693 final_memory = get_memory_mb()
1694 total_increase = final_memory - memory_tracker["initial"]
1696 if verbose:
1697 print("\n" + "=" * 80)
1698 print("MEMORY USAGE SUMMARY")
1699 print("=" * 80)
1700 print(f"Initial memory: {memory_tracker['initial']:>8.1f} MB")
1701 print(f"Final memory: {final_memory:>8.1f} MB")
1702 print(f"Net increase: {total_increase:>+8.1f} MB")
1704 if memory_tracker["checkpoints"]:
1705 print("\nCleanup operations:")
1706 for cp in memory_tracker["checkpoints"]:
1707 if cp.get("freed_mb", 0) > 0:
1708 print(
1709 f" {cp['label']:<40} freed {cp['freed_mb']:>7.1f} MB "
1710 f"(after: {cp['memory_mb']:.1f} MB)"
1711 )
1712 print("=" * 80)
1714 return results
1717def update_model_registry(
1718 model_name: str, results: List[BenchmarkResult], use_hf_reference: bool = False
1719) -> bool:
1720 """Update the model registry with benchmark results.
1722 Args:
1723 model_name: The model that was benchmarked
1724 results: List of benchmark results
1725 use_hf_reference: Whether the run numerically compared against an HF
1726 reference. Defaults to False so an unstated reference state records
1727 a passing run as PROVISIONAL, never VERIFIED.
1729 Returns:
1730 True if registry was updated successfully
1731 """
1732 from transformer_lens.tools.model_registry.registry_io import (
1733 STATUS_FAILED,
1734 STATUS_PROVISIONAL,
1735 add_verification_record,
1736 extract_phase_scores,
1737 pass_status,
1738 update_model_status,
1739 )
1741 # Threshold/note logic shared with verify_models so the two paths can't drift.
1742 from transformer_lens.tools.model_registry.verify_models import (
1743 _build_verified_note,
1744 _check_phase_scores,
1745 _extract_prompt_profile,
1746 _sanitize_note,
1747 )
1749 phase_scores = extract_phase_scores(results)
1751 score_error = _check_phase_scores(phase_scores, results)
1752 if score_error:
1753 status = STATUS_FAILED
1754 note = score_error
1755 else:
1756 status = pass_status(use_hf_reference)
1757 note = _build_verified_note(phase_scores, results)
1758 if status == STATUS_PROVISIONAL:
1759 note = f"Structural only (no HF reference): {note}"
1761 # Try to determine architecture
1762 architecture_id = "Unknown"
1763 try:
1764 from transformers import AutoConfig
1766 config = AutoConfig.from_pretrained(model_name, token=_hf_token())
1767 archs = getattr(config, "architectures", []) or []
1768 if archs: 1768 ↛ 1773line 1768 didn't jump to line 1773 because the condition on line 1768 was always true
1769 architecture_id = archs[0]
1770 except Exception:
1771 pass
1773 updated = update_model_status(
1774 model_id=model_name,
1775 arch_id=architecture_id,
1776 status=status,
1777 phase_scores=phase_scores,
1778 note=note,
1779 sanitize_fn=_sanitize_note,
1780 prompt_profile=_extract_prompt_profile(results),
1781 )
1783 # No history record for provisional runs — VerificationHistory.is_verified()
1784 # treats any record as verified, which would bypass the provisional gate.
1785 if status != STATUS_PROVISIONAL:
1786 add_verification_record(
1787 model_id=model_name,
1788 arch_id=architecture_id,
1789 notes=note,
1790 verified_by="main_benchmark",
1791 sanitize_fn=_sanitize_note,
1792 )
1794 label = {STATUS_FAILED: "FAILED", STATUS_PROVISIONAL: "PROVISIONAL"}.get(status, "VERIFIED")
1795 score_parts = ", ".join(f"P{p}={s}%" for p, s in sorted(phase_scores.items()))
1796 print(f"Updated registry for {model_name} ({label}): {score_parts or 'no phase results'}")
1797 return updated
1800def main():
1801 """Run benchmarks from command line."""
1802 import argparse
1804 parser = argparse.ArgumentParser(description="Run TransformerBridge benchmarks")
1805 parser.add_argument(
1806 "--model",
1807 type=str,
1808 default="gpt2",
1809 help="Model name to benchmark (default: gpt2)",
1810 )
1811 parser.add_argument(
1812 "--device",
1813 type=str,
1814 default="cpu",
1815 help="Device to run on (default: cpu)",
1816 )
1817 parser.add_argument(
1818 "--no-hf-reference",
1819 action="store_true",
1820 help="Disable HuggingFace reference comparison",
1821 )
1822 parser.add_argument(
1823 "--no-compat",
1824 action="store_true",
1825 help="Disable compatibility mode",
1826 )
1827 parser.add_argument(
1828 "--quiet",
1829 action="store_true",
1830 help="Suppress verbose output",
1831 )
1832 parser.add_argument(
1833 "--update-registry",
1834 action="store_true",
1835 help="Update model registry with benchmark results (default: false)",
1836 )
1837 parser.add_argument(
1838 "--trust-remote-code",
1839 action="store_true",
1840 help="Trust remote code for custom architectures (e.g., OpenELM)",
1841 )
1842 args = parser.parse_args()
1844 results = run_benchmark_suite(
1845 model_name=args.model,
1846 device=args.device,
1847 use_hf_reference=not args.no_hf_reference,
1848 enable_compatibility_mode=not args.no_compat,
1849 verbose=not args.quiet,
1850 trust_remote_code=args.trust_remote_code,
1851 )
1853 if args.update_registry:
1854 # Same requested-reference state verify_models feeds pass_status(): a
1855 # --no-hf-reference run can only mint PROVISIONAL, never VERIFIED.
1856 update_model_registry(args.model, results, use_hf_reference=not args.no_hf_reference)
1859if __name__ == "__main__": 1859 ↛ 1860line 1859 didn't jump to line 1860 because the condition on line 1859 was never true
1860 main()