diff --git a/tests/integration/model_bridge/generalized_components/test_joint_qkv_attention_bridge_integration.py b/tests/integration/model_bridge/generalized_components/test_joint_qkv_attention_bridge_integration.py index ebe24b4efb..94e89fac89 100644 --- a/tests/integration/model_bridge/generalized_components/test_joint_qkv_attention_bridge_integration.py +++ b/tests/integration/model_bridge/generalized_components/test_joint_qkv_attention_bridge_integration.py @@ -49,9 +49,9 @@ def dummy_hook(x, hook): assert not hook_point.has_hooks() def test_architecture_imports(self): - """Test that architecture files can be imported and reference JointQKVAttentionBridge.""" + """Test that architecture files can be imported and use appropriate attention bridges.""" # Test that we can import the architecture files without errors - # Test that JointQKVAttentionBridge is referenced in the source files + # Test that appropriate attention bridges are referenced in the source files import inspect from transformer_lens.model_bridge.supported_architectures import ( @@ -65,15 +65,18 @@ def test_architecture_imports(self): "JointQKVAttentionBridge" in gpt2_source ), "GPT-2 architecture should reference JointQKVAttentionBridge" + # BLOOM uses BloomAttentionBridge instead of JointQKVAttentionBridge + # because it requires alibi bias and residual connections bloom_source = inspect.getsource(bloom) assert ( - "JointQKVAttentionBridge" in bloom_source - ), "BLOOM architecture should reference JointQKVAttentionBridge" + "BloomAttentionBridge" in bloom_source + ), "BLOOM architecture should reference BloomAttentionBridge" + # NeoX uses JointQKVPositionEmbeddingsAttentionBridge for rotary embeddings neox_source = inspect.getsource(neox) assert ( - "JointQKVAttentionBridge" in neox_source - ), "NeoX architecture should reference JointQKVAttentionBridge" + "JointQKVPositionEmbeddingsAttentionBridge" in neox_source + ), "NeoX architecture should reference JointQKVPositionEmbeddingsAttentionBridge" @pytest.mark.slow def test_distilgpt2_integration(self): diff --git a/tests/integration/model_bridge/test_weight_processing_integration.py b/tests/integration/model_bridge/test_weight_processing_integration.py index c652438294..5379a583a5 100644 --- a/tests/integration/model_bridge/test_weight_processing_integration.py +++ b/tests/integration/model_bridge/test_weight_processing_integration.py @@ -88,9 +88,7 @@ def head_ablation_hook(value: Float[torch.Tensor, "batch pos head_index d_head"] # =========================================== print("\n3. Loading TransformerBridge without processing...") try: - bridge_unprocessed = TransformerBridge.boot_transformers( - model_name, device=device, apply_weight_processing=False - ) + bridge_unprocessed = TransformerBridge.boot_transformers(model_name, device=device) print("\n Testing baseline performance...") bridge_unprocessed_baseline = bridge_unprocessed(tokens, return_type="loss") @@ -122,9 +120,9 @@ def head_ablation_hook(value: Float[torch.Tensor, "batch pos head_index d_head"] # =========================================== print("\n4. Loading TransformerBridge with processing...") try: - bridge_processed = TransformerBridge.boot_transformers( - model_name, device=device, apply_weight_processing=True - ) + bridge_processed = TransformerBridge.boot_transformers(model_name, device=device) + + bridge_processed.process_weights() print("\n Testing baseline performance...") bridge_processed_baseline = bridge_processed(tokens, return_type="loss") @@ -288,10 +286,9 @@ def head_ablation_hook(value: Float[torch.Tensor, "batch pos head_index d_head"] if overall_success: print("\nšŸŽ‰šŸŽ‰šŸŽ‰ FULL INTEGRATION COMPATIBILITY ACHIEVED! šŸŽ‰šŸŽ‰šŸŽ‰") print("TransformerBridge is fully compatible with HookedTransformer!") - return True else: print("\nāš ļø Integration compatibility issues detected") - return False + pytest.fail("Integration compatibility issues detected") @pytest.mark.skip( diff --git a/tests/integration/model_bridge/test_weight_processing_perfect_match.py b/tests/integration/model_bridge/test_weight_processing_perfect_match.py index ac8a9d2e38..052254e93c 100644 --- a/tests/integration/model_bridge/test_weight_processing_perfect_match.py +++ b/tests/integration/model_bridge/test_weight_processing_perfect_match.py @@ -127,10 +127,9 @@ def head_ablation_hook(value, hook): print("\nšŸŽ‰šŸŽ‰šŸŽ‰ PERFECT MATCH ACHIEVED! šŸŽ‰šŸŽ‰šŸŽ‰") print("The corrected processing matches HookedTransformer exactly!") print("This solution can be applied to TransformerBridge for perfect ablation matching.") - return True else: print("\nāš ļø Not quite perfect yet, but very close!") - return False + pytest.fail("Not quite perfect yet, but very close!") if __name__ == "__main__": diff --git a/transformer_lens/benchmarks/component_outputs.py b/transformer_lens/benchmarks/component_outputs.py index c5cf9e8864..c631ea706f 100644 --- a/transformer_lens/benchmarks/component_outputs.py +++ b/transformer_lens/benchmarks/component_outputs.py @@ -265,13 +265,17 @@ def benchmark_all_components( results: List[ComponentTestResult] = [] + # Block-type components that need to be tested recursively by layer + # (they are ModuleLists that don't have direct forward methods) + block_components = {"blocks", "encoder_blocks", "decoder_blocks"} + # Test top-level components (embed, pos_embed, ln_final, unembed) for comp_name, component in component_mapping.items(): if comp_name in skip_components: continue - if comp_name == "blocks": - # Handle blocks separately + if comp_name in block_components: + # Handle blocks separately - test their subcomponents by layer continue result = self._test_component(comp_name, component, test_inputs) @@ -279,17 +283,26 @@ def benchmark_all_components( results.append(result) # Test block components recursively - if "blocks" in component_mapping and "blocks" not in skip_components: - blocks_component = component_mapping["blocks"] - n_layers = self.cfg.n_layers - - for layer_idx in range(n_layers): - # Recursively test each subcomponent and its nested subcomponents - for subcomp_name, subcomponent in blocks_component.submodules.items(): - comp_path = f"blocks.{layer_idx}.{subcomp_name}" - self._test_component_recursive( - comp_path, subcomponent, test_inputs, results, skip_components - ) + for block_type in block_components: + if block_type in component_mapping and block_type not in skip_components: + blocks_component = component_mapping[block_type] + n_layers = self.cfg.n_layers + + for layer_idx in range(n_layers): + # Recursively test each subcomponent and its nested subcomponents + for subcomp_name, subcomponent in blocks_component.submodules.items(): + comp_path = f"{block_type}.{layer_idx}.{subcomp_name}" + self._test_component_recursive( + comp_path, subcomponent, test_inputs, results, skip_components + ) + + # Clean up test inputs to free memory + if test_inputs is not None: + for key in list(test_inputs.keys()): + tensor = test_inputs[key] + if tensor is not None and isinstance(tensor, torch.Tensor): + del tensor + test_inputs.clear() # Create report passed = sum(1 for r in results if r.passed) @@ -333,6 +346,58 @@ def _test_component_recursive( if component_path in skip_components: return + # Skip MLP components that don't exist as separate modules in HF (name=None) + # These are virtual components where fc1/fc2 are directly on the layer + # Component testing doesn't work for these because get_component returns the parent layer + if "mlp" in component_path and hasattr(component, "name") and component.name is None: + return + + # Skip MLP components with custom forward signatures (e.g., BLOOM requires residual) + # These can't be tested in isolation without full model context + if "mlp" in component_path and hasattr(component, "hf_component"): + import inspect + + try: + sig = inspect.signature(component.hf_component.forward) + params = list(sig.parameters.keys()) + # Standard MLP only needs hidden_states (or self + hidden_states) + # If there are more required params, skip testing + if len(params) > 2: # self + hidden_states + other required params + return + except Exception: + # If we can't inspect, proceed with testing + pass + + # Skip attention components that require position embeddings in Phase 3 + # These can't be tested in isolation without full model context for position embeddings + if ( + "attn" in component_path + and hasattr(component, "requires_position_embeddings") + and component.requires_position_embeddings + ): + return + + # Skip attention components that use native HF attention (maintain_native_attention=True) + # These have custom forward signatures (e.g., BLOOM requires residual, alibi, attention_mask) + # and can't be tested in isolation without full model context + if ( + "attn" in component_path + and hasattr(component, "maintain_native_attention") + and component.maintain_native_attention + ): + return + + # Skip BLOOM and T5 attention and MLP components - they have custom signatures that require + # residual connections, alibi bias, or cache_position from the full model context + if "attn" in component_path or "mlp" in component_path: + # Check if this is a BLOOM or T5 model by looking at the HF model config + hf_model_config = getattr(self.hf_model, "config", None) + if hf_model_config and hasattr(hf_model_config, "model_type"): + # BLOOM requires residual and alibi bias + # T5 requires cache_position for relative position embeddings + if hf_model_config.model_type in ["bloom", "t5"]: + return + # Skip components that require specific shaped inputs from their parent modules # These components expect intermediate outputs from their parent attention/MLP # modules and can't be tested with generic hidden state inputs @@ -402,7 +467,10 @@ def _test_component( """ try: # Get bridge component - bridge_component = self.adapter.get_component(self.bridge_model, component_path) + # The adapter returns nn.Module, but for bridge models it's actually GeneralizedComponent + bridge_component = cast( + GeneralizedComponent, self.adapter.get_component(self.bridge_model, component_path) + ) # Get HuggingFace component hf_component = self.adapter.get_component(self.hf_model, component_path) @@ -412,7 +480,14 @@ def _test_component( if test_input is None: return None - # For embedding components, generate token indices once to use for both + # Get input args/kwargs from the Bridge component + # All bridge components inherit from GeneralizedComponent and have get_dummy_inputs() + batch, seq_len, _ = test_input.shape + pos_indices = ( + torch.arange(seq_len, device=test_input.device).unsqueeze(0).expand(batch, -1) + ) + + # For embedding components, generate token indices once shared_token_indices = None if component_path == "embed": batch, seq_len, _ = test_input.shape @@ -490,13 +565,28 @@ def _test_component( bridge_tensor, hf_tensor ) + # Get output shape before deleting tensors + output_shape = tuple(bridge_tensor.shape) + + # Clean up output tensors immediately to free memory + del bridge_output, hf_output, bridge_tensor, hf_tensor + if shared_inputs is not None: + # Clean up shared inputs + for key in list(shared_inputs.keys()): + val = shared_inputs[key] + if val is not None and isinstance(val, torch.Tensor): + del val + shared_inputs[key] = None + if shared_token_indices is not None: + del shared_token_indices + return ComponentTestResult( component_path=component_path, component_type=type(component).__name__, passed=passed, max_diff=max_diff, mean_diff=mean_diff, - output_shape=tuple(bridge_tensor.shape), + output_shape=output_shape, percentile_diffs=percentile_diffs, ) diff --git a/transformer_lens/benchmarks/forward_pass.py b/transformer_lens/benchmarks/forward_pass.py index d82d06a60f..a66b43a03d 100644 --- a/transformer_lens/benchmarks/forward_pass.py +++ b/transformer_lens/benchmarks/forward_pass.py @@ -14,6 +14,29 @@ from transformer_lens.model_bridge import TransformerBridge +def _is_encoder_decoder(model: torch.nn.Module) -> bool: + """Check if a model is an encoder-decoder architecture.""" + config = getattr(model, "config", None) + if config is None: + return False + return getattr(config, "is_encoder_decoder", False) + + +def _get_decoder_input_ids(model: torch.nn.Module, batch_size: int = 1) -> torch.Tensor: + """Get decoder_input_ids for encoder-decoder models. + + Args: + model: The model to get decoder_start_token_id from + batch_size: Batch size for the decoder_input_ids + + Returns: + Tensor of shape [batch_size, 1] with decoder_start_token_id + """ + config = getattr(model, "config", None) + decoder_start_token_id = getattr(config, "decoder_start_token_id", 0) if config else 0 + return torch.tensor([[decoder_start_token_id]] * batch_size) + + def benchmark_forward_pass( bridge: TransformerBridge, test_text: str, @@ -34,8 +57,20 @@ def benchmark_forward_pass( BenchmarkResult with comparison details """ try: + # Check if this is an encoder-decoder model + is_enc_dec = _is_encoder_decoder(bridge.original_model) + + # Prepare extra kwargs for encoder-decoder models + extra_kwargs = {} + if is_enc_dec: + tokens = bridge.to_tokens(test_text) + batch_size = tokens.shape[0] + decoder_input_ids = _get_decoder_input_ids(bridge.original_model, batch_size) + decoder_input_ids = decoder_input_ids.to(tokens.device) + extra_kwargs["decoder_input_ids"] = decoder_input_ids + # Run bridge forward pass - bridge_output = bridge(test_text, return_type="logits") + bridge_output = bridge(test_text, return_type="logits", **extra_kwargs) if reference_model is None: # No reference model - just verify output shape and validity @@ -69,7 +104,14 @@ def benchmark_forward_pass( # HuggingFace model tokens = bridge.to_tokens(test_text) with torch.no_grad(): - hf_output = reference_model(tokens) + if is_enc_dec: + # Encoder-decoder models need decoder_input_ids + batch_size = tokens.shape[0] + decoder_input_ids = _get_decoder_input_ids(reference_model, batch_size) + decoder_input_ids = decoder_input_ids.to(tokens.device) + hf_output = reference_model(tokens, decoder_input_ids=decoder_input_ids) + else: + hf_output = reference_model(tokens) reference_output = hf_output.logits return compare_tensors( diff --git a/transformer_lens/benchmarks/granular_weight_processing.py b/transformer_lens/benchmarks/granular_weight_processing.py index b50d6f30c5..bee376d085 100644 --- a/transformer_lens/benchmarks/granular_weight_processing.py +++ b/transformer_lens/benchmarks/granular_weight_processing.py @@ -39,12 +39,12 @@ def __str__(self) -> str: return "+".join(flags) if flags else "none" -# Define all weight processing configurations to test -# NOTE: Centering operations (center_writing_weights, center_unembed) are only valid -# when fold_ln=True, as they rely on LayerNorm ignoring the mean. Testing them without -# fold_ln produces invalid/misleading results. -WEIGHT_PROCESSING_CONFIGS = [ - # Test each operation in complete isolation with fold_ln +# Phase 4: Individual weight processing operations (test each flag in isolation) +# NOTE: Centering operations (center_writing_weights, center_unembed) require fold_ln=True +# as they rely on LayerNorm ignoring the mean. Testing them without fold_ln produces +# invalid/misleading results, so we test them with fold_ln enabled. +INDIVIDUAL_CONFIGS = [ + # Test fold_ln alone WeightProcessingConfig( name="only_fold_ln", fold_ln=True, @@ -53,30 +53,37 @@ def __str__(self) -> str: fold_value_biases=False, refactor_factored_attn_matrices=False, ), + # Test center_writing_weights (requires fold_ln) WeightProcessingConfig( name="only_center_weights", - fold_ln=True, + fold_ln=False, center_writing_weights=True, center_unembed=False, fold_value_biases=False, refactor_factored_attn_matrices=False, ), + # Test center_unembed (requires fold_ln) WeightProcessingConfig( name="only_center_unembed", - fold_ln=True, + fold_ln=False, center_writing_weights=False, center_unembed=True, fold_value_biases=False, refactor_factored_attn_matrices=False, ), + # Test fold_value_biases alone WeightProcessingConfig( name="only_fold_value_biases", - fold_ln=True, + fold_ln=False, center_writing_weights=False, center_unembed=False, fold_value_biases=True, refactor_factored_attn_matrices=False, ), +] + +# Phase 5: Combinations of weight processing operations +COMBINATION_CONFIGS = [ # Two-way combinations (fold_ln + one other) WeightProcessingConfig( name="fold_ln+center_weights", @@ -174,11 +181,12 @@ def run_granular_weight_processing_benchmarks( test_text: str, verbose: bool = True, include_refactor_tests: bool = False, + phase: int | None = None, ) -> Dict[str, List[BenchmarkResult]]: """Run benchmarks with each weight processing configuration. - This function tests each weight processing flag individually and in combination - to identify which specific processing steps cause issues. + This function tests each weight processing flag individually (Phase 4) and + in combination (Phase 5) to identify which specific processing steps cause issues. Args: model_name: Name of the model to benchmark @@ -186,6 +194,7 @@ def run_granular_weight_processing_benchmarks( test_text: Test text for generation/inference verbose: Whether to print detailed output include_refactor_tests: Whether to include experimental refactor_factored_attn_matrices tests + phase: Optional phase number (4 for individual, 5 for combinations). If None, runs both. Returns: Dictionary mapping config name to list of benchmark results @@ -234,20 +243,37 @@ def run_granular_weight_processing_benchmarks( all_results["skipped"] = [skip_result] return all_results - # Determine which configurations to test - configs_to_test = WEIGHT_PROCESSING_CONFIGS.copy() + # Determine which configurations to test based on phase + configs_to_test = [] + phase_name = "" + + if phase is None or phase == 4: + configs_to_test.extend(INDIVIDUAL_CONFIGS) + if phase == 4: + phase_name = "PHASE 4: Individual Weight Processing Flags" + + if phase is None or phase == 5: + configs_to_test.extend(COMBINATION_CONFIGS) + if phase == 5: + phase_name = "PHASE 5: Combined Weight Processing Flags" + + if phase is None: + phase_name = "PHASE 4 & 5: Granular Weight Processing" + if include_refactor_tests: configs_to_test.extend(REFACTOR_ATTN_CONFIGS) if verbose: print("\n" + "=" * 80) - print("GRANULAR WEIGHT PROCESSING BENCHMARKS") + print(phase_name) print(f"Model: {model_name}") print(f"Testing {len(configs_to_test)} configurations") + if phase is None or phase == 4: + print(f" Individual flags: {len(INDIVIDUAL_CONFIGS)}") + if phase is None or phase == 5: + print(f" Combinations: {len(COMBINATION_CONFIGS)}") if include_refactor_tests: - print( - f" ({len(WEIGHT_PROCESSING_CONFIGS)} standard + {len(REFACTOR_ATTN_CONFIGS)} refactor tests)" - ) + print(f" Refactor tests: {len(REFACTOR_ATTN_CONFIGS)}") print("=" * 80) for config in configs_to_test: @@ -346,7 +372,13 @@ def run_granular_weight_processing_benchmarks( # Clean up del bridge del ht_ref - torch.cuda.empty_cache() if torch.cuda.is_available() else None + # Force garbage collection (multiple passes to break circular references) + import gc + + for _ in range(3): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() except Exception as e: # Record failure diff --git a/transformer_lens/benchmarks/main_benchmark.py b/transformer_lens/benchmarks/main_benchmark.py index 6a42910805..091a878737 100644 --- a/transformer_lens/benchmarks/main_benchmark.py +++ b/transformer_lens/benchmarks/main_benchmark.py @@ -9,10 +9,10 @@ """ import gc -from typing import List, Optional +from typing import Dict, List, Optional, Union import torch -from transformers import AutoModelForCausalLM +from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeq2SeqLM from transformer_lens import HookedTransformer from transformer_lens.benchmarks.activation_cache import ( @@ -63,6 +63,55 @@ ) from transformer_lens.model_bridge import TransformerBridge +# Architecture names that indicate encoder-decoder models +ENCODER_DECODER_ARCHITECTURES = [ + "T5ForConditionalGeneration", + "BartForConditionalGeneration", + "MBartForConditionalGeneration", + "MT5ForConditionalGeneration", + "PegasusForConditionalGeneration", + "BlenderbotForConditionalGeneration", + "MarianMTModel", +] + + +def is_encoder_decoder_model(model_name: str) -> bool: + """Check if a model is an encoder-decoder architecture. + + Args: + model_name: The HuggingFace model name or path + + Returns: + True if the model is encoder-decoder (like T5), False otherwise + """ + try: + config = AutoConfig.from_pretrained(model_name) + # Check config attribute first + if getattr(config, "is_encoder_decoder", False): + return True + # Fallback to architecture check + architectures = getattr(config, "architectures", []) or [] + return any(arch in ENCODER_DECODER_ARCHITECTURES for arch in architectures) + except Exception: + return False + + +def get_auto_model_class(model_name: str): + """Determine the correct AutoModel class for a given model. + + Some models (like T5) are encoder-decoder and need AutoModelForSeq2SeqLM + instead of AutoModelForCausalLM. + + Args: + model_name: The HuggingFace model name or path + + Returns: + The appropriate AutoModel class (AutoModelForCausalLM or AutoModelForSeq2SeqLM) + """ + if is_encoder_decoder_model(model_name): + return AutoModelForSeq2SeqLM + return AutoModelForCausalLM + def run_comparison_benchmarks( bridge_model: TransformerBridge, @@ -477,16 +526,18 @@ def run_benchmark_suite( verbose: bool = True, track_memory: bool = False, test_weight_processing_individually: bool = False, + phases: list[int] | None = None, ) -> List[BenchmarkResult]: """Run comprehensive benchmark suite for TransformerBridge. - This function implements an optimized 4-phase approach to minimize model reloading: + This function implements an optimized 5-phase approach to minimize model reloading: Phase 1: HF + Bridge (unprocessed) - Compare against raw HuggingFace model Phase 2: Bridge (unprocessed) + HT (unprocessed) - Compare unprocessed models Phase 3: Bridge (processed) + HT (processed) - Full compatibility mode testing - Phase 4: Granular Weight Processing Tests (optional) + Phase 4: Individual Weight Processing Flags (optional) + Phase 5: Combined Weight Processing Flags (optional) - When test_weight_processing_individually=True, an additional Phase 4 runs after + When test_weight_processing_individually=True, Phases 4 & 5 run after Phase 3, testing each weight processing flag individually and in combinations. Args: @@ -500,6 +551,7 @@ def run_benchmark_suite( track_memory: Whether to track and report memory usage (requires psutil) test_weight_processing_individually: Whether to run granular weight processing tests that check each processing flag individually (default: False) + phases: Optional list of phase numbers to run (e.g., [1, 2, 3]). If None, runs all phases. Returns: List of BenchmarkResult objects @@ -540,8 +592,62 @@ def get_memory_mb(): print(f"Device: {device}") print(f"{'='*80}\n") + # Early exit if only running Phase 4/5 (they load their own models independently) + if phases is not None and all(p in [4, 5] for p in phases): + if verbose: + print(f"Skipping Phase 1-3 (only running Phase {', '.join(map(str, sorted(phases)))})") + print("Phase 4/5 load their own models independently\n") + + # Jump directly to Phase 4/5 + # Jump to granular testing + from transformer_lens.benchmarks.granular_weight_processing import ( + run_granular_weight_processing_benchmarks, + ) + + if 4 in phases and test_weight_processing_individually and enable_compatibility_mode: + phase4_results = run_granular_weight_processing_benchmarks( + model_name=model_name, + device=device, + test_text=test_text, + verbose=verbose, + phase=4, + ) + for config_name, config_results in phase4_results.items(): + for result in config_results: + result.phase = 4 + results.append(result) + if verbose: + result.print_immediate() + + if 5 in phases and test_weight_processing_individually and enable_compatibility_mode: + phase5_results = run_granular_weight_processing_benchmarks( + model_name=model_name, + device=device, + test_text=test_text, + verbose=verbose, + phase=5, + ) + for config_name, config_results in phase5_results.items(): + for result in config_results: + result.phase = 5 + results.append(result) + if verbose: + result.print_immediate() + + return results + + # Track current phase for result tagging + current_phase: List[Optional[int]] = [None] # Use list to allow modification in nested function + + def should_run_phase(phase_num: int) -> bool: + """Check if a phase should run based on the phases filter.""" + return phases is None or phase_num in phases + def add_result(result: BenchmarkResult) -> None: """Add a result and optionally print it immediately.""" + # Tag result with current phase + if current_phase[0] is not None and result.phase is None: + result.phase = current_phase[0] results.append(result) if verbose: result.print_immediate() @@ -701,6 +807,7 @@ def cleanup_model(model, model_name_str: str): # ======================================================================== # PHASE 1: HuggingFace + Bridge (unprocessed) # ======================================================================== + current_phase[0] = 1 if verbose: print(f"\n{'='*80}") print("PHASE 1: HuggingFace + TransformerBridge (unprocessed)") @@ -722,8 +829,9 @@ def cleanup_model(model, model_name_str: str): attn_implementation = bridge_config_only.adapter.cfg.attn_implementation if verbose and attn_implementation: print(f"āœ“ Detected attn_implementation={attn_implementation}") - # Clean up config-only bridge + # Clean up config-only bridge immediately to free memory del bridge_config_only + gc.collect() # Force garbage collection immediately except Exception as e: if verbose: print(f"āš  Could not detect config (will use defaults): {str(e)}") @@ -734,12 +842,19 @@ def cleanup_model(model, model_name_str: str): if verbose: print("Loading HuggingFace reference model...") # Match attn_implementation from bridge to ensure numerical consistency - hf_kwargs = {"device_map": device} + hf_kwargs = { + "device_map": device, + "low_cpu_mem_usage": True, # Reduce memory spikes during loading + } if attn_implementation is not None: hf_kwargs["attn_implementation"] = attn_implementation if verbose: print(f"Using attn_implementation={attn_implementation}") - hf_model = AutoModelForCausalLM.from_pretrained(model_name, **hf_kwargs) # type: ignore[arg-type] + # Use appropriate AutoModel class (e.g., AutoModelForSeq2SeqLM for T5) + auto_model_class = get_auto_model_class(model_name) + if verbose and auto_model_class != AutoModelForCausalLM: + print(f"Using {auto_model_class.__name__} for encoder-decoder model") + hf_model = auto_model_class.from_pretrained(model_name, **hf_kwargs) # type: ignore[arg-type] hf_model.eval() # Detect dtype from HF model try: @@ -779,7 +894,7 @@ def cleanup_model(model, model_name_str: str): return results # Run Phase 1 benchmarks - if hf_model and bridge_unprocessed: + if should_run_phase(1) and hf_model and bridge_unprocessed: if verbose: print("Running Phase 1 benchmarks...\n") @@ -792,6 +907,10 @@ def cleanup_model(model, model_name_str: str): if verbose: status = "āœ“" if component_result.passed else "āœ—" print(f"{status} {component_result.message}\n") + # Clean up after component testing + gc.collect() + if device != "cpu" and torch.cuda.is_available(): + torch.cuda.empty_cache() except Exception as e: if verbose: print(f"āœ— Component benchmark failed: {e}\n") @@ -818,6 +937,7 @@ def cleanup_model(model, model_name_str: str): # ======================================================================== # PHASE 2: Bridge (unprocessed) + HookedTransformer (unprocessed) # ======================================================================== + current_phase[0] = 2 if verbose: print(f"\n{'='*80}") print("PHASE 2: TransformerBridge (unprocessed) + HookedTransformer (unprocessed)") @@ -825,37 +945,69 @@ def cleanup_model(model, model_name_str: str): # OPTIMIZATION: Run generation benchmarks first (only bridge in memory) # Then cleanup bridge before loading HT to reduce peak memory - if bridge_unprocessed: + if should_run_phase(2) and bridge_unprocessed: if verbose: print("Running Phase 2 benchmarks...\n") # Generation benchmarks (unprocessed only) - RUN FIRST + # Skip for encoder-decoder models (T5, etc.) which require different generation API + is_enc_dec = is_encoder_decoder_model(model_name) if verbose: print("1. Generation Benchmarks (unprocessed)") - try: - add_result(benchmark_generation(bridge_unprocessed, test_text, max_new_tokens=10)) + if is_enc_dec: + if verbose: + print("ā­ļø Skipped (encoder-decoder model - requires decoder_input_ids)\n") add_result( - benchmark_generation_with_kv_cache(bridge_unprocessed, test_text, max_new_tokens=10) + BenchmarkResult( + name="generation", + severity=BenchmarkSeverity.INFO, + passed=True, + message="Skipped (encoder-decoder model)", + ) ) add_result( - benchmark_multiple_generation_calls( - bridge_unprocessed, - test_prompts=[ - "The quick brown fox", - "Hello world", - "Machine learning is", - ], - max_new_tokens=5, + BenchmarkResult( + name="generation_with_kv_cache", + severity=BenchmarkSeverity.INFO, + passed=True, + message="Skipped (encoder-decoder model)", ) ) - gc.collect() # Force cleanup after generation benchmarks - except Exception as e: - if verbose: - print(f"āœ— Generation benchmark failed: {e}\n") + add_result( + BenchmarkResult( + name="multiple_generation_calls", + severity=BenchmarkSeverity.INFO, + passed=True, + message="Skipped (encoder-decoder model)", + ) + ) + else: + try: + add_result(benchmark_generation(bridge_unprocessed, test_text, max_new_tokens=10)) + add_result( + benchmark_generation_with_kv_cache( + bridge_unprocessed, test_text, max_new_tokens=10 + ) + ) + add_result( + benchmark_multiple_generation_calls( + bridge_unprocessed, + test_prompts=[ + "The quick brown fox", + "Hello world", + "Machine learning is", + ], + max_new_tokens=5, + ) + ) + gc.collect() # Force cleanup after generation benchmarks + except Exception as e: + if verbose: + print(f"āœ— Generation benchmark failed: {e}\n") # Load HookedTransformer for comparison (after generation benchmarks) ht_model_unprocessed = None - if use_ht_reference: + if should_run_phase(2) and use_ht_reference: try: if verbose: print("Loading HookedTransformer (unprocessed) for comparison...") @@ -875,7 +1027,7 @@ def cleanup_model(model, model_name_str: str): print(f"āœ— Could not load unprocessed HookedTransformer: {str(e)}\n") # Run Phase 2 comparison benchmarks using unified function - if bridge_unprocessed: + if should_run_phase(2) and bridge_unprocessed: if verbose: print("2. Running Unprocessed Model Comparison Benchmarks\n") @@ -888,6 +1040,10 @@ def cleanup_model(model, model_name_str: str): verbose=verbose, gpt2_reference=None, # No cross-model ref for Phase 2 ) + # Tag all phase 2 results with phase number + for result in phase2_results: + if result.phase is None: + result.phase = 2 results.extend(phase2_results) # Generation benchmarks already run above (before loading HT) @@ -903,6 +1059,7 @@ def cleanup_model(model, model_name_str: str): # ======================================================================== # PHASE 3: Bridge (processed) + HookedTransformer (processed) # ======================================================================== + current_phase[0] = 3 if not enable_compatibility_mode: if verbose: print("\nāš  Compatibility mode disabled - skipping Phase 3\n") @@ -910,6 +1067,18 @@ def cleanup_model(model, model_name_str: str): print("\n" + format_results(results)) return results + if not should_run_phase(3): + if verbose: + print("\nāš  Phase 3 skipped (not in phases list)\n") + return results + + # Skip Phase 3 for encoder-decoder models - weight processing is designed for decoder-only models + if is_encoder_decoder_model(model_name): + if verbose: + print("\nāš  Phase 3 skipped (encoder-decoder model - weight processing not supported)\n") + print("\n" + format_results(results)) + return results + if verbose: print(f"\n{'='*80}") print("PHASE 3: TransformerBridge (processed) + HookedTransformer (processed)") @@ -1042,6 +1211,10 @@ def cleanup_model(model, model_name_str: str): verbose=verbose, gpt2_reference=gpt2_reference, # Use GPT-2 cross-model ref if no same-arch HT ) + # Tag all phase 3 results with phase number + for result in phase3_results: + if result.phase is None: + result.phase = 3 results.extend(phase3_results) # Clean up Phase 3 models before reporting memory @@ -1107,12 +1280,46 @@ def cleanup_model(model, model_name_str: str): print("BENCHMARK SUMMARY") print("=" * 80) + # Group results by phase + results_by_phase: Dict[Union[int, str], List[BenchmarkResult]] = {} + for r in results: + phase = r.phase if r.phase is not None else "Other" + if phase not in results_by_phase: + results_by_phase[phase] = [] + results_by_phase[phase].append(r) + + # Print phase-by-phase summary + for phase in sorted( + results_by_phase.keys(), key=lambda x: x if isinstance(x, int) else 999 + ): + phase_results = results_by_phase[phase] + phase_name = f"Phase {phase}" if isinstance(phase, int) else phase + + phase_passed = sum( + 1 for r in phase_results if r.passed and r.severity != BenchmarkSeverity.SKIPPED + ) + phase_failed = sum( + 1 for r in phase_results if not r.passed and r.severity != BenchmarkSeverity.SKIPPED + ) + phase_skipped = sum(1 for r in phase_results if r.severity == BenchmarkSeverity.SKIPPED) + phase_total = len(phase_results) + phase_run = phase_total - phase_skipped + + print(f"\n{phase_name}: {phase_run} tests run") + if phase_run > 0: + print(f" Passed: {phase_passed}/{phase_run} ({phase_passed/phase_run*100:.1f}%)") + print(f" Failed: {phase_failed}/{phase_run} ({phase_failed/phase_run*100:.1f}%)") + if phase_skipped > 0: + print(f" Skipped: {phase_skipped}") + + # Overall summary passed = sum(1 for r in results if r.passed and r.severity != BenchmarkSeverity.SKIPPED) failed = sum(1 for r in results if not r.passed and r.severity != BenchmarkSeverity.SKIPPED) skipped = sum(1 for r in results if r.severity == BenchmarkSeverity.SKIPPED) total = len(results) run_tests = total - skipped + print(f"\nOverall:") print(f"Total: {total} tests") if skipped > 0: print(f"Run: {run_tests} tests") diff --git a/transformer_lens/benchmarks/utils.py b/transformer_lens/benchmarks/utils.py index 37f3e9658e..c11c16066d 100644 --- a/transformer_lens/benchmarks/utils.py +++ b/transformer_lens/benchmarks/utils.py @@ -26,6 +26,7 @@ class BenchmarkResult: message: str details: Optional[Dict[str, Any]] = None passed: bool = True + phase: Optional[int] = None # Phase number (1, 2, 3, etc.) def __str__(self) -> str: """Format result for console output.""" diff --git a/transformer_lens/config/TransformerBridgeConfig.py b/transformer_lens/config/TransformerBridgeConfig.py index 8b5495f55f..aaf1b16e6c 100644 --- a/transformer_lens/config/TransformerBridgeConfig.py +++ b/transformer_lens/config/TransformerBridgeConfig.py @@ -22,8 +22,8 @@ def __init__( d_head: int, n_layers: int, n_ctx: int, + n_heads: int = -1, # Add n_heads to signature so it's not filtered out by from_dict d_vocab: int = -1, - n_heads: int = -1, architecture: Optional[str] = None, tokenizer_prepends_bos: bool = True, default_padding_side: Optional[str] = None, @@ -83,6 +83,7 @@ def __init__( NTK_by_parts_high_freq_factor: float = 4.0, NTK_by_parts_factor: float = 8.0, eps_attr: str = "eps", + rmsnorm_uses_offset: bool = False, attn_implementation: Optional[str] = None, **kwargs, ): @@ -163,6 +164,7 @@ def __init__( self.NTK_by_parts_high_freq_factor = NTK_by_parts_high_freq_factor self.NTK_by_parts_factor = NTK_by_parts_factor self.eps_attr = eps_attr + self.rmsnorm_uses_offset = rmsnorm_uses_offset self.attn_implementation = attn_implementation self.__post_init__() diff --git a/transformer_lens/conversion_utils/conversion_steps/arithmetic_tensor_conversion.py b/transformer_lens/conversion_utils/conversion_steps/arithmetic_tensor_conversion.py index d0ecd00fef..74afa54dc8 100644 --- a/transformer_lens/conversion_utils/conversion_steps/arithmetic_tensor_conversion.py +++ b/transformer_lens/conversion_utils/conversion_steps/arithmetic_tensor_conversion.py @@ -39,5 +39,26 @@ def handle_conversion(self, input_value, *full_context): case OperationTypes.DIVISION: return input_value / self.value + def revert(self, input_value, *full_context): + """Revert the arithmetic operation (apply inverse operation).""" + # Apply input filter if present + input_value = ( + self.input_filter(input_value) if self.input_filter is not None else input_value + ) + + # Apply inverse operation + match self.operation: + case OperationTypes.ADDITION: + output = input_value - self.value + case OperationTypes.SUBTRACTION: + output = input_value + self.value + case OperationTypes.MULTIPLICATION: + output = input_value / self.value + case OperationTypes.DIVISION: + output = input_value * self.value + + # Apply output filter if present + return self.output_filter(output) if self.output_filter is not None else output + def __repr__(self): return f"Is the following arithmetic operation: {self.operation} and value: {self.value}" diff --git a/transformer_lens/model_bridge/component_setup.py b/transformer_lens/model_bridge/component_setup.py index 24bf244d45..82cba7b68d 100644 --- a/transformer_lens/model_bridge/component_setup.py +++ b/transformer_lens/model_bridge/component_setup.py @@ -10,6 +10,7 @@ from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, ) +from transformer_lens.model_bridge.generalized_components.symbolic import SymbolicBridge from transformer_lens.model_bridge.types import RemoteModel if TYPE_CHECKING: @@ -75,6 +76,20 @@ def setup_submodules( replace_remote_component(bridged_list, submodule.name, original_model) # Add to real_components mapping component.real_components[module_name] = (submodule.name, list(bridged_list)) + elif isinstance(submodule, SymbolicBridge): + # SymbolicBridge: Set up the symbolic component and its submodules + # The symbolic bridge itself doesn't wrap a real component, but its + # submodules are set up using the parent's original_model as context + setup_submodules(submodule, architecture_adapter, original_model) + + # Add the symbolic bridge as a module (for structural access like blocks[i].mlp.in) + if module_name not in component._modules: + component.add_module(module_name, submodule) + + # Add symbolic bridge's real_components to parent's mapping with prefixed keys + for sub_name, (sub_path, sub_comp) in submodule.real_components.items(): + prefixed_key = f"{module_name}.{sub_name}" + component.real_components[prefixed_key] = (sub_path, sub_comp) else: # Set up original_component if not already set if submodule.original_component is None: diff --git a/transformer_lens/model_bridge/generalized_components/__init__.py b/transformer_lens/model_bridge/generalized_components/__init__.py index 77971c1ec8..8f56121cea 100644 --- a/transformer_lens/model_bridge/generalized_components/__init__.py +++ b/transformer_lens/model_bridge/generalized_components/__init__.py @@ -15,6 +15,9 @@ from transformer_lens.model_bridge.generalized_components.joint_qkv_attention import ( JointQKVAttentionBridge, ) +from transformer_lens.model_bridge.generalized_components.joint_qkv_position_embeddings_attention import ( + JointQKVPositionEmbeddingsAttentionBridge, +) from transformer_lens.model_bridge.generalized_components.position_embeddings_attention import ( PositionEmbeddingsAttentionBridge, ) @@ -24,12 +27,21 @@ from transformer_lens.model_bridge.generalized_components.joint_gate_up_mlp import ( JointGateUpMLPBridge, ) +from transformer_lens.model_bridge.generalized_components.symbolic import SymbolicBridge from transformer_lens.model_bridge.generalized_components.unembedding import UnembeddingBridge from transformer_lens.model_bridge.generalized_components.t5_block import T5BlockBridge +from transformer_lens.model_bridge.generalized_components.bloom_block import BloomBlockBridge +from transformer_lens.model_bridge.generalized_components.bloom_attention import ( + BloomAttentionBridge, +) +from transformer_lens.model_bridge.generalized_components.bloom_mlp import BloomMLPBridge __all__ = [ "AttentionBridge", "BlockBridge", + "BloomBlockBridge", + "BloomAttentionBridge", + "BloomMLPBridge", "Conv1DBridge", "EmbeddingBridge", "RotaryEmbeddingBridge", @@ -37,12 +49,14 @@ "NormalizationBridge", "RMSNormalizationBridge", "JointQKVAttentionBridge", + "JointQKVPositionEmbeddingsAttentionBridge", "JointGateUpMLPBridge", "LinearBridge", "MLPBridge", "GatedMLPBridge", "MoEBridge", "PositionEmbeddingsAttentionBridge", + "SymbolicBridge", "UnembeddingBridge", "T5BlockBridge", ] diff --git a/transformer_lens/model_bridge/generalized_components/attention.py b/transformer_lens/model_bridge/generalized_components/attention.py index ae6be503a8..61a3632d9b 100644 --- a/transformer_lens/model_bridge/generalized_components/attention.py +++ b/transformer_lens/model_bridge/generalized_components/attention.py @@ -54,6 +54,7 @@ def __init__( maintain_native_attention: bool = False, requires_position_embeddings: bool = False, requires_attention_mask: bool = False, + attention_mask_4d: bool = False, ): """Initialize the attention bridge. @@ -71,6 +72,8 @@ def __init__( (e.g., Gemma-3 with dual RoPE). Defaults to False. requires_attention_mask: If True, this attention requires attention_mask argument (e.g., GPTNeoX/Pythia). Defaults to False. + attention_mask_4d: If True, generate 4D attention_mask [batch, 1, tgt_len, src_len] + instead of 2D [batch, seq_len]. Required for OPT. Defaults to False. """ if conversion_rule is None: conversion_rule = AttentionAutoConversion(config) @@ -95,6 +98,7 @@ def __init__( self.maintain_native_attention = maintain_native_attention self.requires_position_embeddings = requires_position_embeddings self.requires_attention_mask = requires_attention_mask + self.attention_mask_4d = attention_mask_4d def setup_hook_compatibility(self) -> None: """Setup hook compatibility transformations to match HookedTransformer behavior. @@ -157,7 +161,14 @@ def get_random_inputs( sin = torch.zeros(1, seq_len, rotary_ndims, device=device, dtype=dtype) inputs["position_embeddings"] = (cos, sin) if self.requires_attention_mask: - inputs["attention_mask"] = torch.ones(batch_size, seq_len, device=device) + if self.attention_mask_4d: + # Generate 4D attention mask [batch, 1, tgt_len, src_len] for models like OPT + inputs["attention_mask"] = torch.ones( + batch_size, 1, seq_len, seq_len, device=device + ) + else: + # Generate 2D attention mask [batch, seq_len] for most models + inputs["attention_mask"] = torch.ones(batch_size, seq_len, device=device) return inputs def _setup_qkv_hook_reshaping(self) -> None: @@ -312,18 +323,23 @@ def forward(self, *args: Any, **kwargs: Any) -> Any: args = args[1:] output = self.original_component(*args, **kwargs) if isinstance(output, tuple) and len(output) >= 2: - # output[0] is attention output, output[1] is attention weights (pattern) + # output[0] is attention output + # output[1] may be attention weights (pattern) or position_bias (T5) + # Additional elements may include position_bias, attention weights, etc. attn_output = self.hook_out(output[0]) - attn_weights = output[1] + second_element = output[1] - # Fire hook_pattern if attention weights are present - if isinstance(attn_weights, torch.Tensor): - attn_weights = self.hook_pattern(attn_weights) + # Fire hook_pattern if the second element is attention weights (4D tensor) + # For T5, second element is position_bias which should be passed through + if isinstance(second_element, torch.Tensor) and second_element.dim() == 4: + # This looks like attention weights [batch, heads, seq, seq] + second_element = self.hook_pattern(second_element) # Also store for potential hook_attn_scores (before softmax) - # Note: Most HF implementations return post-softmax weights, so pattern == attn_scores for them - self.hook_attn_scores(attn_weights) + # Note: Most HF implementations return post-softmax weights + self.hook_attn_scores(second_element) - output = (attn_output, attn_weights) + # Preserve all output elements (important for T5 position_bias and other models) + output = (attn_output, second_element) + output[2:] elif isinstance(output, tuple) and len(output) == 1: output = (self.hook_out(output[0]),) else: @@ -339,10 +355,19 @@ def W_Q(self) -> torch.Tensor: ) # 2D: [d_model, n_heads*d_head] for Conv1D or [n_heads*d_head, d_model] for Linear if weight.ndim == 2 and self.config is not None: n_heads = self.config.n_heads if hasattr(self.config, "n_heads") else self.config.n_head - # Assume Conv1D format [d_model, n_heads*d_head] since postprocess_weights should have handled Linear models - return einops.rearrange( - weight, "d_model (n_heads d_head) -> n_heads d_model d_head", n_heads=n_heads - ) + # Detect format based on weight shape + # Linear format: [(n_heads*d_head), d_model] + # Conv1D format: [d_model, (n_heads*d_head)] + if weight.shape[0] % n_heads == 0: + # Linear format - first dimension is (n_heads*d_head) + return einops.rearrange( + weight, "(n_heads d_head) d_model -> n_heads d_model d_head", n_heads=n_heads + ) + else: + # Conv1D format - second dimension is (n_heads*d_head) + return einops.rearrange( + weight, "d_model (n_heads d_head) -> n_heads d_model d_head", n_heads=n_heads + ) return weight @property @@ -358,10 +383,19 @@ def W_K(self) -> torch.Tensor: self.config.n_heads if hasattr(self.config, "n_heads") else self.config.n_head ) ) - # Assume Conv1D format [d_model, n_heads*d_head] since postprocess_weights should have handled Linear models - return einops.rearrange( - weight, "d_model (n_heads d_head) -> n_heads d_model d_head", n_heads=n_heads - ) + # Detect format based on weight shape + # Linear format: [(n_heads*d_head), d_model] + # Conv1D format: [d_model, (n_heads*d_head)] + if weight.shape[0] % n_heads == 0: + # Linear format - first dimension is (n_heads*d_head) + return einops.rearrange( + weight, "(n_heads d_head) d_model -> n_heads d_model d_head", n_heads=n_heads + ) + else: + # Conv1D format - second dimension is (n_heads*d_head) + return einops.rearrange( + weight, "d_model (n_heads d_head) -> n_heads d_model d_head", n_heads=n_heads + ) return weight @property @@ -377,10 +411,19 @@ def W_V(self) -> torch.Tensor: self.config.n_heads if hasattr(self.config, "n_heads") else self.config.n_head ) ) - # Assume Conv1D format [d_model, n_heads*d_head] since postprocess_weights should have handled Linear models - return einops.rearrange( - weight, "d_model (n_heads d_head) -> n_heads d_model d_head", n_heads=n_heads - ) + # Detect format based on weight shape + # Linear format: [(n_heads*d_head), d_model] + # Conv1D format: [d_model, (n_heads*d_head)] + if weight.shape[0] % n_heads == 0: + # Linear format - first dimension is (n_heads*d_head) + return einops.rearrange( + weight, "(n_heads d_head) d_model -> n_heads d_model d_head", n_heads=n_heads + ) + else: + # Conv1D format - second dimension is (n_heads*d_head) + return einops.rearrange( + weight, "d_model (n_heads d_head) -> n_heads d_model d_head", n_heads=n_heads + ) return weight @property diff --git a/transformer_lens/model_bridge/generalized_components/base.py b/transformer_lens/model_bridge/generalized_components/base.py index 8308f0d472..db270a6444 100644 --- a/transformer_lens/model_bridge/generalized_components/base.py +++ b/transformer_lens/model_bridge/generalized_components/base.py @@ -33,6 +33,7 @@ def __init__( config: Optional[Any] = None, submodules: Optional[Dict[str, "GeneralizedComponent"]] = None, conversion_rule: Optional[BaseTensorConversion] = None, + hook_alias_overrides: Optional[Dict[str, str]] = None, ): """Initialize the generalized component. @@ -41,6 +42,9 @@ def __init__( config: Optional configuration object for the component submodules: Dictionary of GeneralizedComponent submodules to register conversion_rule: Optional conversion rule for this component's hooks + hook_alias_overrides: Optional dictionary to override default hook aliases. + For example, {"hook_attn_out": "ln1_post.hook_out"} will make hook_attn_out + point to ln1_post.hook_out instead of the default value in self.hook_aliases. """ super().__init__() self.name = name @@ -59,6 +63,12 @@ def __init__( self.hook_in.hook_conversion = self.conversion_rule self.hook_out.hook_conversion = self.conversion_rule + # Copy class-level hook_aliases and apply any overrides + if hook_alias_overrides is not None: + # Make a copy of class-level aliases and update with overrides + self.hook_aliases = self.__class__.hook_aliases.copy() + self.hook_aliases.update(hook_alias_overrides) + def _register_hook(self, name: str, hook: HookPoint) -> None: """Register a hook in the component's hook registry.""" hook.name = name diff --git a/transformer_lens/model_bridge/generalized_components/block.py b/transformer_lens/model_bridge/generalized_components/block.py index b97ad95e60..a6423fbb5b 100644 --- a/transformer_lens/model_bridge/generalized_components/block.py +++ b/transformer_lens/model_bridge/generalized_components/block.py @@ -42,6 +42,7 @@ def __init__( name: str, config: Optional[Any] = None, submodules: Optional[Dict[str, GeneralizedComponent]] = None, + hook_alias_overrides: Optional[Dict[str, str]] = None, ): """Initialize the block bridge. @@ -49,12 +50,33 @@ def __init__( name: The name of the component in the model config: Optional configuration (unused for BlockBridge) submodules: Dictionary of submodules to register + hook_alias_overrides: Optional dictionary to override default hook aliases. + For example, {"hook_attn_out": "ln1_post.hook_out"} will make hook_attn_out + point to ln1_post.hook_out instead of the default attn.hook_out. """ - super().__init__(name, config, submodules=submodules if submodules is not None else {}) + # Apply automatic aliases based on submodules before calling parent + # This allows submodule-based aliases to be combined with explicit overrides + auto_overrides = {} + if submodules is not None: + # If ln1_post exists, hook_attn_out should point to it instead of attn.hook_out + # This matches HookedTransformer behavior where ln1_post is applied before hook_attn_out + if "ln1_post" in submodules: + auto_overrides["hook_attn_out"] = "ln1_post.hook_out" + # If ln2_post exists, hook_mlp_out should point to it instead of mlp.hook_out + if "ln2_post" in submodules: + auto_overrides["hook_mlp_out"] = "ln2_post.hook_out" + + # Merge automatic and explicit overrides (explicit takes precedence) + merged_overrides = {**auto_overrides, **(hook_alias_overrides or {})} + + # Call parent with merged overrides + super().__init__( + name, + config, + submodules=submodules if submodules is not None else {}, + hook_alias_overrides=merged_overrides if merged_overrides else None, + ) self._original_block_forward: Optional[Callable[..., Any]] = None - if submodules is not None and "ln2_post" in submodules: - self.hook_aliases = self.__class__.hook_aliases.copy() - self.hook_aliases["hook_mlp_out"] = "ln2_post.hook_out" def forward(self, *args: Any, **kwargs: Any) -> Any: """Forward pass through the block bridge. diff --git a/transformer_lens/model_bridge/generalized_components/bloom_attention.py b/transformer_lens/model_bridge/generalized_components/bloom_attention.py new file mode 100644 index 0000000000..96dfa583c5 --- /dev/null +++ b/transformer_lens/model_bridge/generalized_components/bloom_attention.py @@ -0,0 +1,111 @@ +"""BLOOM-specific attention bridge component. + +BLOOM attention requires special arguments (residual, alibi, attention_mask) that standard +JointQKVAttentionBridge doesn't handle. This custom component passes these arguments through. +""" +from typing import Any, Callable, Dict, Optional + +import torch + +from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import ( + BaseTensorConversion, +) +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) +from transformer_lens.model_bridge.generalized_components.joint_qkv_attention import ( + JointQKVAttentionBridge, +) + + +class BloomAttentionBridge(JointQKVAttentionBridge): + """Attention bridge for BLOOM models that handles residual connections and ALiBi. + + BLOOM attention has a unique forward signature that requires: + - residual: The residual connection tensor from before the attention layer + - alibi: ALiBi positional encoding bias + - attention_mask: Attention mask for padding/causality + + This bridge ensures these arguments are properly passed through to the original component. + """ + + def __init__( + self, + name: str, + config: Any, + split_qkv_matrix: Optional[Callable] = None, + submodules: Optional[Dict[str, GeneralizedComponent]] = None, + qkv_conversion_rule: Optional[BaseTensorConversion] = None, + attn_conversion_rule: Optional[BaseTensorConversion] = None, + pattern_conversion_rule: Optional[BaseTensorConversion] = None, + ): + """Initialize the BLOOM attention bridge. + + Args: + name: The name of this component + config: Model configuration + split_qkv_matrix: Function to split the qkv matrix into q, k, and v + submodules: Dictionary of submodules to register + qkv_conversion_rule: Optional conversion rule for q, k, v matrices + attn_conversion_rule: Optional conversion rule for attention output + pattern_conversion_rule: Optional conversion rule for attention patterns + """ + # BLOOM attention doesn't require attention_mask as a constructor arg, + # but it DOES require it in forward(), so we don't set requires_attention_mask=True + super().__init__( + name=name, + config=config, + split_qkv_matrix=split_qkv_matrix, + submodules=submodules, + qkv_conversion_rule=qkv_conversion_rule, + attn_conversion_rule=attn_conversion_rule, + pattern_conversion_rule=pattern_conversion_rule, + requires_position_embeddings=False, + requires_attention_mask=False, + ) + + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Forward pass through BLOOM attention with hooks. + + BLOOM attention requires these arguments: + - hidden_states (first positional arg) + - residual (second positional arg) + - alibi, attention_mask, layer_past, etc. (keyword args) + + Args: + *args: Input arguments (hidden_states, residual) + **kwargs: Additional keyword arguments including alibi, attention_mask + + Returns: + Output from BLOOM attention (tuple of hidden_states and optionally attention_weights) + """ + if self.original_component is None: + raise RuntimeError( + f"Original component not set for {self.name}. Call set_original_component() first." + ) + + # Apply hook_in to hidden_states (first positional argument) + if len(args) > 0 and isinstance(args[0], torch.Tensor): + hooked_input = self.hook_in(args[0]) + args = (hooked_input,) + args[1:] + elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor): + kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"]) + + # BLOOM attention requires residual as second positional arg + # The original BLOOM block passes it, so we just pass everything through + # No need to validate since the original component will handle it + + # Call the original BLOOM attention component with all arguments + # BLOOM attention returns (hidden_states,) or (hidden_states, attention_weights) + output = self.original_component(*args, **kwargs) + + # Apply hook_out to the hidden_states (first element of tuple) + if isinstance(output, tuple) and len(output) > 0: + processed_output = list(output) + if isinstance(output[0], torch.Tensor): + processed_output[0] = self.hook_out(output[0]) + output = tuple(processed_output) + elif isinstance(output, torch.Tensor): + output = self.hook_out(output) + + return output diff --git a/transformer_lens/model_bridge/generalized_components/bloom_block.py b/transformer_lens/model_bridge/generalized_components/bloom_block.py new file mode 100644 index 0000000000..ccb29bc03e --- /dev/null +++ b/transformer_lens/model_bridge/generalized_components/bloom_block.py @@ -0,0 +1,213 @@ +"""BLOOM-specific block bridge component. + +BLOOM blocks require special arguments (alibi, attention_mask, residual) that standard +BlockBridge doesn't handle. This custom component generates and passes these arguments. +""" +from typing import Any, Dict, Optional + +import torch + +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) +from transformer_lens.model_bridge.generalized_components.block import BlockBridge + + +class BloomBlockBridge(BlockBridge): + """Block bridge for BLOOM models that handles ALiBi positional encoding. + + BLOOM uses ALiBi (Attention with Linear Biases) instead of standard positional + embeddings. This requires generating an alibi tensor and passing it to each block + along with the attention_mask. + """ + + def __init__( + self, + name: str, + config: Optional[Any] = None, + submodules: Optional[Dict[str, GeneralizedComponent]] = None, + hook_alias_overrides: Optional[Dict[str, str]] = None, + ): + """Initialize the BLOOM block bridge. + + Args: + name: The name of the component in the model + config: Model configuration (used to get n_heads for ALiBi) + submodules: Dictionary of submodules to register + hook_alias_overrides: Optional dictionary to override default hook aliases + """ + super().__init__(name, config, submodules, hook_alias_overrides) + self.config = config + self._alibi_cache: Optional[torch.Tensor] = None + + @staticmethod + def build_alibi_tensor( + attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype + ) -> torch.Tensor: + """Build ALiBi tensor for attention biasing. + + This is adapted from the HuggingFace BLOOM implementation. + + Args: + attention_mask: Attention mask of shape [batch_size, seq_length] + num_heads: Number of attention heads + dtype: Data type for the tensor + + Returns: + ALiBi tensor of shape [num_heads, 1, seq_length] + """ + batch_size, seq_length = attention_mask.shape + closest_power_of_2 = 2 ** torch.floor( + torch.log2(torch.tensor(num_heads, dtype=torch.float32)) + ) + base = torch.tensor( + 2 ** (-(2 ** -(torch.log2(closest_power_of_2) - 3))), dtype=torch.float32 + ) + powers = torch.arange(1, 1 + closest_power_of_2, dtype=torch.int32) + slopes = torch.pow(base, powers) + + if closest_power_of_2 != num_heads: + extra_base = torch.tensor( + 2 + ** ( + -( + 2 + ** -( + torch.log2(torch.tensor(2 * closest_power_of_2, dtype=torch.float32)) + - 3 + ) + ) + ), + dtype=torch.float32, + ) + num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2) + extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, dtype=torch.int32) + slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0) + + # Shape: [num_heads, 1, seq_length] + arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :] + alibi = slopes[..., None, None] * arange_tensor + return alibi.to(dtype) + + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Forward pass through the BLOOM block. + + BLOOM blocks require `alibi` and `attention_mask` arguments. If the HF model's + BloomModel.forward() is being called, it will generate these and pass them through. + If they're missing (e.g., when called standalone), we generate them here. + + Args: + *args: Positional arguments (first should be hidden_states) + **kwargs: Keyword arguments + + Returns: + Output from the original BLOOM block + """ + # Debug: Check if alibi is being passed + # print(f"BloomBlockBridge.forward() called with kwargs keys: {list(kwargs.keys())}") + + if self.original_component is None: + raise RuntimeError( + f"Original component not set for {self.name}. Call set_original_component() first." + ) + + # Check if we should stop before executing this block + if hasattr(self, "_stop_at_layer_idx") and self._stop_at_layer_idx is not None: + import re + + if self.name is not None: + match = ( + re.search(r"blocks\.(\d+)", self.name) + or re.search(r"\.h\.(\d+)", self.name) + or re.search(r"\.layers\.(\d+)", self.name) + ) + else: + match = None + if match: + layer_idx = int(match.group(1)) + if layer_idx == self._stop_at_layer_idx: + if len(args) > 0 and isinstance(args[0], torch.Tensor): + input_tensor = args[0] + elif "hidden_states" in kwargs and isinstance( + kwargs["hidden_states"], torch.Tensor + ): + input_tensor = kwargs["hidden_states"] + else: + raise ValueError(f"Cannot find input tensor to stop at layer {layer_idx}") + input_tensor = self.hook_in(input_tensor) + from transformer_lens.model_bridge.exceptions import ( + StopAtLayerException, + ) + + raise StopAtLayerException(input_tensor) + + # Apply hook_in to hidden_states + if len(args) > 0 and isinstance(args[0], torch.Tensor): + hooked_input = self.hook_in(args[0]) + args = (hooked_input,) + args[1:] + elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor): + kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"]) + + # BLOOM blocks require 'alibi' and 'attention_mask' arguments. + # If HF's BloomModel.forward() is calling us, these will already be present. + # Only generate them if they're missing (e.g., standalone block testing). + if "alibi" not in kwargs or kwargs["alibi"] is None: + # Get hidden_states to determine shape + if len(args) > 0 and isinstance(args[0], torch.Tensor): + hidden_states = args[0] + elif "hidden_states" in kwargs: + hidden_states = kwargs["hidden_states"] + else: + raise ValueError("Could not find hidden_states in args or kwargs") + + batch_size, seq_length, _ = hidden_states.shape + device = hidden_states.device + dtype = hidden_states.dtype + + # Generate attention_mask if missing + if "attention_mask" not in kwargs or kwargs["attention_mask"] is None: + # Create default attention mask (all ones) + attention_mask = torch.ones(batch_size, seq_length, dtype=torch.long, device=device) + else: + attention_mask = kwargs["attention_mask"] + # Ensure it's 2D [batch, seq_length] for ALiBi generation + if attention_mask.dim() == 4: + # If 4D, we need 2D version for ALiBi generation + # Extract the last row which tells us which positions are valid + attention_mask_2d = attention_mask[:, 0, -1, :].long() + elif attention_mask.dim() == 2: + attention_mask_2d = attention_mask + else: + raise ValueError( + f"Unexpected attention_mask dimensions: {attention_mask.dim()}" + ) + + # Generate ALiBi bias + if self.config and hasattr(self.config, "n_heads"): + num_heads = self.config.n_heads + else: + # Fallback: try to infer from model + num_heads = 16 # BLOOM-560M has 16 heads + + # Generate alibi + alibi = self.build_alibi_tensor(attention_mask_2d, num_heads, dtype) + + # Add alibi to kwargs + kwargs["alibi"] = alibi + # else: alibi is already present from HF, don't overwrite it! + + # Call original component + output = self.original_component(*args, **kwargs) + + # Apply hook_out + if isinstance(output, tuple) and len(output) > 0: + first = output[0] + if isinstance(first, torch.Tensor): + first = self.hook_out(first) + if len(output) == 1: + return first + output = (first,) + output[1:] + return output + if isinstance(output, torch.Tensor): + output = self.hook_out(output) + return output diff --git a/transformer_lens/model_bridge/generalized_components/bloom_mlp.py b/transformer_lens/model_bridge/generalized_components/bloom_mlp.py new file mode 100644 index 0000000000..eb45dada1f --- /dev/null +++ b/transformer_lens/model_bridge/generalized_components/bloom_mlp.py @@ -0,0 +1,77 @@ +"""BLOOM-specific MLP bridge component. + +BLOOM MLP requires a special 'residual' argument that standard MLPBridge doesn't handle. +This custom component passes the residual argument through to the original component. +""" +from typing import Any, Dict, Optional + +import torch + +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) +from transformer_lens.model_bridge.generalized_components.mlp import MLPBridge + + +class BloomMLPBridge(MLPBridge): + """MLP bridge for BLOOM models that handles residual connections. + + BLOOM MLP has a unique forward signature that requires: + - hidden_states (first positional arg) + - residual (keyword arg): The residual connection tensor + + This bridge ensures the residual argument is properly passed through. + """ + + def __init__( + self, + name: Optional[str], + config: Optional[Any] = None, + submodules: Optional[Dict[str, GeneralizedComponent]] = None, + ): + """Initialize the BLOOM MLP bridge. + + Args: + name: The name of the component in the model + config: Optional configuration + submodules: Dictionary of submodules to register (e.g., dense_h_to_4h, dense_4h_to_h) + """ + super().__init__(name, config, submodules or {}) + + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Forward pass through BLOOM MLP with hooks. + + BLOOM MLP requires these arguments: + - hidden_states (first positional arg) + - residual (second positional arg) + + Args: + *args: Input arguments (hidden_states, residual) + **kwargs: Additional keyword arguments (if any) + + Returns: + Output tensor from BLOOM MLP + """ + if self.original_component is None: + raise RuntimeError( + f"Original component not set for {self.name}. Call set_original_component() first." + ) + + # Apply hook_in to hidden_states (first positional argument) + if len(args) > 0 and isinstance(args[0], torch.Tensor): + hooked_input = self.hook_in(args[0]) + args = (hooked_input,) + args[1:] + elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor): + kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"]) + + # BLOOM MLP requires residual as second positional arg + # The original BLOOM block passes it, so we just pass everything through + # No need to validate since the original component will handle it + + # Call the original BLOOM MLP component with all arguments + output = self.original_component(*args, **kwargs) + + # Apply hook_out + output = self.hook_out(output) + + return output diff --git a/transformer_lens/model_bridge/generalized_components/joint_qkv_position_embeddings_attention.py b/transformer_lens/model_bridge/generalized_components/joint_qkv_position_embeddings_attention.py new file mode 100644 index 0000000000..6fbcd0f7f5 --- /dev/null +++ b/transformer_lens/model_bridge/generalized_components/joint_qkv_position_embeddings_attention.py @@ -0,0 +1,277 @@ +"""Joint QKV attention bridge with position embeddings support. + +This module provides an attention bridge for models that use both: +1. Fused QKV matrices (like Pythia) +2. Position embeddings like RoPE (Rotary Position Embeddings) +""" +from __future__ import annotations + +from typing import Any, Callable, Dict, Optional + +import torch + +from transformer_lens.hook_points import HookPoint +from transformer_lens.model_bridge.generalized_components.joint_qkv_attention import ( + JointQKVAttentionBridge, +) + + +class JointQKVPositionEmbeddingsAttentionBridge(JointQKVAttentionBridge): + """Attention bridge for models with fused QKV and position embeddings (e.g., Pythia). + + This combines the functionality of JointQKVAttentionBridge (splitting fused QKV matrices) + with position embeddings support (for models using RoPE). + + The position_embeddings are generated by calling the model's rotary_emb + component with dummy Q/K tensors and position_ids. + """ + + def __init__( + self, + name: str, + config: Any, + split_qkv_matrix: Optional[Callable] = None, + submodules: Optional[Dict[str, Any]] = None, + **kwargs, + ): + """Initialize Joint QKV Position Embeddings attention bridge. + + Args: + name: Component name + config: Model configuration + split_qkv_matrix: Optional function to split the qkv matrix + submodules: Dictionary of subcomponents + **kwargs: Additional arguments passed to JointQKVAttentionBridge + """ + # Ensure position embeddings are required + kwargs["requires_position_embeddings"] = True + super().__init__( + name=name, + config=config, + split_qkv_matrix=split_qkv_matrix, + submodules=submodules, + **kwargs, + ) + self._rotary_emb = None + # Add hooks for cos and sin to match HookedTransformer pattern + self.hook_cos = HookPoint() + self.hook_sin = HookPoint() + + def set_rotary_emb(self, rotary_emb: Any) -> None: + """Set reference to the model's rotary embedding component. + + Args: + rotary_emb: The model's rotary_emb component (from model.gpt_neox.layers[0].attention.rotary_emb) + """ + self._rotary_emb = rotary_emb + + def _apply_position_embedding_hooks(self, position_embeddings): + """Apply hooks to position embeddings (cos, sin tuple). + + Extracts cos and sin from the position_embeddings tuple and passes them + through hook_cos and hook_sin to match HookedTransformer's behavior. + + Args: + position_embeddings: Tuple of (cos, sin) tensors + + Returns: + Tuple of (hooked_cos, hooked_sin) tensors + """ + if isinstance(position_embeddings, tuple) and len(position_embeddings) == 2: + cos, sin = position_embeddings + # Apply hooks to match HookedTransformer's rotary_cos/rotary_sin pattern + hooked_cos = self.hook_cos(cos) + hooked_sin = self.hook_sin(sin) + return (hooked_cos, hooked_sin) + return position_embeddings + + def get_random_inputs( + self, + batch_size: int = 2, + seq_len: int = 8, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> Dict[str, Any]: + """Generate random inputs for component testing. + + For models using RoPE, position_embeddings are generated by calling rotary_emb + which returns a tuple of (cos, sin) tensors. + + Args: + batch_size: Batch size for generated inputs + seq_len: Sequence length for generated inputs + device: Device to place tensors on + dtype: Dtype for generated tensors + + Returns: + Dictionary with keys: hidden_states, position_embeddings, attention_mask + """ + if device is None: + device = torch.device("cpu") + if dtype is None: + dtype = torch.float32 + + d_model = self.config.d_model if self.config and hasattr(self.config, "d_model") else 512 + inputs: Dict[str, Any] = { + "hidden_states": torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype) + } + + num_heads = ( + self.config.num_attention_heads + if self.config and hasattr(self.config, "num_attention_heads") + else self.config.n_heads + if self.config and hasattr(self.config, "n_heads") + else 8 + ) + head_dim = ( + self.config.head_dim + if self.config and hasattr(self.config, "head_dim") + else (d_model // num_heads) + ) + + # Generate position_embeddings using rotary_emb if available + if self._rotary_emb is not None: + try: + # For GPT-NeoX/Pythia rotary embeddings + # rotary_emb expects (seq_len, device) and returns (cos, sin) + position_embeddings = self._rotary_emb(seq_len, device=device) + inputs["position_embeddings"] = position_embeddings + except Exception: + # Fallback: create dummy cos/sin tensors + cos = torch.ones(1, seq_len, head_dim, device=device, dtype=dtype) + sin = torch.zeros(1, seq_len, head_dim, device=device, dtype=dtype) + inputs["position_embeddings"] = (cos, sin) + else: + # Fallback: create dummy cos/sin tensors + cos = torch.ones(1, seq_len, head_dim, device=device, dtype=dtype) + sin = torch.zeros(1, seq_len, head_dim, device=device, dtype=dtype) + inputs["position_embeddings"] = (cos, sin) + + inputs["attention_mask"] = None + return inputs + + def _apply_rotary_pos_emb( + self, q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Apply rotary position embeddings to query and key tensors. + + This implements the same logic as HuggingFace's apply_rotary_pos_emb. + + Args: + q: Query tensor [batch, heads, seq_len, head_dim] + k: Key tensor [batch, heads, seq_len, head_dim] + cos: Cosine values [batch, seq_len, head_dim] or [1, seq_len, head_dim] + sin: Sine values [batch, seq_len, head_dim] or [1, seq_len, head_dim] + + Returns: + Tuple of (rotated_q, rotated_k) + """ + # Add head dimension for broadcasting: [batch, 1, seq_len, head_dim] + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + + # Apply rotary embeddings + # Split into rotary and passthrough dimensions + rotary_dim = cos.shape[-1] + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + + # Apply rotation: q_embed = (q * cos) + (rotate_half(q) * sin) + def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin) + k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin) + + # Concatenate with passthrough dimensions + q_final = torch.cat([q_embed, q_pass], dim=-1) if q_pass.numel() > 0 else q_embed + k_final = torch.cat([k_embed, k_pass], dim=-1) if k_pass.numel() > 0 else k_embed + + return q_final, k_final + + def _reconstruct_attention( + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs + ) -> tuple: + """Manual attention computation with rotary position embeddings. + + This overrides the parent class to apply rotary embeddings to Q and K + before computing attention scores. + """ + original_component = self.original_component + assert original_component is not None + assert self.config is not None + num_heads = self.config.n_heads + + # Reshape Q, K, V to [batch, heads, seq_len, head_dim] + if len(q.shape) == 3: + batch_size, seq_len, hidden_size = q.shape + head_dim: int = hidden_size // num_heads + q = q.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2) + k = k.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2) + v = v.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2) + elif len(q.shape) == 4: + batch_size, seq_len, num_heads_tensor, head_dim = q.shape + assert ( + num_heads_tensor == num_heads + ), f"Expected {num_heads} heads, got {num_heads_tensor}" + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + seq_len = q.shape[2] # Update seq_len from transposed tensor + else: + raise ValueError(f"Unexpected Q tensor shape: {q.shape}. Expected 3D or 4D tensor.") + + # Apply rotary position embeddings if provided + position_embeddings = kwargs.get("position_embeddings", None) + if position_embeddings is not None and isinstance(position_embeddings, tuple): + cos, sin = position_embeddings + # Apply hooks to position embeddings + hooked_position_embeddings = self._apply_position_embedding_hooks(position_embeddings) + cos, sin = hooked_position_embeddings + # Apply rotary embeddings to Q and K + q, k = self._apply_rotary_pos_emb(q, k, cos, sin) + + # Compute attention scores + scale = head_dim ** (-0.5) + attn_scores = torch.matmul(q, k.transpose(-2, -1)) * scale + + # Apply causal mask + causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=q.device)) + attn_scores = attn_scores.masked_fill(causal_mask == 0, float("-inf")) + + # Apply attention mask if provided + attention_mask = kwargs.get("attention_mask", None) + if attention_mask is not None: + if attention_mask.shape[-1] != seq_len: + attention_mask = attention_mask[..., :seq_len] + if attention_mask.shape[-2] != seq_len: + attention_mask = attention_mask[..., :seq_len, :] + attn_scores = attn_scores + attention_mask + + # Apply hook to attention scores + attn_scores = self.hook_attn_scores(attn_scores) + + # Compute attention weights + attn_weights = torch.nn.functional.softmax(attn_scores, dim=-1) + + # Apply dropout if present + if hasattr(original_component, "attn_dropout"): + attn_weights = original_component.attn_dropout(attn_weights) # type: ignore[operator] + + # Apply hook to attention pattern + attn_weights = self.hook_pattern(attn_weights) + + # Compute attention output + attn_output = torch.matmul(attn_weights, v) + attn_output = attn_output.transpose(1, 2).contiguous() + final_hidden_size: int = num_heads * head_dim + attn_output = attn_output.view(batch_size, seq_len, final_hidden_size) + + # Apply output projection if present + if hasattr(self, "o") and self.o is not None: + attn_output = self.o(attn_output) + + return (attn_output, attn_weights) diff --git a/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py b/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py index 2302718490..42fa1c94ff 100644 --- a/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py +++ b/transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py @@ -1,13 +1,15 @@ -"""Gemma-3 specialized attention bridge. +"""Gemma-2/3 specialized attention bridge. -Gemma-3 uses a unique dual RoPE (Rotary Position Embedding) system that requires -position_embeddings to be generated from the model's rotary_emb component rather -than using the standard (cos, sin) tuple format. +Gemma-2/3 use Rotary Position Embeddings (RoPE) which requires special handling +to fire hook_rot_q and hook_rot_k with the correct post-rotation Q/K values. + +This is achieved by wrapping HuggingFace's eager_attention_forward function +to intercept the query and key tensors after rotary embeddings have been applied. """ from __future__ import annotations -import warnings -from typing import Any, Dict, Optional +import weakref +from typing import Any, Callable, Dict, Optional import torch import transformers.models.gemma2.modeling_gemma2 as gemma2_module @@ -17,6 +19,74 @@ AttentionBridge, ) +# Global registry mapping HF attention modules to their bridge instances +# Uses WeakValueDictionary to avoid preventing garbage collection of bridges +_ATTENTION_BRIDGE_REGISTRY: weakref.WeakValueDictionary = weakref.WeakValueDictionary() + +# Track whether we've already wrapped eager_attention_forward +_EAGER_ATTENTION_WRAPPED = False + +# Store the original function for restoration +_ORIGINAL_EAGER_ATTENTION_FORWARD: Optional[Callable] = None + + +def _setup_eager_attention_hook_wrapper() -> None: + """Wrap gemma2's eager_attention_forward to fire hook_rot_q and hook_rot_k. + + This function monkey-patches the module-level eager_attention_forward function + to intercept query and key tensors (which have already had rotary embeddings applied) + and fire the corresponding hooks on the registered bridge instance. + + This is safe to call multiple times - it will only wrap once. + """ + global _EAGER_ATTENTION_WRAPPED, _ORIGINAL_EAGER_ATTENTION_FORWARD + + if _EAGER_ATTENTION_WRAPPED: + return + + # Store the original function + _ORIGINAL_EAGER_ATTENTION_FORWARD = gemma2_module.eager_attention_forward + + def hooked_eager_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + **kwargs: Any, + ) -> tuple: + """Wrapped eager_attention_forward that fires rotary hooks. + + Args: + module: The HF attention module (used to look up the bridge) + query: Query tensor AFTER rotary embeddings applied + key: Key tensor AFTER rotary embeddings applied + value: Value tensor + attention_mask: Attention mask + **kwargs: Additional arguments (dropout, scaling, etc.) + + Returns: + Tuple of (attn_output, attn_weights) + """ + # Look up the bridge instance for this attention module + bridge = _ATTENTION_BRIDGE_REGISTRY.get(id(module)) + + if bridge is not None: + # Fire hook_rot_q and hook_rot_k with the post-rotary Q/K + if hasattr(bridge, "hook_rot_q"): + query = bridge.hook_rot_q(query) + if hasattr(bridge, "hook_rot_k"): + key = bridge.hook_rot_k(key) + + # Call the original function + return _ORIGINAL_EAGER_ATTENTION_FORWARD( + module, query, key, value, attention_mask, **kwargs + ) + + # Replace the module-level function + gemma2_module.eager_attention_forward = hooked_eager_attention_forward # type: ignore[assignment] + _EAGER_ATTENTION_WRAPPED = True + class PositionEmbeddingsAttentionBridge(AttentionBridge): """Attention bridge for models that require position embeddings (e.g., Gemma-3). @@ -49,7 +119,7 @@ def __init__( self.hook_cos = HookPoint() self.hook_sin = HookPoint() - def set_rotary_emb(self, rotary_emb): + def set_rotary_emb(self, rotary_emb: Any) -> None: """Set reference to the model's rotary embedding component. Args: @@ -57,6 +127,24 @@ def set_rotary_emb(self, rotary_emb): """ self._rotary_emb = rotary_emb + def set_original_component(self, component: torch.nn.Module) -> None: + """Set the original HF component and register for rotary hook firing. + + This overrides the base class method to also: + 1. Register this bridge in the global registry (for hook_rot_q/hook_rot_k) + 2. Set up the eager_attention_forward wrapper if not already done + + Args: + component: The HuggingFace attention module + """ + super().set_original_component(component) + + # Register this bridge instance so the wrapped eager_attention_forward can find it + _ATTENTION_BRIDGE_REGISTRY[id(component)] = self + + # Ensure the wrapper is set up + _setup_eager_attention_hook_wrapper() + def _apply_position_embedding_hooks(self, position_embeddings): """Apply hooks to position embeddings (cos, sin tuple). @@ -128,209 +216,3 @@ def get_random_inputs( inputs["position_embeddings"] = (cos, sin) inputs["attention_mask"] = None return inputs - - def setup_hook_compatibility(self) -> None: - """Setup hook compatibility transformations to match HookedTransformer behavior. - - For models requiring position embeddings (like Gemma with RoPE), this wraps the - HF attention forward to inject required position_embeddings and attention_mask. - This is necessary because some code paths may call the HF attention directly, - bypassing the bridge's forward method. - - Also sets up rotary embedding hooks (hook_rot_q, hook_rot_k) by monkey-patching - the HF attention's forward to intercept Q/K after rotary position embeddings are applied. - - Also sets up Q/K/V/Z hook reshaping like the base AttentionBridge. - - This is called during Bridge.__init__ and should always be run. - Note: This method is idempotent - can be called multiple times safely. - """ - if self.original_component is None: - return - original_hf_forward = self.original_component.forward - - def wrapped_forward(*args, **kwargs): - """Wrapper that injects position_embeddings, attention_mask, and rotary hooks.""" - if "position_embeddings" not in kwargs or kwargs["position_embeddings"] is None: - hidden_states = kwargs.get("hidden_states") or (args[0] if len(args) > 0 else None) - if hidden_states is not None: - batch_size, seq_len, _ = hidden_states.shape - device = hidden_states.device - dtype = hidden_states.dtype - if self._rotary_emb is not None: - head_dim = ( - self.config.head_dim - if self.config and hasattr(self.config, "head_dim") - else 256 - ) - num_heads = ( - self.config.num_attention_heads - if self.config and hasattr(self.config, "num_attention_heads") - else 4 - ) - dummy_qk = torch.randn( - 1, seq_len, num_heads, head_dim, device=device, dtype=dtype - ) - position_ids = torch.arange(seq_len, device=device).unsqueeze(0) - try: - position_embeddings = self._rotary_emb(dummy_qk, position_ids) - # Apply hooks to cos and sin - kwargs["position_embeddings"] = self._apply_position_embedding_hooks( - position_embeddings - ) - except Exception as e: - warnings.warn(f"Rotary embedding call failed: {e}, using fallback") - cos = torch.ones(1, seq_len, head_dim, device=device, dtype=dtype) - sin = torch.zeros(1, seq_len, head_dim, device=device, dtype=dtype) - # Apply hooks to fallback cos/sin - kwargs["position_embeddings"] = self._apply_position_embedding_hooks( - (cos, sin) - ) - else: - head_dim = ( - self.config.head_dim - if self.config and hasattr(self.config, "head_dim") - else 256 - ) - cos = torch.ones(1, seq_len, head_dim, device=device, dtype=dtype) - sin = torch.zeros(1, seq_len, head_dim, device=device, dtype=dtype) - # Apply hooks to fallback cos/sin - kwargs["position_embeddings"] = self._apply_position_embedding_hooks( - (cos, sin) - ) - if "attention_mask" not in kwargs: - kwargs["attention_mask"] = None - original_apply_rotary = gemma2_module.apply_rotary_pos_emb - - def hooked_apply_rotary(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): - """Wrapper that applies rotary hooks BEFORE rotation (matching HookedTransformer).""" - if hasattr(self, "hook_rot_q"): - q = self.hook_rot_q(q) - if hasattr(self, "hook_rot_k"): - k = self.hook_rot_k(k) - rotated_q, rotated_k = original_apply_rotary( - q, k, cos, sin, position_ids, unsqueeze_dim - ) - return (rotated_q, rotated_k) - - gemma2_module.apply_rotary_pos_emb = hooked_apply_rotary - try: - result = original_hf_forward(*args, **kwargs) - finally: - gemma2_module.apply_rotary_pos_emb = original_apply_rotary - return result - - self.original_component.forward = wrapped_forward - self._hf_forward_wrapped = True - if hasattr(self.config, "n_heads"): - self._setup_qkv_hook_reshaping() - - def forward(self, *args: Any, **kwargs: Any) -> Any: - """Simplified forward pass - minimal wrapping around original component. - - This override strips back the complex processing in AttentionBridge._process_output() - to do minimal wrapping: hook_in ā†’ delegate ā†’ hook_out. - - This ensures we match HuggingFace's exact output without any intermediate - pattern extraction or re-application logic. - - Args: - *args: Input arguments to pass to the original component - **kwargs: Input keyword arguments to pass to the original component - - Returns: - The output from the original component, with only input/output hooks applied - """ - if self.original_component is None: - raise RuntimeError( - f"Original component not set for {self.name}. Call set_original_component() first." - ) - target_dtype = None - try: - target_dtype = next(self.original_component.parameters()).dtype - except StopIteration: - pass - if "hidden_states" in kwargs: - hooked = self.hook_in(kwargs["hidden_states"]) - if ( - target_dtype is not None - and isinstance(hooked, torch.Tensor) - and hooked.is_floating_point() - ): - hooked = hooked.to(dtype=target_dtype) - kwargs["hidden_states"] = hooked - elif len(args) > 0: - hooked = self.hook_in(args[0]) - if ( - target_dtype is not None - and isinstance(hooked, torch.Tensor) - and hooked.is_floating_point() - ): - hooked = hooked.to(dtype=target_dtype) - args = (hooked,) + args[1:] - kwargs["hidden_states"] = args[0] - args = args[1:] - else: - raise ValueError("Gemma3Attention requires hidden_states argument") - if "position_embeddings" not in kwargs or kwargs["position_embeddings"] is None: - hidden_states = kwargs["hidden_states"] - batch_size, seq_len, _ = hidden_states.shape - device = hidden_states.device - dtype = hidden_states.dtype - if self._rotary_emb is not None: - head_dim = ( - self.config.head_dim - if self.config and hasattr(self.config, "head_dim") - else 256 - ) - num_heads = ( - self.config.num_attention_heads - if self.config and hasattr(self.config, "num_attention_heads") - else 4 - ) - dummy_qk = torch.randn(1, seq_len, num_heads, head_dim, device=device, dtype=dtype) - position_ids = torch.arange(seq_len, device=device).unsqueeze(0) - try: - position_embeddings = self._rotary_emb(dummy_qk, position_ids) - # Apply hooks to cos and sin - kwargs["position_embeddings"] = self._apply_position_embedding_hooks( - position_embeddings - ) - except Exception: - cos = torch.ones(1, seq_len, head_dim, device=device, dtype=dtype) - sin = torch.zeros(1, seq_len, head_dim, device=device, dtype=dtype) - # Apply hooks to fallback cos/sin - kwargs["position_embeddings"] = self._apply_position_embedding_hooks((cos, sin)) - else: - head_dim = ( - self.config.head_dim - if self.config and hasattr(self.config, "head_dim") - else 256 - ) - cos = torch.ones(1, seq_len, head_dim, device=device, dtype=dtype) - sin = torch.zeros(1, seq_len, head_dim, device=device, dtype=dtype) - # Apply hooks to fallback cos/sin - kwargs["position_embeddings"] = self._apply_position_embedding_hooks((cos, sin)) - if "attention_mask" not in kwargs: - kwargs["attention_mask"] = None - kwargs["output_attentions"] = True - output = self.original_component(*args, **kwargs) - if isinstance(output, tuple) and len(output) >= 2: - attn_weights = output[1] - if attn_weights is not None: - _ = self.hook_pattern(attn_weights) - with torch.no_grad(): - d_head = ( - self.config.head_dim - if self.config and hasattr(self.config, "head_dim") - else 256 - ) - attn_scores_approx = torch.log(attn_weights + 1e-10) * d_head**0.5 - _ = self.hook_attn_scores(attn_scores_approx) - if isinstance(output, tuple) and len(output) >= 2: - output = (self.hook_out(output[0]), output[1]) - elif isinstance(output, tuple) and len(output) == 1: - output = (self.hook_out(output[0]),) - else: - output = self.hook_out(output) - return output diff --git a/transformer_lens/model_bridge/generalized_components/rotary_embedding.py b/transformer_lens/model_bridge/generalized_components/rotary_embedding.py index 4b9442559e..c3bb813780 100644 --- a/transformer_lens/model_bridge/generalized_components/rotary_embedding.py +++ b/transformer_lens/model_bridge/generalized_components/rotary_embedding.py @@ -123,3 +123,31 @@ def forward(self, *args: Any, **kwargs: Any) -> Tuple[torch.Tensor, torch.Tensor else: # For unexpected tuple lengths, just pass through return output + + def get_dummy_inputs( + self, test_input: torch.Tensor, **kwargs: Any + ) -> tuple[tuple[Any, ...], dict[str, Any]]: + """Generate dummy inputs for rotary embedding forward method. + + Rotary embeddings typically expect (x, position_ids) where: + - x: input tensor [batch, seq, d_model] + - position_ids: position indices [batch, seq] + + Args: + test_input: Base test input tensor [batch, seq, d_model] + **kwargs: Additional context including position_ids + + Returns: + Tuple of (args, kwargs) for the rotary embedding forward method + """ + batch, seq_len, _ = test_input.shape + + # Get position_ids from kwargs, or generate default + position_ids = kwargs.get("position_ids") + if position_ids is None: + position_ids = ( + torch.arange(seq_len, device=test_input.device).unsqueeze(0).expand(batch, -1) + ) + + # Rotary embeddings expect (x, position_ids) + return (test_input, position_ids), {} diff --git a/transformer_lens/model_bridge/generalized_components/symbolic.py b/transformer_lens/model_bridge/generalized_components/symbolic.py new file mode 100644 index 0000000000..2825c5719a --- /dev/null +++ b/transformer_lens/model_bridge/generalized_components/symbolic.py @@ -0,0 +1,71 @@ +"""Symbolic bridge component. + +This module contains a bridge component that acts as a structural placeholder +when a model doesn't have a corresponding container for grouped subcomponents. +""" +from typing import Any, Dict, Optional + +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) + + +class SymbolicBridge(GeneralizedComponent): + """A placeholder bridge component for maintaining TransformerLens structure. + + This bridge is used when a model doesn't have a container component that exists + in the TransformerLens standard structure. For example, OPT has fc1/fc2 layers + directly on the block rather than inside an MLP container. + + When the model is set up, the subcomponents defined in this SymbolicBridge + are promoted to the parent component, allowing the TransformerLens structure + to be maintained while correctly mapping to the underlying model's architecture. + + Example usage: + # OPT doesn't have an "mlp" container - fc1/fc2 are on the block directly + "mlp": SymbolicBridge( + submodules={ + "in": LinearBridge(name="fc1"), + "out": LinearBridge(name="fc2"), + }, + ) + + # During setup, "in" and "out" will be accessible as: + # - blocks[i].mlp.in (pointing to blocks[i].fc1) + # - blocks[i].mlp.out (pointing to blocks[i].fc2) + + Attributes: + is_symbolic: Always True, indicates this is a structural placeholder. + """ + + is_symbolic: bool = True + + def __init__( + self, + submodules: Optional[Dict[str, GeneralizedComponent]] = None, + config: Optional[Any] = None, + ): + """Initialize the SymbolicBridge. + + Args: + submodules: Dictionary of submodules to register. These will be set up + using the parent's original_component as their context. + config: Optional configuration object + """ + # SymbolicBridge always has name=None since it doesn't map to a real component + super().__init__(name=None, config=config, submodules=submodules or {}) + + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Forward pass is not supported for SymbolicBridge. + + SymbolicBridge is a structural placeholder and should not be called directly. + The actual computation should go through the subcomponents which are set up + on the parent. + + Raises: + RuntimeError: Always, since SymbolicBridge should not be called directly. + """ + raise RuntimeError( + "SymbolicBridge is a structural placeholder and should not be called directly. " + "Use the subcomponents (e.g., 'in', 'out') instead." + ) diff --git a/transformer_lens/model_bridge/generalized_components/t5_block.py b/transformer_lens/model_bridge/generalized_components/t5_block.py index 1348742468..fcacd1679a 100644 --- a/transformer_lens/model_bridge/generalized_components/t5_block.py +++ b/transformer_lens/model_bridge/generalized_components/t5_block.py @@ -82,6 +82,7 @@ def patched_forward( use_cache=False, output_attentions=False, return_dict=True, + cache_position=None, **kwargs, ): """Patched T5 block forward with hooks.""" @@ -112,9 +113,12 @@ def patched_forward( past_key_value=self_attn_past_key_value, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, ) hidden_states = self_attention_outputs[0] - present_key_value_state = self_attention_outputs[1] if use_cache else None + # Keep self-attention outputs and relative position weights + # attention_outputs contains: (position_bias,) or (position_bias, attn_weights) + attention_outputs = self_attention_outputs[1:] hidden_states = self.hook_resid_mid(hidden_states) if is_decoder_block and encoder_hidden_states is not None: cross_attention_outputs = layers[1]( @@ -126,24 +130,25 @@ def patched_forward( past_key_value=cross_attn_past_key_value, use_cache=use_cache, output_attentions=output_attentions, + cache_position=cache_position, ) hidden_states = cross_attention_outputs[0] if hasattr(self, "hook_resid_mid2"): hidden_states = self.hook_resid_mid2(hidden_states) - if use_cache: - present_key_value_state = present_key_value_state + cross_attention_outputs[1] + # Keep cross-attention outputs and relative position weights + attention_outputs = attention_outputs + cross_attention_outputs[1:] ff_layer_idx = 2 if is_decoder_block else 1 feed_forward_outputs = layers[ff_layer_idx](hidden_states) - hidden_states = feed_forward_outputs[0] + # T5LayerFF returns a tensor, not a tuple + if isinstance(feed_forward_outputs, tuple): + hidden_states = feed_forward_outputs[0] + else: + hidden_states = feed_forward_outputs hidden_states = self.hook_out(hidden_states) outputs: tuple[Any, ...] = (hidden_states,) - if use_cache: - outputs = outputs + (present_key_value_state,) - if output_attentions: - outputs = outputs + (self_attention_outputs[2],) - if is_decoder_block and encoder_hidden_states is not None: - outputs = outputs + (cross_attention_outputs[2],) - return outputs + # Return: hidden-states, (self-attention position bias), (self-attention weights), + # (cross-attention position bias), (cross-attention weights) + return outputs + attention_outputs self.original_component.forward = types.MethodType(patched_forward, self.original_component) diff --git a/transformer_lens/model_bridge/generalized_components/unembedding.py b/transformer_lens/model_bridge/generalized_components/unembedding.py index dd0d820091..ca17cab501 100644 --- a/transformer_lens/model_bridge/generalized_components/unembedding.py +++ b/transformer_lens/model_bridge/generalized_components/unembedding.py @@ -101,6 +101,15 @@ def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor: ): hidden_states = hidden_states.to(dtype=target_dtype) output = self.original_component(hidden_states, **kwargs) + + # Apply logit soft-capping if configured (e.g., for Gemma-2) + if self.config is not None and hasattr(self.config, "output_logits_soft_cap"): + output_logits_soft_cap = self.config.output_logits_soft_cap + if output_logits_soft_cap is not None and output_logits_soft_cap > 0: + output = output / output_logits_soft_cap + output = torch.tanh(output) + output = output * output_logits_soft_cap + output = self.hook_out(output) return output diff --git a/transformer_lens/model_bridge/supported_architectures/bloom.py b/transformer_lens/model_bridge/supported_architectures/bloom.py index 997e2237e3..87984c1301 100644 --- a/transformer_lens/model_bridge/supported_architectures/bloom.py +++ b/transformer_lens/model_bridge/supported_architectures/bloom.py @@ -10,11 +10,11 @@ ) from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( - BlockBridge, + BloomAttentionBridge, + BloomBlockBridge, + BloomMLPBridge, EmbeddingBridge, - JointQKVAttentionBridge, LinearBridge, - MLPBridge, NormalizationBridge, UnembeddingBridge, ) @@ -71,12 +71,13 @@ def __init__(self, cfg: Any) -> None: "embed_ln": NormalizationBridge( name="transformer.word_embeddings_layernorm", config=self.cfg ), - "blocks": BlockBridge( + "blocks": BloomBlockBridge( name="transformer.h", + config=self.cfg, submodules={ "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg), "ln2": NormalizationBridge(name="post_attention_layernorm", config=self.cfg), - "attn": JointQKVAttentionBridge( + "attn": BloomAttentionBridge( name="self_attention", config=self.cfg, split_qkv_matrix=self.split_qkv_matrix, @@ -85,7 +86,7 @@ def __init__(self, cfg: Any) -> None: "o": LinearBridge(name="dense"), }, ), - "mlp": MLPBridge( + "mlp": BloomMLPBridge( name="mlp", submodules={ "in": LinearBridge(name="dense_h_to_4h"), diff --git a/transformer_lens/model_bridge/supported_architectures/gemma1.py b/transformer_lens/model_bridge/supported_architectures/gemma1.py index c1318ffc28..2b36fa408a 100644 --- a/transformer_lens/model_bridge/supported_architectures/gemma1.py +++ b/transformer_lens/model_bridge/supported_architectures/gemma1.py @@ -2,17 +2,24 @@ from typing import Any -from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion +from transformer_lens.conversion_utils.conversion_steps import ( + ArithmeticTensorConversion, + RearrangeTensorConversion, + TransposeTensorConversion, +) +from transformer_lens.conversion_utils.conversion_steps.arithmetic_tensor_conversion import ( + OperationTypes, +) from transformer_lens.conversion_utils.param_processing_conversion import ( ParamProcessingConversion, ) from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( - AttentionBridge, BlockBridge, EmbeddingBridge, GatedMLPBridge, LinearBridge, + PositionEmbeddingsAttentionBridge, RMSNormalizationBridge, RotaryEmbeddingBridge, UnembeddingBridge, @@ -36,16 +43,17 @@ def __init__(self, cfg: Any) -> None: # Gemma models were not trained with BOS tokens self.cfg.default_prepend_bos = False self.cfg.uses_rms_norm = True + # Gemma models use (1.0 + weight) in RMSNorm instead of just weight + # See: https://github.com/huggingface/transformers/pull/29402 + self.cfg.rmsnorm_uses_offset = True self.weight_processing_conversions = { - # Gemma1 scales embeddings by sqrt(d_model) - "embed.e": ParamProcessingConversion( - tensor_conversion=RearrangeTensorConversion( - "d_vocab d_model -> d_vocab d_model", - scale=self.cfg.d_model**0.5, - ), - source_key="model.embed_tokens.weight", - ), + # Note: Gemma1 scales embeddings by sqrt(d_model) in the forward pass. + # This is handled in setup_hook_compatibility() which applies the scaling + # to hook_embed output at runtime, matching HuggingFace's behavior. + # We do NOT scale the stored weights here. + # + # Attention weight conversions "blocks.{i}.attn.q.weight": ParamProcessingConversion( tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=self.cfg.n_heads), ), @@ -58,6 +66,31 @@ def __init__(self, cfg: Any) -> None: "blocks.{i}.attn.o.weight": ParamProcessingConversion( tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=self.cfg.n_heads), ), + # RMSNorm weight conversions - Gemma adds 1.0 to weights before applying + # See: https://github.com/huggingface/transformers/pull/29402 + "blocks.{i}.ln1.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + "blocks.{i}.ln2.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + "ln_final.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + # MLP weight conversions - transpose from [out, in] to [in, out] + "blocks.{i}.mlp.gate.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), + "blocks.{i}.mlp.in.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), + "blocks.{i}.mlp.out.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), + # Unembed weight conversion - transpose from [vocab, d_model] to [d_model, vocab] + "unembed.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), } self.component_mapping = { @@ -68,7 +101,7 @@ def __init__(self, cfg: Any) -> None: submodules={ "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), - "attn": AttentionBridge( + "attn": PositionEmbeddingsAttentionBridge( name="self_attn", config=self.cfg, submodules={ @@ -77,6 +110,8 @@ def __init__(self, cfg: Any) -> None: "v": LinearBridge(name="v_proj"), "o": LinearBridge(name="o_proj"), }, + requires_attention_mask=True, + requires_position_embeddings=True, ), "mlp": GatedMLPBridge( name="mlp", @@ -92,3 +127,73 @@ def __init__(self, cfg: Any) -> None: "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), "unembed": UnembeddingBridge(name="lm_head"), } + + def setup_hook_compatibility(self, bridge: Any) -> None: + """Setup hook compatibility for Gemma1 models. + + Gemma1 scales embeddings by sqrt(d_model) in its forward pass, + but the HuggingFace embed_tokens layer doesn't include this scaling. + We need to apply it to hook_embed to match HookedTransformer behavior. + + Args: + bridge: The TransformerBridge instance + """ + from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import ( + BaseTensorConversion, + ) + + class EmbeddingScaleConversion(BaseTensorConversion): + """Scale embeddings by sqrt(d_model) for Gemma models.""" + + def __init__(self, scale: float): + super().__init__() + self.scale = scale + + def handle_conversion(self, input_value: Any, *full_context: Any) -> Any: + """Scale the embedding output.""" + return input_value * self.scale + + def revert(self, input_value: Any, *full_context: Any) -> Any: + """Unscale the embedding output (for user modifications).""" + return input_value / self.scale + + # Apply scaling to embed.hook_out + if hasattr(bridge, "embed") and hasattr(bridge.embed, "hook_out"): + scale_factor = self.cfg.d_model**0.5 + bridge.embed.hook_out.hook_conversion = EmbeddingScaleConversion(scale_factor) + + def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: + """Set up rotary embedding references for Gemma1 component testing. + + Gemma1 uses RoPE (Rotary Position Embeddings). We set the rotary_emb reference + on all attention bridge instances for component testing. + + Args: + hf_model: The HuggingFace Gemma1 model instance + bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances) + """ + # Get rotary embedding instance from the model + rotary_emb = hf_model.model.rotary_emb + + # Force HF model to use "eager" attention to match bridge implementation + # Bridge uses "eager" to support output_attentions for hook compatibility + # SDPA and eager are mathematically equivalent but have numerical differences + if hasattr(hf_model, "config") and hasattr(hf_model.config, "_attn_implementation"): + hf_model.config._attn_implementation = "eager" + + # Also set on all attention layers + if hasattr(hf_model, "model") and hasattr(hf_model.model, "layers"): + for layer in hf_model.model.layers: + if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "config"): + layer.self_attn.config._attn_implementation = "eager" + + # Set rotary_emb on actual bridge instances in bridge_model if available + if bridge_model is not None and hasattr(bridge_model, "blocks"): + # Set on each layer's actual attention bridge instance + for block in bridge_model.blocks: + if hasattr(block, "attn"): + block.attn.set_rotary_emb(rotary_emb) + + # Also set on the template for get_generalized_component() calls + attn_bridge = self.get_generalized_component("blocks.0.attn") + attn_bridge.set_rotary_emb(rotary_emb) diff --git a/transformer_lens/model_bridge/supported_architectures/gemma2.py b/transformer_lens/model_bridge/supported_architectures/gemma2.py index f6a3e5154c..685f2ee13e 100644 --- a/transformer_lens/model_bridge/supported_architectures/gemma2.py +++ b/transformer_lens/model_bridge/supported_architectures/gemma2.py @@ -1,8 +1,18 @@ """Gemma2 architecture adapter.""" -from typing import Any +from typing import TYPE_CHECKING, Any -from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion +if TYPE_CHECKING: + pass + +from transformer_lens.conversion_utils.conversion_steps import ( + ArithmeticTensorConversion, + RearrangeTensorConversion, + TransposeTensorConversion, +) +from transformer_lens.conversion_utils.conversion_steps.arithmetic_tensor_conversion import ( + OperationTypes, +) from transformer_lens.conversion_utils.param_processing_conversion import ( ParamProcessingConversion, ) @@ -36,18 +46,25 @@ def __init__(self, cfg: Any) -> None: # Gemma models were not trained with BOS tokens # self.cfg.default_prepend_bos = False self.cfg.uses_rms_norm = True + # Gemma models use (1.0 + weight) in RMSNorm instead of just weight + # See: https://github.com/huggingface/transformers/pull/29402 + self.cfg.rmsnorm_uses_offset = True + + # Gemma2 uses logit softcapping + if hasattr(self.cfg, "final_logit_softcapping"): + self.cfg.output_logits_soft_cap = self.cfg.final_logit_softcapping + if hasattr(self.cfg, "attn_logit_softcapping"): + self.cfg.attn_scores_soft_cap = self.cfg.attn_logit_softcapping # Note: n_key_value_heads is now automatically mapped from num_key_value_heads # by map_default_transformer_lens_config() in sources/transformers.py self.weight_processing_conversions = { - # Gemma2 scales embeddings by sqrt(d_model) - "embed.e": ParamProcessingConversion( - tensor_conversion=RearrangeTensorConversion( - "d_vocab d_model -> d_vocab d_model", - scale=self.cfg.d_model**0.5, + # Embedding weight scaling - Gemma models scale embeddings by sqrt(d_model) + "embed.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion( + OperationTypes.MULTIPLICATION, self.cfg.d_model**0.5 ), - source_key="model.embed_tokens.weight", ), "blocks.{i}.attn.q.weight": ParamProcessingConversion( tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=self.cfg.n_heads), @@ -67,6 +84,37 @@ def __init__(self, cfg: Any) -> None: "blocks.{i}.attn.o.weight": ParamProcessingConversion( tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=self.cfg.n_heads), ), + # RMSNorm weight conversions - Gemma adds 1.0 to weights before applying + # See: https://github.com/huggingface/transformers/pull/29402 + "blocks.{i}.ln1.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + "blocks.{i}.ln1_post.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + "blocks.{i}.ln2.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + "blocks.{i}.ln2_post.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + "ln_final.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + # MLP weight conversions - transpose from [out, in] to [in, out] + "blocks.{i}.mlp.gate.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), + "blocks.{i}.mlp.in.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), + "blocks.{i}.mlp.out.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), + # # Unembed weight conversion - transpose from [vocab, d_model] to [d_model, vocab] + "unembed.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), } self.component_mapping = { @@ -74,6 +122,7 @@ def __init__(self, cfg: Any) -> None: "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), "blocks": BlockBridge( name="model.layers", + config=self.cfg, submodules={ # Gemma 2 uses RMSNorm for all normalization layers "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), @@ -96,6 +145,8 @@ def __init__(self, cfg: Any) -> None: "v": LinearBridge(name="v_proj"), "o": LinearBridge(name="o_proj"), }, + requires_attention_mask=True, + requires_position_embeddings=True, ), "mlp": GatedMLPBridge( name="mlp", @@ -109,60 +160,25 @@ def __init__(self, cfg: Any) -> None: }, ), "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), - "unembed": UnembeddingBridge(name="lm_head"), + "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), } def setup_hook_compatibility(self, bridge: Any) -> None: """Setup hook compatibility for Gemma2 models. - Gemma2 scales embeddings by sqrt(d_model) in its forward pass, - but the HuggingFace embed_tokens layer doesn't include this scaling. - We need to apply it to hook_embed to match HookedTransformer behavior. + Gemma2 scales embeddings by sqrt(d_model). The weights are pre-scaled via + preprocess_weights(), but we still need to apply the scaling conversion to + the hook output for proper hook functionality (so user modifications are + correctly scaled/unscaled). Args: bridge: The TransformerBridge instance """ - from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import ( - BaseTensorConversion, - ) - - class EmbeddingScaleConversion(BaseTensorConversion): - """Scale embeddings by sqrt(d_model) for Gemma models. - - This only applies when NOT using processed weights, since processed - weights have the scaling baked into the embedding matrix itself. - """ - - def __init__(self, scale: float, embed_component): - super().__init__() - self.scale = scale - self.embed_component = embed_component - - def handle_conversion(self, input_value, *full_context): - """Scale the embedding output if not using processed weights.""" - # Skip scaling if using processed weights (they're already scaled) - if ( - hasattr(self.embed_component, "_use_processed_weights") - and self.embed_component._use_processed_weights - ): - return input_value - return input_value * self.scale - - def revert(self, input_value, *full_context): - """Unscale the embedding output (for user modifications).""" - # Skip unscaling if using processed weights - if ( - hasattr(self.embed_component, "_use_processed_weights") - and self.embed_component._use_processed_weights - ): - return input_value - return input_value / self.scale - - # Apply scaling to embed.hook_out + # Apply embedding scaling conversion to hook output if hasattr(bridge, "embed") and hasattr(bridge.embed, "hook_out"): scale_factor = self.cfg.d_model**0.5 - bridge.embed.hook_out.hook_conversion = EmbeddingScaleConversion( - scale_factor, bridge.embed + bridge.embed.hook_out.hook_conversion = ArithmeticTensorConversion( + OperationTypes.MULTIPLICATION, scale_factor ) def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: @@ -205,27 +221,3 @@ def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> No # Also set on the template for get_generalized_component() calls attn_bridge = self.get_generalized_component("blocks.0.attn") attn_bridge.set_rotary_emb(rotary_emb) - - def preprocess_weights(self, state_dict: dict[str, Any]) -> dict[str, Any]: - """Apply Gemma2-specific weight transformations before ProcessWeights. - - Gemma2 models scale embeddings by sqrt(d_model) in their forward pass. - We bake this scaling into the embedding weights to avoid applying it every forward pass. - - Args: - state_dict: The state dictionary with HuggingFace format keys - - Returns: - The modified state dictionary with scaled embedding weights - """ - from transformer_lens.weight_processing import ProcessWeights - - # Get the HF key for the embedding weight - embed_key = ProcessWeights._get_param_key("embed.W_E", self) - - if embed_key in state_dict: - # Scale embeddings by sqrt(d_model) to match Gemma2's forward pass behavior - scale_factor = self.cfg.d_model**0.5 - state_dict[embed_key] = state_dict[embed_key] * scale_factor - - return state_dict diff --git a/transformer_lens/model_bridge/supported_architectures/gemma3.py b/transformer_lens/model_bridge/supported_architectures/gemma3.py index f906b1b36c..76ee59b3b8 100644 --- a/transformer_lens/model_bridge/supported_architectures/gemma3.py +++ b/transformer_lens/model_bridge/supported_architectures/gemma3.py @@ -3,7 +3,14 @@ from typing import Any -from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion +from transformer_lens.conversion_utils.conversion_steps import ( + ArithmeticTensorConversion, + RearrangeTensorConversion, + TransposeTensorConversion, +) +from transformer_lens.conversion_utils.conversion_steps.arithmetic_tensor_conversion import ( + OperationTypes, +) from transformer_lens.conversion_utils.param_processing_conversion import ( ParamProcessingConversion, ) @@ -42,14 +49,11 @@ def __init__(self, cfg: Any) -> None: self.cfg.attn_implementation = "eager" self.weight_processing_conversions = { - # Gemma3 scales embeddings by sqrt(d_model) - "embed": ParamProcessingConversion( - tensor_conversion=RearrangeTensorConversion( - "d_vocab d_model -> d_vocab d_model", - scale=self.cfg.d_model**0.5, - ), - source_key="model.embed_tokens.weight", - ), + # Note: Gemma3 scales embeddings by sqrt(d_model) in the forward pass. + # This is handled in setup_hook_compatibility() which applies the scaling + # to hook_embed output at runtime, matching HuggingFace's behavior. + # We do NOT scale the stored weights here. + # # Q/K/V weight conversions "blocks.{i}.attn.q.weight": ParamProcessingConversion( tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=self.cfg.n_heads), @@ -77,6 +81,44 @@ def __init__(self, cfg: Any) -> None: "blocks.{i}.attn.o.weight": ParamProcessingConversion( tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=self.cfg.n_heads), ), + # RMSNorm weight conversions - Gemma adds 1.0 to weights before applying + # See: https://github.com/huggingface/transformers/pull/29402 + "blocks.{i}.ln1.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + "blocks.{i}.ln1_post.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + "blocks.{i}.ln2.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + "blocks.{i}.ln2_post.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + "ln_final.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + # Gemma-3 also has q_norm and k_norm in attention + "blocks.{i}.attn.q_norm.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + "blocks.{i}.attn.k_norm.weight": ParamProcessingConversion( + tensor_conversion=ArithmeticTensorConversion(OperationTypes.ADDITION, 1.0), + ), + # MLP weight conversions - transpose from [out, in] to [in, out] + "blocks.{i}.mlp.gate.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), + "blocks.{i}.mlp.in.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), + "blocks.{i}.mlp.out.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), + # Unembed weight conversion - transpose from [vocab, d_model] to [d_model, vocab] + "unembed.weight": ParamProcessingConversion( + tensor_conversion=TransposeTensorConversion(), + ), # Note: Gemma-3 does NOT have biases on attention projections (q/k/v/o_proj.bias are all None) # No bias conversions needed } @@ -84,10 +126,8 @@ def __init__(self, cfg: Any) -> None: # Set up component mapping with actual bridge instances self.component_mapping = { "embed": EmbeddingBridge(name="model.embed_tokens"), - "rotary_emb_local": RotaryEmbeddingBridge( - name="model.rotary_emb_local", config=self.cfg - ), - "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), + "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), + "rotary_emb_local": RotaryEmbeddingBridge(name="model.rotary_emb_local"), "blocks": BlockBridge( name="model.layers", submodules={ @@ -129,6 +169,40 @@ def __init__(self, cfg: Any) -> None: "unembed": UnembeddingBridge(name="lm_head"), } + def setup_hook_compatibility(self, bridge: Any) -> None: + """Setup hook compatibility for Gemma3 models. + + Gemma3 scales embeddings by sqrt(d_model) in its forward pass, + but the HuggingFace embed_tokens layer doesn't include this scaling. + We need to apply it to hook_embed to match HookedTransformer behavior. + + Args: + bridge: The TransformerBridge instance + """ + from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import ( + BaseTensorConversion, + ) + + class EmbeddingScaleConversion(BaseTensorConversion): + """Scale embeddings by sqrt(d_model) for Gemma models.""" + + def __init__(self, scale: float): + super().__init__() + self.scale = scale + + def handle_conversion(self, input_value: Any, *full_context: Any) -> Any: + """Scale the embedding output.""" + return input_value * self.scale + + def revert(self, input_value: Any, *full_context: Any) -> Any: + """Unscale the embedding output (for user modifications).""" + return input_value / self.scale + + # Apply scaling to embed.hook_out + if hasattr(bridge, "embed") and hasattr(bridge.embed, "hook_out"): + scale_factor = self.cfg.d_model**0.5 + bridge.embed.hook_out.hook_conversion = EmbeddingScaleConversion(scale_factor) + def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: """Set up rotary embedding references and native autograd for Gemma-3 component testing. diff --git a/transformer_lens/model_bridge/supported_architectures/llama.py b/transformer_lens/model_bridge/supported_architectures/llama.py index 4690e6562a..7c9ae0a720 100644 --- a/transformer_lens/model_bridge/supported_architectures/llama.py +++ b/transformer_lens/model_bridge/supported_architectures/llama.py @@ -8,11 +8,11 @@ ) from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( - AttentionBridge, BlockBridge, EmbeddingBridge, GatedMLPBridge, LinearBridge, + PositionEmbeddingsAttentionBridge, RMSNormalizationBridge, RotaryEmbeddingBridge, UnembeddingBridge, @@ -99,7 +99,7 @@ def __init__(self, cfg: Any) -> None: submodules={ "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), - "attn": AttentionBridge( + "attn": PositionEmbeddingsAttentionBridge( name="self_attn", config=self.cfg, submodules={ @@ -108,6 +108,8 @@ def __init__(self, cfg: Any) -> None: "v": LinearBridge(name="v_proj"), "o": LinearBridge(name="o_proj"), }, + requires_attention_mask=True, + requires_position_embeddings=True, ), "mlp": GatedMLPBridge( name="mlp", @@ -121,5 +123,29 @@ def __init__(self, cfg: Any) -> None: }, ), "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), - "unembed": UnembeddingBridge(name="lm_head"), + "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), } + + def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: + """Set up rotary embedding references for Llama component testing. + + Llama uses RoPE (Rotary Position Embeddings). We set the rotary_emb reference + on all attention bridge instances for component testing. + + Args: + hf_model: The HuggingFace Llama model instance + bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances) + """ + # Get rotary embedding instance from the model + rotary_emb = hf_model.model.rotary_emb + + # Set rotary_emb on actual bridge instances in bridge_model if available + if bridge_model is not None and hasattr(bridge_model, "blocks"): + # Set on each layer's actual attention bridge instance + for block in bridge_model.blocks: + if hasattr(block, "attn"): + block.attn.set_rotary_emb(rotary_emb) + + # Also set on the template for get_generalized_component() calls + attn_bridge = self.get_generalized_component("blocks.0.attn") + attn_bridge.set_rotary_emb(rotary_emb) diff --git a/transformer_lens/model_bridge/supported_architectures/neox.py b/transformer_lens/model_bridge/supported_architectures/neox.py index dae9e42ce7..634afbeaef 100644 --- a/transformer_lens/model_bridge/supported_architectures/neox.py +++ b/transformer_lens/model_bridge/supported_architectures/neox.py @@ -18,10 +18,11 @@ from transformer_lens.model_bridge.generalized_components import ( BlockBridge, EmbeddingBridge, - JointQKVAttentionBridge, + JointQKVPositionEmbeddingsAttentionBridge, LinearBridge, MLPBridge, NormalizationBridge, + RotaryEmbeddingBridge, UnembeddingBridge, ) @@ -133,18 +134,17 @@ def __init__(self, cfg: Any) -> None: self.component_mapping = { "embed": EmbeddingBridge(name="gpt_neox.embed_in"), - "rotary_emb": EmbeddingBridge(name="gpt_neox.rotary_emb"), + "rotary_emb": RotaryEmbeddingBridge(name="gpt_neox.rotary_emb"), "blocks": BlockBridge( name="gpt_neox.layers", submodules={ "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg), "ln2": NormalizationBridge(name="post_attention_layernorm", config=self.cfg), - "attn": JointQKVAttentionBridge( + "attn": JointQKVPositionEmbeddingsAttentionBridge( name="attention", config=self.cfg, split_qkv_matrix=self.split_qkv_matrix, - requires_attention_mask=True, # GPTNeoX/Pythia requires attention_mask - requires_position_embeddings=True, # GPTNeoX/Pythia requires position_embeddings (rotary) + requires_attention_mask=True, # GPTNeoX/StableLM requires attention_mask submodules={ "qkv": LinearBridge(name="query_key_value"), "o": LinearBridge(name="dense"), @@ -167,49 +167,82 @@ def split_qkv_matrix( self, original_attention_component: Any ) -> tuple[torch.nn.Linear, torch.nn.Linear, torch.nn.Linear]: """Split the QKV matrix into separate linear transformations. + + GPT-NeoX/StableLM uses an interleaved QKV format where the weights are stored as + [Q_h0, K_h0, V_h0, Q_h1, K_h1, V_h1, ...] - i.e., Q, K, V are interleaved per head. + + The weight shape is [n_heads * 3 * d_head, d_model] and the output is reshaped + by HuggingFace as [batch, seq, n_heads, 3*d_head] then split on the last dim. + Args: - attention_component: The original attention layer component + original_attention_component: The original attention layer component + Returns: Tuple of nn.Linear modules for Q, K, and V transformations """ - - # Keep mypy happy assert original_attention_component is not None assert original_attention_component.query_key_value is not None qkv_weights = original_attention_component.query_key_value.weight - - # Keep mypy happy assert isinstance(qkv_weights, torch.Tensor) - # Original qkv_weights shape: [3 * d_model, d_model] -> Transposed to [d_model, 3 * d_model] - # Split into three equal parts along dimension 1 to get Q, K, V weights - W_Q, W_K, W_V = torch.tensor_split(qkv_weights.T, 3, dim=1) + n_heads = self.cfg.n_heads + d_head = self.cfg.d_head + d_model = self.cfg.d_model - qkv_bias = original_attention_component.query_key_value.bias + # Weight shape: [n_heads * 3 * d_head, d_model] + # Reshape to [n_heads, 3 * d_head, d_model] to access Q, K, V per head + W_reshaped = qkv_weights.view(n_heads, 3 * d_head, d_model) - # Keep mypy happy + # Extract Q, K, V weights for all heads and flatten back + W_Q = W_reshaped[:, :d_head, :].reshape(n_heads * d_head, d_model) + W_K = W_reshaped[:, d_head : 2 * d_head, :].reshape(n_heads * d_head, d_model) + W_V = W_reshaped[:, 2 * d_head :, :].reshape(n_heads * d_head, d_model) + + # Handle bias - same interleaved format + qkv_bias = original_attention_component.query_key_value.bias assert isinstance(qkv_bias, torch.Tensor) - # Original qkv_bias shape: [n_heads * 3 * d_head] - # Reshape to [3, n_heads * d_head] to split by Q, K, V - qkv_bias = qkv_bias.reshape(3, self.cfg.n_heads * self.cfg.d_head) - b_Q, b_K, b_V = qkv_bias[0, :], qkv_bias[1, :], qkv_bias[2, :] + # Bias shape: [n_heads * 3 * d_head] + # Reshape to [n_heads, 3 * d_head] to access Q, K, V per head + b_reshaped = qkv_bias.view(n_heads, 3 * d_head) + b_Q = b_reshaped[:, :d_head].reshape(n_heads * d_head) + b_K = b_reshaped[:, d_head : 2 * d_head].reshape(n_heads * d_head) + b_V = b_reshaped[:, 2 * d_head :].reshape(n_heads * d_head) # Create nn.Linear modules - # After tensor_split, W_Q, W_K, W_V shapes are [d_model, d_model] ([in_features, out_features]) - # nn.Linear expects weight shape [out_features, in_features] - # So we need to transpose the weights - W_Q_transformation = torch.nn.Linear(W_Q.shape[0], W_Q.shape[1], bias=True) - W_Q_transformation.weight = torch.nn.Parameter(W_Q.T) + # Weight shape for nn.Linear is [out_features, in_features] + W_Q_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True) + W_Q_transformation.weight = torch.nn.Parameter(W_Q) W_Q_transformation.bias = torch.nn.Parameter(b_Q) - W_K_transformation = torch.nn.Linear(W_K.shape[0], W_K.shape[1], bias=True) - W_K_transformation.weight = torch.nn.Parameter(W_K.T) + W_K_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True) + W_K_transformation.weight = torch.nn.Parameter(W_K) W_K_transformation.bias = torch.nn.Parameter(b_K) - W_V_transformation = torch.nn.Linear(W_V.shape[0], W_V.shape[1], bias=True) - W_V_transformation.weight = torch.nn.Parameter(W_V.T) + W_V_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True) + W_V_transformation.weight = torch.nn.Parameter(W_V) W_V_transformation.bias = torch.nn.Parameter(b_V) return W_Q_transformation, W_K_transformation, W_V_transformation + + def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: + """Set up rotary embedding references for GPT-NeoX/StableLM component testing. + + GPT-NeoX models use RoPE (Rotary Position Embeddings) which need to be + set on all attention bridge instances for component testing. + + Args: + hf_model: The HuggingFace GPT-NeoX model instance + bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances) + """ + # Get rotary embedding instance from model level + # In GPT-NeoX/StableLM, rotary_emb is at the model level + rotary_emb = hf_model.gpt_neox.rotary_emb + + # Set rotary_emb on actual bridge instances in bridge_model if available + if bridge_model is not None and hasattr(bridge_model, "blocks"): + # Set on each layer's actual attention bridge instance + for block in bridge_model.blocks: + if hasattr(block, "attn"): + block.attn.set_rotary_emb(rotary_emb) diff --git a/transformer_lens/model_bridge/supported_architectures/opt.py b/transformer_lens/model_bridge/supported_architectures/opt.py index bd19600dfb..e90c031e60 100644 --- a/transformer_lens/model_bridge/supported_architectures/opt.py +++ b/transformer_lens/model_bridge/supported_architectures/opt.py @@ -12,9 +12,9 @@ BlockBridge, EmbeddingBridge, LinearBridge, - MLPBridge, NormalizationBridge, PosEmbedBridge, + SymbolicBridge, UnembeddingBridge, ) @@ -60,6 +60,8 @@ def __init__(self, cfg: Any) -> None: "attn": AttentionBridge( name="self_attn", config=self.cfg, + requires_attention_mask=True, # OPT requires attention_mask + attention_mask_4d=True, # OPT expects 4D mask [batch, 1, tgt_len, src_len] submodules={ "q": LinearBridge(name="q_proj"), "k": LinearBridge(name="k_proj"), @@ -68,9 +70,10 @@ def __init__(self, cfg: Any) -> None: }, ), "ln2": NormalizationBridge(name="final_layer_norm", config=self.cfg), - "mlp": MLPBridge( - name=None, # No MLP container; fc1/fc2 are on block - config=self.cfg, # Pass config for activation function + # OPT has fc1/fc2 directly on the block, not in an MLP container. + # Use SymbolicBridge to maintain TransformerLens structure while + # correctly mapping to the underlying architecture. + "mlp": SymbolicBridge( submodules={ "in": LinearBridge(name="fc1"), "out": LinearBridge(name="fc2"), diff --git a/transformer_lens/model_bridge/supported_architectures/phi.py b/transformer_lens/model_bridge/supported_architectures/phi.py index ae2ad141ce..4925d1deaf 100644 --- a/transformer_lens/model_bridge/supported_architectures/phi.py +++ b/transformer_lens/model_bridge/supported_architectures/phi.py @@ -8,12 +8,13 @@ ) from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( - AttentionBridge, BlockBridge, EmbeddingBridge, LinearBridge, MLPBridge, NormalizationBridge, + PositionEmbeddingsAttentionBridge, + RotaryEmbeddingBridge, UnembeddingBridge, ) @@ -82,12 +83,12 @@ def __init__(self, cfg: Any) -> None: # Set up component mapping self.component_mapping = { "embed": EmbeddingBridge(name="model.embed_tokens"), - "rotary_emb": EmbeddingBridge(name="model.rotary_emb"), + "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), "blocks": BlockBridge( name="model.layers", submodules={ "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg), - "attn": AttentionBridge( + "attn": PositionEmbeddingsAttentionBridge( name="self_attn", config=self.cfg, submodules={ @@ -96,6 +97,8 @@ def __init__(self, cfg: Any) -> None: "v": LinearBridge(name="v_proj"), "o": LinearBridge(name="dense"), }, + requires_attention_mask=True, + requires_position_embeddings=True, ), "mlp": MLPBridge( name="mlp", @@ -109,3 +112,44 @@ def __init__(self, cfg: Any) -> None: "ln_final": NormalizationBridge(name="model.final_layernorm", config=self.cfg), "unembed": UnembeddingBridge(name="lm_head"), } + + def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: + """Set up rotary embedding references for Phi component testing. + + Phi uses RoPE (Rotary Position Embeddings). We set the rotary_emb reference + on all attention bridge instances for component testing. + + Args: + hf_model: The HuggingFace Phi model instance + bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances) + """ + # Get rotary embedding instance from the model + # Phi models have rotary_emb at model.model.rotary_emb + if hasattr(hf_model, "model") and hasattr(hf_model.model, "rotary_emb"): + rotary_emb = hf_model.model.rotary_emb + else: + # Fallback: try to get from first layer + if hasattr(hf_model, "model") and hasattr(hf_model.model, "layers"): + if len(hf_model.model.layers) > 0: + first_layer = hf_model.model.layers[0] + if hasattr(first_layer, "self_attn") and hasattr( + first_layer.self_attn, "rotary_emb" + ): + rotary_emb = first_layer.self_attn.rotary_emb + else: + return # Can't find rotary_emb + else: + return + else: + return + + # Set rotary_emb on actual bridge instances in bridge_model if available + if bridge_model is not None and hasattr(bridge_model, "blocks"): + # Set on each layer's actual attention bridge instance + for block in bridge_model.blocks: + if hasattr(block, "attn"): + block.attn.set_rotary_emb(rotary_emb) + + # Also set on the template for get_generalized_component() calls + attn_bridge = self.get_generalized_component("blocks.0.attn") + attn_bridge.set_rotary_emb(rotary_emb) diff --git a/transformer_lens/model_bridge/supported_architectures/pythia.py b/transformer_lens/model_bridge/supported_architectures/pythia.py index 5a6dd9afc3..3f9c2c89aa 100644 --- a/transformer_lens/model_bridge/supported_architectures/pythia.py +++ b/transformer_lens/model_bridge/supported_architectures/pythia.py @@ -18,10 +18,11 @@ from transformer_lens.model_bridge.generalized_components import ( BlockBridge, EmbeddingBridge, - JointQKVAttentionBridge, + JointQKVPositionEmbeddingsAttentionBridge, LinearBridge, MLPBridge, NormalizationBridge, + RotaryEmbeddingBridge, UnembeddingBridge, ) @@ -124,16 +125,15 @@ def __init__(self, cfg: Any) -> None: ), } - # NOTE: rotary_emb is not included in component_mapping because it's an internal helper - # used by attention, not a standalone testable component self.component_mapping = { "embed": EmbeddingBridge(name="gpt_neox.embed_in"), + "rotary_emb": RotaryEmbeddingBridge(name="gpt_neox.rotary_emb", config=self.cfg), "blocks": BlockBridge( name="gpt_neox.layers", submodules={ "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg), "ln2": NormalizationBridge(name="post_attention_layernorm", config=self.cfg), - "attn": JointQKVAttentionBridge( + "attn": JointQKVPositionEmbeddingsAttentionBridge( name="attention", config=self.cfg, split_qkv_matrix=self.split_qkv_matrix, @@ -160,49 +160,87 @@ def split_qkv_matrix( self, original_attention_component: Any ) -> tuple[torch.nn.Linear, torch.nn.Linear, torch.nn.Linear]: """Split the QKV matrix into separate linear transformations. + + GPT-NeoX/Pythia uses an interleaved QKV format where the weights are stored as + [Q_h0, K_h0, V_h0, Q_h1, K_h1, V_h1, ...] - i.e., Q, K, V are interleaved per head. + + The weight shape is [n_heads * 3 * d_head, d_model] and the output is reshaped + by HuggingFace as [batch, seq, n_heads, 3*d_head] then split on the last dim. + Args: - attention_component: The original attention layer component + original_attention_component: The original attention layer component + Returns: Tuple of nn.Linear modules for Q, K, and V transformations """ - - # Keep mypy happy assert original_attention_component is not None assert original_attention_component.query_key_value is not None qkv_weights = original_attention_component.query_key_value.weight - - # Keep mypy happy assert isinstance(qkv_weights, torch.Tensor) - # Original qkv_weights shape: [3 * d_model, d_model] -> Transposed to [d_model, 3 * d_model] - # Split into three equal parts along dimension 1 to get Q, K, V weights - W_Q, W_K, W_V = torch.tensor_split(qkv_weights.T, 3, dim=1) + n_heads = self.cfg.n_heads + d_head = self.cfg.d_head + d_model = self.cfg.d_model - qkv_bias = original_attention_component.query_key_value.bias + # Weight shape: [n_heads * 3 * d_head, d_model] + # Reshape to [n_heads, 3 * d_head, d_model] to access Q, K, V per head + W_reshaped = qkv_weights.view(n_heads, 3 * d_head, d_model) - # Keep mypy happy + # Extract Q, K, V weights for all heads and flatten back + W_Q = W_reshaped[:, :d_head, :].reshape(n_heads * d_head, d_model) + W_K = W_reshaped[:, d_head : 2 * d_head, :].reshape(n_heads * d_head, d_model) + W_V = W_reshaped[:, 2 * d_head :, :].reshape(n_heads * d_head, d_model) + + # Handle bias - same interleaved format + qkv_bias = original_attention_component.query_key_value.bias assert isinstance(qkv_bias, torch.Tensor) - # Original qkv_bias shape: [n_heads * 3 * d_head] - # Reshape to [3, n_heads * d_head] to split by Q, K, V - qkv_bias = qkv_bias.reshape(3, self.cfg.n_heads * self.cfg.d_head) - b_Q, b_K, b_V = qkv_bias[0, :], qkv_bias[1, :], qkv_bias[2, :] + # Bias shape: [n_heads * 3 * d_head] + # Reshape to [n_heads, 3 * d_head] to access Q, K, V per head + b_reshaped = qkv_bias.view(n_heads, 3 * d_head) + b_Q = b_reshaped[:, :d_head].reshape(n_heads * d_head) + b_K = b_reshaped[:, d_head : 2 * d_head].reshape(n_heads * d_head) + b_V = b_reshaped[:, 2 * d_head :].reshape(n_heads * d_head) # Create nn.Linear modules - # After tensor_split, W_Q, W_K, W_V shapes are [d_model, d_model] ([in_features, out_features]) - # nn.Linear expects weight shape [out_features, in_features] - # So we need to transpose the weights - W_Q_transformation = torch.nn.Linear(W_Q.shape[0], W_Q.shape[1], bias=True) - W_Q_transformation.weight = torch.nn.Parameter(W_Q.T) + # Weight shape for nn.Linear is [out_features, in_features] + W_Q_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True) + W_Q_transformation.weight = torch.nn.Parameter(W_Q) W_Q_transformation.bias = torch.nn.Parameter(b_Q) - W_K_transformation = torch.nn.Linear(W_K.shape[0], W_K.shape[1], bias=True) - W_K_transformation.weight = torch.nn.Parameter(W_K.T) + W_K_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True) + W_K_transformation.weight = torch.nn.Parameter(W_K) W_K_transformation.bias = torch.nn.Parameter(b_K) - W_V_transformation = torch.nn.Linear(W_V.shape[0], W_V.shape[1], bias=True) - W_V_transformation.weight = torch.nn.Parameter(W_V.T) + W_V_transformation = torch.nn.Linear(d_model, n_heads * d_head, bias=True) + W_V_transformation.weight = torch.nn.Parameter(W_V) W_V_transformation.bias = torch.nn.Parameter(b_V) return W_Q_transformation, W_K_transformation, W_V_transformation + + def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: + """Set up rotary embedding references for Pythia component testing. + + Pythia uses RoPE (Rotary Position Embeddings) in the GPT-NeoX architecture. + We need to set the rotary_emb reference on all attention bridge instances + for component testing. + + Args: + hf_model: The HuggingFace Pythia model instance + bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances) + """ + # Get rotary embedding instance from model level + # In GPT-NeoX/Pythia, rotary_emb is at the model level + rotary_emb = hf_model.gpt_neox.rotary_emb + + # Set rotary_emb on actual bridge instances in bridge_model if available + if bridge_model is not None and hasattr(bridge_model, "blocks"): + # Set on each layer's actual attention bridge instance + for block in bridge_model.blocks: + if hasattr(block, "attn"): + block.attn.set_rotary_emb(rotary_emb) + + # Also set on the template for get_generalized_component() calls + attn_bridge = self.get_generalized_component("blocks.0.attn") + attn_bridge.set_rotary_emb(rotary_emb) diff --git a/transformer_lens/model_bridge/supported_architectures/qwen2.py b/transformer_lens/model_bridge/supported_architectures/qwen2.py index 66f672ead0..fbe94fe775 100644 --- a/transformer_lens/model_bridge/supported_architectures/qwen2.py +++ b/transformer_lens/model_bridge/supported_architectures/qwen2.py @@ -8,12 +8,13 @@ ) from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( - AttentionBridge, BlockBridge, EmbeddingBridge, LinearBridge, MLPBridge, + PositionEmbeddingsAttentionBridge, RMSNormalizationBridge, + RotaryEmbeddingBridge, UnembeddingBridge, ) @@ -76,12 +77,14 @@ def __init__(self, cfg: Any) -> None: } self.component_mapping = { "embed": EmbeddingBridge(name="model.embed_tokens"), + "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), "blocks": BlockBridge( name="model.layers", + config=self.cfg, submodules={ "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), - "attn": AttentionBridge( + "attn": PositionEmbeddingsAttentionBridge( name="self_attn", config=self.cfg, submodules={ @@ -90,6 +93,8 @@ def __init__(self, cfg: Any) -> None: "v": LinearBridge(name="v_proj"), "o": LinearBridge(name="o_proj"), }, + requires_attention_mask=True, + requires_position_embeddings=True, ), "mlp": MLPBridge( name="mlp", @@ -103,5 +108,29 @@ def __init__(self, cfg: Any) -> None: }, ), "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), - "unembed": UnembeddingBridge(name="lm_head"), + "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), } + + def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: + """Set up rotary embedding references for Qwen2 component testing. + + Qwen2 uses RoPE (Rotary Position Embeddings). We set the rotary_emb reference + on all attention bridge instances for component testing. + + Args: + hf_model: The HuggingFace Qwen2 model instance + bridge_model: The TransformerBridge model (if available, set rotary_emb on actual instances) + """ + # Get rotary embedding instance from the model + rotary_emb = hf_model.model.rotary_emb + + # Set rotary_emb on actual bridge instances in bridge_model if available + if bridge_model is not None and hasattr(bridge_model, "blocks"): + # Set on each layer's actual attention bridge instance + for block in bridge_model.blocks: + if hasattr(block, "attn"): + block.attn.set_rotary_emb(rotary_emb) + + # Also set on the template for get_generalized_component() calls + attn_bridge = self.get_generalized_component("blocks.0.attn") + attn_bridge.set_rotary_emb(rotary_emb) diff --git a/transformer_lens/weight_processing.py b/transformer_lens/weight_processing.py index 19a701eaec..2e407ab056 100644 --- a/transformer_lens/weight_processing.py +++ b/transformer_lens/weight_processing.py @@ -283,6 +283,21 @@ def extract_attention_tensors_for_folding( ln1_b_key = ProcessWeights._get_param_key(f"blocks.{layer}.ln1.b", adapter) ln1_w_key = ProcessWeights._get_param_key(f"blocks.{layer}.ln1.w", adapter) + # For GQA models, Q, K and V weights may use underscore prefix (_W_Q, _W_K, _W_V) + # Check if standard keys exist, otherwise update to use underscore-prefixed versions + if W_Q_key not in state_dict: + W_Q_key = W_Q_key.replace(".W_Q", "._W_Q") + if W_K_key not in state_dict: + W_K_key = W_K_key.replace(".W_K", "._W_K") + if W_V_key not in state_dict: + W_V_key = W_V_key.replace(".W_V", "._W_V") + if b_Q_key not in state_dict: + b_Q_key = b_Q_key.replace(".b_Q", "._b_Q") + if b_K_key not in state_dict: + b_K_key = b_K_key.replace(".b_K", "._b_K") + if b_V_key not in state_dict: + b_V_key = b_V_key.replace(".b_V", "._b_V") + wq_tensor: Optional[torch.Tensor] = state_dict.get(W_Q_key) wk_tensor: Optional[torch.Tensor] = state_dict.get(W_K_key) wv_tensor: Optional[torch.Tensor] = state_dict.get(W_V_key) @@ -662,6 +677,12 @@ def _fold_unembed_layer_norm( unembed_W_U_key = ProcessWeights._get_param_key("unembed.W_U", adapter) ln_final_b_key = ProcessWeights._get_param_key("ln_final.b", adapter) ln_final_w_key = ProcessWeights._get_param_key("ln_final.w", adapter) + + # Skip layer norm folding if ln_final doesn't exist + # (e.g., encoder-decoder models like T5 have encoder_ln_final/decoder_ln_final instead) + if ln_final_w_key not in state_dict: + return state_dict + has_unembed_bias = unembed_b_U_key in state_dict unembed_weight = ProcessWeights.convert_tensor_to_tl_format( unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), cfg, adapter, None