fix: honor in-place edits on native normalization hooks

tests/integration/model_bridge/test_native_norm_inplace_hooks.py

@@ -0,0 +1,60 @@
+"""Native normalization interventions through a tiny real Llama forward, without downloads."""
+
+import pytest
+import torch
+from transformers import LlamaConfig, LlamaForCausalLM
+
+from transformer_lens.model_bridge.sources import build_bridge_from_module
+
+
+@pytest.fixture(scope="module")
+def tiny_llama():
+    with torch.random.fork_rng():
+        torch.manual_seed(3)
+        config = LlamaConfig(
+            vocab_size=32,
+            hidden_size=16,
+            intermediate_size=32,
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            max_position_embeddings=32,
+        )
+        config._attn_implementation = "eager"
+        hf = LlamaForCausalLM(config).eval()
+        tokens = torch.tensor([[2, 3, 4]])
+        with torch.no_grad():
+            reference = hf(tokens).logits
+        bridge = build_bridge_from_module(hf, "LlamaForCausalLM", hf_config=config, device="cpu")
+    return bridge, tokens, reference
+
+
+@pytest.mark.parametrize(
+    "mode", [torch.no_grad, torch.inference_mode], ids=["no-grad", "inference"]
+)
+@pytest.mark.parametrize(
+    "name",
+    [
+        "ln_final.hook_normalized",
+        "blocks.0.ln1.hook_normalized",
+        "ln_final.hook_scale",
+        "blocks.0.ln1.hook_scale",
+    ],
+)
+def test_llama_inplace_norm_intervention_matches_replacement(tiny_llama, name, mode):
+    bridge, tokens, reference = tiny_llama
+    with mode():
+        assert torch.equal(bridge(tokens), reference)
+        observed, _ = bridge.run_with_cache(tokens)
+        assert torch.equal(observed, reference)
+        with pytest.warns(UserWarning, match="reconstructed from the hooked values"):
+            expected = bridge.run_with_hooks(
+                tokens, fwd_hooks=[(name, lambda value, hook: value * 2)]
+            )
+        with pytest.warns(UserWarning, match="reconstructed from the hooked values"):
+            actual = bridge.run_with_hooks(
+                tokens, fwd_hooks=[(name, lambda value, hook: value.mul_(2))]
+            )
+        torch.testing.assert_close(actual, expected)
+        assert not torch.equal(actual, reference)
+        assert torch.equal(bridge(tokens), reference)

tests/unit/model_bridge/generalized_components/test_normalization_hook_semantics.py

@@ -271,3 +271,90 @@ def test_layernorm_edit_uses_weight_and_bias_despite_shared_rms_offset():
         torch.ones_like(x) * bridge.original_component.weight + bridge.original_component.bias
     )
     torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.parametrize("kind", ["layernorm", "rmsnorm", "offset"])
