Conversation
* Refactored `transformers` so that bridge configuration code is source-agnostic
* Initial breakup of Bridge into core and bridge only components. Setup remote driver in prep
* Initial setup of the vllm source
Cherry-picked from 98bbb835 on feature/vllm-integration. Resolved against
dev-4.x's Phase A type split:
- bridge.py: kept the shim (TransformerBridge content lives in
transformer_bridge.py now).
- sources/transformers.py: accepted the scaffold's package restructure
(transformers.py → transformers/{__init__,helpers,source}.py +
_hf_format.py).
- sources/transformers/source.py: re-applied Phase A's explicit
TransformersDriver construction in boot().
sources/vllm/source.py's boot_vllm currently sets bridge._forward_impl,
which Phase A removed. The function is dead until the next chunk lifts
it onto VLLMDriver + RemoteBridge — landing as-is so the scaffold tree is
on dev-4.x for incremental Phase B work.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* Revision round 1 of the vLLM system
* Full vLLM integration initial completion
* jupiter patch
* dynamo patching
* pass through batch size
* add max model length to remote bridge boot
* Updating config construction to properly pass through to vLLM
* Wiring TLWorkerExtension
* Additional testing + diagnostic
* Another spot check
* Remove logging
* Fix issue with logit generation on hook_out
* Adjust to use token_ids directly instead of logprobs
* Updating Step 5 and driver bridge cleanup
* Add decoder layer for materializing the residual stream
* Updating tests due to changes on `dev`
* Fix pip install path
* Testing ln_final divergence
* documenting differences in ln_final
* Format cleanup
* Fix bug with beartype
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* Refactored `transformers` so that bridge configuration code is source-agnostic
* Initial breakup of Bridge into core and bridge only components. Setup remote driver in prep
* Initial setup of the vllm source
Cherry-picked from 98bbb835 on feature/vllm-integration. Resolved against
dev-4.x's Phase A type split:
- bridge.py: kept the shim (TransformerBridge content lives in
transformer_bridge.py now).
- sources/transformers.py: accepted the scaffold's package restructure
(transformers.py → transformers/{__init__,helpers,source}.py +
_hf_format.py).
- sources/transformers/source.py: re-applied Phase A's explicit
TransformersDriver construction in boot().
sources/vllm/source.py's boot_vllm currently sets bridge._forward_impl,
which Phase A removed. The function is dead until the next chunk lifts
it onto VLLMDriver + RemoteBridge — landing as-is so the scaffold tree is
on dev-4.x for incremental Phase B work.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* Revision round 1 of the vLLM system
* Full vLLM integration initial completion
* jupiter patch
* dynamo patching
* pass through batch size
* add max model length to remote bridge boot
* Updating config construction to properly pass through to vLLM
* Wiring TLWorkerExtension
* Additional testing + diagnostic
* Another spot check
* Remove logging
* Fix issue with logit generation on hook_out
* Adjust to use token_ids directly instead of logprobs
* Updating Step 5 and driver bridge cleanup
* Add decoder layer for materializing the residual stream
* Updating tests due to changes on `dev`
* Fix pip install path
* Testing ln_final divergence
* documenting differences in ln_final
* Format cleanup
* Initial batching test
* Add version print
* Setup batching for vllm
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* Refactored `transformers` so that bridge configuration code is source-agnostic
* Initial breakup of Bridge into core and bridge only components. Setup remote driver in prep
* Initial setup of the vllm source
Cherry-picked from 98bbb835 on feature/vllm-integration. Resolved against
dev-4.x's Phase A type split:
- bridge.py: kept the shim (TransformerBridge content lives in
transformer_bridge.py now).
- sources/transformers.py: accepted the scaffold's package restructure
(transformers.py → transformers/{__init__,helpers,source}.py +
_hf_format.py).
- sources/transformers/source.py: re-applied Phase A's explicit
TransformersDriver construction in boot().
sources/vllm/source.py's boot_vllm currently sets bridge._forward_impl,
which Phase A removed. The function is dead until the next chunk lifts
it onto VLLMDriver + RemoteBridge — landing as-is so the scaffold tree is
on dev-4.x for incremental Phase B work.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* Revision round 1 of the vLLM system
* Full vLLM integration initial completion
* jupiter patch
* dynamo patching
* pass through batch size
* add max model length to remote bridge boot
* Updating config construction to properly pass through to vLLM
* Wiring TLWorkerExtension
* Additional testing + diagnostic
* Another spot check
* Remove logging
* Fix issue with logit generation on hook_out
* Adjust to use token_ids directly instead of logprobs
* Updating Step 5 and driver bridge cleanup
* Add decoder layer for materializing the residual stream
* Updating tests due to changes on `dev`
* Fix pip install path
* Testing ln_final divergence
* documenting differences in ln_final
* Format cleanup
* Initial batching test
* Add version print
* Setup batching for vllm
* Additional documentation and coverage
* CI Failure cleanup
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* INitial Inspect Driver setup * Additional improvements to inspect * Cleaning up inspect driver flaws * Adding additional coverage to ensure verify interventions & CUDA support * Improving architecture verification process * extended inspect tests to cover more model types * Clean up leaking rand change * Fix inspect memoize bug * Additional inspect improvements * Initial setup of inspect vllm provider * Additional inspect driver feature tightening * migrating dev 4.x changes * format checks * docstring cleanup * Fixing CI issues
* notebook fix * Rebuilding real logits * improve buffers * Fixed softcap bug * per-position interventino * Inspect HF provider servers head-split attention hooks * Resolving final gaps in interp skills
* Initial bug fixes from sweep * cleanup duplication and comments * vllm extra
* multi-gpu setup * Fixing typing errors * parity report improvements * improve variance * test splitting * disabled cache * adjusted tests to be divisible by 2 gpu * fixed bug * tensor decoding * Documented successfuly test * Add pipeline parallelism * bug fix for 0'ed layers * bug fix – dynamic ownership * Update cleanup
* Add dedicated GraniteMoe adapter tests (#1302) (#1524) granite_moe had substantive test coverage, but embedded inside test_granite_adapter.py rather than in its own file — inconsistent with every other architecture listed in #1302 (including granite_moe_hybrid), which each get a dedicated test_<name>_adapter.py. Moves that coverage into test_granite_moe_adapter.py and adds the one piece that was missing: an explicit config-flag test class (GraniteMoe doesn't override config setup, so it inherits dense Granite's flags, but nothing asserted that directly before). test_granite_adapter.py goes back to covering only GraniteArchitectureAdapter, matching its own docstring. * Updating documentation for future deprecation (#1523) * feat: add RWKV-7 (Goose) TransformerBridge adapter (RWKV7ForCausalLM) (#1521) * feat: add RWKV-7 (Goose) TransformerBridge adapter (RWKV7ForCausalLM) Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * refactor: use OpaqueBlockBridge for RWKV-7; introduce OpaqueBlockBridge base class Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * added hook_out_is_single_residual_stream default to OpaqueBlockBridge * style: apply black formatting to opaque_block, ssm_block, __init__ Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: fix import order in generalized_components/__init__.py (isort) Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: fix import order in generalized_components/__init__.py (isort) Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: fix import order in generalized_components/__init__.py (isort) Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * refactor: replace SSM2MixerBridge with GeneralizedComponent for RWKV-7 attn/ffn sublayers Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: fix black formatting in generate_report.py Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> --------- Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> Co-authored-by: jlarson4 <jonahalarson@comcast.net> * feat: add RavenForCausalLM (Huginn) depth-recurrent adapter (#1520) * Add RavenForCausalLM (Huginn) TransformerBridge adapter Adds a TransformerBridge architecture adapter for RavenForCausalLM (tomg-group-umd/huginn-0125), a depth-recurrent decoder with a prelude / weight-tied recurrent core / coda structure, resolving #1469. - raven.py delegates the recurrence to the remote-code HF forward and maps the three physical block lists (prelude / core_block / coda) via SSMBlockBridge, with combined-QKV native attention and a gated MLP. Sets applicable_phases=[] (a random initial latent state and post-residual sandwich norms diverge from the verify_models phases) and supports_fold_ln=False (ln_f is reused mid-network). prepare_loading patches Huginn's remote code for transformers v5 (tied-weights-keys dict form + a weight re-init guard). - Registers in the adapter factory, the model registry (canonical author and description) and supported_models.json; surfaces the recurrence-shape config via both _HF_PASSTHROUGH_ATTRS lists. - Adds synthetic-config unit tests and CI-gated integration tests. * fix: pass num_steps as int to iterate_forward (0-d tensor has no len) * fix: remove unused torch import and use setattr for dynamic cfg attrs * style: apply black formatting to raven.py Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * refactor: introduce OpaqueBlockBridge; use it for Raven instead of SSMBlockBridge Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: apply black formatting to opaque_block, ssm_block, __init__ Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: fix import order in generalized_components/__init__.py (isort) Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * fix: strip dead hook_q/k/v aliases from Raven AttentionBridge Raven uses a combined Wqkv projection (no separate q/k/v submodules), so AttentionBridge's default hook_q/hook_k/hook_v aliases (which target q.hook_out / k.hook_out / v.hook_out) are unresolvable. The upstream test_every_hook_alias_resolves_to_hookpoint audit catches these as 9 dead aliases across prelude / core_block / coda. Strip them by setting an instance-level hook_aliases that omits those three keys, leaving only the aliases that have real HookPoint targets. Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * chore: merge upstream/dev — add RWKV-7 passthrough attrs and registry entry Brings in the RWKV-7 additions from #1521 that landed in dev: - _bridge_builder.py: added RWKV-7 passthrough attrs (num_heads, value_dim, decay/gate/a/v_low_rank_dim, norm_first, norm_bias, fuse_norm, attn_mode, hidden_act) - generate_report.py: added RWKV7ForCausalLM description entry Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> --------- Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * Fix Tracr demo output_label: compiled labels are auto-numbered (#1525) The Tracr demo passes output_label="reverse", but Tracr builds labels as f"{name}_{unique_id}", so the compiled expression's residual labels are reverse_1:* and the demo raises ValueError at the state-dict cell. Pass reverse.label instead of hardcoding the name, and correct the same wrong label shape in the unit-test fixture (real Tracr never emits "reverse:1") and in the infer_tracr_output_label docstring example. Co-authored-by: Claude Fable 5 <noreply@anthropic.com> * Add TransformerBridge adapter for lightweight decoder-only pretraining models (#1519) * Add TransformerLens pretrain bridge adapter * Add TransformerLens pretrain bridge adapter * Setup registry for hookedtransformer deprecation * bug fixes and rebase to 4.x * Hook management improvements * implemented `stop_at_layer` on TransformerBridge * input_to_embed and pos_slice added to bridge * prevent corruption, check_hooks_to_add placed in all locations that need it * Adjusted head detector to be HookedTransformer agnostic, added key value injection to bridge * Add specific migration differences section * Attempting to let pytest run in parallel * Attempting to fix the OOM on the multi-device Full Coverage run --------- Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> Co-authored-by: Delaida Muminovic <delaida.muminovic.22@size.ba> Co-authored-by: Mukund Pandey <mukund.pandey@gmail.com> Co-authored-by: Joseph Quevedo <41499530+dewstend@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Cacapice <kombrellaro@gmail.com>
* Add dedicated GraniteMoe adapter tests (#1302) (#1524) granite_moe had substantive test coverage, but embedded inside test_granite_adapter.py rather than in its own file — inconsistent with every other architecture listed in #1302 (including granite_moe_hybrid), which each get a dedicated test_<name>_adapter.py. Moves that coverage into test_granite_moe_adapter.py and adds the one piece that was missing: an explicit config-flag test class (GraniteMoe doesn't override config setup, so it inherits dense Granite's flags, but nothing asserted that directly before). test_granite_adapter.py goes back to covering only GraniteArchitectureAdapter, matching its own docstring. * Updating documentation for future deprecation (#1523) * feat: add RWKV-7 (Goose) TransformerBridge adapter (RWKV7ForCausalLM) (#1521) * feat: add RWKV-7 (Goose) TransformerBridge adapter (RWKV7ForCausalLM) Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * refactor: use OpaqueBlockBridge for RWKV-7; introduce OpaqueBlockBridge base class Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * added hook_out_is_single_residual_stream default to OpaqueBlockBridge * style: apply black formatting to opaque_block, ssm_block, __init__ Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: fix import order in generalized_components/__init__.py (isort) Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: fix import order in generalized_components/__init__.py (isort) Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: fix import order in generalized_components/__init__.py (isort) Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * refactor: replace SSM2MixerBridge with GeneralizedComponent for RWKV-7 attn/ffn sublayers Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: fix black formatting in generate_report.py Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> --------- Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> Co-authored-by: jlarson4 <jonahalarson@comcast.net> * feat: add RavenForCausalLM (Huginn) depth-recurrent adapter (#1520) * Add RavenForCausalLM (Huginn) TransformerBridge adapter Adds a TransformerBridge architecture adapter for RavenForCausalLM (tomg-group-umd/huginn-0125), a depth-recurrent decoder with a prelude / weight-tied recurrent core / coda structure, resolving #1469. - raven.py delegates the recurrence to the remote-code HF forward and maps the three physical block lists (prelude / core_block / coda) via SSMBlockBridge, with combined-QKV native attention and a gated MLP. Sets applicable_phases=[] (a random initial latent state and post-residual sandwich norms diverge from the verify_models phases) and supports_fold_ln=False (ln_f is reused mid-network). prepare_loading patches Huginn's remote code for transformers v5 (tied-weights-keys dict form + a weight re-init guard). - Registers in the adapter factory, the model registry (canonical author and description) and supported_models.json; surfaces the recurrence-shape config via both _HF_PASSTHROUGH_ATTRS lists. - Adds synthetic-config unit tests and CI-gated integration tests. * fix: pass num_steps as int to iterate_forward (0-d tensor has no len) * fix: remove unused torch import and use setattr for dynamic cfg attrs * style: apply black formatting to raven.py Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * refactor: introduce OpaqueBlockBridge; use it for Raven instead of SSMBlockBridge Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: apply black formatting to opaque_block, ssm_block, __init__ Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * style: fix import order in generalized_components/__init__.py (isort) Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * fix: strip dead hook_q/k/v aliases from Raven AttentionBridge Raven uses a combined Wqkv projection (no separate q/k/v submodules), so AttentionBridge's default hook_q/hook_k/hook_v aliases (which target q.hook_out / k.hook_out / v.hook_out) are unresolvable. The upstream test_every_hook_alias_resolves_to_hookpoint audit catches these as 9 dead aliases across prelude / core_block / coda. Strip them by setting an instance-level hook_aliases that omits those three keys, leaving only the aliases that have real HookPoint targets. Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * chore: merge upstream/dev — add RWKV-7 passthrough attrs and registry entry Brings in the RWKV-7 additions from #1521 that landed in dev: - _bridge_builder.py: added RWKV-7 passthrough attrs (num_heads, value_dim, decay/gate/a/v_low_rank_dim, norm_first, norm_bias, fuse_norm, attn_mode, hidden_act) - generate_report.py: added RWKV7ForCausalLM description entry Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> --------- Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> * Fix Tracr demo output_label: compiled labels are auto-numbered (#1525) The Tracr demo passes output_label="reverse", but Tracr builds labels as f"{name}_{unique_id}", so the compiled expression's residual labels are reverse_1:* and the demo raises ValueError at the state-dict cell. Pass reverse.label instead of hardcoding the name, and correct the same wrong label shape in the unit-test fixture (real Tracr never emits "reverse:1") and in the infer_tracr_output_label docstring example. Co-authored-by: Claude Fable 5 <noreply@anthropic.com> * Add TransformerBridge adapter for lightweight decoder-only pretraining models (#1519) * Add TransformerLens pretrain bridge adapter * Add TransformerLens pretrain bridge adapter * Setup registry for hookedtransformer deprecation * bug fixes and rebase to 4.x * Hook management improvements * implemented `stop_at_layer` on TransformerBridge * input_to_embed and pos_slice added to bridge * prevent corruption, check_hooks_to_add placed in all locations that need it * Adjusted head detector to be HookedTransformer agnostic, added key value injection to bridge * Add specific migration differences section * Attempting to let pytest run in parallel * Attempting to fix the OOM on the multi-device Full Coverage run * Big comment sweep * Comment fix --------- Signed-off-by: Mukund Pandey <mukund.pandey@gmail.com> Co-authored-by: Delaida Muminovic <delaida.muminovic.22@size.ba> Co-authored-by: Mukund Pandey <mukund.pandey@gmail.com> Co-authored-by: Joseph Quevedo <41499530+dewstend@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Cacapice <kombrellaro@gmail.com>
* Fixing tests that were broken by the migration from dev to 4.x * Updating docs for latest information, adding support for vision benchmarking * Clean up dead code in architecture adapters, properly documenting all our slow tests, fixing some test code configs * Test and architecture DRY cleanup * more DRY tests * Refactoring repeated architecture functions into helpers * Reduced duplication across adapters in component mapping and other setup components
* fix(bridge): make tokenizer assignment re-run wiring logic * pipeline fix * pipeline fix * pipeline fix
Co-authored-by: Guan Tong <tguan@gmail.com>
* rebased on dev-4.x * change base branch
…ort (#1580) * feat(native): add param-free pre-norm (LNPre/RMSPre) and fold_ln support * pipeline fix * review changes
* wrappign up final pre-deprecation items * Eliminating Hooked* classes from test suites * Updating goldens revision for latest data * Final demo updates * Codified hookedrootmodule * prep for removal of Hooked* system * Format cleanup * test fix
* HookedTransformer Deprecation * Additional hooked deprecation prep
…dge (#1758) In #1740 the import in LIT_Integration_Demo.ipynb was updated to TransformerBridge, but the model loading cell still called HookedTransformer.from_pretrained, which raised a NameError since HookedTransformer was no longer imported. Update the cell to load via TransformerBridge.boot_transformers with enable_compatibility_mode(). Co-authored-by: Jonah Larson <jonahalarson@comcast.net>
* Fix IOIDataset BOS handling * Document IOIDataset BOS behavior
tiny_model was built from the unseeded global RNG, so its weights varied with xdist scheduling. The linear-LN approximation error is weight-dependent (~17% of random inits exceed the 1e-3 tolerance), making test_correctness_against_actual_ln_forward flaky. Seed the fixture and correct the docstrings that claimed folding LN makes the approximation exact.
…ly Bridges (GPT-2, Pythia, GPT-NeoX) (#1778) * refactor(backward_lens): derive MLP weight layout from Bridge oracle Replace the GPT2ArchitectureAdapter isinstance check and the Conv1D-only projection guard with the Bridge's own weight_layout_in_out oracle, so support is decided per MLP projection instead of by model class. Dense MLPBridge input/output projections now resolve to "in_out" (Conv1D, e.g. GPT-2) or "out_in" (torch.nn.Linear, e.g. Pythia/GPT-NeoX), and the resolved layout threads through gradient-factor capture instead of the previous hardcoded "in_out". A projection whose wrapped module the oracle cannot orient is rejected with a clear error. Rename _require_raw_gpt2_bridge, _get_gpt2_mlp_projections, _capture_gpt2_mlp_gradient_factors, _GPT2GradientCapture, and _GPT2LayerGradientFactors to their model-agnostic names, and update call sites and docstrings accordingly. Repurpose the non-Conv1D rejection test to assert rejection of an unorientable component, since torch.nn.Linear projections are now accepted. Behavior-preserving for GPT-2: unit and integration suites for this module pass unchanged. * test(backward_lens): cover nn.Linear layout + Pythia-70m reconstruction Add model-free unit tests for dense-MLP projection discovery that exercise the Bridge weight-layout oracle directly, using real MLPBridge/LinearBridge instances rather than a booted model: a torch.nn.Linear MLP resolves to the "out_in" layout and accepts its transposed weight shape, a Conv1D MLP resolves to "in_out" with GPT-2-parity shapes, an unorientable wrapped module is rejected with a clear error, and a gated MLP is rejected regardless of orientation. Parametrize the core gradient-reconstruction integration test over a raw GPT-2 Bridge and a raw Pythia-70m Bridge sharing the same prompt and target token, asserting the same reconstruction tolerance bands for both and that each model's projections resolve to its own weight layout without any model-class conditional. Extend the weight/hook state-preservation and cleanup-on-failure integration tests to both models as well. * docs(backward_lens): state dense-MLP contract; generalize demo note Update the Backward Lens doc and demo notebook to describe the generalized dense-MLP decoder-only contract (GPT-2 and Pythia/GPT-NeoX) instead of the prior GPT-2-only framing. Explain both weight layouts (Conv1D [in, out] and torch.nn.Linear [out, in]) and that BackwardLens reads the layout from the Bridge projection component rather than the model class. Add a minimal Pythia-70m reconstruction example to the doc and a corresponding single-layer demo cell in the notebook.
* fix: reject fabricated MoE weights in SVDInterpreter * test: sparse MoE layers are omitted from get_params, not zero-filled
As a newcomer to these analysis methods, I found there was no single guide explaining which tool fits which research question. Add a question-led overview of the available analysis tools, their requirements, costs, and interpretation limits, and link it from the documentation index.
…rmerBridge (#1781) * feat(attribution_patching): per-head hook requirements on the Bridge Edges into and out of attention heads read attn.hook_result (writer, a head's contribution before the sum into resid) and the split attn.hook_q_input/hook_k_input/hook_v_input (reader, each head's separate Q/K/V input). Both hook families exist on the Bridge unconditionally but only fire once cfg.use_attn_result and cfg.use_split_qkv_input are on. Add an edge-granularity variant of the required hook set, a helper that turns on both flags before caching, and generalize the missing-hook ValueError enumerate_nodes already raises into a shared check so any granularity's cache can be validated against it instead of duplicating the check per graph. * feat(attribution_patching): edge enumeration in the typed graph Add enumerate_edges(model, cache), producing every writer -> reader pair in the residual-stream graph. Extend NodeKind with the reader kinds q_input/k_input/v_input/mlp_in and extend Node's invariants and hook_name accordingly. At a fixed sequence position the residual stream is a running sum, so a reader is fed by every writer enumerated ahead of it; the graph is built position-by-position, tracking that "available" writer set. A writer feeding both a direct edge and a through-MLP edge yields two distinct writer-reader pairs rather than a summed one; a dedicated uniqueness guard enforces that invariant. Edges into the MLP additionally need the MLP-entry hook point (hook_mlp_in per layer), so edge enumeration validates its presence alongside the per-head hook set already required for edge granularity. * feat(attribution_patching): edge scoring in attribution_patch Wire granularity="edge" through attribution_patch: enumerate the writer -> reader edges from the cached activations and gradients, score each with (a_clean[writer] - a_corrupt[writer]) . d(metric)/d(input of reader), and populate AttributionResult.edge_scores. node_scores for an edge sweep becomes each writer's aggregate over its own outgoing edge scores rather than a direct node measurement. Remove the granularity == "edge" raise from EdgeAttributionConfig; the ig_steps > 1 raise is untouched. top_edges() ranks by absolute score and returns the ranked edge list instead of raising. Enable use_hook_mlp_in alongside the existing per-head flags so MLP-input edges have a populated reader hook. * test(attribution_patching): exact-patch parity and mutation-checked reconstruction Add the two Risk-1 guards for edge scoring: a genuine single-edge activation patch (adding one writer's clean-minus-corrupt delta into a reader's cached input, everything else held at corrupt) matches _edge_effects' sign and, on this fully linear toy Bridge, its magnitude too; and perturbing one writer's captured contribution changes only the edges that writer feeds, leaving every other edge's score (and sibling heads/positions in the same cached tensor) untouched. Also add a sum-to-total check for the edge form: a reader's incoming edge scores sum to the same value as scoring its own cached input-delta against its own gradient directly, an identity that holds unconditionally from the residual stream's additive construction. generic_activation_patch is not used for the exact-patch check: its model: HookedTransformer parameter is enforced at runtime by this repo's jaxtyping/beartype pytest configuration, which rejects any argument that is not actually a HookedTransformer instance, including a real TransformerBridge. The patch is driven directly through the Bridge's own hooks(), the same mechanism generic_activation_patch uses internally. * feat(attribution_patching): scope edge hook-flag mutation with save/restore Add _edge_hook_flags, a context manager that enables the Bridge flags an edge sweep needs (use_attn_result, use_split_qkv_input, use_hook_mlp_in) and restores the caller's prior flag state on exit, including on error. The prior one-way helper only turned flags on, so every forward after an edge sweep kept materializing the per-head tensors, and a caller who arrived with use_attn_in set hit the use_split_qkv_input exclusivity error and was still left mutated. The context manager snapshots all four flags, clears use_attn_in before enabling the split input, and restores each flag in a finally block so a raise mid-sweep cannot leave the model mutated. Extend the edge-hook toy Bridge with use_attn_in and its setter (mutually exclusive with use_split_qkv_input) and cover enable-in-scope, restore on normal exit and on raise, already-enabled flags, and the use_attn_in caller path. * feat(attribution_patching): add terminal logits reader and parallel_attn_mlp handling to edge enumeration * feat(attribution_patching): batch edge scoring and scope the sweep's hook-flag mutation * test(attribution_patching): score edge parity and mutation at the readout position and cover multi-pair averaging and head-to-logits ranking * test(attribution_patching): construct edge toy bridges in eval mode
…1785) * 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. * 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. * 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. * 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. * 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. * 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. * 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(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. * 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. * 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). * refactor(model_bridge): gated-MLP relevance rules without recompute * fix(model_bridge): LN-rule treats a missing norm weight as identity * fix(model_bridge): relevance-rule mount coverage
* fix: raise on stop_at_layer when no 'blocks' stack is registered Replace the stop_at_layer reject-list with an allowlist guard and check self._modules instead of hasattr for both stop_at_layer and start_at_layer. Addresses #1769 * docs: note NotImplementedError in stop_at_layer docstring test: cover a wrapped model exposing its own .blocks Addresses #1769 * test: cover a wrapped model exposing its own .blocks Addresses #1769 --------- Co-authored-by: jlarson4 <jonahalarson@comcast.net>
4.x Merge into `dev`
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
18 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.
Description
The Next Major Release
Type of change
Checklist: