Coverage for transformer_lens/model_bridge/generalized_components/t5_block.py: 72%
108 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""T5-specific block bridge component.
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
8import types
9from typing import Any, Callable, Dict, Optional
11import torch
13from transformer_lens.hook_points import HookPoint
14from transformer_lens.model_bridge.generalized_components.base import (
15 GeneralizedComponent,
16)
19def _clamp_fp16_inf(hidden_states: torch.Tensor) -> torch.Tensor:
20 """HF T5Block clamps fp16 hidden states after every sublayer; the patched
21 forward must too, or fp16 runs overflow exactly where HF's do not."""
22 if hidden_states.dtype == torch.float16:
23 clamp_value = torch.where(
24 torch.isinf(hidden_states).any(),
25 torch.finfo(hidden_states.dtype).max - 1000,
26 torch.finfo(hidden_states.dtype).max,
27 )
28 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
29 return hidden_states
32class T5BlockBridge(GeneralizedComponent):
33 """Bridge component for T5 transformer blocks.
35 T5 has two types of blocks:
36 - Encoder blocks: 2 layers (self-attention, feed-forward)
37 - Decoder blocks: 3 layers (self-attention, cross-attention, feed-forward)
39 This bridge handles both types based on the presence of cross-attention.
40 """
42 is_list_item: bool = True
43 hook_aliases = {"hook_resid_pre": "hook_in", "hook_resid_post": "hook_out"}
45 def __init__(
46 self,
47 name: str,
48 config: Optional[Any] = None,
49 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
50 is_decoder: bool = False,
51 ):
52 """Initialize the T5 block bridge.
54 Args:
55 name: The name of the component in the model
56 config: Optional configuration
57 submodules: Dictionary of submodules to register
58 is_decoder: Whether this is a decoder block (has cross-attention)
59 """
60 super().__init__(name, config, submodules=submodules or {})
61 # T5 adds the residual INSIDE each sublayer, so the wrapped module's
62 # output already IS the added contribution. Keyed to what this adapter
63 # declared, since names differ across the family.
64 self.hook_aliases = dict(type(self).hook_aliases)
65 declared = submodules or {}
66 if "attn" in declared:
67 self.hook_aliases["hook_attn_out"] = "attn.hook_out"
68 elif "self_attn" in declared: 68 ↛ 70line 68 didn't jump to line 70 because the condition on line 68 was always true
69 self.hook_aliases["hook_attn_out"] = "self_attn.hook_out"
70 if "cross_attn" in declared:
71 self.hook_aliases["hook_cross_attn_out"] = "cross_attn.hook_out"
72 if "mlp" in declared: 72 ↛ 74line 72 didn't jump to line 74 because the condition on line 72 was always true
73 self.hook_aliases["hook_mlp_out"] = "mlp.hook_out"
74 self.is_decoder = is_decoder
75 self.hook_resid_mid = HookPoint()
76 self._register_hook("hook_resid_mid", self.hook_resid_mid)
77 if is_decoder:
78 self.hook_resid_mid2 = HookPoint()
79 self._register_hook("hook_resid_mid2", self.hook_resid_mid2)
80 self._original_block_forward: Optional[Callable[..., Any]] = None
82 def set_original_component(self, component: torch.nn.Module):
83 """Set the original component and monkey-patch its forward method.
85 Args:
86 component: The original PyTorch module to wrap
87 """
88 super().set_original_component(component)
89 self._patch_t5_block_forward()
91 def _patch_t5_block_forward(self):
92 """Monkey-patch the T5 block's forward method to insert hooks."""
93 if self.original_component is None: 93 ↛ 94line 93 didn't jump to line 94 because the condition on line 93 was never true
94 return
95 self._original_block_forward = self.original_component.forward
97 def patched_forward(
98 block_self,
99 hidden_states,
100 attention_mask=None,
101 position_bias=None,
102 encoder_hidden_states=None,
103 encoder_attention_mask=None,
104 encoder_decoder_position_bias=None,
105 layer_head_mask=None,
106 cross_attn_layer_head_mask=None,
107 past_key_value=None,
108 use_cache=False,
109 output_attentions=False,
110 return_dict=True,
111 cache_position=None,
112 **kwargs,
113 ):
114 """Patched T5 block forward with hooks."""
115 import inspect
117 hidden_states = self.hook_in(hidden_states)
118 if not hasattr(block_self, "layer"): 118 ↛ 119line 118 didn't jump to line 119 because the condition on line 118 was never true
119 raise RuntimeError(f"T5 block {block_self} does not have 'layer' attribute")
120 layers = block_self.layer
121 is_decoder_block = len(layers) == 3
123 # Check which parameters are accepted by the layer forward methods
124 # (Transformers v5 removed past_key_value, use_cache, layer_head_mask)
125 self_attn_params = set(inspect.signature(layers[0].forward).parameters.keys())
127 if "past_key_value" in self_attn_params and past_key_value is not None: 127 ↛ 128line 127 didn't jump to line 128 because the condition on line 127 was never true
128 if not is_decoder_block:
129 expected_num_past_key_values = 0
130 else:
131 expected_num_past_key_values = 2
132 if len(past_key_value) != expected_num_past_key_values:
133 raise ValueError(
134 f"There should be {expected_num_past_key_values} past states. Got {len(past_key_value)}."
135 )
136 self_attn_past_key_value = past_key_value[:2] if is_decoder_block else None
137 cross_attn_past_key_value = past_key_value[2:4] if is_decoder_block else None
138 else:
139 self_attn_past_key_value = None
140 cross_attn_past_key_value = None
141 self_attn_kwargs = dict(
142 hidden_states=hidden_states,
143 attention_mask=attention_mask,
144 position_bias=position_bias,
145 output_attentions=output_attentions,
146 cache_position=cache_position,
147 )
148 # Conditionally pass parameters removed in Transformers v5
149 if "past_key_value" in self_attn_params: 149 ↛ 150line 149 didn't jump to line 150 because the condition on line 149 was never true
150 self_attn_kwargs["past_key_value"] = self_attn_past_key_value
151 if "use_cache" in self_attn_params:
152 self_attn_kwargs["use_cache"] = use_cache
153 if "layer_head_mask" in self_attn_params: 153 ↛ 154line 153 didn't jump to line 154 because the condition on line 153 was never true
154 self_attn_kwargs["layer_head_mask"] = layer_head_mask
155 self_attention_outputs = layers[0](**self_attn_kwargs)
156 hidden_states = _clamp_fp16_inf(self_attention_outputs[0])
157 # Keep self-attention outputs and relative position weights
158 # attention_outputs contains: (position_bias,) or (position_bias, attn_weights)
159 attention_outputs = self_attention_outputs[1:]
160 hidden_states = self.hook_resid_mid(hidden_states)
161 if is_decoder_block and encoder_hidden_states is not None:
162 cross_attn_params = set(inspect.signature(layers[1].forward).parameters.keys())
163 cross_attn_kwargs = dict(
164 hidden_states=hidden_states,
165 key_value_states=encoder_hidden_states,
166 attention_mask=encoder_attention_mask,
167 position_bias=encoder_decoder_position_bias,
168 output_attentions=output_attentions,
169 cache_position=cache_position,
170 )
171 if "past_key_value" in cross_attn_params: 171 ↛ 172line 171 didn't jump to line 172 because the condition on line 171 was never true
172 cross_attn_kwargs["past_key_value"] = cross_attn_past_key_value
173 if "use_cache" in cross_attn_params: 173 ↛ 174line 173 didn't jump to line 174 because the condition on line 173 was never true
174 cross_attn_kwargs["use_cache"] = use_cache
175 if "layer_head_mask" in cross_attn_params: 175 ↛ 176line 175 didn't jump to line 176 because the condition on line 175 was never true
176 cross_attn_kwargs["layer_head_mask"] = cross_attn_layer_head_mask
177 cross_attention_outputs = layers[1](**cross_attn_kwargs)
178 hidden_states = _clamp_fp16_inf(cross_attention_outputs[0])
179 if hasattr(self, "hook_resid_mid2"): 179 ↛ 182line 179 didn't jump to line 182 because the condition on line 179 was always true
180 hidden_states = self.hook_resid_mid2(hidden_states)
181 # Keep cross-attention outputs and relative position weights
182 attention_outputs = attention_outputs + cross_attention_outputs[1:]
183 ff_layer_idx = 2 if is_decoder_block else 1
184 feed_forward_outputs = layers[ff_layer_idx](hidden_states)
185 # T5LayerFF returns a tensor, not a tuple
186 if isinstance(feed_forward_outputs, tuple): 186 ↛ 187line 186 didn't jump to line 187 because the condition on line 186 was never true
187 hidden_states = feed_forward_outputs[0]
188 else:
189 hidden_states = feed_forward_outputs
190 hidden_states = _clamp_fp16_inf(hidden_states)
191 hidden_states = self.hook_out(hidden_states)
192 outputs: tuple[Any, ...] = (hidden_states,)
193 # Return: hidden-states, (self-attention position bias), (self-attention weights),
194 # (cross-attention position bias), (cross-attention weights)
195 return outputs + attention_outputs
197 self.original_component.forward = types.MethodType(patched_forward, self.original_component)
199 def forward(self, *args: Any, **kwargs: Any) -> Any:
200 """Forward pass through the block bridge.
202 Args:
203 *args: Input arguments
204 **kwargs: Input keyword arguments
206 Returns:
207 The output from the original component
208 """
209 if self.original_component is None: 209 ↛ 210line 209 didn't jump to line 210 because the condition on line 209 was never true
210 raise RuntimeError(
211 f"Original component not set for {self.name}. Call set_original_component() first."
212 )
213 output = self.original_component(*args, **kwargs)
214 return output
216 def get_expected_parameter_names(self, prefix: str = "") -> list[str]:
217 """Get the expected TransformerLens parameter names for this block.
219 Args:
220 prefix: Prefix to add to parameter names (e.g., "blocks.0")
222 Returns:
223 List of expected parameter names in TransformerLens format
224 """
225 param_names = []
226 for sub_name, sub_component in self.submodules.items():
227 sub_prefix = f"{prefix}.{sub_name}" if prefix else sub_name
228 param_names.extend(sub_component.get_expected_parameter_names(sub_prefix))
229 return param_names
231 def get_list_size(self) -> int:
232 """Get the number of transformer blocks.
234 Returns:
235 Number of layers in the model
236 """
237 if self.config is None:
238 return 0
239 return getattr(self.config, "n_layers", 0)