Coverage for transformer_lens/model_bridge/generalized_components/base.py: 85%
233 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Base class for generalized transformer components."""
2from __future__ import annotations
4import contextlib
5import inspect
6import warnings
7from collections.abc import Callable
8from typing import Any, Dict, List, Optional, Union
10import torch
11import torch.nn as nn
12from accelerate.utils import align_module_device
14from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import (
15 BaseTensorConversion,
16)
17from transformer_lens.hook_points import HookPoint
20def align_offloaded_subtree(module: nn.Module) -> contextlib.ExitStack:
21 """Materialize every Accelerate-offloaded descendant of ``module`` for the
22 caller's duration (an ``ExitStack`` of ``align_module_device`` contexts).
24 Accelerate attaches offload hooks at leaf level - whichever submodule
25 directly owns the Parameter (e.g. ``c_attn``, ``c_proj``) - not on
26 container modules like an attention block as a whole. ``align_module_device``
27 on a container alone is therefore a no-op even though its descendants are
28 offloaded. Walking every descendant and entering each one's
29 ``align_module_device`` (a cheap no-op for any module that has no hook of
30 its own) covers both a leaf ``original_component`` and a multi-level
31 container uniformly, without needing to know in advance which specific
32 descendant a given architecture adapter's code actually reads from.
33 """
34 stack = contextlib.ExitStack()
35 for submodule in module.modules():
36 stack.enter_context(align_module_device(submodule))
37 return stack
40class CloneOutputUnderGradMixin(nn.Module):
41 """Clone the forward output so HF's in-place mutation cannot corrupt it.
43 Under grad, autograd forbids in-place writes to backward-hook views; under
44 no_grad, cached hook_out tensors alias the storage HF then rewrites.
45 Mix in ahead of a bridge class: ``class X(CloneOutputUnderGradMixin, LinearBridge)``.
46 """
48 def forward(self, *args: Any, **kwargs: Any) -> Any:
49 out = super().forward(*args, **kwargs)
50 if isinstance(out, torch.Tensor): 50 ↛ 52line 50 didn't jump to line 52 because the condition on line 50 was always true
51 out = out.clone()
52 return out
55# Bumped whenever any component rebinds its hook_aliases (MoEBridge's dense/
56# sparse dispatch, OlmoHybrid's per-layer selection). Alias caches key on it:
57# a rebind can leave the hook REGISTRY unchanged while changing what the
58# aliases point at, so size-based cache keys cannot see it.
59_ALIAS_GENERATION = 0
62def alias_generation() -> int:
63 """Current global alias-rebind generation."""
64 return _ALIAS_GENERATION
67def _bump_alias_generation() -> None:
68 global _ALIAS_GENERATION
69 _ALIAS_GENERATION += 1
72class GeneralizedComponent(nn.Module):
73 """Base class for generalized transformer components.
75 This class provides a standardized interface for transformer components
76 and handles hook registration and execution.
77 """
79 is_list_item: bool = False
80 hook_out_is_single_residual_stream: bool = False
81 compatibility_mode: bool = False
82 disable_warnings: bool = False
83 hook_aliases: Dict[str, Union[str, List[str]]] = {}
85 # Projection-hook protocol between container bridges (MLPBridge) and the
86 # projection bridges they wrap (LinearBridge / Conv1DBridge): the wrapper
87 # records that its hook_out fired so the container only re-fires it as a
88 # bypass fallback, and the container suppresses the wrapper's next hook_in
89 # after pre-firing it itself. See MLPBridge.forward.
90 _fired_hook_out: bool = False
91 _suppress_next_hook_in: bool = False
92 property_aliases: Dict[str, str] = {}
94 def __init__(
95 self,
96 name: Optional[str],
97 config: Optional[Any] = None,
98 submodules: Optional[Dict[str, "GeneralizedComponent"]] = None,
99 conversion_rule: Optional[BaseTensorConversion] = None,
100 hook_alias_overrides: Optional[Dict[str, str]] = None,
101 optional: bool = False,
102 ):
103 """Initialize the generalized component.
105 Args:
106 name: The name of this component (None if component has no container in remote model)
107 config: Optional configuration object for the component
108 submodules: Dictionary of GeneralizedComponent submodules to register
109 conversion_rule: Optional conversion rule for this component's hooks
110 hook_alias_overrides: Optional dictionary to override default hook aliases.
111 For example, {"hook_attn_out": "ln1_post.hook_out"} will make hook_attn_out
112 point to ln1_post.hook_out instead of the default value in self.hook_aliases.
113 optional: If True, setup skips this subtree when absent (hybrid architectures).
114 """
115 super().__init__()
116 self.name = name
117 self.config = config
118 self.submodules = submodules or {}
119 self.conversion_rule = conversion_rule
120 self.optional = optional
121 self._hook_registry: Dict[str, HookPoint] = {}
122 self._hook_alias_registry: Dict[str, Union[str, List[str]]] = {}
123 self._property_alias_registry: Dict[str, str] = {}
124 self.hook_in = HookPoint()
125 self.hook_out = HookPoint()
126 # real_components maps TL keys to (remote_path, actual_instance) tuples
127 # For list components, actual_instance will be a list of component instances
128 self.real_components: Dict[str, tuple] = {}
129 if self.conversion_rule is not None:
130 self.hook_in.hook_conversion = self.conversion_rule
131 self.hook_out.hook_conversion = self.conversion_rule
133 # Copy class-level hook_aliases and apply any overrides
134 if hook_alias_overrides is not None:
135 self.hook_aliases = self.__class__.hook_aliases.copy()
136 self.hook_aliases.update(hook_alias_overrides)
138 def __call__(self, *args: Any, **kwargs: Any) -> Any:
139 """Run forward(), materializing the wrapped component's params first if offloaded.
141 Bridge components read the wrapped module's raw parameters directly
142 (self.weight / self.original_component.bias / etc. via __getattr__ or
143 direct attribute access) rather than exclusively calling the wrapped
144 module's own forward(). Under an Accelerate CPU/disk device_map,
145 Accelerate only swaps a meta placeholder for the real, materialized
146 tensor around the wrapped module's OWN forward() call (its pre/post
147 forward hooks) - a raw attribute read outside that window silently
148 sees the meta placeholder instead, with no error. align_module_device
149 wraps this call in the same pre/post-forward hook Accelerate would
150 have run, so every read during this call - whichever way the code
151 reaches it - gets real data. It's a no-op (bare yield) when the
152 wrapped module has no offload hook, which covers the common case of
153 no device_map, a single device, or a multi-GPU split with no CPU/disk
154 offload involved.
156 Deliberately just original_component, not align_offloaded_subtree's
157 whole-subtree walk: Accelerate hooks the leaf modules that actually own
158 Parameters (e.g. NormalizationBridge/LinearBridge wrap one directly), so
159 materializing exactly that leaf for exactly its own call keeps the same
160 one-leaf-at-a-time memory footprint Accelerate's native per-module hooks
161 would give a plain (non-bridge) forward pass. A component whose
162 original_component is a container of several separately-hooked leaves
163 (e.g. an attention module wrapping distinct q/k/v/o projections) doesn't
164 need this to also materialize here - each of ITS own sub-bridges calls
165 into its own leaf the same way. align_offloaded_subtree is for the
166 narrower case of code that reads a specific descendant's raw params
167 directly during setup, without going through that descendant's own
168 bridge __call__ at all (see JointQKVAttentionBridge.set_original_component).
169 """
170 original_component = self._modules.get("_original_component")
171 if original_component is None:
172 return super().__call__(*args, **kwargs)
173 with align_module_device(original_component):
174 return super().__call__(*args, **kwargs)
176 def _register_hook(self, name: str, hook: HookPoint) -> None:
177 """Register a hook in the component's hook registry."""
178 hook.name = name
179 self._hook_registry[name] = hook
181 def _register_aliases(self) -> None:
182 """Register aliases from class-level dictionaries.
184 Called unconditionally at bridge init (see bridge.py); compatibility mode
185 additionally re-registers after weight processing.
186 It creates actual Python attributes/properties that directly reference the target objects.
188 Note: Re-registration expects to run after
189 weight processing is complete to ensure property aliases point to processed weights.
190 """
191 if self.hook_aliases:
192 self._hook_alias_registry.update(self.hook_aliases)
193 if self.property_aliases:
194 self._property_alias_registry.update(self.property_aliases)
195 for alias_name, target_path in self._hook_alias_registry.items():
196 resolved = False
197 if isinstance(target_path, list): 197 ↛ 198line 197 didn't jump to line 198 because the condition on line 197 was never true
198 for single_target in target_path:
199 try:
200 target_obj = self
201 for part in single_target.split("."):
202 target_obj = getattr(target_obj, part)
203 object.__setattr__(self, alias_name, target_obj)
204 resolved = True
205 break
206 except AttributeError:
207 continue
208 else:
209 try:
210 target_obj = self
211 for part in target_path.split("."):
212 target_obj = getattr(target_obj, part)
213 object.__setattr__(self, alias_name, target_obj)
214 resolved = True
215 except AttributeError:
216 pass
217 if not resolved:
218 # Surface drops instead of silently swallowing — some aliases are
219 # legitimately conditional on optional submodules, but an author
220 # needs to see which ones dropped at bridge-init.
221 warnings.warn(
222 f"Hook alias '{alias_name}' -> '{target_path}' on "
223 f"{type(self).__name__}(name={getattr(self, 'name', None)!r}) "
224 f"did not resolve; this hook will not be accessible.",
225 stacklevel=2,
226 )
227 for alias_name, target_path in self._property_alias_registry.items():
228 try:
229 target_obj = self
230 for part in target_path.split("."):
231 target_obj = getattr(target_obj, part)
232 object.__setattr__(self, alias_name, target_obj)
233 except AttributeError:
234 pass
236 def get_hooks(self) -> Dict[str, HookPoint]:
237 """Get all hooks registered in this component."""
238 hooks = self._hook_registry.copy()
239 if self.compatibility_mode and self._hook_alias_registry:
240 for alias_name in self._hook_alias_registry.keys():
241 if hasattr(self, alias_name): 241 ↛ 240line 241 didn't jump to line 240 because the condition on line 241 was always true
242 target_hook = getattr(self, alias_name)
243 if isinstance(target_hook, HookPoint): 243 ↛ 240line 243 didn't jump to line 240 because the condition on line 243 was always true
244 hooks[alias_name] = target_hook
245 return hooks
247 def _is_getattr_called_internally(self) -> bool:
248 """This function checks if the __getattr__ method was being called internally
249 (e.g by the setup process or run_with_cache).
250 """
251 for frame_info in inspect.stack():
252 if "setup_components" in frame_info.function or "run_with_cache" in frame_info.function:
253 return True
254 return False
256 def set_original_component(self, original_component: nn.Module) -> None:
257 """Set the original component that this bridge wraps.
259 Args:
260 original_component: The original transformer component to wrap
261 """
262 self.add_module("_original_component", original_component)
263 # An opaque wrapper (created with config=None) shadows the wrapped
264 # module's own config. HF forwards sometimes read a submodule's config
265 # directly (e.g. Qwen2Audio's forward does create_bidirectional_mask(
266 # config=self.audio_tower.config)), so inherit the real config to avoid
267 # exposing None. Components given an explicit config keep it.
268 if self.config is None:
269 self.config = getattr(original_component, "config", None)
271 @property
272 def original_component(self) -> Optional[nn.Module]:
273 """Get the original component."""
274 return self._modules.get("_original_component", None)
276 def add_hook(self, hook_fn: Callable[..., torch.Tensor], hook_name: str = "output") -> None:
277 """Add a hook function (HookedTransformer-compatible interface).
279 Args:
280 hook_fn: Function to call at this hook point
281 hook_name: Name of the hook point (defaults to "output")
282 """
283 if hook_name == "output":
284 self.hook_out.add_hook(hook_fn)
285 elif hook_name == "input":
286 self.hook_in.add_hook(hook_fn)
287 else:
288 raise ValueError(
289 f"Hook name '{hook_name}' not supported. Supported names are 'output' and 'input'."
290 )
292 def remove_hooks(self, hook_name: str | None = None) -> None:
293 """Remove hooks (HookedTransformer-compatible interface).
295 Args:
296 hook_name: Name of the hook point to remove. If None, removes all hooks.
297 """
298 if hook_name is None:
299 self.hook_in.remove_hooks(dir="both")
300 self.hook_out.remove_hooks(dir="both")
301 elif hook_name == "output":
302 self.hook_out.remove_hooks(dir="both")
303 elif hook_name == "input":
304 self.hook_in.remove_hooks(dir="both")
305 else:
306 raise ValueError(
307 f"Hook name '{hook_name}' not supported. Supported names are 'output' and 'input'."
308 )
310 def set_processed_weights(
311 self, weights: Dict[str, torch.Tensor], verbose: bool = False
312 ) -> None:
313 """Set the processed weights for use in compatibility mode.
315 This method stores processed weights as attributes on the component so they can be
316 used directly in the forward pass without modifying the original component.
318 Components should override this method to handle their specific weight structure.
319 The weights dict contains keys like "weight", "bias", "W_in", "W_out", etc.
321 If this component has submodules, this method will automatically distribute the
322 weights to those subcomponents using ProcessWeights.distribute_weights_to_components.
324 Args:
325 weights: Dictionary of processed weight tensors
326 verbose: If True, print detailed information about weight setting
327 """
328 if verbose: 328 ↛ 329line 328 didn't jump to line 329 because the condition on line 328 was never true
329 print(
330 f"\n set_processed_weights: {self.__class__.__name__} (name={getattr(self, 'name', 'unknown')})"
331 )
332 print(f" Received {len(weights)} weight keys")
334 # First, handle single-part keys (keys without ".") by setting them as parameters
335 # on the original component
336 if self.original_component is not None: 336 ↛ 365line 336 didn't jump to line 365 because the condition on line 336 was always true
337 for key, weight_tensor in weights.items():
338 # Only process keys without "." (single-part keys)
339 if "." not in key:
340 # Try to set the parameter on the original component
341 if hasattr(self.original_component, key):
342 param = getattr(self.original_component, key)
343 if param is not None and isinstance(param, torch.nn.Parameter):
344 if param.shape != weight_tensor.shape: 344 ↛ 345line 344 didn't jump to line 345 because the condition on line 344 was never true
345 raise ValueError(
346 f"Shape mismatch when setting weight '{key}' in {type(self.original_component).__name__}: "
347 f"existing param shape {param.shape} != new tensor shape {weight_tensor.shape}"
348 )
349 if verbose: 349 ↛ 350line 349 didn't jump to line 350 because the condition on line 349 was never true
350 print(f" Setting weight: {key} (shape: {weight_tensor.shape})")
351 # break tying by creating a new param
352 new_param = nn.Parameter(weight_tensor)
353 setattr(self.original_component, key, new_param)
354 elif param is None: 354 ↛ 337line 354 didn't jump to line 337 because the condition on line 354 was always true
355 # Parameter exists but is None (e.g., bias=False in nn.Linear)
356 # Create a new parameter from the weight tensor
357 if verbose: 357 ↛ 358line 357 didn't jump to line 358 because the condition on line 357 was never true
358 print(
359 f" Creating weight: {key} (shape: {weight_tensor.shape}) - was None"
360 )
361 new_param = nn.Parameter(weight_tensor)
362 setattr(self.original_component, key, new_param)
364 # If this component has submodules, distribute weights to them
365 if self.real_components:
366 from transformer_lens.weight_processing import ProcessWeights
368 if verbose: 368 ↛ 369line 368 didn't jump to line 369 because the condition on line 368 was never true
369 print(f" Has {len(self.real_components)} subcomponents, distributing weights...")
371 ProcessWeights.distribute_weights_to_components(
372 state_dict=weights,
373 component_mapping=self.real_components,
374 verbose=verbose,
375 )
377 def forward(self, *args: Any, **kwargs: Any) -> Any:
378 """Generic forward pass for bridge components with input/output hooks."""
379 original_component = self._modules.get("_original_component", None)
380 if original_component is None: 380 ↛ 381line 380 didn't jump to line 381 because the condition on line 380 was never true
381 raise RuntimeError(
382 f"Original component not set for {self.name}. Call set_original_component() first."
383 )
384 input_arg_names = [
385 "input",
386 "hidden_states",
387 "input_ids",
388 "query_input",
389 "x",
390 "inputs_embeds",
391 ]
392 input_found = False
393 for name in input_arg_names:
394 if name in kwargs:
395 hooked = self.hook_in(kwargs[name])
396 kwargs[name] = hooked
397 input_found = True
398 break
399 if not input_found and len(args) > 0 and isinstance(args[0], torch.Tensor):
400 hooked_input = self.hook_in(args[0])
401 args = (hooked_input,) + args[1:]
402 input_found = True
403 output = original_component(*args, **kwargs)
404 if isinstance(output, tuple):
405 hooked_first = self.hook_out(output[0])
406 output = (hooked_first,) + output[1:]
407 elif not isinstance(output, torch.Tensor) and isinstance(
408 getattr(output, "last_hidden_state", None), torch.Tensor
409 ):
410 # ModelOutput-returning components (e.g. vision/audio towers).
411 output.last_hidden_state = self.hook_out(output.last_hidden_state)
412 else:
413 output = self.hook_out(output)
414 return output
416 def __getattr__(self, name: str) -> Any:
417 modules = object.__getattribute__(self, "__dict__").get("_modules")
418 if modules is not None and name in modules:
419 return modules[name]
420 if name == "original_component": 420 ↛ 421line 420 didn't jump to line 421 because the condition on line 420 was never true
421 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
422 submodules = object.__getattribute__(self, "__dict__").get("submodules")
423 if submodules is not None and name in submodules:
424 # Don't return submodule here - it should be accessed via _modules after add_module()
425 # Raising AttributeError allows PyTorch's add_module() to work correctly
426 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
427 if modules is not None: 427 ↛ 441line 427 didn't jump to line 441 because the condition on line 427 was always true
428 original_component = modules.get("_original_component")
429 if original_component is not None:
430 try:
431 if "." in name: 431 ↛ 432line 431 didn't jump to line 432 because the condition on line 431 was never true
432 name_split = name.split(".")
433 current = getattr(original_component, name_split[0])
434 for part in name_split[1:]:
435 current = getattr(current, part)
436 return current
437 else:
438 return getattr(original_component, name)
439 except AttributeError:
440 pass
441 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
443 def __setattr__(self, name: str, value: Any) -> None:
444 """Set attribute, with passthrough to original component for compatibility."""
445 if isinstance(value, HookPoint):
446 self._register_hook(name, value)
447 super().__setattr__(name, value)
448 return
449 if name == "hook_aliases":
450 # Any alias rebind invalidates alias caches downstream.
451 _bump_alias_generation()
452 if name.startswith("_") or name in [
453 "name",
454 "config",
455 "submodules",
456 "conversion_rule",
457 "compatibility_mode",
458 "disable_warnings",
459 "optional",
460 # Components rebind these per layer at bind time (MoEBridge's
461 # dense/sparse dispatch). Without the carve-out the assignment is
462 # forwarded to the wrapped HF module whenever it happens to expose
463 # the attribute — the rebind then silently vanishes.
464 "hook_aliases",
465 "property_aliases",
466 # train()/eval() set self.training; redirecting it to the original
467 # component leaves the wrapper stuck in training mode (dropout at
468 # inference). Recursion still reaches the original via _modules.
469 "training",
470 ]:
471 super().__setattr__(name, value)
472 return
473 class_attr = getattr(type(self), name, None)
474 if class_attr is not None and isinstance(class_attr, property):
475 if class_attr.fset is not None:
476 super().__setattr__(name, value)
477 return
478 if hasattr(self, "_modules") and "_original_component" in self._modules:
479 original_component = self._modules["_original_component"]
480 if hasattr(original_component, name):
481 try:
482 setattr(original_component, name, value)
483 return
484 except AttributeError:
485 pass
486 super().__setattr__(name, value)