diff --git a/tests/unit/model_bridge/test_boot_native.py b/tests/unit/model_bridge/test_boot_native.py index ec998d8a6..835e920b0 100644 --- a/tests/unit/model_bridge/test_boot_native.py +++ b/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, diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py index b5c6a677a..9ba09bd35 100644 --- a/transformer_lens/model_bridge/transformer_bridge.py +++ b/transformer_lens/model_bridge/transformer_bridge.py @@ -3494,15 +3494,54 @@ 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: @@ -3510,10 +3549,7 @@ def load_state_dict(self, state_dict, strict=True, assign=False): 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.