Coverage for transformer_lens/model_bridge/sources/tl_legacy.py: 53%
106 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Loader for legacy TransformerLens-format HF repos.
3NeelNanda/*, ArthurConmy/*, and Baidicoot/* repos predate the HF model format:
4no ``model_type`` in config.json (AutoConfig refuses them), weights stored as
5``*.pth`` state dicts in HookedTransformer property format (or older layouts
6converted below), and training checkpoints as ``checkpoints/*_<label>.pth``
7files rather than revisions. This module derives a TransformerBridgeConfig from
8the repo's TL-style config.json, fetches and normalizes the state dict, and
9loads it into a ``boot_native`` bridge via ``convert_tl_checkpoint``.
11The two legacy layout converters are ports of the HookedTransformer loaders
12(pretrained/weight_conversions/{neel_solu_old,mingpt}.py), rehomed here so this
13path survives the 4.0 deletion of the legacy loading stack.
14"""
16from __future__ import annotations
18import json
19from typing import Any, Optional, Union
21import einops
22import torch
24from transformer_lens.config import TransformerBridgeConfig
25from transformer_lens.model_bridge.bridge import TransformerBridge
26from transformer_lens.utilities.tl_checkpoint_conversion import convert_tl_checkpoint
28TL_LEGACY_PREFIXES = ("NeelNanda/", "ArthurConmy/", "Baidicoot/")
31def _fetch_json(repo_id: str, filename: str) -> dict:
32 from huggingface_hub import hf_hub_download
34 with open(hf_hub_download(repo_id, filename)) as f:
35 return json.load(f)
38def derive_tl_legacy_config(repo_id: str) -> TransformerBridgeConfig:
39 """TransformerBridgeConfig from a legacy TL repo's config.json."""
40 cfg_json = _fetch_json(repo_id, "config.json")
41 architecture = cfg_json.get(
42 "architecture", "neel" if "_old" not in repo_id else "neel-solu-old"
43 )
44 normalization = cfg_json.get("normalization", cfg_json.get("normalization_type"))
45 if cfg_json.get("shortformer_pos", False): 45 ↛ 46line 45 didn't jump to line 46 because the condition on line 45 was never true
46 raise NotImplementedError(
47 f"{repo_id} uses shortformer positional embeddings, which the native "
48 "bridge does not implement."
49 )
50 cfg = TransformerBridgeConfig(
51 d_model=cfg_json["d_model"],
52 n_layers=cfg_json["n_layers"],
53 d_mlp=cfg_json["d_mlp"],
54 d_head=cfg_json["d_head"],
55 n_heads=cfg_json["n_heads"],
56 n_ctx=cfg_json["n_ctx"],
57 d_vocab=cfg_json["d_vocab"],
58 act_fn=cfg_json["act_fn"],
59 attn_only=cfg_json["attn_only"],
60 final_rms=cfg_json.get("final_rms", False),
61 normalization_type=normalization,
62 positional_embedding_type="standard",
63 tokenizer_name=cfg_json.get("tokenizer_name"),
64 architecture="TransformerLensNative",
65 )
66 cfg.original_architecture = architecture # type: ignore[attr-defined]
67 return cfg
70def _convert_neel_solu_old_weights(state_dict: dict, cfg: TransformerBridgeConfig) -> dict:
71 """Old-layout SoLU repos ('*_old'): left-facing weights below 8L, and 8L's
72 W_pos alone left-facing. Port of the HookedTransformer converter."""
73 reverse_pos = cfg.n_layers <= 8
74 reverse_weights = cfg.n_layers <= 6
75 new_state_dict = {}
76 for k, v in state_dict.items():
77 k = k.replace("norm", "ln")
78 if k.startswith("ln."):
79 k = k.replace("ln.", "ln_final.")
80 new_state_dict[k] = v
81 if reverse_pos:
82 new_state_dict["pos_embed.W_pos"] = new_state_dict["pos_embed.W_pos"].T
83 if reverse_weights:
84 for k, v in new_state_dict.items():
85 if "W_" in k and "W_pos" not in k:
86 new_state_dict[k] = v.transpose(-2, -1)
87 return new_state_dict
90def _convert_mingpt_weights(old_state_dict: dict, cfg: TransformerBridgeConfig) -> dict:
91 """minGPT layout (Baidicoot/Othello-GPT): unconcatenated QKV. Port of the
92 HookedTransformer converter."""
93 state_dict = {
94 "embed.W_E": old_state_dict["tok_emb.weight"],
95 "pos_embed.W_pos": old_state_dict["pos_emb"].squeeze(),
96 }
97 for l in range(cfg.n_layers):
98 state_dict[f"blocks.{l}.ln1.w"] = old_state_dict[f"blocks.{l}.ln1.weight"]
99 state_dict[f"blocks.{l}.ln1.b"] = old_state_dict[f"blocks.{l}.ln1.bias"]
100 for name, hf in (("Q", "query"), ("K", "key"), ("V", "value")):
101 w = einops.rearrange(
102 old_state_dict[f"blocks.{l}.attn.{hf}.weight"], "(i h) m->i m h", i=cfg.n_heads
103 )
104 b = einops.rearrange(
105 old_state_dict[f"blocks.{l}.attn.{hf}.bias"], "(i h)->i h", i=cfg.n_heads
106 )
107 state_dict[f"blocks.{l}.attn.W_{name}"] = w
108 state_dict[f"blocks.{l}.attn.b_{name}"] = b
109 state_dict[f"blocks.{l}.attn.W_O"] = einops.rearrange(
110 old_state_dict[f"blocks.{l}.attn.proj.weight"], "m (i h)->i h m", i=cfg.n_heads
111 )
112 state_dict[f"blocks.{l}.attn.b_O"] = old_state_dict[f"blocks.{l}.attn.proj.bias"]
113 state_dict[f"blocks.{l}.ln2.w"] = old_state_dict[f"blocks.{l}.ln2.weight"]
114 state_dict[f"blocks.{l}.ln2.b"] = old_state_dict[f"blocks.{l}.ln2.bias"]
115 state_dict[f"blocks.{l}.mlp.W_in"] = old_state_dict[f"blocks.{l}.mlp.0.weight"].T
116 state_dict[f"blocks.{l}.mlp.b_in"] = old_state_dict[f"blocks.{l}.mlp.0.bias"]
117 state_dict[f"blocks.{l}.mlp.W_out"] = old_state_dict[f"blocks.{l}.mlp.2.weight"].T
118 state_dict[f"blocks.{l}.mlp.b_out"] = old_state_dict[f"blocks.{l}.mlp.2.bias"]
119 state_dict["ln_final.w"] = old_state_dict["ln_f.weight"]
120 state_dict["ln_final.b"] = old_state_dict["ln_f.bias"]
121 state_dict["unembed.W_U"] = old_state_dict["head.weight"].T
122 return state_dict
125def fetch_tl_legacy_state_dict(
126 repo_id: str,
127 cfg: TransformerBridgeConfig,
128 checkpoint_value: Optional[int] = None,
129 dtype: torch.dtype = torch.float32,
130) -> dict:
131 """Download and normalize a legacy repo's state dict to TL property format."""
132 from huggingface_hub import HfApi, hf_hub_download
134 repo_files = HfApi().list_repo_files(repo_id)
135 suffix = f"{checkpoint_value}.pth" if checkpoint_value is not None else "final.pth"
136 matches = [f for f in repo_files if f.endswith(suffix)]
137 if not matches: 137 ↛ 138line 137 didn't jump to line 138 because the condition on line 137 was never true
138 raise FileNotFoundError(f"No '*{suffix}' file in {repo_id}; files: {repo_files[:8]}...")
139 state_dict = torch.load(
140 hf_hub_download(repo_id, matches[0]), map_location="cpu", weights_only=True
141 )
142 state_dict = {k: v.to(dtype) for k, v in state_dict.items()}
144 original_architecture = getattr(cfg, "original_architecture", None)
145 if original_architecture == "neel-solu-old": 145 ↛ 146line 145 didn't jump to line 146 because the condition on line 145 was never true
146 state_dict = _convert_neel_solu_old_weights(state_dict, cfg)
147 elif original_architecture == "mingpt": 147 ↛ 148line 147 didn't jump to line 148 because the condition on line 147 was never true
148 state_dict = _convert_mingpt_weights(state_dict, cfg)
149 return state_dict
152def boot(
153 model_name: str,
154 checkpoint_index: Optional[int] = None,
155 checkpoint_value: Optional[int] = None,
156 device: Optional[Union[str, torch.device]] = None,
157 dtype: torch.dtype = torch.float32,
158 tokenizer: Optional[Any] = None,
159) -> TransformerBridge:
160 """Build a bridge for a legacy TransformerLens-format HF repo.
162 ``checkpoint_index`` / ``checkpoint_value`` select a training checkpoint
163 (``checkpoints/*_<label>.pth``); by default the final weights load. The
164 resolved values are stamped on ``cfg.checkpoint_index`` /
165 ``cfg.checkpoint_value``, mirroring the legacy loader.
166 """
167 if not model_name.startswith(TL_LEGACY_PREFIXES):
168 raise ValueError(
169 f"{model_name!r} is not a known legacy TransformerLens repo family "
170 f"{TL_LEGACY_PREFIXES}. Use TransformerBridge.boot_transformers for "
171 "HuggingFace-format models."
172 )
173 cfg = derive_tl_legacy_config(model_name)
175 if checkpoint_index is not None or checkpoint_value is not None:
176 from transformer_lens.tools.model_registry.checkpoints import (
177 get_checkpoint_labels,
178 )
180 labels, _ = get_checkpoint_labels(model_name)
181 if checkpoint_value is None: 181 ↛ 191line 181 didn't jump to line 191 because the condition on line 181 was always true
182 assert checkpoint_index is not None
183 # Negative indices count from the end (checkpoint_index=-1 is the
184 # final checkpoint), matching the legacy loader's list indexing.
185 if not -len(labels) <= checkpoint_index < len(labels):
186 raise ValueError(
187 f"checkpoint_index={checkpoint_index} out of range "
188 f"[-{len(labels)}, {len(labels)}) for {model_name!r}."
189 )
190 checkpoint_value = labels[checkpoint_index]
191 elif checkpoint_value not in labels:
192 raise ValueError(
193 f"checkpoint_value={checkpoint_value} not in available checkpoints for "
194 f"{model_name!r} ({len(labels)} labels, {labels[0]}..{labels[-1]})."
195 )
196 cfg.checkpoint_index = labels.index(checkpoint_value) # type: ignore[attr-defined]
197 cfg.checkpoint_value = checkpoint_value # type: ignore[attr-defined]
198 else:
199 cfg.checkpoint_index = None # type: ignore[attr-defined]
200 cfg.checkpoint_value = None # type: ignore[attr-defined]
202 if tokenizer is None and cfg.tokenizer_name is not None: 202 ↛ 207line 202 didn't jump to line 207 because the condition on line 202 was always true
203 from transformers import AutoTokenizer
205 tokenizer = AutoTokenizer.from_pretrained(cfg.tokenizer_name)
207 from transformer_lens.model_bridge.sources.native import boot as _boot_native
209 bridge = _boot_native(
210 cfg, tokenizer=tokenizer, device=device, dtype=dtype, model_name=model_name
211 )
212 legacy_sd = fetch_tl_legacy_state_dict(
213 model_name, cfg, checkpoint_value=checkpoint_value, dtype=dtype
214 )
215 converted = convert_tl_checkpoint(legacy_sd, cfg)
216 # Several legacy repos ship no unembed bias; the legacy loader zero-filled
217 # missing params, so mirror that for exactly this key rather than failing
218 # a strict load or silently accepting arbitrary gaps.
219 if "unembed.bias" not in converted: 219 ↛ 221line 219 didn't jump to line 221 because the condition on line 219 was always true
220 converted["unembed.bias"] = torch.zeros(cfg.d_vocab, dtype=dtype)
221 result = bridge.load_state_dict(converted, strict=True)
222 assert not result.missing_keys and not result.unexpected_keys
223 return bridge
226setattr(TransformerBridge, "boot_tl_legacy", staticmethod(boot))