From f42c2772ad59865977314492f159ecccb33968fc Mon Sep 17 00:00:00 2001 From: Hetens Date: Tue, 21 Jul 2026 18:49:49 -0400 Subject: [PATCH 1/3] Fix silent hook no-op on NormalizationBridge native-autograd path (#1526) hook_scale / hook_normalized were fired on detached observation values and their return values discarded, so forward edits silently no-oped and backward hooks never fired on the 13 architectures using use_native_layernorm_autograd. - backward hooks attached: fall back to the python-norm path (grad-connected) - forward hook edits (detected by identity: returning non-None replaces the tensor): reconstruct output from the hooked values, grad-connected - observation-only hooks (run_with_cache) and hook-free forwards: unchanged, bit-identical native HF output - both fallbacks emit a UserWarning since numerics differ at rounding scale --- .../test_normalization_hook_semantics.py | 167 ++++++++++++++++++ .../generalized_components/normalization.py | 117 ++++++++---- 2 files changed, 249 insertions(+), 35 deletions(-) create mode 100644 tests/unit/model_bridge/generalized_components/test_normalization_hook_semantics.py diff --git a/tests/unit/model_bridge/generalized_components/test_normalization_hook_semantics.py b/tests/unit/model_bridge/generalized_components/test_normalization_hook_semantics.py new file mode 100644 index 000000000..108b00457 --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_normalization_hook_semantics.py @@ -0,0 +1,167 @@ +"""Hook semantics on NormalizationBridge's native-autograd path. + +Regression coverage for #1526: with ``use_native_layernorm_autograd=True`` the +bridge fired ``hook_scale`` / ``hook_normalized`` on detached observation values +and discarded their return values — forward edits silently no-oped and backward +hooks never fired, while the default python-norm path honored both. +""" + +import warnings + +import pytest +import torch +import torch.nn as nn + +# Warnings are matched by message substring, not imported constants, so this file +# collects (and fails red) on unfixed code where the constants don't exist yet. +from transformer_lens.model_bridge.generalized_components.normalization import ( + NormalizationBridge, +) + + +class _Cfg: + def __init__(self, uses_rms_norm: bool = False, eps: float = 1e-5): + self.uses_rms_norm = uses_rms_norm + self.eps = eps + + +class _TinyRMSNorm(nn.Module): + """Minimal RMSNorm mirroring LlamaRMSNorm's forward.""" + + def __init__(self, d: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(d)) + self.variance_epsilon = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + variance = x.pow(2).mean(-1, keepdim=True) + return self.weight * x * torch.rsqrt(variance + self.variance_epsilon) + + +def _layernorm(d: int) -> nn.LayerNorm: + layer = nn.LayerNorm(d, eps=1e-5) + nn.init.normal_(layer.weight, std=0.1) + nn.init.normal_(layer.bias, std=0.1) + layer.eval() + return layer + + +def _make_bridge(native: bool, rms: bool = False, d: int = 16) -> NormalizationBridge: + layer: nn.Module = _TinyRMSNorm(d) if rms else _layernorm(d) + bridge = NormalizationBridge( + name="ln", + config=_Cfg(uses_rms_norm=rms), + use_native_layernorm_autograd=native, + ) + bridge.set_original_component(layer) + return bridge + + +def _zero_hook(tensor, hook=None): + return torch.zeros_like(tensor) + + +@pytest.mark.parametrize("rms", [False, True], ids=["layernorm", "rmsnorm"]) +def test_native_path_normalized_edit_propagates(rms): + """Editing hook_normalized must change the output on the native-autograd path.""" + bridge = _make_bridge(native=True, rms=rms) + x = torch.randn(2, 5, 16) + baseline = bridge(x) + bridge.hook_normalized.add_hook(_zero_hook) + with pytest.warns(UserWarning, match="reconstructed from the hooked values"): + patched = bridge(x) + bridge.hook_normalized.remove_hooks() + assert not torch.allclose(baseline, patched) + + +def test_native_path_scale_edit_propagates(): + """Editing hook_scale must change the output on the native-autograd path.""" + bridge = _make_bridge(native=True) + x = torch.randn(2, 5, 16) + baseline = bridge(x) + bridge.hook_scale.add_hook(lambda t, hook=None: t * 2.0) + with pytest.warns(UserWarning, match="reconstructed from the hooked values"): + patched = bridge(x) + bridge.hook_scale.remove_hooks() + assert not torch.allclose(baseline, patched) + + +def test_native_path_bwd_hook_fires(): + """Backward hooks on hook_normalized must fire on the native-autograd path.""" + bridge = _make_bridge(native=True) + x = torch.randn(2, 5, 16, requires_grad=True) + fired = {} + bridge.hook_normalized.add_hook( + lambda grad, hook=None: fired.setdefault("bwd", True) and grad, dir="bwd" + ) + with pytest.warns(UserWarning, match="Backward hooks"): + out = bridge(x) + out.sum().backward() + bridge.hook_normalized.remove_hooks() + assert fired.get("bwd", False) + + +def test_native_path_edited_output_is_grad_connected(): + """An edited forward must stay differentiable back to the input.""" + bridge = _make_bridge(native=True) + x = torch.randn(2, 5, 16, requires_grad=True) + bridge.hook_scale.add_hook(lambda t, hook=None: t * 2.0) + with pytest.warns(UserWarning): + out = bridge(x) + bridge.hook_scale.remove_hooks() + out.sum().backward() + assert x.grad is not None + assert torch.any(x.grad != 0) + + +def test_native_path_observation_hooks_keep_native_numerics(): + """run_with_cache-style hooks (return None) must not change output numerics. + + Guards the fallback dispatch: routing on "any hook attached" would silently + switch caching runs from HF's native forward to the python-norm path. + """ + bridge = _make_bridge(native=True) + x = torch.randn(2, 5, 16) + baseline = bridge(x) + cache = {} + + def observe(tensor, hook=None): + cache["normalized"] = tensor.detach() + return None + + bridge.hook_normalized.add_hook(observe) + with warnings.catch_warnings(): + warnings.simplefilter("error") # any fallback warning fails the test + observed = bridge(x) + bridge.hook_normalized.remove_hooks() + assert torch.equal(baseline, observed) + assert "normalized" in cache + + +def test_native_path_no_hooks_matches_original_component(): + """Without hooks, the native path is exactly the wrapped module's output.""" + bridge = _make_bridge(native=True) + x = torch.randn(2, 5, 16) + assert torch.equal(bridge(x), bridge.original_component(x)) + + +def test_default_path_edit_still_propagates(): + """Control: the python-norm path's editing semantics are unchanged.""" + bridge = _make_bridge(native=False) + x = torch.randn(2, 5, 16) + baseline = bridge(x) + bridge.hook_normalized.add_hook(_zero_hook) + patched = bridge(x) + bridge.hook_normalized.remove_hooks() + assert not torch.allclose(baseline, patched) + + +def test_default_path_emits_no_fallback_warnings(): + """Control: python-norm path edits never warn — warnings are native-path only.""" + bridge = _make_bridge(native=False) + x = torch.randn(2, 5, 16) + bridge.hook_normalized.add_hook(_zero_hook) + with warnings.catch_warnings(): + warnings.simplefilter("error") + bridge(x) + bridge.hook_normalized.remove_hooks() diff --git a/transformer_lens/model_bridge/generalized_components/normalization.py b/transformer_lens/model_bridge/generalized_components/normalization.py index bed44dafc..ed28fa826 100644 --- a/transformer_lens/model_bridge/generalized_components/normalization.py +++ b/transformer_lens/model_bridge/generalized_components/normalization.py @@ -1,5 +1,7 @@ """Normalization bridge component implementation.""" -from typing import Any, Dict, Optional, cast +import contextlib +import warnings +from typing import Any, ContextManager, Dict, Optional, cast import torch @@ -8,6 +10,20 @@ GeneralizedComponent, ) +# The native-autograd path returns HF's own output, so hook edits and backward hooks +# can only be honored by switching to the python-norm computation, whose numerics +# differ from HF's at float-rounding scale. +NATIVE_PATH_BWD_FALLBACK_WARNING = ( + "Backward hooks on hook_scale/hook_normalized require grad-connected hook tensors; " + "falling back from the native-autograd path to the python-norm path. Output numerics " + "may differ from the unhooked forward at float-rounding scale." +) +NATIVE_PATH_EDIT_FALLBACK_WARNING = ( + "A forward hook edited hook_scale/hook_normalized on the native-autograd path; the " + "output is reconstructed from the hooked values instead of HF's native forward. " + "Output numerics may differ from the unhooked forward at float-rounding scale." +) + class NormalizationBridge(GeneralizedComponent): """Normalization bridge that wraps transformer normalization layers but implements the calculation from scratch. @@ -83,48 +99,69 @@ def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor: elif hasattr(self.config, "layer_norm_folding") and self.config.layer_norm_folding: result = self._hf_autograd_forward_with_hooks(hidden_states) else: - uses_rms_norm = self.uses_rms_norm - # Upcast to float32 for normalization precision (matches HT's RMSNorm behavior) - input_dtype = hidden_states.dtype - if input_dtype not in (torch.float32, torch.float64): - hidden_states = hidden_states.float() - if not uses_rms_norm: - hidden_states = hidden_states - hidden_states.mean(-1, keepdim=True) - scale = self.hook_scale( - ( - hidden_states.pow(2).mean(-1, keepdim=True) + getattr(self.config, "eps", 1e-05) - ).sqrt() - ) - hidden_states = self.hook_normalized(hidden_states / scale) - # Apply weight/bias in float32 before casting back (matches HF precision). - if uses_rms_norm: - hidden_states = hidden_states * self.weight - else: - hidden_states = hidden_states * self.weight - if ( - hasattr(self.original_component, "bias") - and self.original_component.bias is not None - ): - hidden_states = hidden_states + cast(torch.Tensor, self.original_component.bias) - result = hidden_states.to(input_dtype) + result = self._python_norm_forward(hidden_states) output = self.hook_out(result) return output + def _python_norm_forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """From-scratch normalization with live hooks: edits propagate, gradients flow.""" + # Upcast to float32 for normalization precision (matches HT's RMSNorm behavior) + input_dtype = hidden_states.dtype + if input_dtype not in (torch.float32, torch.float64): + hidden_states = hidden_states.float() + if not self.uses_rms_norm: + hidden_states = hidden_states - hidden_states.mean(-1, keepdim=True) + scale = self.hook_scale( + ( + hidden_states.pow(2).mean(-1, keepdim=True) + getattr(self.config, "eps", 1e-05) + ).sqrt() + ) + hidden_states = self.hook_normalized(hidden_states / scale) + return self._apply_weight_and_bias(hidden_states, input_dtype) + + def _apply_weight_and_bias( + self, hidden_states: torch.Tensor, input_dtype: torch.dtype + ) -> torch.Tensor: + """Apply weight/bias in float32 before casting back (matches HF precision).""" + hidden_states = hidden_states * self.weight + component = self.original_component + if ( + not self.uses_rms_norm + and component is not None + and hasattr(component, "bias") + and component.bias is not None + ): + hidden_states = hidden_states + cast(torch.Tensor, component.bias) + return hidden_states.to(input_dtype) + def _hf_autograd_forward_with_hooks(self, x: torch.Tensor) -> torch.Tensor: """Forward pass that preserves HF's autograd while firing intermediate hooks. - This method calls HF's LayerNorm for the final result (to preserve exact gradients), - but also computes intermediate values to fire hook_scale and hook_normalized. + When hooks only observe (return ``None``, e.g. ``run_with_cache``), the result is + HF's own forward — bit-identical numerics and exact autograd. When a forward hook + edits ``hook_scale`` / ``hook_normalized``, the output is reconstructed from the + hooked values so the edit propagates; when backward hooks are attached, the whole + computation takes the python-norm path so hook tensors stay in the autograd graph. + Both fallbacks warn, since their numerics differ from HF's at rounding scale. Args: x: Input tensor Returns: - Normalized output tensor from HF's LayerNorm + Normalized output tensor """ if self.original_component is None: raise RuntimeError(f"Original component not set for {self.name}") - with torch.no_grad(): + if self.hook_scale.bwd_hooks or self.hook_normalized.bwd_hooks: + warnings.warn(NATIVE_PATH_BWD_FALLBACK_WARNING) + return self._python_norm_forward(x) + has_fwd_hooks = bool(self.hook_scale.fwd_hooks or self.hook_normalized.fwd_hooks) + # No hooks: skip building a graph for observation-only intermediates. With hooks, + # keep grad so an edited value stays connected to the input. + grad_ctx: ContextManager[Any] = ( + contextlib.nullcontext() if has_fwd_hooks else torch.no_grad() + ) + with grad_ctx: # Upcast to float32 for hook precision (matches HT's RMSNorm/LayerNorm behavior) x_float = x.float() if x.dtype not in (torch.float32, torch.float64) else x if not self.uses_rms_norm: @@ -149,10 +186,20 @@ def _hf_autograd_forward_with_hooks(self, x: torch.Tensor) -> torch.Tensor: # (LlamaRMSNorm uses x * rsqrt(variance + eps)). Keep scale as sqrt # for hook_scale (denominator convention used by HookedTransformer). x_normalized = x_centered * inv_rms - _ = self.hook_scale(scale) - _ = self.hook_normalized(x_normalized) + hooked_scale = self.hook_scale(scale) + if hooked_scale is not scale: + # Edited scale: recompute with the denominator convention so the edit + # feeds hook_normalized, mirroring the python-norm path's ordering. + x_normalized = x_centered / hooked_scale + hooked_normalized = self.hook_normalized(x_normalized) input_dtype = x.dtype - result = self.original_component(x) - if result.dtype != input_dtype: - result = result.to(input_dtype) - return result + # A hook returning None keeps the original tensor object (see HookPoint), so + # identity is the edit signal. Note in-place mutation of the hook value without + # returning it is NOT detected — return the tensor from the hook to edit. + if hooked_scale is scale and hooked_normalized is x_normalized: + result = self.original_component(x) + if result.dtype != input_dtype: + result = result.to(input_dtype) + return result + warnings.warn(NATIVE_PATH_EDIT_FALLBACK_WARNING) + return self._apply_weight_and_bias(hooked_normalized, input_dtype) From 7e2a1672d6d875ab9a8d618b4d24fac3221dbea2 Mon Sep 17 00:00:00 2001 From: Hetens Date: Wed, 22 Jul 2026 13:36:46 -0400 Subject: [PATCH 2/3] Apply Gemma-style (1 + weight) RMSNorm offset in fallback weight application --- .../test_normalization_hook_semantics.py | 59 +++++++++++++++++-- .../generalized_components/normalization.py | 8 ++- 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/tests/unit/model_bridge/generalized_components/test_normalization_hook_semantics.py b/tests/unit/model_bridge/generalized_components/test_normalization_hook_semantics.py index 108b00457..5a7640f71 100644 --- a/tests/unit/model_bridge/generalized_components/test_normalization_hook_semantics.py +++ b/tests/unit/model_bridge/generalized_components/test_normalization_hook_semantics.py @@ -20,9 +20,15 @@ class _Cfg: - def __init__(self, uses_rms_norm: bool = False, eps: float = 1e-5): + def __init__( + self, + uses_rms_norm: bool = False, + eps: float = 1e-5, + rmsnorm_uses_offset: bool = False, + ): self.uses_rms_norm = uses_rms_norm self.eps = eps + self.rmsnorm_uses_offset = rmsnorm_uses_offset class _TinyRMSNorm(nn.Module): @@ -38,6 +44,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.weight * x * torch.rsqrt(variance + self.variance_epsilon) +class _TinyGemmaRMSNorm(nn.Module): + """Minimal Gemma-style RMSNorm: weight is stored as an offset from 1.""" + + def __init__(self, d: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.full((d,), 0.5)) + self.variance_epsilon = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + variance = x.pow(2).mean(-1, keepdim=True) + x_normed = x * torch.rsqrt(variance + self.variance_epsilon) + return x_normed * (1.0 + self.weight) + + def _layernorm(d: int) -> nn.LayerNorm: layer = nn.LayerNorm(d, eps=1e-5) nn.init.normal_(layer.weight, std=0.1) @@ -46,11 +66,19 @@ def _layernorm(d: int) -> nn.LayerNorm: return layer -def _make_bridge(native: bool, rms: bool = False, d: int = 16) -> NormalizationBridge: - layer: nn.Module = _TinyRMSNorm(d) if rms else _layernorm(d) +def _make_bridge( + native: bool, rms: bool = False, d: int = 16, offset: bool = False +) -> NormalizationBridge: + layer: nn.Module + if offset: + layer = _TinyGemmaRMSNorm(d) + elif rms: + layer = _TinyRMSNorm(d) + else: + layer = _layernorm(d) bridge = NormalizationBridge( name="ln", - config=_Cfg(uses_rms_norm=rms), + config=_Cfg(uses_rms_norm=rms or offset, rmsnorm_uses_offset=offset), use_native_layernorm_autograd=native, ) bridge.set_original_component(layer) @@ -145,6 +173,29 @@ def test_native_path_no_hooks_matches_original_component(): assert torch.equal(bridge(x), bridge.original_component(x)) +def test_native_path_edit_fallback_applies_rmsnorm_offset(): + """Gemma-family RMSNorm multiplies by (1 + weight); the edit fallback must too.""" + bridge = _make_bridge(native=True, offset=True) + x = torch.randn(2, 5, 16) + bridge.hook_normalized.add_hook(lambda t, hook=None: torch.ones_like(t)) + with pytest.warns(UserWarning, match="reconstructed from the hooked values"): + patched = bridge(x) + bridge.hook_normalized.remove_hooks() + expected = torch.ones(2, 5, 16) * (1.0 + bridge.original_component.weight) + assert torch.allclose(patched, expected) + + +def test_native_path_bwd_fallback_applies_rmsnorm_offset(): + """The python-norm fallback for backward hooks must match Gemma-style output.""" + bridge = _make_bridge(native=True, offset=True) + x = torch.randn(2, 5, 16, requires_grad=True) + bridge.hook_normalized.add_hook(lambda grad, hook=None: grad, dir="bwd") + with pytest.warns(UserWarning, match="Backward hooks"): + out = bridge(x) + bridge.hook_normalized.remove_hooks() + assert torch.allclose(out, bridge.original_component(x), atol=1e-6) + + def test_default_path_edit_still_propagates(): """Control: the python-norm path's editing semantics are unchanged.""" bridge = _make_bridge(native=False) diff --git a/transformer_lens/model_bridge/generalized_components/normalization.py b/transformer_lens/model_bridge/generalized_components/normalization.py index ed28fa826..e720884ac 100644 --- a/transformer_lens/model_bridge/generalized_components/normalization.py +++ b/transformer_lens/model_bridge/generalized_components/normalization.py @@ -123,7 +123,13 @@ def _apply_weight_and_bias( self, hidden_states: torch.Tensor, input_dtype: torch.dtype ) -> torch.Tensor: """Apply weight/bias in float32 before casting back (matches HF precision).""" - hidden_states = hidden_states * self.weight + # Gemma-family RMSNorm stores weight as an offset from 1 (output uses 1 + weight). + weight = ( + (1.0 + self.weight) + if getattr(self.config, "rmsnorm_uses_offset", False) + else self.weight + ) + hidden_states = hidden_states * weight component = self.original_component if ( not self.uses_rms_norm From 08982fc4200f225d751358bc399aafda7588a2e5 Mon Sep 17 00:00:00 2001 From: Hetens Date: Wed, 22 Jul 2026 14:21:57 -0400 Subject: [PATCH 3/3] Retrigger CI after HF Hub outage