Coverage for transformer_lens/model_bridge/generalized_components/opaque_block.py: 67%

1"""Generic opaque block bridge for non-SSM, non-standard-transformer architectures.""" 

2 

3from __future__ import annotations 

4 

5import re 

6from typing import Any, Dict, Optional 

7 

8import torch 

9 

10from transformer_lens.model_bridge.exceptions import StopAtLayerException 

11from transformer_lens.model_bridge.generalized_components.base import ( 

12 GeneralizedComponent, 

13) 

14 

15 

16class OpaqueBlockBridge(GeneralizedComponent): 

17 """Generic block bridge that delegates a full block forward unchanged. 

18 

19 Exposes ``hook_resid_pre`` / ``hook_resid_post`` (aliases of ``hook_in`` 

20 / ``hook_out``) on the residual stream only — no assumptions are made 

21 about the block's internal structure. 

22 

23 Use this for architectures whose block internals do not follow a standard 

24 SSM or transformer pre-norm flow, e.g.: 

25 

26 - Post-residual (sandwich-norm) blocks like Raven / Huginn. 

27 - Attention-free recurrent blocks like RWKV-7. 

28 - Any custom block where ``hook_mixer_in`` / ``hook_mixer_out`` would be 

29 semantically wrong or structurally absent. 

30 

31 For SSM architectures (Mamba, Falcon-H1) use :class:`SSMBlockBridge`, 

32 which extends this class and adds the ``hook_mixer_in`` / 

33 ``hook_mixer_out`` aliases pointing at ``mixer.hook_in`` / 

34 ``mixer.hook_out``. 

35 """ 

36 

37 is_list_item: bool = True 

38 hook_out_is_single_residual_stream: bool = True 

39 hook_aliases = { 

40 "hook_resid_pre": "hook_in", 

41 "hook_resid_post": "hook_out", 

42 } 

43 

44 def __init__( 

45 self, 

46 name: str, 

47 config: Optional[Any] = None, 

48 submodules: Optional[Dict[str, GeneralizedComponent]] = None, 

49 hook_alias_overrides: Optional[Dict[str, str]] = None, 

50 ): 

51 super().__init__( 

52 name, 

53 config, 

54 submodules=submodules if submodules is not None else {}, 

55 hook_alias_overrides=hook_alias_overrides, 

56 ) 

57 

58 def forward(self, *args: Any, **kwargs: Any) -> Any: 

59 """Delegate to the HF block with hook_in/hook_out wrapped around it.""" 

60 if self.original_component is None: 60 ↛ 61line 60 didn't jump to line 61 because the condition on line 60 was never true

61 raise RuntimeError( 

62 f"Original component not set for {self.name}. " 

63 "Call set_original_component() first." 

64 ) 

65 

66 self._check_stop_at_layer(*args, **kwargs) 

67 args, kwargs = self._hook_input_hidden_states(args, kwargs) 

68 output = self.original_component(*args, **kwargs) 

69 return self._apply_output_hook(output) 

70 

71 def _apply_output_hook(self, output: Any) -> Any: 

72 """Hook the primary output tensor, preserving tuple structure if present.""" 

73 if isinstance(output, tuple) and len(output) > 0: 73 ↛ 74line 73 didn't jump to line 74 because the condition on line 73 was never true

74 first = output[0] 

75 if isinstance(first, torch.Tensor): 

76 first = self.hook_out(first) 

77 return (first,) + output[1:] 

78 return output 

79 if isinstance(output, torch.Tensor): 79 ↛ 81line 79 didn't jump to line 81 because the condition on line 79 was always true

80 return self.hook_out(output) 

81 return output 

82 

83 def _hook_input_hidden_states(self, args: tuple, kwargs: dict) -> tuple[tuple, dict]: 

84 """Hook the hidden_states input whether it arrives positionally or by name.""" 

85 if len(args) > 0 and isinstance(args[0], torch.Tensor): 85 ↛ 88line 85 didn't jump to line 88 because the condition on line 85 was always true

86 hooked = self.hook_in(args[0]) 

87 args = (hooked,) + args[1:] 

88 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor): 

89 kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"]) 

90 return args, kwargs 

91 

92 def _check_stop_at_layer(self, *args: Any, **kwargs: Any) -> None: 

93 """Raise StopAtLayerException when the configured stop index matches this block.""" 

94 if not (hasattr(self, "_stop_at_layer_idx") and self._stop_at_layer_idx is not None): 

95 return 

96 if self.name is None: 96 ↛ 97line 96 didn't jump to line 97 because the condition on line 96 was never true

97 return 

98 # Mamba uses `.layers.{i}`; `blocks.{i}` is the fallback TL convention. 

99 match = re.search(r"\.layers\.(\d+)", self.name) or re.search(r"blocks\.(\d+)", self.name) 

100 if not match: 100 ↛ 101line 100 didn't jump to line 101 because the condition on line 100 was never true

101 return 

102 layer_idx = int(match.group(1)) 

103 if layer_idx != self._stop_at_layer_idx: 

104 return 

105 if len(args) > 0 and isinstance(args[0], torch.Tensor): 105 ↛ 107line 105 didn't jump to line 107 because the condition on line 105 was always true

106 input_tensor = args[0] 

107 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor): 

108 input_tensor = kwargs["hidden_states"] 

109 else: 

110 raise ValueError(f"Cannot find input tensor to stop at layer {layer_idx}") 

111 input_tensor = self.hook_in(input_tensor) 

112 raise StopAtLayerException(input_tensor)