fix(bridge): preserve split-component views under load_state_dict(assign=True)

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)

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):