Coverage for transformer_lens/model_bridge/generalized_components/t5_block.py: 81%

1"""T5-specific block bridge component. 

2 

3This module contains the bridge component for T5 blocks, which have a different 

4structure than standard transformer blocks (3 layers in decoder vs 2 layers). 

5""" 

6from __future__ import annotations 

7 

8import inspect 

9import types 

10from typing import Any, Callable, Dict, Optional 

11 

12import torch 

13 

14from transformer_lens.hook_points import HookPoint 

15from transformer_lens.model_bridge.generalized_components.base import ( 

16 GeneralizedComponent, 

17) 

18 

19 

20def _clamp_fp16_inf(hidden_states: torch.Tensor) -> torch.Tensor: 

21 """HF T5Block clamps fp16 hidden states after every sublayer; the patched 

22 forward must too, or fp16 runs overflow exactly where HF's do not.""" 

23 if hidden_states.dtype == torch.float16: 

24 clamp_value = torch.where( 

25 torch.isinf(hidden_states).any(), 

26 torch.finfo(hidden_states.dtype).max - 1000, 

27 torch.finfo(hidden_states.dtype).max, 

28 ) 

29 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) 

30 return hidden_states 

31 

32 

33class T5BlockBridge(GeneralizedComponent): 

34 """Bridge component for T5 transformer blocks. 

35 

36 T5 has two types of blocks: 

37 - Encoder blocks: 2 layers (self-attention, feed-forward) 

38 - Decoder blocks: 3 layers (self-attention, cross-attention, feed-forward) 

39 

40 This bridge handles both types based on the presence of cross-attention. 

41 """ 

42 

43 is_list_item: bool = True 

44 hook_aliases = {"hook_resid_pre": "hook_in", "hook_resid_post": "hook_out"} 

45 

46 def __init__( 

47 self, 

48 name: str, 

49 config: Optional[Any] = None, 

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

51 is_decoder: bool = False, 

52 ): 

53 """Initialize the T5 block bridge. 

54 

55 Args: 

56 name: The name of the component in the model 

57 config: Optional configuration 

58 submodules: Dictionary of submodules to register 

59 is_decoder: Whether this is a decoder block (has cross-attention) 

60 """ 

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

62 # T5 adds the residual INSIDE each sublayer, so the wrapped module's 

63 # output already IS the added contribution. Keyed to what this adapter 

64 # declared, since names differ across the family. 

65 self.hook_aliases = dict(type(self).hook_aliases) 

66 declared = submodules or {} 

67 if "attn" in declared: 

68 self.hook_aliases["hook_attn_out"] = "attn.hook_out" 

69 elif "self_attn" in declared: 

70 self.hook_aliases["hook_attn_out"] = "self_attn.hook_out" 

71 if "cross_attn" in declared: 

72 self.hook_aliases["hook_cross_attn_out"] = "cross_attn.hook_out" 

73 if "mlp" in declared: 

74 self.hook_aliases["hook_mlp_out"] = "mlp.hook_out" 

75 self.is_decoder = is_decoder 

76 self.hook_resid_mid = HookPoint() 

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

78 if is_decoder: 

79 self.hook_resid_mid2 = HookPoint() 

80 self._register_hook("hook_resid_mid2", self.hook_resid_mid2) 

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

82 

83 def set_original_component(self, component: torch.nn.Module): 

84 """Set the original component and monkey-patch its forward method. 

85 

86 Args: 

87 component: The original PyTorch module to wrap 

88 """ 

89 super().set_original_component(component) 

90 self._patch_t5_block_forward() 

91 

92 def _patch_t5_block_forward(self): 

93 """Monkey-patch the T5 block's forward method to insert hooks.""" 

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

95 return 

96 self._original_block_forward = self.original_component.forward 

97 # Transformers >= 5.15 T5Block.forward dropped output_attentions and returns a fixed 

98 # (hidden, self_bias, cross_bias); LongT5Block keeps the older layout the stack expects. 

99 returns_optional_attentions = ( 

100 "output_attentions" in inspect.signature(self._original_block_forward).parameters 

101 ) 

102 

103 def patched_forward( 

104 block_self, 

105 hidden_states, 

106 attention_mask=None, 

107 position_bias=None, 

108 encoder_hidden_states=None, 

109 encoder_attention_mask=None, 

110 encoder_decoder_position_bias=None, 

111 layer_head_mask=None, 

112 cross_attn_layer_head_mask=None, 

113 past_key_values=None, 

114 use_cache=False, 

115 output_attentions=False, 

116 return_dict=True, 

117 cache_position=None, 

118 **kwargs, 

119 ): 

120 """Patched T5 block forward with hooks.""" 

121 hidden_states = self.hook_in(hidden_states) 

122 if not hasattr(block_self, "layer"): 122 ↛ 123line 122 didn't jump to line 123 because the condition on line 122 was never true

123 raise RuntimeError(f"T5 block {block_self} does not have 'layer' attribute") 

124 layers = block_self.layer 

125 is_decoder_block = len(layers) == 3 

