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

44 """Initialize the normalization bridge.

46 Args:

47 name: The name of this component

48 config: Optional configuration

49 submodules: Dictionary of GeneralizedComponent submodules to register

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

51 autograd for exact gradient matching. If False,

52 use custom implementation. Defaults to False.

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

54 then ``config.uses_rms_norm``.

55 """

56 super().__init__(name, config, submodules=submodules)

57 self.hook_normalized = HookPoint()

58 self.hook_scale = HookPoint()

59 self.use_native_layernorm_autograd = use_native_layernorm_autograd

60 self._uses_rms_norm_override = uses_rms_norm

62 @property

63 def uses_rms_norm(self) -> bool:

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

66 Override > module introspection > config. Introspection guards against

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

68 a real ``nn.LayerNorm``.

69 """

70 if self._uses_rms_norm_override is not None:

71 return self._uses_rms_norm_override

72 component = self.original_component

73 if component is not None:

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

75 return False

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

77 return True

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

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

81 """Forward pass through the normalization bridge.

83 Args:

84 hidden_states: Input hidden states

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

87 Returns:

88 Normalized output

89 """

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

91 raise RuntimeError(

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

93 )

94 assert self.config is not None

95 hidden_states = self.hook_in(hidden_states)

96 self._last_input_before_norm = hidden_states

97 if self.use_native_layernorm_autograd:

98 result = self._hf_autograd_forward_with_hooks(hidden_states)

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

100 result = self._hf_autograd_forward_with_hooks(hidden_states)

101 else:

102 result = self._python_norm_forward(hidden_states)

103 output = self.hook_out(result)

104 return output

105

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

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

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

109 input_dtype = hidden_states.dtype

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

111 hidden_states = hidden_states.float()

112 if not self.uses_rms_norm:

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

114 scale = self.hook_scale(

115 (

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

117 ).sqrt()

118 )

119 hidden_states = self.hook_normalized(hidden_states / scale)

120 return self._apply_weight_and_bias(hidden_states, input_dtype)

121

122 def _apply_weight_and_bias(

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

124 ) -> torch.Tensor:

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

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

127 weight = (

128 (1.0 + self.weight)

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

130 else self.weight

131 )

132 hidden_states = hidden_states * weight

133 component = self.original_component

134 if (

135 not self.uses_rms_norm

136 and component is not None

137 and hasattr(component, "bias")

138 and component.bias is not None

139 ):

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

141 return hidden_states.to(input_dtype)

142

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

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

145

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

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

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

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

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

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

152

153 Args:

154 x: Input tensor

155

156 Returns:

157 Normalized output tensor

158 """

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

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

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

162 warnings.warn(NATIVE_PATH_BWD_FALLBACK_WARNING)

163 return self._python_norm_forward(x)

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

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

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

167 grad_ctx: ContextManager[Any] = (

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

169 )

170 with grad_ctx:

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

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

173 if not self.uses_rms_norm:

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

175 else:

176 x_centered = x_float

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

178 if eps_tensor is None:

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

180 if eps_tensor is None:

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

182 else:

183 eps_value = eps_tensor

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

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

186 inv_rms = torch.rsqrt(variance + eps_value)

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

188 else:

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

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

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

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

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

194 x_normalized = x_centered * inv_rms

195 hooked_scale = self.hook_scale(scale)

196 if hooked_scale is not scale:

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

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

199 x_normalized = x_centered / hooked_scale

200 hooked_normalized = self.hook_normalized(x_normalized)

201 input_dtype = x.dtype

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

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

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

205 if hooked_scale is scale and hooked_normalized is x_normalized:

206 result = self.original_component(x)

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

208 result = result.to(input_dtype)

209 return result

210 warnings.warn(NATIVE_PATH_EDIT_FALLBACK_WARNING)

211 return self._apply_weight_and_bias(hooked_normalized, input_dtype)

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

94 statements