Fix recursive TransformerBridge state dict composition (#1661)

* Fix recursive TransformerBridge state dict composition

* Clarify filtered checkpoint loading

Author: emerardd

tests/unit/model_bridge/test_state_dict_composition.py

@@ -0,0 +1,171 @@
+"""Regression tests for recursive TransformerBridge checkpoint composition (#1655)."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers import GPT2Config, GPT2LMHeadModel
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.model_bridge.generalized_components import (
+    JointGateUpMLPBridge,
+    JointQKVAttentionBridge,
+    LinearBridge,
+)
+from transformer_lens.model_bridge.sources import build_bridge_from_module
+
+
+def _native_bridge() -> TransformerBridge:
+    cfg = TransformerBridgeConfig(
+        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,
+    )
+    return TransformerBridge.boot_native(cfg)
+
+
+def _parent_with_bridge(bridge: TransformerBridge) -> torch.nn.Module:
+    parent = torch.nn.Module()
+    parent.add_module("bridge", bridge)
+    return parent
+
+
+def test_state_dict_with_destination_and_prefix_uses_recursive_semantics() -> None:
+    bridge = _native_bridge()
+    sentinel = torch.tensor(1)
+    destination: OrderedDict[str, torch.Tensor] = OrderedDict({"sentinel": sentinel})
+
+    returned = bridge.state_dict(destination=destination, prefix="nested.bridge.")
+
+    assert returned is destination
+    assert destination["sentinel"] is sentinel
+    recursive_keys = set(destination) - {"sentinel"}
+    assert recursive_keys
+    assert all(key.startswith("nested.bridge.") for key in recursive_keys)
+
+
+def test_parent_state_dict_strict_round_trip() -> None:
+    parent = _parent_with_bridge(_native_bridge())
+    checkpoint = {key: value.clone() for key, value in parent.state_dict().items()}
+
+    with torch.no_grad():
+        for parameter in parent.parameters():
+            parameter.zero_()
+
+    result = parent.load_state_dict(checkpoint, strict=True)
+
+    assert result.missing_keys == []
+    assert result.unexpected_keys == []
+    reloaded = parent.state_dict()
+    for key, value in checkpoint.items():
+        assert torch.equal(reloaded[key], value), f"{key} did not round-trip"
+
+
+def test_parent_registration_is_stable_across_first_forward() -> None:
+    bridge = _native_bridge()
+    parent = _parent_with_bridge(bridge)
+
+    for block in bridge.blocks:
+        assert block.attn._ln1_module is block.ln1.original_component
+
+    keys_before = tuple(parent.state_dict())
+    assert not any("._ln1_module." in key for key in keys_before)
+    with torch.no_grad():
+        bridge(torch.randint(0, bridge.cfg.d_vocab, (1, 4)))
+    keys_after = tuple(parent.state_dict())
+
+    assert keys_after == keys_before
+    assert not any("._ln1_module." in key for key in keys_after)
+
+
+def test_nested_joint_qkv_bridge_strict_round_trip() -> None:
+    cfg = GPT2Config(
+        vocab_size=32,
+        n_positions=16,
+        n_embd=16,
+        n_layer=1,
+        n_head=2,
+        n_inner=32,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+    )
+    bridge = build_bridge_from_module(
+        GPT2LMHeadModel(cfg),
+        architecture="GPT2LMHeadModel",
+        hf_config=cfg,
+    )
+    parent = _parent_with_bridge(bridge)
+
+    checkpoint = {key: value.clone() for key, value in parent.state_dict().items()}
+    assert not any(".qkv." in key for key in checkpoint)
+    with torch.no_grad():
+        for parameter in parent.parameters():
+            parameter.zero_()
+
+    result = parent.load_state_dict(checkpoint, strict=True)
+
+    assert result.missing_keys == []
+    assert result.unexpected_keys == []
+    reloaded = parent.state_dict()
+    for key, value in checkpoint.items():
+        assert torch.equal(reloaded[key], value), f"{key} did not round-trip"
+
+
+def _filtered_joint_component(kind: str) -> torch.nn.Module:
+    filtered_child = LinearBridge(name=kind)
+    filtered_child.set_original_component(torch.nn.Linear(4, 8))
+    cfg = SimpleNamespace(n_heads=2, d_head=4)
+
+    if kind == "qkv":
+        qkv_component = JointQKVAttentionBridge(
+            name="attn",
+            config=cfg,
+            submodules={"qkv": filtered_child},
+        )
+        for child_name in ("q", "k", "v"):
+            getattr(qkv_component, child_name).set_original_component(torch.nn.Linear(4, 4))
+        return qkv_component
+    gate_up_component = JointGateUpMLPBridge(
+        name="mlp",
+        config=cfg,
+        submodules={"gate_up": filtered_child},
+    )
+    gate_up_component.add_module("gate_up", filtered_child)
+    gate_up_component.gate.set_original_component(torch.nn.Linear(4, 4))
+    getattr(gate_up_component, "in").set_original_component(torch.nn.Linear(4, 4))
+    return gate_up_component
+
+
+@pytest.mark.parametrize("filtered_child_name", ["qkv", "gate_up"])
+def test_filtered_joint_component_strict_round_trip(filtered_child_name: str) -> None:
+    component = _filtered_joint_component(filtered_child_name)
+    filtered_child = component.get_submodule(filtered_child_name)
+    checkpoint = {key: value.clone() for key, value in component.state_dict().items()}
+
+    assert checkpoint
+    assert not any(key.startswith(f"{filtered_child_name}.") for key in checkpoint)
+    with torch.no_grad():
+        for parameter in component.parameters():
+            parameter.zero_()
+
+    result = component.load_state_dict(checkpoint, strict=True)
+
+    assert result.missing_keys == []
+    assert result.unexpected_keys == []
+    reloaded = component.state_dict()
+    for key, value in checkpoint.items():
+        assert torch.equal(reloaded[key], value), f"{key} did not round-trip"
+    for parameter in filtered_child.parameters():
+        assert torch.count_nonzero(parameter) == 0

transformer_lens/model_bridge/component_setup.py

@@ -303,6 +303,8 @@ def setup_blocks_bridge(
         block_bridge.name = f"{blocks_template.name}.{i}"
         block_bridge.set_original_component(original_block)
         setup_submodules(block_bridge, architecture_adapter, original_block)
+        if hasattr(block_bridge, "_wire_ln1_module"):
+            block_bridge._wire_ln1_module()
         bridged_blocks.append(block_bridge)
     replace_remote_component(bridged_blocks, blocks_template.name, original_model)
     return bridged_blocks

transformer_lens/model_bridge/generalized_components/block.py

@@ -110,6 +110,28 @@ def __init__(
         # Fires pre-ln2 when use_hook_mlp_in is set. See #1317.
         self.hook_mlp_in = HookPoint()
 
+    def _wire_ln1_module(self) -> None:
+        """Keep the raw ln1 execution reference outside the ownership tree."""
+        from transformer_lens.model_bridge.generalized_components.attention import (
+            AttentionBridge,
+        )
+
+        ln1 = self.submodules.get("ln1") if self.submodules else None
+        attn = self.submodules.get("attn") if self.submodules else None
+        if not isinstance(attn, AttentionBridge):
+            return
+
+        ln1_module = None
+        if (
+            ln1 is not None
+            and getattr(attn, "supports_split_qkv_fork", False)
+            and getattr(ln1, "original_component", None) is not None
+        ):
+            ln1_module = ln1.original_component
+
+        attn._modules.pop("_ln1_module", None)
+        object.__setattr__(attn, "_ln1_module", ln1_module)
+
     def _maybe_wire_pre_ln_capture(self) -> None:
         """Install ln1/ln2 forward_pre_hooks that feed the bridge's pre-LN hooks (#1317).
 
@@ -118,6 +140,7 @@ def _maybe_wire_pre_ln_capture(self) -> None:
         forward never calls the raw module, so a hook there would silently miss
         on most adapters. Idempotent.
         """
+        self._wire_ln1_module()
         if self._pre_ln_capture_wired:
             return
         from transformer_lens.model_bridge.generalized_components.attention import (
@@ -140,7 +163,6 @@ def _capture_pre_ln1(_module: torch.nn.Module, args: tuple) -> None:
 
             handle = ln1.register_forward_pre_hook(_capture_pre_ln1)
             self._pre_ln_capture_handles.append(handle)
-            attn._ln1_module = ln1.original_component
 
         ln2 = self.submodules.get("ln2") if self.submodules else None
         if ln2 is not None and getattr(ln2, "original_component", None) is not None:

transformer_lens/model_bridge/generalized_components/joint_gate_up_mlp.py

@@ -60,6 +60,9 @@ def __init__(
         self._activation_fn: Any = None
 
         self._register_state_dict_hook(JointGateUpMLPBridge._filter_gate_up_state_dict)
+        self.register_load_state_dict_pre_hook(
+            JointGateUpMLPBridge._restore_filtered_gate_up_state_dict
+        )
 
     @staticmethod
     def _filter_gate_up_state_dict(
@@ -74,6 +77,29 @@ def _filter_gate_up_state_dict(
         for k in keys_to_remove:
             del state_dict[k]
 
+    @staticmethod
+    def _restore_filtered_gate_up_state_dict(
+        module: torch.nn.Module,
+        state_dict: Dict[str, Any],
+        prefix: str,
+        local_metadata: Dict[str, Any],
+        strict: bool,
+        missing_keys: list[str],
+        unexpected_keys: list[str],
+        error_msgs: list[str],
+    ) -> None:
+        """Insert current combined weights only to satisfy strict key matching.
+
+        Production checkpoints restore authoritative values through the unfiltered
+        Hugging Face ``_original_component`` path.
+        """
+        del local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+        gate_up = module._modules.get("gate_up")
+        if gate_up is None:
+            return
+        for key, value in gate_up.state_dict(prefix=f"{prefix}gate_up.").items():
+            state_dict.setdefault(key, value)
+
     @staticmethod
     def _default_split_gate_up(
         original_mlp_component: Any,

transformer_lens/model_bridge/generalized_components/joint_qkv_attention.py

@@ -100,6 +100,9 @@ def __init__(
 
         # Exclude stale qkv combined weights from state_dict after splitting.
         self._register_state_dict_hook(JointQKVAttentionBridge._filter_qkv_state_dict)
+        self.register_load_state_dict_pre_hook(
+            JointQKVAttentionBridge._restore_filtered_qkv_state_dict
+        )
 
     def __deepcopy__(self, memo):
         """Share split_qkv_matrix and config across clones instead of copying.
@@ -143,6 +146,29 @@ def _filter_qkv_state_dict(
         for k in keys_to_remove:
             del state_dict[k]
 
+    @staticmethod
+    def _restore_filtered_qkv_state_dict(
+        module: torch.nn.Module,
+        state_dict: Dict[str, Any],
+        prefix: str,
+        local_metadata: Dict[str, Any],
+        strict: bool,
+        missing_keys: list[str],
+        unexpected_keys: list[str],
+        error_msgs: list[str],
+    ) -> None:
+        """Insert current combined weights only to satisfy strict key matching.
+
+        Production checkpoints restore authoritative values through the unfiltered
+        Hugging Face ``_original_component`` path.
+        """
+        del local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+        qkv = module._modules.get("qkv")
+        if qkv is None:
+            return
+        for key, value in qkv.state_dict(prefix=f"{prefix}qkv.").items():
+            state_dict.setdefault(key, value)
+
     def _create_qkv_conversion_rule(self) -> BaseTensorConversion:
         """Create the appropriate conversion rule for the individual q, k, and v matrices.
 

transformer_lens/model_bridge/transformer_bridge.py

@@ -3824,24 +3824,28 @@ def state_dict(self, destination=None, prefix="", keep_vars=False):
         Converts HuggingFace format keys to TransformerLens format and filters out
         _original_component references and nested HuggingFace components.
 
-        This returns a clean state dict with only bridge component paths converted to TL format,
-        excluding nested HF components (like c_fc, c_proj, c_attn) that exist inside
-        original_component modules.
+        A direct no-argument call returns a clean state dict with bridge component
+        paths converted to TL format. Calls that supply ``destination`` or
+        ``prefix`` use standard ``nn.Module`` recursive semantics so a Bridge can
+        compose inside a parent module.
 
         Args:
             destination: Optional dict to store state dict in
             prefix: Optional prefix to add to all keys
             keep_vars: Whether to keep variables as Variables instead of tensors
 
         Returns:
-            Dict containing the state dict with TransformerLens format keys
+            Direct calls return TransformerLens-format keys; recursive calls
+            return the supplied destination with standard module-tree keys.
         """
-        if destination is not None:
-            raw_state_dict = self.original_model.state_dict(
-                destination=destination, prefix=prefix, keep_vars=keep_vars
+        if destination is not None or prefix:
+            return super().state_dict(
+                destination=destination,
+                prefix=prefix,
+                keep_vars=keep_vars,
             )
-        else:
-            raw_state_dict = self.original_model.state_dict(prefix=prefix, keep_vars=keep_vars)
+
+        raw_state_dict = self.original_model.state_dict(keep_vars=keep_vars)
 
         # Clean _original_component references and convert to TL format
         # Also filter out nested HuggingFace components that are wrapped by bridge components