transformer_lens.model_bridge.generalized_components.normalization module

Normalization bridge component implementation.

class transformer_lens.model_bridge.generalized_components.normalization.NormalizationBridge(name: str, config: Any, submodules: Dict[str, GeneralizedComponent] | None = {}, use_native_layernorm_autograd: bool = False, uses_rms_norm: bool | None = None)

Bases: GeneralizedComponent

Normalization bridge that wraps transformer normalization layers but implements the calculation from scratch.

This component provides standardized input/output hooks.

__init__(name: str, config: Any, submodules: Dict[str, GeneralizedComponent] | None = {}, use_native_layernorm_autograd: bool = False, uses_rms_norm: bool | None = None)

Initialize the normalization bridge.

Parameters:

  • name – The name of this component

  • config – Optional configuration

  • submodules – Dictionary of GeneralizedComponent submodules to register

  • use_native_layernorm_autograd – If True, use HuggingFace’s native LayerNorm autograd for exact gradient matching. If False, use custom implementation. Defaults to False.

  • uses_rms_norm – Force RMSNorm vs LayerNorm; None defers to introspection then config.uses_rms_norm.

forward(hidden_states: Tensor, **kwargs: Any) → Tensor

Forward pass through the normalization bridge.

Parameters:

  • hidden_states – Input hidden states

  • **kwargs – Additional arguments to pass to the original component

Returns:

Normalized output

property_aliases: Dict[str, str] = {'b': 'bias', 'w': 'weight'}
real_components: Dict[str, tuple]
training: bool
property uses_rms_norm: bool

Whether this bridge treats the wrapped module as RMSNorm.

Override > module introspection > config. Introspection guards against a shared config (RMSNorm LM + LayerNorm vision tower) misclassifying a real nn.LayerNorm.