From b4cea22949cac9e60338f832e9302db918a30006 Mon Sep 17 00:00:00 2001 From: lightyagami6669 Date: Thu, 13 Aug 2026 03:32:55 +0530 Subject: [PATCH] fix(bridge): preserve split-component views under load_state_dict(assign=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 #1637 --- .../test_state_dict_round_trip.py | 77 +++++++++++++++++++ .../model_bridge/transformer_bridge.py | 51 +++++++++++- 2 files changed, 126 insertions(+), 2 deletions(-) diff --git a/tests/unit/model_bridge/test_state_dict_round_trip.py b/tests/unit/model_bridge/test_state_dict_round_trip.py index b0d707d89..e6a965750 100644 --- a/tests/unit/model_bridge/test_state_dict_round_trip.py +++ b/tests/unit/model_bridge/test_state_dict_round_trip.py @@ -4,6 +4,11 @@ load_state_dict() only matched raw native parameter names, so a state_dict() -> load_state_dict() round trip silently loaded nothing and strict=True was silently downgraded to strict=False. + +Also covers "copy-split staleness" (#1637): load_state_dict(..., assign=True) +used to replace view-backed split-component parameters (e.g. gpt2's q/k/v, +which are torch.tensor_split views into the combined c_attn weight) wholesale, +desyncing them from the combined weight they share storage with. """ from __future__ import annotations @@ -12,6 +17,7 @@ from transformer_lens.config import TransformerBridgeConfig from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.model_bridge.transformer_bridge import _is_view_backed def _native_cfg(**overrides) -> TransformerBridgeConfig: @@ -143,6 +149,38 @@ def test_native_clean_key_dict_with_partial_aliases_does_not_raise_strict(): ), f"{actual_key} (alias of {clean_key}) did not round-trip" +def test_is_view_backed_detects_tensor_split_views(): + """Core detection helper for #1637: a torch.tensor_split view shares storage + with a larger source tensor, even after nn.Parameter wrapping (which does not + reliably preserve Tensor._base view tracking, so this must not rely on it).""" + combined = torch.nn.Parameter(torch.randn(12, 4)) + a, b, c = torch.tensor_split(combined, 3, dim=0) + split_param = torch.nn.Parameter(a) + assert split_param.untyped_storage().data_ptr() == combined.untyped_storage().data_ptr() + assert _is_view_backed(split_param) + + assert not _is_view_backed(combined) + assert not _is_view_backed(torch.nn.Parameter(torch.randn(4, 4))) + + +def test_native_assign_true_round_trip_no_split_components(): + """boot_native's components are independent parameters (no split/view + components), so assign=True should take the ordinary passthrough path and + round-trip exactly like assign=False does.""" + bridge = TransformerBridge.boot_native(_native_cfg()) + + sd = {k: v.clone() for k, v in bridge.state_dict().items()} + with torch.no_grad(): + for p in bridge.parameters(): + p.zero_() + + bridge.load_state_dict(sd, strict=True, assign=True) + + reloaded = bridge.state_dict() + for key, value in sd.items(): + assert torch.equal(reloaded[key], value), f"{key} did not round-trip under assign=True" + + @pytest.mark.slow def test_boot_transformers_round_trip_matches_forward_pass(): """GPT-2's Conv1D-combined attention makes the bridge's q/k/v components @@ -194,3 +232,42 @@ def test_boot_transformers_clean_key_dict_does_not_raise_strict(): result = bridge.load_state_dict(clean_sd, strict=True) assert result.missing_keys == [] assert result.unexpected_keys == [] + + +@pytest.mark.slow +def test_boot_transformers_assign_true_does_not_leave_combined_weight_stale(): + """#1637 repro: assign=True on a split QKV component used to replace the + parameter object instead of copying into it, breaking the view relationship + with c_attn -- the bridge itself read the new value, but + original_model.state_dict() (what save_pretrained() exports) silently kept + the pre-load data for the combined weight.""" + 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) + + assert torch.equal( + bridge.blocks[0].attn.q.original_component.weight, mutated["blocks.0.attn.q.weight"] + ) + + 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 + assert torch.allclose(c_attn_w[:, :d_model].T, mutated["blocks.0.attn.q.weight"]) + + +@pytest.mark.slow +def test_boot_transformers_assign_true_shape_mismatch_raises_clear_error(): + """A view-backed split component can only be loaded under assign=True via an + in-place copy, which requires a matching shape -- fail loudly instead of a + confusing error surfacing from deep inside copy_, or silently corrupting data.""" + bridge = TransformerBridge.boot_transformers("gpt2", device="cpu") + + sd = dict(bridge.state_dict()) + sd["blocks.0.attn.q.weight"] = sd["blocks.0.attn.q.weight"][:-1] + + with pytest.raises(RuntimeError, match="view sharing storage"): + bridge.load_state_dict(sd, strict=True, assign=True) diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index 331bde2e2..db918bdc9 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -72,6 +72,18 @@ def _resolve_attr_path(obj: nn.Module, attr_path: str) -> Optional[torch.Tensor] return cast(torch.Tensor, result) +def _is_view_backed(tensor: torch.Tensor) -> bool: + """Whether ``tensor`` only occupies part of the storage it points to. + + True for e.g. a split QKV/gate-up component's parameter, which is a + ``torch.tensor_split`` view into a larger combined weight (``c_attn``, + ``gate_up_proj``) even after being wrapped in ``nn.Parameter`` -- that + wrapping doesn't reliably preserve PyTorch's own ``Tensor._base`` view + tracking, so storage-size comparison is used instead of ``_base``. + """ + return tensor.numel() * tensor.element_size() < tensor.untyped_storage().nbytes() + + class TransformerBridge(BridgeCore, HookIntrospectionMixin, nn.Module): """Torch-backed bridge: HF, vLLM-via-torch, anything that wraps an ``nn.Module``. @@ -3913,7 +3925,7 @@ def load_state_dict(self, state_dict, strict=True, assign=False): Returns: NamedTuple with missing_keys and unexpected_keys fields """ - current_state_dict = self.original_model.state_dict() + current_state_dict = self.original_model.state_dict(keep_vars=True) clean_to_actual = {} for actual_key in current_state_dict.keys(): if actual_key != "_original_component": @@ -3964,7 +3976,42 @@ def load_state_dict(self, state_dict, strict=True, assign=False): ) ) - result = self.original_model.load_state_dict(mapped_state_dict, strict=False, assign=assign) + if not assign: + result = self.original_model.load_state_dict( + mapped_state_dict, strict=False, assign=False + ) + return type(result)(missing_keys=missing_keys, unexpected_keys=unexpected_keys) + + # assign=True normally makes nn.Module.load_state_dict *replace* each + # target parameter/buffer with the incoming tensor rather than copying + # into existing storage. For a view-backed target (e.g. a split QKV/ + # gate-up component sharing storage with a combined weight like + # c_attn), replacing the object desyncs it from whatever it's a view + # of -- the bridge itself would keep reading the new, correct value + # (LinearBridge.forward reads the split component directly), but the + # combined weight -- and anything reading original_model's state + # independent of the bridge, e.g. save_pretrained() -- would silently + # keep the stale pre-load data. Route those specific keys through an + # explicit in-place .data.copy_() instead, regardless of the caller's + # assign=True, so the view relationship survives. + 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): + passthrough_items[key] = value + continue + if tuple(target.shape) != tuple(value.shape): + raise RuntimeError( + f"Cannot load {key!r} with assign=True: the current parameter is a " + "view sharing storage with another parameter (e.g. a split QKV/" + "gate-up component's view into a combined weight), so it can only " + "be loaded via an in-place copy, which requires a matching shape. " + f"Got {tuple(value.shape)}, expected {tuple(target.shape)}." + ) + with torch.no_grad(): + target.data.copy_(value) + + result = self.original_model.load_state_dict(passthrough_items, strict=False, assign=True) return type(result)(missing_keys=missing_keys, unexpected_keys=unexpected_keys) def get_params(self):