fix(bridge): return W_in/W_out/W_gate in TL orientation for nn.Linear-backed models - #1558
fix(bridge): return W_in/W_out/W_gate in TL orientation for nn.Linear-backed models#1558MdSadiqMd wants to merge 7 commits into
Conversation
|
Hey @jlarson4, the pipeline passed. Please review it when you're free |
|
Will do, I probably will not have time to review it today but will get reviewing all of these new PRs first thing tomorrow |
|
Sure, thank you for maintaining this repo |
|
Of course! Happy to do it! Full review coming shortly |
jlarson4
left a comment
jlarson4
left a comment
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.
|
Hey @jlarson4 implemented the changes. Let me know if anything is required from my side |
jlarson4
left a comment
jlarson4
left a comment
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: |
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") |
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?
2 participants
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: