Coverage for transformer_lens/loading_from_pretrained.py: 70%

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 math 

11import os 

12import re 

13from pathlib import Path 

14from typing import Any 

15 

16import torch 

17from huggingface_hub import HfApi 

18from transformers import ( 

19 AutoConfig, 

20 AutoModel, 

21 AutoModelForCausalLM, 

22 BertForPreTraining, 

23 HubertModel, 

24 T5ForConditionalGeneration, 

25 Wav2Vec2Model, 

26) 

27from transformers.utils.quantization_config import Mxfp4Config 

28 

29import transformer_lens.utilities as utils 

30from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig 

31from transformer_lens.pretrained.weight_conversions import ( 

32 convert_apertus_weights, 

33 convert_bert_weights, 

34 convert_bloom_weights, 

35 convert_coder_weights, 

36 convert_gemma_weights, 

37 convert_gpt2_weights, 

38 convert_gpt_oss_weights, 

39 convert_gptj_weights, 

40 convert_hubert_weights, 

41 convert_llama_weights, 

42 convert_mingpt_weights, 

43 convert_mistral_weights, 

44 convert_mixtral_weights, 

45 convert_neel_solu_old_weights, 

46 convert_neo_weights, 

47 convert_neox_weights, 

48 convert_olmo2_weights, 

49 convert_olmo3_weights, 

50 convert_olmo_weights, 

51 convert_olmoe_weights, 

52 convert_opt_weights, 

53 convert_phi3_weights, 

54 convert_phi_weights, 

55 convert_qwen2_weights, 

56 convert_qwen3_weights, 

57 convert_qwen_weights, 

58 convert_t5_weights, 

59) 

60from transformer_lens.supported_models import MODEL_ALIASES, OFFICIAL_MODEL_NAMES 

61from transformer_lens.utilities.architectures import POST_NORM_ARCHITECTURES 

62from transformer_lens.utilities.heterogeneous_config import het_safe_view 

63from transformer_lens.utilities.hf_utils import get_rotary_pct_from_config 

64from transformer_lens.utilities.quantization import ( 

65 quantization_method, 

66 unreadable_weight_reason, 

67) 

68 

69NON_HF_HOSTED_MODEL_NAMES = [ 

70 "llama-7b-hf", 

71 "llama-13b-hf", 

72 "llama-30b-hf", 

73 "llama-65b-hf", 

74] 

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

76 

77NEED_REMOTE_CODE_MODELS = ( 

78 "bigcode/santacoder", 

79 "Qwen/Qwen-", 

80 "Qwen/Qwen3-", 

81 "microsoft/phi-2", 

82 "microsoft/phi-4", 

83 "apple/OpenELM", 

84 "openai/gpt-oss-", 

85 "swiss-ai/Apertus-", 

86) 

87 

88 

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

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

91 

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

93 hf_config.rope_parameters['rope_theta']. 

94 """ 

95 # Try direct attribute first (transformers < 5.0) 

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

97 if rope_theta is not None: 

98 return rope_theta 

99 # Try rope_parameters dict (transformers >= 5.0) 

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

101 if rope_params is not None and isinstance(rope_params, dict): 

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

103 return default 

104 

105 

106def _apply_llama3_rope_scaling(cfg_dict: dict[str, Any], hf_config: Any) -> None: 

107 """Populate the NTK-by-parts fields when a config requests llama3 rope scaling.""" 

108 rope_scaling = getattr(hf_config, "rope_scaling", None) 

109 if not rope_scaling: 

110 return 

111 rope_type = (rope_scaling.get("type") or rope_scaling.get("rope_type") or "").lower() 

112 if rope_type != "llama3": 112 ↛ 113line 112 didn't jump to line 113 because the condition on line 112 was never true

113 return 

114 cfg_dict["use_NTK_by_parts_rope"] = True 

115 cfg_dict["NTK_original_ctx_len"] = rope_scaling.get( 

116 "original_max_position_embeddings", hf_config.max_position_embeddings 

117 ) 

118 cfg_dict["NTK_by_parts_low_freq_factor"] = rope_scaling.get("low_freq_factor", 1.0) 

119 cfg_dict["NTK_by_parts_high_freq_factor"] = rope_scaling.get("high_freq_factor", 4.0) 

120 cfg_dict["NTK_by_parts_factor"] = rope_scaling.get("factor", 1.0) 

121 

122 

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

124 """ 

