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..e76887e18 100644 --- a/tests/unit/model_bridge/test_state_dict_round_trip.py +++ b/tests/unit/model_bridge/test_state_dict_round_trip.py @@ -4,14 +4,23 @@ 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 parameters that share underlying storage with something else +(a split QKV/gate-up component's view into its combined weight, the combined +weight itself, or a tied pair like embed/unembed, #1725) wholesale, desyncing +them from whatever they share storage with. """ from __future__ import annotations import pytest import torch +from transformers import GPT2Config, GPT2LMHeadModel, Phi3Config, Phi3ForCausalLM from transformer_lens.config import TransformerBridgeConfig from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.model_bridge.sources import build_bridge_from_module +from transformer_lens.model_bridge.transformer_bridge import _storage_group_keys def _native_cfg(**overrides) -> TransformerBridgeConfig: @@ -31,6 +40,49 @@ def _native_cfg(**overrides) -> TransformerBridgeConfig: return TransformerBridgeConfig(**base) +def _tiny_gpt2_bridge() -> TransformerBridge: + """Real, randomly-initialized tiny GPT-2 bridge (JointQKVAttentionBridge path). + + Built from a config, not a download -- fast enough for the default test + tier, unlike the boot_transformers("gpt2") tests below. + """ + cfg = GPT2Config( + vocab_size=32, + n_positions=16, + n_embd=16, + n_layer=1, + n_head=2, + n_inner=32, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + ) + torch.manual_seed(0) + return build_bridge_from_module( + GPT2LMHeadModel(cfg), architecture="GPT2LMHeadModel", hf_config=cfg + ) + + +def _tiny_phi3_bridge() -> TransformerBridge: + """Real, randomly-initialized tiny Phi-3 bridge (JointGateUpMLPBridge path).""" + cfg = Phi3Config( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + max_position_embeddings=16, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + ) + torch.manual_seed(0) + return build_bridge_from_module( + Phi3ForCausalLM(cfg), architecture="Phi3ForCausalLM", hf_config=cfg + ) + + def test_native_round_trip_overwrites_params_not_a_noop(): bridge = TransformerBridge.boot_native(_native_cfg()) @@ -143,6 +195,219 @@ 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_storage_group_keys_detects_shared_storage(): + """Core detection helper for #1637/#1725: any key whose tensor shares + underlying storage with another key's belongs to the same group, whether + it's the smaller view, the larger owner, or a same-size tied duplicate -- + even though nn.Parameter wrapping doesn't reliably preserve Tensor._base, + so this can't rely on that.""" + combined = torch.nn.Parameter(torch.randn(12, 4)) + view = torch.nn.Parameter(torch.tensor_split(combined, 3, dim=0)[0]) + tied_a = torch.nn.Parameter(torch.randn(4, 4)) + tied_b = torch.nn.Parameter(tied_a.data) + independent = torch.nn.Parameter(torch.randn(4, 4)) + + state_dict = { + "combined": combined, + "view": view, + "tied_a": tied_a, + "tied_b": tied_b, + "independent": independent, + } + grouped = _storage_group_keys(state_dict) + assert grouped == {"combined", "view", "tied_a", "tied_b"} + + +def test_storage_group_keys_excludes_meta_tensors(): + """Every meta tensor reports untyped_storage().data_ptr() == 0 (no real + backing memory), so comparing meta tensors by data pointer would + spuriously group every unrelated offloaded parameter in the model + together. Meta-target keys must never be flagged as a storage group.""" + a = torch.nn.Parameter(torch.empty(4, 4, device="meta")) + b = torch.nn.Parameter(torch.empty(4, 4, device="meta")) + assert a.untyped_storage().data_ptr() == b.untyped_storage().data_ptr() == 0 + + grouped = _storage_group_keys({"a": a, "b": b}) + assert grouped == set() + + +def test_native_assign_true_round_trip_no_split_components(): + """boot_native's components are independent parameters (no split/view or + tied components), so assign=True should take the ordinary passthrough + path and round-trip exactly like assign=False does -- and, unlike + assign=False, actually adopt the incoming tensor objects rather than + copying into the existing ones, since that's the point of assign=True.""" + 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" + + first_key = next(iter(sd)) + assert ( + bridge.state_dict()[first_key].data_ptr() == sd[first_key].data_ptr() + ), "non-shared key should be assigned (adopt the incoming tensor object), not copied into" + + +def test_tiny_gpt2_assign_true_does_not_leave_combined_weight_stale(): + """Fast (no-download) CI-covered version of the #1637 repro on the QKV + split path: assign=True on a split component must not desync it from the + combined weight (c_attn) it shares storage with.""" + bridge = _tiny_gpt2_bridge() + + 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"]) + + +def test_tiny_phi3_assign_true_does_not_leave_combined_weight_stale(): + """Fast (no-download) CI-covered version of the #1637 repro on the + gate/up split path (JointGateUpMLPBridge).""" + bridge = _tiny_phi3_bridge() + + sd = {k: v.clone() for k, v in bridge.state_dict().items()} + gate_key = next(k for k in sd if k.endswith("mlp.gate.weight")) + mutated = dict(sd) + mutated[gate_key] = sd[gate_key] + 100.0 + + bridge.load_state_dict(mutated, strict=True, assign=True) + + gate_component = bridge + for part in gate_key.split(".")[:-1]: + gate_component = ( + getattr(gate_component, part) if not part.isdigit() else gate_component[int(part)] + ) + assert torch.equal(gate_component.original_component.weight, mutated[gate_key]) + + raw_sd = bridge.original_model.state_dict() + combined_key = next( + k + for k in raw_sd + if "gate_up_proj" in k and k.endswith("weight") and "0" in k.split(".")[:3] + ) + d_mlp = bridge.cfg.d_mlp + combined_w = raw_sd[combined_key] + assert torch.allclose(combined_w[:d_mlp, :], mutated[gate_key]) + + +def test_tiny_gpt2_assign_true_raw_combined_key_does_not_orphan_views(): + """Reviewer-flagged gap (jlarson4, PR #1660): loading the combined + weight's own *raw* key directly (not through a split alias) under + assign=True must not orphan the q/k/v views that share its storage -- + otherwise the bridge's own forward pass would keep reading stale q/k/v + data while original_model.state_dict() shows the new c_attn value.""" + bridge = _tiny_gpt2_bridge() + + raw_sd = bridge.original_model.state_dict() + c_attn_key = next( + k for k in raw_sd if "h.0.attn" in k and k.endswith("c_attn._original_component.weight") + ) + new_c_attn = raw_sd[c_attn_key].clone() + 100.0 + + bridge.load_state_dict({c_attn_key: new_c_attn}, strict=False, assign=True) + + d_model = bridge.cfg.d_model + expected_q = new_c_attn[:, :d_model].T + assert torch.equal( + bridge.blocks[0].attn.q.original_component.weight, expected_q + ), "q view was orphaned by a direct assign=True write to the combined c_attn key" + + +def test_assign_true_dtype_mismatch_raises_clear_error(): + """Reviewer-flagged gap (jlarson4): only shape was guarded; a dtype + mismatch on a storage-shared target would otherwise silently + copy-convert instead of erroring, e.g. producing a silently mixed-dtype + model from one call.""" + bridge = _tiny_gpt2_bridge() + + sd = dict(bridge.state_dict()) + sd["blocks.0.attn.q.weight"] = sd["blocks.0.attn.q.weight"].to(torch.float64) + + with pytest.raises(RuntimeError, match="dtype"): + bridge.load_state_dict(sd, strict=True, assign=True) + + +def test_assign_true_device_mismatch_raises_clear_error(): + """Reviewer-flagged gap (jlarson4): a device mismatch on a storage-shared + target must also be rejected explicitly, not silently accepted.""" + bridge = _tiny_gpt2_bridge() + + sd = dict(bridge.state_dict()) + sd["blocks.0.attn.q.weight"] = sd["blocks.0.attn.q.weight"].to("meta") + + with pytest.raises(RuntimeError, match="device"): + bridge.load_state_dict(sd, strict=True, assign=True) + + +def test_assign_true_meta_target_materializes_via_passthrough(): + """koriyoshi2041's review finding (PR #1660): copy_() onto a meta target + silently no-ops rather than raising, so a naive view-backed branch could + report a successful load while the parameter stays meta forever. Meta + targets are excluded from the storage-group check (see + test_storage_group_keys_excludes_meta_tensors) and instead take the + ordinary assign=True path, which correctly materializes them -- the + standard way to load real weights onto a meta-initialized model.""" + bridge = _tiny_gpt2_bridge() + + current = bridge.original_model.state_dict(keep_vars=True) + q_key = next(k for k in current if "h.0.attn.q" in k and k.endswith("weight")) + q_param = current[q_key] + assert not q_param.is_meta + + # A plain `.data = ...to("meta")` reassignment is rejected by PyTorch + # (incompatible tensor type), so simulate an offloaded/meta parameter by + # replacing the owning submodule's Parameter object outright, the same + # way accelerate's own offload hooks do it. + owner_path = ".".join(q_key.split(".")[:-1]) + owner = bridge.original_model.get_submodule(owner_path) + owner.weight = torch.nn.Parameter(q_param.data.to("meta"), requires_grad=q_param.requires_grad) + assert bridge.original_model.state_dict(keep_vars=True)[q_key].is_meta + + new_value = torch.randn(q_param.shape) + bridge.load_state_dict({"blocks.0.attn.q.weight": new_value}, strict=False, assign=True) + + reloaded = bridge.original_model.state_dict(keep_vars=True)[q_key] + assert not reloaded.is_meta, "meta target was not materialized by assign=True" + assert torch.equal(reloaded, new_value) + + +def test_assign_true_validates_all_keys_before_copying_any(): + """Reviewer-flagged gap (jlarson4): the loop must validate every + storage-shared key before applying any copy_, not copy-as-it-validates -- + otherwise a failure partway through leaves some keys already written and + others untouched instead of the whole call atomically failing.""" + bridge = _tiny_gpt2_bridge() + + sd = dict(bridge.state_dict()) + original_k = sd["blocks.0.attn.k.weight"].clone() + sd["blocks.0.attn.k.weight"] = sd["blocks.0.attn.k.weight"] + 100.0 # valid + sd["blocks.0.attn.q.weight"] = sd["blocks.0.attn.q.weight"][:-1] # shape mismatch + + with pytest.raises(RuntimeError): + bridge.load_state_dict(sd, strict=True, assign=True) + + assert torch.equal( + bridge.blocks[0].attn.k.original_component.weight, original_k + ), "a valid key was written before the whole call raised on a different key's mismatch" + + @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 +459,69 @@ 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 storage-shared target 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="shares storage"): + bridge.load_state_dict(sd, strict=True, assign=True) + + +@pytest.mark.slow +def test_boot_transformers_assign_true_tied_embed_unembed_stays_in_sync(): + """#1725 repro, fixed as a consequence of the same storage-group + generalization: gpt2 ties embed/unembed weights (separate nn.Parameter + objects sharing storage, confirmed via data_ptr equality). Loading only + embed.weight under assign=True must not leave unembed reading stale + pre-load data -- previously this broke the bridge's own forward pass, + not just the exported checkpoint.""" + bridge = TransformerBridge.boot_transformers("gpt2", device="cpu") + bridge.eval() + tokens = torch.randint(0, 1000, (1, 6)) + with torch.no_grad(): + logits_before = bridge(tokens).clone() + + sd = {k: v.clone() for k, v in bridge.state_dict().items()} + new_embed = sd["embed.weight"] + 3.0 + bridge.load_state_dict({"embed.weight": new_embed}, strict=False, assign=True) + + assert torch.equal(bridge.unembed.original_component.weight, new_embed) + + with torch.no_grad(): + logits_after = bridge(tokens).clone() + assert not torch.allclose( + logits_before, logits_after + ), "loading a new embedding matrix should change the forward pass output" diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index 64c53ce8f..77ac1dd7c 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -81,6 +81,32 @@ def _resolve_attr_path(obj: nn.Module, attr_path: str) -> Optional[torch.Tensor] return cast(torch.Tensor, result) +def _storage_group_keys(state_dict: dict[str, torch.Tensor]) -> set[str]: + """Keys whose current tensor shares underlying storage with another key's. + + Covers every direction of the desync #1637 reported: a partial view onto + a combined weight (e.g. a split QKV/gate-up component's ``torch.tensor_split`` + view into ``c_attn``/``gate_up_proj``), the combined weight itself (writing + it directly would silently orphan the views that share its storage), and a + same-size tied pair (e.g. tied embed/unembed weights, #1725) -- none of + which reliably show up via ``Tensor._base``, since wrapping in + ``nn.Parameter`` doesn't preserve that tracking here. + + Meta tensors are excluded from the comparison: every meta tensor reports + ``untyped_storage().data_ptr() == 0`` (it has no real backing memory), so + comparing meta tensors by data pointer would spuriously group every + unrelated offloaded parameter in the model together. A key whose current + target is meta falls through to ordinary ``assign=True`` handling instead + (the standard way to materialize a meta tensor from a real one). + """ + by_storage: dict[int, list[str]] = {} + for key, tensor in state_dict.items(): + if tensor.is_meta: + continue + by_storage.setdefault(tensor.untyped_storage().data_ptr(), []).append(key) + return {key for keys in by_storage.values() if len(keys) > 1 for key in keys} + + class TransformerBridge(BridgeCore, HookIntrospectionMixin, nn.Module): """Torch-backed bridge: HF, vLLM-via-torch, anything that wraps an ``nn.Module``. @@ -4341,7 +4367,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": @@ -4392,9 +4418,73 @@ 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 assign: - refresh_container_state_owners(self) + 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. Any key whose current tensor shares storage + # with another key -- a split QKV/gate-up component's view into a + # combined weight, the combined weight itself (writing it directly + # would silently orphan the views), or a tied pair like embed/unembed + # -- would desync from whatever it shares storage with: the bridge + # might keep reading a stale value, or a live view not even part of + # this load would silently keep the pre-load data while + # original_model.state_dict() (what save_pretrained() exports) shows + # the new one. Route every key in a storage-sharing group through an + # explicit in-place .data.copy_() instead, regardless of the caller's + # assign=True, so those relationships survive. + shared_keys = _storage_group_keys(current_state_dict) + + copy_items: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + passthrough_items = {} + errors = [] + for key, value in mapped_state_dict.items(): + target = current_state_dict.get(key) + if target is None or key not in shared_keys: + passthrough_items[key] = value + continue + problems = [] + if target.is_meta: + # copy_() onto a meta tensor silently no-ops rather than + # raising, so a meta target would otherwise report a + # successful load while never actually materializing. + problems.append( + "the current parameter is a meta tensor; an in-place copy " + "can't materialize it, and this key shares storage with " + "another parameter so ordinary assign=True replacement " + "isn't safe here either" + ) + else: + if tuple(target.shape) != tuple(value.shape): + problems.append(f"shape {tuple(value.shape)} != expected {tuple(target.shape)}") + if target.dtype != value.dtype: + problems.append(f"dtype {value.dtype} != expected {target.dtype}") + if target.device != value.device: + problems.append(f"device {value.device} != expected {target.device}") + if problems: + errors.append(f'"{key}": ' + "; ".join(problems)) + else: + copy_items[key] = (target, value) + + if errors: + raise RuntimeError( + "Cannot load the following key(s) with assign=True: each shares " + "storage with another parameter/buffer (e.g. a split QKV/" + "gate-up component's view into a combined weight, or a tied " + "pair like embed/unembed), so it can only be loaded via an " + "in-place copy, which requires an exact match.\n\t" + "\n\t".join(errors) + ) + + for target, value in copy_items.values(): + with torch.no_grad(): + target.data.copy_(value) + + result = self.original_model.load_state_dict(passthrough_items, strict=False, assign=True) + refresh_container_state_owners(self) return type(result)(missing_keys=missing_keys, unexpected_keys=unexpected_keys) def get_params(self):