fix(bridge): make native TransformerBridge state_dict()/load_state_dict() true inverses - #1598
Conversation
…ct() true inverses state_dict() emits TL-renamed keys, but load_state_dict() only matched raw native names, so a round trip silently loaded nothing and strict=True was silently downgraded to strict=False. Adds the inverse key mapping (including aliased parameters reachable via multiple attribute paths, e.g. GPT-2's split q/k/v views into c_attn) and proper missing/unexpected key accounting that raises under strict=True. Fixes TransformerLensOrg#1587
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for taking a look at this @LightWork666! My intention is to take this implementation as our primary solution, and keep a couple tests and other small elements from #1591 and #1595. I do have a couple change requests, please take a look when you have a moment. Your hard work on this is much appreciated
| clean_key = actual_key.replace("._original_component", "") | ||
| clean_to_actual[clean_key] = actual_key | ||
| actual_to_clean[actual_key] = clean_key | ||
| clean_to_actual[actual_key.replace("._original_component", "")] = actual_key |
There was a problem hiding this comment.
A complete raw-HF-format gpt2 checkpoint raises under strict=True (337 missing) even though strict=False restores the forward pass exactly. The clean branch maps one alias while required_actual_keys demands the union (lines 3546-3549). Can the strict accounting treat a required key as satisfied when any alias of its tensor is written?
There was a problem hiding this comment.
You're right, thanks for catching this. The bug was that required_actual_keys demanded every alias in a TL key's group be present, but the clean-key branch only ever writes one alias per key. Since aliases of the same TL key share the underlying storage, writing any one of them already updates what forward() reads for all of them, so treating the group as satisfied by any single alias is correct. Fixed in the latest commit: missing_keys now checks any(k in mapped_state_dict for k in actual_keys) per group instead of requiring the full union. Verified against your exact repro (gpt2, clean-key dict, strict=True): 0 missing, 0 unexpected now.
| assert torch.equal(reloaded_raw[key], value), f"{key} did not round-trip" | ||
|
|
||
|
|
||
| @pytest.mark.slow |
There was a problem hiding this comment.
This is the only test in any of the three competing PRs that exercises real HF-key conversion, but slow is deselected by make unit-test and the MPS job only runs for main-targeted PRs. Can you add an unmarked boot_native test that builds a multi-alias clean-key dict at strict=True?
There was a problem hiding this comment.
Added test_native_clean_key_dict_with_partial_aliases_does_not_raise_strict. No HF download needed; it uses boot_native's own _original_component wrapping, which turns out to alias internally too (not just gpt2's c_attn split), so it exercises the same bug without needing the slow tier. Also added a second slow test (test_boot_transformers_clean_key_dict_does_not_raise_strict) that's your exact repro on real gpt2, so there's a permanent regression test for the specific case you hit too.
…s written load_state_dict's strict=True missing-keys accounting required every actual-key alias of a shared-storage TL key to be present in the input, even though writing any one alias already updates what forward() reads for all of them (they're views onto the same Parameter). A complete raw-HF-format checkpoint using clean keys - which map to exactly one alias per key - triggered false "missing key" errors under strict=True despite loading correctly. Reported by @jlarson4 on gpt2 (337 false missing keys); reproduced and fixed here, plus a fast boot_native regression test that doesn't need a real model download.
|
@LightWork666 thanks for the updates! Looks good, merging now |
2 participants
Fixes #1587.
The bug
On a native
TransformerBridge,state_dict()returns TL-renamed keys (embed.weight,blocks.0.attn.q.weight, ...), butload_state_dict()only knew how to match raw native parameter names. Sobridge.load_state_dict(bridge.state_dict())silently did nothing — no error, no warning, params just stayed whatever they were before. Worse,strict=Truewas silently downgraded tostrict=Falsewhenever the key counts didn't line up, so there was no way to even notice the round trip had failed.The fix
load_state_dictnow builds the inverse of the TL-key renaming thatstate_dict()applies (_tl_key_to_actual_keys), so it can map TL-format keys back to the underlying native parameter paths before handing them to the wrapped model's ownload_state_dict. One wrinkle: some bridge components expose the same underlying parameter through more than one attribute path — e.g. GPT-2'sq/k/vare views into the wrapped module's combinedc_attnweight, reachable both through a block-level shortcut and through the nested_original_componentchain. Writing to only one of those paths leaves the model's actual forward pass untouched even thoughstate_dict()looks fine, so the mapping keeps every alias for a given TL key, not just the first one found.The silent
strict=True → Falsedowngrade is gone. Missing/unexpected keys are now computed properly (scoped to what the TL state dict actually needs) and raise a realRuntimeErrorunderstrict=True, matching howHookedTransformeralready behaves.The raw-key loading path used by
tracr(make_tracr_transformer_bridge_state_dict) still works — raw native keys are matched directly before falling through to the TL-key path.Testing
tests/unit/model_bridge/test_state_dict_round_trip.py: round-trip actually overwrites params (not a no-op),strict=Trueraises on both missing and unexpected keys,strict=Falsedoesn't raise on a partial dict, the tracr raw-key path still loads, and a real GPT-2 case checking forward-pass logits match after a zero-and-reload cycle (this last one is what caught the aliasing issue above — a naive key-rename fix passes the round-trip-key-equality check but still produces different logits, because it never touches the aliased storage).uv run mypy .— clean.tests/unit/model_bridge/suite (207 files) in isolated batches; everything passes except a handful of pre-existinggenerate()/KV-cache crashes already tracked as an upstream PyTorch/HF bug on Apple Silicon intests/QUARANTINES.md— confirmed viagit stashthat those reproduce identically on unmodifieddev-4.x, unrelated to this change.Note on #1595
I noticed after finishing this that #1595 is already open for the same issue, taking a related but different approach (it detects aliasing by comparing live tensor identity rather than by TL-key-name collisions, which is arguably a more principled check). I'm submitting this anyway since it was already done and transparency seemed better than not mentioning it. Happy to have the maintainers pick whichever they prefer, or close this if #1595 is the better fix.