Coverage for transformer_lens/loading_from_pretrained.py: 53%

1"""Loading Pretrained Models Utilities. 

2 

3This module contains functions for loading pretrained models from the Hugging Face Hub. 

4""" 

5 

6from __future__ import annotations 

7 

8import dataclasses 

9import logging 

10import os 

11import re 

12from pathlib import Path 

13from typing import Any 

14 

15import torch 

16from huggingface_hub import HfApi 

17from transformers import ( 

18 AutoConfig, 

19 AutoModel, 

20 AutoModelForCausalLM, 

21 BertForPreTraining, 

22 HubertModel, 

23 T5ForConditionalGeneration, 

24 Wav2Vec2Model, 

25) 

26 

27import transformer_lens.utilities as utils 

28from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig 

29from transformer_lens.pretrained.weight_conversions import ( 

30 convert_apertus_weights, 

31 convert_bert_weights, 

32 convert_bloom_weights, 

33 convert_coder_weights, 

34 convert_gemma_weights, 

35 convert_gpt2_weights, 

36 convert_gpt_oss_weights, 

37 convert_gptj_weights, 

38 convert_hubert_weights, 

39 convert_llama_weights, 

40 convert_mingpt_weights, 

41 convert_mistral_weights, 

42 convert_mixtral_weights, 

43 convert_neel_solu_old_weights, 

44 convert_neo_weights, 

45 convert_neox_weights, 

46 convert_olmo2_weights, 

47 convert_olmo3_weights, 

48 convert_olmo_weights, 

49 convert_olmoe_weights, 

50 convert_opt_weights, 

51 convert_phi3_weights, 

52 convert_phi_weights, 

53 convert_qwen2_weights, 

54 convert_qwen3_weights, 

55 convert_qwen_weights, 

56 convert_t5_weights, 

57) 

58from transformer_lens.supported_models import MODEL_ALIASES, OFFICIAL_MODEL_NAMES 

59from transformer_lens.utilities.hf_utils import get_rotary_pct_from_config 

60 

61NON_HF_HOSTED_MODEL_NAMES = [ 

62 "llama-7b-hf", 

63 "llama-13b-hf", 

64 "llama-30b-hf", 

65 "llama-65b-hf", 

66] 

67"""Official model names for models not hosted on HuggingFace.""" 

68 

69NEED_REMOTE_CODE_MODELS = ( 

70 "bigcode/santacoder", 

71 "Qwen/Qwen-", 

72 "Qwen/Qwen3-", 

73 "microsoft/phi-2", 

74 "microsoft/phi-4", 

75 "apple/OpenELM", 

76 "openai/gpt-oss-", 

77 "swiss-ai/Apertus-", 

78) 

79 

80 

81def _get_rope_theta(hf_config: Any, default: float = 10000.0) -> float | int: 

82 """Extract rope_theta from a HuggingFace config, handling both old and new formats. 

83 

84 In transformers v5+, rope_theta moved from a top-level attribute to 

85 hf_config.rope_parameters['rope_theta']. 

86 """ 

87 # Try direct attribute first (transformers < 5.0) 

88 rope_theta = getattr(hf_config, "rope_theta", None) 

89 if rope_theta is not None: 

90 return rope_theta 

91 # Try rope_parameters dict (transformers >= 5.0) 

92 rope_params = getattr(hf_config, "rope_parameters", None) 

93 if rope_params is not None and isinstance(rope_params, dict): 93 ↛ 95line 93 didn't jump to line 95 because the condition on line 93 was always true

94 return rope_params.get("rope_theta", default) 

95 return default 

96 

97 

98def make_model_alias_map() -> dict[str, str]: 

99 """ 

100 Converts OFFICIAL_MODEL_NAMES (the list of actual model names on 

101 HuggingFace) and MODEL_ALIASES (a dictionary mapping official model names to 

102 aliases) into a dictionary mapping all aliases to the official model name. 

103 """ 

104 model_alias_map = {} 

105 for official_model_name in OFFICIAL_MODEL_NAMES: 

106 aliases = MODEL_ALIASES.get(official_model_name, []) 

107 for alias in aliases: 

108 model_alias_map[alias.lower()] = official_model_name 

109 model_alias_map[official_model_name.lower()] = official_model_name 

110 return model_alias_map 

111 

112 

113def get_official_model_name(model_name: str) -> str: 

114 """ 

115 Returns the official model name for a given model name (or alias). 

116 """ 

117 model_alias_map = make_model_alias_map() 

118 official_model_name = model_alias_map.get(model_name.lower()) 

119 if official_model_name is None: 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true

120 raise ValueError( 

121 f"{model_name} not found. Valid official model names (excl aliases): {OFFICIAL_MODEL_NAMES}" 

122 ) 

123 return official_model_name 

124 

125 

126def convert_hf_model_config(model_name: str, **kwargs: Any) -> dict[str, Any]: 

127 """ 

128 Returns the model config for a HuggingFace model, converted to a dictionary 

129 in the HookedTransformerConfig format. 

130 

131 Takes the official_model_name as an input. 

132 """ 

