[Proposal] One-time converter for TL-convention checkpoints (blocks.N.attn.W_Q, embed.W_E) to bridge format

[jlarson4]

Proposal

Base on dev-4.x. Add a standalone conversion utility, some thing like convert_tl_checkpoint(state_dict, cfg) -> dict in a flat transformer_lens/utilities/ module (sibling of utilities/tracr.py). It maps a TL-property-format state dict (blocks.N.attn.W_Q, embed.W_E, ln_final.w, unembed.W_U) to the key/tensor format TransformerBridge.load_state_dict accepts natively. Deliberately not an extension of load_state_dict itself: the legacy key format should be convertible once and re-saved, not kept alive as a permanent bridge load path.

Motivation

Deprecating HookedTransformer (removal at 4.0) strands published toy checkpoints like OthelloGPT, grokking, and ARENA stored in TL convention. These are historical data artifacts (finished training runs that can't be re-run), so a data on-ramp is warranted, but it should be a one-time converter, keeping load_state_dict native-only. Users convert, re-save in bridge format, and move on.

Pitch

• Invert TL property keys to native/HF module keys: mirror translate_transformer_lens_path (architecture_adapter.py:616) for the TL→HF direction, and invert convert_hf_key_to_tl_key (architecture_adapter.py:880) for HF→TL (reuse the round-trip inverse-key machinery).
• Reverse the read-only attention reshape so per-head [n_heads, d_head, ...] tensors land in flat nn.Linear q/k/v/o slots.
• Wire end-to-end so a native bridge consumes a converted OthelloGPT-style TL checkpoint via the standard load_state_dict (demos/Othello_GPT.ipynb is the acceptance vehicle).
• Document the two-step in the migration guide: convert once → bridge.load_state_dict(converted) → re-save with bridge.state_dict().

Acceptance:

• Converter maps a TL-format state dict; native bridge loads the result via the unchanged load_state_dict and matches source forward outputs
• W_Q/W_K/W_V/W_O land in correct native q/k/v/o slots
• load_state_dict (transformer_bridge.py:3478-3510) itself is unchanged — no TL-property key handling added to it
• New test under tests/unit/model_bridge/ round-trips convert→load→parity
• make unit-test and uv run mypy . clean

Mirror: the weight-space reverse reshape already exists as ProcessWeights.convert_tensor_to_hf_format (transformer_lens/weight_processing.py:2003, via ParamProcessingConversion.revert, transformer_lens/conversion_utils/param_processing_conversion.py:75, class at :13) reshapes TL to HF and is invoked throughout process_weights. Reuse it (and _reshape_qkv/_reshape_o at transformer_bridge.py:937-949 for reference).

Note: _setup_qkv_hook_reshaping (attention.py:265-360) is the activation-space hook reshape, not the weight reshape, so it is not the thing to invert. Also note bridge.state_dict() emits TL component paths with torch parameter names (blocks.N.attn.q.weight), not TL property names (W_Q), so this converter handles an external legacy-convention dict and is distinct from the export round-trip (issue #1587).

Additional context

Must wait for #1587 to land so the converted output has a stable target format. Project mirrors HT into TransformerBridge.

Checklist

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

[jlarson4]

Explicit Restatement:

This issue is contingent on the completion of #1587, do not attempt to complete until 1587 is at least available via a PR branch, so you can base your work on that. If you are attempting to work on this but did not complete 1587, coordinate with that author to make sure both fixes land accurately

[LightWork666]

Hey @jlarson4, I'd like to take this on dev-4.x. I just put up #1598 for #1587 (there's also a competing #1595 open for the same issue), so I have the inverse-key-mapping machinery fresh in mind for this converter. Will base the work on whichever of those lands, and coordinate if needed.

[jlarson4]

@LightWork666 Thanks for taking this on! I am currently reviewing the three competing implementations of #1587 (#1591 is also a parallel fix), and will get back to you on this PR once I've determined how to address those three contributions.