transformer_lens.model_bridge.generalized_components.rms_normalization module

RMS Normalization bridge component implementation.

RMSNorm (Root Mean Square Layer Normalization) is used in models like T5, LLaMA, Mistral, etc. Unlike LayerNorm, RMSNorm doesn’t center the inputs (no mean subtraction) and has no bias.

class transformer_lens.model_bridge.generalized_components.rms_normalization.RMSNormalizationBridge(name: str, config: Any, submodules: Dict[str, 'GeneralizedComponent'] | None = None, use_native_layernorm_autograd: bool = True)

Bases: NormalizationBridge

RMS Normalization bridge for models that use RMSNorm (T5, LLaMA, etc).

RMSNorm differs from LayerNorm in two ways: 1. No mean centering (no subtraction of mean) 2. No bias term (only weight/scale parameter)

This bridge does a simple pass-through to the original HuggingFace component with hooks on input and output.

__init__(name: str, config: Any, submodules: Dict[str, 'GeneralizedComponent'] | None = None, use_native_layernorm_autograd: bool = True)

Initialize the RMS normalization bridge.

Parameters:

  • name – The name of this component

  • config – Configuration object

  • submodules – Dictionary of GeneralizedComponent submodules to register

  • use_native_layernorm_autograd – Use HF’s RMSNorm implementation for exact numerical match

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