125 Converts OFFICIAL_MODEL_NAMES (the list of actual model names on 

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

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

128 """ 

129 model_alias_map = {} 

130 for official_model_name in OFFICIAL_MODEL_NAMES: 

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

132 for alias in aliases: 

133 model_alias_map[alias.lower()] = official_model_name 

134 model_alias_map[official_model_name.lower()] = official_model_name 

135 return model_alias_map 

136 

137 

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

139 """ 

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

141 """ 

142 model_alias_map = make_model_alias_map() 

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

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

145 raise ValueError( 

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

147 ) 

148 return official_model_name 

149 

150 

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

152 """ 

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

154 in the HookedTransformerConfig format. 

155 

156 Takes the official_model_name as an input. 

157 """ 

158 # In case the user passed in an alias 

159 if (Path(model_name) / "config.json").exists(): 

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

161 official_model_name = model_name 

162 else: 

163 official_model_name = get_official_model_name(model_name) 

164 

165 # Load HuggingFace model config. Stays None on the name-based branches 

166 # below, which infer the architecture from the model name and never fetch. 

167 hf_config: Any = None 

168 if "llama" in official_model_name.lower(): 

169 architecture = "LlamaForCausalLM" 

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

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

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

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

174 architecture = "Gemma3ForCausalLM" 

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

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

177 architecture = "Gemma3ForCausalLM" 

178 else: 

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

180 architecture = "Gemma3ForConditionalGeneration" 

181 elif "gemma-2-" in official_model_name.lower(): 

182 architecture = "Gemma2ForCausalLM" 

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

184 architecture = "GemmaForCausalLM" 

185 else: 

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

187 hf_config = AutoConfig.from_pretrained( 

188 official_model_name, 

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

190 **kwargs, 

191 ) 

192 # het view: transformers>=5.15 per-layer fields raise (not 

193 # AttributeError) on global reads, so every hasattr/getattr probe in 

194 # the branch chain below is a crash site without it. 

195 hf_config = het_safe_view(hf_config) 

196 architecture = hf_config.architectures[0] 

197 

198 cfg_dict: dict[str, Any] 

199 if official_model_name.startswith( 

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

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

202 cfg_dict = { 

203 "d_model": 4096, 

204 "d_head": 4096 // 32, 

205 "n_heads": 32, 

206 "d_mlp": 11008, 

207 "n_layers": 32, 

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

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

210 "d_vocab": 32000, 

211 "act_fn": "silu", 

212 "normalization_type": "RMS", 

213 "positional_embedding_type": "rotary", 

214 "rotary_adjacent_pairs": False, 

215 "rotary_dim": 4096 // 32, 

216 "final_rms": True, 

217 "gated_mlp": True, 

218 } 

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

220 cfg_dict = { 

221 "d_model": 4096, 

222 "d_head": 4096 // 32, 

223 "n_heads": 32, 

224 "d_mlp": 11008, 

225 "n_layers": 32, 

226 "n_ctx": 4096, 

227 "eps": 1e-5, 

228 "d_vocab": 32016, 

229 "act_fn": "silu", 

230 "normalization_type": "RMS", 

231 "positional_embedding_type": "rotary", 

232 "rotary_dim": 4096 // 32, 

233 "final_rms": True, 

234 "gated_mlp": True, 

235 "rotary_base": 1000000, 

236 } 

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

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

239 cfg_dict["d_vocab"] = 32000 

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

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

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

243 cfg_dict = { 

244 "d_model": 5120, 

245 "d_head": 5120 // 40, 

246 "n_heads": 40, 

247 "d_mlp": 13824, 

248 "n_layers": 40, 

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

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

251 "d_vocab": 32000, 

252 "act_fn": "silu", 

253 "normalization_type": "RMS", 

254 "positional_embedding_type": "rotary", 

255 "rotary_adjacent_pairs": False, 

256 "rotary_dim": 5120 // 40, 

257 "final_rms": True, 

258 "gated_mlp": True, 

259 } 

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

261 cfg_dict = { 

262 "d_model": 6656, 

263 "d_head": 6656 // 52, 

264 "n_heads": 52, 

265 "d_mlp": 17920, 

266 "n_layers": 60, 

267 "n_ctx": 2048, 

268 "eps": 1e-6, 

269 "d_vocab": 32000, 

270 "act_fn": "silu", 

271 "normalization_type": "RMS", 

272 "positional_embedding_type": "rotary", 

273 "rotary_adjacent_pairs": False, 

274 "rotary_dim": 6656 // 52, 

275 "final_rms": True, 

276 "gated_mlp": True, 

277 } 

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

279 cfg_dict = { 

280 "d_model": 8192, 

281 "d_head": 8192 // 64, 

282 "n_heads": 64, 

283 "d_mlp": 22016, 

284 "n_layers": 80, 

285 "n_ctx": 2048, 

286 "eps": 1e-6, 

287 "d_vocab": 32000, 

288 "act_fn": "silu", 

289 "normalization_type": "RMS", 

290 "positional_embedding_type": "rotary", 

291 "rotary_dim": 8192 // 64, 

292 "rotary_adjacent_pairs": False, 

293 "final_rms": True, 

294 "gated_mlp": True, 

295 } 

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

297 cfg_dict = { 

298 "d_model": 8192, 

299 "d_head": 128, 

300 "n_heads": 64, 

301 "d_mlp": 28672, 

302 "n_layers": 80, 

303 "n_ctx": 4096, 

304 "eps": 1e-5, 

305 "d_vocab": 32000, 

306 "act_fn": "silu", 

307 "n_key_value_heads": 8, 

308 "normalization_type": "RMS", 

309 "positional_embedding_type": "rotary", 

310 "rotary_adjacent_pairs": False, 

311 "rotary_dim": 128, 

312 "final_rms": True, 

313 "gated_mlp": True, 

314 } 

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

316 cfg_dict = { 

317 "d_model": 4096, 

318 "d_head": 128, 

319 "n_heads": 32, 

320 "d_mlp": 14336, 

321 "n_layers": 32, 

322 "n_ctx": 8192, 

323 "eps": 1e-5, 

324 "d_vocab": 128256, 

325 "act_fn": "silu", 

326 "n_key_value_heads": 8, 

327 "normalization_type": "RMS", 

328 "positional_embedding_type": "rotary", 

329 "rotary_adjacent_pairs": False, 

330 "rotary_dim": 128, 

331 "final_rms": True, 

332 "gated_mlp": True, 

333 "rotary_base": 500000.0, 

334 } 

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

336 cfg_dict = { 

337 "d_model": 8192, 

338 "d_head": 128, 

339 "n_heads": 64, 

340 "d_mlp": 28672, 

341 "n_layers": 80, 

342 "n_ctx": 8192, 

343 "eps": 1e-5, 

344 "d_vocab": 128256, 

345 "act_fn": "silu", 

346 "n_key_value_heads": 8, 

347 "normalization_type": "RMS", 

348 "positional_embedding_type": "rotary", 

349 "rotary_adjacent_pairs": False, 

350 "rotary_dim": 128, 

351 "final_rms": True, 

352 "gated_mlp": True, 

353 "rotary_base": 500000.0, 

354 } 

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

356 cfg_dict = { 

357 "d_model": 2048, 

358 "d_head": 64, 

359 "n_heads": 32, 

360 "d_mlp": 8192, 

361 "n_layers": 16, 

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

363 "eps": 1e-5, 

364 "d_vocab": 128256, 

365 "act_fn": "silu", 

366 "n_key_value_heads": 8, 

367 "normalization_type": "RMS", 

368 "positional_embedding_type": "rotary", 

369 "rotary_adjacent_pairs": False, 

370 "rotary_dim": 64, 

371 "final_rms": True, 

372 "gated_mlp": True, 

373 "rotary_base": 500000.0, 

374 "use_NTK_by_parts_rope": True, 

375 "NTK_by_parts_low_freq_factor": 1.0, 

376 "NTK_by_parts_high_freq_factor": 4.0, 

377 "NTK_by_parts_factor": 32.0, 

378 "NTK_original_ctx_len": 8192, 

379 } 

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

381 cfg_dict = { 

382 "d_model": 3072, 

383 "d_head": 128, 

384 "n_heads": 24, 

385 "d_mlp": 8192, 

386 "n_layers": 28, 

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

388 "eps": 1e-5, 

389 "d_vocab": 128256, 

390 "act_fn": "silu", 

391 "n_key_value_heads": 8, 

392 "normalization_type": "RMS", 

393 "positional_embedding_type": "rotary", 

394 "rotary_adjacent_pairs": False, 

395 "rotary_dim": 128, 

396 "final_rms": True, 

397 "gated_mlp": True, 

398 "rotary_base": 500000.0, 

399 "use_NTK_by_parts_rope": True, 

400 "NTK_by_parts_low_freq_factor": 1.0, 

401 "NTK_by_parts_high_freq_factor": 4.0, 

402 "NTK_by_parts_factor": 32.0, 

403 "NTK_original_ctx_len": 8192, 

404 } 

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

406 cfg_dict = { 

407 "d_model": 8192, 

408 "d_head": 128, 

409 "n_heads": 64, 

410 "d_mlp": 28672, 

411 "n_layers": 80, 

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

413 "eps": 1e-5, 

414 "d_vocab": 128256, 

415 "act_fn": "silu", 

416 "n_key_value_heads": 8, 

417 "normalization_type": "RMS", 

418 "positional_embedding_type": "rotary", 

419 "rotary_adjacent_pairs": False, 

420 "rotary_dim": 128, 

421 "final_rms": True, 

422 "gated_mlp": True, 

423 "rotary_base": 500000.0, 

424 "use_NTK_by_parts_rope": True, 

425 "NTK_by_parts_low_freq_factor": 1.0, 

426 "NTK_by_parts_high_freq_factor": 4.0, 

427 "NTK_by_parts_factor": 8.0, 

428 "NTK_original_ctx_len": 8192, 

429 } 

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

431 cfg_dict = { 

432 "d_model": 4096, 

433 "d_head": 128, 

434 "n_heads": 32, 

435 "d_mlp": 14336, 

436 "n_layers": 32, 

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

438 "eps": 1e-5, 

439 "d_vocab": 128256, 

440 "act_fn": "silu", 

441 "n_key_value_heads": 8, 

442 "normalization_type": "RMS", 

443 "positional_embedding_type": "rotary", 

444 "rotary_adjacent_pairs": False, 

445 "rotary_dim": 128, 

446 "final_rms": True, 

447 "gated_mlp": True, 

448 "rotary_base": 500000.0, 

449 "use_NTK_by_parts_rope": True, 

450 "NTK_by_parts_low_freq_factor": 1.0, 

451 "NTK_by_parts_high_freq_factor": 4.0, 

452 "NTK_by_parts_factor": 8.0, 

453 "NTK_original_ctx_len": 8192, 

454 } 

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

456 cfg_dict = { 

457 "d_model": 8192, 

458 "d_head": 128, 

459 "n_heads": 64, 

460 "d_mlp": 28672, 

461 "n_layers": 80, 

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

463 "eps": 1e-5, 

464 "d_vocab": 128256, 

465 "act_fn": "silu", 

466 "n_key_value_heads": 8, 

467 "normalization_type": "RMS", 

468 "positional_embedding_type": "rotary", 

469 "rotary_adjacent_pairs": False, 

470 "rotary_dim": 128, 

471 "final_rms": True, 

472 "gated_mlp": True, 

473 "rotary_base": 500000.0, 

474 "use_NTK_by_parts_rope": True, 

475 "NTK_by_parts_low_freq_factor": 1.0, 

476 "NTK_by_parts_high_freq_factor": 4.0, 

477 "NTK_by_parts_factor": 8.0, 

478 "NTK_original_ctx_len": 8192, 

479 } 

480 elif architecture == "GPTNeoForCausalLM": 

481 cfg_dict = { 

482 "d_model": hf_config.hidden_size, 

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

484 "n_heads": hf_config.num_heads, 

485 "d_mlp": hf_config.hidden_size * 4, 

486 "n_layers": hf_config.num_layers, 

487 "n_ctx": hf_config.max_position_embeddings, 

488 "eps": hf_config.layer_norm_epsilon, 

489 "d_vocab": hf_config.vocab_size, 

490 "attn_types": hf_config.attention_layers, 

491 "act_fn": hf_config.activation_function, 

492 "use_attn_scale": False, 

493 "use_local_attn": True, 

494 "window_size": hf_config.window_size, 

495 "scale_attn_by_inverse_layer_idx": False, 

496 "normalization_type": "LN", 

497 } 

498 elif architecture == "GPT2LMHeadModel": 

499 cfg_dict = { 

500 "d_model": hf_config.n_embd, 

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

502 "n_heads": hf_config.n_head, 

503 "d_mlp": hf_config.n_embd * 4, 

504 "n_layers": hf_config.n_layer, 

505 "n_ctx": hf_config.n_ctx, 

506 "eps": hf_config.layer_norm_epsilon, 

507 "d_vocab": hf_config.vocab_size, 

508 "act_fn": hf_config.activation_function, 

509 "use_attn_scale": getattr(hf_config, "scale_attn_weights", True), 

510 "use_local_attn": False, 

511 "scale_attn_by_inverse_layer_idx": hf_config.scale_attn_by_inverse_layer_idx, 

512 "normalization_type": "LN", 

513 } 

514 elif architecture == "OPTForCausalLM": 

515 cfg_dict = { 

516 "d_model": hf_config.hidden_size, 

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

518 "n_heads": hf_config.num_attention_heads, 

519 "d_mlp": hf_config.ffn_dim, 

520 "n_layers": hf_config.num_hidden_layers, 

521 "n_ctx": hf_config.max_position_embeddings, 

522 "eps": 1e-5, 

523 "d_vocab": hf_config.vocab_size, 

524 "act_fn": hf_config.activation_function, 

525 "use_attn_scale": True, 

526 "use_local_attn": False, 

527 "scale_attn_by_inverse_layer_idx": False, 

528 "normalization_type": "LN", 

529 } 

530 elif architecture == "GPTJForCausalLM": 

531 cfg_dict = { 

532 "d_model": hf_config.n_embd, 

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

534 "n_heads": hf_config.n_head, 

535 "d_mlp": 4 * hf_config.n_embd, 

536 "n_layers": hf_config.n_layer, 

537 "n_ctx": hf_config.n_positions, 

538 "eps": 1e-5, 

539 "d_vocab": hf_config.vocab_size, 

540 "act_fn": hf_config.activation_function, 

541 "use_attn_scale": True, 

542 "use_local_attn": False, 

543 "scale_attn_by_inverse_layer_idx": False, 

544 "parallel_attn_mlp": True, 

545 "positional_embedding_type": "rotary", 

546 "rotary_dim": hf_config.rotary_dim, 

547 "rotary_adjacent_pairs": True, 

548 "normalization_type": "LN", 

549 } 

550 elif architecture == "GPTNeoXForCausalLM": 

551 cfg_dict = { 

552 "d_model": hf_config.hidden_size, 

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

554 "n_heads": hf_config.num_attention_heads, 

555 "d_mlp": hf_config.intermediate_size, 

556 "n_layers": hf_config.num_hidden_layers, 

557 "n_ctx": hf_config.max_position_embeddings, 

558 "eps": hf_config.layer_norm_eps, 

559 "d_vocab": hf_config.vocab_size, 

560 "act_fn": hf_config.hidden_act, 

561 "use_attn_scale": True, 

562 "use_local_attn": False, 

563 "scale_attn_by_inverse_layer_idx": False, 

564 # GPTNeoX ships sequential variants too (use_parallel_residual=False). 

565 "parallel_attn_mlp": getattr(hf_config, "use_parallel_residual", True), 

566 "positional_embedding_type": "rotary", 

567 "rotary_adjacent_pairs": False, 

568 "normalization_type": "LN", 

569 "default_prepend_bos": False, 

570 } 

571 rotary_pct = get_rotary_pct_from_config(hf_config) 

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

573 elif architecture == "HubertModel": 

574 # Basic transformer configuration 

575 cfg_dict = { 

576 "d_model": hf_config.hidden_size, 

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

578 "n_heads": hf_config.num_attention_heads, 

579 "d_mlp": hf_config.intermediate_size, 

580 "n_layers": hf_config.num_hidden_layers, 

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

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

583 "eps": hf_config.layer_norm_eps, 

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

585 "attention_dir": "bidirectional", 

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

587 } 

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

589 # Basic transformer configuration 

590 cfg_dict = { 

591 "d_model": hf_config.hidden_size, 

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

593 "n_heads": hf_config.num_attention_heads, 

594 "d_mlp": hf_config.intermediate_size, 

595 "n_layers": hf_config.num_hidden_layers, 

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

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

598 "eps": hf_config.layer_norm_eps, 

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

600 "attention_dir": "bidirectional", 

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

602 } 

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

604 # Basic transformer configuration 

605 cfg_dict = { 

606 "d_model": hf_config.hidden_size, 

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

608 "n_heads": hf_config.num_attention_heads, 

609 "d_mlp": hf_config.intermediate_size, 

610 "n_layers": hf_config.num_hidden_layers, 

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

612 "eps": hf_config.layer_norm_eps, 

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

614 "attention_dir": "bidirectional", 

615 # For CTC models: 

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

617 } 

618 elif architecture == "BertForMaskedLM": 

619 # All supported Bert architectures have the same config, 

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

621 cfg_dict = { 

622 "d_model": hf_config.hidden_size, 

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

624 "n_heads": hf_config.num_attention_heads, 

625 "d_mlp": hf_config.intermediate_size, 

626 "n_layers": hf_config.num_hidden_layers, 

627 "n_ctx": hf_config.max_position_embeddings, 

628 "eps": hf_config.layer_norm_eps, 

629 "d_vocab": hf_config.vocab_size, 

630 "act_fn": "gelu", 

631 "attention_dir": "bidirectional", 

632 } 

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

634 use_local_attn = True if hf_config.sliding_window else False 

635 cfg_dict = { 

636 "d_model": hf_config.hidden_size, 

637 "d_head": ( 

638 hf_config.head_dim 

639 if hasattr(hf_config, "head_dim") 

640 and hf_config.head_dim is not None 

641 and hf_config.head_dim > 0 

642 else hf_config.hidden_size // hf_config.num_attention_heads 

643 ), 

644 "n_heads": hf_config.num_attention_heads, 

645 "d_mlp": hf_config.intermediate_size, 

646 "n_layers": hf_config.num_hidden_layers, 

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

648 "d_vocab": hf_config.vocab_size, 

649 "act_fn": hf_config.hidden_act, 

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

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

652 "eps": hf_config.rms_norm_eps, 

653 "rotary_base": _get_rope_theta(hf_config), 

654 "n_key_value_heads": hf_config.num_key_value_heads, 

655 "use_local_attn": use_local_attn, 

656 "normalization_type": "RMS", 

657 "positional_embedding_type": "rotary", 

658 "gated_mlp": True, 

659 } 

660 elif architecture == "MixtralForCausalLM": 

661 cfg_dict = { 

662 "dtype": torch.bfloat16, 

663 "d_model": hf_config.hidden_size, 

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

665 "n_heads": hf_config.num_attention_heads, 

666 "d_mlp": hf_config.intermediate_size, 

667 "n_layers": hf_config.num_hidden_layers, 

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

669 "d_vocab": hf_config.vocab_size, 

670 "act_fn": hf_config.hidden_act, 

671 "normalization_type": "RMS", 

672 "positional_embedding_type": "rotary", 

673 "rotary_base": _get_rope_theta(hf_config), 

674 # None on the released 8x7B, but a variant that sets it must window. 

675 "window_size": hf_config.sliding_window, 

676 "attn_types": ( 

677 ["local" if hf_config.sliding_window else "global"] * hf_config.num_hidden_layers 

678 ), 

679 "eps": hf_config.rms_norm_eps, 

680 "n_key_value_heads": hf_config.num_key_value_heads, 

681 "gated_mlp": True, 

682 "use_local_attn": bool(hf_config.sliding_window), 

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

684 "num_experts": hf_config.num_local_experts, 

685 "experts_per_token": hf_config.num_experts_per_tok, 

686 # MixtralTopKRouter renormalizes the top-k weights unconditionally 

687 # (modeling_mixtral.py: `router_top_value /= router_top_value.sum(...)`), 

688 # and MixtralConfig has no norm_topk_prob field to read it from — so 

689 # this is pinned to HF's behavior rather than sourced from the config. 

690 # TL's MoE skips the renormalization unless this is set, which would 

691 # leave routing weights unnormalized and the outputs silently wrong. 

692 "norm_topk_prob": True, 

693 } 

694 elif architecture == "GptOssForCausalLM": 

695 cfg_dict = { 

696 "dtype": torch.bfloat16, 

697 "d_model": hf_config.hidden_size, 

698 "d_head": hf_config.head_dim, 

699 "n_heads": hf_config.num_attention_heads, 

700 "d_mlp": hf_config.intermediate_size, 

701 "n_layers": hf_config.num_hidden_layers, 

702 "n_ctx": hf_config.max_position_embeddings, 

703 "d_vocab": hf_config.vocab_size, 

704 "act_fn": hf_config.hidden_act, 

705 "normalization_type": "RMS", 

706 "positional_embedding_type": "rotary", 

707 "rotary_base": _get_rope_theta(hf_config), 

708 "eps": hf_config.rms_norm_eps, 

709 "n_key_value_heads": hf_config.num_key_value_heads, 

710 "gated_mlp": True, 

711 "final_rms": True, 

712 "rotary_dim": hf_config.head_dim, 

713 "num_experts": hf_config.num_local_experts, 

714 "experts_per_token": hf_config.num_experts_per_tok, 

715 "use_attention_sinks": True, 

716 # Alternating sliding_attention / full_attention layers; HT's local 

717 # attention mask (last window_size keys) matches HF's sliding window. 

718 "use_local_attn": True, 

719 "window_size": hf_config.sliding_window, 

720 "attn_types": [ 

721 "local" if layer_type == "sliding_attention" else "global" 

722 for layer_type in hf_config.layer_types 

723 ], 

724 } 

725 rope_params = getattr(hf_config, "rope_parameters", None) or {} 

726 if rope_params.get("rope_type") == "yarn": 726 ↛ 1642line 726 didn't jump to line 1642 because the condition on line 726 was always true

727 yarn_factor = rope_params["factor"] 

728 attention_factor = rope_params.get("attention_factor") 

729 if attention_factor is None: 729 ↛ 732line 729 didn't jump to line 732 because the condition on line 729 was always true

730 # HF's default: get_mscale(factor) = 0.1 * ln(factor) + 1 

731 attention_factor = 0.1 * math.log(yarn_factor) + 1.0 if yarn_factor > 1 else 1.0 

732 cfg_dict.update( 

733 { 

734 "use_yarn_rope": True, 

735 "yarn_factor": float(yarn_factor), 

736 "yarn_attention_factor": float(attention_factor), 

737 "yarn_beta_fast": float(rope_params.get("beta_fast") or 32.0), 

738 "yarn_beta_slow": float(rope_params.get("beta_slow") or 1.0), 

739 "yarn_original_max_position_embeddings": rope_params[ 

740 "original_max_position_embeddings" 

741 ], 

742 "yarn_truncate": rope_params.get("truncate", True), 

743 } 

744 ) 

745 elif architecture == "BloomForCausalLM": 

746 cfg_dict = { 

747 "d_model": hf_config.hidden_size, 

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

749 "n_heads": hf_config.n_head, 

750 "d_mlp": hf_config.hidden_size * 4, 

751 "n_layers": hf_config.n_layer, 

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

753 "d_vocab": hf_config.vocab_size, 

754 "act_fn": "gelu_fast", 

755 "eps": hf_config.layer_norm_epsilon, 

756 "normalization_type": "LN", 

757 "post_embedding_ln": True, 

758 "positional_embedding_type": "alibi", 

759 "default_prepend_bos": False, 

760 } 

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

762 # santacoder 

763 cfg_dict = { 

764 "d_model": hf_config.n_embd, 

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

766 "n_heads": hf_config.n_head, 

767 "d_mlp": hf_config.n_embd * 4, 

768 "n_layers": hf_config.n_layer, 

769 "n_ctx": hf_config.n_positions, 

770 "eps": hf_config.layer_norm_epsilon, 

771 "d_vocab": hf_config.vocab_size, 

772 "act_fn": hf_config.activation_function, 

773 "use_attn_scale": True, 

774 "use_local_attn": False, 

775 "trust_remote_code": "santacoder" 

776 in official_model_name, # Only santacoder needs trust_remote_code 

777 "scale_attn_by_inverse_layer_idx": hf_config.scale_attn_by_inverse_layer_idx, 

778 "normalization_type": "LN", 

779 } 

780 elif architecture == "LlamaForCausalLM": 

781 cfg_dict = { 

782 "d_model": hf_config.hidden_size, 

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

784 "n_heads": hf_config.num_attention_heads, 

785 "d_mlp": hf_config.intermediate_size, 

786 "n_layers": hf_config.num_hidden_layers, 

787 "n_ctx": hf_config.max_position_embeddings, 

788 "eps": hf_config.rms_norm_eps, 

789 "d_vocab": hf_config.vocab_size, 

790 "act_fn": hf_config.hidden_act, 

791 "n_key_value_heads": ( 

792 hf_config.num_key_value_heads 

793 if hf_config.num_key_value_heads != hf_config.num_attention_heads 

794 else None 

795 ), 

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

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

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

799 "normalization_type": "RMS", 

800 "positional_embedding_type": "rotary", 

801 "rotary_adjacent_pairs": False, 

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

803 # Llama-arch checkpoints without a name-matched branch above (Yi ships 

804 # 5e6) reach here; the config default would silently be 10000. 

805 "rotary_base": _get_rope_theta(hf_config), 

806 "final_rms": True, 

807 "gated_mlp": True, 

808 } 

809 _apply_llama3_rope_scaling(cfg_dict, hf_config) 

810 elif architecture == "QWenLMHeadModel": 

811 cfg_dict = { 

812 "d_model": hf_config.hidden_size, 

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

814 "n_heads": hf_config.num_attention_heads, 

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

816 "n_layers": hf_config.num_hidden_layers, 

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

818 "n_ctx": hf_config.seq_length, 

819 "use_logn_attn": getattr(hf_config, "use_logn_attn", False), 

820 "use_dynamic_ntk_rope": getattr(hf_config, "use_dynamic_ntk", False), 

821 "train_seq_length": hf_config.seq_length, 

822 "eps": hf_config.layer_norm_epsilon, 

823 "d_vocab": hf_config.vocab_size, 

824 "act_fn": "silu", 

825 "use_attn_scale": hf_config.scale_attn_weights, 

826 "initializer_range": hf_config.initializer_range, 

827 "normalization_type": "RMS", 

828 "positional_embedding_type": "rotary", 

829 "rotary_dim": hf_config.kv_channels, 

830 "rotary_adjacent_pairs": False, 

831 "tokenizer_prepends_bos": True, 

832 "trust_remote_code": True, 

833 "final_rms": True, 

834 "gated_mlp": True, 

835 "default_prepend_bos": False, 

836 }