diff --git a/docs/source/content/migrating_to_v3.md b/docs/source/content/migrating_to_v3.md index 2fb4c3a78..7bea3601d 100644 --- a/docs/source/content/migrating_to_v3.md +++ b/docs/source/content/migrating_to_v3.md @@ -320,6 +320,30 @@ for position in range(tokens.shape[1]): HuggingFace model. It is not a `TransformerLensKeyValueCache`, and code should not depend on the latter's layout or methods. +### Load a legacy TL-format checkpoint + +Historical training-run checkpoints (OthelloGPT, grokking demos, ARENA +content) were saved via `HookedTransformer.state_dict()` before the bridge +existed, using property-style keys (`blocks.0.attn.W_Q`, `embed.W_E`, ...) +and per-head tensor shapes that `bridge.load_state_dict` doesn't recognize +natively. `convert_tl_checkpoint` is a one-time converter for exactly this: +convert once, load, then re-save in bridge format — `load_state_dict` itself +stays native-only rather than carrying a second, permanent key convention. + +```python +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.utilities.tl_checkpoint_conversion import convert_tl_checkpoint + +cfg = TransformerBridgeConfig(...) # same hyperparameters the checkpoint was trained under +legacy_state_dict = torch.load("othello_gpt.pth") + +bridge = TransformerBridge.boot_native(cfg) +bridge.load_state_dict(convert_tl_checkpoint(legacy_state_dict, cfg), strict=True) + +torch.save(bridge.state_dict(), "othello_gpt_bridge_format.pth") # re-save once, done +``` + ### Type helpers for both model classes Use the structural protocol when a helper should accept either a diff --git a/tests/unit/model_bridge/test_tl_checkpoint_conversion.py b/tests/unit/model_bridge/test_tl_checkpoint_conversion.py new file mode 100644 index 000000000..3cc9d9304 --- /dev/null +++ b/tests/unit/model_bridge/test_tl_checkpoint_conversion.py @@ -0,0 +1,211 @@ +"""Tests for the legacy TL-property-format checkpoint converter (#1588). + +Historical HookedTransformer checkpoints (OthelloGPT, grokking, ARENA content) +are saved with the old 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 loaded once and re-saved in bridge +format without teaching `load_state_dict` a second key convention. +""" +from __future__ import annotations + +import pytest +import torch + +from transformer_lens import HookedTransformer +from transformer_lens.config import HookedTransformerConfig, TransformerBridgeConfig +from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.utilities.tl_checkpoint_conversion import convert_tl_checkpoint + + +def _cfg_kwargs(**overrides): + base = dict( + d_model=32, + d_head=16, + n_heads=2, + n_layers=2, + n_ctx=8, + d_vocab=16, + d_mlp=64, + act_fn="gelu", + normalization_type="LN", + seed=0, + ) + base.update(overrides) + return base + + +def _ht_and_bridge_cfg(**overrides): + kwargs = _cfg_kwargs(**overrides) + return HookedTransformerConfig(**kwargs), TransformerBridgeConfig(**kwargs) + + +def test_convert_tl_checkpoint_loads_strict_into_native_bridge(): + ht_cfg, bridge_cfg = _ht_and_bridge_cfg() + ht = HookedTransformer(ht_cfg) + + converted = convert_tl_checkpoint(ht.state_dict(), bridge_cfg) + + bridge = TransformerBridge.boot_native(bridge_cfg) + result = bridge.load_state_dict(converted, strict=True) + + assert result.missing_keys == [] + assert result.unexpected_keys == [] + + +def test_convert_tl_checkpoint_matches_source_forward_pass(): + ht_cfg, bridge_cfg = _ht_and_bridge_cfg() + ht = HookedTransformer(ht_cfg) + + converted = convert_tl_checkpoint(ht.state_dict(), bridge_cfg) + bridge = TransformerBridge.boot_native(bridge_cfg) + bridge.load_state_dict(converted, strict=True) + + tokens = torch.randint(0, ht_cfg.d_vocab, (1, 4)) + with torch.no_grad(): + ht_logits = ht(tokens) + bridge_logits = bridge(tokens) + + torch.testing.assert_close(bridge_logits, ht_logits, atol=1e-4, rtol=1e-4) + + +def test_convert_tl_checkpoint_places_qkvo_in_correct_head_slots(): + """Independent check that per-head Q/K/V/O land in the right slots: read + the converted+loaded bridge back out through its own W_Q/W_K/W_V/W_O + properties (implemented separately from the converter) and compare + directly against the source HookedTransformer's per-head weights, rather + than trusting the converter's own reshape math.""" + ht_cfg, bridge_cfg = _ht_and_bridge_cfg() + ht = HookedTransformer(ht_cfg) + + converted = convert_tl_checkpoint(ht.state_dict(), bridge_cfg) + bridge = TransformerBridge.boot_native(bridge_cfg) + bridge.load_state_dict(converted, strict=True) + + torch.testing.assert_close(bridge.W_Q, ht.W_Q) + torch.testing.assert_close(bridge.W_K, ht.W_K) + torch.testing.assert_close(bridge.W_V, ht.W_V) + torch.testing.assert_close(bridge.W_O, ht.W_O) + torch.testing.assert_close(bridge.b_Q, ht.b_Q) + torch.testing.assert_close(bridge.b_K, ht.b_K) + torch.testing.assert_close(bridge.b_V, ht.b_V) + torch.testing.assert_close(bridge.b_O, ht.b_O) + + +def test_convert_tl_checkpoint_raises_on_cfg_mismatch(): + """A wrong cfg can't be caught by a downstream shape-mismatch error -- + merging per-head dims produces a validly-shaped result for any head + count, since d_model == n_heads * d_head for any factoring of it. The + converter must catch this itself.""" + ht_cfg, _ = _ht_and_bridge_cfg() + ht = HookedTransformer(ht_cfg) + + wrong_cfg = TransformerBridgeConfig( + **_cfg_kwargs(n_heads=4, d_head=8) + ) # same d_model, wrong split + + with pytest.raises(ValueError, match="attn.W_Q"): + convert_tl_checkpoint(ht.state_dict(), wrong_cfg) + + +def test_convert_tl_checkpoint_raises_on_unrecognized_key(): + _, bridge_cfg = _ht_and_bridge_cfg() + with pytest.raises(ValueError, match="not a recognized"): + convert_tl_checkpoint({"blocks.0.attn.totally_unknown_param": torch.zeros(1)}, bridge_cfg) + + +def test_convert_tl_checkpoint_supports_gqa(): + ht_cfg, bridge_cfg = _ht_and_bridge_cfg(n_heads=4, d_head=8, n_key_value_heads=2) + ht = HookedTransformer(ht_cfg) + + converted = convert_tl_checkpoint(ht.state_dict(), bridge_cfg) + bridge = TransformerBridge.boot_native(bridge_cfg) + result = bridge.load_state_dict(converted, strict=True) + + assert result.missing_keys == [] + assert result.unexpected_keys == [] + + tokens = torch.randint(0, ht_cfg.d_vocab, (1, 4)) + with torch.no_grad(): + ht_logits = ht(tokens) + bridge_logits = bridge(tokens) + torch.testing.assert_close(bridge_logits, ht_logits, atol=1e-4, rtol=1e-4) + + +def test_convert_tl_checkpoint_supports_lnpre(): + """OthelloGPT (this converter's motivating use case, #1588) uses + normalization_type="LNPre" -- param-free pre-norm, so ln1/ln2/ln_final + have no weight/bias keys at all in the state dict for this converter to + handle; this just confirms the round trip still works end to end.""" + ht_cfg, bridge_cfg = _ht_and_bridge_cfg(normalization_type="LNPre") + ht = HookedTransformer(ht_cfg) + + converted = convert_tl_checkpoint(ht.state_dict(), bridge_cfg) + bridge = TransformerBridge.boot_native(bridge_cfg) + result = bridge.load_state_dict(converted, strict=True) + + assert result.missing_keys == [] + assert result.unexpected_keys == [] + + tokens = torch.randint(0, ht_cfg.d_vocab, (1, 4)) + with torch.no_grad(): + ht_logits = ht(tokens) + bridge_logits = bridge(tokens) + torch.testing.assert_close(bridge_logits, ht_logits, atol=1e-4, rtol=1e-4) + + +def test_convert_tl_checkpoint_supports_attn_only(): + ht_cfg, bridge_cfg = _ht_and_bridge_cfg(attn_only=True) + ht = HookedTransformer(ht_cfg) + + converted = convert_tl_checkpoint(ht.state_dict(), bridge_cfg) + bridge = TransformerBridge.boot_native(bridge_cfg) + result = bridge.load_state_dict(converted, strict=True) + + assert result.missing_keys == [] + assert result.unexpected_keys == [] + + tokens = torch.randint(0, ht_cfg.d_vocab, (1, 4)) + with torch.no_grad(): + ht_logits = ht(tokens) + bridge_logits = bridge(tokens) + torch.testing.assert_close(bridge_logits, ht_logits, atol=1e-4, rtol=1e-4) + + +def test_convert_tl_checkpoint_drops_attention_buffers(): + ht_cfg, bridge_cfg = _ht_and_bridge_cfg() + ht = HookedTransformer(ht_cfg) + + converted = convert_tl_checkpoint(ht.state_dict(), bridge_cfg) + + assert not any(key.endswith((".mask", ".IGNORE")) for key in converted) + + +def test_convert_tl_checkpoint_supports_gated_mlp_and_rms_norm(): + """The native bridge's gated MLP has no bias parameter at all for + gate/in/out (matching how real gated-MLP HF architectures like Llama are + built) while HookedTransformer's gated MLP keeps live b_in/b_out + parameters (pre-existing mismatch between the two implementations, + unrelated to this converter). convert_tl_checkpoint still faithfully + translates those keys since they're real HT parameters; load_state_dict + is the one that should refuse them under strict=True. Here they're + exactly zero (freshly constructed, untrained model) so dropping them via + strict=False is lossless and the forward pass still matches exactly. + """ + ht_cfg, bridge_cfg = _ht_and_bridge_cfg(gated_mlp=True, normalization_type="RMS", act_fn="silu") + ht = HookedTransformer(ht_cfg) + + converted = convert_tl_checkpoint(ht.state_dict(), bridge_cfg) + bridge = TransformerBridge.boot_native(bridge_cfg) + result = bridge.load_state_dict(converted, strict=False) + + assert result.missing_keys == [] + assert set(result.unexpected_keys) == { + f"blocks.{i}.mlp.{part}.bias" for i in range(ht_cfg.n_layers) for part in ("in", "out") + } + + tokens = torch.randint(0, ht_cfg.d_vocab, (1, 4)) + with torch.no_grad(): + ht_logits = ht(tokens) + bridge_logits = bridge(tokens) + torch.testing.assert_close(bridge_logits, ht_logits, atol=1e-4, rtol=1e-4) diff --git a/transformer_lens/utilities/tl_checkpoint_conversion.py b/transformer_lens/utilities/tl_checkpoint_conversion.py new file mode 100644 index 000000000..81a6d2f37 --- /dev/null +++ b/transformer_lens/utilities/tl_checkpoint_conversion.py @@ -0,0 +1,159 @@ +"""One-time converter for legacy TL-property-format checkpoints (#1588). + +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()``. +""" + +from __future__ import annotations + +from typing import Callable, Optional + +import einops +import torch + +from transformer_lens.config.transformer_bridge_config import TransformerBridgeConfig + +# Buffers that live on HookedTransformer's attention blocks but have no +# Parameter counterpart on the bridge side (causal mask, IGNORE sentinel). +_DROPPED_BUFFER_SUFFIXES = (".mask", ".IGNORE") + +TensorConvert = Callable[[torch.Tensor, TransformerBridgeConfig, str], torch.Tensor] + + +def _validate_shape(tensor: torch.Tensor, expected: tuple[int, ...], key: str) -> None: + # Merging/splitting per-head dims (unlike a plain transpose) produces a + # validly-shaped result for *any* head count, since d_model == n_heads * + # d_head for any factoring of it — a wrong cfg silently mis-groups heads + # without ever tripping a downstream shape-mismatch error. Check the + # untouched per-head shape explicitly before reshaping. + if tuple(tensor.shape) != expected: + raise ValueError( + f"convert_tl_checkpoint: {key!r} has shape {tuple(tensor.shape)}, " + f"expected {expected} for the given cfg. The checkpoint may not " + "match this cfg (n_heads/n_key_value_heads/d_head/d_model)." + ) + + +def _kv_heads(cfg: TransformerBridgeConfig) -> int: + return cfg.n_key_value_heads or cfg.n_heads + + +def _convert_w_q(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: + _validate_shape(t, (cfg.n_heads, cfg.d_model, cfg.d_head), key) + return einops.rearrange(t, "n_heads d_model d_head -> (n_heads d_head) d_model") + + +def _convert_w_kv(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: + _validate_shape(t, (_kv_heads(cfg), cfg.d_model, cfg.d_head), key) + return einops.rearrange(t, "n_heads d_model d_head -> (n_heads d_head) d_model") + + +def _convert_w_o(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: + _validate_shape(t, (cfg.n_heads, cfg.d_head, cfg.d_model), key) + return einops.rearrange(t, "n_heads d_head d_model -> d_model (n_heads d_head)") + + +def _convert_b_q(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: + _validate_shape(t, (cfg.n_heads, cfg.d_head), key) + return einops.rearrange(t, "n_heads d_head -> (n_heads d_head)") + + +def _convert_b_kv(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: + _validate_shape(t, (_kv_heads(cfg), cfg.d_head), key) + return einops.rearrange(t, "n_heads d_head -> (n_heads d_head)") + + +def _identity(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: + return t + + +def _transpose(t: torch.Tensor, cfg: TransformerBridgeConfig, key: str) -> torch.Tensor: + return t.T.contiguous() + + +# Old TL-property suffix -> (new bridge-key suffix, tensor conversion). +# Checked in order, longest/most-specific first, so e.g. ".b_Q" is matched +# before the generic ".b" LayerNorm-bias fallback. +_SUFFIX_CONVERSIONS: list[tuple[str, str, TensorConvert]] = [ + (".W_Q", ".q.weight", _convert_w_q), + # GQA stores K/V under a leading-underscore name (the raw Parameter); + # plain ".W_K"/".W_V" become expanding (non-Parameter) properties instead. + ("._W_K", ".k.weight", _convert_w_kv), + ("._W_V", ".v.weight", _convert_w_kv), + (".W_K", ".k.weight", _convert_w_kv), + (".W_V", ".v.weight", _convert_w_kv), + (".W_O", ".o.weight", _convert_w_o), + (".b_Q", ".q.bias", _convert_b_q), + ("._b_K", ".k.bias", _convert_b_kv), + ("._b_V", ".v.bias", _convert_b_kv), + (".b_K", ".k.bias", _convert_b_kv), + (".b_V", ".v.bias", _convert_b_kv), + (".b_O", ".o.bias", _identity), + (".W_in", ".in.weight", _transpose), + (".b_in", ".in.bias", _identity), + (".W_out", ".out.weight", _transpose), + (".b_out", ".out.bias", _identity), + (".W_gate", ".gate.weight", _transpose), + (".b_gate", ".gate.bias", _identity), + (".W_U", ".weight", _transpose), + (".b_U", ".bias", _identity), + (".W_E", ".weight", _identity), + (".W_pos", ".weight", _identity), + (".w", ".weight", _identity), + (".b", ".bias", _identity), +] + + +def _convert_key_and_tensor( + key: str, tensor: torch.Tensor, cfg: TransformerBridgeConfig +) -> Optional[tuple[str, torch.Tensor]]: + for old_suffix, new_suffix, convert in _SUFFIX_CONVERSIONS: + if key.endswith(old_suffix): + new_key = key[: -len(old_suffix)] + new_suffix + return new_key, convert(tensor, cfg, key) + return None + + +def convert_tl_checkpoint( + state_dict: dict[str, torch.Tensor], + cfg: TransformerBridgeConfig, +) -> dict[str, torch.Tensor]: + """Convert a legacy TL-property-format state dict to the key/tensor + format ``TransformerBridge.boot_native(cfg).load_state_dict`` accepts. + + Args: + 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)``. + """ + converted: dict[str, torch.Tensor] = {} + for key, tensor in state_dict.items(): + if key.endswith(_DROPPED_BUFFER_SUFFIXES): + continue + result = _convert_key_and_tensor(key, tensor, cfg) + if result is None: + raise ValueError( + f"convert_tl_checkpoint: don't know how to convert key {key!r} " + "(not a recognized TL-property parameter or buffer suffix)." + ) + new_key, new_tensor = result + converted[new_key] = new_tensor + return converted