Coverage for transformer_lens/model_bridge/generalized

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 ):

45 """Initialize the normalization bridge.

47 Args:

48 name: The name of this component

49 config: Optional configuration

50 submodules: Dictionary of GeneralizedComponent submodules to register

51 use_native_layernorm_autograd: If True, use HuggingFace's native LayerNorm

52 autograd for exact gradient matching. If False,

53 use custom implementation. Defaults to False.

54 uses_rms_norm: Force RMSNorm vs LayerNorm; None defers to introspection

55 then ``config.uses_rms_norm``.

56 optional: If True, setup skips this subtree when absent (hybrid architectures).

57 """

58 super().__init__(name, config, submodules=submodules, optional=optional)

59 self.hook_normalized = HookPoint()

60 self.hook_scale = HookPoint()

61 self.use_native_layernorm_autograd = use_native_layernorm_autograd

62 self._uses_rms_norm_override = uses_rms_norm

64 @property

65 def uses_rms_norm(self) -> bool:

66 """Whether this bridge treats the wrapped module as RMSNorm.

68 Override > module introspection > config. Introspection guards against

69 a shared config (RMSNorm LM + LayerNorm vision tower) misclassifying

70 a real ``nn.LayerNorm``.

71 """

72 if self._uses_rms_norm_override is not None:

73 return self._uses_rms_norm_override

74 component = self.original_component

75 if component is not None:

76 if isinstance(component, torch.nn.LayerNorm):

77 return False

78 if "RMSNorm" in type(component).__name__:

79 return True

80 return bool(getattr(self.config, "uses_rms_norm", False))

82 def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor:

83 """Forward pass through the normalization bridge.

85 Args:

86 hidden_states: Input hidden states

87 **kwargs: Additional arguments to pass to the original component

89 Returns:

90 Normalized output

91 """

92 if self.original_component is None: 92 ↛ 93line 92 didn't jump to line 93 because the condition on line 92 was never true

93 raise RuntimeError(

94 f"Original component not set for {self.name}. Call set_original_component() first."

95 )

96 assert self.config is not None

97 hidden_states = self.hook_in(hidden_states)

98 if self.use_native_layernorm_autograd:

99 result = self._hf_autograd_forward_with_hooks(hidden_states)

100 elif hasattr(self.config, "layer_norm_folding") and self.config.layer_norm_folding: 100 ↛ 101line 100 didn't jump to line 101 because the condition on line 100 was never true

101 result = self._hf_autograd_forward_with_hooks(hidden_states)

102 else:

103 result = self._python_norm_forward(hidden_states)

104 output = self.hook_out(result)

105 return output

106

107 def _python_norm_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

108 """From-scratch normalization with live hooks: edits propagate, gradients flow."""

109 # Upcast to float32 for normalization precision (matches HT's RMSNorm behavior)

110 input_dtype = hidden_states.dtype

111 if input_dtype not in (torch.float32, torch.float64): 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true

112 hidden_states = hidden_states.float()

113 if not self.uses_rms_norm:

114 hidden_states = hidden_states - hidden_states.mean(-1, keepdim=True)

115 scale = self.hook_scale(

116 (

117 hidden_states.pow(2).mean(-1, keepdim=True) + getattr(self.config, "eps", 1e-05)

118 ).sqrt()

119 )

120 hidden_states = self.hook_normalized(hidden_states / scale)

121 return self._apply_weight_and_bias(hidden_states, input_dtype)

122

123 def _apply_weight_and_bias(

124 self, hidden_states: torch.Tensor, input_dtype: torch.dtype

125 ) -> torch.Tensor:

126 """Apply weight/bias in float32 before casting back (matches HF precision)."""

127 # Gemma-family RMSNorm stores weight as an offset from 1 (output uses 1 + weight).

128 weight = (

129 (1.0 + self.weight)

130 if getattr(self.config, "rmsnorm_uses_offset", False)

131 else self.weight

132 )

133 hidden_states = hidden_states * weight

134 component = self.original_component

135 if (

136 not self.uses_rms_norm

137 and component is not None

138 and hasattr(component, "bias")

139 and component.bias is not None

140 ):

141 hidden_states = hidden_states + cast(torch.Tensor, component.bias)

142 result = hidden_states.to(input_dtype)

143 if not self.uses_rms_norm and not result.is_contiguous():

144 # F.layer_norm materializes a contiguous output while these

145 # pointwise ops preserve the input's strides; downstream HF code

146 # may .view() the result (e.g. Idefics3 pixel_shuffle).

147 result = result.contiguous()