+@pytest.mark.parametrize("hook_name", ["hook_scale", "hook_normalized"])
+@pytest.mark.parametrize("mode", ["grad", "no_grad", "inference"])
+@pytest.mark.parametrize("return_none", [False, True])
+def test_native_inplace_edits_match_replacements(kind, hook_name, mode, return_none):
+    bridge = _make_bridge(native=True, rms=kind == "rmsnorm", offset=kind == "offset")
+    x = torch.randn(2, 3, 16, requires_grad=mode == "grad")
+    context = {
+        "grad": torch.enable_grad,
+        "no_grad": torch.no_grad,
+        "inference": torch.inference_mode,
+    }[mode]
+    point = getattr(bridge, hook_name)
+    with context():
+        point.add_hook(lambda value, hook: value * 2)
+        with pytest.warns(UserWarning, match="reconstructed from the hooked values"):
+            expected = bridge(x)
+        point.remove_hooks()
+
+        def edit(value, hook):
+            value.mul_(2)
+            return None if return_none else value
+
+        point.add_hook(edit)
+        with pytest.warns(UserWarning, match="reconstructed from the hooked values"):
+            actual = bridge(x)
+        point.remove_hooks()
+        torch.testing.assert_close(actual, expected)
+    if mode == "grad":
+        parameters = (x, *bridge.original_component.parameters())
+        probe = torch.randn_like(actual)
+        expected_grads = torch.autograd.grad(expected, parameters, probe)
+        actual_grads = torch.autograd.grad(actual, parameters, probe)
+        for actual_grad, expected_grad in zip(actual_grads, expected_grads):
+            torch.testing.assert_close(actual_grad, expected_grad)
+
+
+@pytest.mark.parametrize("mode", ["grad", "no_grad", "inference"])
+@pytest.mark.parametrize("return_input", [False, True])
+def test_native_observers_keep_exact_output_and_gradients(mode, return_input):
+    bridge = _make_bridge(native=True)
+    x = torch.randn(2, 3, 16, requires_grad=mode == "grad")
+    context = {
+        "grad": torch.enable_grad,
+        "no_grad": torch.no_grad,
+        "inference": torch.inference_mode,
+    }[mode]
+    with context():
+        expected = bridge.original_component(x)
+
+        def observe(value, hook):
+            return value if return_input else None
+
+        bridge.hook_scale.add_hook(observe)
+        bridge.hook_normalized.add_hook(observe)
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            actual = bridge(x)
+        bridge.hook_scale.remove_hooks()
+        bridge.hook_normalized.remove_hooks()
+        assert torch.equal(actual, expected)
+    if mode == "grad":
+        parameters = (x, *bridge.original_component.parameters())
+        probe = torch.randn_like(actual)
+        expected_grads = torch.autograd.grad(expected, parameters, probe)
+        actual_grads = torch.autograd.grad(actual, parameters, probe)
+        for actual_grad, expected_grad in zip(actual_grads, expected_grads):
+            assert torch.equal(actual_grad, expected_grad)
+
+
+@pytest.mark.parametrize("invalid", [float("nan"), float("inf")])
+def test_native_inference_observers_do_not_treat_nonfinite_values_as_edits(invalid):
+    bridge = _make_bridge(native=True)
+    with torch.inference_mode():
+        x = torch.randn(2, 3, 16)
+        x[..., 0] = invalid
+        expected = bridge.original_component(x)
+        bridge.hook_scale.add_hook(lambda value, hook: None)
+        bridge.hook_normalized.add_hook(lambda value, hook: None)
+        with warnings.catch_warnings():
+            warnings.simplefilter("error")
+            actual = bridge(x)
+        bridge.hook_scale.remove_hooks()
+        bridge.hook_normalized.remove_hooks()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0, equal_nan=True)

transformer_lens/model_bridge/generalized_components/normalization.py

@@ -25,6 +25,22 @@
 )
 
 
+def _run_intermediate_hook(point: HookPoint, value: torch.Tensor) -> tuple[torch.Tensor, bool]:
+    """Detect replacement and tensor writes while protecting autograd's saved intermediates."""
+    if not point.fwd_hooks:
+        result = point(value)
+        return result, result is not value
+    hook_input = value.clone()
+    # Inference tensors have no version counter; compare a snapshot only on that path.
+    before = hook_input.clone() if torch.is_inference(hook_input) else hook_input._version
+    result = point(hook_input)
+    if isinstance(before, torch.Tensor):
+        mutated = not torch.allclose(hook_input, before, rtol=0, atol=0, equal_nan=True)
+    else:
+        mutated = hook_input._version != before
+    return result, result is not hook_input or mutated
+
+
 class NormalizationBridge(GeneralizedComponent):
     """Normalization bridge that wraps transformer normalization layers but implements the calculation from scratch.
 
@@ -159,8 +175,8 @@ def _hf_autograd_forward_with_hooks(self, x: torch.Tensor) -> torch.Tensor:
 
         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
+        edits ``hook_scale`` / ``hook_normalized`` (including in-place edits), the output is
+        reconstructed from the hooked values so the edit propagates. With backward hooks, 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.
 
@@ -212,17 +228,16 @@ 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
-            hooked_scale = self.hook_scale(scale)
-            if hooked_scale is not scale:
+            hooked_scale, scale_edited = _run_intermediate_hook(self.hook_scale, scale)
+            if scale_edited:
                 # 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)
+            hooked_normalized, normalized_edited = _run_intermediate_hook(
+                self.hook_normalized, x_normalized
+            )
         input_dtype = x.dtype
-        # 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:
+        if not scale_edited and not normalized_edited:
             result = self.original_component(x)
             if result.dtype != input_dtype:
                 result = result.to(input_dtype)