[Bug Report] TransformerBridge is not composable as an nn.Module: parent traversal misses live parameters

[emerardd]

Describe the bug

TransformerBridge behaves like a standard torch.nn.Module when used directly, but it is not composable as a child of another nn.Module.

TransformerBridge.__init__() stores the live source model with:

self.__dict__["original_model"] = model

This intentionally bypasses nn.Module.__setattr__() and keeps original_model outside the registered _modules tree. Bridge then restores several root-level operations through explicit overrides, including parameters(), named_parameters(), to(), train(), state_dict(), and load_state_dict().

Those overrides make direct calls on the Bridge work, but a parent module does not use them for ordinary recursive traversal:

• parameter discovery walks registered modules and parameters;
• parent.to(...), .half(), and .bfloat16() recurse through _apply();
• parent.requires_grad_(...), optimizers, and zero_grad() use the parent's recursively discovered parameters;
• parent checkpointing passes a shared destination and prefix through recursive state_dict() calls and ignores a child's replacement return object.

The Bridge component tree re-registers only the source modules represented by the adapter mapping. Unmapped parts of the hidden source model are still used by forward, but are invisible to the parent module. As a result, parent optimizers and lifecycle operations can silently omit live parameters, parent dtype conversion can create a mixed-dtype model that crashes, and nested checkpoint round-trips fail.

Code example

This reproduction constructs tiny random models from config only and performs no downloads:

import torch
from transformers import BertConfig, BertForMaskedLM

from transformer_lens.model_bridge.sources import build_bridge_from_module


class Parent(torch.nn.Module):
    def __init__(self, bridge):
        super().__init__()
        self.bridge = bridge

    def forward(self, tokens):
        return self.bridge(tokens)


def make_bridge():
    config = BertConfig(
        vocab_size=32,
        hidden_size=16,
        intermediate_size=32,
        num_hidden_layers=2,
        num_attention_heads=4,
        max_position_embeddings=16,
    )
    return build_bridge_from_module(
        BertForMaskedLM(config).eval(),
        "BertForMaskedLM",
        hf_config=config,
        dtype=torch.float32,
        device="cpu",
        model_name="tiny-parent-module-repro",
    )


bridge = make_bridge()
parent = Parent(bridge)

bridge_ids = {id(parameter) for parameter in bridge.parameters()}
parent_ids = {id(parameter) for parameter in parent.parameters()}
missing_ids = bridge_ids - parent_ids

print(len(bridge_ids), len(parent_ids), len(missing_ids))
print(
    [
        name
        for name, parameter in bridge.original_model.named_parameters()
        if id(parameter) in missing_ids
    ]
)

parent.to(torch.float64)
print({parameter.dtype for parameter in bridge.parameters()})
parent(torch.tensor([[1, 2, 3]]))

On current dev-4.x (ac4f7f134b12), the output is:

42 37 5
[
  'bert.embeddings.token_type_embeddings.weight',
  'bert.embeddings.LayerNorm.weight',
  'bert.embeddings.LayerNorm.bias',
  'cls.predictions.transform.dense.weight',
  'cls.predictions.transform.dense.bias',
]

{torch.float64, torch.float32}
RuntimeError: mixed dtype (CPU): all inputs must share same datatype.

The direct controls work correctly:

bridge.to(torch.float64)
  all 42 parameters are float64
  forward returns float64 logits

bridge.requires_grad_(False)
  no source-model parameter still requires gradients

The corresponding parent operations miss the same five parameters:

parent.to(torch.float64)
  37 parameters are float64
   5 parameters remain float32

parent.requires_grad_(False)
  the same 5 source-model parameters still require gradients

There is also a nested checkpoint symptom. With a tiny GPT-2 Bridge:

destination = {}
returned = bridge.state_dict(destination=destination, prefix="bridge.")

produces:

returned is destination = False
destination first key   = bridge.transformer.wte._original_component.weight
returned first key      = bridge.transformer.wte.weight

