fix(bridge): make native TransformerBridge state_dict()/load_state_dict() true inverses

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 == []

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.