133 # In case the user passed in an alias 

134 if (Path(model_name) / "config.json").exists(): 134 ↛ 135line 134 didn't jump to line 135 because the condition on line 134 was never true

135 logging.info("Loading model config from local directory") 

136 official_model_name = model_name 

137 else: 

138 official_model_name = get_official_model_name(model_name) 

139 

140 # Load HuggingFace model config 

141 if "llama" in official_model_name.lower(): 141 ↛ 142line 141 didn't jump to line 142 because the condition on line 141 was never true

142 architecture = "LlamaForCausalLM" 

143 elif "gemma-3" in official_model_name.lower() or "medgemma" in official_model_name.lower(): 

144 # Gemma 3: 270M and 1B are text-only (CausalLM), 4B+ are multimodal (ConditionalGeneration) 

145 # Exception: medgemma-27b-text-it is text-only 

146 if "270m" in official_model_name.lower() or "1b" in official_model_name.lower(): 

147 architecture = "Gemma3ForCausalLM" 

148 elif "medgemma-27b-text" in official_model_name.lower(): 

149 # medgemma-27b-text-it is text-only variant 

150 architecture = "Gemma3ForCausalLM" 

151 else: 

152 # 4B, 12B, 27B and medgemma are multimodal 

153 architecture = "Gemma3ForConditionalGeneration" 

154 elif "gemma-2-" in official_model_name.lower(): 154 ↛ 155line 154 didn't jump to line 155 because the condition on line 154 was never true

155 architecture = "Gemma2ForCausalLM" 

156 elif "gemma" in official_model_name.lower(): 156 ↛ 157line 156 didn't jump to line 157 because the condition on line 156 was never true

157 architecture = "GemmaForCausalLM" 

158 else: 

159 huggingface_token = os.environ.get("HF_TOKEN", "") 

160 hf_config = AutoConfig.from_pretrained( 

161 official_model_name, 

162 token=huggingface_token if len(huggingface_token) > 0 else None, 

163 **kwargs, 

164 ) 

165 architecture = hf_config.architectures[0] 

166 

167 cfg_dict: dict[str, Any] 

168 if official_model_name.startswith( 168 ↛ 171line 168 didn't jump to line 171 because the condition on line 168 was never true

169 ("llama-7b", "meta-llama/Llama-2-7b") 

170 ): # same architecture for LLaMA and Llama-2 

171 cfg_dict = { 

172 "d_model": 4096, 

173 "d_head": 4096 // 32, 

174 "n_heads": 32, 

175 "d_mlp": 11008, 

176 "n_layers": 32, 

177 "n_ctx": 2048 if official_model_name.startswith("llama-7b") else 4096, 

178 "eps": 1e-6 if official_model_name.startswith("llama-7b") else 1e-5, 

179 "d_vocab": 32000, 

180 "act_fn": "silu", 

181 "normalization_type": "RMS", 

182 "positional_embedding_type": "rotary", 

183 "rotary_adjacent_pairs": False, 

184 "rotary_dim": 4096 // 32, 

185 "final_rms": True, 

186 "gated_mlp": True, 

187 } 

188 elif official_model_name.startswith("codellama"): # same architecture CodeLlama and Llama-2 188 ↛ 189line 188 didn't jump to line 189 because the condition on line 188 was never true

189 cfg_dict = { 

190 "d_model": 4096, 

191 "d_head": 4096 // 32, 

192 "n_heads": 32, 

193 "d_mlp": 11008, 

194 "n_layers": 32, 

195 "n_ctx": 4096, 

196 "eps": 1e-5, 

197 "d_vocab": 32016, 

198 "act_fn": "silu", 

199 "normalization_type": "RMS", 

200 "positional_embedding_type": "rotary", 

201 "rotary_dim": 4096 // 32, 

202 "final_rms": True, 

203 "gated_mlp": True, 

204 "rotary_base": 1000000, 

205 } 

206 if "python" in official_model_name.lower(): 

207 # The vocab size of python version of CodeLlama-7b is 32000 

208 cfg_dict["d_vocab"] = 32000 

209 elif official_model_name.startswith( 209 ↛ 212line 209 didn't jump to line 212 because the condition on line 209 was never true

210 ("llama-13b", "meta-llama/Llama-2-13b") 

211 ): # same architecture for LLaMA and Llama-2 

212 cfg_dict = { 

213 "d_model": 5120, 

214 "d_head": 5120 // 40, 

215 "n_heads": 40, 

216 "d_mlp": 13824, 

217 "n_layers": 40, 

218 "n_ctx": 2048 if official_model_name.startswith("llama-13b") else 4096, 

219 "eps": 1e-6 if official_model_name.startswith("llama-13b") else 1e-5, 

220 "d_vocab": 32000, 

221 "act_fn": "silu", 

222 "normalization_type": "RMS", 

223 "positional_embedding_type": "rotary", 

224 "rotary_adjacent_pairs": False, 

225 "rotary_dim": 5120 // 40, 

226 "final_rms": True, 

227 "gated_mlp": True, 

228 } 

