diff --git a/tests/integration/test_relevance_lens.py b/tests/integration/test_relevance_lens.py new file mode 100644 index 0000000000..e790c7d17f --- /dev/null +++ b/tests/integration/test_relevance_lens.py @@ -0,0 +1,222 @@ +"""End-to-end relevance-rule backend on tiny, offline HF fixtures. + +Per-component unit tests wrap one bridge component directly in a hand-built +single-block harness, which only ever registers a canonical mount name +(``ln1``, ``mlp``) once. That harness cannot see what happens on a real, +fully assembled ``TransformerBridge``, where the same component is *also* +reachable through the raw HF module tree under its own HF attribute name +(for example ``blocks.0._original_component.input_layernorm``). These tests +build a tiny random Qwen2 (the opaque gated-MLP native-forward path) and a tiny +random Phi-3 (``JointGateUpMLPBridge``'s already-reconstructed forward) +fully offline -- random weights from a programmatic HF config, no network +access, no checkpoint download -- and exercise all three rules together on +the real block stack: forward stays bit-identical, canonical mounts are +actually found and installed, and the gradient each rule-active node passes +upstream matches its closed-form VJP given the gradient it actually +received downstream in the real graph. +""" + +import pytest +import torch +from transformers import AutoConfig, AutoModelForCausalLM + +from transformer_lens.model_bridge._relevance_rules import ( + RelevanceRules, + half_rule, + identity_rule, + ln_rule_grad, + use_relevance_rules, +) +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.generalized_components.gated_mlp import ( + resolve_activation_fn, +) +from transformer_lens.model_bridge.sources import build_bridge_config_from_hf +from transformer_lens.model_bridge.supported_architectures.phi3 import ( + Phi3ArchitectureAdapter, +) +from transformer_lens.model_bridge.supported_architectures.qwen2 import ( + Qwen2ArchitectureAdapter, +) + +TOKENS = torch.tensor([[1, 5, 7, 42, 9]]) +N_LAYERS = 2 +TINY_DIMS = dict( + vocab_size=97, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=N_LAYERS, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=64, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, +) + + +class _MockTokenizer: + """Stand-in to satisfy TransformerBridge(tokenizer=...).""" + + +def _build_qwen2_bridge() -> TransformerBridge: + hf_config = AutoConfig.for_model("qwen2", **TINY_DIMS) + torch.manual_seed(0) + hf_model = AutoModelForCausalLM.from_config(hf_config, attn_implementation="eager").eval() + bridge_config = build_bridge_config_from_hf( + hf_model.config, "Qwen2ForCausalLM", "qwen2-tiny", torch.float32 + ) + adapter = Qwen2ArchitectureAdapter(bridge_config) + return TransformerBridge(model=hf_model, adapter=adapter, tokenizer=_MockTokenizer()) + + +def _build_phi3_bridge() -> TransformerBridge: + hf_config = AutoConfig.for_model("phi3", **TINY_DIMS) + torch.manual_seed(0) + hf_model = AutoModelForCausalLM.from_config(hf_config, attn_implementation="eager").eval() + bridge_config = build_bridge_config_from_hf( + hf_model.config, "Phi3ForCausalLM", "phi3-tiny", torch.float32 + ) + adapter = Phi3ArchitectureAdapter(bridge_config) + return TransformerBridge(model=hf_model, adapter=adapter, tokenizer=_MockTokenizer()) + + +# Qwen2 exercises GatedMLPBridge's opaque native-forward path (the primary path +# per real usage), where the rules attach to the live HF activation and down +# projection; Phi-3 exercises JointGateUpMLPBridge's already-reconstructed +# forward, isolating rule bugs from Bridge-integration bugs on the +# fused-projection family. +FIXTURE_BUILDERS = { + "qwen2": _build_qwen2_bridge, + "phi3": _build_phi3_bridge, +} + + +@pytest.fixture(scope="module", params=sorted(FIXTURE_BUILDERS), ids=sorted(FIXTURE_BUILDERS)) +def bridge(request: pytest.FixtureRequest) -> TransformerBridge: + return FIXTURE_BUILDERS[request.param]() + + +def _expected_canonical_mounts() -> set[str]: + ln_mounts = {f"blocks.{i}.ln1" for i in range(N_LAYERS)} | { + f"blocks.{i}.ln2" for i in range(N_LAYERS) + } + mlp_mounts = {f"blocks.{i}.mlp" for i in range(N_LAYERS)} + return ln_mounts | mlp_mounts + + +class TestForwardIdentity: + def test_active_forward_matches_baseline_under_all_three_rules( + self, bridge: TransformerBridge + ) -> None: + with torch.no_grad(): + baseline = bridge(TOKENS) + with use_relevance_rules( + bridge, RelevanceRules(normalization=True, activation=True, multiplicative_gate=True) + ): + with torch.no_grad(): + active = bridge(TOKENS) + assert torch.equal(active, baseline) + + +class TestCoverage: + def test_every_canonical_mount_is_installed_and_none_skipped( + self, bridge: TransformerBridge + ) -> None: + with use_relevance_rules( + bridge, RelevanceRules(normalization=True, activation=True, multiplicative_gate=True) + ) as coverage: + pass + assert set(coverage.installed) == _expected_canonical_mounts() + assert coverage.skipped == () + + +class TestGradientMatchesClosedFormOracle: + """Each rule-active node's local VJP, checked against the gradient it + actually receives in the real graph -- not a hand-rederived whole-model + oracle. The rule Functions compute a purely local closed-form VJP (already + proven against analytic oracles in the primitive and single-component + tests), so the only thing a real multi-block graph can newly break is the + wiring: the wrong node's weights, a disconnected graph path, or a mount + that silently never resolves so its rule never activates. Tapping + hook_in/hook_out -- outside the LN-rule's fail-closed + hook_scale/hook_normalized guard -- exposes exactly the gradient each node + passes upstream and the gradient it receives from downstream, with no + need to reconstruct attention or the rest of the stack by hand. + """ + + def test_ln_rule_and_gated_mlp_rule_match_local_oracles( + self, bridge: TransformerBridge + ) -> None: + ln1 = bridge.blocks[0].ln1 + mlp = bridge.blocks[0].mlp + captured: dict[str, torch.Tensor] = {} + + def _capture(key: str): + def _hook(tensor: torch.Tensor, hook=None) -> None: + captured[key] = tensor.detach().clone() + + return _hook + + ln1.hook_in.add_hook(_capture("ln_x")) + ln1.hook_in.add_hook(_capture("ln_grad_in"), dir="bwd") + ln1.hook_out.add_hook(_capture("ln_grad_out"), dir="bwd") + mlp.hook_in.add_hook(_capture("mlp_x")) + mlp.hook_in.add_hook(_capture("mlp_grad_in"), dir="bwd") + mlp.hook_out.add_hook(_capture("mlp_grad_out"), dir="bwd") + ln1_weight_grad = None + try: + with use_relevance_rules( + bridge, + RelevanceRules(normalization=True, activation=True, multiplicative_gate=True), + ): + logits = bridge(TOKENS) + logits.sum().backward() + ln1_weight_grad = ln1.weight.grad.clone() + finally: + ln1.hook_in.remove_hooks(dir="both") + ln1.hook_out.remove_hooks(dir="both") + mlp.hook_in.remove_hooks(dir="both") + mlp.hook_out.remove_hooks(dir="both") + for parameter in bridge.parameters(): + parameter.grad = None + + eps = getattr(ln1.original_component, "variance_epsilon", 1e-6) + weight = ln1.weight.detach() + denom = (captured["ln_x"].pow(2).mean(-1, keepdim=True) + eps).sqrt() + expected_ln_grad_in = ln_rule_grad(captured["ln_grad_out"] * weight, denom) + torch.testing.assert_close( + captured["ln_grad_in"], expected_ln_grad_in, atol=1e-5, rtol=1e-4 + ) + # weight keeps its ordinary gradient: the rule only redefines the x-path + # VJP, not d(output)/d(weight), given the same grad_out the rule-active + # forward actually produced. + expected_weight_grad = (captured["ln_grad_out"] * (captured["ln_x"] / denom)).sum( + dim=(0, 1) + ) + torch.testing.assert_close(ln1_weight_grad, expected_weight_grad, atol=1e-5, rtol=1e-4) + + x = captured["mlp_x"].detach().requires_grad_(True) + w_gate, w_in, w_out = mlp.W_gate.detach(), mlp.W_in.detach(), mlp.W_out.detach() + b_gate = getattr(mlp.gate, "bias", None) + b_in = getattr(getattr(mlp, "in"), "bias", None) + b_out = getattr(mlp.out, "bias", None) + act_fn = resolve_activation_fn(mlp.config) + with torch.enable_grad(): + gate_output = x @ w_gate + if b_gate is not None: + gate_output = gate_output + b_gate + up_output = x @ w_in + if b_in is not None: + up_output = up_output + b_in + activated = identity_rule(gate_output, act_fn) + gated = half_rule(activated, up_output) + down = gated @ w_out + if b_out is not None: + down = down + b_out + (expected_mlp_grad_in,) = torch.autograd.grad( + down, x, grad_outputs=captured["mlp_grad_out"] + ) + torch.testing.assert_close( + captured["mlp_grad_in"], expected_mlp_grad_in, atol=1e-5, rtol=1e-4 + ) diff --git a/tests/unit/model_bridge/generalized_components/test_gated_mlp_relevance_rule.py b/tests/unit/model_bridge/generalized_components/test_gated_mlp_relevance_rule.py new file mode 100644 index 0000000000..0d2ab27d36 --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_gated_mlp_relevance_rule.py @@ -0,0 +1,427 @@ +"""Identity-/Half-rule integration on GatedMLPBridge's opaque native-forward path. + +The raw (non-fused) gated-MLP path keeps the HF module's own forward intact and +installs the rules on its live submodules for a ``use_relevance_rules`` scope: the +Identity-rule by swapping the module's activation callable, and the Half-rule by a +forward-pre-hook that halves the gradient entering the down projection (the gate*up +product). Because the real forward runs unchanged, anything it does beyond the core +``down(act(gate) * up)`` shape -- a post-product multiplier, activation sparsity, the +module's own hooks -- survives into the backward graph, and the rule VJP is taken +through that real forward rather than a reconstruction of the core shape. The rules +attach regardless of the projection weight class, so an opaque backing that is neither +``nn.Linear`` nor ``Conv1D`` is still covered. Only the relu-family Identity-rule +exclusion and the missing-activation-callable case remain unsupported. +""" + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.pytorch_utils import Conv1D + +from transformer_lens.model_bridge._relevance_rules import ( + RelevanceRules, + RelevanceRuleUnsupportedError, + half_rule, + identity_rule, + use_relevance_rules, +) +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) +from transformer_lens.model_bridge.generalized_components.gated_mlp import ( + GatedMLPBridge, +) +from transformer_lens.model_bridge.generalized_components.linear import LinearBridge + + +class _Cfg: + def __init__(self, hidden_act: str = "silu"): + self.hidden_act = hidden_act + + +class _ReluSquared(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.relu(x).square() + + +_ACTIVATIONS = {"silu": nn.SiLU, "relu2": _ReluSquared} + + +class _OpaqueProj(nn.Module): + """Weight-backed projection that is neither nn.Linear nor Conv1D. + + Stands in for a backing class the retired weight-orientation allowlist would + have refused; the rules now attach without reading its weights at all. + """ + + def __init__(self, d_in: int, d_out: int, bias: bool = True): + super().__init__() + self.weight = nn.Parameter(torch.randn(d_out, d_in)) + self.bias = nn.Parameter(torch.randn(d_out)) if bias else None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + out = x @ self.weight.T + return out if self.bias is None else out + self.bias + + +class _TinyGatedMLP(nn.Module): + """Mirrors the Qwen2/Llama/Gemma gated-MLP: one opaque call over its own submodules. + + The activation is an ``nn.Module`` attribute the forward calls, matching the + ``ACT2FN`` shape the Identity-rule wrap targets. + """ + + def __init__(self, gate_proj, up_proj, down_proj, hidden_act: str = "silu"): + super().__init__() + self.gate_proj = gate_proj + self.up_proj = up_proj + self.down_proj = down_proj + self.act_fn = _ACTIVATIONS[hidden_act]() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class _MultiplierGatedMLP(_TinyGatedMLP): + """A gated MLP whose forward does strictly more than ``down(act(gate) * up)``. + + The constant post-product multiplier stands in for the Falcon-H1 / Gemma3n + families whose native forward applies extra scaling. A rule VJP taken through a + reconstruction of only the core gated shape would omit the multiplier and + disagree with the VJP through this real forward. + """ + + def __init__(self, gate_proj, up_proj, down_proj, multiplier: float = 1.7): + super().__init__(gate_proj, up_proj, down_proj) + self.multiplier = multiplier + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return super().forward(x) * self.multiplier + + +class _Block(nn.Module): + """Mounts a gated-MLP bridge at the canonical mlp position.""" + + def __init__(self, mlp: GeneralizedComponent): + super().__init__() + self.mlp = mlp + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.mlp(x) + + +def _make_projections(backing_class: str, d_model: int = 4, d_mlp: int = 8, bias: bool = True): + torch.manual_seed(0) + if backing_class == "nn.Linear": + return ( + nn.Linear(d_model, d_mlp, bias=bias), + nn.Linear(d_model, d_mlp, bias=bias), + nn.Linear(d_mlp, d_model, bias=bias), + ) + if backing_class == "Conv1D": + return ( + Conv1D(d_mlp, d_model), + Conv1D(d_mlp, d_model), + Conv1D(d_model, d_mlp), + ) + return ( + _OpaqueProj(d_model, d_mlp, bias=bias), + _OpaqueProj(d_model, d_mlp, bias=bias), + _OpaqueProj(d_mlp, d_model, bias=bias), + ) + + +def _wire_bridge(hf_mlp: nn.Module, config: _Cfg) -> GatedMLPBridge: + bridge = GatedMLPBridge(name="mlp", config=config) + gate_bridge = LinearBridge(name="gate_proj") + in_bridge = LinearBridge(name="up_proj") + out_bridge = LinearBridge(name="down_proj") + bridge.add_module("gate", gate_bridge) + bridge.add_module("in", in_bridge) + bridge.add_module("out", out_bridge) + bridge.set_original_component(hf_mlp) + gate_bridge.set_original_component(hf_mlp.gate_proj) + in_bridge.set_original_component(hf_mlp.up_proj) + out_bridge.set_original_component(hf_mlp.down_proj) + return bridge + + +def _make_bridge( + backing_class: str, bias: bool = True, hidden_act: str = "silu" +) -> tuple[_Block, _TinyGatedMLP]: + gate_proj, up_proj, down_proj = _make_projections(backing_class, bias=bias) + hf_mlp = _TinyGatedMLP(gate_proj, up_proj, down_proj, hidden_act=hidden_act) + bridge = _wire_bridge(hf_mlp, _Cfg(hidden_act)) + return _Block(bridge), hf_mlp + + +def _make_multiplier_bridge(multiplier: float = 1.7) -> tuple[_Block, _MultiplierGatedMLP]: + gate_proj, up_proj, down_proj = _make_projections("nn.Linear", bias=True) + hf_mlp = _MultiplierGatedMLP(gate_proj, up_proj, down_proj, multiplier=multiplier) + bridge = _wire_bridge(hf_mlp, _Cfg("silu")) + return _Block(bridge), hf_mlp + + +def _oracle_grads(hf_mlp, x, activation_active, gate_active, multiplier: float = 1.0): + # hf_mlp's parameters are shared with the bridge under test, so a prior backward + # already left gradients on them; reset first or this second backward would + # accumulate on top instead of producing an independently comparable oracle. The + # oracle runs the module's real submodules -- including the multiplier -- so it is + # the VJP through the actual forward, not through the core gated shape alone. + for p in hf_mlp.parameters(): + p.grad = None + x_oracle = x.detach().clone().requires_grad_(True) + gate_output = hf_mlp.gate_proj(x_oracle) + up_output = hf_mlp.up_proj(x_oracle) + activated = ( + identity_rule(gate_output, hf_mlp.act_fn) + if activation_active + else hf_mlp.act_fn(gate_output) + ) + gated = half_rule(activated, up_output) if gate_active else activated * up_output + down = hf_mlp.down_proj(gated) * multiplier + down.sum().backward() + grads = {n: p.grad.clone() for n, p in hf_mlp.named_parameters()} + return x_oracle.grad.clone(), grads + + +BACKING_CLASSES = ["nn.Linear", "Conv1D", "opaque"] + + +class TestGatedMLPRelevanceRuleCapability: + @pytest.mark.parametrize("backing_class", BACKING_CLASSES) + def test_capable_of_both_kinds_regardless_of_weight_backing(self, backing_class): + block, _ = _make_bridge(backing_class) + assert set(block.mlp._relevance_rule_kinds) == {"activation", "multiplicative_gate"} + + def test_identity_rule_unsupported_for_relu_squared_activation(self): + block, _ = _make_bridge("nn.Linear", hidden_act="relu2") + assert block.mlp._relevance_rule_kinds == ("multiplicative_gate",) + + with pytest.raises(RelevanceRuleUnsupportedError, match="mlp"): + with use_relevance_rules(block, RelevanceRules(activation=True)): + pass + + def test_activation_unsupported_when_no_activation_callable_is_exposed(self): + # A gated MLP whose forward uses a bare function has no callable activation + # attribute for the Identity-rule to wrap, so requesting it must raise rather + # than install a silent no-op. The Half-rule needs no activation access and + # stays available. + class _FunctionalActMLP(nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(4, 8) + self.up_proj = nn.Linear(4, 8) + self.down_proj = nn.Linear(8, 4) + + def forward(self, x): + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + torch.manual_seed(0) + bridge = _wire_bridge(_FunctionalActMLP(), _Cfg("silu")) + block = _Block(bridge) + + assert bridge._relevance_rule_kinds == ("multiplicative_gate",) + with pytest.raises(RelevanceRuleUnsupportedError, match="mlp"): + with use_relevance_rules(block, RelevanceRules(activation=True)): + pass + + def test_half_rule_remains_available_under_relu_squared_activation(self): + block, _ = _make_bridge("nn.Linear", hidden_act="relu2") + x = torch.randn(2, 4) + baseline = block(x) + with use_relevance_rules(block, RelevanceRules(multiplicative_gate=True)) as coverage: + assert coverage.installed == ("mlp",) + assert torch.equal(block(x), baseline) + + +class TestGatedMLPRelevanceRuleForwardIdentity: + @pytest.mark.parametrize("backing_class", BACKING_CLASSES) + def test_forward_identical_while_rule_active(self, backing_class): + block, _ = _make_bridge(backing_class) + x = torch.randn(3, 4) + baseline = block(x) + with use_relevance_rules(block, RelevanceRules(activation=True, multiplicative_gate=True)): + active = block(x) + assert torch.equal(active, baseline) + + def test_forward_identical_while_rule_active_with_extra_forward_ops(self): + block, _ = _make_multiplier_bridge() + x = torch.randn(3, 4) + baseline = block(x) + with use_relevance_rules(block, RelevanceRules(activation=True, multiplicative_gate=True)): + active = block(x) + assert torch.equal(active, baseline) + + @pytest.mark.parametrize("backing_class", BACKING_CLASSES) + def test_forward_unchanged_when_rule_inactive(self, backing_class): + block, _ = _make_bridge(backing_class) + x = torch.randn(3, 4) + assert torch.equal(block(x), block(x)) + + def test_activation_callable_restored_after_scope(self): + block, hf_mlp = _make_bridge("nn.Linear") + original_act = hf_mlp._modules["act_fn"] + with use_relevance_rules(block, RelevanceRules(activation=True)): + assert hf_mlp._modules["act_fn"] is not original_act + assert hf_mlp._modules["act_fn"] is original_act + + +class TestGatedMLPRelevanceRuleVJP: + @pytest.mark.parametrize("backing_class", BACKING_CLASSES) + def test_both_rules_active_matches_real_forward_oracle(self, backing_class): + block, hf_mlp = _make_bridge(backing_class) + x = torch.randn(3, 4, requires_grad=True) + + with use_relevance_rules(block, RelevanceRules(activation=True, multiplicative_gate=True)): + out = block(x) + out.sum().backward() + + grad_x = x.grad.clone() + grads = {n: p.grad.clone() for n, p in hf_mlp.named_parameters()} + + expected_grad_x, expected_grads = self._expected(hf_mlp, x, True, True) + torch.testing.assert_close(grad_x, expected_grad_x, atol=1e-5, rtol=1e-5) + for name, grad in grads.items(): + # The down projection's own weight gradient stays ordinary: the Half-rule + # only halves the gradient reaching the product, not the parameter grads. + torch.testing.assert_close(grad, expected_grads[name], atol=1e-5, rtol=1e-5) + + def test_both_rules_match_vjp_through_extra_forward_ops(self): + block, hf_mlp = _make_multiplier_bridge() + x = torch.randn(3, 4, requires_grad=True) + + with use_relevance_rules(block, RelevanceRules(activation=True, multiplicative_gate=True)): + out = block(x) + out.sum().backward() + + grad_x = x.grad.clone() + grads = {n: p.grad.clone() for n, p in hf_mlp.named_parameters()} + + expected_grad_x, expected_grads = _oracle_grads( + hf_mlp, x, activation_active=True, gate_active=True, multiplier=hf_mlp.multiplier + ) + torch.testing.assert_close(grad_x, expected_grad_x, atol=1e-5, rtol=1e-5) + for name, grad in grads.items(): + torch.testing.assert_close(grad, expected_grads[name], atol=1e-5, rtol=1e-5) + + @pytest.mark.parametrize("backing_class", BACKING_CLASSES) + def test_activation_only_leaves_gate_split_ordinary(self, backing_class): + block, hf_mlp = _make_bridge(backing_class) + x = torch.randn(3, 4, requires_grad=True) + + with use_relevance_rules(block, RelevanceRules(activation=True)): + out = block(x) + out.sum().backward() + + grad_x = x.grad.clone() + expected_grad_x, _ = self._expected(hf_mlp, x, True, False) + torch.testing.assert_close(grad_x, expected_grad_x, atol=1e-5, rtol=1e-5) + + @pytest.mark.parametrize("backing_class", BACKING_CLASSES) + def test_multiplicative_gate_only_leaves_activation_ordinary(self, backing_class): + block, hf_mlp = _make_bridge(backing_class) + x = torch.randn(3, 4, requires_grad=True) + + with use_relevance_rules(block, RelevanceRules(multiplicative_gate=True)): + out = block(x) + out.sum().backward() + + grad_x = x.grad.clone() + expected_grad_x, _ = self._expected(hf_mlp, x, False, True) + torch.testing.assert_close(grad_x, expected_grad_x, atol=1e-5, rtol=1e-5) + + def test_rule_inactive_gradients_are_ordinary(self): + block, hf_mlp = _make_bridge("nn.Linear") + x = torch.randn(3, 4, requires_grad=True) + + out = block(x) + out.sum().backward() + grad_x = x.grad.clone() + grads = {n: p.grad.clone() for n, p in hf_mlp.named_parameters()} + + for p in hf_mlp.parameters(): + p.grad = None + x_plain = x.detach().clone().requires_grad_(True) + hf_mlp(x_plain).sum().backward() + + torch.testing.assert_close(grad_x, x_plain.grad) + for name, param in hf_mlp.named_parameters(): + torch.testing.assert_close(grads[name], param.grad) + + def test_bias_free_projections_are_handled(self): + block, hf_mlp = _make_bridge("nn.Linear", bias=False) + x = torch.randn(3, 4, requires_grad=True) + + with use_relevance_rules(block, RelevanceRules(activation=True, multiplicative_gate=True)): + out = block(x) + out.sum().backward() + + grad_x = x.grad.clone() + expected_grad_x, expected_grads = self._expected(hf_mlp, x, True, True) + torch.testing.assert_close(grad_x, expected_grad_x, atol=1e-5, rtol=1e-5) + for name, param in hf_mlp.named_parameters(): + torch.testing.assert_close(param.grad, expected_grads[name], atol=1e-5, rtol=1e-5) + + @staticmethod + def _expected(hf_mlp, x, activation_active, gate_active): + return _oracle_grads(hf_mlp, x, activation_active, gate_active, multiplier=1.0) + + +class TestGatedMLPRelevanceRulePreservesRealForward: + def test_modules_own_backward_hook_fires_during_rule_active_backward(self): + # The rules attach to the live module and leave its native forward in the + # graph, so a backward hook registered on an HF submodule still fires during + # a rule-active backward rather than being bypassed by a separate graph. + block, hf_mlp = _make_bridge("nn.Linear") + fired = {"count": 0} + + def _record(module, grad_input, grad_output): + fired["count"] += 1 + + handle = hf_mlp.down_proj.register_full_backward_hook(_record) + x = torch.randn(3, 4, requires_grad=True) + try: + with use_relevance_rules( + block, RelevanceRules(activation=True, multiplicative_gate=True) + ): + block(x).sum().backward() + finally: + handle.remove() + + assert fired["count"] > 0 + + +class TestGatedMLPRelevanceRuleCompatibilityMode: + """The processed-weights (compatibility) path reconstructs the forward from folded + weights, so it applies the rules inline; before this it skipped them entirely.""" + + def _make_compat_bridge(self, bias: bool = True): + block, hf_mlp = _make_bridge("nn.Linear", bias=bias) + bridge = block.mlp + bridge._use_processed_weights = True + # Plain (non-Parameter) tensors so the bridge's custom __getattr__ resolves + # them; nn.Linear stores weight as [out, in], the layout functional linear + # expects, so the folded-weight forward reproduces the native projections. + bridge._processed_W_gate = hf_mlp.gate_proj.weight.detach() + bridge._processed_b_gate = None if not bias else hf_mlp.gate_proj.bias.detach() + bridge._processed_W_in = hf_mlp.up_proj.weight.detach() + bridge._processed_b_in = None if not bias else hf_mlp.up_proj.bias.detach() + bridge._processed_W_out = hf_mlp.down_proj.weight.detach() + bridge._processed_b_out = None if not bias else hf_mlp.down_proj.bias.detach() + return block, hf_mlp + + def test_compat_forward_identical_and_rule_vjp_matches_oracle(self): + block, hf_mlp = self._make_compat_bridge() + x = torch.randn(3, 4, requires_grad=True) + + baseline = block(x.detach()) + with use_relevance_rules(block, RelevanceRules(activation=True, multiplicative_gate=True)): + active = block(x) + assert torch.equal(active.detach(), baseline) + + active.sum().backward() + grad_x = x.grad.clone() + expected_grad_x, _ = _oracle_grads(hf_mlp, x, True, True) + torch.testing.assert_close(grad_x, expected_grad_x, atol=1e-5, rtol=1e-5) diff --git a/tests/unit/model_bridge/generalized_components/test_joint_gate_up_relevance_rule.py b/tests/unit/model_bridge/generalized_components/test_joint_gate_up_relevance_rule.py new file mode 100644 index 0000000000..f8bee5e693 --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_joint_gate_up_relevance_rule.py @@ -0,0 +1,164 @@ +"""Identity-/Half-rule integration on JointGateUpMLPBridge's reconstructed forward. + +Unlike the raw GatedMLPBridge path (opaque native forward, tested separately), +JointGateUpMLPBridge already reconstructs its forward in Python as +``act_fn(gate_output) * up_output`` through separate gate/up LinearBridge +submodules, so the rules attach directly at that multiplication inline, off the +same boolean flags -- no live-submodule hook is installed here. This covers that +the reconstructed +forward is unaffected by an inactive rule, that both rules apply correctly +together and independently, and that gradients match the oracle produced by the +already-tested Identity-/Half-rule primitives directly. +""" + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from transformer_lens.model_bridge._relevance_rules import ( + RelevanceRules, + half_rule, + identity_rule, + use_relevance_rules, +) +from transformer_lens.model_bridge.generalized_components.joint_gate_up_mlp import ( + JointGateUpMLPBridge, +) +from transformer_lens.model_bridge.generalized_components.linear import LinearBridge + + +class _Cfg: + hidden_act = "silu" + + +class _TinyPhi3MLP(nn.Module): + """Mirrors Phi-3/GLM's fused gate_up_proj structure.""" + + def __init__(self, d_model: int = 4, d_mlp: int = 8, bias: bool = False): + super().__init__() + self.gate_up_proj = nn.Linear(d_model, 2 * d_mlp, bias=bias) + self.down_proj = nn.Linear(d_mlp, d_model, bias=bias) + self.activation_fn = nn.SiLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up = self.gate_up_proj(x) + gate, up = gate_up.chunk(2, dim=-1) + return self.down_proj(self.activation_fn(gate) * up) + + +class _Block(nn.Module): + """Mounts a joint gate-up bridge at the canonical mlp position.""" + + def __init__(self, mlp: JointGateUpMLPBridge): + super().__init__() + self.mlp = mlp + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.mlp(x) + + +def _make_bridge(bias: bool = False) -> tuple[_Block, _TinyPhi3MLP]: + torch.manual_seed(0) + hf_mlp = _TinyPhi3MLP(bias=bias) + bridge = JointGateUpMLPBridge(name="mlp", config=_Cfg(), submodules={}) + out_bridge = LinearBridge(name="down_proj") + bridge.add_module("out", out_bridge) + bridge.set_original_component(hf_mlp) + out_bridge.set_original_component(hf_mlp.down_proj) + return _Block(bridge), hf_mlp + + +def _oracle_grads(hf_mlp, gate_proj, up_proj, x, activation_active: bool, gate_active: bool): + for p in hf_mlp.parameters(): + p.grad = None + x_oracle = x.detach().clone().requires_grad_(True) + gate_output = F.linear(x_oracle, gate_proj.weight, gate_proj.bias) + up_output = F.linear(x_oracle, up_proj.weight, up_proj.bias) + activated = identity_rule(gate_output, F.silu) if activation_active else F.silu(gate_output) + gated = half_rule(activated, up_output) if gate_active else activated * up_output + down = hf_mlp.down_proj(gated) + down.sum().backward() + return x_oracle.grad.clone() + + +class TestJointGateUpRelevanceRuleCapability: + def test_capable_of_both_kinds(self): + block, _ = _make_bridge() + assert set(block.mlp._relevance_rule_kinds) == {"activation", "multiplicative_gate"} + + +class TestJointGateUpRelevanceRuleForwardIdentity: + def test_forward_identical_while_rule_active(self): + block, _ = _make_bridge() + x = torch.randn(3, 4) + baseline = block(x) + with use_relevance_rules(block, RelevanceRules(activation=True, multiplicative_gate=True)): + active = block(x) + assert torch.equal(active, baseline) + + def test_forward_unchanged_when_rule_inactive(self): + block, _ = _make_bridge() + x = torch.randn(3, 4) + before = block(x) + after = block(x) + assert torch.equal(before, after) + + +class TestJointGateUpRelevanceRuleVJP: + @pytest.mark.parametrize( + ("activation_active", "gate_active"), + [(True, True), (True, False), (False, True)], + ) + def test_matches_manually_composed_oracle(self, activation_active, gate_active): + block, hf_mlp = _make_bridge() + x = torch.randn(3, 4, requires_grad=True) + + with use_relevance_rules( + block, + RelevanceRules( + activation=activation_active, + multiplicative_gate=gate_active, + ), + ): + out = block(x) + out.sum().backward() + + grad_x = x.grad.clone() + + gate_proj = block.mlp.gate.original_component + up_proj = getattr(block.mlp, "in").original_component + expected_grad_x = _oracle_grads( + hf_mlp, gate_proj, up_proj, x, activation_active, gate_active + ) + torch.testing.assert_close(grad_x, expected_grad_x, atol=1e-5, rtol=1e-5) + + def test_rule_inactive_gradients_are_ordinary(self): + block, hf_mlp = _make_bridge() + x = torch.randn(3, 4, requires_grad=True) + + out = block(x) + out.sum().backward() + grad_x = x.grad.clone() + + for p in hf_mlp.parameters(): + p.grad = None + x_plain = x.detach().clone().requires_grad_(True) + plain = hf_mlp(x_plain) + plain.sum().backward() + + torch.testing.assert_close(grad_x, x_plain.grad) + + def test_bias_free_and_biased_projections_are_both_handled(self): + block, hf_mlp = _make_bridge(bias=True) + x = torch.randn(3, 4, requires_grad=True) + + with use_relevance_rules(block, RelevanceRules(activation=True, multiplicative_gate=True)): + out = block(x) + out.sum().backward() + + grad_x = x.grad.clone() + gate_proj = block.mlp.gate.original_component + up_proj = getattr(block.mlp, "in").original_component + expected_grad_x = _oracle_grads(hf_mlp, gate_proj, up_proj, x, True, True) + torch.testing.assert_close(grad_x, expected_grad_x, atol=1e-5, rtol=1e-5) diff --git a/tests/unit/model_bridge/generalized_components/test_normalization_relevance_rule.py b/tests/unit/model_bridge/generalized_components/test_normalization_relevance_rule.py new file mode 100644 index 0000000000..bfd6dcd374 --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_normalization_relevance_rule.py @@ -0,0 +1,415 @@ +"""LN-rule integration on NormalizationBridge's native-autograd path. + +Covers the commit-5 contract: the rule-wrapped native forward is bit-identical to +today's native forward by construction (the wrapping calls ``original_component(x)`` +itself rather than reproducing its numerics), the backward follows the LN-rule +(denominator treated as constant) while weight/bias keep their ordinary gradient, +targeting is positional so an ln1/ln2 mount that never reaches the native-autograd +branch reports as skipped rather than silently leaving ordinary gradients in place, +and a hook that would otherwise silently fall back instead raises while the rule is +active. +""" + + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from transformer_lens.model_bridge._relevance_rules import ( + RelevanceRuleConflictError, + RelevanceRules, + use_relevance_rules, +) +from transformer_lens.model_bridge.generalized_components.normalization import ( + LayerNormPreBridge, + NormalizationBridge, + RMSNormPreBridge, +) + + +class _Cfg: + def __init__( + self, + uses_rms_norm: bool = False, + eps: float = 1e-5, + rmsnorm_uses_offset: bool = False, + layer_norm_folding: bool = False, + ): + self.uses_rms_norm = uses_rms_norm + self.eps = eps + self.rmsnorm_uses_offset = rmsnorm_uses_offset + self.layer_norm_folding = layer_norm_folding + + +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.randn(d) * 0.1 + 1.0) + 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) + + +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.randn(d) * 0.1) + 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) + + +class _TinyOlmoLayerNorm(nn.Module): + """Param-free centered LayerNorm mirroring OLMo's OlmoLayerNorm: no weight, no bias.""" + + def __init__(self, d: int, eps: float = 1e-5): + super().__init__() + self.normalized_shape = (d,) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.layer_norm(x, self.normalized_shape, None, None, self.eps) + + +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) + return layer + + +class _Block(nn.Module): + """Mounts a normalization bridge at the canonical ln1 position.""" + + def __init__(self, norm: NormalizationBridge): + super().__init__() + self.ln1 = norm + + +def _make_bridge( + native: bool, + rms: bool = False, + offset: bool = False, + d: int = 16, + layer_norm_folding: bool = False, +) -> NormalizationBridge: + layer: nn.Module + if offset: + layer = _TinyGemmaRMSNorm(d) + elif rms: + layer = _TinyRMSNorm(d) + else: + layer = _layernorm(d) + bridge = NormalizationBridge( + name="ln1", + config=_Cfg( + uses_rms_norm=rms or offset, + rmsnorm_uses_offset=offset, + layer_norm_folding=layer_norm_folding, + ), + use_native_layernorm_autograd=native, + ) + bridge.set_original_component(layer) + return bridge + + +def _denom_detached_oracle( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + uses_rms: bool, + offset: bool, + eps: float, +) -> torch.Tensor: + """Manual recompute with the denominator detached: the LN-rule's target VJP.""" + x_centered = x if uses_rms else x - x.mean(-1, keepdim=True) + denom = (x_centered.pow(2).mean(-1, keepdim=True) + eps).sqrt().detach() + w_eff = (1.0 + weight) if offset else weight + out = (x_centered / denom) * w_eff + if bias is not None: + out = out + bias + return out + + +class TestForwardIdentity: + @pytest.mark.parametrize("rms", [False, True], ids=["layernorm", "rmsnorm"]) + def test_active_forward_matches_inactive_forward(self, rms): + bridge = _make_bridge(native=True, rms=rms) + block = _Block(bridge) + x = torch.randn(2, 5, 16) + baseline = bridge(x) + with use_relevance_rules(block, RelevanceRules(normalization=True)): + active = bridge(x) + assert torch.equal(active, baseline) + + def test_active_forward_matches_original_component_directly(self): + bridge = _make_bridge(native=True) + block = _Block(bridge) + x = torch.randn(2, 5, 16) + with use_relevance_rules(block, RelevanceRules(normalization=True)): + active = bridge(x) + assert torch.equal(active, bridge.original_component(x)) + + +class TestVJPMatchesDetachedDenomOracle: + def test_layernorm(self): + bridge = _make_bridge(native=True) + block = _Block(bridge) + x = torch.randn(2, 5, 16, requires_grad=True) + with use_relevance_rules(block, RelevanceRules(normalization=True)): + y = bridge(x) + y.sum().backward() + grad_rule = x.grad.clone() + + x_oracle = x.detach().clone().requires_grad_(True) + y_oracle = _denom_detached_oracle( + x_oracle, + bridge.original_component.weight, + bridge.original_component.bias, + uses_rms=False, + offset=False, + eps=1e-5, + ) + y_oracle.sum().backward() + torch.testing.assert_close(grad_rule, x_oracle.grad) + + @pytest.mark.parametrize("offset", [False, True], ids=["plain_rms", "gemma_offset"]) + def test_rmsnorm(self, offset): + bridge = _make_bridge(native=True, rms=not offset, offset=offset) + block = _Block(bridge) + x = torch.randn(2, 5, 16, requires_grad=True) + with use_relevance_rules(block, RelevanceRules(normalization=True)): + y = bridge(x) + y.sum().backward() + grad_rule = x.grad.clone() + + x_oracle = x.detach().clone().requires_grad_(True) + y_oracle = _denom_detached_oracle( + x_oracle, + bridge.original_component.weight, + None, + uses_rms=True, + offset=offset, + eps=1e-5, + ) + y_oracle.sum().backward() + torch.testing.assert_close(grad_rule, x_oracle.grad) + + def test_rule_grad_disagrees_with_ordinary_autodiff(self): + """The whole point of the rule: it must actually change the gradient.""" + bridge = _make_bridge(native=True) + block = _Block(bridge) + x = torch.randn(2, 5, 16, requires_grad=True) + + with use_relevance_rules(block, RelevanceRules(normalization=True)): + y_rule = bridge(x) + (grad_rule,) = torch.autograd.grad(y_rule.sum(), x) + + x_plain = x.detach().clone().requires_grad_(True) + y_plain = bridge(x_plain) + (grad_plain,) = torch.autograd.grad(y_plain.sum(), x_plain) + assert not torch.allclose(grad_rule, grad_plain) + + +class TestParameterGradientsPreserved: + def test_weight_and_bias_receive_ordinary_gradient(self): + bridge = _make_bridge(native=True) + block = _Block(bridge) + x = torch.randn(2, 5, 16, requires_grad=True) + with use_relevance_rules(block, RelevanceRules(normalization=True)): + y = bridge(x) + y.sum().backward() + weight_grad = bridge.original_component.weight.grad + bias_grad = bridge.original_component.bias.grad + assert weight_grad is not None and torch.isfinite(weight_grad).all() + assert bias_grad is not None and torch.isfinite(bias_grad).all() + + # Ordinary gradient: d(output)/d(weight) = normalized value (pre-weight), + # independent of the rule's treatment of the x-path; d(output)/d(bias) = 1. + x_oracle = x.detach().clone() + normalized = _denom_detached_oracle( + x_oracle, + torch.ones_like(bridge.original_component.weight), + None, + uses_rms=False, + offset=False, + eps=1e-5, + ) + reduce_dims = tuple(range(normalized.dim() - 1)) + expected_weight_grad = normalized.sum(dim=reduce_dims) + expected_bias_grad = torch.full_like(bias_grad, x.shape[0] * x.shape[1]) + torch.testing.assert_close(weight_grad, expected_weight_grad, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(bias_grad, expected_bias_grad, atol=1e-4, rtol=1e-4) + + +class TestRuleInactiveRegression: + @pytest.mark.parametrize("native", [True, False], ids=["native_autograd", "python_norm"]) + def test_forward_byte_identical_to_ordinary_call(self, native): + bridge = _make_bridge(native=native) + x = torch.randn(2, 5, 16) + expected = bridge(x) + # No use_relevance_rules context at all: today's behavior, unconditionally. + actual = bridge(x) + assert torch.equal(actual, expected) + + def test_non_native_path_rule_request_is_skipped_and_forward_unaffected(self): + """A python-norm-path bridge (no native autograd, no folding) never reaches + the branch the LN-rule wraps, so it must be reported skipped rather than + silently leaving ordinary gradients in place under a claimed rule.""" + bridge = _make_bridge(native=False) + block = _Block(bridge) + x = torch.randn(2, 5, 16) + baseline = bridge(x) + with use_relevance_rules(block, RelevanceRules(normalization=True)) as coverage: + assert coverage.installed == () + assert coverage.skipped == ("ln1",) + active = bridge(x) + assert torch.equal(active, baseline) + + def test_layer_norm_folding_config_flag_makes_the_rule_installable(self): + """layer_norm_folding also dispatches through the native-autograd branch, + so the rule must be installable there even with use_native_layernorm_autograd + left False.""" + bridge = _make_bridge(native=False, layer_norm_folding=True) + block = _Block(bridge) + x = torch.randn(2, 5, 16) + baseline = bridge(x) + with use_relevance_rules(block, RelevanceRules(normalization=True)) as coverage: + assert coverage.installed == ("ln1",) + active = bridge(x) + assert torch.equal(active, baseline) + + @pytest.mark.parametrize( + "bridge_cls", [LayerNormPreBridge, RMSNormPreBridge], ids=["ln_pre", "rms_pre"] + ) + def test_param_free_pre_norm_is_skipped_and_forward_unaffected(self, bridge_cls): + """LNPre/RMSPre always take the python-norm path regardless of the native + flag, so they must never be reported installed.""" + bridge = bridge_cls(name="ln1", config=_Cfg()) + bridge.set_original_component(nn.LayerNorm(16, eps=1e-5, elementwise_affine=False)) + block = _Block(bridge) + x = torch.randn(2, 5, 16) + baseline = bridge(x) + with use_relevance_rules(block, RelevanceRules(normalization=True)) as coverage: + assert coverage.installed == () + assert coverage.skipped == ("ln1",) + active = bridge(x) + assert torch.equal(active, baseline) + + +class TestMissingWeightTreatedAsIdentity: + """A native-autograd bridge over a parameter-free norm (OLMo's OlmoLayerNorm has + no ``weight``) reports the LN-rule installed, so its rule-active forward and + backward must treat the missing weight as a unit scale instead of dereferencing + ``self.weight`` and raising AttributeError.""" + + def _make_param_free_bridge(self, d: int = 16) -> NormalizationBridge: + bridge = NormalizationBridge( + name="ln1", + config=_Cfg(uses_rms_norm=False), + use_native_layernorm_autograd=True, + ) + bridge.set_original_component(_TinyOlmoLayerNorm(d)) + return bridge + + def test_active_forward_matches_native_forward(self): + bridge = self._make_param_free_bridge() + block = _Block(bridge) + x = torch.randn(2, 5, 16) + baseline = bridge.original_component(x) + with use_relevance_rules(block, RelevanceRules(normalization=True)) as coverage: + assert coverage.installed == ("ln1",) + active = bridge(x) + assert torch.equal(active, baseline) + + def test_vjp_matches_detached_denom_oracle_with_unit_weight(self): + bridge = self._make_param_free_bridge() + block = _Block(bridge) + d = 16 + x = torch.randn(2, 5, d, requires_grad=True) + with use_relevance_rules(block, RelevanceRules(normalization=True)): + y = bridge(x) + y.sum().backward() + grad_rule = x.grad.clone() + + x_oracle = x.detach().clone().requires_grad_(True) + y_oracle = _denom_detached_oracle( + x_oracle, + torch.ones(d), + None, + uses_rms=False, + offset=False, + eps=1e-5, + ) + y_oracle.sum().backward() + torch.testing.assert_close(grad_rule, x_oracle.grad) + + +class TestFailClosedHookPrecedence: + def test_bwd_hook_raises_while_rule_active(self): + bridge = _make_bridge(native=True) + block = _Block(bridge) + x = torch.randn(2, 5, 16, requires_grad=True) + bridge.hook_normalized.add_hook(lambda g, hook=None: g, dir="bwd") + with pytest.raises(RelevanceRuleConflictError, match="Backward hooks"): + with use_relevance_rules(block, RelevanceRules(normalization=True)): + bridge(x) + bridge.hook_normalized.remove_hooks() + + def test_forward_edit_raises_while_rule_active(self): + bridge = _make_bridge(native=True) + block = _Block(bridge) + x = torch.randn(2, 5, 16) + bridge.hook_scale.add_hook(lambda t, hook=None: t * 2.0) + with pytest.raises(RelevanceRuleConflictError, match="forward hook edited"): + with use_relevance_rules(block, RelevanceRules(normalization=True)): + bridge(x) + bridge.hook_scale.remove_hooks() + + def test_observation_only_forward_hook_does_not_raise_while_rule_active(self): + """A hook that only observes (returns None) is not an edit and must not + trip the fail-closed check.""" + bridge = _make_bridge(native=True) + block = _Block(bridge) + 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 use_relevance_rules(block, RelevanceRules(normalization=True)): + active = bridge(x) + bridge.hook_normalized.remove_hooks() + assert torch.equal(active, baseline) + assert "normalized" in cache + + def test_bwd_hook_still_falls_back_with_warning_when_rule_inactive(self): + """Control: unchanged behavior when no rule is active (regression guard).""" + bridge = _make_bridge(native=True) + x = torch.randn(2, 5, 16, requires_grad=True) + bridge.hook_normalized.add_hook(lambda g, hook=None: g, dir="bwd") + with pytest.warns(UserWarning, match="Backward hooks"): + bridge(x) + bridge.hook_normalized.remove_hooks() + + def test_edit_still_falls_back_with_warning_when_rule_inactive(self): + """Control: unchanged behavior when no rule is active (regression guard).""" + bridge = _make_bridge(native=True) + x = torch.randn(2, 5, 16) + bridge.hook_scale.add_hook(lambda t, hook=None: t * 2.0) + with pytest.warns(UserWarning, match="reconstructed from the hooked values"): + bridge(x) + bridge.hook_scale.remove_hooks() diff --git a/tests/unit/model_bridge/supported_architectures/test_relevance_rule_mount_placeholders.py b/tests/unit/model_bridge/supported_architectures/test_relevance_rule_mount_placeholders.py new file mode 100644 index 0000000000..4e9d800089 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_relevance_rule_mount_placeholders.py @@ -0,0 +1,128 @@ +"""Relevance-rule coverage classification for architectures that mount q/k/v-norm +or sandwich-norm placeholders instead of a rule-capable ln1/ln2. + +``use_relevance_rules`` targets the "normalization" rule kind only at the ln1/ln2 +mount name (see ``_CANONICAL_MOUNTS`` in ``transformer_lens.model_bridge._relevance_rules``), +matched on the last segment of a live module's registered name. These tests inspect +each adapter's declared ``component_mapping`` directly (no HF weights, no live module +wiring) to lock in the structural facts that make the mechanism's "not applicable" +behavior correct for these three documented edge cases, without needing to build a +working forward-capable module tree per adapter: + +- Gemma 3n: every per-block norm (including the per-head q/k/v norms) is a plain + ``GeneralizedComponent`` placeholder, and none of them is keyed "ln1"/"ln2". +- Gemma 4: ln1/ln2 exist (sandwich norms) but are plain ``GeneralizedComponent`` + placeholders, not ``NormalizationBridge`` -- rule-incapable even though the mount + name matches. The per-head q/k/v norms are, like Gemma 3n, keyed under "self_attn", + never "ln1"/"ln2". +- StableLM: ln1/ln2 ARE real ``NormalizationBridge`` instances (a genuine rule + target), but the per-head norms are keyed "q_norm"/"k_norm" under "attn" even + though the wrapped HF module is named "q_layernorm"/"k_layernorm" -- the mount + key, not the wrapped module's own name, is what the canonical-mount check reads. +""" + +from typing import Any + +from tests.unit.model_bridge.supported_architectures.helpers import make_bridge_cfg +from transformer_lens.model_bridge._relevance_rules import _RelevanceRuleCapable +from transformer_lens.model_bridge.generalized_components import NormalizationBridge +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) +from transformer_lens.model_bridge.supported_architectures.gemma3n import ( + Gemma3nArchitectureAdapter, +) +from transformer_lens.model_bridge.supported_architectures.gemma4 import ( + Gemma4ArchitectureAdapter, +) +from transformer_lens.model_bridge.supported_architectures.stablelm import ( + StableLmArchitectureAdapter, +) + + +def _block_submodules(adapter: Any) -> dict: + return dict(adapter.component_mapping["blocks"].submodules) + + +def _attn_submodules(block_submodules: dict, attn_key: str) -> dict: + return dict(block_submodules[attn_key].submodules) + + +class TestGemma3nPlaceholders: + def _adapter(self) -> Gemma3nArchitectureAdapter: + cfg = make_bridge_cfg("Gemma3nForConditionalGeneration", n_key_value_heads=2, d_head=8) + return Gemma3nArchitectureAdapter(cfg) + + def test_no_ln1_or_ln2_key_at_block_level(self): + block_submodules = _block_submodules(self._adapter()) + assert "ln1" not in block_submodules + assert "ln2" not in block_submodules + + def test_qkv_norm_placeholders_are_plain_and_rule_incapable(self): + block_submodules = _block_submodules(self._adapter()) + attn_submodules = _attn_submodules(block_submodules, "self_attn") + for key in ("q_norm", "k_norm", "v_norm"): + component = attn_submodules[key] + assert type(component) is GeneralizedComponent + assert not isinstance(component, NormalizationBridge) + assert not isinstance(component, _RelevanceRuleCapable) + + +class TestGemma4Placeholders: + def _adapter(self) -> Gemma4ArchitectureAdapter: + from types import SimpleNamespace + + cfg = make_bridge_cfg("Gemma4ForConditionalGeneration", n_key_value_heads=1, d_head=8) + cfg.vision_config = SimpleNamespace( + hidden_size=32, num_hidden_layers=2, num_attention_heads=4 + ) + cfg.vision_soft_tokens_per_image = 4 + return Gemma4ArchitectureAdapter(cfg) + + def test_ln1_ln2_are_plain_sandwich_placeholders_not_normalization_bridge(self): + block_submodules = _block_submodules(self._adapter()) + for key in ("ln1", "ln2"): + component = block_submodules[key] + assert type(component) is GeneralizedComponent + assert not isinstance(component, NormalizationBridge) + assert not isinstance(component, _RelevanceRuleCapable) + + def test_qkv_norm_placeholders_are_not_keyed_ln1_or_ln2(self): + block_submodules = _block_submodules(self._adapter()) + attn_submodules = _attn_submodules(block_submodules, "attn") + for key in ("q_norm", "k_norm", "v_norm"): + assert key not in ("ln1", "ln2") + component = attn_submodules[key] + assert type(component) is GeneralizedComponent + assert not isinstance(component, _RelevanceRuleCapable) + + +class TestStableLmPlaceholders: + def _adapter(self) -> StableLmArchitectureAdapter: + cfg = make_bridge_cfg("StableLmForCausalLM", n_key_value_heads=2, d_head=8) + return StableLmArchitectureAdapter(cfg) + + def test_ln1_and_ln2_are_genuine_rule_targets(self): + """Unlike Gemma 3n/4, StableLM's block-level norms ARE rule-capable -- + this is the genuine positive case the other two are contrasted against.""" + block_submodules = _block_submodules(self._adapter()) + for key in ("ln1", "ln2"): + component = block_submodules[key] + assert isinstance(component, NormalizationBridge) + assert isinstance(component, _RelevanceRuleCapable) + assert component._relevance_rule_kinds == ("normalization",) + + def test_per_head_norms_are_keyed_q_norm_k_norm_not_ln1_ln2(self): + """The wrapped HF module is named q_layernorm/k_layernorm, but the mount + KEY the canonical-mount check reads is q_norm/k_norm -- distinct either + way from ln1/ln2, so these stay invisible to the "normalization" rule + regardless of which naming convention the underlying HF module uses.""" + block_submodules = _block_submodules(self._adapter()) + attn_submodules = _attn_submodules(block_submodules, "attn") + assert attn_submodules["q_norm"].name == "q_layernorm" + assert attn_submodules["k_norm"].name == "k_layernorm" + for key in ("q_norm", "k_norm"): + component = attn_submodules[key] + assert key not in ("ln1", "ln2") + assert type(component) is GeneralizedComponent + assert not isinstance(component, _RelevanceRuleCapable) diff --git a/tests/unit/model_bridge/test_relevance_rules.py b/tests/unit/model_bridge/test_relevance_rules.py new file mode 100644 index 0000000000..57cca39c25 --- /dev/null +++ b/tests/unit/model_bridge/test_relevance_rules.py @@ -0,0 +1,369 @@ +"""Tests for the scoped relevance-rule context: forward-identity, cleanup, and nesting. + +The fixture component below installs the real LN-rule primitive +(``transformer_lens.model_bridge._relevance_rules.ln_rule``) through the protocol +``use_relevance_rules`` relies on to find and toggle rule-capable components. That keeps +these tests focused on the scoping mechanics -- forward identity, gradient restoration, +nested contexts, exception safety, and positional (not class-based) targeting -- rather +than on any concrete NormalizationBridge or gated-MLP integration, which land in later +commits. + +The scoped context does not exist yet, so this fails collection with a single +ImportError -- the expected red state before the context is implemented. +""" + +import dataclasses + +import pytest +import torch +import torch.nn as nn + +from transformer_lens.model_bridge._relevance_rules import ( + RelevanceRuleCoverage, + RelevanceRules, + RelevanceRuleUnsupportedError, + ln_rule, + use_relevance_rules, +) + + +class _FakeNormComponent(nn.Module): + """A minimal ln1/ln2-style target: installs the real LN-rule primitive on request.""" + + _relevance_rule_kinds = ("normalization",) + + def __init__(self, eps: float = 1e-2): + super().__init__() + # Composing this division with a plain second normalization and a linear + # readout makes the LN-rule's denom-as-constant correction shrink linearly + # with eps, so a realistic normalization epsilon (1e-6) would leave the + # correction too small for torch.allclose to detect reliably end to end. + self.eps = eps + self._rule_active = False + + def _enable_relevance_rule(self, kind: str) -> None: + self._rule_active = True + + def _disable_relevance_rule(self, kind: str) -> None: + self._rule_active = False + + def forward(self, x: torch.Tensor) -> torch.Tensor: + denom = x.abs().mean(dim=-1, keepdim=True) + self.eps + if self._rule_active: + return ln_rule(x, denom) + return x / denom + + +class _PlainMount(nn.Module): + """Occupies a targeted mount name but implements no relevance-rule protocol.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + +class _TinyBlock(nn.Module): + """Mimics a block's ln1/ln2 mount points plus a same-class q_norm and an MLP. + + ``q_norm`` uses the identical fake-normalization class as ``ln1``/``ln2`` so tests + can assert that targeting is positional (by mount name) rather than class-based. + """ + + def __init__(self): + super().__init__() + self.ln1: nn.Module = _FakeNormComponent() + self.q_norm = _FakeNormComponent() + self.ln2: nn.Module = _FakeNormComponent() + self.mlp = nn.Linear(4, 4) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.ln1(x) + x = self.q_norm(x) + x = self.ln2(x) + return self.mlp(x) + + +def _tiny_block() -> _TinyBlock: + block = _TinyBlock() + with torch.no_grad(): + block.mlp.weight.copy_(torch.eye(4) * 0.5) + block.mlp.bias.zero_() + return block + + +class _SandwichBlock(nn.Module): + """Mimics a sandwich-norm block: pre-norms ln1/ln2 plus post-norms ln1_post/ln2_post. + + Sandwich-norm architectures mount a second normalization after attention and after + the MLP at ``ln1_post``/``ln2_post``, so the LN-rule must reach those mounts as well + as the pre-norms rather than leaving them silently out of coverage. + """ + + def __init__(self): + super().__init__() + self.ln1: nn.Module = _FakeNormComponent() + self.ln1_post: nn.Module = _FakeNormComponent() + self.ln2: nn.Module = _FakeNormComponent() + self.ln2_post: nn.Module = _FakeNormComponent() + self.mlp = nn.Linear(4, 4) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.ln1(x) + x = self.ln1_post(x) + x = self.ln2(x) + x = self.ln2_post(x) + return self.mlp(x) + + +class _FakeGatedMLPComponent(nn.Module): + """A minimal mlp-mount target that answers to two rule kinds independently. + + A real gated-MLP node answers to both "activation" (Identity-rule on its + activation function) and "multiplicative_gate" (Half-rule on its gate*up + product) at the same mount, and either can be requested without the other, + so this fixture tracks the two kinds as separate booleans rather than one. + """ + + _relevance_rule_kinds = ("activation", "multiplicative_gate") + + def __init__(self): + super().__init__() + self._activation_rule_active = False + self._gate_rule_active = False + + def _enable_relevance_rule(self, kind: str) -> None: + if kind == "activation": + self._activation_rule_active = True + elif kind == "multiplicative_gate": + self._gate_rule_active = True + + def _disable_relevance_rule(self, kind: str) -> None: + if kind == "activation": + self._activation_rule_active = False + elif kind == "multiplicative_gate": + self._gate_rule_active = False + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + +class _Boom(Exception): + """Marker exception raised inside a context to test cleanup on failure.""" + + +def test_relevance_rules_defaults_to_no_rules(): + rules = RelevanceRules() + assert rules.normalization is False + assert rules.activation is False + assert rules.multiplicative_gate is False + assert rules.attention is False + + +def test_relevance_rules_is_frozen(): + rules = RelevanceRules(normalization=True) + with pytest.raises(dataclasses.FrozenInstanceError): + rules.normalization = False # type: ignore[misc] + + +def test_forward_is_identical_while_rule_active(): + block = _tiny_block() + x = torch.randn(2, 3, 4) + baseline = block(x) + with use_relevance_rules(block, RelevanceRules(normalization=True)): + active = block(x) + assert torch.equal(active, baseline) + + +def test_positional_targeting_excludes_same_class_component_at_other_mount(): + block = _tiny_block() + with use_relevance_rules(block, RelevanceRules(normalization=True)) as coverage: + assert block.ln1._rule_active is True + assert block.ln2._rule_active is True + assert block.q_norm._rule_active is False + assert set(coverage.installed) == {"ln1", "ln2"} + assert block.ln1._rule_active is False + assert block.ln2._rule_active is False + + +def test_normalization_rule_installs_on_sandwich_post_norm_mounts(): + block = _SandwichBlock() + with use_relevance_rules(block, RelevanceRules(normalization=True)) as coverage: + assert block.ln1._rule_active is True + assert block.ln1_post._rule_active is True + assert block.ln2._rule_active is True + assert block.ln2_post._rule_active is True + assert set(coverage.installed) == {"ln1", "ln1_post", "ln2", "ln2_post"} + assert coverage.skipped == () + assert block.ln1_post._rule_active is False + assert block.ln2_post._rule_active is False + + +def test_sandwich_post_norm_on_python_path_is_reported_skipped_never_absent(): + block = _SandwichBlock() + # A post-norm mount whose occupant takes the python-norm path (no rule protocol) + # must surface as skipped, never vanish from coverage as it did before the mount + # names were recognized. + block.ln2_post = _PlainMount() + with use_relevance_rules(block, RelevanceRules(normalization=True)) as coverage: + assert set(coverage.installed) == {"ln1", "ln1_post", "ln2"} + assert set(coverage.skipped) == {"ln2_post"} + + +def test_requesting_a_kind_with_no_canonical_mount_raises(): + block = _tiny_block() + # "attention" is a valid RelevanceRules field but has no canonical mount, so the + # request must raise rather than install nothing and return empty coverage. + with pytest.raises(ValueError, match="attention"): + with use_relevance_rules(block, RelevanceRules(attention=True)): + pass + + +def test_unsupported_component_at_targeted_mount_is_skipped(): + block = _tiny_block() + block.ln2 = _PlainMount() + with use_relevance_rules(block, RelevanceRules(normalization=True)) as coverage: + assert isinstance(coverage, RelevanceRuleCoverage) + assert set(coverage.installed) == {"ln1"} + assert set(coverage.skipped) == {"ln2"} + + +def test_requesting_unsupported_kind_on_capable_component_raises(): + block = _tiny_block() + block.mlp = _FakeGatedMLPComponent() + # This mount implements the protocol and currently only supports "activation", + # but -- unlike a mount that never deals with "multiplicative_gate" at all -- + # names it in _relevance_rule_unsupported_kinds as one it is expected to honor + # here and currently cannot, so requesting it raises instead of being skipped. + block.mlp._relevance_rule_kinds = ("activation",) + block.mlp._relevance_rule_unsupported_kinds = ("multiplicative_gate",) + with pytest.raises(RelevanceRuleUnsupportedError, match="mlp"): + with use_relevance_rules(block, RelevanceRules(multiplicative_gate=True)): + pass + assert block.mlp._gate_rule_active is False + + +def test_no_rules_requested_installs_nothing(): + block = _tiny_block() + with use_relevance_rules(block, RelevanceRules()) as coverage: + assert coverage.installed == () + assert coverage.skipped == () + assert block.ln1._rule_active is False + assert block.ln2._rule_active is False + + +def test_ordinary_gradients_restored_after_exit(): + block = _tiny_block() + x = torch.randn(2, 3, 4, requires_grad=True) + + with use_relevance_rules(block, RelevanceRules(normalization=True)): + y_rule = block(x) + (grad_rule,) = torch.autograd.grad(y_rule.sum(), x) + + assert block.ln1._rule_active is False + assert block.ln2._rule_active is False + + x_plain = x.detach().clone().requires_grad_(True) + y_plain = block(x_plain) + (grad_plain,) = torch.autograd.grad(y_plain.sum(), x_plain) + + # The LN-rule treats the denominator as constant, so its VJP differs from + # ordinary autodiff through the same division wherever the denominator + # actually depends on the input -- true for every row of this fixture. + assert not torch.allclose(grad_rule, grad_plain) + + # A second plain pass confirms the exit left no residual rule state: it must + # reproduce grad_plain exactly rather than drifting toward grad_rule. + x_plain_again = x.detach().clone().requires_grad_(True) + y_plain_again = block(x_plain_again) + (grad_plain_again,) = torch.autograd.grad(y_plain_again.sum(), x_plain_again) + torch.testing.assert_close(grad_plain_again, grad_plain) + + +def test_nested_contexts_restore_outer_state_on_inner_exit(): + block = _tiny_block() + x = torch.randn(2, 3, 4) + baseline = block(x) + + with use_relevance_rules(block, RelevanceRules(normalization=True)) as outer_coverage: + assert block.ln1._rule_active is True + with use_relevance_rules(block, RelevanceRules(normalization=True)) as inner_coverage: + assert block.ln1._rule_active is True + assert torch.equal(block(x), baseline) + # The inner exit must not disable the rule the outer context still needs. + assert block.ln1._rule_active is True + assert block.ln2._rule_active is True + assert torch.equal(block(x), baseline) + + assert block.ln1._rule_active is False + assert block.ln2._rule_active is False + assert set(outer_coverage.installed) == {"ln1", "ln2"} + assert set(inner_coverage.installed) == {"ln1", "ln2"} + + +def test_exception_inside_context_leaves_no_rule_state(): + block = _tiny_block() + x = torch.randn(2, 3, 4) + baseline = block(x) + + with pytest.raises(_Boom): + with use_relevance_rules(block, RelevanceRules(normalization=True)): + assert block.ln1._rule_active is True + raise _Boom("failure inside the scoped context") + + assert block.ln1._rule_active is False + assert block.ln2._rule_active is False + assert torch.equal(block(x), baseline) + + +def test_exception_during_nested_context_restores_outer_state(): + block = _tiny_block() + x = torch.randn(2, 3, 4) + baseline = block(x) + + with use_relevance_rules(block, RelevanceRules(normalization=True)): + with pytest.raises(_Boom): + with use_relevance_rules(block, RelevanceRules(normalization=True)): + raise _Boom("failure inside the nested scoped context") + # The outer context is still active after the inner one unwinds. + assert block.ln1._rule_active is True + assert torch.equal(block(x), baseline) + + assert block.ln1._rule_active is False + assert block.ln2._rule_active is False + + +def test_one_kind_requested_on_a_two_kind_mount_leaves_the_other_kind_inactive(): + block = _tiny_block() + block.mlp = _FakeGatedMLPComponent() + with use_relevance_rules(block, RelevanceRules(multiplicative_gate=True)) as coverage: + assert block.mlp._gate_rule_active is True + assert block.mlp._activation_rule_active is False + assert coverage.installed == ("mlp",) + assert block.mlp._gate_rule_active is False + + +def test_both_kinds_requested_together_both_activate_on_the_same_mount(): + block = _tiny_block() + block.mlp = _FakeGatedMLPComponent() + with use_relevance_rules( + block, RelevanceRules(activation=True, multiplicative_gate=True) + ) as coverage: + assert block.mlp._activation_rule_active is True + assert block.mlp._gate_rule_active is True + assert set(coverage.installed) == {"mlp"} + assert block.mlp._activation_rule_active is False + assert block.mlp._gate_rule_active is False + + +def test_nested_scopes_over_different_kinds_on_one_mount_refcount_independently(): + block = _tiny_block() + block.mlp = _FakeGatedMLPComponent() + with use_relevance_rules(block, RelevanceRules(activation=True)): + assert block.mlp._activation_rule_active is True + with use_relevance_rules(block, RelevanceRules(multiplicative_gate=True)): + assert block.mlp._activation_rule_active is True + assert block.mlp._gate_rule_active is True + # The inner (multiplicative_gate-only) scope's exit must not disable + # the outer scope's independently refcounted activation rule. + assert block.mlp._activation_rule_active is True + assert block.mlp._gate_rule_active is False + assert block.mlp._activation_rule_active is False diff --git a/tests/unit/tools/test_relevance_rules.py b/tests/unit/tools/test_relevance_rules.py new file mode 100644 index 0000000000..c8a6e35046 --- /dev/null +++ b/tests/unit/tools/test_relevance_rules.py @@ -0,0 +1,444 @@ +"""Analytic closed-form tests for the LN-, Identity-, and Half-relevance-rule primitives. + +Model-free: every primitive is a plain ``torch.autograd.Function`` exercised on +synthetic tensors, so no model, backward hook, or ``.data`` access is involved. Each +test asserts two properties independently: the forward value is exactly the native +(ordinary-autograd) value, and the backward value matches the rule's closed-form VJP +rather than what ordinary autodiff would produce. +""" + +import math +from functools import partial + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.pytorch_utils import Conv1D + +from transformer_lens.model_bridge._relevance_rules import ( + half_rule, + identity_rule, + ln_rule, +) +from transformer_lens.model_bridge.generalized_components.mlp import ( + normalize_mlp_weight, + weight_layout_in_out, +) + +DTYPES = [torch.float32, torch.float64] + + +def _leaf(values: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return values.to(dtype).clone().requires_grad_(True) + + +def _sample_rows(dtype: torch.dtype) -> torch.Tensor: + """A batch covering an all-zero, all-negative, mixed-sign, and positive row.""" + values = torch.tensor( + [ + [0.0, 0.0, 0.0, 0.0, 0.0], + [-1.0, -2.0, -0.5, -3.0, -0.25], + [1.0, -2.0, 3.0, -4.0, 0.5], + [2.0, 4.0, 6.0, 8.0, 10.0], + ] + ) + return values.to(dtype) + + +def _normal_cdf(x: torch.Tensor) -> torch.Tensor: + return 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0))) + + +class TestLNRule: + """Forward equals ``numerator / denom``; the VJP treats ``denom`` as constant.""" + + @pytest.mark.parametrize("dtype", DTYPES) + def test_forward_matches_plain_division(self, dtype): + numerator = _sample_rows(dtype) + denom = numerator.abs().mean(dim=-1, keepdim=True) + 1e-6 + assert torch.equal(ln_rule(numerator, denom), numerator / denom) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_backward_treats_denom_as_constant(self, dtype): + eps = 1e-6 + base = _sample_rows(dtype) + grad_out = torch.ones_like(base) + + x_rule = _leaf(base, dtype) + denom_rule = x_rule.abs().mean(dim=-1, keepdim=True) + eps + y_rule = ln_rule(x_rule, denom_rule) + (grad_rule,) = torch.autograd.grad(y_rule, x_rule, grad_outputs=grad_out) + + # The closed-form VJP of this rule: grad_x = grad_out / denom, with no + # contribution from d(denom)/dx. + expected = grad_out / denom_rule.detach() + torch.testing.assert_close(grad_rule, expected) + + # Plain autodiff through the same expression differentiates the denom too, + # so it disagrees with the rule everywhere the denom actually depends on x + # (every row here, since eps alone would zero out that dependency). + x_plain = _leaf(base, dtype) + denom_plain = x_plain.abs().mean(dim=-1, keepdim=True) + eps + y_plain = x_plain / denom_plain + (grad_plain,) = torch.autograd.grad(y_plain, x_plain, grad_outputs=grad_out) + assert not torch.allclose(grad_plain, grad_rule) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_backward_zero_row_has_no_division_by_zero(self, dtype): + eps = 1e-6 + x = _leaf(_sample_rows(dtype)[:1], dtype) + denom = x.abs().mean(dim=-1, keepdim=True) + eps + y = ln_rule(x, denom) + (grad,) = torch.autograd.grad(y, x, grad_outputs=torch.ones_like(x)) + assert torch.isfinite(grad).all() + torch.testing.assert_close(grad, torch.full_like(x, 1.0 / eps)) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_forward_and_backward_on_non_contiguous_input(self, dtype): + base = _sample_rows(dtype).t() + assert not base.is_contiguous() + numerator = base.clone().requires_grad_(True) + denom = numerator.abs().mean(dim=-1, keepdim=True) + 1e-6 + + y = ln_rule(numerator, denom) + assert torch.equal(y, numerator.detach() / denom.detach()) + + grad_out = torch.ones_like(numerator) + (grad,) = torch.autograd.grad(y, numerator, grad_outputs=grad_out) + torch.testing.assert_close(grad, grad_out / denom.detach()) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_batched_leading_dims(self, dtype): + numerator = _leaf(torch.randn(2, 3, 4), dtype) + denom = numerator.abs().mean(dim=-1, keepdim=True) + 1e-6 + y = ln_rule(numerator, denom) + assert y.shape == numerator.shape + assert torch.equal(y, numerator.detach() / denom.detach()) + + (grad,) = torch.autograd.grad(y, numerator, grad_outputs=torch.ones_like(numerator)) + torch.testing.assert_close(grad, torch.ones_like(numerator) / denom.detach()) + + +class TestIdentityRule: + """Forward equals the native activation; the VJP is ``grad_out * phi(x)``. + + ``phi`` has a closed form for each activation tested here: ``sigmoid`` for SiLU + and the Gaussian CDF for exact GELU. Both are, by construction, the removable- + singularity limit of ``f(x) / x`` at ``x == 0``, which is the direct oracle used + for the tanh-approximate GELU, where no simpler closed form applies. + """ + + @pytest.mark.parametrize("dtype", DTYPES) + def test_forward_matches_native_silu(self, dtype): + x = _sample_rows(dtype) + assert torch.equal(identity_rule(x, F.silu), F.silu(x)) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_backward_matches_sigmoid_for_silu(self, dtype): + x = _leaf(_sample_rows(dtype), dtype) + y = identity_rule(x, F.silu) + grad_out = torch.ones_like(x) + (grad,) = torch.autograd.grad(y, x, grad_outputs=grad_out) + torch.testing.assert_close(grad, grad_out * torch.sigmoid(x)) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_forward_matches_native_gelu_exact(self, dtype): + x = _sample_rows(dtype) + act_fn = partial(F.gelu, approximate="none") + assert torch.equal(identity_rule(x, act_fn), act_fn(x)) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_backward_matches_gaussian_cdf_for_exact_gelu(self, dtype): + x = _leaf(_sample_rows(dtype), dtype) + act_fn = partial(F.gelu, approximate="none") + y = identity_rule(x, act_fn) + grad_out = torch.ones_like(x) + (grad,) = torch.autograd.grad(y, x, grad_outputs=grad_out) + torch.testing.assert_close(grad, grad_out * _normal_cdf(x)) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_forward_matches_native_gelu_approx(self, dtype): + x = _sample_rows(dtype) + act_fn = partial(F.gelu, approximate="tanh") + assert torch.equal(identity_rule(x, act_fn), act_fn(x)) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_backward_matches_ratio_with_zero_limit_for_approx_gelu(self, dtype): + x = _leaf(_sample_rows(dtype), dtype) + act_fn = partial(F.gelu, approximate="tanh") + y = identity_rule(x, act_fn) + grad_out = torch.ones_like(x) + (grad,) = torch.autograd.grad(y, x, grad_outputs=grad_out) + + with torch.no_grad(): + safe_x = torch.where(x == 0, torch.ones_like(x), x) + expected_phi = torch.where(x == 0, torch.full_like(x, 0.5), act_fn(x.detach()) / safe_x) + assert torch.isfinite(grad).all() + torch.testing.assert_close(grad, grad_out * expected_phi) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_forward_and_backward_on_non_contiguous_input(self, dtype): + base = _sample_rows(dtype).t() + assert not base.is_contiguous() + x = base.clone().requires_grad_(True) + + y = identity_rule(x, F.silu) + assert torch.equal(y, F.silu(x.detach())) + + grad_out = torch.ones_like(x) + (grad,) = torch.autograd.grad(y, x, grad_outputs=grad_out) + torch.testing.assert_close(grad, grad_out * torch.sigmoid(x.detach())) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_batched_leading_dims(self, dtype): + x = _leaf(torch.randn(2, 3, 4), dtype) + y = identity_rule(x, F.silu) + assert y.shape == x.shape + assert torch.equal(y, F.silu(x.detach())) + + (grad,) = torch.autograd.grad(y, x, grad_outputs=torch.ones_like(x)) + torch.testing.assert_close(grad, torch.sigmoid(x.detach())) + + +class TestHalfRule: + """Forward equals ``u * v``; the VJP halves each ordinary product-rule term.""" + + @pytest.mark.parametrize("dtype", DTYPES) + def test_forward_matches_plain_product(self, dtype): + u = _sample_rows(dtype) + v = _sample_rows(dtype).flip(0) + assert torch.equal(half_rule(u, v), u * v) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_backward_halves_each_operand_gradient(self, dtype): + u = _leaf(_sample_rows(dtype), dtype) + v = _leaf(_sample_rows(dtype).flip(0), dtype) + grad_out = torch.ones_like(u) + + y = half_rule(u, v) + grad_u, grad_v = torch.autograd.grad(y, (u, v), grad_outputs=grad_out) + torch.testing.assert_close(grad_u, 0.5 * grad_out * v.detach()) + torch.testing.assert_close(grad_v, 0.5 * grad_out * u.detach()) + + # Ordinary autodiff of u * v would give the full (unhalved) product-rule + # terms, so the rule must disagree with it. + u_plain = _leaf(_sample_rows(dtype), dtype) + v_plain = _leaf(_sample_rows(dtype).flip(0), dtype) + y_plain = u_plain * v_plain + grad_u_plain, grad_v_plain = torch.autograd.grad( + y_plain, (u_plain, v_plain), grad_outputs=grad_out + ) + assert not torch.allclose(grad_u_plain, grad_u) + assert not torch.allclose(grad_v_plain, grad_v) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_forward_and_backward_on_non_contiguous_input(self, dtype): + u = _sample_rows(dtype).t().clone().requires_grad_(True) + v = _sample_rows(dtype).flip(0).t().clone().requires_grad_(True) + assert not u.is_contiguous() + assert not v.is_contiguous() + + y = half_rule(u, v) + assert torch.equal(y, u.detach() * v.detach()) + + grad_out = torch.ones_like(u) + grad_u, grad_v = torch.autograd.grad(y, (u, v), grad_outputs=grad_out) + torch.testing.assert_close(grad_u, 0.5 * grad_out * v.detach()) + torch.testing.assert_close(grad_v, 0.5 * grad_out * u.detach()) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_batched_leading_dims(self, dtype): + u = _leaf(torch.randn(2, 3, 4), dtype) + v = _leaf(torch.randn(2, 3, 4), dtype) + y = half_rule(u, v) + assert y.shape == u.shape + assert torch.equal(y, u.detach() * v.detach()) + + grad_out = torch.ones_like(u) + grad_u, grad_v = torch.autograd.grad(y, (u, v), grad_outputs=grad_out) + torch.testing.assert_close(grad_u, 0.5 * grad_out * v.detach()) + torch.testing.assert_close(grad_v, 0.5 * grad_out * u.detach()) + + +class _Proj: + """Minimal stand-in for a projection bridge: carries only what + ``weight_layout_in_out``/``normalize_mlp_weight`` read (``original_component``).""" + + def __init__(self, original_component): + self.original_component = original_component + + +class TestGatedMLPWeightOrientation: + """``MLPBridge``'s ``W_gate``/``W_in``/``W_out`` read the underlying projection + through ``weight_layout_in_out``/``normalize_mlp_weight``, so an orientation bug in + those helpers would silently transpose a gate/up/down weight for one backing class. + Covers both HF module classes: ``nn.Linear`` (weight stored ``[out, in]``, + transposed to TL orientation) and ``Conv1D`` (weight stored ``[in, out]``, already + TL-oriented). + """ + + @pytest.fixture(params=["nn.Linear", "Conv1D"]) + def backing_class(self, request): + return request.param + + def _make_gate_up_down(self, backing_class: str, d_model: int = 3, d_mlp: int = 5): + torch.manual_seed(0) + if backing_class == "nn.Linear": + gate_proj = nn.Linear(d_model, d_mlp) + up_proj = nn.Linear(d_model, d_mlp) + down_proj = nn.Linear(d_mlp, d_model) + else: + gate_proj = Conv1D(d_mlp, d_model) + up_proj = Conv1D(d_mlp, d_model) + down_proj = Conv1D(d_model, d_mlp) + return gate_proj, up_proj, down_proj + + def _tl_weight(self, proj: torch.nn.Module, pattern: str) -> torch.Tensor: + wrapper = _Proj(proj) + layout = weight_layout_in_out(wrapper) + return normalize_mlp_weight(proj.weight, layout, wrapper, pattern=pattern) + + def test_tl_oriented_matmul_reproduces_native_projection(self, backing_class): + gate_proj, up_proj, down_proj = self._make_gate_up_down(backing_class) + x = torch.randn(2, 3) + + w_gate = self._tl_weight(gate_proj, pattern="in") + w_in = self._tl_weight(up_proj, pattern="in") + w_out = self._tl_weight(down_proj, pattern="out") + + assert torch.allclose(x @ w_gate + gate_proj.bias, gate_proj(x), atol=1e-6) + assert torch.allclose(x @ w_in + up_proj.bias, up_proj(x), atol=1e-6) + hidden = torch.randn(2, 5) + assert torch.allclose(hidden @ w_out + down_proj.bias, down_proj(hidden), atol=1e-6) + + +class TestPinnedReferenceParity: + """Tolerant parity against ``FarnoushRJ/RelP`` pinned at + ``8219d6dc417c3fd7f318342cf61cd2a0c20b7250``. + + That repository vendors an unrelated pre-Bridge TransformerLens fork, so its + rule formulas are reimplemented here directly from the pinned commit's + component diffs rather than imported: + + - LN-rule (``transformer_lens/components/rms_norm.py``): ``x / scale.detach()``. + - Identity-rule (``transformer_lens/utilities/activation_functions.py``, + class ``ModifiedAct``): ``zp = stabilize(x); zp * (act_fn(x) / zp).detach()``, + where ``stabilize(z) = z + ((z == 0) + sign(z)) * 1e-6`` + (``transformer_lens/lrp_utils.py``). + - Half-rule (``transformer_lens/components/mlps/gated_mlp.py``): + ``z = u * v; z / 2 + (z / 2).detach()``. + + The LN- and Half-rule reference formulas produce the same VJP as this module's + primitives to floating-point precision. The Identity-rule reference formula + does not: its epsilon stabilizer only approximates the paper-defined factor + away from ``x == 0``, and collapses to exactly zero at ``x == 0`` where the + paper-defined factor's removable-singularity limit is ``0.5``. + """ + + @staticmethod + def _reference_stabilize(z: torch.Tensor) -> torch.Tensor: + return z + ((z == 0).to(z.dtype) + torch.sign(z)) * 1e-6 + + @classmethod + def _reference_ln_rule_grad(cls, x: torch.Tensor, denom_fn) -> torch.Tensor: + x = x.clone().requires_grad_(True) + denom = denom_fn(x) + y = x / denom.detach() + (grad,) = torch.autograd.grad(y, x, grad_outputs=torch.ones_like(x)) + return grad + + @classmethod + def _reference_identity_rule_grad(cls, x: torch.Tensor, act_fn) -> torch.Tensor: + x = x.clone().requires_grad_(True) + z = act_fn(x) + zp = cls._reference_stabilize(x) + y = zp * (z / zp).detach() + (grad,) = torch.autograd.grad(y, x, grad_outputs=torch.ones_like(x)) + return grad + + @classmethod + def _reference_half_rule_grad(cls, u: torch.Tensor, v: torch.Tensor): + u = u.clone().requires_grad_(True) + v = v.clone().requires_grad_(True) + z = u * v + y = z / 2 + (z / 2).detach() + grad_u, grad_v = torch.autograd.grad(y, (u, v), grad_outputs=torch.ones_like(u)) + return grad_u, grad_v + + @pytest.mark.parametrize("dtype", DTYPES) + def test_ln_rule_matches_reference_grad(self, dtype): + def denom_fn(t: torch.Tensor) -> torch.Tensor: + return (t.pow(2).mean(-1, keepdim=True) + 1e-6).sqrt() + + x = _sample_rows(dtype) + + x_rule = _leaf(x, dtype) + denom_rule = denom_fn(x_rule) + y_rule = ln_rule(x_rule, denom_rule) + (grad_rule,) = torch.autograd.grad(y_rule, x_rule, grad_outputs=torch.ones_like(x_rule)) + + grad_reference = self._reference_ln_rule_grad(x, denom_fn) + torch.testing.assert_close(grad_rule, grad_reference) + + @pytest.mark.parametrize("dtype", DTYPES) + @pytest.mark.parametrize( + "act_fn", + [F.silu, partial(F.gelu, approximate="none"), partial(F.gelu, approximate="tanh")], + ids=["silu", "gelu_exact", "gelu_tanh"], + ) + def test_identity_rule_matches_reference_away_from_zero(self, dtype, act_fn): + x = _sample_rows(dtype) + nonzero_mask = x != 0 + + x_rule = _leaf(x, dtype) + y_rule = identity_rule(x_rule, act_fn) + (grad_rule,) = torch.autograd.grad(y_rule, x_rule, grad_outputs=torch.ones_like(x_rule)) + + grad_reference = self._reference_identity_rule_grad(x, act_fn) + + torch.testing.assert_close( + grad_rule[nonzero_mask], + grad_reference[nonzero_mask], + rtol=1e-4, + atol=1e-5, + ) + + @pytest.mark.parametrize("dtype", DTYPES) + @pytest.mark.parametrize( + "act_fn", + [F.silu, partial(F.gelu, approximate="none"), partial(F.gelu, approximate="tanh")], + ids=["silu", "gelu_exact", "gelu_tanh"], + ) + def test_identity_rule_exact_zero_discrepancy_is_documented(self, dtype, act_fn): + """At ``x == 0`` this module's Identity-rule uses the paper-defined + removable-singularity limit ``0.5``, while the pinned reference's epsilon + stabilizer yields exactly ``0``. Assert both values explicitly, rather than + letting a tolerance absorb the gap, so a change to either side's zero + handling is caught instead of silently passing. + """ + x = torch.zeros(3, dtype=dtype) + + x_rule = _leaf(x, dtype) + y_rule = identity_rule(x_rule, act_fn) + (grad_rule,) = torch.autograd.grad(y_rule, x_rule, grad_outputs=torch.ones_like(x_rule)) + torch.testing.assert_close(grad_rule, torch.full_like(x, 0.5)) + + grad_reference = self._reference_identity_rule_grad(x, act_fn) + torch.testing.assert_close(grad_reference, torch.zeros_like(x)) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_half_rule_matches_reference_grad(self, dtype): + u = _sample_rows(dtype) + v = _sample_rows(dtype).flip(0) + + u_rule = _leaf(u, dtype) + v_rule = _leaf(v, dtype) + y_rule = half_rule(u_rule, v_rule) + grad_u_rule, grad_v_rule = torch.autograd.grad( + y_rule, (u_rule, v_rule), grad_outputs=torch.ones_like(u_rule) + ) + + grad_u_reference, grad_v_reference = self._reference_half_rule_grad(u, v) + torch.testing.assert_close(grad_u_rule, grad_u_reference) + torch.testing.assert_close(grad_v_rule, grad_v_reference) diff --git a/transformer_lens/model_bridge/_relevance_rules.py b/transformer_lens/model_bridge/_relevance_rules.py new file mode 100644 index 0000000000..5da1bbf80d --- /dev/null +++ b/transformer_lens/model_bridge/_relevance_rules.py @@ -0,0 +1,348 @@ +"""Forward-equivalent relevance-rule primitives, plus the scoped context that installs them. + +Each primitive is a ``torch.autograd.Function`` that reproduces its native forward +value exactly while replacing the backward pass with the rule's closed-form VJP. +``use_relevance_rules`` installs these rules on a model's canonical mount points only +for the duration of a ``with`` block, targeting components positionally (by mount +name, never by class) and reporting which mounts were installed versus skipped. +""" + +import dataclasses +from contextlib import contextmanager +from typing import ( + Any, + Callable, + Dict, + Iterator, + List, + Mapping, + Protocol, + Tuple, + runtime_checkable, +) + +import torch +import torch.nn as nn + + +def ln_rule_grad(grad_output: torch.Tensor, denom: torch.Tensor) -> torch.Tensor: + """Core LN-rule VJP: divide by ``denom`` without differentiating through it. + + Shared by the ``ln_rule`` primitive below and by any integration (such as + ``NormalizationBridge``) that wraps a component's own native forward call + instead of reproducing the division itself. + """ + return grad_output / denom + + +class _LNRule(torch.autograd.Function): + """LN-rule: forward is ``numerator / denom``; the VJP treats ``denom`` as constant.""" + + @staticmethod + def forward(ctx: Any, numerator: torch.Tensor, denom: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(denom) + return numerator / denom + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]: + (denom,) = ctx.saved_tensors + return ln_rule_grad(grad_output, denom), None + + +def ln_rule(numerator: torch.Tensor, denom: torch.Tensor) -> torch.Tensor: + """Apply the LN-rule: native division forward, denom-as-constant backward.""" + result: torch.Tensor = _LNRule.apply(numerator, denom) + return result + + +class _IdentityRule(torch.autograd.Function): + """Identity-rule: forward is the native activation; the VJP is ``grad_out * phi(x)``. + + ``phi`` is ``f(x) / x``, the removable singularity at ``x == 0`` filled in with its + limit ``0.5``. This holds for any elementwise activation with ``f(0) == 0`` and a + well-defined derivative at zero, which covers SiLU and both GELU variants. + """ + + @staticmethod + def forward( + ctx: Any, x: torch.Tensor, act_fn: Callable[[torch.Tensor], torch.Tensor] + ) -> torch.Tensor: + y = act_fn(x) + ctx.save_for_backward(x, y) + return y + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]: + x, y = ctx.saved_tensors + safe_x = torch.where(x == 0, torch.ones_like(x), x) + phi = torch.where(x == 0, torch.full_like(x, 0.5), y / safe_x) + return grad_output * phi, None + + +def identity_rule(x: torch.Tensor, act_fn: Callable[[torch.Tensor], torch.Tensor]) -> torch.Tensor: + """Apply the Identity-rule for an elementwise activation with ``f(0) == 0``.""" + result: torch.Tensor = _IdentityRule.apply(x, act_fn) + return result + + +class _HalfRule(torch.autograd.Function): + """Half-rule: forward is ``u * v``; the VJP halves each ordinary product-rule term.""" + + @staticmethod + def forward(ctx: Any, u: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + ctx.save_for_backward(u, v) + return u * v + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + u, v = ctx.saved_tensors + return 0.5 * grad_output * v, 0.5 * grad_output * u + + +def half_rule(u: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """Apply the Half-rule: native product forward, evenly split backward.""" + result: torch.Tensor = _HalfRule.apply(u, v) + return result + + +class _ScaleGradient(torch.autograd.Function): + """Identity forward; the VJP scales the incoming gradient by a constant factor.""" + + @staticmethod + def forward(ctx: Any, x: torch.Tensor, factor: float) -> torch.Tensor: + ctx.factor = factor + return x + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]: + return ctx.factor * grad_output, None + + +def scale_gradient(x: torch.Tensor, factor: float) -> torch.Tensor: + """Pass ``x`` through unchanged while scaling its gradient by ``factor``. + + The Half-rule on a product ``u * v`` halves each ordinary product-rule term, + which is the same as halving the single gradient that enters the product before + it splits. When the product is computed inside an opaque module the bridge cannot + reach term by term (its native forward is called as one unit), scaling the + gradient entering the product by ``0.5`` reproduces the Half-rule at that point + without altering the native forward value. + """ + result: torch.Tensor = _ScaleGradient.apply(x, factor) + return result + + +class RelevanceRuleConflictError(RuntimeError): + """A hook would silently break a rule-active forward/backward invariant. + + Raised instead of the ordinary warn-and-fall-back a component would use when + no rule is active, since falling back while a rule is active would compose the + rule with the hook edit and break the bit-identical-forward guarantee. + """ + + +class RelevanceRuleUnsupportedError(RuntimeError): + """A requested relevance rule cannot be installed on an otherwise-capable component. + + Raised at ``use_relevance_rules`` entry, before any forward or backward pass, when + a component reports the requested kind in its own ``_relevance_rule_unsupported_kinds`` + -- for example a gated-MLP recompute path backed by an unrecognized weight-orientation + class, or an activation form the Identity-rule does not support. Distinct from a + kind that is simply absent from ``_relevance_rule_kinds`` without being named there + (reported ``skipped``, not raised): that covers a component not implementing the + protocol at all, or one whose mount genuinely never deals with the kind (for example + normalization on a dispatch path the LN-rule does not wrap), both benign + non-applicability rather than a rule request the component was expected to honor. + """ + + +@dataclasses.dataclass(frozen=True) +class RelevanceRules: + """Which relevance rules to request for the duration of a ``use_relevance_rules`` scope. + + Each field names a rule kind. Setting it ``True`` requests that rule wherever a + component at that kind's canonical mount point implements ``_RelevanceRuleCapable``. + Unset fields (the default) leave the corresponding components untouched. + """ + + normalization: bool = False + activation: bool = False + multiplicative_gate: bool = False + attention: bool = False + + +@dataclasses.dataclass +class RelevanceRuleCoverage: + """Which canonical mounts a ``use_relevance_rules`` scope installed versus skipped. + + ``installed`` holds the dotted path of every mount where a requested rule kind was + actually enabled. ``skipped`` holds the dotted path of every mount that matched a + requested kind's canonical mount name but did not implement the relevance-rule + protocol there, so no rule could be installed. + """ + + installed: Tuple[str, ...] + skipped: Tuple[str, ...] + + +@runtime_checkable +class _RelevanceRuleCapable(Protocol): + """Structural contract a component must satisfy to accept a relevance rule. + + ``_relevance_rule_kinds`` names every ``RelevanceRules`` field the component + answers to at its current mount -- a gated-MLP node answers to both + "activation" (Identity-rule on its activation function) and + "multiplicative_gate" (Half-rule on its gate*up product) independently, since + either can be requested without the other. ``_enable_relevance_rule``/ + ``_disable_relevance_rule`` take the specific kind being toggled and touch + only that kind's state, without touching model configuration, so the + component's own state is the only thing that changes and only for the + scope's duration. + + A component may optionally also define ``_relevance_rule_unsupported_kinds`` + (a ``Tuple[str, ...]``, not part of this structural protocol so components that + omit it stay isinstance-compatible) naming kinds it is expected to honor at its + mount but currently cannot -- ``use_relevance_rules`` raises + ``RelevanceRuleUnsupportedError`` for those instead of reporting them skipped. + """ + + _relevance_rule_kinds: Tuple[str, ...] + + def _enable_relevance_rule(self, kind: str) -> None: + ... + + def _disable_relevance_rule(self, kind: str) -> None: + ... + + +# Canonical mount name per rule kind. Targeting is positional: a component is only +# considered for a kind when it sits at that kind's mount name, never by isinstance, +# so a same-class component mounted elsewhere (for example a q_norm sharing +# NormalizationBridge's class) is left untouched. The normalization kind lists both +# the pre-norm mounts (ln1, ln2) and the sandwich post-norm mounts (ln1_post, +# ln2_post) so the LN-rule reaches the post-attention/post-MLP norms that +# sandwich-norm architectures mount there, matching the pinned RelP reference. +_CANONICAL_MOUNTS: Mapping[str, Tuple[str, ...]] = { + "normalization": ("ln1", "ln2", "ln1_post", "ln2_post"), + "activation": ("mlp",), + "multiplicative_gate": ("mlp",), +} + + +def _iter_canonical_mount_candidates( + model: nn.Module, mount_names: Tuple[str, ...] +) -> Iterator[Tuple[str, _RelevanceRuleCapable]]: + """Yield each distinct module reachable at one of ``mount_names``, once. + + A bridge component reachable at a canonical mount name (for example + ``blocks.0.ln1``) is also reachable, under the same parent, through the + raw HF module tree the bridge wraps in place (for example + ``blocks.0._original_component.input_layernorm``) -- both names resolve to + the identical object. ``nn.Module.named_modules()`` deduplicates by object + identity and keeps only whichever path it visits first, which is the raw + HF-attribute path (registered before the canonical alias), so on a real + assembled model the canonical name is silently never seen. Walking with + ``remove_duplicate=False`` restores every path so the canonical name is + visible, and picking the fewest-dot-separated-segments path per object + (breaking a tie between two paths that both happen to end in a mount name, + such as ``mlp``, which HF's own attribute name also frequently matches) + reports the shallower, canonical-looking path rather than an internal one. + """ + best_by_id: Dict[int, Tuple[str, _RelevanceRuleCapable]] = {} + for name, module in model.named_modules(remove_duplicate=False): + if name.rsplit(".", 1)[-1] not in mount_names: + continue + existing = best_by_id.get(id(module)) + if existing is None or name.count(".") < existing[0].count("."): + best_by_id[id(module)] = (name, module) + yield from best_by_id.values() + + +def _acquire_rule(module: _RelevanceRuleCapable, kind: str) -> None: + """Enable ``module``'s ``kind`` rule only on the outermost scope that requests it. + + Refcounted per kind, not per module: a gated-MLP node can have its + "activation" rule and "multiplicative_gate" rule independently nested to + different depths, so one kind's inner exit must never disable the other. + """ + counts: Dict[str, int] = getattr(module, "_relevance_rule_refcounts", None) or {} + count = counts.get(kind, 0) + if count == 0: + module._enable_relevance_rule(kind) + counts[kind] = count + 1 + setattr(module, "_relevance_rule_refcounts", counts) + + +def _release_rule(module: _RelevanceRuleCapable, kind: str) -> None: + """Disable ``module``'s ``kind`` rule only once its innermost scope exits.""" + counts: Dict[str, int] = getattr(module, "_relevance_rule_refcounts", None) or {} + count = counts.get(kind, 0) - 1 + counts[kind] = max(count, 0) + setattr(module, "_relevance_rule_refcounts", counts) + if count <= 0: + module._disable_relevance_rule(kind) + + +@contextmanager +def use_relevance_rules(model: nn.Module, rules: RelevanceRules) -> Iterator[RelevanceRuleCoverage]: + """Install the requested relevance rules on ``model`` only for this scope. + + Targeting is positional: a component is considered for a rule kind only when it + sits at that kind's canonical mount name (never by class). A canonical mount + occupied by a component that does not implement ``_RelevanceRuleCapable``, or + whose ``_relevance_rule_kinds`` simply omits the requested kind, is reported as + skipped -- both are benign non-applicability, covering a structurally different + architecture or a mount whose current dispatch path the rule does not wrap. A + component that additionally names the requested kind in its own + ``_relevance_rule_unsupported_kinds`` raises ``RelevanceRuleUnsupportedError`` + instead: that names a kind the component is expected to honor at this mount but + cannot given its current configuration, so silently skipping it would let + analysis proceed as if the caller had never asked. Requesting a kind that has no + canonical mount at all (for example "attention", a valid ``RelevanceRules`` field + with no mount defined) raises ``ValueError`` before the install loop, since there + is no mount to target and the scope would otherwise return empty coverage as + though the request had succeeded. Scopes over the same model + are reference-counted, so an inner scope's exit never disables a rule an outer + scope still needs. No model configuration is mutated; the only state that + changes lives on the participating components, and only for the scope's + duration. + """ + requested_kinds = [ + field.name for field in dataclasses.fields(rules) if getattr(rules, field.name) + ] + + unmapped_kinds = [kind for kind in requested_kinds if kind not in _CANONICAL_MOUNTS] + if unmapped_kinds: + joined = ", ".join(repr(kind) for kind in unmapped_kinds) + raise ValueError( + f"No canonical mount is defined for relevance-rule kind(s) {joined}. " + "Such a kind has no mount to target, so it would install nothing and " + "silently return an empty coverage report instead of applying the rule." + ) + + installed: List[Tuple[str, _RelevanceRuleCapable, str]] = [] + skipped: List[str] = [] + for kind in requested_kinds: + mount_names = _CANONICAL_MOUNTS.get(kind, ()) + for name, module in _iter_canonical_mount_candidates(model, mount_names): + if isinstance(module, _RelevanceRuleCapable) and kind in module._relevance_rule_kinds: + installed.append((name, module, kind)) + continue + unsupported_kinds = getattr(module, "_relevance_rule_unsupported_kinds", ()) + if isinstance(module, _RelevanceRuleCapable) and kind in unsupported_kinds: + raise RelevanceRuleUnsupportedError( + f"{name!r} ({type(module).__name__}) cannot install the {kind!r} " + "relevance rule: unsupported configuration for this component." + ) + skipped.append(name) + + for _, module, kind in installed: + _acquire_rule(module, kind) + try: + yield RelevanceRuleCoverage( + installed=tuple(name for name, _, _ in installed), + skipped=tuple(skipped), + ) + finally: + for _, module, kind in installed: + _release_rule(module, kind) diff --git a/transformer_lens/model_bridge/generalized_components/gated_mlp.py b/transformer_lens/model_bridge/generalized_components/gated_mlp.py index c6d969baf4..ccc533c726 100644 --- a/transformer_lens/model_bridge/generalized_components/gated_mlp.py +++ b/transformer_lens/model_bridge/generalized_components/gated_mlp.py @@ -2,28 +2,57 @@ This module contains the bridge component for gated MLP layers (e.g., LLaMA, Gemma). """ -from typing import Any, Callable, Dict, Mapping, Optional +from typing import Any, Callable, Dict, Mapping, Optional, Tuple, cast import torch +import torch.nn as nn +from transformer_lens.model_bridge._relevance_rules import ( + half_rule, + identity_rule, + scale_gradient, +) from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, ) from transformer_lens.model_bridge.generalized_components.mlp import MLPBridge +def _resolve_activation_fn_name(config: Any) -> Optional[str]: + """The raw activation-name attribute a config exposes, in adapter priority order.""" + if config is None: + return None + for attr in ("activation_function", "hidden_activation", "hidden_act", "act_fn"): + name = getattr(config, attr, None) + if name is not None: + return str(name) + return None + + +_IDENTITY_RULE_UNSUPPORTED_ACTIVATIONS = {"relu", "relu2", "relu_2", "relu_squared"} + + +def identity_rule_supports_activation(config: Any) -> bool: + """Whether the config's resolved activation form is safe for the Identity-rule. + + The Identity-rule's backward multiplier is ``f(x) / x`` (the removable-singularity + limit filled in at zero) rather than the ordinary derivative -- the correct + LRP-style rule for SiLU and both GELU variants, but not for the relu family: + relu-squared's ratio reduces to ``relu(x)``, not its true derivative + ``2 * relu(x)``, and plain relu has no smooth two-sided derivative for the ratio + to represent at the removable singularity either. Both are therefore excluded + rather than silently applying a rule that does not hold for them. + """ + return _resolve_activation_fn_name(config) not in _IDENTITY_RULE_UNSUPPORTED_ACTIVATIONS + + def resolve_activation_fn(config: Any) -> Callable: """Resolve activation function from a model config. Checks config attributes in order: activation_function, hidden_activation, hidden_act, act_fn. Maps common aliases to torch.nn.functional callables. """ - act_fn_name = None - if config is not None: - for attr in ("activation_function", "hidden_activation", "hidden_act", "act_fn"): - act_fn_name = getattr(config, attr, None) - if act_fn_name is not None: - break + act_fn_name = _resolve_activation_fn_name(config) if act_fn_name is None or act_fn_name in ("silu", "swish"): return torch.nn.functional.silu @@ -46,6 +75,25 @@ def relu_squared(x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.silu +class _IdentityRuleActivation(nn.Module): + """Route a wrapped activation through the Identity-rule for a scope's duration. + + The opaque gated-MLP path keeps the HF module's own forward intact and installs + the Identity-rule by swapping the module's activation callable for this wrapper. + Its forward returns ``act_fn(x)`` unchanged, so the native forward value is + preserved, while the backward follows the Identity-rule VJP. Storing the wrapped + activation as an attribute registers it as a child module when it is itself an + ``nn.Module`` (the common ``ACT2FN`` case), so its parameters, if any, stay live. + """ + + def __init__(self, wrapped: Callable[[torch.Tensor], torch.Tensor]): + super().__init__() + self._wrapped_activation = wrapped + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return identity_rule(x, self._wrapped_activation) + + class GatedMLPBridge(MLPBridge): """Bridge component for gated MLP layers. @@ -84,6 +132,188 @@ def __init__( optional: If True, setup skips this bridge when absent (hybrid architectures). """ super().__init__(name, config, submodules=submodules or {}, optional=optional) + self._relevance_rule_activation_active = False + self._relevance_rule_gate_active = False + # Opaque-path rule installers hold their teardown state here. The activation + # wrap records (attr_name, original_value, was_child_module) so the swapped + # activation callable can be restored exactly; the gate handle is the + # forward-pre-hook that scales the gradient entering the down projection. + self._relevance_activation_wrap: Optional[Tuple[str, Any, bool]] = None + self._relevance_gate_hook_handle: Optional[Any] = None + + def _is_gated_mlp_shaped(self) -> bool: + """Whether this instance has the gate/up/down submodules a gated MLP needs. + + A container missing one of these was never wired up as a gated-MLP node at + all (a different architecture at this mount), which is benign + non-applicability rather than an unsupported configuration of a gated-MLP + node -- unlike an activation form the Identity-rule cannot honor, which + occupies exactly this node's shape but cannot be honored correctly. + """ + if self.original_component is None: + return False + gate_module = getattr(self, "gate", None) + in_module = getattr(self, "in", None) + out_module = getattr(self, "out", None) + return gate_module is not None and in_module is not None and out_module is not None + + def _find_activation_attr(self) -> Optional[str]: + """The attribute name under which the HF module holds its activation callable. + + The opaque path installs the Identity-rule by swapping this attribute, so the + activation must be reachable as a callable attribute the native forward calls + (the ``ACT2FN`` module the gated-MLP families store as ``act_fn``). Returns + ``None`` when no such attribute exists, in which case the Identity-rule cannot + be wrapped in and ``"activation"`` is reported unsupported rather than + installed as a silent no-op. + """ + component = self.original_component + if component is None: + return None + for attr in ("act_fn", "activation_fn", "act", "activation"): + if callable(getattr(component, attr, None)): + return attr + return None + + def _activation_rule_installable(self) -> bool: + """Whether the Identity-rule can be installed on this instance's activation. + + Requires both a config activation form the ratio rule is valid for (the + relu family is excluded) and, on the opaque path, an activation callable the + bridge can wrap in place. Subclasses that reconstruct the forward themselves + override this, since they call the activation directly and never wrap it. + """ + return ( + identity_rule_supports_activation(self.config) + and self._find_activation_attr() is not None + ) + + @property + def _relevance_rule_kinds(self) -> Tuple[str, ...]: + """The relevance-rule kinds this instance can currently honor. + + Empty when this is not a gated-MLP-shaped node. Otherwise always includes + ``"multiplicative_gate"`` (the Half-rule is a gradient scale at the gate*up + product and needs no weight or activation access) and includes + ``"activation"`` only when the Identity-rule can be installed on the + configured activation, so a relu-family activation, or one the bridge cannot + reach to wrap, is excluded. + """ + if not self._is_gated_mlp_shaped(): + return () + kinds: Tuple[str, ...] = ("multiplicative_gate",) + if self._activation_rule_installable(): + kinds = ("activation",) + kinds + return kinds + + @property + def _relevance_rule_unsupported_kinds(self) -> Tuple[str, ...]: + """Kinds this gated-MLP node is expected to honor but currently cannot. + + Unlike a kind simply absent from ``_relevance_rule_kinds`` because this is + not a gated-MLP-shaped node at all (benign non-applicability, reported + skipped), a gated-MLP node whose activation form or activation callable the + Identity-rule cannot honor is exactly the kind of component a caller expects + the rule to work on. Requesting ``"activation"`` there raises instead of + silently reporting the mount skipped. The Half-rule applies to every + gated-MLP node, so ``"multiplicative_gate"`` is never reported unsupported. + """ + if not self._is_gated_mlp_shaped(): + return () + if "activation" not in self._relevance_rule_kinds: + return ("activation",) + return () + + def _install_activation_rule(self) -> None: + """Swap the HF module's activation callable for the Identity-rule wrapper. + + No-op when the activation callable cannot be located; requesting the + activation rule in that case is refused earlier through + ``_relevance_rule_unsupported_kinds``. The original value and whether it was + a registered child module are recorded so teardown restores it exactly. + """ + component = self.original_component + attr = self._find_activation_attr() + if component is None or attr is None: + return + was_child_module = attr in component._modules + original = component._modules[attr] if was_child_module else getattr(component, attr, None) + # _find_activation_attr only returns an attribute whose value is callable. + wrapper = _IdentityRuleActivation(cast(Callable[[torch.Tensor], torch.Tensor], original)) + if not was_child_module: + component.__dict__.pop(attr, None) + component._modules[attr] = wrapper + self._relevance_activation_wrap = (attr, original, was_child_module) + + def _teardown_activation_rule(self) -> None: + """Restore the activation callable swapped in by ``_install_activation_rule``.""" + if self._relevance_activation_wrap is None: + return + attr, original, was_child_module = self._relevance_activation_wrap + component = self.original_component + if component is not None: + component._modules.pop(attr, None) + if was_child_module: + component._modules[attr] = original + else: + component.__dict__[attr] = original + self._relevance_activation_wrap = None + + def _install_gate_rule(self) -> None: + """Halve the gradient entering the down projection to reproduce the Half-rule. + + The gate*up product is the down projection's input, so a forward-pre-hook + that routes that input through ``scale_gradient(..., 0.5)`` halves the single + gradient feeding the product before it splits, which matches halving both + product-rule terms. The native forward value is unchanged, and the down + projection's own weight gradient stays ordinary because it is taken against + the unscaled downstream gradient. + """ + out_module = getattr(self, "out", None) + down_component = getattr(out_module, "original_component", None) + if down_component is None: + return + + def _scale_product_gradient( + module: nn.Module, args: Tuple[Any, ...] + ) -> Optional[Tuple[Any, ...]]: + if not args: + return None + return (scale_gradient(args[0], 0.5),) + tuple(args[1:]) + + self._relevance_gate_hook_handle = down_component.register_forward_pre_hook( + _scale_product_gradient + ) + + def _teardown_gate_rule(self) -> None: + """Remove the down-projection gradient-scale hook.""" + if self._relevance_gate_hook_handle is not None: + self._relevance_gate_hook_handle.remove() + self._relevance_gate_hook_handle = None + + def _enable_relevance_rule(self, kind: str) -> None: + """Activate the named rule and install its opaque-path hook. + + The boolean flag drives the reconstructed forward paths (compatibility mode + here, and the inline forward of subclasses that override it). The install + step additionally attaches the rule to the live HF submodules for the opaque + native forward, which a flag alone cannot alter. + """ + if kind == "activation": + self._relevance_rule_activation_active = True + self._install_activation_rule() + elif kind == "multiplicative_gate": + self._relevance_rule_gate_active = True + self._install_gate_rule() + + def _disable_relevance_rule(self, kind: str) -> None: + """Deactivate the named rule and tear down its opaque-path hook.""" + if kind == "activation": + self._teardown_activation_rule() + self._relevance_rule_activation_active = False + elif kind == "multiplicative_gate": + self._teardown_gate_rule() + self._relevance_rule_gate_active = False def forward(self, *args, **kwargs) -> torch.Tensor: """Forward pass through the gated MLP bridge. @@ -120,8 +350,16 @@ def forward(self, *args, **kwargs) -> torch.Tensor: if in_module is not None and hasattr(in_module, "hook_out"): linear_output = in_module.hook_out(linear_output) # type: ignore[misc] act_fn = resolve_activation_fn(self.config) - activated = act_fn(gate_output) - hidden = activated * linear_output + activated = ( + identity_rule(gate_output, act_fn) + if self._relevance_rule_activation_active + else act_fn(gate_output) + ) + hidden = ( + half_rule(activated, linear_output) + if self._relevance_rule_gate_active + else activated * linear_output + ) if hasattr(self, "out") and hasattr(self.out, "hook_in"): hidden = self.out.hook_in(hidden) output = torch.nn.functional.linear( @@ -140,6 +378,10 @@ def forward(self, *args, **kwargs) -> torch.Tensor: hidden_states = args[0] hidden_states = self.hook_in(hidden_states) new_args = (hidden_states,) + args[1:] + # The active relevance rules are attached to the live HF submodules (the + # activation callable and the down projection) by _enable_relevance_rule, + # so the native forward runs unchanged and its own internal hooks, gate + # multipliers, and activation sparsity all remain in the backward graph. output = self.original_component(*new_args, **kwargs) output = self.hook_out(output) return output diff --git a/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py b/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py index 7d40d57afd..28a1981283 100644 --- a/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py +++ b/transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py @@ -6,11 +6,13 @@ import torch +from transformer_lens.model_bridge._relevance_rules import half_rule, identity_rule from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, ) from transformer_lens.model_bridge.generalized_components.gated_mlp import ( GatedMLPBridge, + identity_rule_supports_activation, resolve_activation_fn, ) from transformer_lens.model_bridge.generalized_components.linear import LinearBridge @@ -166,6 +168,29 @@ def _resolve_activation_fn(self) -> Callable: return self._activation_fn return resolve_activation_fn(self.config) + def _activation_rule_installable(self) -> bool: + """The reconstructed forward calls the activation itself, so only the config + activation form gates the Identity-rule; the opaque-path requirement of a + wrappable activation callable does not apply here.""" + return identity_rule_supports_activation(self.config) + + # The reconstructed forward applies both rules inline off the boolean flags set + # by the base ``_enable_relevance_rule``/``_disable_relevance_rule``. The + # opaque-path installers must stay disabled: the gate hook lives on the shared + # down projection this forward also calls, so leaving it active would halve the + # gate*up gradient a second time on top of the inline ``half_rule``. + def _install_activation_rule(self) -> None: + return None + + def _teardown_activation_rule(self) -> None: + return None + + def _install_gate_rule(self) -> None: + return None + + def _teardown_gate_rule(self) -> None: + return None + def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: """Reconstructed gated MLP forward with individual hook access.""" # Delegate to GatedMLPBridge's processed-weights path only when ALL @@ -184,7 +209,16 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: up_output = getattr(self, "in")(hidden_states) act_fn = self._resolve_activation_fn() - gated = act_fn(gate_output) * up_output + activated = ( + identity_rule(gate_output, act_fn) + if self._relevance_rule_activation_active + else act_fn(gate_output) + ) + gated = ( + half_rule(activated, up_output) + if self._relevance_rule_gate_active + else activated * up_output + ) if hasattr(self, "out") and self.out is not None: output = self.out(gated) diff --git a/transformer_lens/model_bridge/generalized_components/normalization.py b/transformer_lens/model_bridge/generalized_components/normalization.py index 418cb3edb9..5aa0b59bc9 100644 --- a/transformer_lens/model_bridge/generalized_components/normalization.py +++ b/transformer_lens/model_bridge/generalized_components/normalization.py @@ -1,11 +1,15 @@ """Normalization bridge component implementation.""" import contextlib import warnings -from typing import Any, ContextManager, Dict, Optional, cast +from typing import Any, ContextManager, Dict, Optional, Tuple, cast import torch from transformer_lens.hook_points import HookPoint +from transformer_lens.model_bridge._relevance_rules import ( + RelevanceRuleConflictError, + ln_rule_grad, +) from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, ) @@ -23,6 +27,83 @@ "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." ) +# While the LN-rule is active, the fallbacks above would silently compose the rule +# with the hook edit and break the bit-identical-forward guarantee, so they raise +# instead of warning. +RULE_ACTIVE_BWD_HOOK_CONFLICT = ( + "Backward hooks on hook_scale/hook_normalized are incompatible with an active " + "LN-rule on '{name}': the rule-wrapped native forward keeps these hook points " + "out of its backward graph, so a backward hook here would silently never fire." +) +RULE_ACTIVE_EDIT_HOOK_CONFLICT = ( + "A forward hook edited hook_scale/hook_normalized while the LN-rule is active on " + "'{name}': honoring the edit would require the python-norm fallback, which would " + "compose the rule with the edit and break the bit-identical-forward guarantee." +) + + +class _NativeLNRuleForward(torch.autograd.Function): + """Wrap a normalization module's own forward call in the LN-rule's VJP. + + Forward returns ``component(x)`` unchanged, so the result is bit-identical to + the native forward by construction. Backward routes the x-path gradient + through the centering op ordinarily but treats ``denom`` as a constant (the + LN-rule), while ``weight`` and ``bias`` receive their ordinary gradient since + the rule only redefines how relevance reaches the input, not parameter + training gradients. A parameter-free norm (``weight`` is ``None``, e.g. + OLMo's ``OlmoLayerNorm``) is treated as a unit scale in both directions. + """ + + @staticmethod + def forward( + ctx: Any, + x_centered: torch.Tensor, + denom: torch.Tensor, + weight: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + x: torch.Tensor, + component: torch.nn.Module, + offset: bool, + input_dtype: torch.dtype, + ) -> torch.Tensor: + ctx.save_for_backward(x_centered, denom, weight) + ctx.has_bias = bias is not None + ctx.bias_requires_grad = bool(bias is not None and bias.requires_grad) + ctx.offset = offset + result = component(x) + if result.dtype != input_dtype: + result = result.to(input_dtype) + return result + + @staticmethod + def backward( + ctx: Any, grad_output: torch.Tensor + ) -> Tuple[ + torch.Tensor, + None, + Optional[torch.Tensor], + Optional[torch.Tensor], + None, + None, + None, + None, + ]: + x_centered, denom, weight = ctx.saved_tensors + if weight is None: + w_eff: torch.Tensor | float = 1.0 + else: + w_eff = (1.0 + weight) if ctx.offset else weight + reduce_dims = tuple(range(grad_output.dim() - 1)) + grad_x_centered = ln_rule_grad(grad_output * w_eff, denom) + grad_weight = ( + (grad_output * (x_centered / denom)).sum(dim=reduce_dims) + if weight is not None and weight.requires_grad + else None + ) + grad_bias = ( + grad_output.sum(dim=reduce_dims) if ctx.has_bias and ctx.bias_requires_grad else None + ) + return grad_x_centered, None, grad_weight, grad_bias, None, None, None, None class NormalizationBridge(GeneralizedComponent): @@ -60,6 +141,31 @@ def __init__( self.hook_scale = HookPoint() self.use_native_layernorm_autograd = use_native_layernorm_autograd self._uses_rms_norm_override = uses_rms_norm + self._relevance_rule_active = False + + @property + def _relevance_rule_kinds(self) -> Tuple[str, ...]: + """``("normalization",)`` only when this instance actually dispatches through + the native-autograd branch the LN-rule wraps (``_hf_autograd_forward_with_hooks``); + empty otherwise, so ``use_relevance_rules`` reports an ln1/ln2 mount that + uses the plain python-norm path (for example ``LayerNormPreBridge`` / + ``RMSNormPreBridge``, or a config without ``layer_norm_folding``) as + skipped rather than silently leaving ordinary gradients in place under a + claimed "installed" rule. + """ + if self.use_native_layernorm_autograd: + return ("normalization",) + if bool(getattr(self.config, "layer_norm_folding", False)): + return ("normalization",) + return () + + def _enable_relevance_rule(self, kind: str) -> None: + """Activate the LN-rule for this instance's native-forward branch only.""" + self._relevance_rule_active = True + + def _disable_relevance_rule(self, kind: str) -> None: + """Deactivate the LN-rule, restoring today's native-forward behavior.""" + self._relevance_rule_active = False @property def uses_rms_norm(self) -> bool: @@ -172,6 +278,10 @@ def _hf_autograd_forward_with_hooks(self, x: torch.Tensor) -> torch.Tensor: _ = self.hook_normalized(x) return x if self.hook_scale.bwd_hooks or self.hook_normalized.bwd_hooks: + if self._relevance_rule_active: + raise RelevanceRuleConflictError( + RULE_ACTIVE_BWD_HOOK_CONFLICT.format(name=self.name) + ) 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) @@ -216,13 +326,64 @@ def _hf_autograd_forward_with_hooks(self, x: torch.Tensor) -> torch.Tensor: # 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 self._relevance_rule_active: + return self._native_forward_with_ln_rule(x, input_dtype) result = self.original_component(x) if result.dtype != input_dtype: result = result.to(input_dtype) return result + if self._relevance_rule_active: + raise RelevanceRuleConflictError(RULE_ACTIVE_EDIT_HOOK_CONFLICT.format(name=self.name)) warnings.warn(NATIVE_PATH_EDIT_FALLBACK_WARNING) return self._apply_weight_and_bias(hooked_normalized, input_dtype) + def _native_forward_with_ln_rule( + self, x: torch.Tensor, input_dtype: torch.dtype + ) -> torch.Tensor: + """Call the native forward unchanged while routing its backward through the LN-rule. + + Recomputes the centered numerator and denominator independently of the + hook-observation pass above (which may run under ``torch.no_grad()``), so + the recompute here stays grad-connected to ``x`` regardless of whether + forward hooks are attached. Autograd then flows from the returned tensor + back through the centering op ordinarily and, at the LN-rule Function + boundary, treats the denominator as a constant -- the same contract as + the shared ``ln_rule`` primitive, without reproducing the native kernel's + own forward numerics. + """ + component = self.original_component + assert component is not None + x_float = x.float() if x.dtype not in (torch.float32, torch.float64) else x + if not self.uses_rms_norm: + x_centered = x_float - x_float.mean(-1, keepdim=True) + else: + x_centered = x_float + eps_tensor = getattr(component, "eps", None) + if eps_tensor is None: + eps_tensor = getattr(component, "variance_epsilon", None) + if eps_tensor is None: + eps_value: float | torch.Tensor = getattr(self.config, "eps", 1e-05) + else: + eps_value = eps_tensor + variance = x_centered.pow(2).mean(-1, keepdim=True) + denom = ( + (variance + eps_value).sqrt() + if isinstance(eps_value, torch.Tensor) + else (variance + float(eps_value)).sqrt() + ) + # A parameter-free norm (OLMo's OlmoLayerNorm) exposes no usable weight; + # pass None so the rule treats the scale as 1 in both forward and backward. + try: + weight: Optional[torch.Tensor] = cast(torch.Tensor, self.weight) + except AttributeError: + weight = None + bias = getattr(component, "bias", None) if not self.uses_rms_norm else None + offset = bool(getattr(self.config, "rmsnorm_uses_offset", False)) + result: torch.Tensor = _NativeLNRuleForward.apply( + x_centered, denom, weight, bias, x, component, offset, input_dtype + ) + return result + class LayerNormPreBridge(NormalizationBridge): """Param-free LayerNorm (LNPre): hook_scale / hook_normalized, no weight or bias."""