148 return result

149

150 def _hf_autograd_forward_with_hooks(self, x: torch.Tensor) -> torch.Tensor:

151 """Forward pass that preserves HF's autograd while firing intermediate hooks.

152

153 When hooks only observe (return ``None``, e.g. ``run_with_cache``), the result is

154 HF's own forward — bit-identical numerics and exact autograd. When a forward hook

155 edits ``hook_scale`` / ``hook_normalized``, the output is reconstructed from the

156 hooked values so the edit propagates; when backward hooks are attached, the whole

157 computation takes the python-norm path so hook tensors stay in the autograd graph.

158 Both fallbacks warn, since their numerics differ from HF's at rounding scale.

159

160 Args:

161 x: Input tensor

162

163 Returns:

164 Normalized output tensor

165 """

166 if self.original_component is None: 166 ↛ 167line 166 didn't jump to line 167 because the condition on line 166 was never true

167 raise RuntimeError(f"Original component not set for {self.name}")

168 if isinstance(self.original_component, torch.nn.Identity):

169 # Non-normalizing slot (ModernBertDecoder layer 0): fire hooks with

170 # pass-through values instead of fabricated LN stats.

171 _ = self.hook_scale(torch.ones_like(x[..., :1]))

172 _ = self.hook_normalized(x)

173 return x

174 if self.hook_scale.bwd_hooks or self.hook_normalized.bwd_hooks:

175 warnings.warn(NATIVE_PATH_BWD_FALLBACK_WARNING)

176 return self._python_norm_forward(x)

177 has_fwd_hooks = bool(self.hook_scale.fwd_hooks or self.hook_normalized.fwd_hooks)

178 # No hooks: skip building a graph for observation-only intermediates. With hooks,

179 # keep grad so an edited value stays connected to the input.

180 grad_ctx: ContextManager[Any] = (

181 contextlib.nullcontext() if has_fwd_hooks else torch.no_grad()

182 )

183 with grad_ctx:

184 # Upcast to float32 for hook precision (matches HT's RMSNorm/LayerNorm behavior)

185 x_float = x.float() if x.dtype not in (torch.float32, torch.float64) else x

186 if not self.uses_rms_norm:

187 x_centered = x_float - x_float.mean(-1, keepdim=True)

188 else:

189 x_centered = x_float

190 eps_tensor = getattr(self.original_component, "eps", None)

191 if eps_tensor is None:

192 eps_tensor = getattr(self.original_component, "variance_epsilon", None)

193 if eps_tensor is None:

194 eps_value: float | torch.Tensor = getattr(self.config, "eps", 1e-05)

195 else:

196 eps_value = eps_tensor

197 variance = x_centered.pow(2).mean(-1, keepdim=True)

198 if isinstance(eps_value, torch.Tensor): 198 ↛ 199line 198 didn't jump to line 199 because the condition on line 198 was never true

199 inv_rms = torch.rsqrt(variance + eps_value)

200 scale = (variance + eps_value).sqrt()

201 else:

202 inv_rms = torch.rsqrt(variance + float(eps_value))

203 scale = (variance + float(eps_value)).sqrt()

204 # Use rsqrt for x_normalized to match HF's actual computation path

205 # (LlamaRMSNorm uses x * rsqrt(variance + eps)). Keep scale as sqrt

206 # for hook_scale (denominator convention used by HookedTransformer).

207 x_normalized = x_centered * inv_rms

208 hooked_scale = self.hook_scale(scale)

209 if hooked_scale is not scale:

210 # Edited scale: recompute with the denominator convention so the edit

211 # feeds hook_normalized, mirroring the python-norm path's ordering.

212 x_normalized = x_centered / hooked_scale

213 hooked_normalized = self.hook_normalized(x_normalized)

214 input_dtype = x.dtype

215 # A hook returning None keeps the original tensor object (see HookPoint), so

216 # identity is the edit signal. Note in-place mutation of the hook value without

217 # returning it is NOT detected — return the tensor from the hook to edit.

218 if hooked_scale is scale and hooked_normalized is x_normalized:

219 result = self.original_component(x)

220 if result.dtype != input_dtype: 220 ↛ 221line 220 didn't jump to line 221 because the condition on line 220 was never true

221 result = result.to(input_dtype)

222 return result

223 warnings.warn(NATIVE_PATH_EDIT_FALLBACK_WARNING)

224 return self._apply_weight_and_bias(hooked_normalized, input_dtype)

Coverage for transformer_lens/model_bridge/generalized_components/normalization.py: 91%

100 statements