229 elif "llama-30b" in official_model_name: 229 ↛ 230line 229 didn't jump to line 230 because the condition on line 229 was never true

230 cfg_dict = { 

231 "d_model": 6656, 

232 "d_head": 6656 // 52, 

233 "n_heads": 52, 

234 "d_mlp": 17920, 

235 "n_layers": 60, 

236 "n_ctx": 2048, 

237 "eps": 1e-6, 

238 "d_vocab": 32000, 

239 "act_fn": "silu", 

240 "normalization_type": "RMS", 

241 "positional_embedding_type": "rotary", 

242 "rotary_adjacent_pairs": False, 

243 "rotary_dim": 6656 // 52, 

244 "final_rms": True, 

245 "gated_mlp": True, 

246 } 

247 elif "llama-65b" in official_model_name: 247 ↛ 248line 247 didn't jump to line 248 because the condition on line 247 was never true

248 cfg_dict = { 

249 "d_model": 8192, 

250 "d_head": 8192 // 64, 

251 "n_heads": 64, 

252 "d_mlp": 22016, 

253 "n_layers": 80, 

254 "n_ctx": 2048, 

255 "eps": 1e-6, 

256 "d_vocab": 32000, 

257 "act_fn": "silu", 

258 "normalization_type": "RMS", 

259 "positional_embedding_type": "rotary", 

260 "rotary_dim": 8192 // 64, 

261 "rotary_adjacent_pairs": False, 

262 "final_rms": True, 

263 "gated_mlp": True, 

264 } 

265 elif "Llama-2-70b" in official_model_name: 265 ↛ 266line 265 didn't jump to line 266 because the condition on line 265 was never true

266 cfg_dict = { 

267 "d_model": 8192, 

268 "d_head": 128, 

269 "n_heads": 64, 

270 "d_mlp": 28672, 

271 "n_layers": 80, 

272 "n_ctx": 4096, 

273 "eps": 1e-5, 

274 "d_vocab": 32000, 

275 "act_fn": "silu", 

276 "n_key_value_heads": 8, 

277 "normalization_type": "RMS", 

278 "positional_embedding_type": "rotary", 

279 "rotary_adjacent_pairs": False, 

280 "rotary_dim": 128, 

281 "final_rms": True, 

282 "gated_mlp": True, 

283 } 

284 elif "Meta-Llama-3-8B" in official_model_name: 284 ↛ 285line 284 didn't jump to line 285 because the condition on line 284 was never true

285 cfg_dict = { 

286 "d_model": 4096, 

287 "d_head": 128, 

288 "n_heads": 32, 

289 "d_mlp": 14336, 

290 "n_layers": 32, 

291 "n_ctx": 8192, 

292 "eps": 1e-5, 

293 "d_vocab": 128256, 

294 "act_fn": "silu", 

295 "n_key_value_heads": 8, 

296 "normalization_type": "RMS", 

297 "positional_embedding_type": "rotary", 

298 "rotary_adjacent_pairs": False, 

299 "rotary_dim": 128, 

300 "final_rms": True, 

301 "gated_mlp": True, 

302 "rotary_base": 500000.0, 

303 } 

304 elif "Meta-Llama-3-70B" in official_model_name: 304 ↛ 305line 304 didn't jump to line 305 because the condition on line 304 was never true

305 cfg_dict = { 

306 "d_model": 8192, 

307 "d_head": 128, 

308 "n_heads": 64, 

309 "d_mlp": 28672, 

310 "n_layers": 80, 

311 "n_ctx": 8192, 

312 "eps": 1e-5, 

313 "d_vocab": 128256, 

314 "act_fn": "silu", 

315 "n_key_value_heads": 8, 

316 "normalization_type": "RMS", 

317 "positional_embedding_type": "rotary", 

318 "rotary_adjacent_pairs": False, 

319 "rotary_dim": 128, 

320 "final_rms": True, 

321 "gated_mlp": True, 

322 "rotary_base": 500000.0, 

323 } 

324 elif "Llama-3.2-1B" in official_model_name: 324 ↛ 325line 324 didn't jump to line 325 because the condition on line 324 was never true

325 cfg_dict = { 

326 "d_model": 2048, 

327 "d_head": 64, 

328 "n_heads": 32, 

329 "d_mlp": 8192, 

330 "n_layers": 16, 

331 "n_ctx": 2048, # capped due to memory issues 

332 "eps": 1e-5, 

333 "d_vocab": 128256, 

334 "act_fn": "silu", 

335 "n_key_value_heads": 8, 

336 "normalization_type": "RMS", 

337 "positional_embedding_type": "rotary", 

338 "rotary_adjacent_pairs": False, 

339 "rotary_dim": 64, 

340 "final_rms": True, 

341 "gated_mlp": True, 

342 "rotary_base": 500000.0, 

343 "use_NTK_by_parts_rope": True, 

344 "NTK_by_parts_low_freq_factor": 1.0, 

345 "NTK_by_parts_high_freq_factor": 4.0, 

346 "NTK_by_parts_factor": 32.0, 

347 "NTK_original_ctx_len": 8192, 

348 } 

