Fix recursive TransformerBridge state dict composition

[emerardd]

Description

This is the state-dict contract half of #1655, following the scope agreed in the maintainer discussion.

When a TransformerBridge is nested inside another nn.Module, PyTorch supplies a shared destination and a child prefix while recursively collecting state. The Bridge override instead delegated to the raw Hugging Face model and returned a different mapping, so parent checkpoints did not follow normal nn.Module composition semantics. Switching to standard recursive traversal also exposed two existing asymmetries: joint QKV/gate-up save hooks removed registered child keys without load-side handling, and the first forward registered _ln1_module, changing the module tree after a checkpoint had been saved.

This PR:

• preserves the existing no-argument bridge.state_dict() TransformerLens key contract;
• uses standard nn.Module.state_dict() semantics whenever destination or prefix is supplied;
• adds load-side handling for the intentionally filtered joint QKV and gate-up children;
• wires the attention layer-norm execution reference during construction as a non-owning reference, keeping the registration tree stable across forward;
• adds download-free regression coverage for shared destinations and prefixes, nested strict round-trips, pre/post-forward registration stability, real GPT-2 joint QKV composition, and filtered-child save/load symmetry with value restoration.

Part of #1655. The BERT adapter mappings, container-buffer ownership, and broader architecture traversal sweep remain isolated to the follow-up PR requested by the maintainer.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Validation

• Full local suite: 5266 passed, 55 skipped, 54 deselected, 10 xfailed in 14m21s
• Focused state-dict/native surface: 51 passed
• mypy .: success across 431 source files
• pycln, isort, Black, and git diff --check: clean

Checklist

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[jlarson4]

Thanks for putting this together @emerardd, looks great! Just a small docstring note and a test logistic comment, everything in the actual solution looks great.

[jlarson4]

Thanks for updating this @emerardd! This is approved, I will merge once it passes CI