[Bug Report] Native bridge state_dict()/load_state_dict() are not inverses

[jlarson4]

Base on dev-4.x

Describe the bug

On a native TransformerBridge, state_dict() emits TL-renamed keys (embed.weight, blocks.0.attn.q.weight, …), but load_state_dict() only matches raw native names, so a state_dict() -> load_state_dict() round-trip silently loads nothing and strict=True is silently downgraded to strict=False. You can reload a checkpoint into a random-init model, get no error, and analyse garbage. This blocks the checkpoint save/reload workflow that HookedTransformer deprecation (removal at 4.0) requires.

Code example

from transformer_lens.model_bridge import TransformerBridge

bridge = TransformerBridge.boot_transformers("gpt2")
import torch
with torch.no_grad():
    for p in bridge.parameters():
        p.zero_()
sd = bridge.state_dict()  # TL-renamed keys, e.g. "embed.weight"
bridge.load_state_dict(sd, strict=True)  # no-op; strict silently downgraded, no error

System Info

Installed from source; any OS; Python 3.10+. (Reproduces on the current dev / bridge.)

Expected behaviour & fix pointer

Round-trip must overwrite params, and strict=True must raise on genuine mismatch — mirroring HookedTransformer's trivially-inverse state_dict/load_state_dict. In transformer_lens/model_bridge/transformer_bridge.py, give load_state_dict (:3478-3510) an inverse of convert_hf_key_to_tl_key (forward at architecture_adapter.py:880) before the clean↔actual mapping, and remove the silent strict downgrade at :3507. The forward direction is the state_dict override at transformer_bridge.py:3424 (TL-key emission at :3470). Compatibility constraint: transformer_lens/utilities/tracr.py make_tracr_transformer_bridge_state_dict (def :99) emits raw native keys (:117-122; documented at :106-108) and relies on raw-name matching, so the inverse-key fix must keep that path loading.

Acceptance:

• load_state_dict(state_dict()) overwrites params (mutate, round-trip, assert equal) — no no-op
• strict=True raises on mismatched state dict
• Existing tracr path (utilities/tracr.py, raw-native names) still loads
• New round-trip test under tests/unit/model_bridge/
• make unit-test and uv run mypy . pass

Checklist

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

[ArS377]

I opened PR #1595 to fix this issue. It adds strict round-trip loading for native bridge state dictionaries and preserves Tracr loading.