fix(bridge): preserve split-component views under load_state_dict(assign=True) - #1660
fix(bridge): preserve split-component views under load_state_dict(assign=True)#1660LightWork666 wants to merge 1 commit into
Conversation
…ign=True) assign=True replaces each target parameter object instead of copying into existing storage, which desyncs view-backed split components (gpt2's q/k/v, gate/up) from the combined weight (c_attn, gate_up_proj) they share storage with -- original_model.state_dict() silently keeps stale data. Route those specific keys through an explicit in-place copy_() instead, and raise a clear error if a shape mismatch makes that unsafe. Fixes TransformerLensOrg#1637
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Great solution to this @LightWork666! Just a couple comments, let me know if you have any questions
| if target is None or not _is_view_backed(target): | ||
| passthrough_items[key] = value | ||
| continue | ||
| if tuple(target.shape) != tuple(value.shape): |
There was a problem hiding this comment.
Only shape is guarded, an fp16 (or MPS-resident) q.weight under assign=True silently copy-converts into the fp32 CPU view while passthrough keys adopt the incoming dtype/device, yielding a silently mixed-dtype or split-device model from one call. Can you extend this guard to target.dtype != value.dtype or target.device != value.device, with the message naming which property mismatched?
| f"Got {tuple(value.shape)}, expected {tuple(target.shape)}." | ||
| ) | ||
| with torch.no_grad(): | ||
| target.data.copy_(value) |
There was a problem hiding this comment.
The loop copies as it validates, so the first mismatch raises with every earlier view-backed key already written and everything else unapplied. PyTorch's loader applies & then aggregates, but in our use case a first pass over mapped_state_dict collecting all mismatches before any copy_ is trivial since the dict is fully materialized.
| passthrough_items = {} | ||
| for key, value in mapped_state_dict.items(): | ||
| target = current_state_dict.get(key) | ||
| if target is None or not _is_view_backed(target): |
There was a problem hiding this comment.
A raw-native-key load of c_attn's actual key takes this passthrough branch, torch replaces the parameter object, and the split views are orphaned: q<->c_attn storage sharing breaks, bridge forward reads stale values while save_pretrained() exports the new ones. A storage-group check (route non-view keys whose storage is shared by a view-backed sibling through copy_ too) would close that gap here.
|
|
||
|
|
||
| @pytest.mark.slow | ||
| def test_boot_transformers_assign_true_does_not_leave_combined_weight_stale(): |
There was a problem hiding this comment.
Can you also add a fast test on a tiny from_config model via build_bridge_from_module? A tiny Phi-3/GLM config would cover the JointGateUpMLPBridge path and a tiny GPT2Config the QKV path, giving us coverage in CI (slow tests don't run in CI because the runners aren't large enough).
| assert not _is_view_backed(torch.nn.Parameter(torch.randn(4, 4))) | ||
|
|
||
|
|
||
| def test_native_assign_true_round_trip_no_split_components(): |
There was a problem hiding this comment.
This test catches dropped-key/no-op loads, but its equality asserts hold whether the passthrough branch assigns or copies, so nothing anywhere pins that non-view keys keep true assign semantics which is the point of not running everything through copy_. One extra assert like bridge.state_dict()[key].data_ptr() == sd[key].data_ptr() would cover that.
2 participants
Fixes #1637.
JointQKVAttentionBridge/JointGateUpMLPBridgebuild their split sub-components once, atset_original_component(), viatorch.tensor_spliton the combined weight (c_attn,gate_up_proj). These are true storage-sharing views, confirmed viauntyped_storage().data_ptr()equality.load_state_dict(state_dict, assign=True)broke that view: PyTorch'sassign=Truereplaces each target parameter object rather than copying into existing storage. The bridge's own forward pass stayed correct (LinearBridge.forwardreads the split component directly), butoriginal_model.state_dict()-- whatsave_pretrained()exports -- silently kept stale data for the combined weight, with no error or warning.Fix (per jlarson4's direction on the issue -- option 3 + option 1 from the three listed there): for
assign=Truecalls, detect view-backed targets structurally (tensor.numel() * tensor.element_size() < tensor.untyped_storage().nbytes()-- more reliable thanTensor._base, which nn.Parameter wrapping doesn't preserve here) and route them through an explicit in-place.data.copy_()instead of letting PyTorch replace the parameter object. If a shape mismatch makes that copy unsafe, raise a clearRuntimeErrorrather than failing opaquely or silently corrupting data.assign=False(the default) was already correct and is untouched.Testing:
_is_view_backedunit test (synthetictorch.tensor_split, no model download)boot_nativetest confirming ordinary (non-split)assign=Trueround-trips unchangedoriginal_model.state_dict()is no longer stale afterassign=Truetest_state_dict_round_trip.pysuite ([Bug Report] Native bridge state_dict()/load_state_dict() are not inverses #1587),generalized_componentssplit-bridge tests, andgpt2/gpt_bigcode/falcon/openai_gptadapter tests all still pass (198 tests, no regressions)uv run mypy .cleanAcceptance criteria from the issue are all met.