feat(model_bridge): shared LN/Identity/Half relevance-rule backend - #1785
Open
janmenjayap wants to merge 10 commits into
Open
janmenjayap wants to merge 10 commits into
janmenjayap wants to merge 10 commits into
Conversation
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.
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.
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.
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.
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.
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.
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.
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.
…es 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.
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).
jlarson4
requested changes
Sep 14, 2026
jlarson4
left a comment
jlarson4
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for splitting the backend out ahead of the fit code, and for the closed-form and RelP-parity tests.
I do have a larger change request for this one than I usually do, as it touches generalized components in a signficant way. I want to specifically focus in on the design of the gated-MLP rules. _GatedMLPRecomputeRule rebuilds down(act(gate(x)) * up(x)) from the weights, so an HF MLP whose forward does anything more gets a wrong gradient while coverage reports the rule installed. Falcon-H1's multipliers and Gemma3n's gate sparsity both hit this today. Running HF's forward inside the Function also takes the MLP's own hooks out of the backward, and the compatibility-mode branch skips the rules entirely. Neither rule needs a recompute. The Half-rule is exactly ordinary autograd with the product's gradient halved, so it can be a gradient-scale hook on mlp.out.hook_in. The Identity-rule can sit on the HF MLP's own act_fn module, applied inline wherever the bridge computes the activation itself, as JointGateUpMLPBridge already does. Could the gated-MLP rules move to that shape, with the recompute dropped?
Generally, where possible, we try to avoid adding recompute if we can. Any areas where we recompute HF elements is a potential failure point on new models, if they have new unexpected behaviors.
| if isinstance(eps_value, torch.Tensor) | ||
| else (variance + float(eps_value)).sqrt() | ||
| ) | ||
| weight = cast(torch.Tensor, self.weight) |
Collaborator
There was a problem hiding this comment.
OLMo's OlmoLayerNorm has no weight, but its ln1/ln2 mounts report the LN-rule installed, and the first rule-active forward then fails with AttributeError on self.weight. Could a missing weight be treated as 1 here and in the backward?
|
|
||
| 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 |
Collaborator
There was a problem hiding this comment.
RelevanceRules(attention=True) has no entry in _CANONICAL_MOUNTS, so the context installs nothing, skips nothing and doesn't raise. A caller asking for attention rules gets ordinary gradients and an empty coverage report. Could requesting a kind with no canonical mount raise instead?
| # 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, ...]] = { |
Collaborator
There was a problem hiding this comment.
The leaf-name match is exact, so the ln1_post/ln2_post norms that Gemma 2/3/4, GLM-4 and several others mount never get the LN-rule and appear nowhere in the coverage report. The pinned RelP reference applies the LN-rule to those post-norms. OLMo 2 already gets it on the same post_attention_layernorm role because its adapter mounts that norm as ln1. Could ln1_post and ln2_post join the normalization mounts?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Lands the scoped custom-VJP backend shared between R-lens and the (separate) RelP issue: LN-,
Identity-, and Half-rule primitives, plus a
use_relevance_rules(model, rules)context thatinstalls them on live
TransformerBridgecomponents by mount position. Every rule-active forwardis
torch.equalto the rule-inactive forward — only backward semantics change. No public API isadded; the module is underscore-private and undocumented until the RelP issue is filed. A reviewer
can evaluate this backend in isolation, before any
RelevanceLens.fitcode exists.Part of #1755 (R-lens). This is PR2 of the 3-PR plan on that issue; PR1
(#1764) landed the
estimator-independent J-lens fit driver this backend will later plug into, PR3 adds
RelevanceLens.fit+ provenance + docs. Note: #1755 is currently closed(
state_reason: completed) even though only PR1 has landed so far. Would it be possible toreopen it (or file a fresh tracking issue) so PR2 and PR3 stay linked to it? Happy to help with
that if useful.
Motivation
R-lens and generic RelP both need forward-equivalent, backward-modified norm/activation/gate
semantics on live Bridge components. Landing the backend as its own PR, ahead of
RelevanceLens.fit, lets a reviewer evaluate the risky part — custom autograd wrappingproduction
NormalizationBridge/GatedMLPBridge/JointGateUpMLPBridgeforward paths — inisolation, before layering the estimator-specific fitting logic on top. This PR addresses five
explicit conditions raised during discussion of the 3-PR plan; each is called out against the
commit that satisfies it below.
What ships (commit by commit)
test(relevance_rules): closed-form VJP + forward-identity primitivesAnalytic tests for LN-, Identity-, and Half-rule against their closed forms, covering zero,
negative, mixed-sign, non-contiguous, and batched inputs at fp32/fp64. Red as expected — the
primitive module doesn't exist yet.
feat(relevance_rules): forward-equivalent rule primitivesImplements each rule as a
torch.autograd.Function: forward reproduces the native valueexactly, backward substitutes the rule's closed-form VJP (LN-rule treats the denominator as
constant; Identity-rule uses
f(x)/xwith the1/2limit at zero; Half-rule splits theproduct-rule gradient evenly).
test(model_bridge): scoped rule context forward-identity + cleanupTests for
use_relevance_rules: forward-identity (standalone and nested), gradient restorationon exit, reference-counted nested scopes, no-op on exception, positional (not class-based)
targeting, and "skipped" reporting for a mount that doesn't implement the rule protocol. Red —
the context doesn't exist yet.
feat(model_bridge): scoped relevance-rule context + coverage reportAdds
RelevanceRules(frozen dataclass:normalization,activation,multiplicative_gate,attention),RelevanceRuleCoverage(installed/skipped), anduse_relevance_rules(...).Installs each requested rule only at that rule kind's canonical mount name — never by
isinstance— with reference-counted nesting and exception-safe cleanup. No model config ismutated.
Satisfies reviewer condition 3 (positional, not class-based, targeting).
feat(model_bridge): native-forward LN-rule on NormalizationBridgeWraps the existing native-autograd no-edit branch's own
original_component(x)call in thecustom
Function, so the rule-active forward is bit-identical to today's by construction —no restructuring of the existing native-autograd / python-norm / hook-fallback dispatch. A
backward hook on
hook_scale/hook_normalized, or a forward hook editing either, now raisesRelevanceRuleConflictErrorwhile the rule is active (checked on every forward call) instead ofthe previous warn-and-fallback, which would have silently composed the rule with the hook edit.
NormalizationBridgeimplements the rule protocol so it's targeted atln1/ln2positionally;any instance that never reaches the native-autograd branch reports its kind as empty, so it's
"skipped" rather than falsely "installed."
Satisfies reviewer conditions 1 (strictly-additive integration) and 2 (fail-closed
hook/rule precedence, checked per forward call).
feat(model_bridge): gated-MLP rule via recompute-from-weights VJPFor the raw
GatedMLPBridge(Qwen2/Llama/Gemma family), wraps the opaqueoriginal_component(x)call and recomputes gate/up/down checkpointing-style from theTL-oriented
W_gate/W_in/W_outproperties (never the compatibility-mode-only_processed_*attributes) to apply Identity-/Half-rule; weight/bias gradients come back off thesame recomputed graph via
torch.autograd.grad. The recompute is allowlisted by the HF moduleclass backing the weights (
nn.Linearvs.Conv1D), not per adapter.JointGateUpMLPBridge(Phi-3/GLM) attaches the rule directly at its already-Python-visible
act_fn(gate) * upmultiplication. A gated-MLP node now answers to both "activation" and "multiplicative_gate"
independently at one mount, so the protocol moved from a single
_relevance_rule_kindto atuple
_relevance_rule_kinds, and refcounting is keyed per(module, kind).Satisfies reviewer condition 4 (gated-MLP allowlist keyed on HF module class, not adapter).
feat(model_bridge): fail-closed coverage on unsupported pathsDistinguishes benign non-applicability (no protocol, or a dispatch path a rule genuinely
doesn't wrap) from a component that's expected to honor a rule but currently can't. Adds an
optional
_relevance_rule_unsupported_kindsattribute;use_relevance_rulesnow raisesRelevanceRuleUnsupportedError, naming the component's dotted path, before any forward/backwardruns.
GatedMLPBridgeuses it for an unrecognized weight-backing class (disqualifies bothkinds) and for relu-family activations (disqualifies only "activation" —
f(x)/xisn't thecorrect derivative for relu/relu2, but Half-rule is unaffected by activation form).
test(relevance_rules): tolerant parity vs pinned FarnoushRJ/RelPCompares all three primitives against formulas ported from
FarnoushRJ/RelPpinned at8219d6dc417c3fd7f318342cf61cd2a0c20b7250, and explicitly asserts the one known discrepancy:this backend's Identity-rule uses the paper-defined
1/2limit atx == 0, while thereference's epsilon stabilizer collapses to exactly
0there. All other cases match tofloating-point precision.
fix(model_bridge): find relevance-rule mounts shadowed by named_modules dedup(not inthe original plan)
named_modules()deduplicates by object identity and keeps only the first-visited path perobject. On a real assembled
TransformerBridge, the raw HF-attribute path(
blocks.N._original_component.input_layernorm) is registered before the canonical alias(
blocks.N.ln1), so the canonical name was never seen — the LN-rule never installed on anyreal model and was silently reported neither installed nor skipped. Fixed by walking with
remove_duplicate=Falseand keeping the shortest-path alias per object. Found and fixed whilebuilding the Commit 10/11 integration fixtures below, before this backend had ever been
exercised end-to-end on an assembled model.
test(integration): backend on tiny Qwen2 + JointGateUp fixturesEnd-to-end, fully offline (no network, no checkpoint download) coverage on a tiny random
Qwen2 (opaque
GatedMLPBridgerecompute path) and a tiny random Phi-3(
JointGateUpMLPBridge), with normalization, activation, and multiplicative_gate all activetogether: forward stays
torch.equalto the rule-inactive baseline, every canonical mount onboth fixtures reports installed with nothing skipped, and each rule-active node's upstream
gradient matches its closed-form VJP in the real graph.
Testing
tests/unit/tools/test_relevance_rules.py+tests/unit/model_bridge/test_relevance_rules.py+tests/unit/model_bridge/generalized_components/test_gated_mlp_relevance_rule.py+test_joint_gate_up_relevance_rule.py+test_normalization_relevance_rule.py+tests/unit/model_bridge/supported_architectures/test_relevance_rule_mount_placeholders.py—118 passed.
tests/integration/test_relevance_lens.py— 6 passed (new backend integration fixtures).tests/integration/test_backward_lens.py+tests/unit/model_bridge/generalized_components/test_normalization_hook_semantics.py— 43passed (defensive gate, untouched by this PR).
uv run mypyon the four touched source files — clean.black --check/isort --check-onlyon all touched files — clean.
The
backward_lensand hook-semantics suites are a defensive gate, not a functional one: this PRadds a strictly-additive branch to
NormalizationBridge/GatedMLPBridge/JointGateUpMLPBridgewithout restructuring existing dispatch, and both suites stay green with the rule context
inactive, which is the state every existing caller runs in today.
Out of scope (later PR on #1755)
RelevanceLens.fit, estimator/rule provenance in the lens artifact contract, merge/registryguards, and docs/notebook integration — all PR3. Source-artifact conversion, pass@k evaluation,
CKA, probe-upper-bound, and MoE/mHC rules are tracked as separate follow-up issues.
API
No public API added or changed.
transformer_lens/model_bridge/_relevance_rules.pyisunderscore-private, not exported from any
__init__.py, and undocumented — per reviewer condition5, it stays that way until the RelP issue is filed.
Checklist
dev; working branchfeat/r-lens-backend.torch.equal) on every claimed path.original_component(x)itself; no restructuring of the existing three-waydispatch; rule-inactive behavior unchanged (regression-covered).
ln1/ln2mount name), neverisinstance;gemma3n/gemma4placeholder mounts and
stablelm'sq_layernorm/k_layernormnaming are covered and classified"not applicable."
W_gate/W_in/W_outproperties, never_processed_*;allowlisted per HF module class, not per adapter.
with the component's dotted path.
and reference-counted per
(module, kind).FarnoushRJ/RelP, with the exact-zero stabilizer discrepancyrecorded rather than tolerance-absorbed.
model.cfgflags; no.data.tests/integration/test_backward_lens.pyandtests/unit/model_bridge/generalized_components/test_normalization_hook_semantics.pygreen,untouched.
make format(black/isort) anduv run mypyclean on touched files.make test-prrun (unit + docstring + acceptance + integration) attached beforesubmit.