Coverage for transformer_lens/model_bridge/generalized_components/altup_block.py: 63%

1"""Block bridge for AltUp (Alternating Updates) decoder layers.""" 

2from __future__ import annotations 

3 

4import re 

5from typing import Any, Dict, Optional 

6 

7import torch 

8 

9from transformer_lens.hook_points import HookPoint 

10from transformer_lens.model_bridge.exceptions import StopAtLayerException 

11from transformer_lens.model_bridge.generalized_components.base import ( 

12 GeneralizedComponent, 

13) 

14 

15 

16class AltUpBlockBridge(GeneralizedComponent): 

17 """Block bridge for a decoder layer that operates on a stacked AltUp residual. 

18 

19 Direct GeneralizedComponent subclass (not BlockBridge) because the layer's residual is a 

20 stacked ``[num_altup_inputs, batch, seq, d_model]`` tensor, not a single stream. 

21 ``hook_in``/``hook_out`` carry the full stack; ``hook_resid_pre``/``hook_resid_post`` expose 

22 the active stream (``altup_active_idx``) as a conventional ``[batch, seq, d_model]`` residual 

23 and are patchable (written back into the stack). 

24 """ 

25 

26 is_list_item: bool = True 

27 hook_aliases = { 

28 "hook_attn_out": "self_attn.hook_out", 

29 "hook_mlp_out": "mlp.hook_out", 

30 } 

31 

32 def __init__( 

33 self, 

34 name: str, 

35 config: Optional[Any] = None, 

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

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

38 ): 

39 # Same post-norm redirect as BlockBridge: the residual contribution is the 

40 # post-sublayer norm's output, not the raw sublayer output. 

41 auto_overrides = {} 

42 if submodules is not None: 

43 if "ln1_post" in submodules: 43 ↛ 45line 43 didn't jump to line 45 because the condition on line 43 was always true

44 auto_overrides["hook_attn_out"] = "ln1_post.hook_out" 

45 if "ln2_post" in submodules: 45 ↛ 47line 45 didn't jump to line 47 because the condition on line 45 was always true

46 auto_overrides["hook_mlp_out"] = "ln2_post.hook_out" 

47 merged_overrides = {**auto_overrides, **(hook_alias_overrides or {})} 

48 super().__init__( 

49 name, 

50 config, 

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

52 hook_alias_overrides=merged_overrides if merged_overrides else None, 

53 ) 

54 self.altup_active_idx = int(getattr(config, "altup_active_idx", 0) or 0) 

55 # Active AltUp stream as a conventional residual (patchable). 

56 self.hook_resid_pre = HookPoint() 

57 self.hook_resid_post = HookPoint() 

58 

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

60 """Delegate to the HF layer, hooking the AltUp stack and the active residual stream.""" 

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

62 raise RuntimeError( 

63 f"Original component not set for {self.name}. Call set_original_component() first." 

64 ) 

65 

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

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

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

69 return self._hook_output(output) 

70 

71 def _patch_active_stream(self, stack: torch.Tensor, hook: HookPoint) -> torch.Tensor: 

72 """Fire ``hook`` on the active AltUp stream, cloning the stack only if it was patched.""" 

73 if ( 73 ↛ 78line 73 didn't jump to line 78 because the condition on line 73 was never true

74 not isinstance(stack, torch.Tensor) 

75 or stack.dim() < 1 

76 or stack.shape[0] <= self.altup_active_idx 

77 ): 

78 return stack 

79 # Capture the view once: indexing returns a fresh object each call, so the identity 

80 # check must compare against this exact view, not a re-indexed one. 

81 active_view = stack[self.altup_active_idx] 

82 active = hook(active_view) 

83 # Common case (no hooks, or read-only hooks that return their input): nothing changed, 

84 # so skip the full [num_altup, batch, seq, d_model] clone. 

85 if active is active_view: 85 ↛ 87line 85 didn't jump to line 87 because the condition on line 85 was always true

86 return stack 

87 stack = stack.clone() 

88 stack[self.altup_active_idx] = active 

89 return stack 

90 

91 def _hook_input(self, args: tuple, kwargs: dict) -> tuple[tuple, dict]: 

92 """Hook the stacked hidden_states then the active residual, positional or by name.""" 

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

94 hidden = self._patch_active_stream(self.hook_in(args[0]), self.hook_resid_pre) 

95 args = (hidden,) + args[1:] 

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

97 kwargs["hidden_states"] = self._patch_active_stream( 

98 self.hook_in(kwargs["hidden_states"]), self.hook_resid_pre 

99 ) 

100 return args, kwargs 

101 

102 def _hook_output(self, output: Any) -> Any: 

103 """Hook the active residual then the stacked output, preserving any tuple structure.""" 

104 primary = output[0] if isinstance(output, tuple) and len(output) > 0 else output 

105 if isinstance(primary, torch.Tensor): 105 ↛ 107line 105 didn't jump to line 107 because the condition on line 105 was always true

106 primary = self.hook_out(self._patch_active_stream(primary, self.hook_resid_post)) 

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

108 return (primary,) + output[1:] 

109 return primary 

110 

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

112 """Raise StopAtLayerException (carrying the AltUp stack) at the configured stop index.""" 

113 if getattr(self, "_stop_at_layer_idx", None) is None or self.name is None: 113 ↛ 115line 113 didn't jump to line 115 because the condition on line 113 was always true

114 return 

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

116 if not match or int(match.group(1)) != self._stop_at_layer_idx: 

117 return 

118 if len(args) > 0 and isinstance(args[0], torch.Tensor): 

119 tensor = args[0] 

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

121 tensor = kwargs["hidden_states"] 

122 else: 

123 raise ValueError(f"Cannot find input tensor to stop at layer {self.name}") 

124 raise StopAtLayerException(self.hook_in(tensor))