126 

127 # Transformers v5 dropped use_cache and layer_head_mask from some sublayers. 

128 self_attn_params = set(inspect.signature(layers[0].forward).parameters.keys()) 

129 self_attn_kwargs = dict( 

130 hidden_states=hidden_states, 

131 attention_mask=attention_mask, 

132 position_bias=position_bias, 

133 output_attentions=output_attentions, 

134 cache_position=cache_position, 

135 ) 

136 # The stack passes the shared Cache as past_key_values; dropping it 

137 # leaves the cache empty and breaks incremental decoding. 

138 if "past_key_values" in self_attn_params: 

139 self_attn_kwargs["past_key_values"] = past_key_values 

140 if "use_cache" in self_attn_params: 

141 self_attn_kwargs["use_cache"] = use_cache 

142 if "layer_head_mask" in self_attn_params: 142 ↛ 143line 142 didn't jump to line 143 because the condition on line 142 was never true

143 self_attn_kwargs["layer_head_mask"] = layer_head_mask 

144 self_attention_outputs = layers[0](**self_attn_kwargs) 

145 cross_attention_outputs = None 

146 hidden_states = _clamp_fp16_inf(self_attention_outputs[0]) 

147 # Keep self-attention outputs and relative position weights 

148 # attention_outputs contains: (position_bias,) or (position_bias, attn_weights) 

149 attention_outputs = self_attention_outputs[1:] 

150 hidden_states = self.hook_resid_mid(hidden_states) 

151 if is_decoder_block and encoder_hidden_states is not None: 

152 cross_attn_params = set(inspect.signature(layers[1].forward).parameters.keys()) 

153 cross_attn_kwargs = dict( 

154 hidden_states=hidden_states, 

155 key_value_states=encoder_hidden_states, 

156 attention_mask=encoder_attention_mask, 

157 position_bias=encoder_decoder_position_bias, 

158 output_attentions=output_attentions, 

159 cache_position=cache_position, 

160 ) 

161 if "past_key_values" in cross_attn_params: 161 ↛ 163line 161 didn't jump to line 163 because the condition on line 161 was always true

162 cross_attn_kwargs["past_key_values"] = past_key_values 

163 if "use_cache" in cross_attn_params: 163 ↛ 164line 163 didn't jump to line 164 because the condition on line 163 was never true

164 cross_attn_kwargs["use_cache"] = use_cache 

165 if "layer_head_mask" in cross_attn_params: 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true

166 cross_attn_kwargs["layer_head_mask"] = cross_attn_layer_head_mask 

167 cross_attention_outputs = layers[1](**cross_attn_kwargs) 

168 hidden_states = _clamp_fp16_inf(cross_attention_outputs[0]) 

169 if hasattr(self, "hook_resid_mid2"): 169 ↛ 172line 169 didn't jump to line 172 because the condition on line 169 was always true

170 hidden_states = self.hook_resid_mid2(hidden_states) 

171 # Keep cross-attention outputs and relative position weights 

172 attention_outputs = attention_outputs + cross_attention_outputs[1:] 

173 ff_layer_idx = 2 if is_decoder_block else 1 

174 feed_forward_outputs = layers[ff_layer_idx](hidden_states) 

175 # T5LayerFF returns a tensor, not a tuple 

176 if isinstance(feed_forward_outputs, tuple): 176 ↛ 177line 176 didn't jump to line 177 because the condition on line 176 was never true

177 hidden_states = feed_forward_outputs[0] 

178 else: 

179 hidden_states = feed_forward_outputs 

180 hidden_states = _clamp_fp16_inf(hidden_states) 

181 hidden_states = self.hook_out(hidden_states) 

182 if not returns_optional_attentions: 

183 # hidden-states, self-attention position bias, cross-attention position bias 

184 return ( 

185 hidden_states, 

186 self_attention_outputs[1], 

187 cross_attention_outputs[1] if cross_attention_outputs is not None else None, 

188 ) 

189 outputs: tuple[Any, ...] = (hidden_states,) 

190 # Return: hidden-states, (self-attention position bias), (self-attention weights), 

191 # (cross-attention position bias), (cross-attention weights) 

192 return outputs + attention_outputs 

193 

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

195 

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

197 """Forward pass through the block bridge. 

198 

199 Args: 

200 *args: Input arguments 

201 **kwargs: Input keyword arguments 

202 

203 Returns: 

204 The output from the original component 

205 """ 

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

207 raise RuntimeError( 

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

209 ) 

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

211 return output 

212 

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

214 """Get the expected TransformerLens parameter names for this block. 

215 

216 Args: 

217 prefix: Prefix to add to parameter names (e.g., "blocks.0") 

218 

219 Returns: 

220 List of expected parameter names in TransformerLens format 

221 """ 

222 param_names = [] 

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

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

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

226 return param_names 

227 

228 def get_list_size(self) -> int: 

229 """Get the number of transformer blocks. 

230 

231 Returns: 

232 Number of layers in the model 

233 """ 

234 if self.config is None: 

235 return 0 

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