349 elif "Llama-3.2-3B" in official_model_name: 349 ↛ 350line 349 didn't jump to line 350 because the condition on line 349 was never true

350 cfg_dict = { 

351 "d_model": 3072, 

352 "d_head": 128, 

353 "n_heads": 24, 

354 "d_mlp": 8192, 

355 "n_layers": 28, 

356 "n_ctx": 2048, # capped due to memory issues 

357 "eps": 1e-5, 

358 "d_vocab": 128256, 

359 "act_fn": "silu", 

360 "n_key_value_heads": 8, 

361 "normalization_type": "RMS", 

362 "positional_embedding_type": "rotary", 

363 "rotary_adjacent_pairs": False, 

364 "rotary_dim": 128, 

365 "final_rms": True, 

366 "gated_mlp": True, 

367 "rotary_base": 500000.0, 

368 "use_NTK_by_parts_rope": True, 

369 "NTK_by_parts_low_freq_factor": 1.0, 

370 "NTK_by_parts_high_freq_factor": 4.0, 

371 "NTK_by_parts_factor": 32.0, 

372 "NTK_original_ctx_len": 8192, 

373 } 

374 elif "Llama-3.3-70B" in official_model_name: 374 ↛ 375line 374 didn't jump to line 375 because the condition on line 374 was never true

375 cfg_dict = { 

376 "d_model": 8192, 

377 "d_head": 128, 

378 "n_heads": 64, 

379 "d_mlp": 28672, 

380 "n_layers": 80, 

381 "n_ctx": 2048, # capped due to memory issues 

382 "eps": 1e-5, 

383 "d_vocab": 128256, 

384 "act_fn": "silu", 

385 "n_key_value_heads": 8, 

386 "normalization_type": "RMS", 

387 "positional_embedding_type": "rotary", 

388 "rotary_adjacent_pairs": False, 

389 "rotary_dim": 128, 

390 "final_rms": True, 

391 "gated_mlp": True, 

392 "rotary_base": 500000.0, 

393 "use_NTK_by_parts_rope": True, 

394 "NTK_by_parts_low_freq_factor": 1.0, 

395 "NTK_by_parts_high_freq_factor": 4.0, 

396 "NTK_by_parts_factor": 8.0, 

397 "NTK_original_ctx_len": 8192, 

398 } 

399 elif "Llama-3.1-8B" in official_model_name: 399 ↛ 400line 399 didn't jump to line 400 because the condition on line 399 was never true

400 cfg_dict = { 

401 "d_model": 4096, 

402 "d_head": 128, 

403 "n_heads": 32, 

404 "d_mlp": 14336, 

405 "n_layers": 32, 

406 "n_ctx": 2048, # capped due to memory issues 

407 "eps": 1e-5, 

408 "d_vocab": 128256, 

409 "act_fn": "silu", 

410 "n_key_value_heads": 8, 

411 "normalization_type": "RMS", 

412 "positional_embedding_type": "rotary", 

413 "rotary_adjacent_pairs": False, 

414 "rotary_dim": 128, 

415 "final_rms": True, 

416 "gated_mlp": True, 

417 "rotary_base": 500000.0, 

418 "use_NTK_by_parts_rope": True, 

419 "NTK_by_parts_low_freq_factor": 1.0, 

420 "NTK_by_parts_high_freq_factor": 4.0, 

421 "NTK_by_parts_factor": 8.0, 

422 "NTK_original_ctx_len": 8192, 

423 } 

424 elif "Llama-3.1-70B" in official_model_name: 424 ↛ 425line 424 didn't jump to line 425 because the condition on line 424 was never true

425 cfg_dict = { 

426 "d_model": 8192, 

427 "d_head": 128, 

428 "n_heads": 64, 

429 "d_mlp": 28672, 

430 "n_layers": 80, 

431 "n_ctx": 2048, # capped due to memory issues 

432 "eps": 1e-5, 

433 "d_vocab": 128256, 

434 "act_fn": "silu", 

435 "n_key_value_heads": 8, 

436 "normalization_type": "RMS", 

437 "positional_embedding_type": "rotary", 

438 "rotary_adjacent_pairs": False, 

439 "rotary_dim": 128, 

440 "final_rms": True, 

441 "gated_mlp": True, 

442 "rotary_base": 500000.0, 

443 "use_NTK_by_parts_rope": True, 

444 "NTK_by_parts_low_freq_factor": 1.0, 

445 "NTK_by_parts_high_freq_factor": 4.0, 

446 "NTK_by_parts_factor": 8.0, 

447 "NTK_original_ctx_len": 8192, 

448 } 

449 elif architecture == "GPTNeoForCausalLM": 

450 cfg_dict = { 

451 "d_model": hf_config.hidden_size, 

452 "d_head": hf_config.hidden_size // hf_config.num_heads, 

453 "n_heads": hf_config.num_heads, 

454 "d_mlp": hf_config.hidden_size * 4, 

455 "n_layers": hf_config.num_layers, 

456 "n_ctx": hf_config.max_position_embeddings, 

457 "eps": hf_config.layer_norm_epsilon, 

458 "d_vocab": hf_config.vocab_size, 

459 "attn_types": hf_config.attention_layers, 

460 "act_fn": hf_config.activation_function, 

461 "use_attn_scale": False, 

462 "use_local_attn": True, 

463 "window_size": hf_config.window_size, 

464 "scale_attn_by_inverse_layer_idx": False, 

465 "normalization_type": "LN", 

466 } 

