From 7d01d2b1dbbbe1b382838b0bcfd5c11477c7bc89 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 14 Sep 2026 12:16:24 +0530 Subject: [PATCH 01/13] test(relevance_rules): closed-form VJP + forward-identity primitives Add analytic tests for the LN-, Identity-, and Half-relevance-rule primitives against their closed forms, asserting both forward-identity and gradient correctness: - LN-rule: forward equals numerator/denom exactly; the VJP treats denom as constant, verified by contrast against plain autodiff (which would also differentiate through denom's dependency on the input). - Identity-rule: forward matches SiLU, exact GELU, and tanh-approximate GELU; the VJP is checked against the sigmoid closed form, the Gaussian CDF closed form, and the direct f(x)/x ratio with its 1/2 limit at x == 0 for the approximate case. - Half-rule: forward equals u * v; the VJP halves each ordinary product-rule term, verified by contrast against plain autodiff. Cases cover zero, negative, mixed-sign, non-contiguous, and batched inputs at both float32 and float64. No model fixtures, backward hooks, or .data access. The primitive module does not exist yet, so this fails collection with a single ModuleNotFoundError -- the expected red state before the primitives are implemented. --- tests/unit/tools/test_relevance_rules.py | 255 +++++++++++++++++++++++ 1 file changed, 255 insertions(+) create mode 100644 tests/unit/tools/test_relevance_rules.py diff --git a/tests/unit/tools/test_relevance_rules.py b/tests/unit/tools/test_relevance_rules.py new file mode 100644 index 0000000000..1f25bce71d --- /dev/null +++ b/tests/unit/tools/test_relevance_rules.py @@ -0,0 +1,255 @@ +"""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.functional as F + +from transformer_lens.model_bridge._relevance_rules import ( + half_rule, + identity_rule, + ln_rule, +) + +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()) From 6ca12f43c01cef89f24e322438c07d0305ea68e5 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 14 Sep 2026 12:25:34 +0530 Subject: [PATCH 02/13] feat(relevance_rules): forward-equivalent rule primitives Implement LN-, Identity-, and Half-rule as torch.autograd.Function primitives. Each reproduces its native forward value exactly and swaps in the rule's closed-form VJP on backward: LN-rule treats the denominator as constant, Identity-rule uses f(x)/x with the 0.5 limit at zero, Half-rule splits the product-rule gradient evenly between operands. --- .../model_bridge/_relevance_rules.py | 79 +++++++++++++++++++ 1 file changed, 79 insertions(+) create mode 100644 transformer_lens/model_bridge/_relevance_rules.py diff --git a/transformer_lens/model_bridge/_relevance_rules.py b/transformer_lens/model_bridge/_relevance_rules.py new file mode 100644 index 0000000000..371cc12d65 --- /dev/null +++ b/transformer_lens/model_bridge/_relevance_rules.py @@ -0,0 +1,79 @@ +"""Forward-equivalent relevance-rule primitives for custom-VJP relevance propagation. + +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. +""" + +from typing import Any, Callable, Tuple + +import torch + + +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 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 From a471b6680db38f45af6a32fb33b03b842cdeff65 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 14 Sep 2026 13:19:01 +0530 Subject: [PATCH 03/13] test(model_bridge): scoped rule context forward-identity + cleanup Add tests for a use_relevance_rules(model, rules) scoped context: - torch.equal forward-identity between native and rule-enabled output on a tiny fixture block, both standalone and nested. - ordinary gradients are restored once the context exits, contrasted against the rule's modified gradient inside the context. - nested contexts are reference-counted: the inner context's exit does not disable a rule the outer context still needs active. - an exception raised inside the context (top-level or nested) leaves no installed rule and no changed forward behavior. - targeting is positional (ln1/ln2), not class-based: a same-class component mounted at q_norm is left untouched. - a component occupying a targeted mount but not implementing the rule protocol is reported as skipped rather than installed or raising. - RelevanceRules defaults every rule kind to False and is frozen. The scoped context does not exist yet, so this fails collection with a single ImportError -- the expected red state before the context is implemented. --- .../unit/model_bridge/test_relevance_rules.py | 226 ++++++++++++++++++ 1 file changed, 226 insertions(+) create mode 100644 tests/unit/model_bridge/test_relevance_rules.py 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..7e56cd5212 --- /dev/null +++ b/tests/unit/model_bridge/test_relevance_rules.py @@ -0,0 +1,226 @@ +"""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, + 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_kind = "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) -> None: + self._rule_active = True + + def _disable_relevance_rule(self) -> 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 _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_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_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 From 6e27ff0efa9d4135c0e0f9aca163e8767edfc0dd Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 14 Sep 2026 13:48:55 +0530 Subject: [PATCH 04/13] feat(model_bridge): scoped relevance-rule context + coverage report Add RelevanceRules (frozen: normalization, activation, multiplicative_gate, attention), RelevanceRuleCoverage (installed, skipped), and use_relevance_rules(model, rules): a scoped context manager that installs each requested rule only on components sitting at that rule kind's canonical mount name, never by isinstance, and reports which mounts were installed versus skipped. Components at a canonical mount that do not implement the relevance-rule protocol are reported skipped rather than raising. Nested scopes over the same model are reference-counted, so an inner scope's exit never disables a rule an outer scope still needs, and state is restored on normal exit and on exception. No model configuration is mutated; the only state that changes lives on the participating components, only for the scope's duration. --- .../model_bridge/_relevance_rules.py | 133 +++++++++++++++++- 1 file changed, 131 insertions(+), 2 deletions(-) diff --git a/transformer_lens/model_bridge/_relevance_rules.py b/transformer_lens/model_bridge/_relevance_rules.py index 371cc12d65..c257db2190 100644 --- a/transformer_lens/model_bridge/_relevance_rules.py +++ b/transformer_lens/model_bridge/_relevance_rules.py @@ -1,12 +1,27 @@ -"""Forward-equivalent relevance-rule primitives for custom-VJP relevance propagation. +"""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. """ -from typing import Any, Callable, Tuple +import dataclasses +from contextlib import contextmanager +from typing import ( + Any, + Callable, + Iterator, + List, + Mapping, + Protocol, + Tuple, + runtime_checkable, +) import torch +import torch.nn as nn class _LNRule(torch.autograd.Function): @@ -77,3 +92,117 @@ 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 + + +@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_kind`` names which ``RelevanceRules`` field the component answers + to; ``_enable_relevance_rule``/``_disable_relevance_rule`` toggle the rule without + touching model configuration, so the component's own state is the only thing that + changes and only for the scope's duration. + """ + + _relevance_rule_kind: str + + def _enable_relevance_rule(self) -> None: + ... + + def _disable_relevance_rule(self) -> 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. +_CANONICAL_MOUNTS: Mapping[str, Tuple[str, ...]] = { + "normalization": ("ln1", "ln2"), +} + + +def _acquire_rule(module: _RelevanceRuleCapable) -> None: + """Enable ``module``'s rule only on the outermost scope that requests it.""" + count = getattr(module, "_relevance_rule_refcount", 0) + if count == 0: + module._enable_relevance_rule() + setattr(module, "_relevance_rule_refcount", count + 1) + + +def _release_rule(module: _RelevanceRuleCapable) -> None: + """Disable ``module``'s rule only once the innermost scope that requested it exits.""" + count = getattr(module, "_relevance_rule_refcount", 0) - 1 + setattr(module, "_relevance_rule_refcount", max(count, 0)) + if count <= 0: + module._disable_relevance_rule() + + +@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 component at a + canonical mount that does not implement ``_RelevanceRuleCapable`` for the + requested kind is reported as skipped rather than installed or raising. 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) + ] + + installed: List[Tuple[str, _RelevanceRuleCapable]] = [] + skipped: List[str] = [] + for kind in requested_kinds: + mount_names = _CANONICAL_MOUNTS.get(kind, ()) + for name, module in model.named_modules(): + if name.rsplit(".", 1)[-1] not in mount_names: + continue + if isinstance(module, _RelevanceRuleCapable) and module._relevance_rule_kind == kind: + installed.append((name, module)) + else: + skipped.append(name) + + for _, module in installed: + _acquire_rule(module) + try: + yield RelevanceRuleCoverage( + installed=tuple(name for name, _ in installed), + skipped=tuple(skipped), + ) + finally: + for _, module in installed: + _release_rule(module) From c97b99271a2da7b4b1f95b8f03dde2e504cab74c Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 14 Sep 2026 14:42:36 +0530 Subject: [PATCH 05/13] feat(model_bridge): native-forward LN-rule on NormalizationBridge Wrap the existing native-autograd no-edit branch's own original_component(x) call in a custom autograd.Function so the rule-active forward stays bit-identical to today's native forward by construction, while backward applies the LN-rule (denominator treated as constant) for the x-path; weight and bias keep their ordinary gradient since the rule only redefines how relevance reaches the input, not parameter training gradients. NormalizationBridge now implements the relevance-rule protocol (_relevance_rule_kind / _enable_relevance_rule / _disable_relevance_rule) so use_relevance_rules can target ln1/ln2 mounts positionally. The reported kind is dynamic: it is empty for any instance that never reaches the native-autograd branch (LayerNormPreBridge/RMSNormPreBridge, or a config without layer_norm_folding), so such a mount is reported skipped instead of silently leaving ordinary gradients in place under a claimed installed rule. A backward hook on hook_scale/hook_normalized, or a forward hook that edits either, now raises RelevanceRuleConflictError while the rule is active instead of the existing warn-and-fallback: falling back would compose the rule with the hook edit and break the bit-identical-forward guarantee. Behavior is unchanged when the rule is inactive. Factor the LN-rule's core VJP (grad_output / denom) into a shared ln_rule_grad helper reused by both the primitive and this integration. --- .../test_normalization_relevance_rule.py | 354 ++++++++++++++++++ .../test_relevance_rule_mount_placeholders.py | 128 +++++++ .../model_bridge/_relevance_rules.py | 21 +- .../generalized_components/normalization.py | 154 +++++++- 4 files changed, 655 insertions(+), 2 deletions(-) create mode 100644 tests/unit/model_bridge/generalized_components/test_normalization_relevance_rule.py create mode 100644 tests/unit/model_bridge/supported_architectures/test_relevance_rule_mount_placeholders.py 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..ae62a2633f --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_normalization_relevance_rule.py @@ -0,0 +1,354 @@ +"""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 + +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) + + +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 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..ff2c44d554 --- /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_kind == "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/transformer_lens/model_bridge/_relevance_rules.py b/transformer_lens/model_bridge/_relevance_rules.py index c257db2190..288c45740b 100644 --- a/transformer_lens/model_bridge/_relevance_rules.py +++ b/transformer_lens/model_bridge/_relevance_rules.py @@ -24,6 +24,16 @@ 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.""" @@ -35,7 +45,7 @@ def forward(ctx: Any, numerator: torch.Tensor, denom: torch.Tensor) -> torch.Ten @staticmethod def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]: (denom,) = ctx.saved_tensors - return grad_output / denom, None + return ln_rule_grad(grad_output, denom), None def ln_rule(numerator: torch.Tensor, denom: torch.Tensor) -> torch.Tensor: @@ -94,6 +104,15 @@ def half_rule(u: torch.Tensor, v: torch.Tensor) -> torch.Tensor: 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. + """ + + @dataclasses.dataclass(frozen=True) class RelevanceRules: """Which relevance rules to request for the duration of a ``use_relevance_rules`` scope. diff --git a/transformer_lens/model_bridge/generalized_components/normalization.py b/transformer_lens/model_bridge/generalized_components/normalization.py index 418cb3edb9..84b0009a87 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,79 @@ "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. + """ + + @staticmethod + def forward( + ctx: Any, + x_centered: torch.Tensor, + denom: torch.Tensor, + weight: 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 + 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.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 +137,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_kind(self) -> 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) -> None: + """Activate the LN-rule for this instance's native-forward branch only.""" + self._relevance_rule_active = True + + def _disable_relevance_rule(self) -> 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 +274,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 +322,59 @@ 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() + ) + weight = cast(torch.Tensor, self.weight) + 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.""" From 2193a281d58985f75dd5abdfb9da900a668bb029 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 14 Sep 2026 15:16:35 +0530 Subject: [PATCH 06/13] feat(model_bridge): gated-MLP rule via recompute-from-weights VJP Add the Identity-/Half-rule integration for gated MLPs, the second real consumer of the relevance-rule backend after NormalizationBridge. For the raw (non-fused) GatedMLPBridge used by the Qwen2/Llama/Gemma family, wrap the existing opaque original_component(x) call in a new _GatedMLPRecomputeRule custom autograd.Function. Forward returns that native call's own output unchanged, so torch.equal holds by construction. Backward has no access to the opaque call's internal gate/up/down intermediates, so it recomputes them checkpointing-style from the TL-oriented W_gate/W_in/W_out (the unconditional MLPBridge property accessors, never the compatibility-mode-only _processed_* attributes) and reapplies the Identity-rule to the activation and/or the Half-rule to the gate*up product, per whichever is independently active. Weight and bias gradients are read back off the same recomputed graph via torch.autograd.grad, so they keep their ordinary form -- the rules only redefine how relevance reaches the input, not parameter training gradients. The recompute is allowlisted by the underlying HF module class backing W_gate/W_in (nn.Linear vs Conv1D, via the existing weight_layout_in_out helper), not per adapter, since several adapters share the same backing class and orientation; an unrecognized backing module reports as skipped rather than silently installed. JointGateUpMLPBridge (Phi-3/GLM) already reconstructs its forward in Python as act_fn(gate_output) * up_output through separate gate/up LinearBridge submodules, so its rules attach directly at that multiplication -- no weights-recompute is needed there. A gated-MLP node answers to both "activation" (Identity-rule) and "multiplicative_gate" (Half-rule) independently at the same mlp mount, which the previous single-kind protocol could not express: _RelevanceRuleCapable now reports a tuple of supported kinds via _relevance_rule_kinds, and _enable_relevance_rule/_disable_relevance_rule take the specific kind being toggled. Refcounting in use_relevance_rules is now keyed per (module, kind) rather than per module, so nesting one kind's scope inside the other's never disables the outer kind early. NormalizationBridge is updated to the same (still single-kind) shape. --- .../test_gated_mlp_relevance_rule.py | 267 ++++++++++++++++++ .../test_joint_gate_up_relevance_rule.py | 163 +++++++++++ .../test_relevance_rule_mount_placeholders.py | 2 +- .../unit/model_bridge/test_relevance_rules.py | 76 ++++- tests/unit/tools/test_relevance_rules.py | 92 ++++++ .../model_bridge/_relevance_rules.py | 67 +++-- .../generalized_components/gated_mlp.py | 181 +++++++++++- .../joint_gate_up_mlp.py | 12 +- .../generalized_components/normalization.py | 16 +- 9 files changed, 835 insertions(+), 41 deletions(-) create mode 100644 tests/unit/model_bridge/generalized_components/test_gated_mlp_relevance_rule.py create mode 100644 tests/unit/model_bridge/generalized_components/test_joint_gate_up_relevance_rule.py 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..e43fcaf2b4 --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_gated_mlp_relevance_rule.py @@ -0,0 +1,267 @@ +"""Identity-/Half-rule integration on GatedMLPBridge's recompute-from-weights path. + +Covers the commit-6 contract for the raw (non-fused) gated-MLP path: the opaque +native forward stays bit-identical to today's forward by construction (the wrapping +returns ``original_component(x)`` itself, never a reproduced numeric value), the +backward recomputes gate/up/down from the TL-oriented ``W_gate``/``W_in``/``W_out`` +and applies the Identity-rule to the activation and/or the Half-rule to the gate*up +product per whichever is independently active, weight and bias keep their ordinary +gradient, and the recompute is allowlisted by the underlying HF module class +(``nn.Linear`` and ``Conv1D``) rather than reported installed for an unrecognized +backing module. +""" + +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, + 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: + hidden_act = "silu" + + +class _TinyGatedMLP(nn.Module): + """Mirrors the Qwen2/Llama/Gemma gated-MLP structure: opaque single call.""" + + def __init__(self, gate_proj, up_proj, down_proj): + super().__init__() + self.gate_proj = gate_proj + self.up_proj = up_proj + self.down_proj = down_proj + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +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), + ) + return ( + Conv1D(d_mlp, d_model), + Conv1D(d_mlp, d_model), + Conv1D(d_model, d_mlp), + ) + + +def _make_bridge(backing_class: str, bias: bool = True) -> tuple[_Block, _TinyGatedMLP]: + gate_proj, up_proj, down_proj = _make_projections(backing_class, bias=bias) + hf_mlp = _TinyGatedMLP(gate_proj, up_proj, down_proj) + + bridge = GatedMLPBridge(name="mlp", config=_Cfg()) + 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(gate_proj) + in_bridge.set_original_component(up_proj) + out_bridge.set_original_component(down_proj) + + return _Block(bridge), hf_mlp + + +BACKING_CLASSES = ["nn.Linear", "Conv1D"] + + +class TestGatedMLPRelevanceRuleCapability: + @pytest.mark.parametrize("backing_class", BACKING_CLASSES) + def test_capable_of_both_kinds_when_allowlisted(self, backing_class): + block, _ = _make_bridge(backing_class) + assert set(block.mlp._relevance_rule_kinds) == {"activation", "multiplicative_gate"} + + def test_not_capable_when_backing_module_is_unrecognized(self): + class _OpaqueProj(nn.Module): + """Neither nn.Linear nor Conv1D -- an unrecognized backing class.""" + + def __init__(self, d_in: int, d_out: int): + super().__init__() + self.weight = nn.Parameter(torch.randn(d_out, d_in)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x @ self.weight.T + + gate_proj = _OpaqueProj(4, 8) + up_proj = _OpaqueProj(4, 8) + down_proj = _OpaqueProj(8, 4) + hf_mlp = _TinyGatedMLP(gate_proj, up_proj, down_proj) + + bridge = GatedMLPBridge(name="mlp", config=_Cfg()) + 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(gate_proj) + in_bridge.set_original_component(up_proj) + out_bridge.set_original_component(down_proj) + block = _Block(bridge) + + assert bridge._relevance_rule_kinds == () + + x = torch.randn(2, 4, requires_grad=True) + baseline = block(x) + with use_relevance_rules( + block, RelevanceRules(activation=True, multiplicative_gate=True) + ) as coverage: + assert set(coverage.skipped) == {"mlp"} + assert coverage.installed == () + active = block(x) + assert torch.equal(active, 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) + + @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) + before = block(x) + after = block(x) + assert torch.equal(before, after) + + +class TestGatedMLPRelevanceRuleVJP: + def _oracle_grads(self, hf_mlp, x, activation_active, gate_active): + # hf_mlp's own parameters are shared with the bridge under test, so a prior + # backward through the bridge already left gradients on them; without + # resetting first, this second backward would accumulate on top instead of + # producing an independently comparable oracle. + 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, 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() + grads = {n: p.grad.clone() for n, p in hf_mlp.named_parameters()} + return x_oracle.grad.clone(), grads + + @pytest.mark.parametrize("backing_class", BACKING_CLASSES) + def test_both_rules_active_matches_manually_composed_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._oracle_grads( + hf_mlp, x, activation_active=True, gate_active=True + ) + 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._oracle_grads( + hf_mlp, x, activation_active=True, gate_active=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._oracle_grads( + hf_mlp, x, activation_active=False, gate_active=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) + plain = hf_mlp(x_plain) + plain.sum().backward() + + torch.testing.assert_close(grad_x, x_plain.grad) + for name, grad in grads.items(): + torch.testing.assert_close(grad, dict(hf_mlp.named_parameters())[name].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._oracle_grads( + hf_mlp, x, activation_active=True, gate_active=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) 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..7834a2d9c2 --- /dev/null +++ b/tests/unit/model_bridge/generalized_components/test_joint_gate_up_relevance_rule.py @@ -0,0 +1,163 @@ +"""Identity-/Half-rule integration on JointGateUpMLPBridge's reconstructed forward. + +Unlike the raw GatedMLPBridge path (recompute-from-weights, 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 -- no +weights-recompute Function is needed 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/supported_architectures/test_relevance_rule_mount_placeholders.py b/tests/unit/model_bridge/supported_architectures/test_relevance_rule_mount_placeholders.py index ff2c44d554..4e9d800089 100644 --- 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 @@ -110,7 +110,7 @@ def test_ln1_and_ln2_are_genuine_rule_targets(self): component = block_submodules[key] assert isinstance(component, NormalizationBridge) assert isinstance(component, _RelevanceRuleCapable) - assert component._relevance_rule_kind == "normalization" + 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 diff --git a/tests/unit/model_bridge/test_relevance_rules.py b/tests/unit/model_bridge/test_relevance_rules.py index 7e56cd5212..adbcc3f002 100644 --- a/tests/unit/model_bridge/test_relevance_rules.py +++ b/tests/unit/model_bridge/test_relevance_rules.py @@ -29,7 +29,7 @@ class _FakeNormComponent(nn.Module): """A minimal ln1/ln2-style target: installs the real LN-rule primitive on request.""" - _relevance_rule_kind = "normalization" + _relevance_rule_kinds = ("normalization",) def __init__(self, eps: float = 1e-2): super().__init__() @@ -40,10 +40,10 @@ def __init__(self, eps: float = 1e-2): self.eps = eps self._rule_active = False - def _enable_relevance_rule(self) -> None: + def _enable_relevance_rule(self, kind: str) -> None: self._rule_active = True - def _disable_relevance_rule(self) -> None: + def _disable_relevance_rule(self, kind: str) -> None: self._rule_active = False def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -89,6 +89,38 @@ def _tiny_block() -> _TinyBlock: return block +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.""" @@ -224,3 +256,41 @@ def test_exception_during_nested_context_restores_outer_state(): 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 index 1f25bce71d..202619abe1 100644 --- a/tests/unit/tools/test_relevance_rules.py +++ b/tests/unit/tools/test_relevance_rules.py @@ -12,13 +12,19 @@ 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] @@ -253,3 +259,89 @@ def test_batched_leading_dims(self, dtype): 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 TestGatedMLPRecomputeOrientation: + """The gated-MLP recompute VJP reads ``W_gate``/``W_in`` through the same + ``weight_layout_in_out``/``normalize_mlp_weight`` helpers ``MLPBridge`` uses, so a + fresh orientation bug in the recompute path (as opposed to the already-tested + accessor) would silently transpose the gate/up product for one backing class. + Covers the two HF module classes the recompute allowlists: ``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) + + def test_recompute_via_tl_weights_matches_native_forward_and_rule_vjp(self, backing_class): + gate_proj, up_proj, down_proj = self._make_gate_up_down(backing_class) + x = torch.randn(2, 3, requires_grad=True) + + 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") + + gate_output = x @ w_gate + gate_proj.bias + up_output = x @ w_in + up_proj.bias + activated = identity_rule(gate_output, F.silu) + gated = half_rule(activated, up_output) + down = gated @ w_out + down_proj.bias + + native_down = down_proj(F.silu(gate_proj(x.detach())) * up_proj(x.detach())) + assert torch.allclose(down, native_down, atol=1e-6) + + (grad_x,) = torch.autograd.grad(down.sum(), x) + + # Oracle: the exact same composition, expressed with ordinary autodiff + # through the rule primitives but built from each class's own native call + # instead of the TL-oriented matmul above -- if the recompute transposed + # either weight incorrectly, this would disagree. + x_oracle = x.detach().clone().requires_grad_(True) + gate_output_oracle = gate_proj(x_oracle) + up_output_oracle = up_proj(x_oracle) + activated_oracle = identity_rule(gate_output_oracle, F.silu) + gated_oracle = half_rule(activated_oracle, up_output_oracle) + down_oracle = down_proj(gated_oracle) + (grad_x_oracle,) = torch.autograd.grad(down_oracle.sum(), x_oracle) + + torch.testing.assert_close(grad_x, grad_x_oracle) diff --git a/transformer_lens/model_bridge/_relevance_rules.py b/transformer_lens/model_bridge/_relevance_rules.py index 288c45740b..a193551abd 100644 --- a/transformer_lens/model_bridge/_relevance_rules.py +++ b/transformer_lens/model_bridge/_relevance_rules.py @@ -12,6 +12,7 @@ from typing import ( Any, Callable, + Dict, Iterator, List, Mapping, @@ -146,18 +147,23 @@ class RelevanceRuleCoverage: class _RelevanceRuleCapable(Protocol): """Structural contract a component must satisfy to accept a relevance rule. - ``_relevance_rule_kind`` names which ``RelevanceRules`` field the component answers - to; ``_enable_relevance_rule``/``_disable_relevance_rule`` toggle the rule without - touching model configuration, so the component's own state is the only thing that - changes and only for the scope's duration. + ``_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. """ - _relevance_rule_kind: str + _relevance_rule_kinds: Tuple[str, ...] - def _enable_relevance_rule(self) -> None: + def _enable_relevance_rule(self, kind: str) -> None: ... - def _disable_relevance_rule(self) -> None: + def _disable_relevance_rule(self, kind: str) -> None: ... @@ -167,23 +173,34 @@ def _disable_relevance_rule(self) -> None: # NormalizationBridge's class) is left untouched. _CANONICAL_MOUNTS: Mapping[str, Tuple[str, ...]] = { "normalization": ("ln1", "ln2"), + "activation": ("mlp",), + "multiplicative_gate": ("mlp",), } -def _acquire_rule(module: _RelevanceRuleCapable) -> None: - """Enable ``module``'s rule only on the outermost scope that requests it.""" - count = getattr(module, "_relevance_rule_refcount", 0) +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() - setattr(module, "_relevance_rule_refcount", count + 1) + module._enable_relevance_rule(kind) + counts[kind] = count + 1 + setattr(module, "_relevance_rule_refcounts", counts) -def _release_rule(module: _RelevanceRuleCapable) -> None: - """Disable ``module``'s rule only once the innermost scope that requested it exits.""" - count = getattr(module, "_relevance_rule_refcount", 0) - 1 - setattr(module, "_relevance_rule_refcount", max(count, 0)) +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() + module._disable_relevance_rule(kind) @contextmanager @@ -203,25 +220,25 @@ def use_relevance_rules(model: nn.Module, rules: RelevanceRules) -> Iterator[Rel field.name for field in dataclasses.fields(rules) if getattr(rules, field.name) ] - installed: List[Tuple[str, _RelevanceRuleCapable]] = [] + installed: List[Tuple[str, _RelevanceRuleCapable, str]] = [] skipped: List[str] = [] for kind in requested_kinds: mount_names = _CANONICAL_MOUNTS.get(kind, ()) for name, module in model.named_modules(): if name.rsplit(".", 1)[-1] not in mount_names: continue - if isinstance(module, _RelevanceRuleCapable) and module._relevance_rule_kind == kind: - installed.append((name, module)) + if isinstance(module, _RelevanceRuleCapable) and kind in module._relevance_rule_kinds: + installed.append((name, module, kind)) else: skipped.append(name) - for _, module in installed: - _acquire_rule(module) + for _, module, kind in installed: + _acquire_rule(module, kind) try: yield RelevanceRuleCoverage( - installed=tuple(name for name, _ in installed), + installed=tuple(name for name, _, _ in installed), skipped=tuple(skipped), ) finally: - for _, module in installed: - _release_rule(module) + 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..6368a05a63 100644 --- a/transformer_lens/model_bridge/generalized_components/gated_mlp.py +++ b/transformer_lens/model_bridge/generalized_components/gated_mlp.py @@ -2,14 +2,18 @@ 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 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.mlp import MLPBridge +from transformer_lens.model_bridge.generalized_components.mlp import ( + MLPBridge, + weight_layout_in_out, +) def resolve_activation_fn(config: Any) -> Callable: @@ -46,6 +50,122 @@ def relu_squared(x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.silu +class _GatedMLPRecomputeRule(torch.autograd.Function): + """Wrap a raw gated-MLP module's own forward call in the Identity-/Half-rule VJP. + + Forward returns ``component(x, ...)`` unchanged, so the result is bit-identical + to the native forward by construction. The opaque call gives no access to its own + internal gate/up/down intermediates, so backward recomputes them from the + TL-oriented ``W_gate``/``W_in``/``W_out`` (checkpointing-style: the values are + rederived here rather than saved from forward) and reapplies the Identity-rule to + the activation and/or the Half-rule to the gate*up product, per whichever of the + two is active. Weight and bias gradients come out of the same recomputed graph, + via ``torch.autograd.grad``, so they keep their ordinary form -- the rules only + redefine how relevance reaches the input, not parameter training gradients. + """ + + @staticmethod + def forward( + ctx: Any, + x: torch.Tensor, + w_gate: torch.Tensor, + b_gate: Optional[torch.Tensor], + w_in: torch.Tensor, + b_in: Optional[torch.Tensor], + w_out: torch.Tensor, + b_out: Optional[torch.Tensor], + act_fn: Callable[[torch.Tensor], torch.Tensor], + activation_rule_active: bool, + gate_rule_active: bool, + component: torch.nn.Module, + extra_args: Tuple[Any, ...], + extra_kwargs: Dict[str, Any], + ) -> torch.Tensor: + ctx.save_for_backward(x, w_gate, w_in, w_out) + ctx.b_gate = b_gate + ctx.b_in = b_in + ctx.b_out = b_out + ctx.act_fn = act_fn + ctx.activation_rule_active = activation_rule_active + ctx.gate_rule_active = gate_rule_active + result: torch.Tensor = component(x, *extra_args, **extra_kwargs) + return result + + @staticmethod + def backward( # type: ignore[override] + ctx: Any, grad_output: torch.Tensor + ) -> Tuple[ + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + None, + None, + None, + None, + None, + None, + ]: + x, w_gate, w_in, w_out = ctx.saved_tensors + b_gate, b_in, b_out = ctx.b_gate, ctx.b_in, ctx.b_out + act_fn = ctx.act_fn + + with torch.enable_grad(): + x_ = x.detach().requires_grad_(True) + w_gate_ = w_gate.detach().requires_grad_(True) + w_in_ = w_in.detach().requires_grad_(True) + w_out_ = w_out.detach().requires_grad_(True) + b_gate_ = b_gate.detach().requires_grad_(True) if b_gate is not None else None + b_in_ = b_in.detach().requires_grad_(True) if b_in is not None else None + b_out_ = b_out.detach().requires_grad_(True) if b_out is not None else None + + 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) + if ctx.activation_rule_active + else act_fn(gate_output) + ) + gated = ( + half_rule(activated, up_output) if ctx.gate_rule_active else activated * up_output + ) + + down = gated @ w_out_ + if b_out_ is not None: + down = down + b_out_ + + leaves = [x_, w_gate_, b_gate_, w_in_, b_in_, w_out_, b_out_] + needed = [leaf for leaf in leaves if leaf is not None] + grads = torch.autograd.grad(down, needed, grad_outputs=grad_output, allow_unused=True) + grad_iter = iter(grads) + results = [next(grad_iter) if leaf is not None else None for leaf in leaves] + grad_x, grad_w_gate, grad_b_gate, grad_w_in, grad_b_in, grad_w_out, grad_b_out = results + + return ( + grad_x, + grad_w_gate, + grad_b_gate, + grad_w_in, + grad_b_in, + grad_w_out, + grad_b_out, + None, + None, + None, + None, + None, + None, + ) + + class GatedMLPBridge(MLPBridge): """Bridge component for gated MLP layers. @@ -84,6 +204,43 @@ 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 + + @property + def _relevance_rule_kinds(self) -> Tuple[str, ...]: + """``("activation", "multiplicative_gate")`` only when the gate/up projections + are backed by an allowlisted HF module class (``nn.Linear`` or ``Conv1D``) in a + recognized orientation; empty otherwise, so an opaque or unrecognized backing + module is reported skipped rather than silently installing a rule the + weights-recompute cannot honor correctly. + """ + if self.original_component is None: + return () + gate_module = getattr(self, "gate", None) + in_module = getattr(self, "in", None) + out_module = getattr(self, "out", None) + if gate_module is None or in_module is None or out_module is None: + return () + if weight_layout_in_out(gate_module) is None or weight_layout_in_out(in_module) is None: + return () + if weight_layout_in_out(out_module) is None: + return () + return ("activation", "multiplicative_gate") + + def _enable_relevance_rule(self, kind: str) -> None: + """Activate the named rule for this instance's recompute path only.""" + if kind == "activation": + self._relevance_rule_activation_active = True + elif kind == "multiplicative_gate": + self._relevance_rule_gate_active = True + + def _disable_relevance_rule(self, kind: str) -> None: + """Deactivate the named rule, restoring today's native-forward behavior.""" + if kind == "activation": + self._relevance_rule_activation_active = False + elif kind == "multiplicative_gate": + self._relevance_rule_gate_active = False def forward(self, *args, **kwargs) -> torch.Tensor: """Forward pass through the gated MLP bridge. @@ -140,7 +297,25 @@ def forward(self, *args, **kwargs) -> torch.Tensor: hidden_states = args[0] hidden_states = self.hook_in(hidden_states) new_args = (hidden_states,) + args[1:] - output = self.original_component(*new_args, **kwargs) + if self._relevance_rule_activation_active or self._relevance_rule_gate_active: + act_fn = resolve_activation_fn(self.config) + output = _GatedMLPRecomputeRule.apply( + hidden_states, + self.W_gate, + getattr(self.gate, "bias", None), + self.W_in, + getattr(getattr(self, "in"), "bias", None), + self.W_out, + getattr(self.out, "bias", None), + act_fn, + self._relevance_rule_activation_active, + self._relevance_rule_gate_active, + self.original_component, + new_args[1:], + kwargs, + ) + else: + 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..a3ebcaf78a 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,6 +6,7 @@ import torch +from transformer_lens.model_bridge._relevance_rules import half_rule, identity_rule from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, ) @@ -184,7 +185,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 84b0009a87..ca088dd0a0 100644 --- a/transformer_lens/model_bridge/generalized_components/normalization.py +++ b/transformer_lens/model_bridge/generalized_components/normalization.py @@ -140,9 +140,9 @@ def __init__( self._relevance_rule_active = False @property - def _relevance_rule_kind(self) -> str: - """ "normalization" only when this instance actually dispatches through the - native-autograd branch the LN-rule wraps (``_hf_autograd_forward_with_hooks``); + 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 @@ -150,16 +150,16 @@ def _relevance_rule_kind(self) -> str: claimed "installed" rule. """ if self.use_native_layernorm_autograd: - return "normalization" + return ("normalization",) if bool(getattr(self.config, "layer_norm_folding", False)): - return "normalization" - return "" + return ("normalization",) + return () - def _enable_relevance_rule(self) -> None: + 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) -> None: + def _disable_relevance_rule(self, kind: str) -> None: """Deactivate the LN-rule, restoring today's native-forward behavior.""" self._relevance_rule_active = False From 0cabf6fe42ecb6039fbc7f806d0e01964ff45942 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 14 Sep 2026 15:48:11 +0530 Subject: [PATCH 07/13] feat(model_bridge): fail-closed coverage on unsupported paths use_relevance_rules previously reported any canonical mount whose component could not honor a requested kind as "skipped", whether the component simply did not implement the relevance-rule protocol at all or implemented it but could not currently honor that specific kind. The two cases need different treatment: a component with no protocol (or a mount whose current dispatch path a rule genuinely does not wrap, such as NormalizationBridge on its python-norm path) is benign non-applicability, but a gated-MLP node at the mlp mount is exactly the kind of component a caller expects either rule to work on, so silently skipping it there would let analysis proceed as if the request had never been made. _RelevanceRuleCapable gains an optional _relevance_rule_unsupported_kinds attribute (not part of the structural protocol, so components that omit it stay isinstance-compatible) naming kinds a component is expected to honor at its mount but currently cannot. use_relevance_rules now raises the new RelevanceRuleUnsupportedError, naming the component's dotted path, for any requested kind found there, before yielding the coverage report and before any forward or backward pass runs. GatedMLPBridge implements the new attribute for two cases: an unrecognized weight-backing class disqualifies both "activation" and "multiplicative_gate" (the recompute cannot orient an opaque module's weights correctly for either rule), and a relu-family activation disqualifies only "activation" -- the Identity-rule's f(x)/x backward multiplier is the correct LRP-style rule for SiLU and both GELU variants, but relu-squared's ratio reduces to relu(x) rather than 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. The Half-rule is unaffected by activation form, so "multiplicative_gate" stays available on a relu-family activation as long as the weight backing is recognized. resolve_activation_fn's config-name lookup is factored into a shared _resolve_activation_fn_name helper reused by the new identity_rule_supports_activation check. --- .../test_gated_mlp_relevance_rule.py | 41 +++++-- .../unit/model_bridge/test_relevance_rules.py | 16 +++ .../model_bridge/_relevance_rules.py | 52 +++++++-- .../generalized_components/gated_mlp.py | 106 ++++++++++++++---- 4 files changed, 176 insertions(+), 39 deletions(-) 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 index e43fcaf2b4..de90784948 100644 --- 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 @@ -19,6 +19,7 @@ from transformer_lens.model_bridge._relevance_rules import ( RelevanceRules, + RelevanceRuleUnsupportedError, half_rule, identity_rule, use_relevance_rules, @@ -33,7 +34,8 @@ class _Cfg: - hidden_act = "silu" + def __init__(self, hidden_act: str = "silu"): + self.hidden_act = hidden_act class _TinyGatedMLP(nn.Module): @@ -75,11 +77,13 @@ def _make_projections(backing_class: str, d_model: int = 4, d_mlp: int = 8, bias ) -def _make_bridge(backing_class: str, bias: bool = True) -> tuple[_Block, _TinyGatedMLP]: +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) - bridge = GatedMLPBridge(name="mlp", config=_Cfg()) + bridge = GatedMLPBridge(name="mlp", config=_Cfg(hidden_act)) gate_bridge = LinearBridge(name="gate_proj") in_bridge = LinearBridge(name="up_proj") out_bridge = LinearBridge(name="down_proj") @@ -134,15 +138,30 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: assert bridge._relevance_rule_kinds == () - x = torch.randn(2, 4, requires_grad=True) + # An opaque backing module implements no relevance-rule kind, so a requested + # rule cannot be installed here -- distinct from a mount with no protocol at + # all (reported skipped), this raises before any forward or backward pass. + with pytest.raises(RelevanceRuleUnsupportedError, match="mlp"): + with use_relevance_rules( + block, RelevanceRules(activation=True, multiplicative_gate=True) + ): + pass + + 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_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(activation=True, multiplicative_gate=True) - ) as coverage: - assert set(coverage.skipped) == {"mlp"} - assert coverage.installed == () - active = block(x) - assert torch.equal(active, baseline) + with use_relevance_rules(block, RelevanceRules(multiplicative_gate=True)) as coverage: + assert coverage.installed == ("mlp",) + assert torch.equal(block(x), baseline) class TestGatedMLPRelevanceRuleForwardIdentity: diff --git a/tests/unit/model_bridge/test_relevance_rules.py b/tests/unit/model_bridge/test_relevance_rules.py index adbcc3f002..394c09fb5f 100644 --- a/tests/unit/model_bridge/test_relevance_rules.py +++ b/tests/unit/model_bridge/test_relevance_rules.py @@ -21,6 +21,7 @@ from transformer_lens.model_bridge._relevance_rules import ( RelevanceRuleCoverage, RelevanceRules, + RelevanceRuleUnsupportedError, ln_rule, use_relevance_rules, ) @@ -168,6 +169,21 @@ def test_unsupported_component_at_targeted_mount_is_skipped(): 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: diff --git a/transformer_lens/model_bridge/_relevance_rules.py b/transformer_lens/model_bridge/_relevance_rules.py index a193551abd..76ff6394b5 100644 --- a/transformer_lens/model_bridge/_relevance_rules.py +++ b/transformer_lens/model_bridge/_relevance_rules.py @@ -114,6 +114,21 @@ class RelevanceRuleConflictError(RuntimeError): """ +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. @@ -156,6 +171,12 @@ class _RelevanceRuleCapable(Protocol): 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, ...] @@ -208,13 +229,20 @@ def use_relevance_rules(model: nn.Module, rules: RelevanceRules) -> Iterator[Rel """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 component at a - canonical mount that does not implement ``_RelevanceRuleCapable`` for the - requested kind is reported as skipped rather than installed or raising. 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. + 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. 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) @@ -229,8 +257,14 @@ def use_relevance_rules(model: nn.Module, rules: RelevanceRules) -> Iterator[Rel continue if isinstance(module, _RelevanceRuleCapable) and kind in module._relevance_rule_kinds: installed.append((name, module, kind)) - else: - skipped.append(name) + 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) diff --git a/transformer_lens/model_bridge/generalized_components/gated_mlp.py b/transformer_lens/model_bridge/generalized_components/gated_mlp.py index 6368a05a63..d3ade94d4f 100644 --- a/transformer_lens/model_bridge/generalized_components/gated_mlp.py +++ b/transformer_lens/model_bridge/generalized_components/gated_mlp.py @@ -16,18 +16,41 @@ ) +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 @@ -207,26 +230,71 @@ def __init__( self._relevance_rule_activation_active = False self._relevance_rule_gate_active = False - @property - def _relevance_rule_kinds(self) -> Tuple[str, ...]: - """``("activation", "multiplicative_gate")`` only when the gate/up projections - are backed by an allowlisted HF module class (``nn.Linear`` or ``Conv1D``) in a - recognized orientation; empty otherwise, so an opaque or unrecognized backing - module is reported skipped rather than silently installing a rule the - weights-recompute cannot honor correctly. + def _is_gated_mlp_shaped(self) -> bool: + """Whether this instance has the gate/up/down submodules the recompute 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 unrecognized weight-backing class or activation form, + which occupy exactly this node's shape but cannot be honored correctly. """ if self.original_component is None: - return () + return False gate_module = getattr(self, "gate", None) in_module = getattr(self, "in", None) out_module = getattr(self, "out", None) - if gate_module is None or in_module is None or out_module is None: - return () - if weight_layout_in_out(gate_module) is None or weight_layout_in_out(in_module) is None: + return gate_module is not None and in_module is not None and out_module is not None + + def _has_recognized_weight_backing(self) -> bool: + """Whether gate/up/down are all backed by an allowlisted HF module class.""" + gate_module = getattr(self, "gate", None) + in_module = getattr(self, "in", None) + out_module = getattr(self, "out", None) + return ( + weight_layout_in_out(gate_module) is not None + and weight_layout_in_out(in_module) is not None + and weight_layout_in_out(out_module) 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, or when the gate/up/down + projections are not backed by an allowlisted HF module class (``nn.Linear`` + or ``Conv1D``) in a recognized orientation, since the weights-recompute + cannot honor either rule correctly on an opaque backing module. Otherwise + always includes ``"multiplicative_gate"`` (the Half-rule does not depend on + the activation function) and includes ``"activation"`` only when the + configured activation form supports the Identity-rule, so a relu-family + activation is excluded independently of weight backing. + """ + if not self._is_gated_mlp_shaped() or not self._has_recognized_weight_backing(): return () - if weight_layout_in_out(out_module) is None: + kinds: Tuple[str, ...] = ("multiplicative_gate",) + if identity_rule_supports_activation(self.config): + 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 with an unrecognized weight-backing class or an + unsupported activation form is exactly the kind of component a caller + expects either rule to work on. Requesting one of these raises instead of + silently reporting the mount skipped. + """ + if not self._is_gated_mlp_shaped(): return () - return ("activation", "multiplicative_gate") + if not self._has_recognized_weight_backing(): + return ("activation", "multiplicative_gate") + if not identity_rule_supports_activation(self.config): + return ("activation",) + return () def _enable_relevance_rule(self, kind: str) -> None: """Activate the named rule for this instance's recompute path only.""" From 90f91c31599083a4df3b575dcbbe78a77877b17d Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 14 Sep 2026 16:00:32 +0530 Subject: [PATCH 08/13] test(relevance_rules): tolerant parity vs pinned FarnoushRJ/RelP Add tests comparing the LN-, Identity-, and Half-relevance-rule primitives against formulas ported directly from FarnoushRJ/RelP pinned at commit 8219d6dc417c3fd7f318342cf61cd2a0c20b7250, reimplemented here since that repository targets an unrelated pre-Bridge TransformerLens fork rather than exposing an importable API: - LN-rule: matches the reference's x / scale.detach() to floating-point precision. - Half-rule: matches the reference's z / 2 + (z / 2).detach() split to floating-point precision. - Identity-rule: matches the reference's epsilon-stabilized ratio away from x == 0 within a tolerance sized to the reference's 1e-6 stabilizer constant, for SiLU, exact GELU, and tanh-approximate GELU. - Identity-rule at x == 0: asserts the known discrepancy explicitly rather than absorbing it into a tolerance -- this module's rule uses the paper-defined removable-singularity limit of 0.5, while the reference's epsilon stabilizer collapses the ratio to exactly 0. All cases pass immediately since the primitives already exist; this commit adds a second, independent oracle rather than driving new implementation. --- tests/unit/tools/test_relevance_rules.py | 131 +++++++++++++++++++++++ 1 file changed, 131 insertions(+) diff --git a/tests/unit/tools/test_relevance_rules.py b/tests/unit/tools/test_relevance_rules.py index 202619abe1..b28ed66326 100644 --- a/tests/unit/tools/test_relevance_rules.py +++ b/tests/unit/tools/test_relevance_rules.py @@ -345,3 +345,134 @@ def test_recompute_via_tl_weights_matches_native_forward_and_rule_vjp(self, back (grad_x_oracle,) = torch.autograd.grad(down_oracle.sum(), x_oracle) torch.testing.assert_close(grad_x, grad_x_oracle) + + +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) From 8ae1d22d7cbe8efe9952f623f40efac02433800e Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 14 Sep 2026 16:32:41 +0530 Subject: [PATCH 09/13] fix(model_bridge): find relevance-rule mounts shadowed by named_modules dedup use_relevance_rules located canonical mounts (ln1, ln2, mlp) by scanning model.named_modules() for a matching leaf name. On a real assembled TransformerBridge, the same bridge component is reachable through two paths: the canonical alias (blocks.N.ln1) and the raw HF module tree the bridge wraps in place (blocks.N._original_component.input_layernorm). nn.Module.named_modules() deduplicates by object identity and keeps only the first-visited path, which is the raw HF-attribute path since it is registered before the canonical alias, so the canonical ln1/ln2 name was never seen. The LN-rule therefore never installed on any real model, and was reported neither installed nor skipped -- silently absent from both. Walk with remove_duplicate=False to recover every path, then keep the fewest-dot-separated-segments path per object so a mount name that happens to match through both the canonical alias and the raw HF attribute (mlp does, on every architecture checked) collapses to a single canonical-looking entry instead of a duplicate. --- .../model_bridge/_relevance_rules.py | 33 +++++++++++++++++-- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/transformer_lens/model_bridge/_relevance_rules.py b/transformer_lens/model_bridge/_relevance_rules.py index 76ff6394b5..ffb5f7749a 100644 --- a/transformer_lens/model_bridge/_relevance_rules.py +++ b/transformer_lens/model_bridge/_relevance_rules.py @@ -199,6 +199,35 @@ def _disable_relevance_rule(self, kind: str) -> None: } +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. @@ -252,9 +281,7 @@ def use_relevance_rules(model: nn.Module, rules: RelevanceRules) -> Iterator[Rel skipped: List[str] = [] for kind in requested_kinds: mount_names = _CANONICAL_MOUNTS.get(kind, ()) - for name, module in model.named_modules(): - if name.rsplit(".", 1)[-1] not in mount_names: - continue + 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 From dd31feddfb861a28476c9dbe716659583f8a3dd1 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Mon, 14 Sep 2026 16:33:05 +0530 Subject: [PATCH 10/13] test(integration): backend on tiny Qwen2 + JointGateUp fixtures End-to-end coverage of the relevance-rule backend on tiny, fully offline HF fixtures built from a programmatic config (no network access, no checkpoint download): a tiny random Qwen2 exercising GatedMLPBridge's opaque recompute-from-weights path, and a tiny random Phi-3 exercising JointGateUpMLPBridge's already-reconstructed forward. With normalization, activation, and multiplicative_gate all active together: forward stays torch.equal to the rule-inactive baseline, every canonical mount across both blocks is reported installed with nothing skipped, 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 (captured via hook_in/hook_out, outside the LN-rule's fail-closed hook_scale/hook_normalized guard). --- tests/integration/test_relevance_lens.py | 221 +++++++++++++++++++++++ 1 file changed, 221 insertions(+) create mode 100644 tests/integration/test_relevance_lens.py diff --git a/tests/integration/test_relevance_lens.py b/tests/integration/test_relevance_lens.py new file mode 100644 index 0000000000..94e62bbe54 --- /dev/null +++ b/tests/integration/test_relevance_lens.py @@ -0,0 +1,221 @@ +"""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 recompute 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 recompute-from-weights path (the +# primary path per real usage); 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 + ) From 9ee03140fde0c9d58de672797b74f4241ec5a740 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 00:53:40 +0530 Subject: [PATCH 11/13] refactor(model_bridge): gated-MLP relevance rules without recompute --- tests/integration/test_relevance_lens.py | 11 +- .../test_gated_mlp_relevance_rule.py | 349 ++++++++++++------ .../test_joint_gate_up_relevance_rule.py | 7 +- tests/unit/tools/test_relevance_rules.py | 48 +-- .../model_bridge/_relevance_rules.py | 27 ++ .../generalized_components/gated_mlp.py | 333 +++++++++-------- .../joint_gate_up_mlp.py | 24 ++ 7 files changed, 479 insertions(+), 320 deletions(-) diff --git a/tests/integration/test_relevance_lens.py b/tests/integration/test_relevance_lens.py index 94e62bbe54..e790c7d17f 100644 --- a/tests/integration/test_relevance_lens.py +++ b/tests/integration/test_relevance_lens.py @@ -6,7 +6,7 @@ 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 recompute path) and a tiny +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 @@ -81,10 +81,11 @@ def _build_phi3_bridge() -> TransformerBridge: return TransformerBridge(model=hf_model, adapter=adapter, tokenizer=_MockTokenizer()) -# Qwen2 exercises GatedMLPBridge's opaque recompute-from-weights path (the -# primary path per real usage); Phi-3 exercises JointGateUpMLPBridge's -# already-reconstructed forward, isolating rule bugs from Bridge-integration -# bugs on the fused-projection family. +# 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, 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 index de90784948..0d2ab27d36 100644 --- 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 @@ -1,14 +1,16 @@ -"""Identity-/Half-rule integration on GatedMLPBridge's recompute-from-weights path. - -Covers the commit-6 contract for the raw (non-fused) gated-MLP path: the opaque -native forward stays bit-identical to today's forward by construction (the wrapping -returns ``original_component(x)`` itself, never a reproduced numeric value), the -backward recomputes gate/up/down from the TL-oriented ``W_gate``/``W_in``/``W_out`` -and applies the Identity-rule to the activation and/or the Half-rule to the gate*up -product per whichever is independently active, weight and bias keep their ordinary -gradient, and the recompute is allowlisted by the underlying HF module class -(``nn.Linear`` and ``Conv1D``) rather than reported installed for an unrecognized -backing module. +"""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 @@ -38,17 +40,64 @@ 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 structure: opaque single call.""" + """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): + 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(F.silu(self.gate_proj(x)) * self.up_proj(x)) + 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): @@ -70,20 +119,21 @@ def _make_projections(backing_class: str, d_model: int = 4, d_mlp: int = 8, 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 ( - Conv1D(d_mlp, d_model), - Conv1D(d_mlp, d_model), - Conv1D(d_model, d_mlp), + _OpaqueProj(d_model, d_mlp, bias=bias), + _OpaqueProj(d_model, d_mlp, bias=bias), + _OpaqueProj(d_mlp, d_model, bias=bias), ) -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) - - bridge = GatedMLPBridge(name="mlp", config=_Cfg(hidden_act)) +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") @@ -91,66 +141,88 @@ def _make_bridge( bridge.add_module("in", in_bridge) bridge.add_module("out", out_bridge) bridge.set_original_component(hf_mlp) - gate_bridge.set_original_component(gate_proj) - in_bridge.set_original_component(up_proj) - out_bridge.set_original_component(down_proj) + 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 -BACKING_CLASSES = ["nn.Linear", "Conv1D"] +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_when_allowlisted(self, backing_class): + 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_not_capable_when_backing_module_is_unrecognized(self): - class _OpaqueProj(nn.Module): - """Neither nn.Linear nor Conv1D -- an unrecognized backing class.""" - - def __init__(self, d_in: int, d_out: int): - super().__init__() - self.weight = nn.Parameter(torch.randn(d_out, d_in)) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return x @ self.weight.T - - gate_proj = _OpaqueProj(4, 8) - up_proj = _OpaqueProj(4, 8) - down_proj = _OpaqueProj(8, 4) - hf_mlp = _TinyGatedMLP(gate_proj, up_proj, down_proj) - - bridge = GatedMLPBridge(name="mlp", config=_Cfg()) - 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(gate_proj) - in_bridge.set_original_component(up_proj) - out_bridge.set_original_component(down_proj) - block = _Block(bridge) - - assert bridge._relevance_rule_kinds == () + 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",) - # An opaque backing module implements no relevance-rule kind, so a requested - # rule cannot be installed here -- distinct from a mount with no protocol at - # all (reported skipped), this raises before any forward or backward pass. with pytest.raises(RelevanceRuleUnsupportedError, match="mlp"): - with use_relevance_rules( - block, RelevanceRules(activation=True, multiplicative_gate=True) - ): + with use_relevance_rules(block, RelevanceRules(activation=True)): pass - 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",) + 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 @@ -174,35 +246,31 @@ def test_forward_identical_while_rule_active(self, backing_class): 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) - before = block(x) - after = block(x) - assert torch.equal(before, after) + 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: - def _oracle_grads(self, hf_mlp, x, activation_active, gate_active): - # hf_mlp's own parameters are shared with the bridge under test, so a prior - # backward through the bridge already left gradients on them; without - # resetting first, this second backward would accumulate on top instead of - # producing an independently comparable oracle. - 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, 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() - grads = {n: p.grad.clone() for n, p in hf_mlp.named_parameters()} - return x_oracle.grad.clone(), grads +class TestGatedMLPRelevanceRuleVJP: @pytest.mark.parametrize("backing_class", BACKING_CLASSES) - def test_both_rules_active_matches_manually_composed_oracle(self, backing_class): + 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) @@ -213,8 +281,26 @@ def test_both_rules_active_matches_manually_composed_oracle(self, backing_class) grad_x = x.grad.clone() grads = {n: p.grad.clone() for n, p in hf_mlp.named_parameters()} - expected_grad_x, expected_grads = self._oracle_grads( - hf_mlp, x, activation_active=True, gate_active=True + 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(): @@ -230,9 +316,7 @@ def test_activation_only_leaves_gate_split_ordinary(self, backing_class): out.sum().backward() grad_x = x.grad.clone() - expected_grad_x, _ = self._oracle_grads( - hf_mlp, x, activation_active=True, gate_active=False - ) + 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) @@ -245,9 +329,7 @@ def test_multiplicative_gate_only_leaves_activation_ordinary(self, backing_class out.sum().backward() grad_x = x.grad.clone() - expected_grad_x, _ = self._oracle_grads( - hf_mlp, x, activation_active=False, gate_active=True - ) + 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): @@ -262,12 +344,11 @@ def test_rule_inactive_gradients_are_ordinary(self): 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() + hf_mlp(x_plain).sum().backward() torch.testing.assert_close(grad_x, x_plain.grad) - for name, grad in grads.items(): - torch.testing.assert_close(grad, dict(hf_mlp.named_parameters())[name].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) @@ -278,9 +359,69 @@ def test_bias_free_projections_are_handled(self): out.sum().backward() grad_x = x.grad.clone() - expected_grad_x, expected_grads = self._oracle_grads( - hf_mlp, x, activation_active=True, gate_active=True - ) + 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 index 7834a2d9c2..f8bee5e693 100644 --- 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 @@ -1,10 +1,11 @@ """Identity-/Half-rule integration on JointGateUpMLPBridge's reconstructed forward. -Unlike the raw GatedMLPBridge path (recompute-from-weights, tested separately), +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 -- no -weights-recompute Function is needed here. This covers that the reconstructed +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. diff --git a/tests/unit/tools/test_relevance_rules.py b/tests/unit/tools/test_relevance_rules.py index b28ed66326..c8a6e35046 100644 --- a/tests/unit/tools/test_relevance_rules.py +++ b/tests/unit/tools/test_relevance_rules.py @@ -269,14 +269,13 @@ def __init__(self, original_component): self.original_component = original_component -class TestGatedMLPRecomputeOrientation: - """The gated-MLP recompute VJP reads ``W_gate``/``W_in`` through the same - ``weight_layout_in_out``/``normalize_mlp_weight`` helpers ``MLPBridge`` uses, so a - fresh orientation bug in the recompute path (as opposed to the already-tested - accessor) would silently transpose the gate/up product for one backing class. - Covers the two HF module classes the recompute allowlists: ``nn.Linear`` (weight - stored ``[out, in]``, transposed to TL orientation) and ``Conv1D`` (weight stored - ``[in, out]``, already TL-oriented). +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"]) @@ -313,39 +312,6 @@ def test_tl_oriented_matmul_reproduces_native_projection(self, backing_class): hidden = torch.randn(2, 5) assert torch.allclose(hidden @ w_out + down_proj.bias, down_proj(hidden), atol=1e-6) - def test_recompute_via_tl_weights_matches_native_forward_and_rule_vjp(self, backing_class): - gate_proj, up_proj, down_proj = self._make_gate_up_down(backing_class) - x = torch.randn(2, 3, requires_grad=True) - - 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") - - gate_output = x @ w_gate + gate_proj.bias - up_output = x @ w_in + up_proj.bias - activated = identity_rule(gate_output, F.silu) - gated = half_rule(activated, up_output) - down = gated @ w_out + down_proj.bias - - native_down = down_proj(F.silu(gate_proj(x.detach())) * up_proj(x.detach())) - assert torch.allclose(down, native_down, atol=1e-6) - - (grad_x,) = torch.autograd.grad(down.sum(), x) - - # Oracle: the exact same composition, expressed with ordinary autodiff - # through the rule primitives but built from each class's own native call - # instead of the TL-oriented matmul above -- if the recompute transposed - # either weight incorrectly, this would disagree. - x_oracle = x.detach().clone().requires_grad_(True) - gate_output_oracle = gate_proj(x_oracle) - up_output_oracle = up_proj(x_oracle) - activated_oracle = identity_rule(gate_output_oracle, F.silu) - gated_oracle = half_rule(activated_oracle, up_output_oracle) - down_oracle = down_proj(gated_oracle) - (grad_x_oracle,) = torch.autograd.grad(down_oracle.sum(), x_oracle) - - torch.testing.assert_close(grad_x, grad_x_oracle) - class TestPinnedReferenceParity: """Tolerant parity against ``FarnoushRJ/RelP`` pinned at diff --git a/transformer_lens/model_bridge/_relevance_rules.py b/transformer_lens/model_bridge/_relevance_rules.py index ffb5f7749a..b3881bb211 100644 --- a/transformer_lens/model_bridge/_relevance_rules.py +++ b/transformer_lens/model_bridge/_relevance_rules.py @@ -105,6 +105,33 @@ def half_rule(u: torch.Tensor, v: torch.Tensor) -> torch.Tensor: 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. diff --git a/transformer_lens/model_bridge/generalized_components/gated_mlp.py b/transformer_lens/model_bridge/generalized_components/gated_mlp.py index d3ade94d4f..ccc533c726 100644 --- a/transformer_lens/model_bridge/generalized_components/gated_mlp.py +++ b/transformer_lens/model_bridge/generalized_components/gated_mlp.py @@ -2,18 +2,20 @@ This module contains the bridge component for gated MLP layers (e.g., LLaMA, Gemma). """ -from typing import Any, Callable, Dict, Mapping, Optional, Tuple +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 +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, - weight_layout_in_out, -) +from transformer_lens.model_bridge.generalized_components.mlp import MLPBridge def _resolve_activation_fn_name(config: Any) -> Optional[str]: @@ -73,120 +75,23 @@ def relu_squared(x: torch.Tensor) -> torch.Tensor: return torch.nn.functional.silu -class _GatedMLPRecomputeRule(torch.autograd.Function): - """Wrap a raw gated-MLP module's own forward call in the Identity-/Half-rule VJP. +class _IdentityRuleActivation(nn.Module): + """Route a wrapped activation through the Identity-rule for a scope's duration. - Forward returns ``component(x, ...)`` unchanged, so the result is bit-identical - to the native forward by construction. The opaque call gives no access to its own - internal gate/up/down intermediates, so backward recomputes them from the - TL-oriented ``W_gate``/``W_in``/``W_out`` (checkpointing-style: the values are - rederived here rather than saved from forward) and reapplies the Identity-rule to - the activation and/or the Half-rule to the gate*up product, per whichever of the - two is active. Weight and bias gradients come out of the same recomputed graph, - via ``torch.autograd.grad``, so they keep their ordinary form -- the rules only - redefine how relevance reaches the input, not parameter training gradients. + 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. """ - @staticmethod - def forward( - ctx: Any, - x: torch.Tensor, - w_gate: torch.Tensor, - b_gate: Optional[torch.Tensor], - w_in: torch.Tensor, - b_in: Optional[torch.Tensor], - w_out: torch.Tensor, - b_out: Optional[torch.Tensor], - act_fn: Callable[[torch.Tensor], torch.Tensor], - activation_rule_active: bool, - gate_rule_active: bool, - component: torch.nn.Module, - extra_args: Tuple[Any, ...], - extra_kwargs: Dict[str, Any], - ) -> torch.Tensor: - ctx.save_for_backward(x, w_gate, w_in, w_out) - ctx.b_gate = b_gate - ctx.b_in = b_in - ctx.b_out = b_out - ctx.act_fn = act_fn - ctx.activation_rule_active = activation_rule_active - ctx.gate_rule_active = gate_rule_active - result: torch.Tensor = component(x, *extra_args, **extra_kwargs) - return result - - @staticmethod - def backward( # type: ignore[override] - ctx: Any, grad_output: torch.Tensor - ) -> Tuple[ - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], - None, - None, - None, - None, - None, - None, - ]: - x, w_gate, w_in, w_out = ctx.saved_tensors - b_gate, b_in, b_out = ctx.b_gate, ctx.b_in, ctx.b_out - act_fn = ctx.act_fn - - with torch.enable_grad(): - x_ = x.detach().requires_grad_(True) - w_gate_ = w_gate.detach().requires_grad_(True) - w_in_ = w_in.detach().requires_grad_(True) - w_out_ = w_out.detach().requires_grad_(True) - b_gate_ = b_gate.detach().requires_grad_(True) if b_gate is not None else None - b_in_ = b_in.detach().requires_grad_(True) if b_in is not None else None - b_out_ = b_out.detach().requires_grad_(True) if b_out is not None else None - - 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_ + def __init__(self, wrapped: Callable[[torch.Tensor], torch.Tensor]): + super().__init__() + self._wrapped_activation = wrapped - activated = ( - identity_rule(gate_output, act_fn) - if ctx.activation_rule_active - else act_fn(gate_output) - ) - gated = ( - half_rule(activated, up_output) if ctx.gate_rule_active else activated * up_output - ) - - down = gated @ w_out_ - if b_out_ is not None: - down = down + b_out_ - - leaves = [x_, w_gate_, b_gate_, w_in_, b_in_, w_out_, b_out_] - needed = [leaf for leaf in leaves if leaf is not None] - grads = torch.autograd.grad(down, needed, grad_outputs=grad_output, allow_unused=True) - grad_iter = iter(grads) - results = [next(grad_iter) if leaf is not None else None for leaf in leaves] - grad_x, grad_w_gate, grad_b_gate, grad_w_in, grad_b_in, grad_w_out, grad_b_out = results - - return ( - grad_x, - grad_w_gate, - grad_b_gate, - grad_w_in, - grad_b_in, - grad_w_out, - grad_b_out, - None, - None, - None, - None, - None, - None, - ) + def forward(self, x: torch.Tensor) -> torch.Tensor: + return identity_rule(x, self._wrapped_activation) class GatedMLPBridge(MLPBridge): @@ -229,15 +134,21 @@ def __init__( 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 the recompute needs. + """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 unrecognized weight-backing class or activation form, - which occupy exactly this node's shape but cannot be honored correctly. + 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 @@ -246,34 +157,52 @@ def _is_gated_mlp_shaped(self) -> bool: 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 _has_recognized_weight_backing(self) -> bool: - """Whether gate/up/down are all backed by an allowlisted HF module class.""" - gate_module = getattr(self, "gate", None) - in_module = getattr(self, "in", None) - out_module = getattr(self, "out", 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 ( - weight_layout_in_out(gate_module) is not None - and weight_layout_in_out(in_module) is not None - and weight_layout_in_out(out_module) is not None + 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, or when the gate/up/down - projections are not backed by an allowlisted HF module class (``nn.Linear`` - or ``Conv1D``) in a recognized orientation, since the weights-recompute - cannot honor either rule correctly on an opaque backing module. Otherwise - always includes ``"multiplicative_gate"`` (the Half-rule does not depend on - the activation function) and includes ``"activation"`` only when the - configured activation form supports the Identity-rule, so a relu-family - activation is excluded independently of weight backing. + 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() or not self._has_recognized_weight_backing(): + if not self._is_gated_mlp_shaped(): return () kinds: Tuple[str, ...] = ("multiplicative_gate",) - if identity_rule_supports_activation(self.config): + if self._activation_rule_installable(): kinds = ("activation",) + kinds return kinds @@ -283,31 +212,107 @@ def _relevance_rule_unsupported_kinds(self) -> Tuple[str, ...]: 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 with an unrecognized weight-backing class or an - unsupported activation form is exactly the kind of component a caller - expects either rule to work on. Requesting one of these raises instead of - silently reporting the mount skipped. + 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 not self._has_recognized_weight_backing(): - return ("activation", "multiplicative_gate") - if not identity_rule_supports_activation(self.config): + 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 for this instance's recompute path only.""" + """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, restoring today's native-forward behavior.""" + """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: @@ -345,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( @@ -365,25 +378,11 @@ def forward(self, *args, **kwargs) -> torch.Tensor: hidden_states = args[0] hidden_states = self.hook_in(hidden_states) new_args = (hidden_states,) + args[1:] - if self._relevance_rule_activation_active or self._relevance_rule_gate_active: - act_fn = resolve_activation_fn(self.config) - output = _GatedMLPRecomputeRule.apply( - hidden_states, - self.W_gate, - getattr(self.gate, "bias", None), - self.W_in, - getattr(getattr(self, "in"), "bias", None), - self.W_out, - getattr(self.out, "bias", None), - act_fn, - self._relevance_rule_activation_active, - self._relevance_rule_gate_active, - self.original_component, - new_args[1:], - kwargs, - ) - else: - output = self.original_component(*new_args, **kwargs) + # 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 a3ebcaf78a..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 @@ -12,6 +12,7 @@ ) 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 @@ -167,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 From de28190779def89f355287ada49ffb5acb211cc3 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 00:59:12 +0530 Subject: [PATCH 12/13] fix(model_bridge): LN-rule treats a missing norm weight as identity --- .../test_normalization_relevance_rule.py | 61 +++++++++++++++++++ .../generalized_components/normalization.py | 19 ++++-- 2 files changed, 75 insertions(+), 5 deletions(-) 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 index ae62a2633f..bfd6dcd374 100644 --- a/tests/unit/model_bridge/generalized_components/test_normalization_relevance_rule.py +++ b/tests/unit/model_bridge/generalized_components/test_normalization_relevance_rule.py @@ -14,6 +14,7 @@ import pytest import torch import torch.nn as nn +import torch.nn.functional as F from transformer_lens.model_bridge._relevance_rules import ( RelevanceRuleConflictError, @@ -68,6 +69,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: 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) @@ -294,6 +307,54 @@ def test_param_free_pre_norm_is_skipped_and_forward_unaffected(self, bridge_cls) 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) diff --git a/transformer_lens/model_bridge/generalized_components/normalization.py b/transformer_lens/model_bridge/generalized_components/normalization.py index ca088dd0a0..5aa0b59bc9 100644 --- a/transformer_lens/model_bridge/generalized_components/normalization.py +++ b/transformer_lens/model_bridge/generalized_components/normalization.py @@ -50,7 +50,8 @@ class _NativeLNRuleForward(torch.autograd.Function): 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. + training gradients. A parameter-free norm (``weight`` is ``None``, e.g. + OLMo's ``OlmoLayerNorm``) is treated as a unit scale in both directions. """ @staticmethod @@ -58,7 +59,7 @@ def forward( ctx: Any, x_centered: torch.Tensor, denom: torch.Tensor, - weight: torch.Tensor, + weight: Optional[torch.Tensor], bias: Optional[torch.Tensor], x: torch.Tensor, component: torch.nn.Module, @@ -88,12 +89,15 @@ def backward( None, ]: x_centered, denom, weight = ctx.saved_tensors - w_eff = (1.0 + weight) if ctx.offset else weight + 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.requires_grad + if weight is not None and weight.requires_grad else None ) grad_bias = ( @@ -367,7 +371,12 @@ def _native_forward_with_ln_rule( if isinstance(eps_value, torch.Tensor) else (variance + float(eps_value)).sqrt() ) - weight = cast(torch.Tensor, self.weight) + # 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( From 709f5e828def147fddafbd464e0c20ed7628dcc1 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 01:14:04 +0530 Subject: [PATCH 13/13] fix(model_bridge): relevance-rule mount coverage --- .../unit/model_bridge/test_relevance_rules.py | 57 +++++++++++++++++++ .../model_bridge/_relevance_rules.py | 22 ++++++- 2 files changed, 76 insertions(+), 3 deletions(-) diff --git a/tests/unit/model_bridge/test_relevance_rules.py b/tests/unit/model_bridge/test_relevance_rules.py index 394c09fb5f..57cca39c25 100644 --- a/tests/unit/model_bridge/test_relevance_rules.py +++ b/tests/unit/model_bridge/test_relevance_rules.py @@ -90,6 +90,30 @@ def _tiny_block() -> _TinyBlock: 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. @@ -160,6 +184,39 @@ def test_positional_targeting_excludes_same_class_component_at_other_mount(): 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() diff --git a/transformer_lens/model_bridge/_relevance_rules.py b/transformer_lens/model_bridge/_relevance_rules.py index b3881bb211..5da1bbf80d 100644 --- a/transformer_lens/model_bridge/_relevance_rules.py +++ b/transformer_lens/model_bridge/_relevance_rules.py @@ -218,9 +218,12 @@ 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. +# 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"), + "normalization": ("ln1", "ln2", "ln1_post", "ln2_post"), "activation": ("mlp",), "multiplicative_gate": ("mlp",), } @@ -294,7 +297,11 @@ def use_relevance_rules(model: nn.Module, rules: RelevanceRules) -> Iterator[Rel ``_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. Scopes over the same model + 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 @@ -304,6 +311,15 @@ def use_relevance_rules(model: nn.Module, rules: RelevanceRules) -> Iterator[Rel 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: