transformer_lens.utilities.tl_checkpoint_conversion module

One-time converter for legacy TL-property-format checkpoints (#1588).

Rotary-model checkpoints are unsupported: their rotary_sin/rotary_cos buffer keys fail loudly as unrecognized keys (drop them first if you need to convert one — the bridge recomputes rotary embeddings from the config).

Historical training runs (OthelloGPT, grokking demos, ARENA content) were saved via HookedTransformer.state_dict() before TransformerBridge existed, using property-style keys (“blocks.0.attn.W_Q”, “embed.W_E”, …) and per-head tensor shapes. convert_tl_checkpoint maps those onto the key/tensor format TransformerBridge.boot_native(cfg).load_state_dict accepts natively, so these checkpoints can be converted once and re-saved in bridge format. This is deliberately a standalone converter rather than a second key convention taught to load_state_dict itself: convert once, bridge.load_state_dict(converted), then re-save with bridge.state_dict().

transformer_lens.utilities.tl_checkpoint_conversion.convert_tl_checkpoint(state_dict: dict[str, Tensor], cfg: TransformerBridgeConfig) → dict[str, Tensor]

Convert a legacy TL-property-format state dict to the key/tensor format TransformerBridge.boot_native(cfg).load_state_dict accepts.

Parameters:

  • state_dict – A state dict in the old HookedTransformer convention (e.g. from HookedTransformer.state_dict()), with keys like "blocks.0.attn.W_Q" and per-head tensor shapes.

  • cfg – The config the checkpoint was trained/saved under. Used both to reshape per-head attention weights and to validate that the checkpoint’s per-head shapes actually match this cfg — a mismatched cfg would otherwise silently mis-group heads without ever tripping a shape error, since d_model == n_heads * d_head holds for any wrong factoring too.

Returns:

A state dict with modern bridge keys (e.g. "blocks.0.attn.q.weight") and flat nn.Linear-oriented tensor shapes, ready for bridge.load_state_dict(converted, strict=True).