467 elif architecture == "GPT2LMHeadModel": 

468 cfg_dict = { 

469 "d_model": hf_config.n_embd, 

470 "d_head": hf_config.n_embd // hf_config.n_head, 

471 "n_heads": hf_config.n_head, 

472 "d_mlp": hf_config.n_embd * 4, 

473 "n_layers": hf_config.n_layer, 

474 "n_ctx": hf_config.n_ctx, 

475 "eps": hf_config.layer_norm_epsilon, 

476 "d_vocab": hf_config.vocab_size, 

477 "act_fn": hf_config.activation_function, 

478 "use_attn_scale": True, 

479 "use_local_attn": False, 

480 "scale_attn_by_inverse_layer_idx": hf_config.scale_attn_by_inverse_layer_idx, 

481 "normalization_type": "LN", 

482 } 

483 elif architecture == "OPTForCausalLM": 

484 cfg_dict = { 

485 "d_model": hf_config.hidden_size, 

486 "d_head": hf_config.hidden_size // hf_config.num_attention_heads, 

487 "n_heads": hf_config.num_attention_heads, 

488 "d_mlp": hf_config.ffn_dim, 

489 "n_layers": hf_config.num_hidden_layers, 

490 "n_ctx": hf_config.max_position_embeddings, 

491 "eps": 1e-5, 

492 "d_vocab": hf_config.vocab_size, 

493 "act_fn": hf_config.activation_function, 

494 "use_attn_scale": True, 

495 "use_local_attn": False, 

496 "scale_attn_by_inverse_layer_idx": False, 

497 "normalization_type": "LN", 

498 } 

499 elif architecture == "GPTJForCausalLM": 

500 cfg_dict = { 

501 "d_model": hf_config.n_embd, 

502 "d_head": hf_config.n_embd // hf_config.n_head, 

503 "n_heads": hf_config.n_head, 

504 "d_mlp": 4 * hf_config.n_embd, 

505 "n_layers": hf_config.n_layer, 

506 "n_ctx": hf_config.n_positions, 

507 "eps": 1e-5, 

508 "d_vocab": hf_config.vocab_size, 

509 "act_fn": hf_config.activation_function, 

510 "use_attn_scale": True, 

511 "use_local_attn": False, 

512 "scale_attn_by_inverse_layer_idx": False, 

513 "parallel_attn_mlp": True, 

514 "positional_embedding_type": "rotary", 

515 "rotary_dim": hf_config.rotary_dim, 

516 "rotary_adjacent_pairs": True, 

517 "normalization_type": "LN", 

518 } 

519 elif architecture == "GPTNeoXForCausalLM": 

520 cfg_dict = { 

521 "d_model": hf_config.hidden_size, 

522 "d_head": hf_config.hidden_size // hf_config.num_attention_heads, 

523 "n_heads": hf_config.num_attention_heads, 

524 "d_mlp": hf_config.intermediate_size, 

525 "n_layers": hf_config.num_hidden_layers, 

526 "n_ctx": hf_config.max_position_embeddings, 

527 "eps": hf_config.layer_norm_eps, 

528 "d_vocab": hf_config.vocab_size, 

529 "act_fn": hf_config.hidden_act, 

530 "use_attn_scale": True, 

531 "use_local_attn": False, 

532 "scale_attn_by_inverse_layer_idx": False, 

533 "parallel_attn_mlp": True, 

534 "positional_embedding_type": "rotary", 

535 "rotary_adjacent_pairs": False, 

536 "normalization_type": "LN", 

537 "default_prepend_bos": False, 

538 } 

539 rotary_pct = get_rotary_pct_from_config(hf_config) 

540 cfg_dict["rotary_dim"] = round(rotary_pct * cfg_dict["d_head"]) 

541 elif architecture == "HubertModel": 

542 # Basic transformer configuration 

543 cfg_dict = { 

544 "d_model": hf_config.hidden_size, 

545 "d_head": hf_config.hidden_size // hf_config.num_attention_heads, 

546 "n_heads": hf_config.num_attention_heads, 

547 "d_mlp": hf_config.intermediate_size, 

548 "n_layers": hf_config.num_hidden_layers, 

549 # HuBERT operates on audio frames, not tokens — n_ctx is flexible 

550 "n_ctx": getattr(hf_config, "max_position_embeddings", 8192), 

551 "eps": hf_config.layer_norm_eps, 

552 "act_fn": getattr(hf_config, "hidden_act", "gelu"), 

553 "attention_dir": "bidirectional", 

554 "d_vocab": -1, # no text vocabulary 

555 } 

556 elif "wav2vec2-base" in official_model_name or "wav2vec2-large" in official_model_name: 556 ↛ 558line 556 didn't jump to line 558 because the condition on line 556 was never true

