Coverage for transformer_lens/model_bridge/generalized_components/normalization.py: 95%
178 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Normalization bridge component implementation."""
2import contextlib
3import warnings
4from typing import Any, ContextManager, Dict, Optional, Tuple, cast
6import torch
8from transformer_lens.hook_points import HookPoint
9from transformer_lens.model_bridge._relevance_rules import (
10 RelevanceRuleConflictError,
11 ln_rule_grad,
12)
13from transformer_lens.model_bridge.generalized_components.base import (
14 GeneralizedComponent,
15)
17# The native-autograd path returns HF's own output, so hook edits and backward hooks
18# can only be honored by switching to the python-norm computation, whose numerics
19# differ from HF's at float-rounding scale.
20NATIVE_PATH_BWD_FALLBACK_WARNING = (
21 "Backward hooks on hook_scale/hook_normalized require grad-connected hook tensors; "
22 "falling back from the native-autograd path to the python-norm path. Output numerics "
23 "may differ from the unhooked forward at float-rounding scale."
24)
25NATIVE_PATH_EDIT_FALLBACK_WARNING = (
26 "A forward hook edited hook_scale/hook_normalized on the native-autograd path; the "
27 "output is reconstructed from the hooked values instead of HF's native forward. "
28 "Output numerics may differ from the unhooked forward at float-rounding scale."
29)
30# While the LN-rule is active, the fallbacks above would silently compose the rule
31# with the hook edit and break the bit-identical-forward guarantee, so they raise
32# instead of warning.
33RULE_ACTIVE_BWD_HOOK_CONFLICT = (
34 "Backward hooks on hook_scale/hook_normalized are incompatible with an active "
35 "LN-rule on '{name}': the rule-wrapped native forward keeps these hook points "
36 "out of its backward graph, so a backward hook here would silently never fire."
37)
38RULE_ACTIVE_EDIT_HOOK_CONFLICT = (
39 "A forward hook edited hook_scale/hook_normalized while the LN-rule is active on "
40 "'{name}': honoring the edit would require the python-norm fallback, which would "
41 "compose the rule with the edit and break the bit-identical-forward guarantee."
42)
45class _NativeLNRuleForward(torch.autograd.Function):
46 """Wrap a normalization module's own forward call in the LN-rule's VJP.
48 Forward returns ``component(x)`` unchanged, so the result is bit-identical to
49 the native forward by construction. Backward routes the x-path gradient
50 through the centering op ordinarily but treats ``denom`` as a constant (the
51 LN-rule), while ``weight`` and ``bias`` receive their ordinary gradient since
52 the rule only redefines how relevance reaches the input, not parameter
53 training gradients. A parameter-free norm (``weight`` is ``None``, e.g.
54 OLMo's ``OlmoLayerNorm``) is treated as a unit scale in both directions.
55 """
57 @staticmethod
58 def forward(
59 ctx: Any,
60 x_centered: torch.Tensor,
61 denom: torch.Tensor,
62 weight: Optional[torch.Tensor],
63 bias: Optional[torch.Tensor],
64 x: torch.Tensor,
65 component: torch.nn.Module,
66 offset: bool,
67 input_dtype: torch.dtype,
68 ) -> torch.Tensor:
69 ctx.save_for_backward(x_centered, denom, weight)
70 ctx.has_bias = bias is not None
71 ctx.bias_requires_grad = bool(bias is not None and bias.requires_grad)
72 ctx.offset = offset
73 result = component(x)
74 if result.dtype != input_dtype: 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true
75 result = result.to(input_dtype)
76 return result
78 @staticmethod
79 def backward(
80 ctx: Any, grad_output: torch.Tensor
81 ) -> Tuple[
82 torch.Tensor,
83 None,
84 Optional[torch.Tensor],
85 Optional[torch.Tensor],
86 None,
87 None,
88 None,
89 None,
90 ]:
91 x_centered, denom, weight = ctx.saved_tensors
92 if weight is None:
93 w_eff: torch.Tensor | float = 1.0
94 else:
95 w_eff = (1.0 + weight) if ctx.offset else weight
96 reduce_dims = tuple(range(grad_output.dim() - 1))
97 grad_x_centered = ln_rule_grad(grad_output * w_eff, denom)
98 grad_weight = (
99 (grad_output * (x_centered / denom)).sum(dim=reduce_dims)
100 if weight is not None and weight.requires_grad
101 else None
102 )
103 grad_bias = (
104 grad_output.sum(dim=reduce_dims) if ctx.has_bias and ctx.bias_requires_grad else None
105 )
106 return grad_x_centered, None, grad_weight, grad_bias, None, None, None, None
109class NormalizationBridge(GeneralizedComponent):
110 """Normalization bridge that wraps transformer normalization layers but implements the calculation from scratch.
112 This component provides standardized input/output hooks.
113 """
115 property_aliases = {"w": "weight", "b": "bias"}
117 def __init__(
118 self,
119 name: str,
120 config: Any,
121 submodules: Optional[Dict[str, GeneralizedComponent]] = {},
122 use_native_layernorm_autograd: bool = False,
123 uses_rms_norm: Optional[bool] = None,
124 optional: bool = False,
125 ):
126 """Initialize the normalization bridge.
128 Args:
129 name: The name of this component
130 config: Optional configuration
131 submodules: Dictionary of GeneralizedComponent submodules to register
132 use_native_layernorm_autograd: If True, use HuggingFace's native LayerNorm
133 autograd for exact gradient matching. If False,
134 use custom implementation. Defaults to False.
135 uses_rms_norm: Force RMSNorm vs LayerNorm; None defers to introspection
136 then ``config.uses_rms_norm``.
137 optional: If True, setup skips this subtree when absent (hybrid architectures).
138 """
139 super().__init__(name, config, submodules=submodules, optional=optional)
140 self.hook_normalized = HookPoint()
141 self.hook_scale = HookPoint()
142 self.use_native_layernorm_autograd = use_native_layernorm_autograd
143 self._uses_rms_norm_override = uses_rms_norm
144 self._relevance_rule_active = False
146 @property
147 def _relevance_rule_kinds(self) -> Tuple[str, ...]:
148 """``("normalization",)`` only when this instance actually dispatches through
149 the native-autograd branch the LN-rule wraps (``_hf_autograd_forward_with_hooks``);
150 empty otherwise, so ``use_relevance_rules`` reports an ln1/ln2 mount that
151 uses the plain python-norm path (for example ``LayerNormPreBridge`` /
152 ``RMSNormPreBridge``, or a config without ``layer_norm_folding``) as
153 skipped rather than silently leaving ordinary gradients in place under a
154 claimed "installed" rule.
155 """
156 if self.use_native_layernorm_autograd:
157 return ("normalization",)
158 if bool(getattr(self.config, "layer_norm_folding", False)):
159 return ("normalization",)
160 return ()
162 def _enable_relevance_rule(self, kind: str) -> None:
163 """Activate the LN-rule for this instance's native-forward branch only."""
164 self._relevance_rule_active = True
166 def _disable_relevance_rule(self, kind: str) -> None:
167 """Deactivate the LN-rule, restoring today's native-forward behavior."""
168 self._relevance_rule_active = False
170 @property
171 def uses_rms_norm(self) -> bool:
172 """Whether this bridge treats the wrapped module as RMSNorm.
174 Override > module introspection > config. Introspection guards against
175 a shared config (RMSNorm LM + LayerNorm vision tower) misclassifying
176 a real ``nn.LayerNorm``.
177 """
178 if self._uses_rms_norm_override is not None:
179 return self._uses_rms_norm_override
180 component = self.original_component
181 if component is not None:
182 if isinstance(component, torch.nn.LayerNorm):
183 return False
184 if "RMSNorm" in type(component).__name__:
185 return True
186 return bool(getattr(self.config, "uses_rms_norm", False))
188 def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor:
189 """Forward pass through the normalization bridge.
191 Args:
192 hidden_states: Input hidden states
193 **kwargs: Additional arguments to pass to the original component
195 Returns:
196 Normalized output
197 """
198 if self.original_component is None: 198 ↛ 199line 198 didn't jump to line 199 because the condition on line 198 was never true
199 raise RuntimeError(
200 f"Original component not set for {self.name}. Call set_original_component() first."
201 )
202 assert self.config is not None
203 hidden_states = self.hook_in(hidden_states)
204 if self.use_native_layernorm_autograd:
205 result = self._hf_autograd_forward_with_hooks(hidden_states)
206 elif hasattr(self.config, "layer_norm_folding") and self.config.layer_norm_folding:
207 result = self._hf_autograd_forward_with_hooks(hidden_states)
208 else:
209 result = self._python_norm_forward(hidden_states)
210 output = self.hook_out(result)
211 return output
213 def _python_norm_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
214 """From-scratch normalization with live hooks: edits propagate, gradients flow."""
215 # Upcast to float32 for normalization precision (matches HT's RMSNorm behavior)
216 input_dtype = hidden_states.dtype
217 if input_dtype not in (torch.float32, torch.float64):
218 hidden_states = hidden_states.float()
219 if not self.uses_rms_norm:
220 hidden_states = hidden_states - hidden_states.mean(-1, keepdim=True)
221 scale = self.hook_scale(
222 (
223 hidden_states.pow(2).mean(-1, keepdim=True) + getattr(self.config, "eps", 1e-05)
224 ).sqrt()
225 )
226 hidden_states = self.hook_normalized(hidden_states / scale)
227 return self._apply_weight_and_bias(hidden_states, input_dtype)
229 def _apply_weight_and_bias(
230 self, hidden_states: torch.Tensor, input_dtype: torch.dtype
231 ) -> torch.Tensor:
232 """Apply weight/bias in float32 before casting back (matches HF precision)."""
233 # Gemma-family RMSNorm stores weight as an offset from 1 (output uses 1 + weight).
234 weight = (
235 (1.0 + self.weight)
236 if getattr(self.config, "rmsnorm_uses_offset", False)
237 else self.weight
238 )
239 hidden_states = hidden_states * weight
240 component = self.original_component
241 if (
242 not self.uses_rms_norm
243 and component is not None
244 and hasattr(component, "bias")
245 and component.bias is not None
246 ):
247 hidden_states = hidden_states + cast(torch.Tensor, component.bias)
248 result = hidden_states.to(input_dtype)
249 if not self.uses_rms_norm and not result.is_contiguous():
250 # F.layer_norm materializes a contiguous output while these
251 # pointwise ops preserve the input's strides; downstream HF code
252 # may .view() the result (e.g. Idefics3 pixel_shuffle).
253 result = result.contiguous()
254 return result
256 def _hf_autograd_forward_with_hooks(self, x: torch.Tensor) -> torch.Tensor:
257 """Forward pass that preserves HF's autograd while firing intermediate hooks.
259 When hooks only observe (return ``None``, e.g. ``run_with_cache``), the result is
260 HF's own forward — bit-identical numerics and exact autograd. When a forward hook
261 edits ``hook_scale`` / ``hook_normalized``, the output is reconstructed from the
262 hooked values so the edit propagates; when backward hooks are attached, the whole
263 computation takes the python-norm path so hook tensors stay in the autograd graph.
264 Both fallbacks warn, since their numerics differ from HF's at rounding scale.
266 Args:
267 x: Input tensor
269 Returns:
270 Normalized output tensor
271 """
272 if self.original_component is None: 272 ↛ 273line 272 didn't jump to line 273 because the condition on line 272 was never true
273 raise RuntimeError(f"Original component not set for {self.name}")
274 if isinstance(self.original_component, torch.nn.Identity):
275 # Non-normalizing slot (ModernBertDecoder layer 0): fire hooks with
276 # pass-through values instead of fabricated LN stats.
277 _ = self.hook_scale(torch.ones_like(x[..., :1]))
278 _ = self.hook_normalized(x)
279 return x
280 if self.hook_scale.bwd_hooks or self.hook_normalized.bwd_hooks:
281 if self._relevance_rule_active:
282 raise RelevanceRuleConflictError(
283 RULE_ACTIVE_BWD_HOOK_CONFLICT.format(name=self.name)
284 )
285 warnings.warn(NATIVE_PATH_BWD_FALLBACK_WARNING)
286 return self._python_norm_forward(x)
287 has_fwd_hooks = bool(self.hook_scale.fwd_hooks or self.hook_normalized.fwd_hooks)
288 # No hooks: skip building a graph for observation-only intermediates. With hooks,
289 # keep grad so an edited value stays connected to the input.
290 grad_ctx: ContextManager[Any] = (
291 contextlib.nullcontext() if has_fwd_hooks else torch.no_grad()
292 )
293 with grad_ctx:
294 # Upcast to float32 for hook precision (matches HT's RMSNorm/LayerNorm behavior)
295 x_float = x.float() if x.dtype not in (torch.float32, torch.float64) else x
296 if not self.uses_rms_norm:
297 x_centered = x_float - x_float.mean(-1, keepdim=True)
298 else:
299 x_centered = x_float
300 eps_tensor = getattr(self.original_component, "eps", None)
301 if eps_tensor is None:
302 eps_tensor = getattr(self.original_component, "variance_epsilon", None)
303 if eps_tensor is None:
304 eps_value: float | torch.Tensor = getattr(self.config, "eps", 1e-05)
305 else:
306 eps_value = eps_tensor
307 variance = x_centered.pow(2).mean(-1, keepdim=True)
308 if isinstance(eps_value, torch.Tensor): 308 ↛ 309line 308 didn't jump to line 309 because the condition on line 308 was never true
309 inv_rms = torch.rsqrt(variance + eps_value)
310 scale = (variance + eps_value).sqrt()
311 else:
312 inv_rms = torch.rsqrt(variance + float(eps_value))
313 scale = (variance + float(eps_value)).sqrt()
314 # Use rsqrt for x_normalized to match HF's actual computation path
315 # (LlamaRMSNorm uses x * rsqrt(variance + eps)). Keep scale as sqrt
316 # for hook_scale (denominator convention used by HookedTransformer).
317 x_normalized = x_centered * inv_rms
318 hooked_scale = self.hook_scale(scale)
319 if hooked_scale is not scale:
320 # Edited scale: recompute with the denominator convention so the edit
321 # feeds hook_normalized, mirroring the python-norm path's ordering.
322 x_normalized = x_centered / hooked_scale
323 hooked_normalized = self.hook_normalized(x_normalized)
324 input_dtype = x.dtype
325 # A hook returning None keeps the original tensor object (see HookPoint), so
326 # identity is the edit signal. Note in-place mutation of the hook value without
327 # returning it is NOT detected — return the tensor from the hook to edit.
328 if hooked_scale is scale and hooked_normalized is x_normalized:
329 if self._relevance_rule_active:
330 return self._native_forward_with_ln_rule(x, input_dtype)
331 result = self.original_component(x)
332 if result.dtype != input_dtype: 332 ↛ 333line 332 didn't jump to line 333 because the condition on line 332 was never true
333 result = result.to(input_dtype)
334 return result
335 if self._relevance_rule_active:
336 raise RelevanceRuleConflictError(RULE_ACTIVE_EDIT_HOOK_CONFLICT.format(name=self.name))
337 warnings.warn(NATIVE_PATH_EDIT_FALLBACK_WARNING)
338 return self._apply_weight_and_bias(hooked_normalized, input_dtype)
340 def _native_forward_with_ln_rule(
341 self, x: torch.Tensor, input_dtype: torch.dtype
342 ) -> torch.Tensor:
343 """Call the native forward unchanged while routing its backward through the LN-rule.
345 Recomputes the centered numerator and denominator independently of the
346 hook-observation pass above (which may run under ``torch.no_grad()``), so
347 the recompute here stays grad-connected to ``x`` regardless of whether
348 forward hooks are attached. Autograd then flows from the returned tensor
349 back through the centering op ordinarily and, at the LN-rule Function
350 boundary, treats the denominator as a constant -- the same contract as
351 the shared ``ln_rule`` primitive, without reproducing the native kernel's
352 own forward numerics.
353 """
354 component = self.original_component
355 assert component is not None
356 x_float = x.float() if x.dtype not in (torch.float32, torch.float64) else x
357 if not self.uses_rms_norm:
358 x_centered = x_float - x_float.mean(-1, keepdim=True)
359 else:
360 x_centered = x_float
361 eps_tensor = getattr(component, "eps", None)
362 if eps_tensor is None:
363 eps_tensor = getattr(component, "variance_epsilon", None)
364 if eps_tensor is None: 364 ↛ 365line 364 didn't jump to line 365 because the condition on line 364 was never true
365 eps_value: float | torch.Tensor = getattr(self.config, "eps", 1e-05)
366 else:
367 eps_value = eps_tensor
368 variance = x_centered.pow(2).mean(-1, keepdim=True)
369 denom = (
370 (variance + eps_value).sqrt()
371 if isinstance(eps_value, torch.Tensor)
372 else (variance + float(eps_value)).sqrt()
373 )
374 # A parameter-free norm (OLMo's OlmoLayerNorm) exposes no usable weight;
375 # pass None so the rule treats the scale as 1 in both forward and backward.
376 try:
377 weight: Optional[torch.Tensor] = cast(torch.Tensor, self.weight)
378 except AttributeError:
379 weight = None
380 bias = getattr(component, "bias", None) if not self.uses_rms_norm else None
381 offset = bool(getattr(self.config, "rmsnorm_uses_offset", False))
382 result: torch.Tensor = _NativeLNRuleForward.apply(
383 x_centered, denom, weight, bias, x, component, offset, input_dtype
384 )
385 return result
388class LayerNormPreBridge(NormalizationBridge):
389 """Param-free LayerNorm (LNPre): hook_scale / hook_normalized, no weight or bias."""
391 property_aliases: Dict[str, str] = {}
393 def __init__(
394 self,
395 name: str,
396 config: Any,
397 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
398 optional: bool = False,
399 ):
400 """Initialize the param-free LayerNorm bridge."""
401 super().__init__(
402 name,
403 config,
404 submodules=submodules or {},
405 use_native_layernorm_autograd=False,
406 uses_rms_norm=False,
407 optional=optional,
408 )
410 def _apply_weight_and_bias(
411 self, hidden_states: torch.Tensor, input_dtype: torch.dtype
412 ) -> torch.Tensor:
413 """No weight or bias to apply — only restore the input dtype."""
414 return hidden_states.to(input_dtype)
417class RMSNormPreBridge(NormalizationBridge):
418 """Param-free RMSNorm (RMSPre): hook_scale / hook_normalized, no learnable scale."""
420 property_aliases: Dict[str, str] = {}
422 def __init__(
423 self,
424 name: str,
425 config: Any,
426 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
427 optional: bool = False,
428 ):
429 """Initialize the param-free RMSNorm bridge."""
430 super().__init__(
431 name,
432 config,
433 submodules=submodules or {},
434 use_native_layernorm_autograd=False,
435 uses_rms_norm=True,
436 optional=optional,
437 )
439 def _apply_weight_and_bias(
440 self, hidden_states: torch.Tensor, input_dtype: torch.dtype
441 ) -> torch.Tensor:
442 """No learnable scale to apply — only restore the input dtype."""
443 return hidden_states.to(input_dtype)