Coverage for transformer_lens/model_bridge/generalized_components/siglip_vision_encoder.py: 79%
50 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""SigLIP Vision Encoder bridge component.
3This module contains the bridge component for SigLIP vision encoder layers
4used in multimodal models like Gemma 3 and MedGemma.
5"""
6from typing import Any, Dict, Optional
8import torch
10from transformer_lens.hook_points import HookPoint
11from transformer_lens.model_bridge.generalized_components.attention import (
12 AttentionBridge,
13)
14from transformer_lens.model_bridge.generalized_components.base import (
15 GeneralizedComponent,
16)
17from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
18from transformer_lens.model_bridge.generalized_components.mlp import MLPBridge
19from transformer_lens.model_bridge.generalized_components.normalization import (
20 NormalizationBridge,
21)
22from transformer_lens.model_bridge.generalized_components.vision_encoder import (
23 vision_attention_config,
24)
27class SiglipVisionEncoderLayerBridge(GeneralizedComponent):
28 """Bridge for a single SigLIP encoder layer.
30 SigLIP encoder layers have:
31 - layer_norm1: LayerNorm
32 - self_attn: SiglipAttention
33 - layer_norm2: LayerNorm
34 - mlp: SiglipMLP
35 """
37 is_list_item: bool = True
38 hook_aliases = {
39 "hook_resid_pre": "hook_in",
40 "hook_resid_post": "hook_out",
41 "hook_attn_in": "attn.hook_in",
42 "hook_attn_out": "attn.hook_out",
43 "hook_mlp_in": "mlp.hook_in",
44 "hook_mlp_out": "mlp.hook_out",
45 }
47 def __init__(
48 self,
49 name: str,
50 config: Optional[Any] = None,
51 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
52 ):
53 """Initialize the SigLIP encoder layer bridge.
55 Args:
56 name: The name of this component (e.g., "encoder.layers")
57 config: Optional configuration object
58 submodules: Dictionary of submodules to register
59 """
60 # use_native_layernorm_autograd makes the bridge run the wrapped LayerNorm's
61 # own forward, so the vision tower's eps applies instead of the LM config's.
62 default_submodules: Dict[str, GeneralizedComponent] = {
63 "ln1": NormalizationBridge(
64 name="layer_norm1", config=config, use_native_layernorm_autograd=True
65 ),
66 "attn": AttentionBridge(
67 name="self_attn",
68 config=vision_attention_config(config),
69 submodules={
70 "q": LinearBridge(name="q_proj"),
71 "k": LinearBridge(name="k_proj"),
72 "v": LinearBridge(name="v_proj"),
73 # SigLIP names the output projection out_proj, not o_proj.
74 "o": LinearBridge(name="out_proj"),
75 },
76 ),
77 "ln2": NormalizationBridge(
78 name="layer_norm2", config=config, use_native_layernorm_autograd=True
79 ),
80 "mlp": MLPBridge(
81 name="mlp",
82 config=config,
83 submodules={
84 "in": LinearBridge(name="fc1"),
85 "out": LinearBridge(name="fc2"),
86 },
87 ),
88 }
89 if submodules: 89 ↛ 90line 89 didn't jump to line 90 because the condition on line 89 was never true
90 default_submodules.update(submodules)
91 super().__init__(name, config, submodules=default_submodules)
93 def forward(
94 self,
95 hidden_states: torch.Tensor,
96 attention_mask: Optional[torch.Tensor] = None,
97 **kwargs: Any,
98 ) -> torch.Tensor:
99 """Forward pass through the vision encoder layer.
101 Args:
102 hidden_states: Input hidden states from previous layer
103 attention_mask: Optional attention mask
104 **kwargs: Additional arguments
106 Returns:
107 Output hidden states
108 """
109 if self.original_component is None: 109 ↛ 110line 109 didn't jump to line 110 because the condition on line 109 was never true
110 raise RuntimeError(
111 f"Original component not set for {self.name}. Call set_original_component() first."
112 )
114 hidden_states = self.hook_in(hidden_states)
115 output = self.original_component(hidden_states, attention_mask=attention_mask, **kwargs)
117 if isinstance(output, tuple): 117 ↛ 118line 117 didn't jump to line 118 because the condition on line 117 was never true
118 output = (self.hook_out(output[0]),) + output[1:]
119 else:
120 output = self.hook_out(output)
122 return output
125class SiglipVisionEncoderBridge(GeneralizedComponent):
126 """Bridge for the complete SigLIP vision encoder.
128 The SigLIP vision tower consists of:
129 - vision_model.embeddings: Patch + position embeddings
130 - vision_model.encoder.layers[]: Stack of encoder layers
131 - post_layernorm: Final layer norm
133 This bridge wraps the entire vision tower to provide hooks for
134 interpretability of the vision processing pipeline.
135 """
137 hook_aliases = {
138 "hook_vision_embed": "embeddings.hook_out",
139 "hook_vision_out": "hook_out",
140 }
142 def __init__(
143 self,
144 name: str,
145 config: Optional[Any] = None,
146 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
147 ):
148 """Initialize the SigLIP vision encoder bridge.
150 Args:
151 name: The name of this component (e.g., "model.vision_tower")
152 config: Optional configuration object
153 submodules: Dictionary of submodules to register
154 """
155 # All submodule names are resolved relative to the parent's
156 # original_component (a SiglipVisionModel) by setup_submodules().
157 # SiglipVisionModel wraps a SiglipVisionTransformer as .vision_model till
158 # transformers version 5.6.0
159 # post_layernorm is nn.LayerNorm; NormalizationBridge introspects the
160 # wrapped module so the RMSNorm-LM config (Gemma 3, LLaVA) doesn't leak,
161 # and the native-autograd flag keeps the vision eps over the LM config's.
162 default_submodules = {
163 "embeddings": GeneralizedComponent(name="vision_model.embeddings"),
164 # Pass config down: without it the layer's attention bridge inherits the
165 # raw HF vision config, which spells its head count num_attention_heads
166 # and so reads as n_heads=1 when the q/k/v/z hooks reshape.
167 "encoder_layers": SiglipVisionEncoderLayerBridge(
168 name="vision_model.encoder.layers", config=config
169 ),
170 "post_layernorm": NormalizationBridge(
171 name="vision_model.post_layernorm",
172 config=config,
173 use_native_layernorm_autograd=True,
174 ),
175 }
177 if submodules: 177 ↛ 178line 177 didn't jump to line 178 because the condition on line 177 was never true
178 default_submodules.update(submodules)
180 super().__init__(name, config, submodules=default_submodules)
182 # Additional hooks for vision-specific processing
183 self.hook_patch_embed = HookPoint() # After patch embedding
184 self.hook_pos_embed = HookPoint() # After position embedding added
186 def forward(
187 self,
188 pixel_values: torch.Tensor,
189 **kwargs: Any,
190 ) -> Any:
191 """Forward pass through the vision encoder.
193 Args:
194 pixel_values: Input image tensor [batch, channels, height, width]
195 **kwargs: Additional arguments
197 Returns:
198 Whatever the wrapped module returns, hooked in place: a
199 ``BaseModelOutput``, a tuple, or a bare tensor of vision embeddings
200 [batch, num_patches, hidden_size]. Callers see HF's own shape, so the
201 return is deliberately not narrowed to a tensor.
202 """
203 if self.original_component is None: 203 ↛ 204line 203 didn't jump to line 204 because the condition on line 203 was never true
204 raise RuntimeError(
205 f"Original component not set for {self.name}. Call set_original_component() first."
206 )
208 pixel_values = self.hook_in(pixel_values)
210 output = self.original_component(pixel_values, **kwargs)
212 # Handle tuple output (some models return (hidden_states, ...))
213 if isinstance(output, tuple): 213 ↛ 214line 213 didn't jump to line 214 because the condition on line 213 was never true
214 output = (self.hook_out(output[0]),) + output[1:]
215 elif hasattr(output, "last_hidden_state"): 215 ↛ 219line 215 didn't jump to line 219 because the condition on line 215 was always true
216 # Handle BaseModelOutput-like returns
217 output.last_hidden_state = self.hook_out(output.last_hidden_state)
218 else:
219 output = self.hook_out(output)
221 return output
223 def set_original_component(self, original_component: torch.nn.Module) -> None:
224 """Set the original component that this bridge wraps.
225 Note that SiglipVisionModel used to wrap a inner object as .vision_model before
226 transformers version 5.6.0, but after that it must directly be used.
227 This is a temporary hack to fix that till the transformers version is bumped.
229 Args:
230 original_component: The original transformer component to wrap
231 """
232 if not hasattr(original_component, "vision_model"):
233 # We should bypass any pytorch module registration.
234 object.__setattr__(original_component, "vision_model", original_component)
235 super().set_original_component(original_component)