557 # Basic transformer configuration 

558 cfg_dict = { 

559 "d_model": hf_config.hidden_size, 

560 "d_head": hf_config.hidden_size // hf_config.num_attention_heads, 

561 "n_heads": hf_config.num_attention_heads, 

562 "d_mlp": hf_config.intermediate_size, 

563 "n_layers": hf_config.num_hidden_layers, 

564 # HuBERT operates on audio frames, not tokens — n_ctx is flexible 

565 "n_ctx": getattr(hf_config, "max_position_embeddings", 8192), 

566 "eps": hf_config.layer_norm_eps, 

567 "act_fn": getattr(hf_config, "hidden_act", "gelu"), 

568 "attention_dir": "bidirectional", 

569 "d_vocab": -1, # no text vocabulary 

570 } 

571 elif architecture == "HubertForCTC": 571 ↛ 573line 571 didn't jump to line 573 because the condition on line 571 was never true

572 # Basic transformer configuration 

573 cfg_dict = { 

574 "d_model": hf_config.hidden_size, 

575 "d_head": hf_config.hidden_size // hf_config.num_attention_heads, 

576 "n_heads": hf_config.num_attention_heads, 

577 "d_mlp": hf_config.intermediate_size, 

578 "n_layers": hf_config.num_hidden_layers, 

579 "n_ctx": getattr(hf_config, "max_position_embeddings", 8192), 

580 "eps": hf_config.layer_norm_eps, 

581 "act_fn": getattr(hf_config, "hidden_act", "gelu"), 

582 "attention_dir": "bidirectional", 

583 # For CTC models: 

584 "d_vocab": hf_config.vocab_size, # text vocab from tokenizer 

585 } 

586 elif architecture == "BertForMaskedLM": 586 ↛ 589line 586 didn't jump to line 589 because the condition on line 586 was never true

587 # All supported Bert architectures have the same config, 

588 # so we can use the BertForMaskedLM config for all of them 

589 cfg_dict = { 

590 "d_model": hf_config.hidden_size, 

591 "d_head": hf_config.hidden_size // hf_config.num_attention_heads, 

592 "n_heads": hf_config.num_attention_heads, 

593 "d_mlp": hf_config.intermediate_size, 

594 "n_layers": hf_config.num_hidden_layers, 

595 "n_ctx": hf_config.max_position_embeddings, 

596 "eps": hf_config.layer_norm_eps, 

597 "d_vocab": hf_config.vocab_size, 

598 "act_fn": "gelu", 

599 "attention_dir": "bidirectional", 

600 } 

601 elif architecture == "MistralForCausalLM": 601 ↛ 602line 601 didn't jump to line 602 because the condition on line 601 was never true

602 use_local_attn = True if hf_config.sliding_window else False 

603 cfg_dict = { 

604 "d_model": hf_config.hidden_size, 

605 "d_head": ( 

606 hf_config.head_dim 

607 if hasattr(hf_config, "head_dim") 

608 and hf_config.head_dim is not None 

609 and hf_config.head_dim > 0 

610 else hf_config.hidden_size // hf_config.num_attention_heads 

611 ), 

612 "n_heads": hf_config.num_attention_heads, 

613 "d_mlp": hf_config.intermediate_size, 

614 "n_layers": hf_config.num_hidden_layers, 

615 "n_ctx": 2048, # Capped due to memory issues 

616 "d_vocab": hf_config.vocab_size, 

617 "act_fn": hf_config.hidden_act, 

618 "window_size": hf_config.sliding_window, # None if no sliding window was used 

619 "attn_types": ["local"] * hf_config.num_hidden_layers if use_local_attn else None, 

620 "eps": hf_config.rms_norm_eps, 

621 "rotary_base": _get_rope_theta(hf_config), 

622 "n_key_value_heads": hf_config.num_key_value_heads, 

623 "use_local_attn": use_local_attn, 

624 "normalization_type": "RMS", 

625 "positional_embedding_type": "rotary", 

626 "gated_mlp": True, 

627 } 

628 elif architecture == "MixtralForCausalLM": 628 ↛ 629line 628 didn't jump to line 629 because the condition on line 628 was never true

629 cfg_dict = { 

630 "dtype": torch.bfloat16, 

631 "d_model": hf_config.hidden_size, 

632 "d_head": hf_config.hidden_size // hf_config.num_attention_heads, 

633 "n_heads": hf_config.num_attention_heads, 

634 "d_mlp": hf_config.intermediate_size, 

635 "n_layers": hf_config.num_hidden_layers, 

636 "n_ctx": hf_config.max_position_embeddings, # Capped due to memory issues 

637 "d_vocab": hf_config.vocab_size, 

638 "act_fn": hf_config.hidden_act, 

639 "normalization_type": "RMS", 

640 "positional_embedding_type": "rotary", 

641 "rotary_base": _get_rope_theta(hf_config), 

642 "window_size": hf_config.sliding_window, # This is None, as no sliding window was used 

643 "attn_types": ["global"] * 32, 

644 "eps": hf_config.rms_norm_eps, 

645 "n_key_value_heads": hf_config.num_key_value_heads, 

646 "gated_mlp": True, 

647 "use_local_attn": False, 

648 "rotary_dim": hf_config.hidden_size // hf_config.num_attention_heads, 

649 "num_experts": hf_config.num_local_experts, 

650 "experts_per_token": hf_config.num_experts_per_tok, 

651 } 

