Merge remote-tracking branch 'origin/dev-3.x' into add_support_for_gpt_oss

Author: bryce13950

tests/integration/model_bridge/generalized_components/test_joint_qkv_attention_integration.py

@@ -0,0 +1,124 @@
+"""Lightweight integration tests for JointQKVAttentionBridge.
+
+Tests the core functionality without loading large models to keep CI fast.
+"""
+
+import pytest
+import torch
+
+import transformer_lens.utils as utils
+
+
+class TestJointQKVAttentionBridgeIntegration:
+    """Minimal integration tests for JointQKVAttentionBridge."""
+
+    def test_hook_alias_resolution(self):
+        """Test that hook aliases are properly resolved."""
+        # Test the hook alias resolution that caused the original issue
+        hook_name = utils.get_act_name("v", 0)
+        assert (
+            hook_name == "blocks.0.attn.hook_v"
+        ), f"Expected 'blocks.0.attn.hook_v', got '{hook_name}'"
+
+        # Test other hook names
+        assert utils.get_act_name("q", 1) == "blocks.1.attn.hook_q"
+        assert utils.get_act_name("k", 2) == "blocks.2.attn.hook_k"
+
+    def test_component_class_exists(self):
+        """Test that JointQKVAttentionBridge class can be imported."""
+        from transformer_lens.model_bridge.generalized_components.joint_qkv_attention import (
+            JointQKVAttentionBridge,
+        )
+
+        # Verify the class exists and has expected methods
+        assert hasattr(JointQKVAttentionBridge, "forward")
+        assert hasattr(JointQKVAttentionBridge, "_reconstruct_attention")
+        assert hasattr(JointQKVAttentionBridge, "_manual_attention_computation")
+
+    def test_hook_point_has_hooks_method(self):
+        """Test that HookPoint.has_hooks method works correctly."""
+        from transformer_lens.hook_points import HookPoint
+
+        hook_point = HookPoint()
+
+        # Test initial state
+        assert not hook_point.has_hooks()
+        assert not hook_point.has_hooks(dir="fwd")
+        assert not hook_point.has_hooks(dir="bwd")
+
+        # Add a hook and test detection
+        def dummy_hook(x, hook):
+            return x
+
+        hook_point.add_hook(dummy_hook)
+        assert hook_point.has_hooks()
+        assert hook_point.has_hooks(dir="fwd")
+        assert not hook_point.has_hooks(dir="bwd")
+
+        # Clean up
+        hook_point.remove_hooks()
+        assert not hook_point.has_hooks()
+
+    def test_architecture_imports(self):
+        """Test that architecture files can be imported and reference JointQKVAttentionBridge."""
+        # Test that we can import the architecture files without errors
+        # Test that JointQKVAttentionBridge is referenced in the source files
+        import inspect
+
+        from transformer_lens.model_bridge.supported_architectures import (
+            bloom,
+            gpt2,
+            neox,
+        )
+
+        gpt2_source = inspect.getsource(gpt2)
+        assert (
+            "JointQKVAttentionBridge" in gpt2_source
+        ), "GPT-2 architecture should reference JointQKVAttentionBridge"
+
+        bloom_source = inspect.getsource(bloom)
+        assert (
+            "JointQKVAttentionBridge" in bloom_source
+        ), "BLOOM architecture should reference JointQKVAttentionBridge"
+
+        neox_source = inspect.getsource(neox)
+        assert (
+            "JointQKVAttentionBridge" in neox_source
+        ), "NeoX architecture should reference JointQKVAttentionBridge"
+
+    @pytest.mark.skip(reason="Requires model loading - too slow for CI")
+    def test_distilgpt2_integration(self):
+        """Full integration test with DistilGPT-2 (skipped in CI)."""
+        # This test would load DistilGPT-2 and test full functionality
+        # but is skipped by default to keep CI fast
+        from transformer_lens.model_bridge import TransformerBridge
+
+        torch.set_grad_enabled(False)
+        model = TransformerBridge.boot_transformers("distilgpt2", device="cpu")
+
+        # Verify JointQKVAttentionBridge usage
+        joint_qkv_modules = [
+            name
+            for name, module in model.named_modules()
+            if "JointQKVAttentionBridge" in getattr(module, "__class__", {}).get("__name__", "")
+        ]
+        assert (
+            len(joint_qkv_modules) == 6
+        ), f"Expected 6 JointQKVAttentionBridge modules, got {len(joint_qkv_modules)}"
+
+        # Test basic functionality
+        tokens = model.to_tokens("Test")
+        with torch.no_grad():
+            loss = model(tokens, return_type="loss")
+            assert torch.isfinite(loss) and loss > 0
+
+        # Test hook integration
+        def v_ablation_hook(value, hook):
+            value[:, :, 0, :] = 0.0  # Ablate first head
+            return value
+
+        original_loss = model(tokens, return_type="loss")
+        hooked_loss = model.run_with_hooks(
+            tokens, return_type="loss", fwd_hooks=[(utils.get_act_name("v", 0), v_ablation_hook)]
+        )
+        assert not torch.isclose(original_loss, hooked_loss, atol=1e-6)

