Coverage for transformer_lens/model_bridge/sources/transformers.py: 81%

1"""Transformers module for TransformerLens. 

2 

3This module provides functionality to load and convert models from HuggingFace to TransformerLens format. 

4""" 

5import contextlib 

6import copy 

7import logging 

8import os 

9import warnings 

10from typing import Any 

11 

12import torch 

13from transformers import ( 

14 AutoConfig, 

15 AutoModelForCausalLM, 

16 AutoModelForMaskedLM, 

17 AutoModelForSeq2SeqLM, 

18 AutoTokenizer, 

19 PreTrainedTokenizerBase, 

20) 

21 

22from transformer_lens.config import TransformerBridgeConfig 

23from transformer_lens.factories.architecture_adapter_factory import ( 

24 SUPPORTED_ARCHITECTURES, 

25 ArchitectureAdapterFactory, 

26) 

27from transformer_lens.model_bridge.bridge import TransformerBridge 

28from transformer_lens.model_bridge.sources._bridge_builder import _HF_PASSTHROUGH_ATTRS 

29from transformer_lens.supported_models import MODEL_ALIASES 

30from transformer_lens.utilities import get_device, get_tokenizer_with_bos 

31from transformer_lens.utilities.heterogeneous_config import ( 

32 het_safe_view, 

33 majority_value, 

34 per_layer_attr_names, 

35 per_layer_values, 

36) 

37 

38# Suppress transformers warnings that go to stderr 

39# This prevents notebook tests from failing due to unexpected stderr output 

40warnings.filterwarnings("ignore", message=".*generation flags.*not valid.*") 

41logging.getLogger("transformers").setLevel(logging.ERROR) 

42 

43 

44def get_effective_text_config(hf_config: Any) -> Any: 

45 """Return the config that owns the language-model forward path.""" 

46 if getattr(hf_config, "text_config", None) is not None: 

47 return hf_config.text_config 

48 decoder = getattr(hf_config, "decoder", None) 

49 if decoder is not None and hasattr(decoder, "hidden_size"): 

50 return decoder 

51 return hf_config 

52 

53 

54def map_default_transformer_lens_config(hf_config): 

55 """Map HuggingFace config fields to TransformerLens config format. 

56 

57 This function provides a standardized mapping from various HuggingFace config 

58 field names to the consistent TransformerLens naming convention. 

59 

60 For multimodal models (LLaVA, Gemma3ForConditionalGeneration), the language 

61 model dimensions are nested under text_config. We extract from text_config 

62 first, then apply the standard mapping. 

63 

64 Args: 

65 hf_config: The HuggingFace config object 

66 

67 Returns: 

68 A copy of hf_config with additional TransformerLens fields 

69 """ 

70 # Extract language model config from text_config for multimodal models 

71 raw_source_config = get_effective_text_config(hf_config) 

72 

73 tl_config = copy.deepcopy(hf_config) 

74 

75 # transformers>=5.15 het configs raise (not AttributeError, so hasattr does 

76 # NOT suppress it) on global reads of per-layer attrs; the view resolves 

77 # them to majority-layer values so every probe below stays safe. 

78 het_attrs = per_layer_attr_names(raw_source_config) 

79 source_config = het_safe_view(raw_source_config) 

80 

81 def legacy_per_layer(base_name: str, global_name: str) -> Any: 

82 """Pre-5.15 Gemma 4 split geometry across <field> (sliding-attention layers) 

83 and global_<field> (full-attention layers); rebuild the per-layer view.""" 

84 if base_name in het_attrs or "layer_types" in het_attrs: 84 ↛ 85line 84 didn't jump to line 85 because the condition on line 84 was never true

85 return None 

86 base = getattr(source_config, base_name, None) 

87 full = getattr(source_config, global_name, None) 

88 layer_types = getattr(source_config, "layer_types", None) 

89 if base is None or full is None or full == base or not layer_types: 

90 return None 

91 return [full if t == "full_attention" else base for t in layer_types] 

92 

93 if hasattr(source_config, "n_embd"): 

94 tl_config.d_model = source_config.n_embd 

95 elif hasattr(source_config, "hidden_size"): 

96 tl_config.d_model = source_config.hidden_size 

97 elif hasattr(source_config, "model_dim"): 97 ↛ 98line 97 didn't jump to line 98 because the condition on line 97 was never true

98 tl_config.d_model = source_config.model_dim 

99 elif hasattr(source_config, "d_model"): 99 ↛ 101line 99 didn't jump to line 101 because the condition on line 99 was always true

100 tl_config.d_model = source_config.d_model 

101 elif hasattr(source_config, "hidden_dim"): 

102 tl_config.d_model = source_config.hidden_dim 

103 if hasattr(source_config, "n_head"): 

104 tl_config.n_heads = source_config.n_head 

105 elif hasattr(source_config, "num_attention_heads"): 

106 n_heads = source_config.num_attention_heads 

107 if isinstance(n_heads, list): 107 ↛ 108line 107 didn't jump to line 108 because the condition on line 107 was never true

108 n_heads = max(n_heads) 

109 tl_config.n_heads = n_heads 

110 elif hasattr(source_config, "num_heads"): 

111 tl_config.n_heads = source_config.num_heads 

112 elif hasattr(source_config, "n_heads"): 

113 tl_config.n_heads = source_config.n_heads 

114 elif hasattr(source_config, "num_query_heads") and isinstance( 114 ↛ 117line 114 didn't jump to line 117 because the condition on line 114 was never true

115 source_config.num_query_heads, list 

116 ): 

117 tl_config.n_heads = max(source_config.num_query_heads) 

118 if "num_key_value_heads" in het_attrs: 

119 per_layer_kv = per_layer_values(source_config, "num_key_value_heads") 

120 elif getattr(source_config, "attention_k_eq_v", True): 

121 # HF applies num_global_key_value_heads only on attention_k_eq_v models; 

122 # the 5.15 per-layer migration gates identically, defaulting True when absent. 

123 per_layer_kv = legacy_per_layer("num_key_value_heads", "num_global_key_value_heads") 

124 else: 

125 per_layer_kv = None 

126 kv_values = [v for v in per_layer_kv if v is not None] if per_layer_kv else [] 

127 if kv_values: 

128 # Heterogeneous KV geometry (e.g. Gemma 4 31B: 16 KV heads on sliding layers, 

129 # 4 on full-attention layers). The scalar keeps the majority-layer value — 

130 # attention math is delegated to HF for these architectures — and the 

131 # per-layer truth is preserved alongside it. 

132 tl_config.per_layer_num_key_value_heads = per_layer_kv 

133 try: 

134 num_kv_heads = int(majority_value(kv_values)) 

135 num_heads = int(getattr(tl_config, "n_heads", 0)) 

136 if num_kv_heads != num_heads: 136 ↛ 183line 136 didn't jump to line 183 because the condition on line 136 was always true

137 tl_config.n_key_value_heads = num_kv_heads 

138 except (TypeError, ValueError): 

139 pass 

140 elif ( 

141 hasattr(source_config, "num_key_value_heads") 

142 and source_config.num_key_value_heads is not None 

143 ): 

144 try: 

145 num_kv_heads = source_config.num_key_value_heads 

146 # Handle per-layer lists (e.g., OpenELM) by taking the max 

147 if isinstance(num_kv_heads, list): 147 ↛ 148line 147 didn't jump to line 148 because the condition on line 147 was never true

148 num_kv_heads = max(num_kv_heads) 

149 if hasattr(num_kv_heads, "item"): 149 ↛ 150line 149 didn't jump to line 150 because the condition on line 149 was never true

150 num_kv_heads = num_kv_heads.item() 

151 num_kv_heads = int(num_kv_heads) 

152 num_heads = tl_config.n_heads 

153 if hasattr(num_heads, "item"): 153 ↛ 154line 153 didn't jump to line 154 because the condition on line 153 was never true

154 num_heads = num_heads.item() 

155 num_heads = int(num_heads) 

156 if num_kv_heads != num_heads: 

157 tl_config.n_key_value_heads = num_kv_heads 

158 except (TypeError, ValueError, AttributeError): 

159 pass 

160 elif hasattr(source_config, "num_kv_heads") and source_config.num_kv_heads is not None: 

161 try: 

162 num_kv_heads = source_config.num_kv_heads 

163 if isinstance(num_kv_heads, list): 163 ↛ 164line 163 didn't jump to line 164 because the condition on line 163 was never true

164 num_kv_heads = max(num_kv_heads) 

165 if hasattr(num_kv_heads, "item"): 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true

166 num_kv_heads = num_kv_heads.item() 

167 num_kv_heads = int(num_kv_heads) 

168 num_heads = tl_config.n_heads 

169 if hasattr(num_heads, "item"): 169 ↛ 170line 169 didn't jump to line 170 because the condition on line 169 was never true

170 num_heads = num_heads.item() 

171 num_heads = int(num_heads) 

172 if num_kv_heads != num_heads: 

173 tl_config.n_key_value_heads = num_kv_heads 

174 except (TypeError, ValueError, AttributeError): 

175 pass 

176 elif hasattr(source_config, "n_kv_heads") and source_config.n_kv_heads is not None: 

177 try: 

178 num_kv_heads = int(source_config.n_kv_heads) 

179 if num_kv_heads != getattr(tl_config, "n_heads", None): 

180 tl_config.n_key_value_heads = num_kv_heads 

181 except (TypeError, ValueError, AttributeError): 

182 pass 

183 if hasattr(source_config, "n_layer"): 

184 tl_config.n_layers = source_config.n_layer 

185 elif hasattr(source_config, "num_hidden_layers"): 

186 tl_config.n_layers = source_config.num_hidden_layers 

187 elif hasattr(source_config, "num_transformer_layers"): 187 ↛ 188line 187 didn't jump to line 188 because the condition on line 187 was never true

188 tl_config.n_layers = source_config.num_transformer_layers 

189 elif hasattr(source_config, "num_layers"): 189 ↛ 190line 189 didn't jump to line 190 because the condition on line 189 was never true

190 tl_config.n_layers = source_config.num_layers 

191 elif hasattr(source_config, "n_layers"): 

192 tl_config.n_layers = source_config.n_layers 

193 elif hasattr(source_config, "n_blocks"): 193 ↛ 194line 193 didn't jump to line 194 because the condition on line 193 was never true

194 tl_config.n_layers = source_config.n_blocks 

195 if hasattr(source_config, "vocab_size") and isinstance(source_config.vocab_size, int): 

196 tl_config.d_vocab = source_config.vocab_size 

197 if hasattr(source_config, "n_positions"): 

198 tl_config.n_ctx = source_config.n_positions 

199 elif hasattr(source_config, "max_position_embeddings"): 

200 tl_config.n_ctx = source_config.max_position_embeddings 

201 elif hasattr(source_config, "max_context_length"): 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true

202 tl_config.n_ctx = source_config.max_context_length 

203 elif hasattr(source_config, "max_length"): 

204 tl_config.n_ctx = source_config.max_length 

205 elif hasattr(source_config, "seq_length"): 

206 tl_config.n_ctx = source_config.seq_length 

207 elif hasattr(source_config, "image_size") and hasattr(source_config, "patch_size"): 

208 # Vision Transformers calculate sequence length dynamically: 

209 # (image_size / patch_size)^2 + 1 (for the CLS token) 

210 image_size = source_config.image_size 

211 patch_size = source_config.patch_size 

212 

213 # HF configs allow these to be integers or tuples/lists 

214 img_h = image_size[0] if isinstance(image_size, (list, tuple)) else image_size 

215 img_w = image_size[1] if isinstance(image_size, (list, tuple)) else image_size 

216 patch_h = patch_size[0] if isinstance(patch_size, (list, tuple)) else patch_size 

217 patch_w = patch_size[1] if isinstance(patch_size, (list, tuple)) else patch_size 

218 

219 tl_config.n_ctx = (img_h // patch_h) * (img_w // patch_w) + 1 

220 elif hasattr(source_config, "max_sequence_length"): 

221 tl_config.n_ctx = source_config.max_sequence_length 

222 else: 

223 # Models like Bloom use ALiBi (no positional embeddings) and have no 

224 # context length field. Default to 2048 as a reasonable fallback. 

225 tl_config.n_ctx = 2048 

226 if hasattr(source_config, "n_inner"): 

227 tl_config.d_mlp = source_config.n_inner 

228 elif "intermediate_size" in het_attrs: 

229 # Same max collapse as the per-layer-list case below. 

230 mlp_values = [v for v in per_layer_values(source_config, "intermediate_size") if v] 

231 tl_config.d_mlp = max(mlp_values) if mlp_values else None 

232 elif hasattr(source_config, "intermediate_size"): 

233 intermediate_size = source_config.intermediate_size 

234 # Gemma 3n exposes a per-layer intermediate_size list (the MatFormer design permits 

235 # variation). All released checkpoints (E2B/E4B) are uniform, and d_mlp is scalar 

236 # metadata (the bridge defers MLP math to HF), so collapse to max — the shared value 

237 # when uniform, an upper bound otherwise. 

238 if isinstance(intermediate_size, (list, tuple)): 

239 intermediate_size = max(intermediate_size) if intermediate_size else None 

240 tl_config.d_mlp = intermediate_size 

241 elif hasattr(source_config, "mlp_hidden_size"): 

242 tl_config.d_mlp = source_config.mlp_hidden_size 

243 elif hasattr(tl_config, "d_model"): 243 ↛ 245line 243 didn't jump to line 245 because the condition on line 243 was always true

244 tl_config.d_mlp = getattr(source_config, "n_inner", 4 * tl_config.d_model) 

245 if "head_dim" in het_attrs: 

246 per_layer_hd = per_layer_values(source_config, "head_dim") 

247 else: 

248 per_layer_hd = legacy_per_layer("head_dim", "global_head_dim") 

249 hd_values = [v for v in per_layer_hd if v is not None] if per_layer_hd else [] 

250 if hd_values: 

251 # Heterogeneous head_dim (e.g. Gemma 4: 256 on sliding layers, 512 on 

252 # full-attention layers). Scalar d_head keeps the majority-layer value; 

253 # the per-layer truth is preserved alongside it. 

254 tl_config.per_layer_head_dim = per_layer_hd 

255 tl_config.d_head = majority_value(hd_values) 

256 elif hasattr(source_config, "head_dim") and source_config.head_dim is not None: 

257 tl_config.d_head = source_config.head_dim 

258 elif hasattr(tl_config, "d_model") and hasattr(tl_config, "n_heads"): 

259 tl_config.d_head = tl_config.d_model // tl_config.n_heads 

260 elif hasattr(tl_config, "d_model"): 260 ↛ 266line 260 didn't jump to line 266 because the condition on line 260 was always true

261 # Models without attention (e.g., Mamba SSMs) have no n_heads or head_dim. 

262 # Set d_head = d_model so TransformerLensConfig.__post_init__ computes 

263 # n_heads = 1. These values are nominal and have no functional meaning 

264 # for attention-less architectures. 

265 tl_config.d_head = tl_config.d_model 

266 if hasattr(source_config, "activation_function"): 

267 tl_config.act_fn = source_config.activation_function 

268 # Gemma family: transformers 5.x exposes only hidden_activation (hidden_act 

269 # was removed); it is authoritative over hidden_act when both exist. 

270 elif getattr(source_config, "hidden_activation", None) is not None: 

271 tl_config.act_fn = source_config.hidden_activation 

272 elif hasattr(source_config, "hidden_act"): 

273 tl_config.act_fn = source_config.hidden_act 

274 elif hasattr(source_config, "activation_type"): 

275 activation_type = source_config.activation_type 

276 tl_config.act_fn = getattr(activation_type, "value", activation_type) 

277 elif getattr(source_config, "activation_fn_name", None) is not None: 

278 # OpenELM spells it activation_fn_name ("swish"); without this the 

279 # cfg keeps the "relu" default and reconstructed FFNs silently 

280 # diverge ~30% from HF. 

281 tl_config.act_fn = source_config.activation_fn_name 

282 if hasattr(source_config, "rope_theta"): 

283 tl_config.rotary_base = source_config.rope_theta 

284 if hasattr(source_config, "weight_tying"): 

285 tl_config.tie_word_embeddings = bool(source_config.weight_tying) 

286 # Layer norm / RMS norm epsilon — HF uses 3 different field names 

287 if hasattr(source_config, "rms_norm_eps"): 

288 tl_config.eps = source_config.rms_norm_eps 

289 elif hasattr(source_config, "layer_norm_eps"): 

290 tl_config.eps = source_config.layer_norm_eps 

291 elif hasattr(source_config, "layer_norm_epsilon"): 

292 tl_config.eps = source_config.layer_norm_epsilon 

293 elif hasattr(source_config, "norm_eps"): 

294 tl_config.eps = source_config.norm_eps 

295 if hasattr(source_config, "num_experts"): 

296 tl_config.num_experts = source_config.num_experts 

297 elif hasattr(source_config, "num_local_experts"): 

298 tl_config.num_experts = source_config.num_local_experts 

299 if hasattr(source_config, "num_experts_per_tok"): 

300 tl_config.experts_per_token = source_config.num_experts_per_tok 

301 if hasattr(source_config, "sliding_window") and source_config.sliding_window is not None: 

302 tl_config.sliding_window = source_config.sliding_window 

303 if getattr(hf_config, "use_parallel_residual", False): 

304 tl_config.parallel_attn_mlp = True 

305 # GPT-J and CodeGen: parallel attn+MLP but missing use_parallel_residual in HF config 

306 arch_classes = getattr(hf_config, "architectures", []) or [] 

307 if any(a in ("GPTJForCausalLM", "CodeGenForCausalLM") for a in arch_classes): 307 ↛ 308line 307 didn't jump to line 308 because the condition on line 307 was never true

308 tl_config.parallel_attn_mlp = True 

309 tl_config.default_prepend_bos = True 

310 return tl_config 

311 

312 

313def determine_architecture_from_hf_config(hf_config): 

314 """Determine the architecture name from HuggingFace config. 

315 

316 Args: 

317 hf_config: The HuggingFace config object 

318 

319 Returns: 

320 str: The architecture name (e.g., "GPT2LMHeadModel", "LlamaForCausalLM") 

321 

322 Raises: 

323 ValueError: If architecture cannot be determined 

324 """ 

325 architectures = [] 

326 if hasattr(hf_config, "original_architecture"): 326 ↛ 327line 326 didn't jump to line 327 because the condition on line 326 was never true

327 architectures.append(hf_config.original_architecture) 

328 if hasattr(hf_config, "architectures") and hf_config.architectures: 

329 architectures.extend(hf_config.architectures) 

330 if hasattr(hf_config, "model_type"): 330 ↛ 430line 330 didn't jump to line 430 because the condition on line 330 was always true

331 model_type = hf_config.model_type 

332 model_type_mappings = { 

333 "afmoe": "AfmoeForCausalLM", 

334 "apertus": "ApertusForCausalLM", 

335 "gpt2": "GPT2LMHeadModel", 

336 "openai-gpt": "OpenAIGPTLMHeadModel", 

337 "hubert": "HubertModel", 

338 "bamba": "BambaForCausalLM", 

339 "bitnet": "BitNetForCausalLM", 

340 "blenderbot": "BlenderbotForConditionalGeneration", 

341 "bart": "BartForConditionalGeneration", 

342 "ernie4_5": "Ernie4_5ForCausalLM", 

343 "ernie4_5_moe": "Ernie4_5_MoeForCausalLM", 

344 "exaone": "ExaoneForCausalLM", 

345 "exaone4": "Exaone4ForCausalLM", 

346 "falcon_mamba": "FalconMambaForCausalLM", 

347 "florence2": "Florence2ForConditionalGeneration", 

348 "longt5": "LongT5ForConditionalGeneration", 

349 "m2m_100": "M2M100ForConditionalGeneration", 

350 "marian": "MarianMTModel", 

351 "mbart": "MBartForConditionalGeneration", 

352 "pegasus": "PegasusForConditionalGeneration", 

353 "seed_oss": "SeedOssForCausalLM", 

354 "starcoder2": "Starcoder2ForCausalLM", 

355 "nemotron": "NemotronForCausalLM", 

356 "idefics3": "Idefics3ForConditionalGeneration", 

357 "qwen2_audio": "Qwen2AudioForConditionalGeneration", 

358 "audioflamingo3": "AudioFlamingo3ForConditionalGeneration", 

359 "musicflamingo": "MusicFlamingoForConditionalGeneration", 

360 "jetmoe": "JetMoeForCausalLM", 

361 "minimax_m2": "MiniMaxM2ForCausalLM", 

362 "led": "LEDForConditionalGeneration", 

363 "llama": "LlamaForCausalLM", 

364 "llama4_text": "Llama4ForCausalLM", 

365 "llama4": "Llama4ForConditionalGeneration", 

366 "llada": "LLaDAModelLM", 

367 "mamba": "MambaForCausalLM", 

368 "mamba2": "Mamba2ForCausalLM", 

369 "mistral": "MistralForCausalLM", 

370 "olmo_hybrid": "OlmoHybridForCausalLM", 

371 "mistral3": "Mistral3ForConditionalGeneration", 

372 "mixtral": "MixtralForCausalLM", 

373 "mpt": "MptForCausalLM", 

374 "gemma": "GemmaForCausalLM", 

375 "gemma2": "Gemma2ForCausalLM", 

376 "gemma3": "Gemma3ForCausalLM", 

377 # gemma3n is tri-modal; the text path loads as the full ForConditionalGeneration 

378 # (vision/audio referenced but unbridged in the text-only adapter). 

379 "gemma3n": "Gemma3nForConditionalGeneration", 

380 # gemma4 is multimodal-only; all released checkpoints load as the full 

381 # ForConditionalGeneration (vision/audio referenced but unbridged). 

382 "gemma4": "Gemma4ForConditionalGeneration", 

383 "gemma4_unified": "Gemma4UnifiedForConditionalGeneration", 

384 "gemma4_text": "Gemma4ForCausalLM", 

385 "glm": "GlmForCausalLM", 

386 "glm4": "Glm4ForCausalLM", 

387 "glm4v": "Glm4vForConditionalGeneration", 

388 "glmasr": "GlmAsrForConditionalGeneration", 

389 "glm4_moe": "Glm4MoeForCausalLM", 

390 "glm4_moe_lite": "Glm4MoeLiteForCausalLM", 

391 "glm_moe_dsa": "GlmMoeDsaForCausalLM", 

392 "bert": "BertForMaskedLM", 

393 "bloom": "BloomForCausalLM", 

394 "codegen": "CodeGenForCausalLM", 

395 "cohere2": "Cohere2ForCausalLM", 

396 "gptj": "GPTJForCausalLM", 

397 "gpt_neo": "GPTNeoForCausalLM", 

398 "gpt_neox": "GPTNeoXForCausalLM", 

399 "opt": "OPTForCausalLM", 

400 "phi": "PhiForCausalLM", 

401 "phi3": "Phi3ForCausalLM", 

402 "qwen": "QwenForCausalLM", 

403 "qwen2": "Qwen2ForCausalLM", 

404 "qwen2_5_vl": "Qwen2_5_VLForConditionalGeneration", 

405 "qwen3_vl": "Qwen3VLForConditionalGeneration", 

406 "qwen3_vl_moe": "Qwen3VLMoeForConditionalGeneration", 

407 "qwen2_moe": "Qwen2MoeForCausalLM", 

408 "qwen3": "Qwen3ForCausalLM", 

409 # qwen3_5 is the top-level multimodal config type; qwen3_5_text is 

410 # the text-only sub-config. Both map to the text-only adapter so 

411 # Qwen3.5 checkpoints (which report qwen3_5 even when loaded as 

412 # text-only) are routed to Qwen3_5ForCausalLM. 

413 "qwen3_5": "Qwen3_5ForCausalLM", 

414 "qwen3_5_text": "Qwen3_5ForCausalLM", 

415 # Same routing convention for the MoE variant. 

416 "qwen3_5_moe": "Qwen3_5MoeForCausalLM", 

417 "qwen3_5_moe_text": "Qwen3_5MoeForCausalLM", 

418 "smollm3": "SmolLM3ForCausalLM", 

419 "openelm": "OpenELMForCausalLM", 

420 "ouro": "OuroForCausalLM", 

421 "stablelm": "StableLmForCausalLM", 

422 "t5": "T5ForConditionalGeneration", 

423 "mt5": "MT5ForConditionalGeneration", 

424 "t5gemma": "T5GemmaForConditionalGeneration", 

425 "t5gemma2": "T5Gemma2ForConditionalGeneration", 

426 } 

427 if model_type in model_type_mappings: 

428 architectures.append(model_type_mappings[model_type]) 

429 

430 for arch in architectures: 430 ↛ 433line 430 didn't jump to line 433 because the loop on line 430 didn't complete

431 if arch in SUPPORTED_ARCHITECTURES: 431 ↛ 430line 431 didn't jump to line 430 because the condition on line 431 was always true

432 return arch 

433 raise ValueError( 

434 f"Could not determine supported architecture from config. Available architectures: {list(SUPPORTED_ARCHITECTURES.keys())}, Config architectures: {architectures}, Model type: {getattr(hf_config, 'model_type', None)}" 

435 ) 

436 

437 

438def get_hf_model_class_for_architecture(architecture: str): 

439 """Determine the correct HuggingFace AutoModel class for loading. 

440 

441 Uses centralized architecture sets from utilities.architectures. 

442 """ 

443 from transformer_lens.utilities.architectures import ( 

444 AUDIO_ARCHITECTURES, 

445 AUDIO_CLASSIFICATION_ARCHITECTURES, 

446 AUDIO_TEXT_ARCHITECTURES, 

447 BASE_AUTOMODEL_ARCHITECTURES, 

448 MASKED_LM_ARCHITECTURES, 

449 MULTIMODAL_ARCHITECTURES, 

450 SEQ2SEQ_ARCHITECTURES, 

451 VISION_ARCHITECTURES, 

452 VISION_CLASSIFICATION_ARCHITECTURES, 

453 ) 

454 

455 if architecture in SEQ2SEQ_ARCHITECTURES or architecture in AUDIO_TEXT_ARCHITECTURES: 

456 return AutoModelForSeq2SeqLM 

457 elif architecture in MASKED_LM_ARCHITECTURES: 

458 return AutoModelForMaskedLM 

459 elif architecture in MULTIMODAL_ARCHITECTURES: 

460 from transformers import AutoModelForImageTextToText 

461 

462 return AutoModelForImageTextToText 

463 elif architecture in BASE_AUTOMODEL_ARCHITECTURES: 463 ↛ 464line 463 didn't jump to line 464 because the condition on line 463 was never true

464 from transformers import AutoModel 

465 

466 return AutoModel 

467 elif architecture in AUDIO_CLASSIFICATION_ARCHITECTURES: 

468 from transformers import AutoModelForAudioClassification 

469 

470 return AutoModelForAudioClassification 

471 elif architecture in AUDIO_ARCHITECTURES: 471 ↛ 472line 471 didn't jump to line 472 because the condition on line 471 was never true

472 if "ForCTC" in architecture: 

473 from transformers import AutoModelForCTC 

474 

475 return AutoModelForCTC 

476 from transformers import AutoModel 

477 

478 return AutoModel 

479 elif architecture in VISION_ARCHITECTURES: 

480 if architecture in VISION_CLASSIFICATION_ARCHITECTURES: 

481 from transformers import AutoModelForImageClassification 

482 

483 return AutoModelForImageClassification 

484 from transformers import AutoModel 

485 

486 return AutoModel 

487 else: 

488 return AutoModelForCausalLM 

489 

490 

491# Known training-checkpoint revision conventions on HF. 

492_CHECKPOINT_REVISION_FORMATS: dict[str, str] = { 

493 "EleutherAI/pythia": "step{value}", 

494 "stanford-crfm": "checkpoint-{value}", 

495} 

496 

497 

498def _resolve_checkpoint_to_revision( 

499 model_name: str, 

500 checkpoint_index: int | None, 

501 checkpoint_value: int | None, 

502) -> str: 

503 """Convert a checkpoint index/value into an HF revision string, validated against ``get_checkpoint_labels``.""" 

504 if checkpoint_index is None and checkpoint_value is None: 

505 raise ValueError("Must specify either checkpoint_index or checkpoint_value.") 

506 

507 format_str: str | None = None 

508 for prefix, fmt in _CHECKPOINT_REVISION_FORMATS.items(): 

509 if model_name.startswith(prefix): 

510 format_str = fmt 

511 break 

512 if format_str is None: 

513 raise ValueError( 

514 f"Model {model_name!r} does not have a known checkpoint revision convention. " 

515 f"Pass revision= directly if your model uses HF revisions. Known checkpoint " 

516 f"families: {list(_CHECKPOINT_REVISION_FORMATS.keys())}." 

517 ) 

518 

519 from transformer_lens.loading_from_pretrained import get_checkpoint_labels 

520 

521 labels, _ = get_checkpoint_labels(model_name) 

522 if checkpoint_value is not None: 

523 if checkpoint_value not in labels: 

524 raise ValueError( 

525 f"checkpoint_value={checkpoint_value} not in available checkpoints for " 

526 f"{model_name!r}. {len(labels)} labels available, " 

527 f"first/last: {labels[0]}..{labels[-1]}." 

528 ) 

529 else: 

530 assert checkpoint_index is not None # narrowed by initial guard 

531 if not 0 <= checkpoint_index < len(labels): 

532 raise ValueError( 

533 f"checkpoint_index={checkpoint_index} out of range [0, {len(labels)}) " 

534 f"for {model_name!r}." 

535 ) 

536 checkpoint_value = labels[checkpoint_index] 

537 return format_str.format(value=checkpoint_value) 

538 

539 

540def boot( 

541 model_name: str, 

542 hf_config_overrides: dict | None = None, 

543 device: str | torch.device | None = None, 

544 dtype: torch.dtype = torch.float32, 

545 tokenizer: PreTrainedTokenizerBase | None = None, 

546 load_weights: bool = True, 

547 trust_remote_code: bool = False, 

548 model_class: Any | None = None, 

549 hf_model: Any | None = None, 

550 n_ctx: int | None = None, 

551 revision: str | None = None, 

552 checkpoint_index: int | None = None, 

553 checkpoint_value: int | None = None, 

554 # Multi-device placement (accelerate-dispatched). GPU-validated 2026-07-16: 

555 # tests/acceptance/model_bridge/test_bridge_multigpu*.py + scripts/bridge_multi_device_parity.py. 

556 device_map: str | dict[str, str | int] | None = None, 

557 n_devices: int | None = None, 

558 max_memory: dict[str | int, str | int] | None = None, 

559) -> TransformerBridge: 

560 """Boot a model from HuggingFace. 

561 

562 Args: 

563 model_name: The name of the model to load. 

564 hf_config_overrides: Optional overrides applied to the HuggingFace config before model load. 

565 device: The device to use. If None, will be determined automatically. Mutually exclusive 

566 with ``device_map``. 

567 dtype: The dtype to use for the model. 

568 tokenizer: Optional pre-initialized tokenizer to use; if not provided one will be created. 

569 load_weights: If False, load model without weights (on meta device) for config inspection only. 

570 model_class: Optional HuggingFace model class to use instead of the default auto-detected 

571 class. When the class name matches a key in SUPPORTED_ARCHITECTURES, the corresponding 

572 adapter is selected automatically (e.g., BertForNextSentencePrediction). 

573 hf_model: Optional pre-loaded HuggingFace model to use instead of loading one. Useful for 

574 models loaded with custom configurations (e.g., quantization via BitsAndBytesConfig). 

575 When provided, load_weights is ignored. 

576 device_map: HuggingFace-style device map (``"auto"``, ``"balanced"``, dict, etc.) for 

577 dispatched inference. Explicit maps may include CPU targets; disk / meta offload 

578 targets are still rejected because Bridge component wrappers need additional 

579 offload-hook routing work. Mutually exclusive with ``device``. 

580 n_devices: Convenience: split the model across this many CUDA devices (translated to a 

581 ``max_memory`` dict internally). Requires CUDA with at least this many visible devices. 

582 max_memory: Optional per-device memory budget for HF's dispatcher. 

583 n_ctx: Optional context length override. The bridge normally uses the model's documented 

584 max context from the HF config. Setting this writes to whichever HF field the model 

585 uses (n_positions / max_position_embeddings / etc.), so callers don't need to know 

586 the field name. If larger than the model's default, a warning is emitted — quality 

587 may degrade past the trained length for rotary models. 

588 revision: Optional HF revision string (branch, tag, or commit). Forwarded to 

589 config, model, and tokenizer loading. 

590 Mutually exclusive with ``checkpoint_index`` and ``checkpoint_value``. 

591 checkpoint_index: Index into the available training checkpoints for the model family. 

592 Convenience over ``revision`` for checkpointed models like EleutherAI/pythia* and 

593 stanford-crfm/*. Resolved to a revision string via the known per-family naming 

594 conventions (``step{value}`` for Pythia, ``checkpoint-{value}`` for stanford-crfm). 

595 checkpoint_value: Training step or token count of the desired checkpoint. Alternative to 

596 ``checkpoint_index``; must be one of the labels returned by ``get_checkpoint_labels``. 

597 

598 Returns: 

599 The bridge to the loaded model. 

600 """ 

601 for official_name, aliases in MODEL_ALIASES.items(): 

602 if model_name in aliases: 

603 logging.warning( 

604 f"DEPRECATED: You are using a deprecated, model_name alias '{model_name}'. TransformerLens will now load the official transformers model name, '{official_name}' instead.\n Please update your code to use the official name by changing model_name from '{model_name}' to '{official_name}'.\nSince TransformerLens v3, all model names should be the official transformers model names.\nThe aliases will be removed in the next version of TransformerLens, so please do the update now." 

605 ) 

606 model_name = official_name 

607 break 

608 if checkpoint_index is not None or checkpoint_value is not None: 

609 if revision is not None: 

610 raise ValueError( 

611 "Specify either revision= or checkpoint_index/checkpoint_value, not both." 

612 ) 

613 revision = _resolve_checkpoint_to_revision(model_name, checkpoint_index, checkpoint_value) 

614 # Pass HF token for gated model access (e.g. meta-llama/*) 

615 from transformer_lens.utilities.hf_utils import ( 

616 autoconfig_with_remote_post_init_compat, 

617 autotokenizer_with_special_token_compat, 

618 get_hf_token, 

619 ) 

620 

621 _hf_token = get_hf_token() 

622 if hf_model is not None: 

623 # Reuse the pre-loaded model's config to avoid a Hub call when model_name 

624 # is a Hub repo ID, but the model is already loaded locally. 

625 hf_config = copy.deepcopy(hf_model.config) 

626 else: 

627 # Compat wrapper: 4.x-era remote-code configs (OpenELM) define an 

628 # argless __post_init__ that 5.x's dataclass machinery calls with the 

629 # class's own fields as kwargs — unloadable without the shim. 

630 hf_config = autoconfig_with_remote_post_init_compat( 

631 model_name, 

632 auto_config=AutoConfig, 

633 output_attentions=True, 

634 trust_remote_code=trust_remote_code, 

635 token=_hf_token, 

636 revision=revision, 

637 ) 

638 _n_ctx_field: str | None = None 

639 if n_ctx is not None: 

640 # Reject non-positive values before doing anything else. 

641 if n_ctx <= 0: 

642 raise ValueError(f"n_ctx must be a positive integer, got n_ctx={n_ctx}.") 

643 # Resolve n_ctx to whichever HF config field this model uses. Mirrors 

644 # the order in map_default_transformer_lens_config so the TL config 

645 # derivation picks up the override. 

646 for _field in ( 646 ↛ 657line 646 didn't jump to line 657 because the loop on line 646 didn't complete

647 "n_positions", 

648 "max_position_embeddings", 

649 "max_context_length", 

650 "max_length", 

651 "seq_length", 

652 "max_sequence_length", 

653 ): 

654 if hasattr(hf_config, _field): 

655 _n_ctx_field = _field 

656 break 

657 if _n_ctx_field is None: 657 ↛ 658line 657 didn't jump to line 658 because the condition on line 657 was never true

658 raise ValueError( 

659 f"Cannot apply n_ctx={n_ctx}: no recognized context-length field on " 

660 f"HF config for {model_name}. Use hf_config_overrides instead." 

661 ) 

662 _default_n_ctx = getattr(hf_config, _n_ctx_field) 

663 if _default_n_ctx is not None and n_ctx > _default_n_ctx: 

664 logging.warning( 

665 "Setting n_ctx=%d which is larger than the model's default " 

666 "context length of %d. The model was not trained on sequences " 

667 "this long and may produce unreliable results (especially for " 

668 "rotary models without RoPE scaling).", 

669 n_ctx, 

670 _default_n_ctx, 

671 ) 

672 # Warn if the caller also set the same field via hf_config_overrides -- explicit n_ctx wins but users should know. 

673 if hf_config_overrides and _n_ctx_field in hf_config_overrides: 

674 _conflicting_value = hf_config_overrides[_n_ctx_field] 

675 if _conflicting_value != n_ctx: 

676 logging.warning( 

677 "Both n_ctx=%d and hf_config_overrides['%s']=%s were provided. " 

678 "The explicit n_ctx takes precedence.", 

679 n_ctx, 

680 _n_ctx_field, 

681 _conflicting_value, 

682 ) 

683 # Explicit n_ctx wins over hf_config_overrides for the resolved field. 

684 hf_config_overrides = dict(hf_config_overrides or {}) 

685 hf_config_overrides[_n_ctx_field] = n_ctx 

686 if hf_config_overrides: 

687 hf_config.__dict__.update(hf_config_overrides) 

688 tl_config = map_default_transformer_lens_config(hf_config) 

689 architecture = determine_architecture_from_hf_config(hf_config) 

690 config_dict = dict(tl_config.__dict__) 

691 # Restore TL attribute names that HF remaps via attribute_map 

692 if "num_local_experts" in config_dict and "num_experts" not in config_dict: 

693 config_dict["num_experts"] = config_dict["num_local_experts"] 

694 bridge_config = TransformerBridgeConfig.from_dict(config_dict) 

695 bridge_config.architecture = architecture 

696 bridge_config.model_name = model_name 

697 bridge_config.dtype = dtype 

698 bridge_config.trust_remote_code = trust_remote_code 

699 # Propagate HF-specific config attributes that adapters may need. 

700 # Canonical list lives in sources/_bridge_builder.py (architecture-agnostic). 

701 effective_config = get_effective_text_config(hf_config) 

702 # Per-layer-registered attrs would raise on global access (transformers>=5.15). 

703 _het_attrs = per_layer_attr_names(effective_config) | per_layer_attr_names(hf_config) 

704 for attr in _HF_PASSTHROUGH_ATTRS: 

705 if attr in _het_attrs: 705 ↛ 706line 705 didn't jump to line 706 because the condition on line 705 was never true

706 continue 

707 val = getattr(effective_config, attr, None) 

708 if val is None and effective_config is not hf_config: 

709 val = getattr(hf_config, attr, None) 

710 if val is not None: 

711 setattr(bridge_config, attr, val) 

712 

713 # Gemma2 softcapping: HF names differ from TL names. het view: per-layer 

714 # fields raise (not AttributeError) on raw getattr. 

715 effective_config = het_safe_view(effective_config) 

716 final_logit_softcapping = getattr(effective_config, "final_logit_softcapping", None) 

717 if final_logit_softcapping is not None: 

718 bridge_config.output_logits_soft_cap = float(final_logit_softcapping) 

719 logits_soft_cap = getattr(effective_config, "logits_soft_cap", None) 

720 if logits_soft_cap is not None: 720 ↛ 721line 720 didn't jump to line 721 because the condition on line 720 was never true

721 bridge_config.output_logits_soft_cap = float(logits_soft_cap) 

722 attn_logit_softcapping = getattr(effective_config, "attn_logit_softcapping", None) 

723 if attn_logit_softcapping is not None: 

724 bridge_config.attn_scores_soft_cap = float(attn_logit_softcapping) 

725 adapter = ArchitectureAdapterFactory.select_architecture_adapter(bridge_config) 

726 # Pre-loaded models carry their own weight placement (possibly set by the caller via 

727 # device_map). Passing device_map / n_devices / max_memory alongside hf_model= is 

728 # ambiguous and would silently be ignored, so fail loudly. 

729 if hf_model is not None and ( 

730 device_map is not None or n_devices is not None or max_memory is not None 

731 ): 

732 raise ValueError( 

733 "device_map / n_devices / max_memory are only supported when the bridge loads " 

734 "the HF model itself. When passing hf_model=..., apply device_map via " 

735 "AutoModel.from_pretrained before handing the model to the bridge." 

736 ) 

737 # Stateful/SSM (e.g. Mamba) models keep a per-layer recurrent cache that must live on 

738 # that layer's device. The bridge currently allocates the stateful cache on a single 

739 # cfg.device, so cross-device splits would silently misplace the cache. Block this 

740 # combination until a v2 addresses per-layer stateful cache placement. 

741 if (n_devices is not None and n_devices > 1) or device_map is not None: 

742 if getattr(bridge_config, "is_stateful", False): 742 ↛ 743line 742 didn't jump to line 743 because the condition on line 742 was never true

743 raise ValueError( 

744 "Multi-device splits are not yet supported for stateful (SSM / Mamba) " 

745 "architectures: the stateful cache allocation is single-device. " 

746 "Load on one device, or wait for v2 support." 

747 ) 

748 # Resolve device_map before defaulting `device` — the two are mutually exclusive, and 

749 # the resolver raises on conflict. If n_devices>1 is passed, it's translated into a 

750 # device_map + max_memory pair here so downstream code only needs to check the 

751 # resolved values. 

752 from transformer_lens.utilities.multi_gpu import ( 

753 MIXED_CPU_GPU_ERROR, 

754 count_unique_devices, 

755 find_embedding_device, 

756 find_misplaced_modules, 

757 is_mixed_cpu_gpu, 

758 resolve_device_map, 

759 ) 

760 

761 resolved_device_map, resolved_max_memory = resolve_device_map( 

762 n_devices, device_map, device, max_memory 

763 ) 

764 if resolved_device_map is None: 

765 if device is None: 

766 device = get_device() 

767 adapter.cfg.device = str(device) 

768 else: 

769 # cfg.device will be set from hf_device_map after the model is loaded. 

770 # Provisionally keep it None; find_embedding_device fills it in below. 

771 adapter.cfg.device = None 

772 if model_class is None: 

773 model_class = get_hf_model_class_for_architecture(architecture) 

774 # Ensure pad_token_id exists (v5 raises AttributeError if missing) 

775 if not hasattr(hf_config, "pad_token_id") or "pad_token_id" not in hf_config.__dict__: 

776 fallback_pad = getattr(hf_config, "eos_token_id", None) 

777 # eos_token_id can be a list (e.g., Gemma3 uses [1, 106]); take the first. 

778 if isinstance(fallback_pad, list): 

779 fallback_pad = fallback_pad[0] if fallback_pad else None 

780 hf_config.pad_token_id = fallback_pad 

781 model_kwargs = {"config": hf_config, "torch_dtype": dtype} 

782 if _hf_token: 782 ↛ 784line 782 didn't jump to line 784 because the condition on line 782 was always true

783 model_kwargs["token"] = _hf_token 

784 if trust_remote_code: 

785 model_kwargs["trust_remote_code"] = True 

786 if revision is not None: 

787 model_kwargs["revision"] = revision 

788 if resolved_device_map is not None: 

789 model_kwargs["device_map"] = resolved_device_map 

790 if resolved_max_memory is not None: 790 ↛ 791line 790 didn't jump to line 791 because the condition on line 790 was never true

791 model_kwargs["max_memory"] = resolved_max_memory 

792 if hasattr(adapter.cfg, "attn_implementation") and adapter.cfg.attn_implementation is not None: 

793 model_kwargs["attn_implementation"] = adapter.cfg.attn_implementation 

794 else: 

795 # Default to eager (required for output_attentions hooks) 

796 model_kwargs["attn_implementation"] = "eager" 

797 adapter.prepare_loading(model_name, model_kwargs) 

798 # Meta device_map targets crash at boot when loading weights 

799 # (NotImplementedError in HF tie_weights, KeyError in Accelerate offload hooks). 

800 # Only accepted with load_weights=False (config inspection; map not applied). 

801 if load_weights and isinstance(resolved_device_map, dict): 

802 _meta_targets = [ 

803 k 

804 for k, v in resolved_device_map.items() 

805 if isinstance(v, str) and v.strip().lower() == "meta" 

806 ] 

807 if _meta_targets: 

808 raise ValueError( 

809 f"device_map contains meta target(s): {_meta_targets}. " 

810 "Meta device_map values crash at boot when loading weights. " 

811 "Set load_weights=False for config inspection only " 

812 "(the map is not applied; parameters load on CPU via from_config)." 

813 ) 

814 if hf_model is not None: 

815 # Use the pre-loaded model as-is (e.g., quantized models with custom device_map) 

816 pass 

817 elif not load_weights: 

818 from_config_kwargs = {} 

819 if trust_remote_code: 819 ↛ 820line 819 didn't jump to line 820 because the condition on line 819 was never true

820 from_config_kwargs["trust_remote_code"] = True 

821 prepared_config = model_kwargs.get("config", hf_config) 

822 with contextlib.redirect_stdout(None):