Address TransformerBridge traversal review feedback

Author: emerardd

tests/integration/model_bridge/test_parent_module_traversal.py

@@ -71,6 +71,35 @@ def _vit_config() -> ViTConfig:
     )
 
 
+def _bart_config() -> BartConfig:
+    return BartConfig(
+        vocab_size=32,
+        d_model=16,
+        encoder_layers=1,
+        decoder_layers=1,
+        encoder_attention_heads=4,
+        decoder_attention_heads=4,
+        encoder_ffn_dim=32,
+        decoder_ffn_dim=32,
+        max_position_embeddings=16,
+    )
+
+
+def _hubert_config() -> HubertConfig:
+    return HubertConfig(
+        vocab_size=32,
+        hidden_size=16,
+        num_hidden_layers=1,
+        num_attention_heads=4,
+        intermediate_size=32,
+        conv_dim=(8,),
+        conv_stride=(2,),
+        conv_kernel=(3,),
+        num_conv_pos_embeddings=4,
+        num_conv_pos_embedding_groups=2,
+    )
+
+
 ARCHITECTURE_CASES = (
     ArchitectureCase(
         "gpt2-joint-qkv",
@@ -163,17 +192,7 @@ def _vit_config() -> ViTConfig:
     ArchitectureCase(
         "bart-encoder-decoder",
         BartForConditionalGeneration,
-        lambda: BartConfig(
-            vocab_size=32,
-            d_model=16,
-            encoder_layers=1,
-            decoder_layers=1,
-            encoder_attention_heads=4,
-            decoder_attention_heads=4,
-            encoder_ffn_dim=32,
-            decoder_ffn_dim=32,
-            max_position_embeddings=16,
-        ),
+        _bart_config,
         "BartForConditionalGeneration",
     ),
     ArchitectureCase(
@@ -202,18 +221,7 @@ def _vit_config() -> ViTConfig:
     ArchitectureCase(
         "hubert-audio",
         HubertForCTC,
-        lambda: HubertConfig(
-            vocab_size=32,
-            hidden_size=16,
-            num_hidden_layers=1,
-            num_attention_heads=4,
-            intermediate_size=32,
-            conv_dim=(8,),
-            conv_stride=(2,),
-            conv_kernel=(3,),
-            num_conv_pos_embeddings=4,
-            num_conv_pos_embedding_groups=2,
-        ),
+        _hubert_config,
         "HubertForCTC",
     ),
     ArchitectureCase(
@@ -234,6 +242,8 @@ def _vit_config() -> ViTConfig:
     ),
 )
 
+ARCHITECTURE_CASE_BY_NAME = {case.name: case for case in ARCHITECTURE_CASES}
+
 
 def _named_identities(named_values: Any) -> dict[int, str]:
     return {id(value): name for name, value in named_values}
@@ -274,17 +284,7 @@ def test_parent_and_direct_traversal_have_identical_state(case: ArchitectureCase
 
 
 def test_parent_dtype_conversion_updates_container_owned_state() -> None:
-    bart_config = BartConfig(
-        vocab_size=32,
-        d_model=16,
-        encoder_layers=1,
-        decoder_layers=1,
-        encoder_attention_heads=4,
-        decoder_attention_heads=4,
-        encoder_ffn_dim=32,
-        decoder_ffn_dim=32,
-        max_position_embeddings=16,
-    )
+    bart_config = _bart_config()
     bridge = build_bridge_from_module(
         BartForConditionalGeneration(bart_config),
         "BartForConditionalGeneration",
@@ -305,17 +305,7 @@ def test_parent_dtype_conversion_updates_container_owned_state() -> None:
 
 
 def test_parent_assign_load_updates_container_owned_state() -> None:
-    bart_config = BartConfig(
-        vocab_size=32,
-        d_model=16,
-        encoder_layers=1,
-        decoder_layers=1,
-        encoder_attention_heads=4,
-        decoder_attention_heads=4,
-        encoder_ffn_dim=32,
-        decoder_ffn_dim=32,
-        max_position_embeddings=16,
-    )
+    bart_config = _bart_config()
     bridge = build_bridge_from_module(
         BartForConditionalGeneration(bart_config),
         "BartForConditionalGeneration",
@@ -336,3 +326,72 @@ def test_parent_assign_load_updates_container_owned_state() -> None:
     assert id(bridge.original_model.final_logits_bias) in {
         id(buffer) for buffer in parent.buffers()
     }
+
+
+@pytest.mark.parametrize(
+    ("case_name", "container_path", "state_name", "state_key"),
+    (
+        ("bart-encoder-decoder", "", "final_logits_bias", "final_logits_bias"),
+        ("hubert-audio", "hubert", "masked_spec_embed", "hubert.masked_spec_embed"),
+    ),
+)
+def test_direct_assign_load_stays_current_after_apply(
+    case_name: str, container_path: str, state_name: str, state_key: str
+) -> None:
+    case = ARCHITECTURE_CASE_BY_NAME[case_name]
+    config = case.config_factory()
+    bridge = build_bridge_from_module(
+        case.model_type(config),
+        case.architecture,
+        hf_config=config,
+        dtype=torch.float32,
+        device="cpu",
+        model_name=f"tiny-{case.name}-direct-assign",
+    )
+    original_container = (
+        bridge.original_model.get_submodule(container_path)
+        if container_path
+        else bridge.original_model
+    )
+    owner_container = (
+        bridge._container_state_owners.get_submodule(container_path)
+        if container_path
+        else bridge._container_state_owners
+    )
+    replacement = torch.full_like(getattr(original_container, state_name), 7)
+
+    bridge.load_state_dict({state_key: replacement}, strict=False, assign=True)
+
+    assert getattr(owner_container, state_name) is getattr(original_container, state_name)
+    bridge.cpu()
+    assert torch.equal(getattr(original_container, state_name), replacement)
+
+
+@pytest.mark.parametrize(
+    ("case_name", "key_fragment", "expected_keys"),
+    (
+        ("bert-nsp", "pooler", {"pooler.weight", "pooler.bias"}),
+        ("vit-bare-pooler", "pooler", {"pooler.weight", "pooler.bias"}),
+        (
+            "ast-audio-classifier",
+            "classifier",
+            {"classifier_ln.weight", "classifier_ln.bias"},
+        ),
+    ),
+)
+def test_task_head_state_dict_keys_are_not_reexpanded(
+    case_name: str, key_fragment: str, expected_keys: set[str]
+) -> None:
+    case = ARCHITECTURE_CASE_BY_NAME[case_name]
+    config = case.config_factory()
+    bridge = build_bridge_from_module(
+        case.model_type(config),
+        case.architecture,
+        hf_config=config,
+        dtype=torch.float32,
+        device="cpu",
+        model_name=f"tiny-{case.name}-state-dict-keys",
+    )
+
+    actual_keys = {key for key in bridge.state_dict() if key_fragment in key}
+    assert actual_keys == expected_keys

tests/unit/model_bridge/supported_architectures/test_vit_adapter.py

@@ -312,7 +312,7 @@ def test_bare_model_maps_pooler_without_root_name_collision(
         self, adapter: ViTArchitectureAdapter
     ) -> None:
         adapter.prepare_model(self._bare_model_with_pooler())
-        assert adapter.component_mapping["vision_pooler"].name == "pooler.dense"
+        assert adapter.component_mapping["pooler"].name == "pooler.dense"
 
     def test_bare_model_does_not_require_encoder_attribute(
         self, adapter: ViTArchitectureAdapter

transformer_lens/model_bridge/component_setup.py

@@ -32,7 +32,15 @@ def _sync_original_container(self) -> None:
         original_container._parameters.update(self._parameters)
         original_container._buffers.update(self._buffers)
 
+    def _refresh_from_original_container(self) -> None:
+        original_container = self.__dict__["_original_container"]
+        for name in self._parameters:
+            self._parameters[name] = original_container._parameters[name]
+        for name in self._buffers:
+            self._buffers[name] = original_container._buffers[name]
+
     def _apply(self, fn: Any, recurse: bool = True) -> "_ContainerStateOwner":
+        self._refresh_from_original_container()
         super()._apply(fn, recurse=recurse)
         self._sync_original_container()
         return self
@@ -59,6 +67,16 @@ def _load_from_state_dict(
         self._sync_original_container()
 
 
+def refresh_container_state_owners(bridge_module: nn.Module) -> None:
+    """Refresh registered container state from the original model tree."""
+    root_owner = bridge_module._modules.get("_container_state_owners")
+    if not isinstance(root_owner, _ContainerStateOwner):
+        return
+    for owner in root_owner.modules():
+        if isinstance(owner, _ContainerStateOwner):
+            owner._refresh_from_original_container()
+
+
 def replace_remote_component(
     replacement_component: nn.Module, remote_path: str, remote_model: RemoteModel
 ) -> None:

transformer_lens/model_bridge/supported_architectures/vit.py

@@ -205,4 +205,4 @@ def prepare_model(self, hf_model: Any) -> None:
             prefix=prefix, with_classifier=with_classifier
         )
         if not with_classifier and getattr(hf_model, "pooler", None) is not None:
-            self.component_mapping["vision_pooler"] = LinearBridge(name="pooler.dense")
+            self.component_mapping["pooler"] = LinearBridge(name="pooler.dense")

transformer_lens/model_bridge/transformer_bridge.py

@@ -37,7 +37,10 @@
 from transformer_lens.hook_points import HookIntrospectionMixin, HookPoint
 from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
 from transformer_lens.model_bridge.bridge_core import BridgeCore
-from transformer_lens.model_bridge.component_setup import set_original_components
+from transformer_lens.model_bridge.component_setup import (
+    refresh_container_state_owners,
+    set_original_components,
+)
 from transformer_lens.model_bridge.composition_scores import CompositionScores
 from transformer_lens.model_bridge.driver_protocol import (
     TensorLike,
@@ -404,6 +407,10 @@ def __getattr__(self, name: str) -> Any:
         # Use __dict__ directly to avoid recursion
         if "_modules" in self.__dict__ and name in self.__dict__["_modules"]:  # type: ignore[arg-type]
             return self.__dict__["_modules"][name]
+        adapter = self.__dict__.get("adapter")
+        component_mapping = getattr(adapter, "component_mapping", None)
+        if component_mapping is not None and name in component_mapping:
+            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
         if "original_model" in self.__dict__ and self.__dict__["original_model"] is not None:
             try:
                 name_split = name.split(".")
@@ -3791,8 +3798,8 @@ def _normalize_bridge_key_to_hf(self, key: str) -> str:
         block_list_names = {"blocks", "L_blocks", "H_blocks", "encoder_blocks", "decoder_blocks"}
         for tl_name, component in component_mapping.items():
             if component.name and tl_name not in block_list_names:
-                # Skip if TL name is already a suffix of the HF path (avoids doubling).
-                if tl_name != component.name and not component.name.endswith("." + tl_name):
+                # Skip if TL name is already a segment of its HF path (avoids doubling).
+                if tl_name != component.name and tl_name not in component.name.split("."):
                     attr_to_hf[tl_name] = component.name
 
         # Map block-level components (ln1, ln2, attn, mlp) for all block lists
@@ -3969,6 +3976,8 @@ 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)
         return type(result)(missing_keys=missing_keys, unexpected_keys=unexpected_keys)
 
     def get_params(self):