tests/integration/model_bridge/test_bridge_root_module_cache_compatibility.py

@@ -11,8 +11,6 @@
     },
 )
 
-# Attention output enabled via hf_config_overrides
-
 act_names_in_cache = [
     # "hook_embed",
     # "hook_pos_embed",

tests/unit/test_hook_points.py

@@ -208,3 +208,264 @@ def test_hook(activation, hook):
     # Since hook returns None, the original input should be returned
     # (HookPoint's forward method returns the input when no valid hook result)
     assert torch.equal(result, test_input)
+
+
+class TestHookPointHasHooks:
+    """Comprehensive test suite for HookPoint.has_hooks method."""
+
+    def setup_method(self):
+        """Set up fresh HookPoint and sample hook for each test."""
+        self.hook_point = HookPoint()
+
+        def sample_hook(activation, hook):
+            return activation
+
+        self.sample_hook = sample_hook
+
+    def test_no_hooks_returns_false(self):
+        """Test that has_hooks returns False when no hooks are present."""
+        assert not self.hook_point.has_hooks()
+        assert not self.hook_point.has_hooks(dir="fwd")
+        assert not self.hook_point.has_hooks(dir="bwd")
+        assert not self.hook_point.has_hooks(dir="both")
+
+    def test_forward_hook_detection(self):
+        """Test detection of forward hooks."""
+        # Add a forward hook
+        self.hook_point.add_hook(self.sample_hook, dir="fwd")
+
+        # Should detect forward hooks
+        assert self.hook_point.has_hooks()
+        assert self.hook_point.has_hooks(dir="fwd")
+        assert self.hook_point.has_hooks(dir="both")
+
+        # Should not detect backward hooks
+        assert not self.hook_point.has_hooks(dir="bwd")
+
+    def test_backward_hook_detection(self):
+        """Test detection of backward hooks."""
+        # Add a backward hook
+        self.hook_point.add_hook(self.sample_hook, dir="bwd")
+
+        # Should detect backward hooks
+        assert self.hook_point.has_hooks()
+        assert self.hook_point.has_hooks(dir="bwd")
+        assert self.hook_point.has_hooks(dir="both")
+
+        # Should not detect forward hooks
+        assert not self.hook_point.has_hooks(dir="fwd")
+
+    def test_both_direction_hooks(self):
+        """Test detection when both forward and backward hooks are present."""
+        # Add both forward and backward hooks
+        self.hook_point.add_hook(self.sample_hook, dir="fwd")
+        self.hook_point.add_hook(self.sample_hook, dir="bwd")
+
+        # All directions should detect hooks
+        assert self.hook_point.has_hooks()
+        assert self.hook_point.has_hooks(dir="fwd")
+        assert self.hook_point.has_hooks(dir="bwd")
+        assert self.hook_point.has_hooks(dir="both")
+
+    def test_permanent_hook_detection(self):
+        """Test detection of permanent hooks."""
+        # Add a permanent forward hook
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", is_permanent=True)
+
+        # Should detect permanent hooks by default
+        assert self.hook_point.has_hooks()
+        assert self.hook_point.has_hooks(including_permanent=True)
+
+        # Should not detect when excluding permanent hooks
+        assert not self.hook_point.has_hooks(including_permanent=False)
+
+    def test_non_permanent_hook_detection(self):
+        """Test detection of non-permanent hooks."""
+        # Add a non-permanent forward hook
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", is_permanent=False)
+
+        # Should detect non-permanent hooks regardless of including_permanent setting
+        assert self.hook_point.has_hooks()
+        assert self.hook_point.has_hooks(including_permanent=True)
+        assert self.hook_point.has_hooks(including_permanent=False)
+
+    def test_mixed_permanent_hooks(self):
+        """Test detection with mix of permanent and non-permanent hooks."""
+        # Add both permanent and non-permanent hooks
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", is_permanent=True)
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", is_permanent=False)
+
+        # Should detect hooks in both cases
+        assert self.hook_point.has_hooks(including_permanent=True)
+        assert self.hook_point.has_hooks(including_permanent=False)
+
+    def test_only_permanent_hooks(self):
+        """Test detection when only permanent hooks are present."""
+        # Add only permanent hooks
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", is_permanent=True)
+        self.hook_point.add_hook(self.sample_hook, dir="bwd", is_permanent=True)
+
+        # Should detect when including permanent
+        assert self.hook_point.has_hooks(including_permanent=True)
+        assert self.hook_point.has_hooks(dir="fwd", including_permanent=True)
+        assert self.hook_point.has_hooks(dir="bwd", including_permanent=True)
+
+        # Should not detect when excluding permanent
+        assert not self.hook_point.has_hooks(including_permanent=False)
+        assert not self.hook_point.has_hooks(dir="fwd", including_permanent=False)
+        assert not self.hook_point.has_hooks(dir="bwd", including_permanent=False)
+
+    def test_context_level_filtering(self):
+        """Test context level filtering functionality."""
+        # Add hooks at different context levels
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", level=0)
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", level=1)
+        self.hook_point.add_hook(self.sample_hook, dir="bwd", level=2)
+
+        # Should detect hooks at specific levels
+        assert self.hook_point.has_hooks(level=0)
+        assert self.hook_point.has_hooks(level=1)
+        assert self.hook_point.has_hooks(level=2)
+
+        # Should not detect hooks at non-existent levels
+        assert not self.hook_point.has_hooks(level=3)
+        assert not self.hook_point.has_hooks(level=-1)
+
+        # Should detect all hooks when level is None
+        assert self.hook_point.has_hooks(level=None)
+
+    def test_context_level_with_direction(self):
+        """Test context level filtering combined with direction filtering."""
+        # Add hooks at different levels and directions
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", level=0)
+        self.hook_point.add_hook(self.sample_hook, dir="bwd", level=1)
+
+        # Test specific combinations
+        assert self.hook_point.has_hooks(dir="fwd", level=0)
+        assert self.hook_point.has_hooks(dir="bwd", level=1)
+
+        # Test non-matching combinations
+        assert not self.hook_point.has_hooks(dir="fwd", level=1)
+        assert not self.hook_point.has_hooks(dir="bwd", level=0)
+
+    def test_context_level_with_permanent_flags(self):
+        """Test context level filtering combined with permanent hook filtering."""
+        # Add permanent and non-permanent hooks at different levels
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", level=0, is_permanent=True)
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", level=1, is_permanent=False)
+
+        # Test combinations
+        assert self.hook_point.has_hooks(level=0, including_permanent=True)
+        assert not self.hook_point.has_hooks(level=0, including_permanent=False)
+        assert self.hook_point.has_hooks(level=1, including_permanent=True)
+        assert self.hook_point.has_hooks(level=1, including_permanent=False)
+
+    def test_all_parameters_combined(self):
+        """Test all parameters combined in various ways."""
+        # Create a complex setup with multiple hooks
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", level=0, is_permanent=True)
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", level=1, is_permanent=False)
+        self.hook_point.add_hook(self.sample_hook, dir="bwd", level=0, is_permanent=False)
+        self.hook_point.add_hook(self.sample_hook, dir="bwd", level=2, is_permanent=True)
+
+        # Test specific combinations
+        assert self.hook_point.has_hooks(dir="fwd", level=0, including_permanent=True)
+        assert not self.hook_point.has_hooks(dir="fwd", level=0, including_permanent=False)
+        assert self.hook_point.has_hooks(dir="fwd", level=1, including_permanent=False)
+        assert self.hook_point.has_hooks(dir="bwd", level=0, including_permanent=False)
+        assert not self.hook_point.has_hooks(dir="bwd", level=1, including_permanent=False)
+        assert self.hook_point.has_hooks(dir="bwd", level=2, including_permanent=True)
+
+    def test_invalid_direction_raises_error(self):
+        """Test that invalid direction parameter raises error (caught by type checking)."""
+        # Note: beartype catches this at the parameter level before reaching the ValueError
+        import pytest
+        from beartype.roar import BeartypeCallHintParamViolation
+
+        with pytest.raises(BeartypeCallHintParamViolation):
+            self.hook_point.has_hooks(dir="invalid")  # type: ignore
+
+    def test_multiple_hooks_same_criteria(self):
+        """Test detection when multiple hooks match the same criteria."""
+        # Add multiple hooks with same criteria
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", level=0, is_permanent=False)
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", level=0, is_permanent=False)
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", level=0, is_permanent=False)
+
+        # Should still detect hooks (method returns True on first match)
+        assert self.hook_point.has_hooks(dir="fwd", level=0, including_permanent=False)
+
+    def test_hook_removal_affects_detection(self):
+        """Test that removing hooks affects detection."""
+        # Add a hook
+        self.hook_point.add_hook(self.sample_hook, dir="fwd")
+        assert self.hook_point.has_hooks()
+
+        # Remove all hooks
+        self.hook_point.remove_hooks(dir="both")
+        assert not self.hook_point.has_hooks()
+
+    def test_default_parameter_values(self):
+        """Test that default parameter values work correctly."""
+        # Add hooks to test defaults
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", is_permanent=True, level=0)
+        self.hook_point.add_hook(self.sample_hook, dir="bwd", is_permanent=False, level=1)
+
+        # Test default behavior (dir="both", including_permanent=True, level=None)
+        assert self.hook_point.has_hooks()
+
+        # This should be equivalent to:
+        assert self.hook_point.has_hooks(dir="both", including_permanent=True, level=None)
+
+    def test_edge_case_empty_after_filtering(self):
+        """Test edge case where hooks exist but are filtered out."""
+        # Add hooks that will be filtered out
+        self.hook_point.add_hook(self.sample_hook, dir="fwd", is_permanent=True, level=5)
+
+        # These should not detect the hook due to filtering
+        assert not self.hook_point.has_hooks(including_permanent=False)
+        assert not self.hook_point.has_hooks(dir="bwd")
+        assert not self.hook_point.has_hooks(level=0)
+        assert not self.hook_point.has_hooks(dir="bwd", level=5, including_permanent=True)
+
+    def test_functional_hook_execution_still_works(self):
+        """Test that has_hooks doesn't interfere with actual hook functionality."""
+        import torch
+
+        results = []
+
+        def test_hook(activation, hook):
+            results.append("hook_called")
+            return activation
+
+        # Add hook and verify detection
+        self.hook_point.add_hook(test_hook, dir="fwd")
+        assert self.hook_point.has_hooks()
+
+        # Execute hook and verify it still works
+        test_input = torch.tensor([1.0, 2.0, 3.0])
+        output = self.hook_point(test_input)
+
+        assert torch.equal(output, test_input)
+        assert "hook_called" in results
+
+    def test_hook_point_with_conversions(self):
+        """Test has_hooks with hook conversions if they exist."""
+        import torch
+
+        # This test ensures has_hooks works even when hook conversions are involved
+        def simple_hook(activation, hook):
+            return activation * 2
+
+        # Add hook
+        self.hook_point.add_hook(simple_hook, dir="fwd")
+
+        # Should detect hook regardless of any internal conversions
+        assert self.hook_point.has_hooks()
+        assert self.hook_point.has_hooks(dir="fwd")
+
+        # Test actual functionality still works
+        test_input = torch.tensor([1.0, 2.0])
+        output = self.hook_point(test_input)
+        expected = torch.tensor([2.0, 4.0])
+        assert torch.allclose(output, expected)