652 elif architecture == "GptOssForCausalLM": 

653 cfg_dict = { 

654 "dtype": torch.bfloat16, 

655 "d_model": hf_config.hidden_size, 

656 "d_head": hf_config.head_dim, 

657 "n_heads": hf_config.num_attention_heads, 

658 "d_mlp": hf_config.intermediate_size, 

659 "n_layers": hf_config.num_hidden_layers, 

660 "n_ctx": hf_config.max_position_embeddings, 

661 "d_vocab": hf_config.vocab_size, 

662 "act_fn": hf_config.hidden_act, 

663 "normalization_type": "RMS", 

664 "positional_embedding_type": "rotary", 

665 "rotary_base": _get_rope_theta(hf_config), 

666 "eps": hf_config.rms_norm_eps, 

667 "n_key_value_heads": hf_config.num_key_value_heads, 

668 "gated_mlp": True, 

669 "final_rms": True, 

670 "use_local_attn": False, 

671 "rotary_dim": hf_config.head_dim, 

672 "num_experts": hf_config.num_local_experts, 

673 "experts_per_token": hf_config.num_experts_per_tok, 

674 } 

675 elif architecture == "BloomForCausalLM": 675 ↛ 676line 675 didn't jump to line 676 because the condition on line 675 was never true

676 cfg_dict = { 

677 "d_model": hf_config.hidden_size, 

678 "d_head": hf_config.hidden_size // hf_config.n_head, 

679 "n_heads": hf_config.n_head, 

680 "d_mlp": hf_config.hidden_size * 4, 

681 "n_layers": hf_config.n_layer, 

682 "n_ctx": 2048, # Capped due to HF Tokenizer Constraints 

683 "d_vocab": hf_config.vocab_size, 

684 "act_fn": "gelu_fast", 

685 "eps": hf_config.layer_norm_epsilon, 

686 "normalization_type": "LN", 

687 "post_embedding_ln": True, 

688 "positional_embedding_type": "alibi", 

689 "default_prepend_bos": False, 

690 } 

691 elif architecture == "GPT2LMHeadCustomModel": 691 ↛ 693line 691 didn't jump to line 693 because the condition on line 691 was never true

692 # santacoder 

693 cfg_dict = { 

694 "d_model": hf_config.n_embd, 

695 "d_head": hf_config.n_embd // hf_config.n_head, 

696 "n_heads": hf_config.n_head, 

697 "d_mlp": hf_config.n_embd * 4, 

698 "n_layers": hf_config.n_layer, 

699 "n_ctx": hf_config.n_positions, 

700 "eps": hf_config.layer_norm_epsilon, 

701 "d_vocab": hf_config.vocab_size, 

702 "act_fn": hf_config.activation_function, 

703 "use_attn_scale": True, 

704 "use_local_attn": False, 

705 "trust_remote_code": "santacoder" 

706 in official_model_name, # Only santacoder needs trust_remote_code 

707 "scale_attn_by_inverse_layer_idx": hf_config.scale_attn_by_inverse_layer_idx, 

708 "normalization_type": "LN", 

709 } 

710 elif architecture == "LlamaForCausalLM": 710 ↛ 711line 710 didn't jump to line 711 because the condition on line 710 was never true

711 cfg_dict = { 

712 "d_model": hf_config.hidden_size, 

713 "d_head": hf_config.hidden_size // hf_config.num_attention_heads, 

714 "n_heads": hf_config.num_attention_heads, 

715 "d_mlp": hf_config.intermediate_size, 

716 "n_layers": hf_config.num_hidden_layers, 

717 "n_ctx": hf_config.max_position_embeddings, 

718 "eps": hf_config.rms_norm_eps, 

719 "d_vocab": hf_config.vocab_size, 

720 "act_fn": hf_config.hidden_act, 

721 "n_key_value_heads": ( 

722 hf_config.num_key_value_heads 

723 if hf_config.num_key_value_heads != hf_config.num_attention_heads 

724 else None 

725 ), 

726 # This is done because the current implementation of GQA will use Grouped-Query Attention if 

727 # n_key_value_heads is not None, but hf_config.num_key_value_heads is sometimes specified as 

728 # the same as hf_config.num_attention_heads, in which case GQA should not be used. 

729 "normalization_type": "RMS", 

730 "positional_embedding_type": "rotary", 

731 "rotary_adjacent_pairs": False, 

732 "rotary_dim": hf_config.hidden_size // hf_config.num_attention_heads, 

733 "final_rms": True, 

734 "gated_mlp": True, 

735 } 

736 elif architecture == "QWenLMHeadModel": 

