Coverage for transformer_lens/model_bridge/generalized_components/base.py: 84%
220 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Base class for generalized transformer components."""
2from __future__ import annotations
4import inspect
5import warnings
6from collections.abc import Callable
7from typing import Any, Dict, List, Optional, Union
9import torch
10import torch.nn as nn
12from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import (
13 BaseTensorConversion,
14)
15from transformer_lens.hook_points import HookPoint
18class CloneOutputUnderGradMixin(nn.Module):
19 """Clone the forward output so HF's in-place mutation cannot corrupt it.
21 Under grad, autograd forbids in-place writes to backward-hook views; under
22 no_grad, cached hook_out tensors alias the storage HF then rewrites.
23 Mix in ahead of a bridge class: ``class X(CloneOutputUnderGradMixin, LinearBridge)``.
24 """
26 def forward(self, *args: Any, **kwargs: Any) -> Any:
27 out = super().forward(*args, **kwargs)
28 if isinstance(out, torch.Tensor): 28 ↛ 30line 28 didn't jump to line 30 because the condition on line 28 was always true
29 out = out.clone()
30 return out
33# Bumped whenever any component rebinds its hook_aliases (MoEBridge's dense/
34# sparse dispatch, OlmoHybrid's per-layer selection). Alias caches key on it:
35# a rebind can leave the hook REGISTRY unchanged while changing what the
36# aliases point at, so size-based cache keys cannot see it.
37_ALIAS_GENERATION = 0
40def alias_generation() -> int:
41 """Current global alias-rebind generation."""
42 return _ALIAS_GENERATION
45def _bump_alias_generation() -> None:
46 global _ALIAS_GENERATION
47 _ALIAS_GENERATION += 1
50class GeneralizedComponent(nn.Module):
51 """Base class for generalized transformer components.
53 This class provides a standardized interface for transformer components
54 and handles hook registration and execution.
55 """
57 is_list_item: bool = False
58 hook_out_is_single_residual_stream: bool = False
59 compatibility_mode: bool = False
60 disable_warnings: bool = False
61 hook_aliases: Dict[str, Union[str, List[str]]] = {}
63 # Projection-hook protocol between container bridges (MLPBridge) and the
64 # projection bridges they wrap (LinearBridge / Conv1DBridge): the wrapper
65 # records that its hook_out fired so the container only re-fires it as a
66 # bypass fallback, and the container suppresses the wrapper's next hook_in
67 # after pre-firing it itself. See MLPBridge.forward.
68 _fired_hook_out: bool = False
69 _suppress_next_hook_in: bool = False
70 property_aliases: Dict[str, str] = {}
72 def __init__(
73 self,
74 name: Optional[str],
75 config: Optional[Any] = None,
76 submodules: Optional[Dict[str, "GeneralizedComponent"]] = None,
77 conversion_rule: Optional[BaseTensorConversion] = None,
78 hook_alias_overrides: Optional[Dict[str, str]] = None,
79 optional: bool = False,
80 ):
81 """Initialize the generalized component.
83 Args:
84 name: The name of this component (None if component has no container in remote model)
85 config: Optional configuration object for the component
86 submodules: Dictionary of GeneralizedComponent submodules to register
87 conversion_rule: Optional conversion rule for this component's hooks
88 hook_alias_overrides: Optional dictionary to override default hook aliases.
89 For example, {"hook_attn_out": "ln1_post.hook_out"} will make hook_attn_out
90 point to ln1_post.hook_out instead of the default value in self.hook_aliases.
91 optional: If True, setup skips this subtree when absent (hybrid architectures).
92 """
93 super().__init__()
94 self.name = name
95 self.config = config
96 self.submodules = submodules or {}
97 self.conversion_rule = conversion_rule
98 self.optional = optional
99 self._hook_registry: Dict[str, HookPoint] = {}
100 self._hook_alias_registry: Dict[str, Union[str, List[str]]] = {}
101 self._property_alias_registry: Dict[str, str] = {}
102 self.hook_in = HookPoint()
103 self.hook_out = HookPoint()
104 # real_components maps TL keys to (remote_path, actual_instance) tuples
105 # For list components, actual_instance will be a list of component instances
106 self.real_components: Dict[str, tuple] = {}
107 if self.conversion_rule is not None:
108 self.hook_in.hook_conversion = self.conversion_rule
109 self.hook_out.hook_conversion = self.conversion_rule
111 # Copy class-level hook_aliases and apply any overrides
112 if hook_alias_overrides is not None:
113 self.hook_aliases = self.__class__.hook_aliases.copy()
114 self.hook_aliases.update(hook_alias_overrides)
116 def _register_hook(self, name: str, hook: HookPoint) -> None:
117 """Register a hook in the component's hook registry."""
118 hook.name = name
119 self._hook_registry[name] = hook
121 def _register_aliases(self) -> None:
122 """Register aliases from class-level dictionaries.
124 This is called ONLY in enable_compatibility_mode() after weight processing.
125 It creates actual Python attributes/properties that directly reference the target objects.
127 Note: This should only be called when compatibility mode is enabled and after
128 weight processing is complete to ensure property aliases point to processed weights.
129 """
130 if self.hook_aliases:
131 self._hook_alias_registry.update(self.hook_aliases)
132 if self.property_aliases:
133 self._property_alias_registry.update(self.property_aliases)
134 for alias_name, target_path in self._hook_alias_registry.items():
135 resolved = False
136 if isinstance(target_path, list): 136 ↛ 137line 136 didn't jump to line 137 because the condition on line 136 was never true
137 for single_target in target_path:
138 try:
139 target_obj = self
140 for part in single_target.split("."):
141 target_obj = getattr(target_obj, part)
142 object.__setattr__(self, alias_name, target_obj)
143 resolved = True
144 break
145 except AttributeError:
146 continue
147 else:
148 try:
149 target_obj = self
150 for part in target_path.split("."):
151 target_obj = getattr(target_obj, part)
152 object.__setattr__(self, alias_name, target_obj)
153 resolved = True
154 except AttributeError:
155 pass
156 if not resolved:
157 # Surface drops instead of silently swallowing — some aliases are
158 # legitimately conditional on optional submodules, but an author
159 # needs to see which ones dropped at bridge-init.
160 warnings.warn(
161 f"Hook alias '{alias_name}' -> '{target_path}' on "
162 f"{type(self).__name__}(name={getattr(self, 'name', None)!r}) "
163 f"did not resolve; this hook will not be accessible.",
164 stacklevel=2,
165 )
166 for alias_name, target_path in self._property_alias_registry.items():
167 try:
168 target_obj = self
169 for part in target_path.split("."):
170 target_obj = getattr(target_obj, part)
171 object.__setattr__(self, alias_name, target_obj)
172 except AttributeError:
173 pass
175 def get_hooks(self) -> Dict[str, HookPoint]:
176 """Get all hooks registered in this component."""
177 hooks = self._hook_registry.copy()
178 if self.compatibility_mode and self._hook_alias_registry:
179 for alias_name in self._hook_alias_registry.keys():
180 if hasattr(self, alias_name): 180 ↛ 179line 180 didn't jump to line 179 because the condition on line 180 was always true
181 target_hook = getattr(self, alias_name)
182 if isinstance(target_hook, HookPoint): 182 ↛ 179line 182 didn't jump to line 179 because the condition on line 182 was always true
183 hooks[alias_name] = target_hook
184 return hooks
186 def _is_getattr_called_internally(self) -> bool:
187 """This function checks if the __getattr__ method was being called internally
188 (e.g by the setup process or run_with_cache).
189 """
190 for frame_info in inspect.stack():
191 if "setup_components" in frame_info.function or "run_with_cache" in frame_info.function:
192 return True
193 return False
195 def set_original_component(self, original_component: nn.Module) -> None:
196 """Set the original component that this bridge wraps.
198 Args:
199 original_component: The original transformer component to wrap
200 """
201 self.add_module("_original_component", original_component)
202 # An opaque wrapper (created with config=None) shadows the wrapped
203 # module's own config. HF forwards sometimes read a submodule's config
204 # directly (e.g. Qwen2Audio's forward does create_bidirectional_mask(
205 # config=self.audio_tower.config)), so inherit the real config to avoid
206 # exposing None. Components given an explicit config keep it.
207 if self.config is None:
208 self.config = getattr(original_component, "config", None)
210 @property
211 def original_component(self) -> Optional[nn.Module]:
212 """Get the original component."""
213 return self._modules.get("_original_component", None)
215 def add_hook(self, hook_fn: Callable[..., torch.Tensor], hook_name: str = "output") -> None:
216 """Add a hook function (HookedTransformer-compatible interface).
218 Args:
219 hook_fn: Function to call at this hook point
220 hook_name: Name of the hook point (defaults to "output")
221 """
222 if hook_name == "output":
223 self.hook_out.add_hook(hook_fn)
224 elif hook_name == "input":
225 self.hook_in.add_hook(hook_fn)
226 else:
227 raise ValueError(
228 f"Hook name '{hook_name}' not supported. Supported names are 'output' and 'input'."
229 )
231 def remove_hooks(self, hook_name: str | None = None) -> None:
232 """Remove hooks (HookedTransformer-compatible interface).
234 Args:
235 hook_name: Name of the hook point to remove. If None, removes all hooks.
236 """
237 if hook_name is None:
238 self.hook_in.remove_hooks(dir="both")
239 self.hook_out.remove_hooks(dir="both")
240 elif hook_name == "output":
241 self.hook_out.remove_hooks(dir="both")
242 elif hook_name == "input":
243 self.hook_in.remove_hooks(dir="both")
244 else:
245 raise ValueError(
246 f"Hook name '{hook_name}' not supported. Supported names are 'output' and 'input'."
247 )
249 def set_processed_weights(
250 self, weights: Dict[str, torch.Tensor], verbose: bool = False
251 ) -> None:
252 """Set the processed weights for use in compatibility mode.
254 This method stores processed weights as attributes on the component so they can be
255 used directly in the forward pass without modifying the original component.
257 Components should override this method to handle their specific weight structure.
258 The weights dict contains keys like "weight", "bias", "W_in", "W_out", etc.
260 If this component has submodules, this method will automatically distribute the
261 weights to those subcomponents using ProcessWeights.distribute_weights_to_components.
263 Args:
264 weights: Dictionary of processed weight tensors
265 verbose: If True, print detailed information about weight setting
266 """
267 if verbose: 267 ↛ 268line 267 didn't jump to line 268 because the condition on line 267 was never true
268 print(
269 f"\n set_processed_weights: {self.__class__.__name__} (name={getattr(self, 'name', 'unknown')})"
270 )
271 print(f" Received {len(weights)} weight keys")
273 # First, handle single-part keys (keys without ".") by setting them as parameters
274 # on the original component
275 if self.original_component is not None: 275 ↛ 304line 275 didn't jump to line 304 because the condition on line 275 was always true
276 for key, weight_tensor in weights.items():
277 # Only process keys without "." (single-part keys)
278 if "." not in key:
279 # Try to set the parameter on the original component
280 if hasattr(self.original_component, key):
281 param = getattr(self.original_component, key)
282 if param is not None and isinstance(param, torch.nn.Parameter):
283 if param.shape != weight_tensor.shape: 283 ↛ 284line 283 didn't jump to line 284 because the condition on line 283 was never true
284 raise ValueError(
285 f"Shape mismatch when setting weight '{key}' in {type(self.original_component).__name__}: "
286 f"existing param shape {param.shape} != new tensor shape {weight_tensor.shape}"
287 )
288 if verbose: 288 ↛ 289line 288 didn't jump to line 289 because the condition on line 288 was never true
289 print(f" Setting weight: {key} (shape: {weight_tensor.shape})")
290 # break tying by creating a new param
291 new_param = nn.Parameter(weight_tensor)
292 setattr(self.original_component, key, new_param)
293 elif param is None: 293 ↛ 276line 293 didn't jump to line 276 because the condition on line 293 was always true
294 # Parameter exists but is None (e.g., bias=False in nn.Linear)
295 # Create a new parameter from the weight tensor
296 if verbose: 296 ↛ 297line 296 didn't jump to line 297 because the condition on line 296 was never true
297 print(
298 f" Creating weight: {key} (shape: {weight_tensor.shape}) - was None"
299 )
300 new_param = nn.Parameter(weight_tensor)
301 setattr(self.original_component, key, new_param)
303 # If this component has submodules, distribute weights to them
304 if self.real_components:
305 from transformer_lens.weight_processing import ProcessWeights
307 if verbose: 307 ↛ 308line 307 didn't jump to line 308 because the condition on line 307 was never true
308 print(f" Has {len(self.real_components)} subcomponents, distributing weights...")
310 ProcessWeights.distribute_weights_to_components(
311 state_dict=weights,
312 component_mapping=self.real_components,
313 verbose=verbose,
314 )
316 def forward(self, *args: Any, **kwargs: Any) -> Any:
317 """Generic forward pass for bridge components with input/output hooks."""
318 original_component = self._modules.get("_original_component", None)
319 if original_component is None: 319 ↛ 320line 319 didn't jump to line 320 because the condition on line 319 was never true
320 raise RuntimeError(
321 f"Original component not set for {self.name}. Call set_original_component() first."
322 )
323 input_arg_names = [
324 "input",
325 "hidden_states",
326 "input_ids",
327 "query_input",
328 "x",
329 "inputs_embeds",
330 ]
331 input_found = False
332 for name in input_arg_names:
333 if name in kwargs:
334 hooked = self.hook_in(kwargs[name])
335 kwargs[name] = hooked
336 input_found = True
337 break
338 if not input_found and len(args) > 0 and isinstance(args[0], torch.Tensor):
339 hooked_input = self.hook_in(args[0])
340 args = (hooked_input,) + args[1:]
341 input_found = True
342 output = original_component(*args, **kwargs)
343 if isinstance(output, tuple):
344 hooked_first = self.hook_out(output[0])
345 output = (hooked_first,) + output[1:]
346 elif not isinstance(output, torch.Tensor) and isinstance(
347 getattr(output, "last_hidden_state", None), torch.Tensor
348 ):
349 # ModelOutput-returning components (e.g. vision/audio towers).
350 output.last_hidden_state = self.hook_out(output.last_hidden_state)
351 else:
352 output = self.hook_out(output)
353 return output
355 def __getattr__(self, name: str) -> Any:
356 modules = object.__getattribute__(self, "__dict__").get("_modules")
357 if modules is not None and name in modules:
358 return modules[name]
359 if name == "original_component": 359 ↛ 360line 359 didn't jump to line 360 because the condition on line 359 was never true
360 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
361 submodules = object.__getattribute__(self, "__dict__").get("submodules")
362 if submodules is not None and name in submodules:
363 # Don't return submodule here - it should be accessed via _modules after add_module()
364 # Raising AttributeError allows PyTorch's add_module() to work correctly
365 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
366 if modules is not None: 366 ↛ 380line 366 didn't jump to line 380 because the condition on line 366 was always true
367 original_component = modules.get("_original_component")
368 if original_component is not None:
369 try:
370 if "." in name: 370 ↛ 371line 370 didn't jump to line 371 because the condition on line 370 was never true
371 name_split = name.split(".")
372 current = getattr(original_component, name_split[0])
373 for part in name_split[1:]:
374 current = getattr(current, part)
375 return current
376 else:
377 return getattr(original_component, name)
378 except AttributeError:
379 pass
380 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
382 def __setattr__(self, name: str, value: Any) -> None:
383 """Set attribute, with passthrough to original component for compatibility."""
384 if isinstance(value, HookPoint):
385 self._register_hook(name, value)
386 super().__setattr__(name, value)
387 return
388 if name == "hook_aliases":
389 # Any alias rebind invalidates alias caches downstream.
390 _bump_alias_generation()
391 if name.startswith("_") or name in [
392 "name",
393 "config",
394 "submodules",
395 "conversion_rule",
396 "compatibility_mode",
397 "disable_warnings",
398 "optional",
399 # Components rebind these per layer at bind time (MoEBridge's
400 # dense/sparse dispatch). Without the carve-out the assignment is
401 # forwarded to the wrapped HF module whenever it happens to expose
402 # the attribute — the rebind then silently vanishes.
403 "hook_aliases",
404 "property_aliases",
405 # train()/eval() set self.training; redirecting it to the original
406 # component leaves the wrapper stuck in training mode (dropout at
407 # inference). Recursion still reaches the original via _modules.
408 "training",
409 ]:
410 super().__setattr__(name, value)
411 return
412 class_attr = getattr(type(self), name, None)
413 if class_attr is not None and isinstance(class_attr, property):
414 if class_attr.fset is not None:
415 super().__setattr__(name, value)
416 return
417 if hasattr(self, "_modules") and "_original_component" in self._modules:
418 original_component = self._modules["_original_component"]
419 if hasattr(original_component, name):
420 try:
421 setattr(original_component, name, value)
422 return
423 except AttributeError:
424 pass
425 super().__setattr__(name, value)