Coverage for transformer_lens/model_bridge/generalized_components/clip_vision_encoder.py: 52%
50 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""CLIP Vision Encoder bridge component.
3This module contains the bridge component for CLIP vision encoder layers
4used in multimodal models like LLava.
5"""
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)
23from transformer_lens.model_bridge.generalized_components.vision_encoder import (
24 vision_attention_config,
25)
28class CLIPVisionEncoderLayerBridge(GeneralizedComponent):
29 """Bridge for a single CLIP encoder layer.
31 CLIP encoder layers have:
32 - layer_norm1: LayerNorm
33 - self_attn: CLIPAttention
34 - layer_norm2: LayerNorm
35 - mlp: CLIPMLP
36 """
38 is_list_item: bool = True
39 hook_aliases = {
40 "hook_resid_pre": "hook_in",
41 "hook_resid_post": "hook_out",
42 "hook_attn_in": "attn.hook_in",
43 "hook_attn_out": "attn.hook_out",
44 "hook_mlp_in": "mlp.hook_in",
45 "hook_mlp_out": "mlp.hook_out",
46 }
48 def __init__(
49 self,
50 name: str,
51 config: Optional[Any] = None,
52 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
53 ):
54 """Initialize the CLIP encoder layer bridge.
56 Args:
57 name: The name of this component (e.g., "encoder.layers")
58 config: Optional configuration object
59 submodules: Dictionary of submodules to register
60 """
61 default_submodules: Dict[str, GeneralizedComponent] = {
62 "ln1": NormalizationBridge(
63 name="layer_norm1", config=config, use_native_layernorm_autograd=True
64 ),
65 "attn": AttentionBridge(
66 name="self_attn",
67 config=vision_attention_config(config),
68 submodules={
69 "q": LinearBridge(name="q_proj"),
70 "k": LinearBridge(name="k_proj"),
71 "v": LinearBridge(name="v_proj"),
72 "o": LinearBridge(name="out_proj"),
73 },
74 ),
75 "ln2": NormalizationBridge(
76 name="layer_norm2", config=config, use_native_layernorm_autograd=True
77 ),
78 "mlp": MLPBridge(
79 name="mlp",
80 config=config,
81 submodules={
82 "in": LinearBridge(name="fc1"),
83 "out": LinearBridge(name="fc2"),
84 },
85 ),
86 }
87 if submodules: 87 ↛ 88line 87 didn't jump to line 88 because the condition on line 87 was never true
88 default_submodules.update(submodules)
89 super().__init__(name, config, submodules=default_submodules)
91 def forward(
92 self,
93 hidden_states: torch.Tensor,
94 attention_mask: Optional[torch.Tensor] = None,
95 causal_attention_mask: Optional[torch.Tensor] = None,
96 **kwargs: Any,
97 ) -> torch.Tensor:
98 """Forward pass through the vision encoder layer.
100 Args:
101 hidden_states: Input hidden states from previous layer
102 attention_mask: Optional attention mask
103 causal_attention_mask: Optional causal attention mask (used by CLIP encoder)
104 **kwargs: Additional arguments
106 Returns:
107 Output hidden states
108 """
109 if self.original_component is None:
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(
116 hidden_states,
117 attention_mask=attention_mask,
118 causal_attention_mask=causal_attention_mask,
119 **kwargs,
120 )
122 if isinstance(output, tuple):
123 output = (self.hook_out(output[0]),) + output[1:]
124 else:
125 output = self.hook_out(output)
127 return output
130class CLIPVisionEncoderBridge(GeneralizedComponent):
131 """Bridge for the complete CLIP vision encoder.
133 The CLIP vision tower consists of:
134 - vision_model.embeddings: Patch + position + CLS token embeddings
135 - vision_model.pre_layrnorm: LayerNorm before encoder layers
136 - vision_model.encoder.layers[]: Stack of encoder layers
137 - vision_model.post_layernorm: Final layer norm
139 This bridge wraps the entire vision tower to provide hooks for
140 interpretability of the vision processing pipeline.
141 """
143 hook_aliases = {
144 "hook_vision_embed": "embeddings.hook_out",
145 "hook_vision_out": "hook_out",
146 }
148 def __init__(
149 self,
150 name: str,
151 config: Optional[Any] = None,
152 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
153 ):
154 """Initialize the CLIP vision encoder bridge.
156 Args:
157 name: The name of this component (e.g., "vision_tower")
158 config: Optional configuration object
159 submodules: Dictionary of submodules to register
160 """
161 default_submodules: Dict[str, GeneralizedComponent] = {
162 "embeddings": GeneralizedComponent(name="vision_model.embeddings"),
163 # Native-autograd runs the wrapped LayerNorm's own forward, so the
164 # vision tower's eps applies instead of the LM config's.
165 "pre_layernorm": NormalizationBridge(
166 name="vision_model.pre_layrnorm",
167 config=config,
168 use_native_layernorm_autograd=True,
169 ),
170 "encoder_layers": CLIPVisionEncoderLayerBridge(
171 name="vision_model.encoder.layers", config=config
172 ),
173 "post_layernorm": NormalizationBridge(
174 name="vision_model.post_layernorm",
175 config=config,
176 use_native_layernorm_autograd=True,
177 ),
178 }
180 if submodules: 180 ↛ 181line 180 didn't jump to line 181 because the condition on line 180 was never true
181 default_submodules.update(submodules)
183 super().__init__(name, config, submodules=default_submodules)
185 # Additional hooks for vision-specific processing
186 self.hook_patch_embed = HookPoint() # After patch embedding
187 self.hook_pos_embed = HookPoint() # After position embedding added
189 def forward(
190 self,
191 pixel_values: torch.Tensor,
192 **kwargs: Any,
193 ) -> torch.Tensor:
194 """Forward pass through the vision encoder.
196 Args:
197 pixel_values: Input image tensor [batch, channels, height, width]
198 **kwargs: Additional arguments
200 Returns:
201 Vision embeddings [batch, num_patches, hidden_size]
202 """
203 if self.original_component is None:
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):
214 output = (self.hook_out(output[0]),) + output[1:]
215 elif hasattr(output, "last_hidden_state"):
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 CLIPVisionModel 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)