tensor_mapping.py 50 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173
  1. from __future__ import annotations
  2. from typing import Sequence
  3. from .constants import MODEL_ARCH, MODEL_TENSOR, MODEL_TENSORS, TENSOR_NAMES
  4. class TensorNameMap:
  5. mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = {
  6. # Token embeddings
  7. MODEL_TENSOR.TOKEN_EMBD: (
  8. "gpt_neox.embed_in", # gptneox
  9. "transformer.wte", # gpt2 gpt-j mpt refact qwen dbrx jais exaone
  10. "transformer.word_embeddings", # falcon
  11. "word_embeddings", # bloom
  12. "model.embed_tokens", # llama-hf nemotron olmoe olmo2 rwkv6qwen2 glm4-0414
  13. "tok_embeddings", # llama-pth
  14. "embeddings.word_embeddings", # bert nomic-bert
  15. "language_model.embedding.word_embeddings", # persimmon
  16. "wte", # gpt2
  17. "transformer.embd.wte", # phi2
  18. "model.tok_embeddings", # internlm2
  19. "model.embedding", # mamba-qbert
  20. "backbone.embedding", # mamba
  21. "backbone.embeddings", # mamba-hf
  22. "transformer.in_out_embed", # Grok
  23. "embedding.word_embeddings", # chatglm
  24. "transformer.token_embeddings", # openelm
  25. "shared", # t5
  26. "rwkv.embeddings", # rwkv6
  27. "model.embeddings", # rwkv7
  28. "model.word_embeddings", # bailingmoe
  29. "language_model.model.embed_tokens", # llama4
  30. ),
  31. # Token type embeddings
  32. MODEL_TENSOR.TOKEN_TYPES: (
  33. "embeddings.token_type_embeddings", # bert nomic-bert
  34. ),
  35. # Normalization of token embeddings
  36. MODEL_TENSOR.TOKEN_EMBD_NORM: (
  37. "word_embeddings_layernorm", # bloom
  38. "embeddings.LayerNorm", # bert
  39. "emb_ln", # nomic-bert
  40. "transformer.norm", # openelm
  41. "rwkv.blocks.0.pre_ln", # rwkv
  42. "rwkv.blocks.0.pre_ln", # rwkv6
  43. "model.pre_ln", # rwkv7
  44. "model.layers.0.pre_norm", # rwkv7
  45. "backbone.norm", # wavtokenizer
  46. ),
  47. # Position embeddings
  48. MODEL_TENSOR.POS_EMBD: (
  49. "transformer.wpe", # gpt2
  50. "embeddings.position_embeddings", # bert
  51. "wpe", # gpt2
  52. ),
  53. # Output
  54. MODEL_TENSOR.OUTPUT: (
  55. "embed_out", # gptneox
  56. "lm_head", # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais nemotron exaone olmoe olmo2 phimoe
  57. "output", # llama-pth bloom internlm2
  58. "word_embeddings_for_head", # persimmon
  59. "lm_head.linear", # phi2
  60. "output_layer", # chatglm
  61. "head", # rwkv
  62. "head.out", # wavtokenizer
  63. "lm_head", # llama4
  64. ),
  65. # Output norm
  66. MODEL_TENSOR.OUTPUT_NORM: (
  67. "gpt_neox.final_layer_norm", # gptneox
  68. "transformer.ln_f", # gpt2 gpt-j falcon jais exaone
  69. "model.norm", # llama-hf baichuan internlm2 olmoe olmo2 phimoe
  70. "norm", # llama-pth
  71. "transformer.norm_f", # mpt dbrx
  72. "ln_f", # refact bloom qwen gpt2
  73. "language_model.encoder.final_layernorm", # persimmon
  74. "model.final_layernorm", # persimmon
  75. "lm_head.ln", # phi2
  76. "model.norm_f", # mamba-qbert
  77. "backbone.norm_f", # mamba
  78. "transformer.rms_norm", # Grok
  79. "encoder.final_layernorm", # chatglm
  80. "transformer.norm", # openelm
  81. "model.norm", # nemotron
  82. "rwkv.ln_out", # rwkv6
  83. "model.ln_out", # rwkv7
  84. "backbone.final_layer_norm", # wavtokenizer
  85. "model.norm", # llama4
  86. ),
  87. # Rope frequencies
  88. MODEL_TENSOR.ROPE_FREQS: (
  89. "rope.freqs", # llama-pth
  90. "rotary_pos_emb.inv_freq", # chatglm
  91. ),
  92. MODEL_TENSOR.ROPE_FACTORS_LONG: (),
  93. MODEL_TENSOR.ROPE_FACTORS_SHORT: (),
  94. MODEL_TENSOR.CONV1D: (
  95. "backbone.embed", # roberta
  96. ),
  97. }
  98. block_mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = {
  99. # Attention norm
  100. MODEL_TENSOR.ATTN_NORM: (
  101. "gpt_neox.layers.{bid}.input_layernorm", # gptneox
  102. "transformer.h.{bid}.ln_1", # gpt2 gpt-j refact qwen jais exaone
  103. "transformer.blocks.{bid}.norm_1", # mpt
  104. "transformer.h.{bid}.input_layernorm", # falcon7b
  105. "h.{bid}.input_layernorm", # bloom
  106. "transformer.h.{bid}.ln_mlp", # falcon40b
  107. "model.layers.{bid}.input_layernorm", # llama-hf nemotron olmoe phimoe
  108. "layers.{bid}.attention_norm", # llama-pth
  109. "language_model.encoder.layers.{bid}.input_layernorm", # persimmon
  110. "model.layers.{bid}.ln1", # yi
  111. "h.{bid}.ln_1", # gpt2
  112. "transformer.h.{bid}.ln", # phi2
  113. "model.layers.layers.{bid}.norm", # plamo
  114. "model.layers.{bid}.attention_norm", # internlm2
  115. "model.layers.{bid}.norm", # mamba-qbert
  116. "backbone.layers.{bid}.norm", # mamba
  117. "transformer.decoder_layer.{bid}.rms_norm", # Grok
  118. "transformer.blocks.{bid}.norm_attn_norm.norm_1", # dbrx
  119. "encoder.layers.{bid}.input_layernorm", # chatglm
  120. "transformer.layers.{bid}.attn_norm", # openelm
  121. "rwkv.blocks.{bid}.ln1", # rwkv6
  122. "model.layers.{bid}.ln1", # rwkv7
  123. "model.layers.{bid}.input_layernorm", # llama4
  124. ),
  125. # Attention norm 2
  126. MODEL_TENSOR.ATTN_NORM_2: (
  127. "transformer.h.{bid}.ln_attn", # falcon40b
  128. "encoder.layer.{bid}.layer_norm_1", # jina-v2-code
  129. "rwkv.blocks.{bid}.ln2", # rwkv6
  130. "model.layers.{bid}.ln2", # rwkv7
  131. ),
  132. # Attention query-key-value
  133. MODEL_TENSOR.ATTN_QKV: (
  134. "gpt_neox.layers.{bid}.attention.query_key_value", # gptneox
  135. "transformer.h.{bid}.attn.c_attn", # gpt2 qwen jais
  136. "transformer.blocks.{bid}.attn.Wqkv", # mpt
  137. "transformer.blocks.{bid}.norm_attn_norm.attn.Wqkv", # dbrx
  138. "transformer.h.{bid}.self_attention.query_key_value", # falcon
  139. "h.{bid}.self_attention.query_key_value", # bloom
  140. "language_model.encoder.layers.{bid}.self_attention.query_key_value", # persimmon
  141. "model.layers.{bid}.self_attn.query_key_value", # persimmon
  142. "h.{bid}.attn.c_attn", # gpt2
  143. "transformer.h.{bid}.mixer.Wqkv", # phi2
  144. "encoder.layers.{bid}.attn.Wqkv", # nomic-bert
  145. "model.layers.{bid}.self_attn.qkv_proj", # phi3
  146. "encoder.layers.{bid}.self_attention.query_key_value", # chatglm
  147. "transformer.layers.{bid}.attn.qkv_proj", # openelm
  148. ),
  149. # Attention query
  150. MODEL_TENSOR.ATTN_Q: (
  151. "model.layers.{bid}.self_attn.q_proj", # llama-hf nemotron olmoe olmo2 phimoe
  152. "model.layers.{bid}.self_attn.q_proj_no_perm", # llama-custom
  153. "layers.{bid}.attention.wq", # llama-pth
  154. "encoder.layer.{bid}.attention.self.query", # bert
  155. "transformer.h.{bid}.attn.q_proj", # gpt-j
  156. "model.layers.layers.{bid}.self_attn.q_proj", # plamo
  157. "model.layers.{bid}.attention.wq", # internlm2
  158. "transformer.decoder_layer.{bid}.multi_head_attention.query",# Grok
  159. "transformer.h.{bid}.attn.attention.q_proj", # exaone
  160. "model.layers.{bid}.self_attn.q_proj", # llama4
  161. ),
  162. # Attention key
  163. MODEL_TENSOR.ATTN_K: (
  164. "model.layers.{bid}.self_attn.k_proj", # llama-hf nemotron olmoe olmo2 phimoe
  165. "model.layers.{bid}.self_attn.k_proj_no_perm", # llama-custom
  166. "layers.{bid}.attention.wk", # llama-pth
  167. "encoder.layer.{bid}.attention.self.key", # bert
  168. "transformer.h.{bid}.attn.k_proj", # gpt-j
  169. "transformer.h.{bid}.attn.k", # refact
  170. "model.layers.layers.{bid}.self_attn.k_proj", # plamo
  171. "model.layers.{bid}.attention.wk", # internlm2
  172. "transformer.decoder_layer.{bid}.multi_head_attention.key",# Grok
  173. "transformer.h.{bid}.attn.attention.k_proj", # exaone
  174. "model.layers.{bid}.self_attn.k_proj", # llama4
  175. ),
  176. # Attention value
  177. MODEL_TENSOR.ATTN_V: (
  178. "model.layers.{bid}.self_attn.v_proj", # llama-hf nemotron olmoe olmo2 phimoe
  179. "layers.{bid}.attention.wv", # llama-pth
  180. "encoder.layer.{bid}.attention.self.value", # bert
  181. "transformer.h.{bid}.attn.v_proj", # gpt-j
  182. "transformer.h.{bid}.attn.v", # refact
  183. "model.layers.layers.{bid}.self_attn.v_proj", # plamo
  184. "model.layers.{bid}.attention.wv", # internlm2
  185. "transformer.decoder_layer.{bid}.multi_head_attention.value",# Grok
  186. "transformer.h.{bid}.attn.attention.v_proj", # exaone
  187. "model.layers.{bid}.self_attn.v_proj", # llama4
  188. ),
  189. # Attention output
  190. MODEL_TENSOR.ATTN_OUT: (
  191. "gpt_neox.layers.{bid}.attention.dense", # gptneox
  192. "transformer.h.{bid}.attn.c_proj", # gpt2 refact qwen jais
  193. "transformer.blocks.{bid}.attn.out_proj", # mpt
  194. "transformer.h.{bid}.self_attention.dense", # falcon
  195. "h.{bid}.self_attention.dense", # bloom
  196. "model.layers.{bid}.self_attn.o_proj", # llama-hf nemotron olmoe olmo2 phimoe
  197. "model.layers.{bid}.self_attn.linear_attn", # deci
  198. "layers.{bid}.attention.wo", # llama-pth
  199. "encoder.layer.{bid}.attention.output.dense", # bert
  200. "transformer.h.{bid}.attn.out_proj", # gpt-j
  201. "language_model.encoder.layers.{bid}.self_attention.dense", # persimmon
  202. "model.layers.{bid}.self_attn.dense", # persimmon
  203. "h.{bid}.attn.c_proj", # gpt2
  204. "transformer.h.{bid}.mixer.out_proj", # phi2
  205. "model.layers.layers.{bid}.self_attn.o_proj", # plamo
  206. "model.layers.{bid}.attention.wo", # internlm2
  207. "encoder.layers.{bid}.attn.out_proj", # nomic-bert
  208. "transformer.decoder_layer.{bid}.multi_head_attention.linear", # Grok
  209. "transformer.blocks.{bid}.norm_attn_norm.attn.out_proj", # dbrx
  210. "encoder.layers.{bid}.self_attention.dense", # chatglm
  211. "transformer.layers.{bid}.attn.out_proj", # openelm
  212. "transformer.h.{bid}.attn.attention.out_proj", # exaone
  213. "model.layers.{bid}.self_attn.o_proj", # llama4
  214. ),
  215. # Attention output norm
  216. MODEL_TENSOR.ATTN_OUT_NORM: (
  217. "encoder.layer.{bid}.attention.output.LayerNorm", # bert
  218. "encoder.layers.{bid}.norm1", # nomic-bert
  219. "transformer.decoder_layer.{bid}.rms_norm_1", # Grok
  220. "transformer.blocks.{bid}.norm_attn_norm.norm_2", # dbrx
  221. ),
  222. MODEL_TENSOR.ATTN_POST_NORM: (
  223. "model.layers.{bid}.post_attention_layernorm", # gemma2 olmo2 # ge
  224. "model.layers.{bid}.post_self_attn_layernorm", # glm-4-0414
  225. ),
  226. # Rotary embeddings
  227. MODEL_TENSOR.ATTN_ROT_EMBD: (
  228. "model.layers.{bid}.self_attn.rotary_emb.inv_freq", # llama-hf
  229. "layers.{bid}.attention.inner_attention.rope.freqs", # llama-pth
  230. "model.layers.layers.{bid}.self_attn.rotary_emb.inv_freq", # plamo
  231. "transformer.h.{bid}.attn.rotary_emb.inv_freq", # codeshell
  232. ),
  233. # Feed-forward norm
  234. MODEL_TENSOR.FFN_NORM: (
  235. "gpt_neox.layers.{bid}.post_attention_layernorm", # gptneox
  236. "transformer.h.{bid}.ln_2", # gpt2 refact qwen jais exaone
  237. "h.{bid}.post_attention_layernorm", # bloom
  238. "transformer.blocks.{bid}.norm_2", # mpt
  239. "model.layers.{bid}.post_attention_layernorm", # llama-hf nemotron olmoe phimoe
  240. "layers.{bid}.ffn_norm", # llama-pth
  241. "language_model.encoder.layers.{bid}.post_attention_layernorm", # persimmon
  242. "model.layers.{bid}.ln2", # yi
  243. "h.{bid}.ln_2", # gpt2
  244. "model.layers.{bid}.ffn_norm", # internlm2
  245. "transformer.decoder_layer.{bid}.rms_norm_2", # Grok
  246. "encoder.layers.{bid}.post_attention_layernorm", # chatglm
  247. "transformer.layers.{bid}.ffn_norm", # openelm
  248. "model.layers.{bid}.post_attention_layernorm", # llama4
  249. ),
  250. # Post feed-forward norm
  251. MODEL_TENSOR.FFN_PRE_NORM: (
  252. "model.layers.{bid}.pre_feedforward_layernorm", # gemma2
  253. ),
  254. # Post feed-forward norm
  255. MODEL_TENSOR.FFN_POST_NORM: (
  256. "model.layers.{bid}.post_feedforward_layernorm", # gemma2 olmo2
  257. "model.layers.{bid}.post_mlp_layernorm", # glm-4-0414
  258. ),
  259. MODEL_TENSOR.FFN_GATE_INP: (
  260. "layers.{bid}.feed_forward.gate", # mixtral
  261. "model.layers.{bid}.block_sparse_moe.gate", # mixtral phimoe
  262. "model.layers.{bid}.mlp.gate", # qwen2moe olmoe
  263. "transformer.decoder_layer.{bid}.router", # Grok
  264. "transformer.blocks.{bid}.ffn.router.layer", # dbrx
  265. "model.layers.{bid}.block_sparse_moe.router.layer", # granitemoe
  266. "model.layers.{bid}.feed_forward.router", # llama4
  267. "encoder.layers.{bid}.mlp.router.layer", # nomic-bert-moe
  268. ),
  269. MODEL_TENSOR.FFN_GATE_INP_SHEXP: (
  270. "model.layers.{bid}.mlp.shared_expert_gate", # qwen2moe
  271. ),
  272. MODEL_TENSOR.FFN_EXP_PROBS_B: (
  273. "model.layers.{bid}.mlp.gate.e_score_correction", # deepseek-v3
  274. ),
  275. # Feed-forward up
  276. MODEL_TENSOR.FFN_UP: (
  277. "gpt_neox.layers.{bid}.mlp.dense_h_to_4h", # gptneox
  278. "transformer.h.{bid}.mlp.c_fc", # gpt2 jais
  279. "transformer.blocks.{bid}.ffn.up_proj", # mpt
  280. "transformer.h.{bid}.mlp.dense_h_to_4h", # falcon
  281. "h.{bid}.mlp.dense_h_to_4h", # bloom
  282. "model.layers.{bid}.mlp.up_proj", # llama-hf refact nemotron olmo2
  283. "layers.{bid}.feed_forward.w3", # llama-pth
  284. "encoder.layer.{bid}.intermediate.dense", # bert
  285. "transformer.h.{bid}.mlp.fc_in", # gpt-j
  286. "transformer.h.{bid}.mlp.linear_3", # refact
  287. "language_model.encoder.layers.{bid}.mlp.dense_h_to_4h", # persimmon
  288. "model.layers.{bid}.mlp.dense_h_to_4h", # persimmon
  289. "transformer.h.{bid}.mlp.w1", # qwen
  290. "h.{bid}.mlp.c_fc", # gpt2
  291. "transformer.h.{bid}.mlp.fc1", # phi2
  292. "model.layers.{bid}.mlp.fc1", # phi2
  293. "model.layers.{bid}.mlp.gate_up_proj", # phi3 glm-4-0414
  294. "model.layers.layers.{bid}.mlp.up_proj", # plamo
  295. "model.layers.{bid}.feed_forward.w3", # internlm2
  296. "encoder.layers.{bid}.mlp.fc11", # nomic-bert
  297. "encoder.layers.{bid}.mlp.fc1", # nomic-bert-moe
  298. "model.layers.{bid}.mlp.c_fc", # starcoder2
  299. "encoder.layer.{bid}.mlp.gated_layers_v", # jina-bert-v2
  300. "model.layers.{bid}.residual_mlp.w3", # arctic
  301. "encoder.layers.{bid}.mlp.dense_h_to_4h", # chatglm
  302. "transformer.h.{bid}.mlp.c_fc_1", # exaone
  303. "model.layers.{bid}.feed_forward.up_proj", # llama4
  304. ),
  305. MODEL_TENSOR.FFN_UP_EXP: (
  306. "layers.{bid}.feed_forward.experts.w3", # mixtral (merged)
  307. "transformer.decoder_layer.{bid}.moe.linear_v", # Grok (merged)
  308. "transformer.blocks.{bid}.ffn.experts.mlp.v1", # dbrx
  309. "model.layers.{bid}.mlp.experts.up_proj", # qwen2moe olmoe (merged)
  310. "model.layers.{bid}.block_sparse_moe.experts.w3", # phimoe (merged)
  311. "model.layers.{bid}.feed_forward.experts.up_proj", # llama4
  312. "encoder.layers.{bid}.mlp.experts.mlp.w1", # nomic-bert-moe
  313. ),
  314. MODEL_TENSOR.FFN_UP_SHEXP: (
  315. "model.layers.{bid}.mlp.shared_expert.up_proj", # qwen2moe
  316. "model.layers.{bid}.mlp.shared_experts.up_proj", # deepseek deepseek2
  317. "model.layers.{bid}.feed_forward.shared_expert.up_proj", # llama4
  318. ),
  319. # AWQ-activation gate
  320. MODEL_TENSOR.FFN_ACT: (
  321. "transformer.blocks.{bid}.ffn.act", # mpt
  322. ),
  323. # Feed-forward gate
  324. MODEL_TENSOR.FFN_GATE: (
  325. "model.layers.{bid}.mlp.gate_proj", # llama-hf refact olmo2
  326. "layers.{bid}.feed_forward.w1", # llama-pth
  327. "transformer.h.{bid}.mlp.w2", # qwen
  328. "transformer.h.{bid}.mlp.c_fc2", # jais
  329. "model.layers.layers.{bid}.mlp.gate_proj", # plamo
  330. "model.layers.{bid}.feed_forward.w1", # internlm2
  331. "encoder.layers.{bid}.mlp.fc12", # nomic-bert
  332. "encoder.layer.{bid}.mlp.gated_layers_w", # jina-bert-v2
  333. "transformer.h.{bid}.mlp.linear_1", # refact
  334. "model.layers.{bid}.residual_mlp.w1", # arctic
  335. "transformer.h.{bid}.mlp.c_fc_0", # exaone
  336. "model.layers.{bid}.feed_forward.gate_proj", # llama4
  337. ),
  338. MODEL_TENSOR.FFN_GATE_EXP: (
  339. "layers.{bid}.feed_forward.experts.w1", # mixtral (merged)
  340. "transformer.decoder_layer.{bid}.moe.linear", # Grok (merged)
  341. "transformer.blocks.{bid}.ffn.experts.mlp.w1", # dbrx
  342. "model.layers.{bid}.mlp.experts.gate_proj", # qwen2moe olmoe (merged)
  343. "model.layers.{bid}.block_sparse_moe.experts.w1", # phimoe (merged)
  344. "model.layers.{bid}.feed_forward.experts.gate_proj", # llama4
  345. ),
  346. MODEL_TENSOR.FFN_GATE_SHEXP: (
  347. "model.layers.{bid}.mlp.shared_expert.gate_proj", # qwen2moe
  348. "model.layers.{bid}.mlp.shared_experts.gate_proj", # deepseek deepseek2
  349. "model.layers.{bid}.feed_forward.shared_expert.gate_proj", # llama4
  350. ),
  351. # Feed-forward down
  352. MODEL_TENSOR.FFN_DOWN: (
  353. "gpt_neox.layers.{bid}.mlp.dense_4h_to_h", # gptneox
  354. "transformer.h.{bid}.mlp.c_proj", # gpt2 refact qwen jais
  355. "transformer.blocks.{bid}.ffn.down_proj", # mpt
  356. "transformer.h.{bid}.mlp.dense_4h_to_h", # falcon
  357. "h.{bid}.mlp.dense_4h_to_h", # bloom
  358. "model.layers.{bid}.mlp.down_proj", # llama-hf nemotron olmo2
  359. "layers.{bid}.feed_forward.w2", # llama-pth
  360. "encoder.layer.{bid}.output.dense", # bert
  361. "transformer.h.{bid}.mlp.fc_out", # gpt-j
  362. "language_model.encoder.layers.{bid}.mlp.dense_4h_to_h", # persimmon
  363. "model.layers.{bid}.mlp.dense_4h_to_h", # persimmon
  364. "h.{bid}.mlp.c_proj", # gpt2
  365. "transformer.h.{bid}.mlp.fc2", # phi2
  366. "model.layers.{bid}.mlp.fc2", # phi2
  367. "model.layers.layers.{bid}.mlp.down_proj", # plamo
  368. "model.layers.{bid}.feed_forward.w2", # internlm2
  369. "encoder.layers.{bid}.mlp.fc2", # nomic-bert
  370. "model.layers.{bid}.mlp.c_proj", # starcoder2
  371. "encoder.layer.{bid}.mlp.wo", # jina-bert-v2
  372. "transformer.layers.{bid}.ffn.proj_2", # openelm
  373. "model.layers.{bid}.residual_mlp.w2", # arctic
  374. "encoder.layer.{bid}.mlp.down_layer", # jina-bert-v2
  375. "encoder.layers.{bid}.mlp.dense_4h_to_h", # chatglm
  376. "model.layers.h.{bid}.mlp.c_proj", # exaone
  377. "model.layers.{bid}.feed_forward.down_proj", # llama4
  378. ),
  379. MODEL_TENSOR.FFN_DOWN_EXP: (
  380. "layers.{bid}.feed_forward.experts.w2", # mixtral (merged)
  381. "transformer.decoder_layer.{bid}.moe.linear_1", # Grok (merged)
  382. "transformer.blocks.{bid}.ffn.experts.mlp.w2", # dbrx
  383. "model.layers.{bid}.mlp.experts.down_proj", # qwen2moe olmoe (merged)
  384. "model.layers.{bid}.block_sparse_moe.output_linear", # granitemoe
  385. "model.layers.{bid}.block_sparse_moe.experts.w2", # phimoe (merged)
  386. "model.layers.{bid}.feed_forward.experts.down_proj", # llama4
  387. "encoder.layers.{bid}.mlp.experts.mlp.w2", # nomic-bert-moe
  388. ),
  389. MODEL_TENSOR.FFN_DOWN_SHEXP: (
  390. "model.layers.{bid}.mlp.shared_expert.down_proj", # qwen2moe
  391. "model.layers.{bid}.mlp.shared_experts.down_proj", # deepseek deepseek2
  392. "model.layers.{bid}.feed_forward.shared_expert.down_proj", # llama4
  393. "model.layers.{bid}.shared_mlp.output_linear", # granitemoe
  394. ),
  395. MODEL_TENSOR.ATTN_Q_NORM: (
  396. "language_model.encoder.layers.{bid}.self_attention.q_layernorm",
  397. "model.layers.{bid}.self_attn.q_layernorm", # persimmon
  398. "model.layers.{bid}.self_attn.q_norm", # cohere olmoe chameleon olmo2
  399. "transformer.blocks.{bid}.attn.q_ln", # sea-lion
  400. "encoder.layer.{bid}.attention.self.layer_norm_q", # jina-bert-v2
  401. "transformer.layers.{bid}.attn.q_norm", # openelm
  402. ),
  403. MODEL_TENSOR.ATTN_K_NORM: (
  404. "language_model.encoder.layers.{bid}.self_attention.k_layernorm",
  405. "model.layers.{bid}.self_attn.k_layernorm", # persimmon
  406. "model.layers.{bid}.self_attn.k_norm", # cohere olmoe chameleon olmo2
  407. "transformer.blocks.{bid}.attn.k_ln", # sea-lion
  408. "encoder.layer.{bid}.attention.self.layer_norm_k", # jina-bert-v2
  409. "transformer.layers.{bid}.attn.k_norm", # openelm
  410. ),
  411. MODEL_TENSOR.ROPE_FREQS: (
  412. "language_model.encoder.layers.{bid}.self_attention.rotary_emb.inv_freq", # persimmon
  413. ),
  414. MODEL_TENSOR.LAYER_OUT_NORM: (
  415. "encoder.layer.{bid}.output.LayerNorm", # bert
  416. "encoder.layers.{bid}.norm2", # nomic-bert
  417. "transformer.decoder_layer.{bid}.rms_norm_3", # Grok
  418. "encoder.layer.{bid}.mlp.layernorm", # jina-bert-v2
  419. "encoder.layer.{bid}.layer_norm_2" # jina-v2-code
  420. ),
  421. MODEL_TENSOR.SSM_IN: (
  422. "model.layers.{bid}.in_proj",
  423. "backbone.layers.{bid}.mixer.in_proj",
  424. ),
  425. MODEL_TENSOR.SSM_CONV1D: (
  426. "model.layers.{bid}.conv1d",
  427. "backbone.layers.{bid}.mixer.conv1d",
  428. ),
  429. MODEL_TENSOR.SSM_X: (
  430. "model.layers.{bid}.x_proj",
  431. "backbone.layers.{bid}.mixer.x_proj",
  432. ),
  433. MODEL_TENSOR.SSM_DT: (
  434. "model.layers.{bid}.dt_proj",
  435. "backbone.layers.{bid}.mixer.dt_proj",
  436. ),
  437. MODEL_TENSOR.SSM_A: (
  438. "model.layers.{bid}.A_log",
  439. "backbone.layers.{bid}.mixer.A_log",
  440. ),
  441. MODEL_TENSOR.SSM_D: (
  442. "model.layers.{bid}.D",
  443. "backbone.layers.{bid}.mixer.D",
  444. ),
  445. MODEL_TENSOR.SSM_OUT: (
  446. "model.layers.{bid}.out_proj",
  447. "backbone.layers.{bid}.mixer.out_proj",
  448. ),
  449. MODEL_TENSOR.TIME_MIX_W0: (
  450. "model.layers.{bid}.attention.w0", # rwkv7
  451. ),
  452. MODEL_TENSOR.TIME_MIX_W1: (
  453. "rwkv.blocks.{bid}.attention.time_maa_w1", # rwkv6
  454. "model.layers.{bid}.self_attn.time_maa_w1", # rwkv6qwen2
  455. "model.layers.{bid}.attention.w1", # rwkv7
  456. ),
  457. MODEL_TENSOR.TIME_MIX_W2: (
  458. "rwkv.blocks.{bid}.attention.time_maa_w2", # rwkv6
  459. "model.layers.{bid}.self_attn.time_maa_w2", # rwkv6qwen2
  460. "model.layers.{bid}.attention.w2", # rwkv7
  461. ),
  462. MODEL_TENSOR.TIME_MIX_A0: (
  463. "model.layers.{bid}.attention.a0", # rwkv7
  464. ),
  465. MODEL_TENSOR.TIME_MIX_A1: (
  466. "model.layers.{bid}.attention.a1", # rwkv7
  467. ),
  468. MODEL_TENSOR.TIME_MIX_A2: (
  469. "model.layers.{bid}.attention.a2", # rwkv7
  470. ),
  471. MODEL_TENSOR.TIME_MIX_V0: (
  472. "model.layers.{bid}.attention.v0", # rwkv7
  473. ),
  474. MODEL_TENSOR.TIME_MIX_V1: (
  475. "model.layers.{bid}.attention.v1", # rwkv7
  476. ),
  477. MODEL_TENSOR.TIME_MIX_V2: (
  478. "model.layers.{bid}.attention.v2", # rwkv7
  479. ),
  480. MODEL_TENSOR.TIME_MIX_G1: (
  481. "model.layers.{bid}.attention.g1", # rwkv7
  482. ),
  483. MODEL_TENSOR.TIME_MIX_G2: (
  484. "model.layers.{bid}.attention.g2", # rwkv7
  485. ),
  486. MODEL_TENSOR.TIME_MIX_K_K: (
  487. "model.layers.{bid}.attention.k_k", # rwkv7
  488. ),
  489. MODEL_TENSOR.TIME_MIX_K_A: (
  490. "model.layers.{bid}.attention.k_a", # rwkv7
  491. ),
  492. MODEL_TENSOR.TIME_MIX_R_K: (
  493. "model.layers.{bid}.attention.r_k", # rwkv7
  494. ),
  495. MODEL_TENSOR.TIME_MIX_LERP_X: (
  496. "rwkv.blocks.{bid}.attention.time_maa_x", # rwkv6
  497. "model.layers.{bid}.self_attn.time_maa_x", # rwkv6qwen2
  498. ),
  499. MODEL_TENSOR.TIME_MIX_LERP_K: (
  500. "rwkv.blocks.{bid}.attention.time_maa_k", # rwkv6
  501. "model.layers.{bid}.self_attn.time_maa_k", # rwkv6qwen2
  502. ),
  503. MODEL_TENSOR.TIME_MIX_LERP_V: (
  504. "rwkv.blocks.{bid}.attention.time_maa_v", # rwkv6
  505. "model.layers.{bid}.self_attn.time_maa_v", # rwkv6qwen2
  506. ),
  507. MODEL_TENSOR.TIME_MIX_LERP_R: (
  508. "rwkv.blocks.{bid}.attention.time_maa_r", # rwkv6
  509. "model.layers.{bid}.self_attn.time_maa_r", # rwkv6qwen2
  510. ),
  511. MODEL_TENSOR.TIME_MIX_LERP_G: (
  512. "rwkv.blocks.{bid}.attention.time_maa_g", # rwkv6
  513. "model.layers.{bid}.self_attn.time_maa_g", # rwkv6qwen2
  514. ),
  515. MODEL_TENSOR.TIME_MIX_LERP_W: (
  516. "rwkv.blocks.{bid}.attention.time_maa_w", # rwkv6
  517. "model.layers.{bid}.self_attn.time_maa_w", # rwkv6qwen2
  518. ),
  519. MODEL_TENSOR.TIME_MIX_FIRST: (
  520. "rwkv.blocks.{bid}.attention.time_faaaa", # rwkv6
  521. ),
  522. MODEL_TENSOR.TIME_MIX_DECAY: (
  523. "rwkv.blocks.{bid}.attention.time_decay", # rwkv6
  524. "model.layers.{bid}.self_attn.time_decay", # rwkv6qwen2
  525. ),
  526. MODEL_TENSOR.TIME_MIX_DECAY_W1: (
  527. "rwkv.blocks.{bid}.attention.time_decay_w1", # rwkv6
  528. "model.layers.{bid}.self_attn.time_decay_w1", # rwkv6qwen2
  529. ),
  530. MODEL_TENSOR.TIME_MIX_DECAY_W2: (
  531. "rwkv.blocks.{bid}.attention.time_decay_w2", # rwkv6
  532. "model.layers.{bid}.self_attn.time_decay_w2", # rwkv6qwen2
  533. ),
  534. MODEL_TENSOR.TIME_MIX_KEY: (
  535. "rwkv.blocks.{bid}.attention.key", # rwkv6
  536. "model.layers.{bid}.self_attn.k_proj", # rwkv6qwen2
  537. "model.layers.{bid}.attention.key", # rwkv7
  538. "model.layers.{bid}.attention.k_proj", # rwkv7
  539. ),
  540. MODEL_TENSOR.TIME_MIX_VALUE: (
  541. "rwkv.blocks.{bid}.attention.value", # rwkv6
  542. "model.layers.{bid}.self_attn.v_proj", # rwkv6qwen2
  543. "model.layers.{bid}.attention.value", # rwkv7
  544. "model.layers.{bid}.attention.v_proj", # rwkv7
  545. ),
  546. MODEL_TENSOR.TIME_MIX_RECEPTANCE: (
  547. "rwkv.blocks.{bid}.attention.receptance", # rwkv6
  548. "model.layers.{bid}.self_attn.q_proj", # rwkv6qwen2
  549. "model.layers.{bid}.attention.receptance", # rwkv7
  550. "model.layers.{bid}.attention.r_proj", # rwkv7
  551. ),
  552. MODEL_TENSOR.TIME_MIX_GATE: (
  553. "rwkv.blocks.{bid}.attention.gate", # rwkv6
  554. "model.layers.{bid}.self_attn.gate", # rwkv6qwen2
  555. ),
  556. MODEL_TENSOR.TIME_MIX_LN: (
  557. "rwkv.blocks.{bid}.attention.ln_x", # rwkv6
  558. "model.layers.{bid}.attention.ln_x" # rwkv7
  559. ),
  560. MODEL_TENSOR.TIME_MIX_OUTPUT: (
  561. "rwkv.blocks.{bid}.attention.output", # rwkv6
  562. "model.layers.{bid}.self_attn.o_proj", # rwkv6qwen2
  563. "model.layers.{bid}.attention.output", # rwkv7
  564. "model.layers.{bid}.attention.o_proj", # rwkv7
  565. ),
  566. MODEL_TENSOR.CHANNEL_MIX_LERP_K: (
  567. "rwkv.blocks.{bid}.feed_forward.time_maa_k", # rwkv6
  568. "model.layers.{bid}.feed_forward.x_k", # rwkv7
  569. ),
  570. MODEL_TENSOR.CHANNEL_MIX_LERP_R: (
  571. "rwkv.blocks.{bid}.feed_forward.time_maa_r", # rwkv6
  572. ),
  573. MODEL_TENSOR.CHANNEL_MIX_KEY: (
  574. "rwkv.blocks.{bid}.feed_forward.key", # rwkv6
  575. "model.layers.{bid}.feed_forward.key", # rwkv7
  576. ),
  577. MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE: (
  578. "rwkv.blocks.{bid}.feed_forward.receptance", # rwkv6
  579. ),
  580. MODEL_TENSOR.CHANNEL_MIX_VALUE: (
  581. "rwkv.blocks.{bid}.feed_forward.value", # rwkv6
  582. "model.layers.{bid}.feed_forward.value", # rwkv7
  583. ),
  584. MODEL_TENSOR.ATTN_Q_A: (
  585. "model.layers.{bid}.self_attn.q_a_proj", # deepseek2
  586. ),
  587. MODEL_TENSOR.ATTN_Q_B: (
  588. "model.layers.{bid}.self_attn.q_b_proj", # deepseek2
  589. ),
  590. MODEL_TENSOR.ATTN_KV_A_MQA: (
  591. "model.layers.{bid}.self_attn.kv_a_proj_with_mqa", # deepseek2
  592. ),
  593. MODEL_TENSOR.ATTN_KV_B: (
  594. "model.layers.{bid}.self_attn.kv_b_proj", # deepseek2
  595. ),
  596. MODEL_TENSOR.ATTN_K_B: (
  597. "model.layers.{bid}.self_attn.k_b_proj", # deepseek2
  598. ),
  599. MODEL_TENSOR.ATTN_V_B: (
  600. "model.layers.{bid}.self_attn.v_b_proj", # deepseek2
  601. ),
  602. MODEL_TENSOR.ATTN_Q_A_NORM: (
  603. "model.layers.{bid}.self_attn.q_a_layernorm", # deepseek2
  604. ),
  605. MODEL_TENSOR.ATTN_KV_A_NORM: (
  606. "model.layers.{bid}.self_attn.kv_a_layernorm", # deepseek2
  607. ),
  608. MODEL_TENSOR.ATTN_SUB_NORM: (
  609. "model.layers.{bid}.self_attn.inner_attn_ln", # bitnet
  610. ),
  611. MODEL_TENSOR.FFN_SUB_NORM: (
  612. "model.layers.{bid}.mlp.ffn_layernorm", # bitnet
  613. ),
  614. MODEL_TENSOR.DEC_ATTN_NORM: (
  615. "decoder.block.{bid}.layer.0.layer_norm", # t5
  616. ),
  617. MODEL_TENSOR.DEC_ATTN_Q: (
  618. "decoder.block.{bid}.layer.0.SelfAttention.q", # t5
  619. ),
  620. MODEL_TENSOR.DEC_ATTN_K: (
  621. "decoder.block.{bid}.layer.0.SelfAttention.k", # t5
  622. ),
  623. MODEL_TENSOR.DEC_ATTN_V: (
  624. "decoder.block.{bid}.layer.0.SelfAttention.v", # t5
  625. ),
  626. MODEL_TENSOR.DEC_ATTN_OUT: (
  627. "decoder.block.{bid}.layer.0.SelfAttention.o", # t5
  628. ),
  629. MODEL_TENSOR.DEC_ATTN_REL_B: (
  630. "decoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5
  631. ),
  632. MODEL_TENSOR.DEC_CROSS_ATTN_NORM: (
  633. "decoder.block.{bid}.layer.1.layer_norm", # t5
  634. ),
  635. MODEL_TENSOR.DEC_CROSS_ATTN_Q: (
  636. "decoder.block.{bid}.layer.1.EncDecAttention.q", # t5
  637. ),
  638. MODEL_TENSOR.DEC_CROSS_ATTN_K: (
  639. "decoder.block.{bid}.layer.1.EncDecAttention.k", # t5
  640. ),
  641. MODEL_TENSOR.DEC_CROSS_ATTN_V: (
  642. "decoder.block.{bid}.layer.1.EncDecAttention.v", # t5
  643. ),
  644. MODEL_TENSOR.DEC_CROSS_ATTN_OUT: (
  645. "decoder.block.{bid}.layer.1.EncDecAttention.o", # t5
  646. ),
  647. MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: (
  648. "decoder.block.{bid}.layer.1.EncDecAttention.relative_attention_bias", # t5
  649. ),
  650. MODEL_TENSOR.DEC_FFN_NORM: (
  651. "decoder.block.{bid}.layer.2.layer_norm", # t5
  652. ),
  653. MODEL_TENSOR.DEC_FFN_GATE: (
  654. "decoder.block.{bid}.layer.2.DenseReluDense.wi_0", # flan-t5
  655. ),
  656. MODEL_TENSOR.DEC_FFN_UP: (
  657. "decoder.block.{bid}.layer.2.DenseReluDense.wi", # t5
  658. "decoder.block.{bid}.layer.2.DenseReluDense.wi_1", # flan-t5
  659. ),
  660. MODEL_TENSOR.DEC_FFN_DOWN: (
  661. "decoder.block.{bid}.layer.2.DenseReluDense.wo", # t5
  662. ),
  663. MODEL_TENSOR.DEC_OUTPUT_NORM: (
  664. "decoder.final_layer_norm", # t5
  665. ),
  666. MODEL_TENSOR.ENC_ATTN_NORM: (
  667. "encoder.block.{bid}.layer.0.layer_norm", # t5
  668. ),
  669. MODEL_TENSOR.ENC_ATTN_Q: (
  670. "encoder.block.{bid}.layer.0.SelfAttention.q", # t5
  671. ),
  672. MODEL_TENSOR.ENC_ATTN_K: (
  673. "encoder.block.{bid}.layer.0.SelfAttention.k", # t5
  674. ),
  675. MODEL_TENSOR.ENC_ATTN_V: (
  676. "encoder.block.{bid}.layer.0.SelfAttention.v", # t5
  677. ),
  678. MODEL_TENSOR.ENC_ATTN_OUT: (
  679. "encoder.block.{bid}.layer.0.SelfAttention.o", # t5
  680. ),
  681. MODEL_TENSOR.ENC_ATTN_REL_B: (
  682. "encoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5
  683. ),
  684. MODEL_TENSOR.ENC_FFN_NORM: (
  685. "encoder.block.{bid}.layer.1.layer_norm", # t5
  686. ),
  687. MODEL_TENSOR.ENC_FFN_GATE: (
  688. "encoder.block.{bid}.layer.1.DenseReluDense.wi_0", # flan-t5
  689. ),
  690. MODEL_TENSOR.ENC_FFN_UP: (
  691. "encoder.block.{bid}.layer.1.DenseReluDense.wi", # t5
  692. "encoder.block.{bid}.layer.1.DenseReluDense.wi_1", # flan-t5
  693. ),
  694. MODEL_TENSOR.ENC_FFN_DOWN: (
  695. "encoder.block.{bid}.layer.1.DenseReluDense.wo", # t5
  696. ),
  697. ############################################################################
  698. # TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg
  699. MODEL_TENSOR.ENC_OUTPUT_NORM: (
  700. "encoder.final_layer_norm", # t5
  701. ),
  702. MODEL_TENSOR.CLS: (
  703. "classifier", # jina
  704. "classifier.dense", # roberta
  705. ),
  706. MODEL_TENSOR.CLS_OUT: (
  707. "classifier.out_proj", # roberta
  708. ),
  709. #############################################################################
  710. MODEL_TENSOR.CONVNEXT_DW: (
  711. "backbone.convnext.{bid}.dwconv", # wavtokenizer
  712. ),
  713. MODEL_TENSOR.CONVNEXT_NORM: (
  714. "backbone.convnext.{bid}.norm", # wavtokenizer
  715. ),
  716. MODEL_TENSOR.CONVNEXT_PW1: (
  717. "backbone.convnext.{bid}.pwconv1", # wavtokenizer
  718. ),
  719. MODEL_TENSOR.CONVNEXT_PW2: (
  720. "backbone.convnext.{bid}.pwconv2", # wavtokenizer
  721. ),
  722. MODEL_TENSOR.CONVNEXT_GAMMA: (
  723. "backbone.convnext.{bid}.gamma", # wavtokenizer
  724. ),
  725. MODEL_TENSOR.POSNET_CONV1: (
  726. "backbone.posnet.{bid}.conv1", # wavtokenizer
  727. ),
  728. MODEL_TENSOR.POSNET_CONV2: (
  729. "backbone.posnet.{bid}.conv2", # wavtokenizer
  730. ),
  731. MODEL_TENSOR.POSNET_NORM: (
  732. "backbone.posnet.{bid}.norm", # wavtokenizer
  733. ),
  734. MODEL_TENSOR.POSNET_NORM1: (
  735. "backbone.posnet.{bid}.norm1", # wavtokenizer
  736. ),
  737. MODEL_TENSOR.POSNET_NORM2: (
  738. "backbone.posnet.{bid}.norm2", # wavtokenizer
  739. ),
  740. MODEL_TENSOR.POSNET_ATTN_NORM: (
  741. "backbone.posnet.{bid}.norm", # wavtokenizer
  742. ),
  743. MODEL_TENSOR.POSNET_ATTN_Q: (
  744. "backbone.posnet.{bid}.q", # wavtokenizer
  745. ),
  746. MODEL_TENSOR.POSNET_ATTN_K: (
  747. "backbone.posnet.{bid}.k", # wavtokenizer
  748. ),
  749. MODEL_TENSOR.POSNET_ATTN_V: (
  750. "backbone.posnet.{bid}.v", # wavtokenizer
  751. ),
  752. MODEL_TENSOR.POSNET_ATTN_OUT: (
  753. "backbone.posnet.{bid}.proj_out", # wavtokenizer
  754. ),
  755. #############################################################################
  756. ## Vision encoder
  757. MODEL_TENSOR.V_MMPROJ: (
  758. "multi_modal_projector.linear_{bid}",
  759. "visual.merger.mlp.{bid}", # qwen2vl
  760. ),
  761. MODEL_TENSOR.V_MMPROJ_FC: (
  762. "model.connector.modality_projection.proj", # SmolVLM
  763. ),
  764. MODEL_TENSOR.V_MMPROJ_MLP: (
  765. "model.mm_projector.mlp.mlp.{bid}",
  766. "mlp1.{bid}", # InternVL
  767. ),
  768. MODEL_TENSOR.V_MMPROJ_PEG: (
  769. "model.mm_projector.peg.peg.{bid}",
  770. ),
  771. MODEL_TENSOR.V_ENC_EMBD_CLS: (
  772. "vision_tower.vision_model.embeddings.class_embedding",
  773. ),
  774. MODEL_TENSOR.V_ENC_EMBD_PATCH: (
  775. "vision_tower.vision_model.embeddings.patch_embedding",
  776. "vpm.embeddings.patch_embedding",
  777. "model.vision_model.embeddings.patch_embedding", # SmolVLM
  778. "vision_tower.patch_conv", # pixtral
  779. "visual.patch_embed.proj", # qwen2vl
  780. ),
  781. MODEL_TENSOR.V_ENC_EMBD_POS: (
  782. "vision_tower.vision_model.embeddings.position_embedding",
  783. "vpm.embeddings.position_embedding",
  784. "model.vision_model.embeddings.position_embedding", # SmolVLM
  785. ),
  786. MODEL_TENSOR.V_ENC_ATTN_Q: (
  787. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.q_proj",
  788. "vpm.encoder.layers.{bid}.self_attn.q_proj",
  789. "model.vision_model.encoder.layers.{bid}.self_attn.q_proj", # SmolVLM
  790. "vision_tower.transformer.layers.{bid}.attention.q_proj", # pixtral
  791. "visual.blocks.{bid}.attn.q", # qwen2vl, generated
  792. ),
  793. MODEL_TENSOR.V_ENC_ATTN_Q_NORM: (
  794. "vision_tower.vision_model.encoder.layers.{bid}.attn.q_norm", # InternVL
  795. ),
  796. MODEL_TENSOR.V_ENC_ATTN_K: (
  797. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.k_proj",
  798. "vpm.encoder.layers.{bid}.self_attn.k_proj",
  799. "model.vision_model.encoder.layers.{bid}.self_attn.k_proj", # SmolVLM
  800. "vision_tower.transformer.layers.{bid}.attention.k_proj", # pixtral
  801. "visual.blocks.{bid}.attn.k", # qwen2vl, generated
  802. ),
  803. MODEL_TENSOR.V_ENC_ATTN_K_NORM: (
  804. "vision_tower.vision_model.encoder.layers.{bid}.attn.k_norm", # InternVL
  805. ),
  806. MODEL_TENSOR.V_ENC_ATTN_V: (
  807. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.v_proj",
  808. "vpm.encoder.layers.{bid}.self_attn.v_proj",
  809. "model.vision_model.encoder.layers.{bid}.self_attn.v_proj", # SmolVLM
  810. "vision_tower.transformer.layers.{bid}.attention.v_proj", # pixtral
  811. "visual.blocks.{bid}.attn.v", # qwen2vl, generated
  812. ),
  813. MODEL_TENSOR.V_ENC_INPUT_NORM: (
  814. "vision_tower.vision_model.encoder.layers.{bid}.layer_norm1",
  815. "vision_tower.vision_model.encoder.layers.{bid}.norm1", # InternVL
  816. "vpm.encoder.layers.{bid}.layer_norm1",
  817. "model.vision_model.encoder.layers.{bid}.layer_norm1", # SmolVLM
  818. "vision_tower.transformer.layers.{bid}.attention_norm", # pixtral
  819. "visual.blocks.{bid}.norm1", # qwen2vl
  820. ),
  821. MODEL_TENSOR.V_ENC_OUTPUT: (
  822. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.out_proj",
  823. "vision_tower.vision_model.encoder.layers.{bid}.attn.proj", # InternVL
  824. "vpm.encoder.layers.{bid}.self_attn.out_proj",
  825. "model.vision_model.encoder.layers.{bid}.self_attn.out_proj", # SmolVLM
  826. "vision_tower.transformer.layers.{bid}.attention.o_proj", # pixtral
  827. "visual.blocks.{bid}.attn.proj", # qwen2vl
  828. ),
  829. MODEL_TENSOR.V_ENC_OUTPUT_NORM: (
  830. "vision_tower.vision_model.encoder.layers.{bid}.layer_norm2",
  831. "vision_tower.vision_model.encoder.layers.{bid}.norm2", # InternVL
  832. "vpm.encoder.layers.{bid}.layer_norm2",
  833. "model.vision_model.encoder.layers.{bid}.layer_norm2", # SmolVLM
  834. "vision_tower.transformer.layers.{bid}.ffn_norm", # pixtral
  835. "visual.blocks.{bid}.norm2", # qwen2vl
  836. ),
  837. MODEL_TENSOR.V_ENC_FFN_UP: (
  838. "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc1",
  839. "vpm.encoder.layers.{bid}.mlp.fc1",
  840. "model.vision_model.encoder.layers.{bid}.mlp.fc1", # SmolVLM, gemma3
  841. "vision_tower.transformer.layers.{bid}.feed_forward.up_proj", # pixtral
  842. "visual.blocks.{bid}.mlp.fc1", # qwen2vl
  843. "visual.blocks.{bid}.mlp.up_proj", # qwen2.5vl
  844. ),
  845. MODEL_TENSOR.V_ENC_FFN_GATE: (
  846. "vision_tower.transformer.layers.{bid}.feed_forward.gate_proj", # pixtral
  847. "visual.blocks.{bid}.mlp.gate_proj", # qwen2.5vl
  848. ),
  849. MODEL_TENSOR.V_ENC_FFN_DOWN: (
  850. "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc2",
  851. "vpm.encoder.layers.{bid}.mlp.fc2",
  852. "model.vision_model.encoder.layers.{bid}.mlp.fc2", # SmolVLM, gemma3
  853. "vision_tower.transformer.layers.{bid}.feed_forward.down_proj", # pixtral
  854. "visual.blocks.{bid}.mlp.fc2", # qwen2vl
  855. "visual.blocks.{bid}.mlp.down_proj", # qwen2.5vl
  856. ),
  857. MODEL_TENSOR.V_LAYER_SCALE_1: (
  858. "vision_tower.vision_model.encoder.layers.{bid}.ls1", # InternVL
  859. ),
  860. MODEL_TENSOR.V_LAYER_SCALE_2: (
  861. "vision_tower.vision_model.encoder.layers.{bid}.ls2", # InternVL
  862. ),
  863. MODEL_TENSOR.V_PRE_NORM: (
  864. "vision_tower.vision_model.pre_layrnorm",
  865. "vision_tower.ln_pre", # pixtral
  866. ),
  867. MODEL_TENSOR.V_POST_NORM: (
  868. "vision_tower.vision_model.post_layernorm",
  869. "model.vision_model.post_layernorm", # SmolVLM
  870. "visual.merger.ln_q", # qwen2vl
  871. ),
  872. MODEL_TENSOR.V_MM_INP_PROJ: (
  873. "multi_modal_projector.mm_input_projection",
  874. ),
  875. MODEL_TENSOR.V_MM_INP_NORM: (
  876. "multi_modal_projector.norm",
  877. ),
  878. MODEL_TENSOR.V_MM_SOFT_EMB_NORM: (
  879. "multi_modal_projector.mm_soft_emb_norm",
  880. ),
  881. MODEL_TENSOR.V_RESMPL_POS_EMBD_K: (
  882. "resampler.pos_embed_k",
  883. ),
  884. MODEL_TENSOR.V_RESMPL_ATTN_Q: (
  885. "resampler.attn.in_proj_q", # tensor generated from resampler.attn.in_proj
  886. ),
  887. MODEL_TENSOR.V_RESMPL_ATTN_K: (
  888. "resampler.attn.in_proj_k", # tensor generated from resampler.attn.in_proj
  889. ),
  890. MODEL_TENSOR.V_RESMPL_ATTN_V: (
  891. "resampler.attn.in_proj_v", # tensor generated from resampler.attn.in_proj
  892. ),
  893. MODEL_TENSOR.V_RESMPL_ATTN_OUT: (
  894. "resampler.attn.out_proj",
  895. ),
  896. MODEL_TENSOR.V_RESMPL_KV: (
  897. "resampler.kv_proj",
  898. ),
  899. MODEL_TENSOR.V_RESMPL_POST_NORM: (
  900. "resampler.ln_post",
  901. ),
  902. MODEL_TENSOR.V_RESMPL_KV_NORM: (
  903. "resampler.ln_kv",
  904. ),
  905. MODEL_TENSOR.V_RESMPL_Q_NORM: (
  906. "resampler.ln_q",
  907. ),
  908. MODEL_TENSOR.V_RESMPL_PROJ: (
  909. "resampler.proj",
  910. ),
  911. MODEL_TENSOR.V_RESMPL_QUERY: (
  912. "resampler.query",
  913. ),
  914. MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK: (
  915. "v.token_embd.img_break", # for pixtral, this is a generated vector
  916. ),
  917. MODEL_TENSOR.V_MM_PATCH_MERGER: (
  918. "multi_modal_projector.patch_merger.merging_layer", # mistral small 3.1
  919. ),
  920. }
  921. # architecture-specific block mappings
  922. arch_block_mappings_cfg: dict[MODEL_ARCH, dict[MODEL_TENSOR, tuple[str, ...]]] = {
  923. MODEL_ARCH.ARCTIC: {
  924. MODEL_TENSOR.FFN_NORM: (
  925. "model.layers.{bid}.residual_layernorm",
  926. ),
  927. MODEL_TENSOR.FFN_NORM_EXP: (
  928. "model.layers.{bid}.post_attention_layernorm",
  929. ),
  930. },
  931. }
  932. mapping: dict[str, tuple[MODEL_TENSOR, str]]
  933. def __init__(self, arch: MODEL_ARCH, n_blocks: int):
  934. self.mapping = {}
  935. for tensor, keys in self.mappings_cfg.items():
  936. if tensor not in MODEL_TENSORS[arch]:
  937. continue
  938. tensor_name = TENSOR_NAMES[tensor]
  939. self.mapping[tensor_name] = (tensor, tensor_name)
  940. for key in keys:
  941. self.mapping[key] = (tensor, tensor_name)
  942. if arch in self.arch_block_mappings_cfg:
  943. self.block_mappings_cfg.update(self.arch_block_mappings_cfg[arch])
  944. for bid in range(n_blocks):
  945. for tensor, keys in self.block_mappings_cfg.items():
  946. if tensor not in MODEL_TENSORS[arch]:
  947. continue
  948. tensor_name = TENSOR_NAMES[tensor].format(bid = bid)
  949. self.mapping[tensor_name] = (tensor, tensor_name)
  950. for key in keys:
  951. key = key.format(bid = bid)
  952. self.mapping[key] = (tensor, tensor_name)
  953. def get_type_and_name(self, key: str, try_suffixes: Sequence[str] = ()) -> tuple[MODEL_TENSOR, str] | None:
  954. result = self.mapping.get(key)
  955. if result is not None:
  956. return result
  957. for suffix in try_suffixes:
  958. if key.endswith(suffix):
  959. result = self.mapping.get(key[:-len(suffix)])
  960. if result is not None:
  961. return result[0], result[1] + suffix
  962. return None
  963. def get_name(self, key: str, try_suffixes: Sequence[str] = ()) -> str | None:
  964. result = self.get_type_and_name(key, try_suffixes = try_suffixes)
  965. if result is None:
  966. return None
  967. return result[1]
  968. def get_type(self, key: str, try_suffixes: Sequence[str] = ()) -> MODEL_TENSOR | None:
  969. result = self.get_type_and_name(key, try_suffixes = try_suffixes)
  970. if result is None:
  971. return None
  972. return result[0]
  973. def __getitem__(self, key: str) -> str:
  974. try:
  975. return self.mapping[key][1]
  976. except KeyError:
  977. raise KeyError(key)
  978. def __contains__(self, key: str) -> bool:
  979. return key in self.mapping
  980. def __repr__(self) -> str:
  981. return repr(self.mapping)
  982. def get_tensor_name_map(arch: MODEL_ARCH, n_blocks: int) -> TensorNameMap:
  983. return TensorNameMap(arch, n_blocks)