transformer_lens.model_bridge.generalized_components.base module

Base class for generalized transformer components.

class transformer_lens.model_bridge.generalized_components.base.CloneOutputUnderGradMixin(*args: Any, **kwargs: Any)

Bases: Module

Clone the forward output so HF’s in-place mutation cannot corrupt it.

Under grad, autograd forbids in-place writes to backward-hook views; under no_grad, cached hook_out tensors alias the storage HF then rewrites. Mix in ahead of a bridge class: class X(CloneOutputUnderGradMixin, LinearBridge).

forward(*args: Any, **kwargs: Any) → Any

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class transformer_lens.model_bridge.generalized_components.base.GeneralizedComponent(name: str | None, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, conversion_rule: BaseTensorConversion | None = None, hook_alias_overrides: Dict[str, str] | None = None, optional: bool = False)

Bases: Module

Base class for generalized transformer components.

This class provides a standardized interface for transformer components and handles hook registration and execution.

__init__(name: str | None, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, conversion_rule: BaseTensorConversion | None = None, hook_alias_overrides: Dict[str, str] | None = None, optional: bool = False)

Initialize the generalized component.

Parameters:

  • name – The name of this component (None if component has no container in remote model)

  • config – Optional configuration object for the component

  • submodules – Dictionary of GeneralizedComponent submodules to register

  • conversion_rule – Optional conversion rule for this component’s hooks

  • hook_alias_overrides – Optional dictionary to override default hook aliases. For example, {“hook_attn_out”: “ln1_post.hook_out”} will make hook_attn_out point to ln1_post.hook_out instead of the default value in self.hook_aliases.

  • optional – If True, setup skips this subtree when absent (hybrid architectures).

add_hook(hook_fn: Callable[[...], Tensor], hook_name: str = 'output') → None

Add a hook function (HookedTransformer-compatible interface).

Parameters:

  • hook_fn – Function to call at this hook point

  • hook_name – Name of the hook point (defaults to “output”)

compatibility_mode: bool = False
disable_warnings: bool = False
forward(*args: Any, **kwargs: Any) → Any

Generic forward pass for bridge components with input/output hooks.

get_hooks() → Dict[str, HookPoint]

Get all hooks registered in this component.

hook_aliases: Dict[str, str | List[str]] = {}
hook_out_is_single_residual_stream: bool = False
is_list_item: bool = False
property original_component: Module | None

Get the original component.

property_aliases: Dict[str, str] = {}
remove_hooks(hook_name: str | None = None) → None

Remove hooks (HookedTransformer-compatible interface).

Parameters:

hook_name – Name of the hook point to remove. If None, removes all hooks.

set_original_component(original_component: Module) → None

Set the original component that this bridge wraps.

Parameters:

original_component – The original transformer component to wrap

set_processed_weights(weights: Dict[str, Tensor], verbose: bool = False) → None

Set the processed weights for use in compatibility mode.

This method stores processed weights as attributes on the component so they can be used directly in the forward pass without modifying the original component.

Components should override this method to handle their specific weight structure. The weights dict contains keys like “weight”, “bias”, “W_in”, “W_out”, etc.

If this component has submodules, this method will automatically distribute the weights to those subcomponents using ProcessWeights.distribute_weights_to_components.

Parameters:

  • weights – Dictionary of processed weight tensors

  • verbose – If True, print detailed information about weight setting

transformer_lens.model_bridge.generalized_components.base.alias_generation() → int

Current global alias-rebind generation.

transformer_lens.model_bridge.generalized_components.base.align_offloaded_subtree(module: Module) → ExitStack

Materialize every Accelerate-offloaded descendant of module for the caller’s duration (an ExitStack of align_module_device contexts).

Accelerate attaches offload hooks at leaf level - whichever submodule directly owns the Parameter (e.g. c_attn, c_proj) - not on container modules like an attention block as a whole. align_module_device on a container alone is therefore a no-op even though its descendants are offloaded. Walking every descendant and entering each one’s align_module_device (a cheap no-op for any module that has no hook of its own) covers both a leaf original_component and a multi-level container uniformly, without needing to know in advance which specific descendant a given architecture adapter’s code actually reads from.