737 cfg_dict = { 

738 "d_model": hf_config.hidden_size, 

739 "d_head": hf_config.hidden_size // hf_config.num_attention_heads, 

740 "n_heads": hf_config.num_attention_heads, 

741 "d_mlp": hf_config.intermediate_size // 2, 

742 "n_layers": hf_config.num_hidden_layers, 

743 # QWenLMHeadModel uses seq_length in its remote-code attention/rotary logic. 

744 "n_ctx": hf_config.seq_length, 

745 "eps": hf_config.layer_norm_epsilon, 

746 "d_vocab": hf_config.vocab_size, 

747 "act_fn": "silu", 

748 "use_attn_scale": hf_config.scale_attn_weights, 

749 "initializer_range": hf_config.initializer_range, 

750 "normalization_type": "RMS", 

751 "positional_embedding_type": "rotary", 

752 "rotary_dim": hf_config.kv_channels, 

753 "rotary_adjacent_pairs": False, 

754 "tokenizer_prepends_bos": True, 

755 "trust_remote_code": True, 

756 "final_rms": True, 

757 "gated_mlp": True, 

758 "default_prepend_bos": False, 

759 } 

760 elif architecture == "Qwen2ForCausalLM": 

761 # Note that Qwen1.5 models have architecture type Qwen2ForCausalLM. 

762 cfg_dict = { 

763 "d_model": hf_config.hidden_size, 

764 "d_head": hf_config.hidden_size // hf_config.num_attention_heads, 

765 "n_heads": hf_config.num_attention_heads, 

766 "n_key_value_heads": hf_config.num_key_value_heads, 

767 "d_mlp": hf_config.intermediate_size, 

768 "n_layers": hf_config.num_hidden_layers, 

769 "n_ctx": hf_config.max_position_embeddings, 

770 "eps": hf_config.rms_norm_eps, 

771 "d_vocab": hf_config.vocab_size, 

772 "act_fn": hf_config.hidden_act, 

773 "use_attn_scale": True, 

774 "initializer_range": hf_config.initializer_range, 

775 "normalization_type": "RMS", 

776 "positional_embedding_type": "rotary", 

777 "rotary_base": int(_get_rope_theta(hf_config)), 

778 "rotary_adjacent_pairs": False, 

779 "rotary_dim": hf_config.hidden_size // hf_config.num_attention_heads, 

780 "tokenizer_prepends_bos": True, 

781 "final_rms": True, 

782 "gated_mlp": True, 

783 "default_prepend_bos": False, 

784 } 

785 elif architecture == "Qwen3ForCausalLM": 785 ↛ 786line 785 didn't jump to line 786 because the condition on line 785 was never true

786 cfg_dict = { 

787 "d_model": hf_config.hidden_size, 

788 "d_head": hf_config.head_dim 

789 if hasattr(hf_config, "head_dim") 

790 and hf_config.head_dim is not None 

791 and hf_config.head_dim > 0 

792 else hf_config.hidden_size // hf_config.num_attention_heads, 

793 "n_heads": hf_config.num_attention_heads, 

794 "n_key_value_heads": ( 

795 hf_config.num_key_value_heads 

796 if hf_config.num_key_value_heads != hf_config.num_attention_heads 

797 else None 

798 ), 

799 "d_mlp": hf_config.intermediate_size, 

800 "n_layers": hf_config.num_hidden_layers, 

801 "n_ctx": 2048, 

802 "eps": hf_config.rms_norm_eps, 

803 "d_vocab": hf_config.vocab_size, 

804 "act_fn": hf_config.hidden_act, 

805 "use_attn_scale": True, 

806 "initializer_range": hf_config.initializer_range, 

807 "normalization_type": "RMS", 

808 "positional_embedding_type": "rotary", 

809 "rotary_base": int(_get_rope_theta(hf_config)), 

810 "rotary_adjacent_pairs": False, 

811 "rotary_dim": hf_config.head_dim 

812 if hasattr(hf_config, "head_dim") and hf_config.head_dim > 0 

813 else hf_config.hidden_size // hf_config.num_attention_heads, 

814 "tokenizer_prepends_bos": True, 

815 "final_rms": True, 

816 "gated_mlp": True, 

817 "default_prepend_bos": False, 

818 "use_qk_norm": True, 

819 "trust_remote_code": True, 

820 } 

821 elif architecture == "PhiForCausalLM": 821 ↛ 823line 821 didn't jump to line 823 because the condition on line 821 was never true

822 # Architecture for microsoft/phi models 

823 cfg_dict = { 

824 "d_model": hf_config.hidden_size, 

825 "d_head": hf_config.hidden_size // hf_config.num_attention_heads, 

826 "n_heads": hf_config.num_attention_heads, 

827 "d_mlp": hf_config.intermediate_size, 

828 "n_layers": hf_config.num_hidden_layers, 

829 "n_ctx": hf_config.max_position_embeddings, 

830 "eps": hf_config.layer_norm_eps, 

831 "d_vocab": hf_config.vocab_size, 

832 "act_fn": hf_config.hidden_act, 

833 "initializer_range": hf_config.initializer_range, 

834 "normalization_type": "LN", 

835 "positional_embedding_type": "rotary", 

836 "trust_remote_code": True,