Coverage for transformer_lens/model_bridge/generalized_components/base.py: 85%
237 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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 mark_pos_dim(self, pos_dim: int) -> None:
139 """Declare the position axis of every HookPoint this component owns.
141 Call it on a component whose activations do not carry position on dim 1
142 of the batched layout (e.g. a norm applied to ``[batch, heads, pos, d_head]``)
143 so ``pos_slice`` slices positions on all of its hooks, including ones a
144 subclass adds (``hook_normalized`` / ``hook_scale`` on a norm). Submodule
145 components are not touched; they declare their own layout.
146 """
147 for child in self.children():
148 if isinstance(child, HookPoint): 148 ↛ 147line 148 didn't jump to line 147 because the condition on line 148 was always true
149 child.pos_dim = pos_dim
151 def __call__(self, *args: Any, **kwargs: Any) -> Any:
152 """Run forward(), materializing the wrapped component's params first if offloaded.
154 Bridge components read the wrapped module's raw parameters directly
155 (self.weight / self.original_component.bias / etc. via __getattr__ or
156 direct attribute access) rather than exclusively calling the wrapped
157 module's own forward(). Under an Accelerate CPU/disk device_map,
158 Accelerate only swaps a meta placeholder for the real, materialized
159 tensor around the wrapped module's OWN forward() call (its pre/post
160 forward hooks) - a raw attribute read outside that window silently
161 sees the meta placeholder instead, with no error. align_module_device
162 wraps this call in the same pre/post-forward hook Accelerate would
163 have run, so every read during this call - whichever way the code
164 reaches it - gets real data. It's a no-op (bare yield) when the
165 wrapped module has no offload hook, which covers the common case of
166 no device_map, a single device, or a multi-GPU split with no CPU/disk
167 offload involved.
169 Deliberately just original_component, not align_offloaded_subtree's
170 whole-subtree walk: Accelerate hooks the leaf modules that actually own
171 Parameters (e.g. NormalizationBridge/LinearBridge wrap one directly), so
172 materializing exactly that leaf for exactly its own call keeps the same
173 one-leaf-at-a-time memory footprint Accelerate's native per-module hooks
174 would give a plain (non-bridge) forward pass. A component whose
175 original_component is a container of several separately-hooked leaves
176 (e.g. an attention module wrapping distinct q/k/v/o projections) doesn't
177 need this to also materialize here - each of ITS own sub-bridges calls
178 into its own leaf the same way. align_offloaded_subtree is for the
179 narrower case of code that reads a specific descendant's raw params
180 directly during setup, without going through that descendant's own
181 bridge __call__ at all (see JointQKVAttentionBridge.set_original_component).
182 """
183 original_component = self._modules.get("_original_component")
184 if original_component is None:
185 return super().__call__(*args, **kwargs)
186 with align_module_device(original_component):
187 return super().__call__(*args, **kwargs)
189 def _register_hook(self, name: str, hook: HookPoint) -> None:
190 """Register a hook in the component's hook registry."""
191 hook.name = name
192 self._hook_registry[name] = hook
194 def _register_aliases(self) -> None:
195 """Register aliases from class-level dictionaries.
197 Called unconditionally at bridge init (see bridge.py); compatibility mode
198 additionally re-registers after weight processing.
199 It creates actual Python attributes/properties that directly reference the target objects.
201 Note: Re-registration expects to run after
202 weight processing is complete to ensure property aliases point to processed weights.
203 """
204 if self.hook_aliases:
205 self._hook_alias_registry.update(self.hook_aliases)
206 if self.property_aliases:
207 self._property_alias_registry.update(self.property_aliases)
208 for alias_name, target_path in self._hook_alias_registry.items():
209 resolved = False
210 if isinstance(target_path, list): 210 ↛ 211line 210 didn't jump to line 211 because the condition on line 210 was never true
211 for single_target in target_path:
212 try:
213 target_obj = self
214 for part in single_target.split("."):
215 target_obj = getattr(target_obj, part)
216 object.__setattr__(self, alias_name, target_obj)
217 resolved = True
218 break
219 except AttributeError:
220 continue
221 else:
222 try:
223 target_obj = self
224 for part in target_path.split("."):
225 target_obj = getattr(target_obj, part)
226 object.__setattr__(self, alias_name, target_obj)
227 resolved = True
228 except AttributeError:
229 pass
230 if not resolved:
231 # Surface drops instead of silently swallowing — some aliases are
232 # legitimately conditional on optional submodules, but an author
233 # needs to see which ones dropped at bridge-init.
234 warnings.warn(
235 f"Hook alias '{alias_name}' -> '{target_path}' on "
236 f"{type(self).__name__}(name={getattr(self, 'name', None)!r}) "
237 f"did not resolve; this hook will not be accessible.",
238 stacklevel=2,
239 )
240 for alias_name, target_path in self._property_alias_registry.items():
241 try:
242 target_obj = self
243 for part in target_path.split("."):
244 target_obj = getattr(target_obj, part)
245 object.__setattr__(self, alias_name, target_obj)
246 except AttributeError:
247 pass
249 def get_hooks(self) -> Dict[str, HookPoint]:
250 """Get all hooks registered in this component."""
251 hooks = self._hook_registry.copy()
252 if self.compatibility_mode and self._hook_alias_registry:
253 for alias_name in self._hook_alias_registry.keys():
254 if hasattr(self, alias_name): 254 ↛ 253line 254 didn't jump to line 253 because the condition on line 254 was always true
255 target_hook = getattr(self, alias_name)
256 if isinstance(target_hook, HookPoint): 256 ↛ 253line 256 didn't jump to line 253 because the condition on line 256 was always true
257 hooks[alias_name] = target_hook
258 return hooks
260 def _is_getattr_called_internally(self) -> bool:
261 """This function checks if the __getattr__ method was being called internally
262 (e.g by the setup process or run_with_cache).
263 """
264 for frame_info in inspect.stack():
265 if "setup_components" in frame_info.function or "run_with_cache" in frame_info.function:
266 return True
267 return False
269 def set_original_component(self, original_component: nn.Module) -> None:
270 """Set the original component that this bridge wraps.
272 Args:
273 original_component: The original transformer component to wrap
274 """
275 self.add_module("_original_component", original_component)
276 # An opaque wrapper (created with config=None) shadows the wrapped
277 # module's own config. HF forwards sometimes read a submodule's config
278 # directly (e.g. Qwen2Audio's forward does create_bidirectional_mask(
279 # config=self.audio_tower.config)), so inherit the real config to avoid
280 # exposing None. Components given an explicit config keep it.
281 if self.config is None:
282 self.config = getattr(original_component, "config", None)
284 @property
285 def original_component(self) -> Optional[nn.Module]:
286 """Get the original component."""
287 return self._modules.get("_original_component", None)
289 def add_hook(self, hook_fn: Callable[..., torch.Tensor], hook_name: str = "output") -> None:
290 """Add a hook function (HookedTransformer-compatible interface).
292 Args:
293 hook_fn: Function to call at this hook point
294 hook_name: Name of the hook point (defaults to "output")
295 """
296 if hook_name == "output":
297 self.hook_out.add_hook(hook_fn)
298 elif hook_name == "input":
299 self.hook_in.add_hook(hook_fn)
300 else:
301 raise ValueError(
302 f"Hook name '{hook_name}' not supported. Supported names are 'output' and 'input'."
303 )
305 def remove_hooks(self, hook_name: str | None = None) -> None:
306 """Remove hooks (HookedTransformer-compatible interface).
308 Args:
309 hook_name: Name of the hook point to remove. If None, removes all hooks.
310 """
311 if hook_name is None:
312 self.hook_in.remove_hooks(dir="both")
313 self.hook_out.remove_hooks(dir="both")
314 elif hook_name == "output":
315 self.hook_out.remove_hooks(dir="both")
316 elif hook_name == "input":
317 self.hook_in.remove_hooks(dir="both")
318 else:
319 raise ValueError(
320 f"Hook name '{hook_name}' not supported. Supported names are 'output' and 'input'."
321 )
323 def set_processed_weights(
324 self, weights: Dict[str, torch.Tensor], verbose: bool = False
325 ) -> None:
326 """Set the processed weights for use in compatibility mode.
328 This method stores processed weights as attributes on the component so they can be
329 used directly in the forward pass without modifying the original component.
331 Components should override this method to handle their specific weight structure.
332 The weights dict contains keys like "weight", "bias", "W_in", "W_out", etc.
334 If this component has submodules, this method will automatically distribute the
335 weights to those subcomponents using ProcessWeights.distribute_weights_to_components.
337 Args:
338 weights: Dictionary of processed weight tensors
339 verbose: If True, print detailed information about weight setting
340 """
341 if verbose: 341 ↛ 342line 341 didn't jump to line 342 because the condition on line 341 was never true
342 print(
343 f"\n set_processed_weights: {self.__class__.__name__} (name={getattr(self, 'name', 'unknown')})"
344 )
345 print(f" Received {len(weights)} weight keys")
347 # First, handle single-part keys (keys without ".") by setting them as parameters
348 # on the original component
349 if self.original_component is not None: 349 ↛ 378line 349 didn't jump to line 378 because the condition on line 349 was always true
350 for key, weight_tensor in weights.items():
351 # Only process keys without "." (single-part keys)
352 if "." not in key:
353 # Try to set the parameter on the original component
354 if hasattr(self.original_component, key):
355 param = getattr(self.original_component, key)
356 if param is not None and isinstance(param, torch.nn.Parameter):
357 if param.shape != weight_tensor.shape: 357 ↛ 358line 357 didn't jump to line 358 because the condition on line 357 was never true
358 raise ValueError(
359 f"Shape mismatch when setting weight '{key}' in {type(self.original_component).__name__}: "
360 f"existing param shape {param.shape} != new tensor shape {weight_tensor.shape}"
361 )
362 if verbose: 362 ↛ 363line 362 didn't jump to line 363 because the condition on line 362 was never true
363 print(f" Setting weight: {key} (shape: {weight_tensor.shape})")
364 # break tying by creating a new param
365 new_param = nn.Parameter(weight_tensor)
366 setattr(self.original_component, key, new_param)
367 elif param is None:
368 # Parameter exists but is None (e.g., bias=False in nn.Linear)
369 # Create a new parameter from the weight tensor
370 if verbose: 370 ↛ 371line 370 didn't jump to line 371 because the condition on line 370 was never true
371 print(
372 f" Creating weight: {key} (shape: {weight_tensor.shape}) - was None"
373 )
374 new_param = nn.Parameter(weight_tensor)
375 setattr(self.original_component, key, new_param)
377 # If this component has submodules, distribute weights to them
378 if self.real_components:
379 from transformer_lens.weight_processing import ProcessWeights
381 if verbose: 381 ↛ 382line 381 didn't jump to line 382 because the condition on line 381 was never true
382 print(f" Has {len(self.real_components)} subcomponents, distributing weights...")
384 ProcessWeights.distribute_weights_to_components(
385 state_dict=weights,
386 component_mapping=self.real_components,
387 verbose=verbose,
388 )
390 def forward(self, *args: Any, **kwargs: Any) -> Any:
391 """Generic forward pass for bridge components with input/output hooks."""
392 original_component = self._modules.get("_original_component", None)
393 if original_component is None: 393 ↛ 394line 393 didn't jump to line 394 because the condition on line 393 was never true
394 raise RuntimeError(
395 f"Original component not set for {self.name}. Call set_original_component() first."
396 )
397 input_arg_names = [
398 "input",
399 "hidden_states",
400 "input_ids",
401 "query_input",
402 "x",
403 "inputs_embeds",
404 ]
405 input_found = False
406 for name in input_arg_names:
407 if name in kwargs:
408 hooked = self.hook_in(kwargs[name])
409 kwargs[name] = hooked
410 input_found = True
411 break
412 if not input_found and len(args) > 0 and isinstance(args[0], torch.Tensor):
413 hooked_input = self.hook_in(args[0])
414 args = (hooked_input,) + args[1:]
415 input_found = True
416 output = original_component(*args, **kwargs)
417 if isinstance(output, tuple):
418 hooked_first = self.hook_out(output[0])
419 output = (hooked_first,) + output[1:]
420 elif not isinstance(output, torch.Tensor) and isinstance(
421 getattr(output, "last_hidden_state", None), torch.Tensor
422 ):
423 # ModelOutput-returning components (e.g. vision/audio towers).
424 output.last_hidden_state = self.hook_out(output.last_hidden_state)
425 else:
426 output = self.hook_out(output)
427 return output
429 def __getattr__(self, name: str) -> Any:
430 modules = object.__getattribute__(self, "__dict__").get("_modules")
431 if modules is not None and name in modules:
432 return modules[name]
433 if name == "original_component": 433 ↛ 434line 433 didn't jump to line 434 because the condition on line 433 was never true
434 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
435 submodules = object.__getattribute__(self, "__dict__").get("submodules")
436 if submodules is not None and name in submodules:
437 # Don't return submodule here - it should be accessed via _modules after add_module()
438 # Raising AttributeError allows PyTorch's add_module() to work correctly
439 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
440 if modules is not None: 440 ↛ 454line 440 didn't jump to line 454 because the condition on line 440 was always true
441 original_component = modules.get("_original_component")
442 if original_component is not None:
443 try:
444 if "." in name: 444 ↛ 445line 444 didn't jump to line 445 because the condition on line 444 was never true
445 name_split = name.split(".")
446 current = getattr(original_component, name_split[0])
447 for part in name_split[1:]:
448 current = getattr(current, part)
449 return current
450 else:
451 return getattr(original_component, name)
452 except AttributeError:
453 pass
454 raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
456 def __setattr__(self, name: str, value: Any) -> None:
457 """Set attribute, with passthrough to original component for compatibility."""
458 if isinstance(value, HookPoint):
459 self._register_hook(name, value)
460 super().__setattr__(name, value)
461 return
462 if name == "hook_aliases":
463 # Any alias rebind invalidates alias caches downstream.
464 _bump_alias_generation()
465 if name.startswith("_") or name in [
466 "name",
467 "config",
468 "submodules",
469 "conversion_rule",
470 "compatibility_mode",
471 "disable_warnings",
472 "optional",
473 # Components rebind these per layer at bind time (MoEBridge's
474 # dense/sparse dispatch). Without the carve-out the assignment is
475 # forwarded to the wrapped HF module whenever it happens to expose
476 # the attribute — the rebind then silently vanishes.
477 "hook_aliases",
478 "property_aliases",
479 # train()/eval() set self.training; redirecting it to the original
480 # component leaves the wrapper stuck in training mode (dropout at
481 # inference). Recursion still reaches the original via _modules.
482 "training",
483 ]:
484 super().__setattr__(name, value)
485 return
486 class_attr = getattr(type(self), name, None)
487 if class_attr is not None and isinstance(class_attr, property):
488 if class_attr.fset is not None:
489 super().__setattr__(name, value)
490 return
491 if hasattr(self, "_modules") and "_original_component" in self._modules:
492 original_component = self._modules["_original_component"]
493 if hasattr(original_component, name):
494 try:
495 setattr(original_component, name, value)
496 return
497 except AttributeError:
498 pass
499 super().__setattr__(name, value)