Drop optional unused aliases on hybrid architectures by KYinXu · Pull Request #1579 · TransformerLensOrg/TransformerLens

KYinXu · 2026-07-31T05:53:22Z

Description

When hybrid layers skip an optional submodule (e.g. attn on GDN / Mamba-2 layers), the shared BlockBridge template still declared HT aliases like hook_attn_out → attn.hook_out. _register_aliases then failed to resolve them and emitted a harmless warning per hybrid layer (notably when loading Qwen3.5; see also #1547).

This PR prunes those inapplicable aliases in setup_submodules when an optional submodule is skipped—per layer, so full-attention siblings keep their HT aliases and real unresolved-alias bugs still warn. Unit tests cover prune behavior and no warning on register. The Jacobian Lens Qwen3.5 demo cell is updated now that boot no longer expects that stderr.

Fixes #1578

Type of change

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
This change requires a documentation update

Screenshots

Warnings no longer appear

Checklist:

I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes
I have not rewritten tests relating to key interfaces which would affect backward compatibility

jlarson4

Nice, tightly-scoped fix @KYinXu!

A couple small requests before merging, detailed below

jlarson4 · 2026-07-31T17:19:06Z

+
+
+def _prune_hook_aliases_for_skipped(component: GeneralizedComponent, skipped: list[str]) -> None:
+    """Drop hook aliases whose target path starts at a skipped optional submodule.


The prune fires for every optional skip across 12 architectures, including native SSM aliases on nemotron_h/zamba2. Can we update the docstring to be more general, and add one non-attn test case?

jlarson4 · 2026-07-31T17:19:50Z

+                to_drop.append(alias_name)
+            elif len(kept) != len(target):
+                aliases[alias_name] = kept
+        elif _first_segment(target) in skipped_set:


On OLMo-Hybrid, BlockBridge retargets hook_mlp_out onto the optional ln2_post (block.py:78), so this branch drops it on linear-attention layers that do have an mlp, and removes the warning that flagged it. Can the prune keep aliases with a live fallback target?

jlarson4 · 2026-07-31T17:20:31Z

+        assert "hook_resid_pre" in blocks[3].hook_aliases
+        assert "hook_mlp_out" in blocks[3].hook_aliases
+
+    def test_pruned_aliases_do_not_warn_on_register(self):


Both new tests call _register_aliases() directly, so nothing asserts a hybrid bridge booting warning-free through TransformerBridge.__init__. Can you add a boot-level assertion using the config-only tiny-bridge fixture in test_qwen3_5_adapter.py:562?

jlarson4 · 2026-07-31T17:21:29Z

+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            blocks[3]._register_aliases()
+        unresolved = [


AttnAdapter declares no ln2/mlp, so block 3 still emits two legitimate unresolved warnings. This forces the substring filter instead of asserting the list is empty. Can the fixture carry those submodules so the assertion can be unconditional?

KYinXu · 2026-08-06T05:07:46Z

@jlarson4 just addressed changes, sorry about delay was busy with other tasks earlier in the week and had to catch up.

jlarson4 · 2026-08-06T13:26:23Z

@KYinXu No worries at all about the delay, thank you for coming back and wrapping those up. I will review shortly

jlarson4 · 2026-08-06T15:00:27Z

Great work @KYinXu merging now!

* Add Jacobian lens fitting guide (#1544) * Quantify Jacobian Lens causal swap success (#1545) Co-authored-by: Dreamer431 <113128214+Dreamer431@users.noreply.github.com> * test(integration): add oracle parity test for JacobianLens (#1539 Tier-1) (#1543) * test(integration): add oracle parity test for JacobianLens (#1539 Tier-1) Compares TransformerBridge JacobianLens.readout() against the reference anthropics/jacobian-lens oracle (pinned to 581d398) on google/gemma-2-2b-it across 75 layer x prompt cells (5 prompts x 15 sampled layers). Pass criteria per the #1539 spec (matching #1505 spike numbers): - Worst-case top-8 token overlap >= 7/8 in every cell - Spearman rank-correlation >= 0.95 on the top-64 logit union per cell The oracle is installed at test time via pip from the pinned commit so the threshold is reproducible independent of upstream drift. Reuses the bridge's original_model + tokenizer to avoid a second model copy in RAM. * style: apply black formatting (line-length=100) * fix: remove unused Dict, Tuple typing imports (pycln) * style: fix black formatting for py310 target (double blank lines) * test: use pytest.importorskip for oracle dep; add oracle-parity CI workflow Replace subprocess pip-install fixture with pytest.importorskip so the test skips gracefully in standard uv venvs (no pip present) and does not mutate the developer environment with no cleanup. Add .github/workflows/oracle-parity.yml — a dedicated workflow that reads ORACLE_COMMIT from the test file (single source of truth) and installs the oracle out-of-band before running the @pytest.mark.slow suite. Triggers on workflow_dispatch and on pushes that touch the test or workflow file, keeping oracle runs opt-in for PR checks. * Visual Encoders (ViT, DeiT) Support Rollout (#1546) * Add ViTArchitectureAdapter to supported architectures * Add ViTArchitectureAdapter to architecture factory * Add ViT and DeiT models to model registry * Add new model descriptions for Vision Transformers and Wav2Vec2 * Add ViTArchitectureAdapter for vision models Implement ViT/DeiT architecture adapter for model bridging. * Create vision_embedings.py * Add VisionClassifierHeadBridge for CLS token classification Implement VisionClassifierHeadBridge to handle CLS token slicing for classification. * Add visual model configuration to transformer bridge * Clarify pixel_values usage for multimodal and vision models Updated documentation for pixel_values parameter to clarify its use with vision models. * Update bridge.py * Update bridge.py * Rename vision_embedings.py to vision_embeddings.py * Define vision model and classification architectures Added vision model architectures and classification heads. * Add support for vision architectures in transformers * Refactor VisionClassifierHeadBridge to use pooled output Updated the VisionClassifierHeadBridge to directly use an already-pooled CLS token instead of slicing from the sequence output. Adjusted the forward method to reflect this change and improved error handling for the original component. * Update vit.py * Add unit tests for ViTArchitectureAdapter This file contains unit tests for the ViTArchitectureAdapter, covering component mapping, configuration flags, weight conversions, and model preparation methods. * Create test_vit_adapter.py * Update transformers.py * Update vit.py * Update vit.py * Update vit.py * Fix type hint for get_remote_component method * Fix type hint for get_remote_component method * Change import of torch to torch.nn in vit.py * Update vit.py * Re-add dummy 'mlp' attribute injection for ViTLayer Reintroduce a patch_layers function to inject a dummy 'mlp' attribute into ViTLayer blocks for MLPBridge compatibility. * Refactor ViTLayer handling by removing patch_layers Removed the patch_layers function and its call, which injected a dummy 'mlp' attribute into ViTLayer blocks. Updated comments for clarity regarding the MLPBridge container. * Add dummy 'mlp' attribute to ViTLayer blocks Inject a dummy 'mlp' attribute into ViTLayer blocks to satisfy hasattr check for TransformerLens. * Update vit.py * Remove TestViTConfigNCtx and related test case Removed deprecated TestViTConfigNCtx class and its test case for n_ctx. * Enhance ViTLayer with MLP wrapper and fix forward method Added a non-circular MLP wrapper to ViTLayer blocks and fixed tuple-chaining bug in forward method. * Refactor ViT layer forward pass handling Refactor forward pass handling for ViT layers to safely unpack tuple outputs and ensure compatibility with the model's internal structure. * Refactor ViTLayer forward pass handling Refactor forward pass handling for ViTLayer to fix tuple-chaining bug and ensure compatibility with HF model outputs. * Fix tuple handling in ViTLayer forward method Modified the forward method to handle tuple inputs and outputs for ViTLayer, ensuring compatibility with Tensor expectations. * Reorder model prefix checks for better clarity * Update vit.py * Detect model class name in prepare_model method Added detection for model class name in prepare_model method. * Simplify prefix determination for ViT models Refactor model prefix detection logic for ViT and DeiT models. * Implement fixture for distilled DeiT model testing Added a fixture to load the distilled DeiT model for testing. * Update DeiT bridge tests for bare model handling Refactor tests for DeiT bridge to accommodate bare model behavior and update assertions accordingly. * Set architecture in Hugging Face model configuration * Support DeiTLayer in patch_layers function * sort * Replace direct attribute assignment with setattr * black fix * fix formatting after merge * Update vit.py * Update ViT adapter test paths for consistency * Remove redundant test for n_ctx in prepare_loading Removed test for prepare_loading not affecting n_ctx. * black sorted * black reorder * Refactor vit_bridge and vit_bare_bridge fixtures * temp support up to transformers 5.8.0 * support transformers 5.13.0 * format fixed. Unit test all passed. Intergration test all passed. should be good to go * Update vit.py * Clarify tokenizer support in ViTArchitectureAdapter Added comment to clarify the lack of tokenizer support for vision models. * Remove head_dim assignment from hf_config Removed unused head_dim assignment from hf_config. * Update vit.py * Update bridge.py * Add VisionEmbeddingsBridge and VisionClassifierHeadBridge * Update test_vit_adapter.py * Update test_vit_adapter.py * Update test_vit_adapter.py * Update test_vit_adapter.py * Improve compatibility mode error and output handling Updated error message for clarity and added handling for last_hidden_state in output. * Refactor test to check output type and shape Update test to verify that the forward method returns a tensor instead of a raw HF output object. Adjust assertions to match the expected behavior after changes in bridge.py. * Update bridge.py * Update test_vit_adapter.py * Remove obsolete tests from TestViTPrepareLoading Removed deprecated tests for prepare_loading() in TestViTPrepareLoading. * Update test_vit_adapter.py * Update test_vit_adapter.py * formatted --------- Co-authored-by: Jonah Larson <jonahalarson@comcast.net> * Add Starcoder2 architecture adapter (#1533) Co-authored-by: jlarson4 <jonahalarson@comcast.net> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> * Docs: Jlens Qwen3.5-4b Demo (#1547) * Run experiments with Qwen-3.5 architecture support * Run notebook cells * Add lfm2 tiny integration test (#1552) * Add lfm2 tiny integration test * Fix formatting * Add AST (Audio Spectrogram Transformer) Adapter (#1484) * chore: save WIP on V3 transformerbridge migration re add ast import and add to factory :wq y wq :wq * feat(ast): migrate AST adapter to V3 TransformerBridge and component_mapping * refactor(ast): resolve PR feedback for docstrings, prefix-awareness, unit tests * test(ast): split parity to integration folder, add load_weights boot test, fix audio classification load path, and add to ARCHITECTURE_DESCRIPTIONS * fix(ast): union audio classification sets, specific boot test assertions and two comment typo fixes * Verificaiton for lapa (#1556) * ViT and AST model verification (#1582) * verified a few models for ViT and AST * improved vision testing for ViT models * fix(bridge): return W_in/W_out/W_gate in TL orientation for nn.Linear-backed models (#1558) * feat: respect prepend_bos and add return_input_tokens flag * fix torch orientation bug * review changes * review changes * pipeline fix * deprecate remaining hooked entry points (#1592) * deprecate remaining hooked entry points * test: account for hooked transformer warning in notebook * fix: address deprecation warning review feedback * fix: correct encoder deprecation warning stacklevel * support loading fit checkpoints in JacobianLens.load() (#1574) * support loading fit checkpoints in JacobianLens.load() * fix: black formatting and update conflicting test for checkpoint load * address jlarson4 review: preserve target_layer, add reference fixture, fix tuned-lens note - Remove "target_layer" from _FIT_RESERVED_KEYS so it survives checkpoint conversion and validate_model() can refuse non-final-target lenses - Add test_load_checkpoint_mirrors_fit_payload_schema: fixture matches the exact keys fit() produces so format drift causes a test failure - Fix tuned-lens note: it is the Jacobian artifact format that has no bias slot, not the tuned-lens format; tuned-lens translators are affine (weight + bias) * address jlarson4 review: read n_done key in _from_checkpoint_payload Real checkpoint writers (reference package) store the prompt count as n_done, not n_prompts. Prefer n_done with n_prompts as fallback so genuine checkpoints are not rejected with n_prompts=0. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NiwNUm3YFj9yAuSBuGDnd8 * address jlarson4 review: align checkpoint loader with reference write_checkpoint() schema - _from_checkpoint_payload: read n_done first (real checkpoints use n_done, not n_prompts) - _from_checkpoint_payload: infer d_model from jacobian_sum matrix shape (real checkpoints have no d_model key) - _from_checkpoint_payload: harvest top-level target_layer into metadata (reference format stores it at top level, not nested) - _from_checkpoint_payload: guard empty jacobian_sum with a clear ValueError before attempting shape derivation - load() docstring: update Fit checkpoint schema to reflect the real 6-key reference format - tests: replace test_load_checkpoint_with_zero_n_prompts_raises with two tests - tests: rewrite test_load_checkpoint_mirrors_fit_payload_schema to use verbatim 6-key reference payload - tests: add test_load_checkpoint_harvests_flat_provenance_and_strips_fit_keys - docs: update schema table — replace n_prompts/d_model with real 6-key format - docs: note d_model inferred from matrix shape - docs: document target_layer as deliberate exception to fit-key stripping --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> * Drop optional unused aliases on hybrid architectures (#1579) * Drop optional unused aliases * feat: assign fallbacks on pruned * Improve unittests * Add kurtosis-profile validation test for JacobianLens (#1539 Tier-1) (#1616) * Add kurtosis-profile validation test for JacobianLens (#1539 Tier-1) Asserts the workspace-band signature as relative structure rather than absolute levels, per the cross-family measurement in #1539: band rise vs the model's own early-third baseline, lens-specificity vs the logit-lens control through the identical code path, a gpt2-small negative control, and final-layer identity-transport agreement between arms. * Fix gpt oss olmo3 parity (#1621) * Resolution for issue 1619 * Updated for 1620 * Add verification script * Fixed 1619 on HF * Verification script repair * fixing per-layer olmo * cleanup --------- Co-authored-by: abhi <abhinavbellapu@berkeley.edu> Co-authored-by: emerardd <113128214+emerardd@users.noreply.github.com> Co-authored-by: Dreamer431 <113128214+Dreamer431@users.noreply.github.com> Co-authored-by: Mukund Pandey <mukund.pandey@gmail.com> Co-authored-by: Jiankun Wei <72998341+david-wei-01001@users.noreply.github.com> Co-authored-by: SanjidMzi <56235075+SanjidMzi@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Kyle Yin Xu <123780557+KYinXu@users.noreply.github.com> Co-authored-by: Syed Adil Ahmed <tensorcruncher@gmail.com> Co-authored-by: Dylan <159935143+dylanberens@users.noreply.github.com> Co-authored-by: Md.Sadiq <mohammadsadiq4950@gmail.com> Co-authored-by: msaule <sau24006@byui.edu> Co-authored-by: Priyanka Bajaj <42418272+priyanka25aug@users.noreply.github.com>

Drop optional unused aliases

586ac39

KYinXu changed the title ~~Drop optional unused aliases~~ Drop optional unused aliases on hybrid architectures Jul 31, 2026

jlarson4 linked an issue Jul 31, 2026 that may be closed by this pull request

[Proposal] Hybrid architectures warn on unresolved optional 'hook_attn' aliases when they are skipped #1578

Closed

1 task

jlarson4 reviewed Jul 31, 2026

View reviewed changes

KYinXu added 2 commits August 6, 2026 01:06

feat: assign fallbacks on pruned

8fe80fa

Improve unittests

95d3562

jlarson4 merged commit a18739d into TransformerLensOrg:dev Aug 6, 2026
25 checks passed

Navigation Menu

Uh oh!

Drop optional unused aliases on hybrid architectures - #1579

Drop optional unused aliases on hybrid architectures#1579
jlarson4 merged 3 commits into
TransformerLensOrg:devfrom
KYinXu:fix/gdn-adaptation

KYinXu commented Jul 31, 2026

Uh oh!

jlarson4 left a comment

Uh oh!

jlarson4 Jul 31, 2026

Uh oh!

jlarson4 Jul 31, 2026

Uh oh!

jlarson4 Jul 31, 2026

Uh oh!

jlarson4 Jul 31, 2026

Uh oh!

KYinXu commented Aug 6, 2026

Uh oh!

jlarson4 commented Aug 6, 2026

Uh oh!

jlarson4 commented Aug 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		def _prune_hook_aliases_for_skipped(component: GeneralizedComponent, skipped: list[str]) -> None:
		"""Drop hook aliases whose target path starts at a skipped optional submodule.

Uh oh!

Conversation

KYinXu commented Jul 31, 2026

Description

Type of change

Screenshots

Checklist:

Uh oh!

jlarson4 left a comment

Choose a reason for hiding this comment

Uh oh!

jlarson4 Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

jlarson4 Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

KYinXu commented Aug 6, 2026

Uh oh!

jlarson4 commented Aug 6, 2026

Uh oh!

jlarson4 commented Aug 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants