transformer_lens.benchmarks package

Submodules

• transformer_lens.benchmarks.activation_cache module
  • benchmark_activation_cache()
  • benchmark_run_with_cache()
• transformer_lens.benchmarks.audio module
  • benchmark_audio_cache()
  • benchmark_audio_ctc_decode()
  • benchmark_audio_feature_extractor()
  • benchmark_audio_forward()
  • benchmark_audio_representation_stability()
  • benchmark_audio_text_forward()
  • run_audio_benchmarks()
• transformer_lens.benchmarks.backward_gradients module
  • benchmark_backward_hooks()
  • benchmark_critical_backward_hooks()
  • benchmark_gradient_computation()
  • gradient_mismatch_is_numerical_noise()
  • gradient_mismatch_stats()
  • needs_fp32_gradients()
• transformer_lens.benchmarks.component_benchmark module
  • benchmark_all_components()
• transformer_lens.benchmarks.component_outputs module
  • BenchmarkReport
    • BenchmarkReport.component_results
    • BenchmarkReport.failed_components
    • BenchmarkReport.get_component_type_summary()
    • BenchmarkReport.get_failure_by_severity()
    • BenchmarkReport.model_name
    • BenchmarkReport.pass_rate
    • BenchmarkReport.passed_components
    • BenchmarkReport.print_detailed_analysis()
    • BenchmarkReport.print_summary()
    • BenchmarkReport.total_components
  • ComponentBenchmarker
    • ComponentBenchmarker.__init__()
    • ComponentBenchmarker.benchmark_all_components()
  • ComponentTestResult
    • ComponentTestResult.component_path
    • ComponentTestResult.component_type
    • ComponentTestResult.error_message
    • ComponentTestResult.get_failure_severity()
    • ComponentTestResult.max_diff
    • ComponentTestResult.mean_diff
    • ComponentTestResult.output_shape
    • ComponentTestResult.passed
    • ComponentTestResult.percentile_diffs
  • benchmark_model()
• transformer_lens.benchmarks.encoder_common module
  • benchmark_encoder_cache()
  • benchmark_encoder_forward()
  • benchmark_encoder_representation_stability()
  • extract_encoder_states()
• transformer_lens.benchmarks.forward_pass module
  • benchmark_forward_pass()
  • benchmark_logits_equivalence()
  • benchmark_loss_equivalence()
• transformer_lens.benchmarks.generation module
  • benchmark_generation()
  • benchmark_generation_with_kv_cache()
  • benchmark_multiple_generation_calls()
  • resolve_text_generator()
• transformer_lens.benchmarks.hook_registration module
  • benchmark_critical_forward_hooks()
  • benchmark_forward_hooks()
  • benchmark_gated_hooks_fire()
  • benchmark_hook_functionality()
  • benchmark_hook_registry()
• transformer_lens.benchmarks.main_benchmark module
  • get_auto_model_class()
  • main()
  • run_benchmark_suite()
  • run_comparison_benchmarks()
  • update_model_registry()
• transformer_lens.benchmarks.multimodal module
  • benchmark_multimodal_cache()
  • benchmark_multimodal_forward()
  • benchmark_multimodal_generation()
• transformer_lens.benchmarks.text_quality module
  • benchmark_text_quality()
  • load_judge()
• transformer_lens.benchmarks.text_quality_profiles module
  • HFSignals
    • HFSignals.languages
    • HFSignals.pipeline_tag
    • HFSignals.tags
  • ProfilePrompt
    • ProfilePrompt.lang
    • ProfilePrompt.prompt
    • ProfilePrompt.reference
  • ProfileSpec
    • ProfileSpec.kind
    • ProfileSpec.lang
    • ProfileSpec.parse()
    • ProfileSpec.src
  • extract_languages()
  • is_default_profile()
  • p4_pass_threshold()
  • profile_from_hf_signals()
  • prompts_for()
  • resolve_profile()
• transformer_lens.benchmarks.utils module
  • BenchmarkResult
    • BenchmarkResult.details
    • BenchmarkResult.message
    • BenchmarkResult.name
    • BenchmarkResult.passed
    • BenchmarkResult.phase
    • BenchmarkResult.print_immediate()
    • BenchmarkResult.severity
  • BenchmarkSeverity
    • BenchmarkSeverity.DANGER
    • BenchmarkSeverity.ERROR
    • BenchmarkSeverity.INFO
    • BenchmarkSeverity.SKIPPED
    • BenchmarkSeverity.WARNING
  • PhaseReferenceData
    • PhaseReferenceData.hf_logits
    • PhaseReferenceData.hf_loss
    • PhaseReferenceData.test_text
  • bridge_self_target_loss()
  • build_modality_input()
  • compare_activation_dicts()
  • compare_scalars()
  • compare_tensors()
  • deterministic_rng()
  • filter_expected_missing_hooks()
  • format_results()
  • is_tiny_test_model()
  • make_capture_hook()
  • make_grad_capture_hook()
  • safe_allclose()
