Repository navigation
fix(bridge): return W_in/W_out/W_gate in TL orientation for nn.Linear-backed models - #1558
Merged
jlarson4 merged 7 commits intoAug 3, 2026
Conversation
Contributor
Author
|
Hey @jlarson4, the pipeline passed. Please review it when you're free |
Collaborator
|
Will do, I probably will not have time to review it today but will get reviewing all of these new PRs first thing tomorrow |
Contributor
Author
|
Sure, thank you for maintaining this repo |
Collaborator
|
Of course! Happy to do it! Full review coming shortly |
jlarson4
reviewed
Jul 31, 2026
jlarson4
left a comment
jlarson4
left a comment
Collaborator
There was a problem hiding this comment.
This looks great! Thanks for tackling it @MdSadiqMd. Just a couple small changes listed below
One side effect of the fix I noticed: but get_params() now disagrees with bridge.W_in on the same object. get_params_util.py:128-129 and :148 still return raw .weight (measured (64,32) vs (32,64)). It should be fixable by sourcing from the properties you just setup.
Contributor
Author
|
Hey @jlarson4 implemented the changes. Let me know if anything is required from my side |
jlarson4
reviewed
Jul 31, 2026
jlarson4
left a comment
jlarson4
left a comment
Collaborator
There was a problem hiding this comment.
Two more small comments on this one. Should be good after these are resolved
| None falls back to shape heuristic. | ||
| pattern: "in" for W_in/W_gate [d_model, d_mlp], "out" for W_out [d_mlp, d_model] | ||
| """ | ||
| if layout is None: |
Collaborator
There was a problem hiding this comment.
GIDD's ScaledLinear is a bare nn.Module and the size-ordering guess silently mis-orients whenever d_mlp < d_model. We assumed "typical" here, but untested. Could this key off the module's declared in_features/out_features instead? If not, can we at least document the blind spot?
| raise AttributeError("No 'in' submodule on this MLP bridge") | ||
| weight = in_module.weight | ||
| layout = self._weight_layout_in_out(in_module) | ||
| return self._normalize_mlp_weight(weight, layout, pattern="in") |
Collaborator
There was a problem hiding this comment.
bridge.get_params()["blocks.0.mlp.W_in"] still returns [d_mlp, d_model] while this property now returns [d_model, d_mlp]. Could the params path reuse these accessors, the way it already reshapes KV weights in get_params_util.py on line 9Z?
jlarson4
added a commit
that referenced
this pull request
Aug 7, 2026
* 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>
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.
Description
On nn.Linear-backed (native) TransformerBridge models,
W_in/W_out/W_gatewere returned transposed relative to TL convention, while Conv1D-backed models (GPT-2) returned TL orientation. Neuron-level analysis migrated from HookedTransformer (Grokking/Othello/superposition demos read.W_in/.W_out) silently computed garbage.Root cause:
TransformerBridge.W_incalled_stack_block_params("mlp.W_in")without a reshape callback, unlike attention weights which pass_reshape_qkv/_reshape_o. MLP weights passed through raw fromoriginal_component.weight.Fix: Add
_weight_layout_in_out()helper to MLPBridge (same pattern as AttentionBridge), replace property_aliases with actual@propertymethods that transpose nn.Linear weights to match TL convention.[n_layers, d_mlp, d_model][n_layers, d_model, d_mlp][n_layers, d_model, d_mlp][n_layers, d_model, d_mlp]Fixes #1554
Type of change
Checklist: