diff --git a/tests/unit/model_bridge/test_state_dict_round_trip.py b/tests/unit/model_bridge/test_state_dict_round_trip.py new file mode 100644 index 000000000..b0d707d89 --- /dev/null +++ b/tests/unit/model_bridge/test_state_dict_round_trip.py @@ -0,0 +1,196 @@ +"""Regression tests for TransformerBridge.state_dict()/load_state_dict() round-tripping (#1587). + +state_dict() emits TL-renamed keys (e.g. "blocks.0.attn.q.weight"), but +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. +""" +from __future__ import annotations + +import pytest +import torch + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge import TransformerBridge + + +def _native_cfg(**overrides) -> TransformerBridgeConfig: + base = dict( + d_model=32, + d_head=16, + n_heads=2, + n_layers=2, + n_ctx=8, + d_vocab=16, + d_mlp=64, + act_fn="gelu", + normalization_type="LN", + seed=0, + ) + base.update(overrides) + return TransformerBridgeConfig(**base) + + +def test_native_round_trip_overwrites_params_not_a_noop(): + bridge = TransformerBridge.boot_native(_native_cfg()) + + sd = {k: v.clone() for k, v in bridge.state_dict().items()} + assert sd, "state_dict() returned no TL-format keys" + + with torch.no_grad(): + for p in bridge.parameters(): + p.zero_() + assert all((p == 0).all() for p in bridge.parameters()) + + bridge.load_state_dict(sd, strict=True) + + # Compare against the snapshot directly rather than asserting "not all + # zero" - LayerNorm bias legitimately initializes to all-zero, so that + # check would pass even for a param that never got reloaded. + reloaded = bridge.state_dict() + for key, value in sd.items(): + assert torch.equal(reloaded[key], value), f"{key} did not round-trip" + + +def test_native_strict_true_raises_on_missing_key(): + bridge = TransformerBridge.boot_native(_native_cfg()) + sd = bridge.state_dict() + incomplete = dict(sd) + incomplete.pop(next(iter(incomplete))) + + with pytest.raises(RuntimeError, match="Missing key"): + bridge.load_state_dict(incomplete, strict=True) + + +def test_native_strict_true_raises_on_unexpected_key(): + bridge = TransformerBridge.boot_native(_native_cfg()) + sd = dict(bridge.state_dict()) + sd["totally.bogus.key"] = torch.zeros(1) + + with pytest.raises(RuntimeError, match="Unexpected key"): + bridge.load_state_dict(sd, strict=True) + + +def test_native_strict_false_does_not_raise_on_partial_dict(): + bridge = TransformerBridge.boot_native(_native_cfg()) + sd = bridge.state_dict() + first_key = next(iter(sd)) + partial = {first_key: sd[first_key]} + + result = bridge.load_state_dict(partial, strict=False) + assert result.unexpected_keys == [] + assert len(result.missing_keys) > 0 + + +def test_native_raw_keys_still_load_tracr_style(): + """boot_native's own raw parameter names must keep loading directly, + mirroring tracr's make_tracr_transformer_bridge_state_dict compatibility + contract (utilities/tracr.py).""" + bridge = TransformerBridge.boot_native(_native_cfg()) + raw_sd = {k: v.clone() for k, v in bridge.original_model.state_dict().items()} + + with torch.no_grad(): + for p in bridge.parameters(): + p.zero_() + + bridge.load_state_dict(raw_sd, strict=True) + + reloaded_raw = bridge.original_model.state_dict() + for key, value in raw_sd.items(): + assert torch.equal(reloaded_raw[key], value), f"{key} did not round-trip" + + +def test_native_clean_key_dict_with_partial_aliases_does_not_raise_strict(): + """A complete raw-HF-format-style state dict (clean keys, _original_component + stripped) writes only one alias per shared-storage TL key -- boot_native's own + wrapping produces this aliasing internally (e.g. "layers.0.ln1.weight" is + reachable via two different _original_component paths onto the same + Parameter), not just gpt2's c_attn split. Since aliases of the same tensor + share storage, writing any one of them is sufficient; strict=True must not + report the other, unwritten aliases as missing.""" + bridge = TransformerBridge.boot_native(_native_cfg()) + + raw_sd = bridge.original_model.state_dict() + clean_to_actuals: dict[str, list[str]] = {} + for actual_key in raw_sd: + if actual_key != "_original_component": + clean_to_actuals.setdefault(actual_key.replace("._original_component", ""), []).append( + actual_key + ) + assert any(len(keys) > 1 for keys in clean_to_actuals.values()), ( + "fixture assumption broken: expected boot_native to have some " + "clean key reachable through more than one actual path" + ) + + # One representative actual key's value per clean key, same shape as a + # real raw-HF-format checkpoint (no duplicate paths for the same param). + clean_sd = {clean_key: raw_sd[keys[0]].clone() for clean_key, keys in clean_to_actuals.items()} + + with torch.no_grad(): + for p in bridge.parameters(): + p.zero_() + + result = bridge.load_state_dict(clean_sd, strict=True) + assert result.missing_keys == [] + assert result.unexpected_keys == [] + + reloaded_raw = bridge.original_model.state_dict() + for clean_key, actual_keys in clean_to_actuals.items(): + value = clean_sd[clean_key] + for actual_key in actual_keys: + assert torch.equal( + reloaded_raw[actual_key], value + ), f"{actual_key} (alias of {clean_key}) did not round-trip" + + +@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 + storage-sharing VIEWS into c_attn, not independent parameters - so this is + the case that actually exercises convert_hf_key_to_tl_key's HF-name + renaming, not just identity passthrough like boot_native does.""" + bridge = TransformerBridge.boot_transformers("gpt2", device="cpu") + bridge.eval() + + torch.manual_seed(0) + tokens = torch.randint(0, 1000, (1, 8)) + with torch.no_grad(): + logits_before = bridge(tokens).clone() + + 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) + + with torch.no_grad(): + logits_after = bridge(tokens).clone() + + max_diff = (logits_before - logits_after).abs().max().item() + assert torch.allclose( + logits_before, logits_after, atol=1e-5 + ), f"round trip did not restore forward-pass output: max diff={max_diff:.3e}" + + +@pytest.mark.slow +def test_boot_transformers_clean_key_dict_does_not_raise_strict(): + """Reported review case on real gpt2: a complete raw-HF-format-style state + dict (clean keys) writes only one alias per shared-storage TL key, since + gpt2's split q/k/v are views into c_attn reachable via multiple actual + paths. strict=True previously raised ~337 false "missing key" errors even + though the load fully restores the forward pass.""" + bridge = TransformerBridge.boot_transformers("gpt2", device="cpu") + + raw_sd = bridge.original_model.state_dict() + clean_sd = { + actual_key.replace("._original_component", ""): value.clone() + for actual_key, value in raw_sd.items() + if actual_key != "_original_component" + } + assert len(clean_sd) < len(raw_sd), "fixture assumption broken: expected some aliasing on gpt2" + + result = bridge.load_state_dict(clean_sd, strict=True) + assert result.missing_keys == [] + assert result.unexpected_keys == [] diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index b5c6a677a..9e550c5a8 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -3481,9 +3481,39 @@ def state_dict(self, destination=None, prefix="", keep_vars=False): return tl_state_dict + def _tl_key_to_actual_keys(self) -> dict[str, list[str]]: + """Inverse of the renaming state_dict() applies: map each TL-format key + back to every raw parameter/buffer path that represents it. + + Mirrors the filtering and key-conversion in state_dict() exactly, except + it keeps every raw key for a given TL key instead of only the first-seen + one. Bridge components frequently expose the same underlying parameter + through more than one attribute path (e.g. GPT-2's split q/k/v weights + are views into the wrapped module's combined c_attn weight, reachable + both via a block-level shortcut and via the nested _original_component + chain) - all of those aliases must be written for the round trip to + actually change what forward() reads, not just what state_dict() shows. + """ + mapping: dict[str, list[str]] = {} + for actual_key in self.original_model.state_dict(): + if actual_key == "_original_component" or actual_key.startswith("_original_component."): + continue + clean_key = actual_key.replace("._original_component", "") + if not self._is_valid_bridge_path(clean_key): + continue + hf_key = self._normalize_bridge_key_to_hf(clean_key) + tl_key = self.adapter.convert_hf_key_to_tl_key(hf_key) + mapping.setdefault(tl_key, []).append(actual_key) + return mapping + def load_state_dict(self, state_dict, strict=True, assign=False): """Load state dict into the model, handling both clean keys and original keys with _original_component references. + Accepts three key formats: TL-format keys as emitted by state_dict() + (e.g. "blocks.0.attn.q.weight"), raw native parameter paths (e.g. for + ``boot_native`` / tracr-style loading), and raw paths with + "_original_component" segments stripped. + Args: state_dict: Dictionary containing a whole state of the module strict: Whether to strictly enforce that the keys in state_dict match the keys returned by this module's state_dict() function @@ -3494,27 +3524,58 @@ def load_state_dict(self, state_dict, strict=True, assign=False): """ current_state_dict = self.original_model.state_dict() clean_to_actual = {} - actual_to_clean = {} for actual_key in current_state_dict.keys(): if actual_key != "_original_component": - clean_key = actual_key.replace("._original_component", "") - clean_to_actual[clean_key] = actual_key - actual_to_clean[actual_key] = clean_key + clean_to_actual[actual_key.replace("._original_component", "")] = actual_key + + tl_to_actual = self._tl_key_to_actual_keys() + mapped_state_dict = {} + unexpected_keys = [] for input_key, value in state_dict.items(): if input_key in current_state_dict: mapped_state_dict[input_key] = value - else: - if input_key in clean_to_actual: - actual_key = clean_to_actual[input_key] + elif input_key in clean_to_actual: + mapped_state_dict[clean_to_actual[input_key]] = value + elif input_key in tl_to_actual: + for actual_key in tl_to_actual[input_key]: mapped_state_dict[actual_key] = value - else: - mapped_state_dict[input_key] = value - effective_strict = strict and len(mapped_state_dict) == len(current_state_dict) - return self.original_model.load_state_dict( - mapped_state_dict, strict=effective_strict, assign=assign + else: + unexpected_keys.append(input_key) + + # A TL key's actual-key aliases share the same underlying storage (see + # _tl_key_to_actual_keys), so writing any one of them already updates + # what forward() reads for all of them. Treat the group as satisfied + # if any alias was written -- e.g. a caller supplying clean/raw keys + # (the branch above maps each clean key to exactly one actual key) + # shouldn't have the *other*, unwritten aliases reported as missing. + missing_keys = sorted( + actual_key + for actual_keys in tl_to_actual.values() + if not any(k in mapped_state_dict for k in actual_keys) + for actual_key in actual_keys ) + if strict and (missing_keys or unexpected_keys): + error_msgs = [] + if unexpected_keys: + error_msgs.append( + "Unexpected key(s) in state_dict: " + + ", ".join(f'"{k}"' for k in sorted(unexpected_keys)) + ) + if missing_keys: + error_msgs.append( + "Missing key(s) in state_dict: " + ", ".join(f'"{k}"' for k in missing_keys) + ) + raise RuntimeError( + "Error(s) in loading state_dict for {}:\n\t{}".format( + type(self.original_model).__name__, "\n\t".join(error_msgs) + ) + ) + + result = self.original_model.load_state_dict(mapped_state_dict, strict=False, assign=assign) + return type(result)(missing_keys=missing_keys, unexpected_keys=unexpected_keys) + def get_params(self): """Access to model parameters in the format expected by SVDInterpreter.