Fix silent hook_normalized/hook_scale no-op on NormalizationBridge native-autograd path

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

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)