When the Bridge is registered under Parent, parent.state_dict() retains the raw recursive entries rather than the returned TL-key dictionary, contains no W_Q key, and parent.load_state_dict(parent.state_dict()) fails under the default strict load.

Expected behavior

A TransformerBridge registered as a child module should participate in standard PyTorch ownership and recursion:

• the parent should discover every unique live parameter used by Bridge forward;
• parent-level dtype/device conversion and freezing should reach all live parameters and buffers;
• optimizers built from parent.parameters() should not silently omit source-model parameters;
• nested state_dict() should honor the supplied destination and prefix and strictly round-trip through the parent.

Related issues and duplicate boundary

• [Bug Report] TransformerBridge parameters() returns non-leaf tensors, breaking PyTorch optimizer compatibility #1141 / Improve TransformerBridge optimizer compatibility via dual PyTorch/TransformerLens parameter access API #1143 changed direct bridge.parameters() and named_parameters() to delegate to the source model. Parent parameter recursion does not call those child overrides, so this remains unfixed.
• Add TransformerBridge adapter for lightweight decoder-only pretraining models #1519 added direct TransformerBridge.train() propagation because the hidden original_model is outside _modules. It confirms the lifecycle-delegation pattern but does not cover parent composition.
• [Bug Report] Native bridge state_dict()/load_state_dict() are not inverses #1587 / fix(bridge): make native TransformerBridge state_dict()/load_state_dict() true inverses #1598 fixed direct bridge.state_dict() / load_state_dict() inversion and strictness. They do not cover a supplied destination, a non-empty prefix, or a Bridge nested under another module.
• [Bug Report] TransformerBridge.load_state_dict(assign=True) desyncs split QKV/gate-up components from the combined weight they view into #1637 concerns load_state_dict(assign=True) breaking split-parameter view relationships. That is independent of the incomplete parent ownership graph reported here.
• Same-day searches for parent modules, nested nn.Module, optimizer traversal, requires_grad, FSDP, original_model, and state-dict destination/prefix found no issue or PR covering this parent-recursion defect.

Suggested fix direction

Adding more root-level delegation methods is not a complete fix. An _apply() override could repair direct dtype/device movement but would not make parent.parameters() complete; a state_dict() patch alone would not repair parent optimizers or freezing.

The implementation needs one authoritative registered ownership graph for the parameters and buffers used by forward. One possible direction is to register the source model as the owning graph and make Bridge components non-owning views/proxies over it. Naively registering both the complete source model and all current Bridge component wrappers may create duplicate serialization paths, so parameter identity, state-dict keys, hooks, and existing raw/TL checkpoint behavior need to be considered together.

It may also be cleaner to keep standard recursive PyTorch state_dict() semantics separate from an explicitly named TransformerLens-format export API, rather than making one method serve both parent recursion and key conversion.

Suggested regression coverage:

• A parent containing a tiny BERT-family Bridge sees the same unique parameter identities as direct bridge.parameters().
• parent.to(device/dtype), .half(), and .bfloat16() reach every live forward parameter and buffer without mixed-dtype execution.
• parent.requires_grad_(False) and parent.zero_grad() reach all source-model parameters.
• An optimizer created from parent.parameters() can update every intended trainable parameter.
• bridge.state_dict(destination=dest, prefix="bridge.") returns dest and writes the intended prefixed keys into it.
• A parent containing a Bridge can strictly load its own state_dict().
• Parent-level save/mutate/load restores the values actually used by forward.
• Existing direct Bridge checkpoint behavior, TransformerLens-format analysis access, and Tracr raw-key loading remain supported.
• If DDP/FSDP composition is intended to be supported, add a lightweight wrapper smoke test based on standard parameter enumeration; the current report does not claim a full distributed reproduction.

System Info

• Installed from source in the repository uv environment
• Windows 11 / PowerShell
• Python 3.12.10
• TransformerLens dev-4.x commit ac4f7f134b12
• CPU-only reproduction; no model or tokenizer download

Additional context

The visible symptom varies by architecture because adapter mappings re-register different fractions of each source-model tree. BERT is a useful control because it leaves five forward parameters outside the Bridge's registered component graph while direct Bridge delegation still sees all of them.

Checklist

• I have checked that there is no similar issue in the repo (required)

[emerardd]

I'd like to take this on. Could you assign me?

Before implementation, I would appreciate maintainer guidance on the intended ownership and checkpoint contract. My current inclination is to establish one authoritative registered owner for every live Hugging Face parameter and buffer, with Bridge components acting as non-owning hook/analysis views. For checkpointing, standard recursive nn.Module.state_dict(destination=..., prefix=...) would retain PyTorch semantics, while TransformerLens-remapped keys would be exposed through a separate explicit export API.

I am hesitant to implement that direction unilaterally because registering both original_model and the current component wrappers could introduce duplicate module/state-dict paths, while changing direct bridge.state_dict() key behavior could affect existing checkpoint consumers and raw-HF/Tracr compatibility.

Would you prefer the ownership/API cleanup above, even if it requires a narrowly documented checkpoint API adjustment, or a compatibility-first design that preserves the current direct bridge.state_dict() behavior and adds parent-module composition separately?

Once the intended contract is confirmed, I can add download-free regression coverage for parent parameter identity, parent .to() and freezing, optimizer enumeration, and strict nested state-dict round-trips before implementing the fix.

[jlarson4]

Hi @emerardd! Thanks for the thorough breakdown of this issue and for providing possible resolutions that I can review.

After digging through and reproducing the issue locally, I have answers to your design question. First and foremost, this issue breaks down into two independent bugs, which changes how we will want to fix it.

Bug 1: Adapter Coverage

Your five missing BERT parameters correspond exactly to three modules absent from the BERT adapter's component_mapping (transformer_lens/model_bridge/supported_architectures/bert.py:88): token_type_embeddings, embeddings.LayerNorm, cls.predictions.transform.dense. Adding these entries into the mapping closes the gap, the parent sees all 42 parameters, parent.to() converts everything, forward and hooks work. As a control, tiny LLaMA, GPT-2, GPT-NeoX, and Mistral bridges show zero parameters or buffers missed by parent traversal and already properly compose for discovery, dtype/device, freezing, and optimizers.

Bug 2: The state_dict contract

Beyond the destination/prefix violation you documented, I tested the obvious fix (standard nn.Module.state_dict semantics over the registered tree) and strict nested round-trip still fails, for two additional reasons:

1. JointQKVAttentionBridge/JointGateUpMLPBridge register save-time hooks that strip combined qkv/gate-up keys with no load-side counterpart (joint_qkv_attention.py:102, joint_gate_up_mlp.py:62)
2. BlockBridge.forward lazily registers attn._ln1_module (block.py:195). One forward pass adds state-dict keys and a checkpoint saved pre-forward won't strict-load post-forward. Any fix needs save/load symmetry and a registration tree that is stable across forward.

As to how we actually want to fix this, lets focus on compatibility first as well as break this into two PRs, one for each bug.

1. State-dict contract PR – honor standard recursive semantics whenever destination/prefix are supplied. Keep direct no-arg bridge.state_dict() returning TL-format keys for now to preserve the fix(bridge): make native TransformerBridge state_dict()/load_state_dict() true inverses #1598 contract. Add symmetric load-side handling for the qkv/gate-up filter hooks, and move _ln1_module wiring to construction time.
2. Adapter-coverage PR – add the three BERT component mapping entries (named to match HookedEncoder conventions), container-buffer ownership, and a parametrized download-free parity test asserting parent-wrapped and direct traversal see identical parameter and buffer identity sets. Scope it to 10–15 structural archetypes (joint-QKV, split-QKV/RoPE, GQA, encoder+task heads including the NSP variant, encoder-decoder, MoE, vision, audio). Also do a one-time best-effort sweep across the registered architectures to enumerate other gapped adapters for this PR.

Your proposed regression list is good, please add post-forward registration stability and save/load hook symmetry to it. You are now assigned, I look forward to reviewing your PRs!