Coverage for transformer_lens/model_bridge/generalized_components/siglip_vision_encoder.py: 81%
56 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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 types import SimpleNamespace
7from typing import Any, Dict, Optional
9import torch
11from transformer_lens.hook_points import HookPoint
12from transformer_lens.model_bridge.generalized_components.attention import (
13 AttentionBridge,
14)
15from transformer_lens.model_bridge.generalized_components.base import (
16 GeneralizedComponent,
17)
18from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
19from transformer_lens.model_bridge.generalized_components.mlp import MLPBridge
20from transformer_lens.model_bridge.generalized_components.normalization import (
21 NormalizationBridge,
22)
25def _vision_attention_config(config: Any) -> Any:
26 """A config view carrying the vision tower's dims, not the language model's.
28 AttentionBridge reshapes its q/k/v/z hooks by ``config.n_heads`` at fire time,
29 so handing it the language config reshapes vision activations by the wrong head
30 count -- granite-docling runs 9 text heads over 576 dims against the tower's 12
31 over 768. Reads there are hasattr-guarded, so the dims are all this needs.
32 """
33 n_heads = getattr(config, "vision_num_heads", None)
34 d_model = getattr(config, "vision_hidden_size", None)
35 if not n_heads or not d_model:
36 return config
37 return SimpleNamespace(n_heads=n_heads, d_model=d_model, d_head=d_model // n_heads)
40class SiglipVisionEncoderLayerBridge(GeneralizedComponent):
41 """Bridge for a single SigLIP encoder layer.
43 SigLIP encoder layers have:
44 - layer_norm1: LayerNorm
45 - self_attn: SiglipAttention
46 - layer_norm2: LayerNorm
47 - mlp: SiglipMLP
48 """
50 is_list_item: bool = True
51 hook_aliases = {
52 "hook_resid_pre": "hook_in",
53 "hook_resid_post": "hook_out",
54 "hook_attn_in": "attn.hook_in",
55 "hook_attn_out": "attn.hook_out",
56 "hook_mlp_in": "mlp.hook_in",
57 "hook_mlp_out": "mlp.hook_out",
58 }
60 def __init__(
61 self,
62 name: str,
63 config: Optional[Any] = None,
64 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
65 ):
66 """Initialize the SigLIP encoder layer bridge.
68 Args:
69 name: The name of this component (e.g., "encoder.layers")
70 config: Optional configuration object
71 submodules: Dictionary of submodules to register
72 """
73 default_submodules: Dict[str, GeneralizedComponent] = {
74 "ln1": NormalizationBridge(name="layer_norm1", config=config),
75 "attn": AttentionBridge(
76 name="self_attn",
77 config=_vision_attention_config(config),
78 submodules={
79 "q": LinearBridge(name="q_proj"),
80 "k": LinearBridge(name="k_proj"),
81 "v": LinearBridge(name="v_proj"),
82 # SigLIP names the output projection out_proj, not o_proj.
83 "o": LinearBridge(name="out_proj"),
84 },
85 ),
86 "ln2": NormalizationBridge(name="layer_norm2", config=config),
87 "mlp": MLPBridge(
88 name="mlp",
89 config=config,
90 submodules={
91 "in": LinearBridge(name="fc1"),
92 "out": LinearBridge(name="fc2"),
93 },
94 ),
95 }
96 if submodules: 96 ↛ 97line 96 didn't jump to line 97 because the condition on line 96 was never true
97 default_submodules.update(submodules)
98 super().__init__(name, config, submodules=default_submodules)
100 def forward(
101 self,
102 hidden_states: torch.Tensor,
103 attention_mask: Optional[torch.Tensor] = None,
104 **kwargs: Any,
105 ) -> torch.Tensor:
106 """Forward pass through the vision encoder layer.
108 Args:
109 hidden_states: Input hidden states from previous layer
110 attention_mask: Optional attention mask
111 **kwargs: Additional arguments
113 Returns:
114 Output hidden states
115 """
116 if self.original_component is None: 116 ↛ 117line 116 didn't jump to line 117 because the condition on line 116 was never true
117 raise RuntimeError(
118 f"Original component not set for {self.name}. Call set_original_component() first."
119 )
121 hidden_states = self.hook_in(hidden_states)
122 output = self.original_component(hidden_states, attention_mask=attention_mask, **kwargs)
124 if isinstance(output, tuple): 124 ↛ 125line 124 didn't jump to line 125 because the condition on line 124 was never true
125 output = (self.hook_out(output[0]),) + output[1:]
126 else:
127 output = self.hook_out(output)
129 return output
132class SiglipVisionEncoderBridge(GeneralizedComponent):
133 """Bridge for the complete SigLIP vision encoder.
135 The SigLIP vision tower consists of:
136 - vision_model.embeddings: Patch + position embeddings
137 - vision_model.encoder.layers[]: Stack of encoder layers
138 - post_layernorm: Final layer norm
140 This bridge wraps the entire vision tower to provide hooks for
141 interpretability of the vision processing pipeline.
142 """
144 hook_aliases = {
145 "hook_vision_embed": "embeddings.hook_out",
146 "hook_vision_out": "hook_out",
147 }
149 def __init__(
150 self,
151 name: str,
152 config: Optional[Any] = None,
153 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
154 ):
155 """Initialize the SigLIP vision encoder bridge.
157 Args:
158 name: The name of this component (e.g., "model.vision_tower")
159 config: Optional configuration object
160 submodules: Dictionary of submodules to register
161 """
162 # All submodule names are resolved relative to the parent's
163 # original_component (a SiglipVisionModel) by setup_submodules().
164 # SiglipVisionModel wraps a SiglipVisionTransformer as .vision_model till
165 # transformers version 5.6.0
166 # post_layernorm is nn.LayerNorm; NormalizationBridge introspects the
167 # wrapped module so the RMSNorm-LM config (Gemma 3, LLaVA) doesn't leak.
168 default_submodules = {
169 "embeddings": GeneralizedComponent(name="vision_model.embeddings"),
170 # Pass config down: without it the layer's attention bridge inherits the
171 # raw HF vision config, which spells its head count num_attention_heads
172 # and so reads as n_heads=1 when the q/k/v/z hooks reshape.
173 "encoder_layers": SiglipVisionEncoderLayerBridge(
174 name="vision_model.encoder.layers", config=config
175 ),
176 "post_layernorm": NormalizationBridge(
177 name="vision_model.post_layernorm", config=config
178 ),
179 }
181 if submodules: 181 ↛ 182line 181 didn't jump to line 182 because the condition on line 181 was never true
182 default_submodules.update(submodules)
184 super().__init__(name, config, submodules=default_submodules)
186 # Additional hooks for vision-specific processing
187 self.hook_patch_embed = HookPoint() # After patch embedding
188 self.hook_pos_embed = HookPoint() # After position embedding added
190 def forward(
191 self,
192 pixel_values: torch.Tensor,
193 **kwargs: Any,
194 ) -> Any:
195 """Forward pass through the vision encoder.
197 Args:
198 pixel_values: Input image tensor [batch, channels, height, width]
199 **kwargs: Additional arguments
201 Returns:
202 Whatever the wrapped module returns, hooked in place: a
203 ``BaseModelOutput``, a tuple, or a bare tensor of vision embeddings
204 [batch, num_patches, hidden_size]. Callers see HF's own shape, so the
205 return is deliberately not narrowed to a tensor.
206 """
207 if self.original_component is None: 207 ↛ 208line 207 didn't jump to line 208 because the condition on line 207 was never true
208 raise RuntimeError(
209 f"Original component not set for {self.name}. Call set_original_component() first."
210 )
212 pixel_values = self.hook_in(pixel_values)
214 output = self.original_component(pixel_values, **kwargs)
216 # Handle tuple output (some models return (hidden_states, ...))
217 if isinstance(output, tuple): 217 ↛ 218line 217 didn't jump to line 218 because the condition on line 217 was never true
218 output = (self.hook_out(output[0]),) + output[1:]
219 elif hasattr(output, "last_hidden_state"): 219 ↛ 223line 219 didn't jump to line 223 because the condition on line 219 was always true
220 # Handle BaseModelOutput-like returns
221 output.last_hidden_state = self.hook_out(output.last_hidden_state)
222 else:
223 output = self.hook_out(output)
225 return output
227 def set_original_component(self, original_component: torch.nn.Module) -> None:
228 """Set the original component that this bridge wraps.
229 Note that SiglipVisionModel used to wrap a inner object as .vision_model before
230 transformers version 5.6.0, but after that it must directly be used.
231 This is a temporary hack to fix that till the transformers version is bumped.
233 Args:
234 original_component: The original transformer component to wrap
235 """
236 if not hasattr(original_component, "vision_model"):
237 # We should bypass any pytorch module registration.
238 object.__setattr__(original_component, "vision_model", original_component)
239 super().set_original_component(original_component)