[Bug Report] TransformerBridge.load_state_dict(assign=True) desyncs split QKV/gate-up components from the combined weight they view into

[LightWork666]

Base on dev-4.x

Describe the bug

JointQKVAttentionBridge/JointGateUpMLPBridge build their split q/k/v (or gate/in) sub-bridges once, at set_original_component(), via torch.tensor_split on the combined weight (c_attn, gate_up_proj) — tensor_split returns views, and wrapping one in nn.Parameter preserves the view (confirmed: q.original_component.weight.untyped_storage().data_ptr() == c_attn.weight.untyped_storage().data_ptr() is True right after construction, for both the QKV and gate/up split paths).

bridge.load_state_dict(state_dict, assign=True) breaks that view. PyTorch's assign=True replaces each target parameter with the tensor from state_dict rather than copying into existing storage (assign=False, the default, does an in-place .copy_() and correctly preserves the view — not a bug there). Once a split sub-bridge's parameter is replaced, it's no longer connected to the combined weight's storage at all — the combined weight keeps whatever it had before the call.

The bridge's own forward pass is unaffected, since LinearBridge.forward() reads self.original_component (the split module) directly, never c_attn itself. The break only surfaces if something reads the wrapped model's own state independent of the bridge — most concretely, bridge.original_model.state_dict() (what save_pretrained() would export).

Code example

import torch
from transformer_lens.model_bridge import TransformerBridge

bridge = TransformerBridge.boot_transformers("gpt2", device="cpu")
sd = {k: v.clone() for k, v in bridge.state_dict().items()}
mutated = dict(sd)
mutated["blocks.0.attn.q.weight"] = sd["blocks.0.attn.q.weight"] + 100.0

bridge.load_state_dict(mutated, strict=True, assign=True)

# Bridge itself: correct.
print(torch.equal(bridge.blocks[0].attn.q.original_component.weight, mutated["blocks.0.attn.q.weight"]))
# -> True

# The wrapped model's own state, independent of the bridge: silently stale.
raw_sd = bridge.original_model.state_dict()
c_attn_w = raw_sd["transformer.h.0.attn._original_component.c_attn._original_component.weight"]
d_model = bridge.cfg.d_model
print(torch.allclose(c_attn_w[:, :d_model].T, mutated["blocks.0.attn.q.weight"]))  # updated? -> False
print(torch.allclose(c_attn_w[:, :d_model].T, sd["blocks.0.attn.q.weight"]))       # stale?   -> True

No error, no warning — bridge.original_model.state_dict() (and by extension bridge.original_model.save_pretrained(...)) silently contains pre-assign=True data for every split component's slice of the combined weight, while the bridge itself reports (and computes with) the correct value. Anyone exporting a bridge-loaded checkpoint after an assign=True load gets a corrupted file with no indication anything went wrong.

Confirmed on both split paths that build sub-components this way: JointQKVAttentionBridge (c_attn → q/k/v, real gpt2 checkpoint) and JointGateUpMLPBridge (gate_up_proj → gate/up, synthetic module).

System Info

• Installed from source, dev-4.x
• OS: macOS (arm64), CPU-only
• Reproduces with boot_transformers("gpt2"); underlying mechanism confirmed directly on JointGateUpMLPBridge too

Additional context

Not a regression from #1615/#1587 — assign=False (the default load_state_dict() path, and everything those PRs tested) is unaffected; verified the combined weight stays correctly in sync in that case. This is specifically about the assign=True parameter PyTorch's own nn.Module.load_state_dict exposes and this bridge passes through untouched.

Expected behaviour & fix pointers

Three directions, roughly in order of how invasive they are — flagging rather than picking, since the right one probably depends on how much assign=True (memory-efficient loading, e.g. onto meta tensors) actually matters for split-component architectures specifically:

1. Disallow assign=True when the target has any split (JointQKVAttentionBridge/JointGateUpMLPBridge-style) components — raise a clear error instead of silently desyncing. Simplest, but removes the memory-efficiency benefit of assign=True for these architectures specifically.
2. After original_model.load_state_dict(..., assign=True) returns, re-sync the combined weight's slice from each split component's (possibly-replaced) parameter. Keeps assign=True working, but is extra bookkeeping specific to every split-component type, and needs to run even when the split components weren't in the state dict being loaded (e.g. a partial load under strict=False).
3. Route assign=True state-dict entries for split components through .data.copy_() explicitly, bypassing PyTorch's own replace-semantics for just this subset of keys, so the view relationship survives regardless of the caller's assign choice. Keeps behavior uniform without a post-hoc resync pass, but means assign=True isn't actually honored (still doing a copy under the hood) for these specific keys — worth documenting if so.

Acceptance:

• bridge.load_state_dict(sd, assign=True) does not leave bridge.original_model.state_dict() stale for any split-component key (QKV and gate/up)
• assign=False (default) behavior and existing [Bug Report] Native bridge state_dict()/load_state_dict() are not inverses #1587/feat(bridge): support disk device_map offload targets #1615 tests are unaffected
• New regression test using boot_native or a synthetic split component (no real-model download needed) exercising assign=True specifically
• make unit-test and uv run mypy . pass

Checklist

• I have checked that there is no similar issue in the repo (required)

[Canonik]

are you already working on it? would be glad to take this on

[LightWork666]

Yes, planning to work on it. Already have the repro and a good handle on the mechanism (it's the same view/storage pattern from #1615). Will wait to hear which fix direction @jlarson4 prefers before starting the actual PR, but appreciate the offer!

[jlarson4]

Hi @LightWork666, sorry for the delay in the response! Thanks for documenting it, I have successfully reproduced it on my end.

Let's go with option 3 + option 1. Option 3 is the primary solution, we will do Option 1 as well as additional coverage to prevent any potential issues.

For the parallel implementation of Option 1, we will want to logic to work like so: if a view-backed target can't faithfully emulate assign, raise a clear error rather than silently casting. That's the primary case where it would produce surprising results instead of identical ones.

Let me know if you have any questions, I look forward to your PR!