• transformer_lens.benchmarks.vision module
  • benchmark_vision_cache()
  • benchmark_vision_classification_decode()
  • benchmark_vision_embeddings()
  • benchmark_vision_forward()
  • benchmark_vision_representation_stability()
  • run_vision_benchmarks()
• transformer_lens.benchmarks.weight_processing module
  • benchmark_attention_output_centering()
  • benchmark_layer_norm_folding()
  • benchmark_mlp_output_centering()
  • benchmark_no_nan_inf()
  • benchmark_unembed_centering()
  • benchmark_value_bias_folding()
  • benchmark_weight_magnitudes()
  • benchmark_weight_modification()

Module contents

Benchmark utilities for TransformerBridge testing.

This module provides reusable benchmark functions for validating TransformerBridge against raw HuggingFace models and pre-computed reference data (golden fixtures).

class transformer_lens.benchmarks.BenchmarkResult(name: str, severity: BenchmarkSeverity, message: str, details: Dict[str, Any] | None = None, passed: bool = True, phase: int | None = None)

Bases: object

Result of a benchmark test.

details: Dict[str, Any] | None = None

message: str

name: str

passed: bool = True

phase: int | None = None

print_immediate() → None

Print this result immediately to console.

severity: BenchmarkSeverity

class transformer_lens.benchmarks.BenchmarkSeverity(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: Enum

Severity levels for benchmark results.

DANGER = 'danger'

ERROR = 'error'

INFO = 'info'

SKIPPED = 'skipped'

WARNING = 'warning'

class transformer_lens.benchmarks.PhaseReferenceData(hf_logits: Tensor | None = None, hf_loss: float | None = None, test_text: str | None = None)

Bases: object

Float32 reference data from Phase 1 for Phase 3 equivalence comparison.

hf_logits: Tensor | None = None

hf_loss: float | None = None

test_text: str | None = None

transformer_lens.benchmarks.benchmark_activation_cache(bridge: TransformerBridge, test_text: str, reference_cache: dict[str, Tensor] | None = None, tolerance: float = 0.001) → BenchmarkResult

Benchmark activation cache values against a reference activation snapshot.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing (must match the snapshot’s prompt)

• reference_cache – Optional reference activations keyed by hook name (e.g. a golden fixture snapshot). Structural self-check only if None.

• tolerance – Tolerance for activation comparison

Returns:

BenchmarkResult with cache value comparison details

transformer_lens.benchmarks.benchmark_backward_hooks(bridge: TransformerBridge, test_text: str, reference_gradients: Dict[str, Tensor] | None = None, abs_tolerance: float = 0.2, rel_tolerance: float = 0.0003) → BenchmarkResult

Benchmark all backward hooks for gradient matching.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing (must match the snapshot’s prompt)

• reference_gradients – Optional reference gradients keyed by hook name (e.g. a golden fixture snapshot). Capture-only self-check if None.

• abs_tolerance – Absolute tolerance for gradient comparison

• rel_tolerance – Relative tolerance for gradient comparison

Returns:

BenchmarkResult with backward hook comparison details

transformer_lens.benchmarks.benchmark_critical_backward_hooks(bridge: TransformerBridge, test_text: str, reference_gradients: Dict[str, Tensor] | None = None, abs_tolerance: float = 0.2, rel_tolerance: float = 0.0003) → BenchmarkResult

Benchmark critical backward hooks for gradient matching.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing (must match the snapshot’s prompt)

• reference_gradients – Optional reference gradients keyed by hook name (e.g. a golden fixture snapshot). Capture-only self-check if None.

• abs_tolerance – Absolute tolerance for gradient comparison

• rel_tolerance – Relative tolerance for gradient comparison

Returns:

BenchmarkResult with critical backward hook comparison details

transformer_lens.benchmarks.benchmark_critical_forward_hooks(bridge: TransformerBridge, test_text: str, reference_activations: Dict[str, Tensor] | None = None, tolerance: float = 0.02) → BenchmarkResult

Benchmark critical forward hooks commonly used in interpretability research.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing (must match the snapshot’s prompt)

• reference_activations – Optional reference activations keyed by hook name (e.g. a golden fixture snapshot). Capture-only self-check if None.

• tolerance – Tolerance for activation comparison

Returns:

BenchmarkResult with critical hook comparison details

transformer_lens.benchmarks.benchmark_forward_hooks(bridge: TransformerBridge, test_text: str, reference_activations: Dict[str, Tensor] | None = None, tolerance: float = 0.5, prepend_bos: bool | None = None) → BenchmarkResult

Benchmark all forward hooks for activation matching.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing (must match the snapshot’s prompt)

• reference_activations – Optional reference activations keyed by hook name (e.g. a golden fixture snapshot). Fire-only self-check if None.

• tolerance – Tolerance for activation matching (fraction of mismatches allowed)

• prepend_bos – Whether to prepend BOS token. If None, uses model default.

Returns:

BenchmarkResult with hook activation comparison details

transformer_lens.benchmarks.benchmark_forward_pass(bridge: TransformerBridge, test_input: str | Tensor, reference_model: Module | None = None, reference_logits: Tensor | None = None, atol: float = 0.001, rtol: float = 0.03) → BenchmarkResult

Benchmark forward pass between TransformerBridge and reference model.

Parameters:

• bridge – TransformerBridge model to test

• test_input – Input text string or audio waveform tensor for testing

• reference_model – Optional live HF reference model (audio / encoder-decoder paths)

• reference_logits – Optional pre-computed reference logits/hidden states tensor (e.g., saved from a prior HF forward pass to avoid needing both models in memory)

• atol – Absolute tolerance for comparison

• rtol – Relative tolerance for comparison

Returns:

BenchmarkResult with comparison details

transformer_lens.benchmarks.benchmark_gated_hooks_fire(bridge: TransformerBridge, test_text: str = 'The quick brown fox', prepend_bos: bool | None = None) → BenchmarkResult

Verify each cfg-gated attention hook fires when its flag is enabled.

Hooks like hook_result, hook_q_input, hook_attn_in exist unconditionally on the attention bridge but are only populated when the corresponding config flag is set (keeping default-path cost at zero). This benchmark toggles each flag in turn, runs a short forward, and asserts at least one layer’s matching hook actually captured an activation.

use_attn_in and use_split_qkv_input are mutually exclusive, so each flag runs in its own forward pass. Plain AttentionBridge (non-PEA/JPEA) adapters raise NotImplementedError from the setter — recorded as skipped rather than failed, since the applicability gate is intentional.

transformer_lens.benchmarks.benchmark_generation(bridge: TransformerBridge, test_text: str, max_new_tokens: int = 10) → BenchmarkResult

Benchmark basic text generation.

transformer_lens.benchmarks.benchmark_generation_with_kv_cache(bridge: TransformerBridge, test_text: str, max_new_tokens: int = 10) → BenchmarkResult

Benchmark text generation with KV caching enabled.

This ensures that the KV cache is properly passed through attention layers during generation, and that the cache update logic works correctly.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for generation

• max_new_tokens – Number of tokens to generate

Returns:

BenchmarkResult with generation details

transformer_lens.benchmarks.benchmark_gradient_computation(bridge: TransformerBridge, test_text: str, reference_loss: float | None = None, atol: float = 0.001) → BenchmarkResult

Benchmark basic gradient computation.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing (must match the reference’s prompt)

• reference_loss – Optional reference last-position summed-logit value (e.g. from a golden fixture or an HF forward). Self-check only if None.

• atol – Absolute tolerance for gradient comparison

Returns:

BenchmarkResult with gradient computation comparison details

transformer_lens.benchmarks.benchmark_hook_functionality(bridge: TransformerBridge, test_text: str, reference_effect: float | None = None, atol: float = 0.002) → BenchmarkResult

Benchmark hook system functionality through ablation effects.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing (must match the reference’s prompt)

• reference_effect – Optional reference ablation loss-delta (e.g. from a golden fixture). Effect-only self-check if None.

• atol – Absolute tolerance for effect comparison

Returns:

BenchmarkResult with hook functionality comparison details

transformer_lens.benchmarks.benchmark_hook_registry(bridge: TransformerBridge, reference_hooks: Iterable[str] | None = None) → BenchmarkResult

Benchmark hook registry completeness.

Parameters:

• bridge – TransformerBridge model to test

• reference_hooks – Optional reference hook-name collection (e.g. the keys of a golden hook manifest). Structural self-check only if None.

Returns:

BenchmarkResult with registry comparison details

transformer_lens.benchmarks.benchmark_logits_equivalence(bridge: TransformerBridge, test_text: str, reference_logits: Tensor | None = None, atol: float = 0.03, rtol: float = 0.03) → BenchmarkResult

Benchmark logits output against a pre-computed reference tensor.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

• reference_logits – Optional pre-computed reference logits tensor (e.g., the HF logits captured in Phase 1, or a golden fixture). Self-check only if None.

• atol – Absolute tolerance for comparison

• rtol – Relative tolerance for comparison

Returns:

BenchmarkResult with comparison details

transformer_lens.benchmarks.benchmark_loss_equivalence(bridge: TransformerBridge, test_text: str, reference_loss: float | None = None, atol: float = 0.001) → BenchmarkResult

Benchmark loss computation against a pre-computed reference value.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

• reference_loss – Optional pre-computed reference loss value (e.g., the HF loss captured in Phase 1, or a golden fixture). Self-check only if None.

• atol – Absolute tolerance for comparison

Returns:

BenchmarkResult with comparison details

transformer_lens.benchmarks.benchmark_multiple_generation_calls(bridge: TransformerBridge, test_prompts: list, max_new_tokens: int = 5) → BenchmarkResult

Benchmark multiple generation calls to ensure KV cache handling is robust.

Parameters:

• bridge – TransformerBridge model to test

• test_prompts – List of input prompts for generation

• max_new_tokens – Number of tokens to generate per prompt

Returns:

BenchmarkResult with multiple generation details

transformer_lens.benchmarks.benchmark_run_with_cache(bridge: TransformerBridge, test_text: str) → BenchmarkResult

Benchmark run_with_cache functionality (structural self-check).

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

Returns:

BenchmarkResult with cache functionality details

transformer_lens.benchmarks.benchmark_text_quality(bridge: Any, profile: str | ProfileSpec = 'continuation', *, max_new_tokens: int | None = None, judge_model: Any | None = None, judge_tokenizer: Any | None = None, model_name: str | None = None) → BenchmarkResult

Benchmark text generation quality with profile prompts and reference-ratio scoring.

Generates from the model’s prompt-profile prompts through the real user path (bridge.generate), then scores each output against the prompt’s reference completion via the pinned judge’s perplexity ratio, with repetition and length penalties.

transformer_lens.benchmarks.benchmark_weight_modification(bridge: TransformerBridge, test_text: str) → BenchmarkResult

Benchmark that weight modifications propagate correctly.

Parameters:

• bridge – TransformerBridge model to test

• test_text – Input text for testing

Returns:

BenchmarkResult with weight modification verification details

transformer_lens.benchmarks.run_benchmark_suite(model_name: str, device: str = 'cpu', dtype: dtype = torch.float32, test_text: str | None = None, use_hf_reference: bool = True, enable_compatibility_mode: bool = True, verbose: bool = True, track_memory: bool = False, phases: list[int] | None = None, trust_remote_code: bool = False, judge_model: PreTrainedModel | None = None, judge_tokenizer: PreTrainedTokenizerBase | None = None, prompt_profile: str | None = None) → List[BenchmarkResult]

Run comprehensive benchmark suite for TransformerBridge.

This function implements an optimized multi-phase approach to minimize model reloading: Phase 1: HF + Bridge (unprocessed) - Compare against raw HuggingFace model Phase 2: Bridge (unprocessed) - Runtime self-checks + HF logits/loss equivalence Phase 3: Bridge (processed) - Compatibility mode + HF logits/loss equivalence Phase 4: Text Quality - profile prompts scored by a pinned judge’s perplexity ratio

Parameters:

• model_name – Name of the model to benchmark (e.g., “gpt2”)

• device – Device to run on (“cpu” or “cuda”)

• dtype – Precision for model loading (default: torch.float32). Use torch.bfloat16 to halve memory for larger models. Phase 2/3 comparisons automatically upcast to float32 for precision.

• test_text – Optional test text (default: standard test prompt)

• use_hf_reference – Whether to compare against HuggingFace model

• enable_compatibility_mode – Whether to enable compatibility mode on bridge

• verbose – Whether to print results to console

• track_memory – Whether to track and report memory usage (requires psutil)

• phases – Optional list of phase numbers to run (e.g., [1, 2, 3]). If None, runs all phases.

• trust_remote_code – Whether to trust remote code for custom architectures.

• judge_model – Optional pre-loaded Phase-4 judge. When provided with judge_tokenizer, avoids reloading for each model in batch.

• judge_tokenizer – Optional pre-loaded tokenizer for the Phase-4 judge.

• prompt_profile – Optional Phase-4 prompt profile (e.g. “chat”, “task:translation@en-de”). Resolved from curation + the registry when None.

Returns:

List of BenchmarkResult objects