Coverage for transformer_lens/model_bridge/generalized_components/t5gemma2_decoder_block.py: 77%

1"""T5Gemma2-specific decoder block bridge. 

2 

3T5Gemma2DecoderLayer replaces T5Gemma's separate self-attention + cross-attention 

4with a single T5Gemma2MergedAttention module. That module computes decoder self 

5queries/keys/values from ``hidden_states`` and cross keys/values from 

6``encoder_hidden_states`` using the *same* q/k/v/o projections, concatenates the 

7self and cross key/value states, and runs a single softmax. As a result the 

8decoder layer has no separate cross-attention module and no cross-attention 

9layernorms — its structure mirrors the encoder layer plus encoder-state input. 

10 

11This bridge monkey-patches the layer forward to fire hook points at the canonical 

12HookedTransformer residual-stream positions while delegating all attention math 

13(QK-norm, RoPE, scaling, merged KV) to the native HF module. 

14""" 

15from __future__ import annotations 

16 

17import types 

18from typing import Any, Callable, Dict, Optional 

19 

20import torch 

21 

22from transformer_lens.hook_points import HookPoint 

23from transformer_lens.model_bridge.generalized_components.base import ( 

24 GeneralizedComponent, 

25) 

26 

27 

28class T5Gemma2DecoderBlockBridge(GeneralizedComponent): 

29 """Bridge for T5Gemma2 decoder layers (merged self+cross attention). 

30 

31 Inserts hook points around the two sub-components of each decoder layer: 

32 - hook_in (hook_resid_pre): residual before self-attention pre-norm 

33 - hook_resid_mid: residual after merged-attention + residual add, before MLP pre-norm 

34 - hook_out (hook_resid_post): residual after MLP + residual add 

35 """ 

36 

37 is_list_item: bool = True 

38 hook_aliases = { 

39 "hook_resid_pre": "hook_in", 

40 "hook_resid_post": "hook_out", 

41 } 

42 

43 def __init__( 

44 self, 

45 name: str, 

46 config: Optional[Any] = None, 

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

48 ): 

49 super().__init__(name, config, submodules=submodules or {}) 

50 self.hook_resid_mid = HookPoint() 

51 self._register_hook("hook_resid_mid", self.hook_resid_mid) 

52 self._original_block_forward: Optional[Callable[..., Any]] = None 

53 

54 def set_original_component(self, component: torch.nn.Module) -> None: 

55 super().set_original_component(component) 

56 self._patch_decoder_layer_forward() 

57 

58 def _patch_decoder_layer_forward(self) -> None: 

59 """Monkey-patch T5Gemma2DecoderLayer.forward to insert hook points. 

60 

61 The patched forward preserves the original residual-stream semantics but 

62 fires hook_in, hook_resid_mid, and hook_out at the canonical 

63 HookedTransformer positions. Attention math stays native. 

64 """ 

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

66 return 

67 self._original_block_forward = self.original_component.forward 

68 

69 hook_in = self.hook_in # fires at hook_resid_pre 

70 hook_resid_mid = self.hook_resid_mid 

71 hook_out = self.hook_out # fires at hook_resid_post 

72 

73 def patched_forward( 

74 layer_self, 

75 hidden_states: torch.Tensor, 

76 position_embeddings=None, 

77 merged_attention_mask=None, 

78 position_ids=None, 

79 past_key_values=None, 

80 use_cache=None, 

81 encoder_hidden_states=None, 

82 **kwargs: Any, 

83 ) -> torch.Tensor: 

84 hidden_states = hook_in(hidden_states) 

85 

86 # --- merged self/cross attention sub-layer --- 

87 residual = hidden_states 

88 hidden_states = layer_self.pre_self_attn_layernorm(hidden_states) 

89 attn_out, _, _ = layer_self.self_attn( 

90 hidden_states=hidden_states, 

91 position_embeddings=position_embeddings, 

92 merged_attention_mask=merged_attention_mask, 

93 encoder_hidden_states=encoder_hidden_states, 

94 past_key_values=past_key_values, 

95 use_cache=use_cache, 

96 **kwargs, 

97 ) 

98 hidden_states = layer_self.post_self_attn_layernorm(attn_out) 

99 hidden_states = residual + layer_self.dropout(hidden_states) 

100 hidden_states = hook_resid_mid(hidden_states) 

101 

102 # --- MLP sub-layer --- 

103 residual = hidden_states 

104 hidden_states = layer_self.pre_feedforward_layernorm(hidden_states) 

105 hidden_states = layer_self.mlp(hidden_states) 

106 hidden_states = layer_self.post_feedforward_layernorm(hidden_states) 

107 hidden_states = residual + layer_self.dropout(hidden_states) 

108 hidden_states = hook_out(hidden_states) 

109 

110 return hidden_states 

111 

112 self.original_component.forward = types.MethodType(patched_forward, self.original_component) 

113 

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

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

116 raise RuntimeError( 

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

118 "Call set_original_component() first." 

119 ) 

120 return self.original_component(*args, **kwargs) 

121 

122 def get_expected_parameter_names(self, prefix: str = "") -> list[str]: 

123 param_names = [] 

124 for sub_name, sub_component in self.submodules.items(): 

125 sub_prefix = f"{prefix}.{sub_name}" if prefix else sub_name 

126 param_names.extend(sub_component.get_expected_parameter_names(sub_prefix)) 

127 return param_names 

128 

129 def get_list_size(self) -> int: 

130 if self.config is None: 

131 return 0 

132 return getattr(self.config, "n_layers", 0)