Coverage for transformer_lens/model_bridge/generalized_components/normalization.py: 93%
118 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +0000
1"""Normalization bridge component implementation."""
2import contextlib
3import warnings
4from typing import Any, ContextManager, Dict, Optional, cast
6import torch
8from transformer_lens.hook_points import HookPoint
9from transformer_lens.model_bridge.generalized_components.base import (
10 GeneralizedComponent,
11)
13# The native-autograd path returns HF's own output, so hook edits and backward hooks
14# can only be honored by switching to the python-norm computation, whose numerics
15# differ from HF's at float-rounding scale.
16NATIVE_PATH_BWD_FALLBACK_WARNING = (
17 "Backward hooks on hook_scale/hook_normalized require grad-connected hook tensors; "
18 "falling back from the native-autograd path to the python-norm path. Output numerics "
19 "may differ from the unhooked forward at float-rounding scale."
20)
21NATIVE_PATH_EDIT_FALLBACK_WARNING = (
22 "A forward hook edited hook_scale/hook_normalized on the native-autograd path; the "
23 "output is reconstructed from the hooked values instead of HF's native forward. "
24 "Output numerics may differ from the unhooked forward at float-rounding scale."
25)
28class NormalizationBridge(GeneralizedComponent):
29 """Normalization bridge that wraps transformer normalization layers but implements the calculation from scratch.
31 This component provides standardized input/output hooks.
32 """
34 property_aliases = {"w": "weight", "b": "bias"}
36 def __init__(
37 self,
38 name: str,
39 config: Any,
40 submodules: Optional[Dict[str, GeneralizedComponent]] = {},
41 use_native_layernorm_autograd: bool = False,
42 uses_rms_norm: Optional[bool] = None,
43 optional: bool = False,
44 uses_weight_offset: Optional[bool] = None,
45 ):
46 """Initialize the normalization bridge.
48 Args:
49 name: The name of this component
50 config: Optional configuration
51 submodules: Dictionary of GeneralizedComponent submodules to register
52 use_native_layernorm_autograd: If True, use HuggingFace's native LayerNorm
53 autograd for exact gradient matching. If False,
54 use custom implementation. Defaults to False.
55 uses_rms_norm: Force RMSNorm vs LayerNorm; None defers to introspection
56 then ``config.uses_rms_norm``.
57 optional: If True, setup skips this subtree when absent (hybrid architectures).
58 uses_weight_offset: Override the affine weight convention independently of
59 RMS vs LayerNorm. None uses the RMSNorm configuration flag.
60 """
61 super().__init__(name, config, submodules=submodules, optional=optional)
62 self.hook_normalized = HookPoint()
63 self.hook_scale = HookPoint()
64 self.use_native_layernorm_autograd = use_native_layernorm_autograd
65 self._uses_rms_norm_override = uses_rms_norm
66 self._uses_weight_offset_override = uses_weight_offset
68 @property
69 def uses_rms_norm(self) -> bool:
70 """Whether this bridge treats the wrapped module as RMSNorm.
72 Override > module introspection > config. Introspection guards against
73 a shared config (RMSNorm LM + LayerNorm vision tower) misclassifying
74 a real ``nn.LayerNorm``.
75 """
76 if self._uses_rms_norm_override is not None:
77 return self._uses_rms_norm_override
78 component = self.original_component
79 if component is not None:
80 if isinstance(component, torch.nn.LayerNorm):
81 return False
82 if "RMSNorm" in type(component).__name__:
83 return True
84 return bool(getattr(self.config, "uses_rms_norm", False))
86 @property
87 def uses_weight_offset(self) -> bool:
88 """Whether affine scaling uses 1 + weight, including explicit LayerNorm1P slots."""
89 if self._uses_weight_offset_override is not None:
90 return self._uses_weight_offset_override
91 return self.uses_rms_norm and bool(getattr(self.config, "rmsnorm_uses_offset", False))
93 def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor:
94 """Forward pass through the normalization bridge.
96 Args:
97 hidden_states: Input hidden states
98 **kwargs: Additional arguments to pass to the original component
100 Returns:
101 Normalized output
102 """
103 if self.original_component is None: 103 ↛ 104line 103 didn't jump to line 104 because the condition on line 103 was never true
104 raise RuntimeError(
105 f"Original component not set for {self.name}. Call set_original_component() first."
106 )
107 assert self.config is not None
108 hidden_states = self.hook_in(hidden_states)
109 if self.use_native_layernorm_autograd:
110 result = self._hf_autograd_forward_with_hooks(hidden_states)
111 elif hasattr(self.config, "layer_norm_folding") and self.config.layer_norm_folding: 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true
112 result = self._hf_autograd_forward_with_hooks(hidden_states)
113 else:
114 result = self._python_norm_forward(hidden_states)
115 output = self.hook_out(result)
116 return output
118 def _python_norm_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
119 """From-scratch normalization with live hooks: edits propagate, gradients flow."""
120 # Upcast to float32 for normalization precision (matches HT's RMSNorm behavior)
121 input_dtype = hidden_states.dtype
122 if input_dtype not in (torch.float32, torch.float64):
123 hidden_states = hidden_states.float()
124 if not self.uses_rms_norm:
125 hidden_states = hidden_states - hidden_states.mean(-1, keepdim=True)
126 scale = self.hook_scale(
127 (
128 hidden_states.pow(2).mean(-1, keepdim=True) + getattr(self.config, "eps", 1e-05)
129 ).sqrt()
130 )
131 hidden_states = self.hook_normalized(hidden_states / scale)
132 return self._apply_weight_and_bias(hidden_states, input_dtype)
134 def _apply_weight_and_bias(
135 self, hidden_states: torch.Tensor, input_dtype: torch.dtype
136 ) -> torch.Tensor:
137 """Apply weight/bias in float32 before casting back (matches HF precision)."""
138 # Gemma RMSNorm and LayerNorm1P store affine weight as an offset from 1.
139 weight = (1.0 + self.weight) if self.uses_weight_offset else self.weight
140 hidden_states = hidden_states * weight
141 component = self.original_component
142 if (
143 not self.uses_rms_norm
144 and component is not None
145 and hasattr(component, "bias")
146 and component.bias is not None
147 ):
148 hidden_states = hidden_states + cast(torch.Tensor, component.bias)
149 result = hidden_states.to(input_dtype)
150 if not self.uses_rms_norm and not result.is_contiguous():
151 # F.layer_norm materializes a contiguous output while these
152 # pointwise ops preserve the input's strides; downstream HF code
153 # may .view() the result (e.g. Idefics3 pixel_shuffle).
154 result = result.contiguous()
155 return result
157 def _hf_autograd_forward_with_hooks(self, x: torch.Tensor) -> torch.Tensor:
158 """Forward pass that preserves HF's autograd while firing intermediate hooks.
160 When hooks only observe (return ``None``, e.g. ``run_with_cache``), the result is
161 HF's own forward — bit-identical numerics and exact autograd. When a forward hook
162 edits ``hook_scale`` / ``hook_normalized``, the output is reconstructed from the
163 hooked values so the edit propagates; when backward hooks are attached, the whole
164 computation takes the python-norm path so hook tensors stay in the autograd graph.
165 Both fallbacks warn, since their numerics differ from HF's at rounding scale.
167 Args:
168 x: Input tensor
170 Returns:
171 Normalized output tensor
172 """
173 if self.original_component is None: 173 ↛ 174line 173 didn't jump to line 174 because the condition on line 173 was never true
174 raise RuntimeError(f"Original component not set for {self.name}")
175 if isinstance(self.original_component, torch.nn.Identity):
176 # Non-normalizing slot (ModernBertDecoder layer 0): fire hooks with
177 # pass-through values instead of fabricated LN stats.
178 _ = self.hook_scale(torch.ones_like(x[..., :1]))
179 _ = self.hook_normalized(x)
180 return x
181 if self.hook_scale.bwd_hooks or self.hook_normalized.bwd_hooks:
182 warnings.warn(NATIVE_PATH_BWD_FALLBACK_WARNING)
183 return self._python_norm_forward(x)
184 has_fwd_hooks = bool(self.hook_scale.fwd_hooks or self.hook_normalized.fwd_hooks)
185 # No hooks: skip building a graph for observation-only intermediates. With hooks,
186 # keep grad so an edited value stays connected to the input.
187 grad_ctx: ContextManager[Any] = (
188 contextlib.nullcontext() if has_fwd_hooks else torch.no_grad()
189 )
190 with grad_ctx:
191 # Upcast to float32 for hook precision (matches HT's RMSNorm/LayerNorm behavior)
192 x_float = x.float() if x.dtype not in (torch.float32, torch.float64) else x
193 if not self.uses_rms_norm:
194 x_centered = x_float - x_float.mean(-1, keepdim=True)
195 else:
196 x_centered = x_float
197 eps_tensor = getattr(self.original_component, "eps", None)
198 if eps_tensor is None:
199 eps_tensor = getattr(self.original_component, "variance_epsilon", None)
200 if eps_tensor is None:
201 eps_value: float | torch.Tensor = getattr(self.config, "eps", 1e-05)
202 else:
203 eps_value = eps_tensor
204 variance = x_centered.pow(2).mean(-1, keepdim=True)
205 if isinstance(eps_value, torch.Tensor): 205 ↛ 206line 205 didn't jump to line 206 because the condition on line 205 was never true
206 inv_rms = torch.rsqrt(variance + eps_value)
207 scale = (variance + eps_value).sqrt()
208 else:
209 inv_rms = torch.rsqrt(variance + float(eps_value))
210 scale = (variance + float(eps_value)).sqrt()
211 # Use rsqrt for x_normalized to match HF's actual computation path
212 # (LlamaRMSNorm uses x * rsqrt(variance + eps)). Keep scale as sqrt
213 # for hook_scale (denominator convention used by HookedTransformer).
214 x_normalized = x_centered * inv_rms
215 hooked_scale = self.hook_scale(scale)
216 if hooked_scale is not scale:
217 # Edited scale: recompute with the denominator convention so the edit
218 # feeds hook_normalized, mirroring the python-norm path's ordering.
219 x_normalized = x_centered / hooked_scale
220 hooked_normalized = self.hook_normalized(x_normalized)
221 input_dtype = x.dtype
222 # A hook returning None keeps the original tensor object (see HookPoint), so
223 # identity is the edit signal. Note in-place mutation of the hook value without
224 # returning it is NOT detected — return the tensor from the hook to edit.
225 if hooked_scale is scale and hooked_normalized is x_normalized:
226 result = self.original_component(x)
227 if result.dtype != input_dtype: 227 ↛ 228line 227 didn't jump to line 228 because the condition on line 227 was never true
228 result = result.to(input_dtype)
229 return result
230 warnings.warn(NATIVE_PATH_EDIT_FALLBACK_WARNING)
231 return self._apply_weight_and_bias(hooked_normalized, input_dtype)
234class LayerNormPreBridge(NormalizationBridge):
235 """Param-free LayerNorm (LNPre): hook_scale / hook_normalized, no weight or bias."""
237 property_aliases: Dict[str, str] = {}
239 def __init__(
240 self,
241 name: str,
242 config: Any,
243 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
244 optional: bool = False,
245 ):
246 """Initialize the param-free LayerNorm bridge."""
247 super().__init__(
248 name,
249 config,
250 submodules=submodules or {},
251 use_native_layernorm_autograd=False,
252 uses_rms_norm=False,
253 optional=optional,
254 )
256 def _apply_weight_and_bias(
257 self, hidden_states: torch.Tensor, input_dtype: torch.dtype
258 ) -> torch.Tensor:
259 """No weight or bias to apply — only restore the input dtype."""
260 return hidden_states.to(input_dtype)
263class RMSNormPreBridge(NormalizationBridge):
264 """Param-free RMSNorm (RMSPre): hook_scale / hook_normalized, no learnable scale."""
266 property_aliases: Dict[str, str] = {}
268 def __init__(
269 self,
270 name: str,
271 config: Any,
272 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
273 optional: bool = False,
274 ):
275 """Initialize the param-free RMSNorm bridge."""
276 super().__init__(
277 name,
278 config,
279 submodules=submodules or {},
280 use_native_layernorm_autograd=False,
281 uses_rms_norm=True,
282 optional=optional,
283 )
285 def _apply_weight_and_bias(
286 self, hidden_states: torch.Tensor, input_dtype: torch.dtype
287 ) -> torch.Tensor:
288 """No learnable scale to apply — only restore the input dtype."""
289 return hidden_states.to(input_dtype)