diff --git a/tests/integration/model_bridge/test_native_norm_inplace_hooks.py b/tests/integration/model_bridge/test_native_norm_inplace_hooks.py new file mode 100644 index 000000000..bd21ffdc8 --- /dev/null +++ b/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) 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 758cf492a..663909075 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 @@ -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) diff --git a/transformer_lens/model_bridge/generalized_components/normalization.py b/transformer_lens/model_bridge/generalized_components/normalization.py index 401332bae..5f14a6517 100644 --- a/transformer_lens/model_bridge/generalized_components/normalization.py +++ b/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)