fix(model-bridge): restore native state dict round trips

tests/unit/model_bridge/test_boot_native.py

@@ -3,6 +3,7 @@
 
 import sys
 
+import pytest
 import torch
 
 from transformer_lens.config import TransformerBridgeConfig
@@ -71,6 +72,42 @@ def test_boot_native_returns_bridge_over_native_model():
     assert isinstance(bridge.original_model, NativeModel)
 
 
+def test_state_dict_round_trip_restores_native_bridge():
+    bridge = TransformerBridge.boot_native(_cfg())
+    expected = {key: value.clone() for key, value in bridge.state_dict().items()}
+
+    with torch.no_grad():
+        for parameter in bridge.parameters():
+            parameter.zero_()
+
+    incompatible_keys = bridge.load_state_dict(expected, strict=True)
+
+    assert incompatible_keys.missing_keys == []
+    assert incompatible_keys.unexpected_keys == []
+    for key, expected_value in expected.items():
+        assert torch.equal(bridge.state_dict()[key], expected_value), key
+
+
+def test_state_dict_strict_load_rejects_missing_tl_key():
+    bridge = TransformerBridge.boot_native(_cfg())
+    state_dict = bridge.state_dict()
+    state_dict.pop("embed.weight")
+
+    with pytest.raises(RuntimeError, match=r"Missing key\(s\)"):
+        bridge.load_state_dict(state_dict, strict=True)
+
+
+def test_load_state_dict_preserves_raw_native_keys():
+    bridge = TransformerBridge.boot_native(_cfg())
+    raw_key = "layers.0.attn.k.weight"
+    actual_key = "layers.0.attn.k._original_component.weight"
+    replacement = torch.full_like(bridge.original_model.state_dict()[actual_key], 0.25)
+
+    bridge.load_state_dict({raw_key: replacement}, strict=False)
+
+    assert torch.equal(bridge.original_model.state_dict()[actual_key], replacement)
+
+
 def test_boot_native_accepts_dict_config():
     cfg_dict = dict(
         d_model=32,

transformer_lens/model_bridge/transformer_bridge.py

@@ -3494,26 +3494,62 @@ 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
+
+        # Build the inverse of the conversion used by state_dict(). A single
+        # TL key can represent several aliased raw keys, all of which must be
+        # populated for strict loading to work on the wrapped module.
+        tl_to_actual: dict[str, list[str]] = {}
+        actual_to_tl: dict[str, str] = {}
+        tensor_to_tl: dict[tuple, str] = {}
+        for actual_key in current_state_dict:
+            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)
+            tl_to_actual.setdefault(tl_key, []).append(actual_key)
+            actual_to_tl[actual_key] = tl_key
+            value = current_state_dict[actual_key]
+            if torch.is_tensor(value):
+                signature = (
+                    value.data_ptr(),
+                    value.dtype,
+                    value.device,
+                    value.shape,
+                    value.stride(),
+                )
+                tensor_to_tl.setdefault(signature, tl_key)
+
+        # Some wrapped modules expose additional raw aliases that are filtered
+        # from state_dict() because they are nested HF implementation details.
+        # Include aliases that point at the same tensor so strict loading still
+        # sees the complete underlying module state.
+        for actual_key, value in current_state_dict.items():
+            if actual_key in actual_to_tl or not torch.is_tensor(value):
+                continue
+            signature = (value.data_ptr(), value.dtype, value.device, value.shape, value.stride())
+            mapped_tl_key = tensor_to_tl.get(signature)
+            if mapped_tl_key is not None:
+                tl_to_actual[mapped_tl_key].append(actual_key)
+
         mapped_state_dict = {}
         for input_key, value in state_dict.items():
-            if input_key in current_state_dict:
+            if input_key in tl_to_actual:
+                for actual_key in tl_to_actual[input_key]:
+                    mapped_state_dict[actual_key] = value
+            elif 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]
                     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
-        )
+        return self.original_model.load_state_dict(mapped_state_dict, strict=strict, assign=assign)
 
     def get_params(self):
         """Access to model parameters in the format expected by SVDInterpreter.