tensor_mapping.py 50 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172
  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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.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. "language_model.model.layers.{bid}.feed_forward.shared_expert.down_proj", # llama4
  393. ),
  394. MODEL_TENSOR.ATTN_Q_NORM: (
  395. "language_model.encoder.layers.{bid}.self_attention.q_layernorm",
  396. "model.layers.{bid}.self_attn.q_layernorm", # persimmon
  397. "model.layers.{bid}.self_attn.q_norm", # cohere olmoe chameleon olmo2
  398. "transformer.blocks.{bid}.attn.q_ln", # sea-lion
  399. "encoder.layer.{bid}.attention.self.layer_norm_q", # jina-bert-v2
  400. "transformer.layers.{bid}.attn.q_norm", # openelm
  401. ),
  402. MODEL_TENSOR.ATTN_K_NORM: (
  403. "language_model.encoder.layers.{bid}.self_attention.k_layernorm",
  404. "model.layers.{bid}.self_attn.k_layernorm", # persimmon
  405. "model.layers.{bid}.self_attn.k_norm", # cohere olmoe chameleon olmo2
  406. "transformer.blocks.{bid}.attn.k_ln", # sea-lion
  407. "encoder.layer.{bid}.attention.self.layer_norm_k", # jina-bert-v2
  408. "transformer.layers.{bid}.attn.k_norm", # openelm
  409. ),
  410. MODEL_TENSOR.ROPE_FREQS: (
  411. "language_model.encoder.layers.{bid}.self_attention.rotary_emb.inv_freq", # persimmon
  412. ),
  413. MODEL_TENSOR.LAYER_OUT_NORM: (
  414. "encoder.layer.{bid}.output.LayerNorm", # bert
  415. "encoder.layers.{bid}.norm2", # nomic-bert
  416. "transformer.decoder_layer.{bid}.rms_norm_3", # Grok
  417. "encoder.layer.{bid}.mlp.layernorm", # jina-bert-v2
  418. "encoder.layer.{bid}.layer_norm_2" # jina-v2-code
  419. ),
  420. MODEL_TENSOR.SSM_IN: (
  421. "model.layers.{bid}.in_proj",
  422. "backbone.layers.{bid}.mixer.in_proj",
  423. ),
  424. MODEL_TENSOR.SSM_CONV1D: (
  425. "model.layers.{bid}.conv1d",
  426. "backbone.layers.{bid}.mixer.conv1d",
  427. ),
  428. MODEL_TENSOR.SSM_X: (
  429. "model.layers.{bid}.x_proj",
  430. "backbone.layers.{bid}.mixer.x_proj",
  431. ),
  432. MODEL_TENSOR.SSM_DT: (
  433. "model.layers.{bid}.dt_proj",
  434. "backbone.layers.{bid}.mixer.dt_proj",
  435. ),
  436. MODEL_TENSOR.SSM_A: (
  437. "model.layers.{bid}.A_log",
  438. "backbone.layers.{bid}.mixer.A_log",
  439. ),
  440. MODEL_TENSOR.SSM_D: (
  441. "model.layers.{bid}.D",
  442. "backbone.layers.{bid}.mixer.D",
  443. ),
  444. MODEL_TENSOR.SSM_OUT: (
  445. "model.layers.{bid}.out_proj",
  446. "backbone.layers.{bid}.mixer.out_proj",
  447. ),
  448. MODEL_TENSOR.TIME_MIX_W0: (
  449. "model.layers.{bid}.attention.w0", # rwkv7
  450. ),
  451. MODEL_TENSOR.TIME_MIX_W1: (
  452. "rwkv.blocks.{bid}.attention.time_maa_w1", # rwkv6
  453. "model.layers.{bid}.self_attn.time_maa_w1", # rwkv6qwen2
  454. "model.layers.{bid}.attention.w1", # rwkv7
  455. ),
  456. MODEL_TENSOR.TIME_MIX_W2: (
  457. "rwkv.blocks.{bid}.attention.time_maa_w2", # rwkv6
  458. "model.layers.{bid}.self_attn.time_maa_w2", # rwkv6qwen2
  459. "model.layers.{bid}.attention.w2", # rwkv7
  460. ),
  461. MODEL_TENSOR.TIME_MIX_A0: (
  462. "model.layers.{bid}.attention.a0", # rwkv7
  463. ),
  464. MODEL_TENSOR.TIME_MIX_A1: (
  465. "model.layers.{bid}.attention.a1", # rwkv7
  466. ),
  467. MODEL_TENSOR.TIME_MIX_A2: (
  468. "model.layers.{bid}.attention.a2", # rwkv7
  469. ),
  470. MODEL_TENSOR.TIME_MIX_V0: (
  471. "model.layers.{bid}.attention.v0", # rwkv7
  472. ),
  473. MODEL_TENSOR.TIME_MIX_V1: (
  474. "model.layers.{bid}.attention.v1", # rwkv7
  475. ),
  476. MODEL_TENSOR.TIME_MIX_V2: (
  477. "model.layers.{bid}.attention.v2", # rwkv7
  478. ),
  479. MODEL_TENSOR.TIME_MIX_G1: (
  480. "model.layers.{bid}.attention.g1", # rwkv7
  481. ),
  482. MODEL_TENSOR.TIME_MIX_G2: (
  483. "model.layers.{bid}.attention.g2", # rwkv7
  484. ),
  485. MODEL_TENSOR.TIME_MIX_K_K: (
  486. "model.layers.{bid}.attention.k_k", # rwkv7
  487. ),
  488. MODEL_TENSOR.TIME_MIX_K_A: (
  489. "model.layers.{bid}.attention.k_a", # rwkv7
  490. ),
  491. MODEL_TENSOR.TIME_MIX_R_K: (
  492. "model.layers.{bid}.attention.r_k", # rwkv7
  493. ),
  494. MODEL_TENSOR.TIME_MIX_LERP_X: (
  495. "rwkv.blocks.{bid}.attention.time_maa_x", # rwkv6
  496. "model.layers.{bid}.self_attn.time_maa_x", # rwkv6qwen2
  497. ),
  498. MODEL_TENSOR.TIME_MIX_LERP_K: (
  499. "rwkv.blocks.{bid}.attention.time_maa_k", # rwkv6
  500. "model.layers.{bid}.self_attn.time_maa_k", # rwkv6qwen2
  501. ),
  502. MODEL_TENSOR.TIME_MIX_LERP_V: (
  503. "rwkv.blocks.{bid}.attention.time_maa_v", # rwkv6
  504. "model.layers.{bid}.self_attn.time_maa_v", # rwkv6qwen2
  505. ),
  506. MODEL_TENSOR.TIME_MIX_LERP_R: (
  507. "rwkv.blocks.{bid}.attention.time_maa_r", # rwkv6
  508. "model.layers.{bid}.self_attn.time_maa_r", # rwkv6qwen2
  509. ),
  510. MODEL_TENSOR.TIME_MIX_LERP_G: (
  511. "rwkv.blocks.{bid}.attention.time_maa_g", # rwkv6
  512. "model.layers.{bid}.self_attn.time_maa_g", # rwkv6qwen2
  513. ),
  514. MODEL_TENSOR.TIME_MIX_LERP_W: (
  515. "rwkv.blocks.{bid}.attention.time_maa_w", # rwkv6
  516. "model.layers.{bid}.self_attn.time_maa_w", # rwkv6qwen2
  517. ),
  518. MODEL_TENSOR.TIME_MIX_FIRST: (
  519. "rwkv.blocks.{bid}.attention.time_faaaa", # rwkv6
  520. ),
  521. MODEL_TENSOR.TIME_MIX_DECAY: (
  522. "rwkv.blocks.{bid}.attention.time_decay", # rwkv6
  523. "model.layers.{bid}.self_attn.time_decay", # rwkv6qwen2
  524. ),
  525. MODEL_TENSOR.TIME_MIX_DECAY_W1: (
  526. "rwkv.blocks.{bid}.attention.time_decay_w1", # rwkv6
  527. "model.layers.{bid}.self_attn.time_decay_w1", # rwkv6qwen2
  528. ),
  529. MODEL_TENSOR.TIME_MIX_DECAY_W2: (
  530. "rwkv.blocks.{bid}.attention.time_decay_w2", # rwkv6
  531. "model.layers.{bid}.self_attn.time_decay_w2", # rwkv6qwen2
  532. ),
  533. MODEL_TENSOR.TIME_MIX_KEY: (
  534. "rwkv.blocks.{bid}.attention.key", # rwkv6
  535. "model.layers.{bid}.self_attn.k_proj", # rwkv6qwen2
  536. "model.layers.{bid}.attention.key", # rwkv7
  537. "model.layers.{bid}.attention.k_proj", # rwkv7
  538. ),
  539. MODEL_TENSOR.TIME_MIX_VALUE: (
  540. "rwkv.blocks.{bid}.attention.value", # rwkv6
  541. "model.layers.{bid}.self_attn.v_proj", # rwkv6qwen2
  542. "model.layers.{bid}.attention.value", # rwkv7
  543. "model.layers.{bid}.attention.v_proj", # rwkv7
  544. ),
  545. MODEL_TENSOR.TIME_MIX_RECEPTANCE: (
  546. "rwkv.blocks.{bid}.attention.receptance", # rwkv6
  547. "model.layers.{bid}.self_attn.q_proj", # rwkv6qwen2
  548. "model.layers.{bid}.attention.receptance", # rwkv7
  549. "model.layers.{bid}.attention.r_proj", # rwkv7
  550. ),
  551. MODEL_TENSOR.TIME_MIX_GATE: (
  552. "rwkv.blocks.{bid}.attention.gate", # rwkv6
  553. "model.layers.{bid}.self_attn.gate", # rwkv6qwen2
  554. ),
  555. MODEL_TENSOR.TIME_MIX_LN: (
  556. "rwkv.blocks.{bid}.attention.ln_x", # rwkv6
  557. "model.layers.{bid}.attention.ln_x" # rwkv7
  558. ),
  559. MODEL_TENSOR.TIME_MIX_OUTPUT: (
  560. "rwkv.blocks.{bid}.attention.output", # rwkv6
  561. "model.layers.{bid}.self_attn.o_proj", # rwkv6qwen2
  562. "model.layers.{bid}.attention.output", # rwkv7
  563. "model.layers.{bid}.attention.o_proj", # rwkv7
  564. ),
  565. MODEL_TENSOR.CHANNEL_MIX_LERP_K: (
  566. "rwkv.blocks.{bid}.feed_forward.time_maa_k", # rwkv6
  567. "model.layers.{bid}.feed_forward.x_k", # rwkv7
  568. ),
  569. MODEL_TENSOR.CHANNEL_MIX_LERP_R: (
  570. "rwkv.blocks.{bid}.feed_forward.time_maa_r", # rwkv6
  571. ),
  572. MODEL_TENSOR.CHANNEL_MIX_KEY: (
  573. "rwkv.blocks.{bid}.feed_forward.key", # rwkv6
  574. "model.layers.{bid}.feed_forward.key", # rwkv7
  575. ),
  576. MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE: (
  577. "rwkv.blocks.{bid}.feed_forward.receptance", # rwkv6
  578. ),
  579. MODEL_TENSOR.CHANNEL_MIX_VALUE: (
  580. "rwkv.blocks.{bid}.feed_forward.value", # rwkv6
  581. "model.layers.{bid}.feed_forward.value", # rwkv7
  582. ),
  583. MODEL_TENSOR.ATTN_Q_A: (
  584. "model.layers.{bid}.self_attn.q_a_proj", # deepseek2
  585. ),
  586. MODEL_TENSOR.ATTN_Q_B: (
  587. "model.layers.{bid}.self_attn.q_b_proj", # deepseek2
  588. ),
  589. MODEL_TENSOR.ATTN_KV_A_MQA: (
  590. "model.layers.{bid}.self_attn.kv_a_proj_with_mqa", # deepseek2
  591. ),
  592. MODEL_TENSOR.ATTN_KV_B: (
  593. "model.layers.{bid}.self_attn.kv_b_proj", # deepseek2
  594. ),
  595. MODEL_TENSOR.ATTN_K_B: (
  596. "model.layers.{bid}.self_attn.k_b_proj", # deepseek2
  597. ),
  598. MODEL_TENSOR.ATTN_V_B: (
  599. "model.layers.{bid}.self_attn.v_b_proj", # deepseek2
  600. ),
  601. MODEL_TENSOR.ATTN_Q_A_NORM: (
  602. "model.layers.{bid}.self_attn.q_a_layernorm", # deepseek2
  603. ),
  604. MODEL_TENSOR.ATTN_KV_A_NORM: (
  605. "model.layers.{bid}.self_attn.kv_a_layernorm", # deepseek2
  606. ),
  607. MODEL_TENSOR.ATTN_SUB_NORM: (
  608. "model.layers.{bid}.self_attn.inner_attn_ln", # bitnet
  609. ),
  610. MODEL_TENSOR.FFN_SUB_NORM: (
  611. "model.layers.{bid}.mlp.ffn_layernorm", # bitnet
  612. ),
  613. MODEL_TENSOR.DEC_ATTN_NORM: (
  614. "decoder.block.{bid}.layer.0.layer_norm", # t5
  615. ),
  616. MODEL_TENSOR.DEC_ATTN_Q: (
  617. "decoder.block.{bid}.layer.0.SelfAttention.q", # t5
  618. ),
  619. MODEL_TENSOR.DEC_ATTN_K: (
  620. "decoder.block.{bid}.layer.0.SelfAttention.k", # t5
  621. ),
  622. MODEL_TENSOR.DEC_ATTN_V: (
  623. "decoder.block.{bid}.layer.0.SelfAttention.v", # t5
  624. ),
  625. MODEL_TENSOR.DEC_ATTN_OUT: (
  626. "decoder.block.{bid}.layer.0.SelfAttention.o", # t5
  627. ),
  628. MODEL_TENSOR.DEC_ATTN_REL_B: (
  629. "decoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5
  630. ),
  631. MODEL_TENSOR.DEC_CROSS_ATTN_NORM: (
  632. "decoder.block.{bid}.layer.1.layer_norm", # t5
  633. ),
  634. MODEL_TENSOR.DEC_CROSS_ATTN_Q: (
  635. "decoder.block.{bid}.layer.1.EncDecAttention.q", # t5
  636. ),
  637. MODEL_TENSOR.DEC_CROSS_ATTN_K: (
  638. "decoder.block.{bid}.layer.1.EncDecAttention.k", # t5
  639. ),
  640. MODEL_TENSOR.DEC_CROSS_ATTN_V: (
  641. "decoder.block.{bid}.layer.1.EncDecAttention.v", # t5
  642. ),
  643. MODEL_TENSOR.DEC_CROSS_ATTN_OUT: (
  644. "decoder.block.{bid}.layer.1.EncDecAttention.o", # t5
  645. ),
  646. MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: (
  647. "decoder.block.{bid}.layer.1.EncDecAttention.relative_attention_bias", # t5
  648. ),
  649. MODEL_TENSOR.DEC_FFN_NORM: (
  650. "decoder.block.{bid}.layer.2.layer_norm", # t5
  651. ),
  652. MODEL_TENSOR.DEC_FFN_GATE: (
  653. "decoder.block.{bid}.layer.2.DenseReluDense.wi_0", # flan-t5
  654. ),
  655. MODEL_TENSOR.DEC_FFN_UP: (
  656. "decoder.block.{bid}.layer.2.DenseReluDense.wi", # t5
  657. "decoder.block.{bid}.layer.2.DenseReluDense.wi_1", # flan-t5
  658. ),
  659. MODEL_TENSOR.DEC_FFN_DOWN: (
  660. "decoder.block.{bid}.layer.2.DenseReluDense.wo", # t5
  661. ),
  662. MODEL_TENSOR.DEC_OUTPUT_NORM: (
  663. "decoder.final_layer_norm", # t5
  664. ),
  665. MODEL_TENSOR.ENC_ATTN_NORM: (
  666. "encoder.block.{bid}.layer.0.layer_norm", # t5
  667. ),
  668. MODEL_TENSOR.ENC_ATTN_Q: (
  669. "encoder.block.{bid}.layer.0.SelfAttention.q", # t5
  670. ),
  671. MODEL_TENSOR.ENC_ATTN_K: (
  672. "encoder.block.{bid}.layer.0.SelfAttention.k", # t5
  673. ),
  674. MODEL_TENSOR.ENC_ATTN_V: (
  675. "encoder.block.{bid}.layer.0.SelfAttention.v", # t5
  676. ),
  677. MODEL_TENSOR.ENC_ATTN_OUT: (
  678. "encoder.block.{bid}.layer.0.SelfAttention.o", # t5
  679. ),
  680. MODEL_TENSOR.ENC_ATTN_REL_B: (
  681. "encoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5
  682. ),
  683. MODEL_TENSOR.ENC_FFN_NORM: (
  684. "encoder.block.{bid}.layer.1.layer_norm", # t5
  685. ),
  686. MODEL_TENSOR.ENC_FFN_GATE: (
  687. "encoder.block.{bid}.layer.1.DenseReluDense.wi_0", # flan-t5
  688. ),
  689. MODEL_TENSOR.ENC_FFN_UP: (
  690. "encoder.block.{bid}.layer.1.DenseReluDense.wi", # t5
  691. "encoder.block.{bid}.layer.1.DenseReluDense.wi_1", # flan-t5
  692. ),
  693. MODEL_TENSOR.ENC_FFN_DOWN: (
  694. "encoder.block.{bid}.layer.1.DenseReluDense.wo", # t5
  695. ),
  696. ############################################################################
  697. # TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg
  698. MODEL_TENSOR.ENC_OUTPUT_NORM: (
  699. "encoder.final_layer_norm", # t5
  700. ),
  701. MODEL_TENSOR.CLS: (
  702. "classifier", # jina
  703. "classifier.dense", # roberta
  704. ),
  705. MODEL_TENSOR.CLS_OUT: (
  706. "classifier.out_proj", # roberta
  707. ),
  708. #############################################################################
  709. MODEL_TENSOR.CONVNEXT_DW: (
  710. "backbone.convnext.{bid}.dwconv", # wavtokenizer
  711. ),
  712. MODEL_TENSOR.CONVNEXT_NORM: (
  713. "backbone.convnext.{bid}.norm", # wavtokenizer
  714. ),
  715. MODEL_TENSOR.CONVNEXT_PW1: (
  716. "backbone.convnext.{bid}.pwconv1", # wavtokenizer
  717. ),
  718. MODEL_TENSOR.CONVNEXT_PW2: (
  719. "backbone.convnext.{bid}.pwconv2", # wavtokenizer
  720. ),
  721. MODEL_TENSOR.CONVNEXT_GAMMA: (
  722. "backbone.convnext.{bid}.gamma", # wavtokenizer
  723. ),
  724. MODEL_TENSOR.POSNET_CONV1: (
  725. "backbone.posnet.{bid}.conv1", # wavtokenizer
  726. ),
  727. MODEL_TENSOR.POSNET_CONV2: (
  728. "backbone.posnet.{bid}.conv2", # wavtokenizer
  729. ),
  730. MODEL_TENSOR.POSNET_NORM: (
  731. "backbone.posnet.{bid}.norm", # wavtokenizer
  732. ),
  733. MODEL_TENSOR.POSNET_NORM1: (
  734. "backbone.posnet.{bid}.norm1", # wavtokenizer
  735. ),
  736. MODEL_TENSOR.POSNET_NORM2: (
  737. "backbone.posnet.{bid}.norm2", # wavtokenizer
  738. ),
  739. MODEL_TENSOR.POSNET_ATTN_NORM: (
  740. "backbone.posnet.{bid}.norm", # wavtokenizer
  741. ),
  742. MODEL_TENSOR.POSNET_ATTN_Q: (
  743. "backbone.posnet.{bid}.q", # wavtokenizer
  744. ),
  745. MODEL_TENSOR.POSNET_ATTN_K: (
  746. "backbone.posnet.{bid}.k", # wavtokenizer
  747. ),
  748. MODEL_TENSOR.POSNET_ATTN_V: (
  749. "backbone.posnet.{bid}.v", # wavtokenizer
  750. ),
  751. MODEL_TENSOR.POSNET_ATTN_OUT: (
  752. "backbone.posnet.{bid}.proj_out", # wavtokenizer
  753. ),
  754. #############################################################################
  755. ## Vision encoder
  756. MODEL_TENSOR.V_MMPROJ: (
  757. "multi_modal_projector.linear_{bid}",
  758. "visual.merger.mlp.{bid}", # qwen2vl
  759. ),
  760. MODEL_TENSOR.V_MMPROJ_FC: (
  761. "model.connector.modality_projection.proj", # SmolVLM
  762. ),
  763. MODEL_TENSOR.V_MMPROJ_MLP: (
  764. "model.mm_projector.mlp.mlp.{bid}",
  765. "mlp1.{bid}", # InternVL
  766. ),
  767. MODEL_TENSOR.V_MMPROJ_PEG: (
  768. "model.mm_projector.peg.peg.{bid}",
  769. ),
  770. MODEL_TENSOR.V_ENC_EMBD_CLS: (
  771. "vision_tower.vision_model.embeddings.class_embedding",
  772. ),
  773. MODEL_TENSOR.V_ENC_EMBD_PATCH: (
  774. "vision_tower.vision_model.embeddings.patch_embedding",
  775. "vpm.embeddings.patch_embedding",
  776. "model.vision_model.embeddings.patch_embedding", # SmolVLM
  777. "vision_tower.patch_conv", # pixtral
  778. "visual.patch_embed.proj", # qwen2vl
  779. ),
  780. MODEL_TENSOR.V_ENC_EMBD_POS: (
  781. "vision_tower.vision_model.embeddings.position_embedding",
  782. "vpm.embeddings.position_embedding",
  783. "model.vision_model.embeddings.position_embedding", # SmolVLM
  784. ),
  785. MODEL_TENSOR.V_ENC_ATTN_Q: (
  786. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.q_proj",
  787. "vpm.encoder.layers.{bid}.self_attn.q_proj",
  788. "model.vision_model.encoder.layers.{bid}.self_attn.q_proj", # SmolVLM
  789. "vision_tower.transformer.layers.{bid}.attention.q_proj", # pixtral
  790. "visual.blocks.{bid}.attn.q", # qwen2vl, generated
  791. ),
  792. MODEL_TENSOR.V_ENC_ATTN_Q_NORM: (
  793. "vision_tower.vision_model.encoder.layers.{bid}.attn.q_norm", # InternVL
  794. ),
  795. MODEL_TENSOR.V_ENC_ATTN_K: (
  796. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.k_proj",
  797. "vpm.encoder.layers.{bid}.self_attn.k_proj",
  798. "model.vision_model.encoder.layers.{bid}.self_attn.k_proj", # SmolVLM
  799. "vision_tower.transformer.layers.{bid}.attention.k_proj", # pixtral
  800. "visual.blocks.{bid}.attn.k", # qwen2vl, generated
  801. ),
  802. MODEL_TENSOR.V_ENC_ATTN_K_NORM: (
  803. "vision_tower.vision_model.encoder.layers.{bid}.attn.k_norm", # InternVL
  804. ),
  805. MODEL_TENSOR.V_ENC_ATTN_V: (
  806. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.v_proj",
  807. "vpm.encoder.layers.{bid}.self_attn.v_proj",
  808. "model.vision_model.encoder.layers.{bid}.self_attn.v_proj", # SmolVLM
  809. "vision_tower.transformer.layers.{bid}.attention.v_proj", # pixtral
  810. "visual.blocks.{bid}.attn.v", # qwen2vl, generated
  811. ),
  812. MODEL_TENSOR.V_ENC_INPUT_NORM: (
  813. "vision_tower.vision_model.encoder.layers.{bid}.layer_norm1",
  814. "vision_tower.vision_model.encoder.layers.{bid}.norm1", # InternVL
  815. "vpm.encoder.layers.{bid}.layer_norm1",
  816. "model.vision_model.encoder.layers.{bid}.layer_norm1", # SmolVLM
  817. "vision_tower.transformer.layers.{bid}.attention_norm", # pixtral
  818. "visual.blocks.{bid}.norm1", # qwen2vl
  819. ),
  820. MODEL_TENSOR.V_ENC_OUTPUT: (
  821. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.out_proj",
  822. "vision_tower.vision_model.encoder.layers.{bid}.attn.proj", # InternVL
  823. "vpm.encoder.layers.{bid}.self_attn.out_proj",
  824. "model.vision_model.encoder.layers.{bid}.self_attn.out_proj", # SmolVLM
  825. "vision_tower.transformer.layers.{bid}.attention.o_proj", # pixtral
  826. "visual.blocks.{bid}.attn.proj", # qwen2vl
  827. ),
  828. MODEL_TENSOR.V_ENC_OUTPUT_NORM: (
  829. "vision_tower.vision_model.encoder.layers.{bid}.layer_norm2",
  830. "vision_tower.vision_model.encoder.layers.{bid}.norm2", # InternVL
  831. "vpm.encoder.layers.{bid}.layer_norm2",
  832. "model.vision_model.encoder.layers.{bid}.layer_norm2", # SmolVLM
  833. "vision_tower.transformer.layers.{bid}.ffn_norm", # pixtral
  834. "visual.blocks.{bid}.norm2", # qwen2vl
  835. ),
  836. MODEL_TENSOR.V_ENC_FFN_UP: (
  837. "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc1",
  838. "vpm.encoder.layers.{bid}.mlp.fc1",
  839. "model.vision_model.encoder.layers.{bid}.mlp.fc1", # SmolVLM, gemma3
  840. "vision_tower.transformer.layers.{bid}.feed_forward.up_proj", # pixtral
  841. "visual.blocks.{bid}.mlp.fc1", # qwen2vl
  842. "visual.blocks.{bid}.mlp.up_proj", # qwen2.5vl
  843. ),
  844. MODEL_TENSOR.V_ENC_FFN_GATE: (
  845. "vision_tower.transformer.layers.{bid}.feed_forward.gate_proj", # pixtral
  846. "visual.blocks.{bid}.mlp.gate_proj", # qwen2.5vl
  847. ),
  848. MODEL_TENSOR.V_ENC_FFN_DOWN: (
  849. "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc2",
  850. "vpm.encoder.layers.{bid}.mlp.fc2",
  851. "model.vision_model.encoder.layers.{bid}.mlp.fc2", # SmolVLM, gemma3
  852. "vision_tower.transformer.layers.{bid}.feed_forward.down_proj", # pixtral
  853. "visual.blocks.{bid}.mlp.fc2", # qwen2vl
  854. "visual.blocks.{bid}.mlp.down_proj", # qwen2.5vl
  855. ),
  856. MODEL_TENSOR.V_LAYER_SCALE_1: (
  857. "vision_tower.vision_model.encoder.layers.{bid}.ls1", # InternVL
  858. ),
  859. MODEL_TENSOR.V_LAYER_SCALE_2: (
  860. "vision_tower.vision_model.encoder.layers.{bid}.ls2", # InternVL
  861. ),
  862. MODEL_TENSOR.V_PRE_NORM: (
  863. "vision_tower.vision_model.pre_layrnorm",
  864. "vision_tower.ln_pre", # pixtral
  865. ),
  866. MODEL_TENSOR.V_POST_NORM: (
  867. "vision_tower.vision_model.post_layernorm",
  868. "model.vision_model.post_layernorm", # SmolVLM
  869. "visual.merger.ln_q", # qwen2vl
  870. ),
  871. MODEL_TENSOR.V_MM_INP_PROJ: (
  872. "multi_modal_projector.mm_input_projection",
  873. ),
  874. MODEL_TENSOR.V_MM_INP_NORM: (
  875. "multi_modal_projector.norm",
  876. ),
  877. MODEL_TENSOR.V_MM_SOFT_EMB_NORM: (
  878. "multi_modal_projector.mm_soft_emb_norm",
  879. ),
  880. MODEL_TENSOR.V_RESMPL_POS_EMBD_K: (
  881. "resampler.pos_embed_k",
  882. ),
  883. MODEL_TENSOR.V_RESMPL_ATTN_Q: (
  884. "resampler.attn.in_proj_q", # tensor generated from resampler.attn.in_proj
  885. ),
  886. MODEL_TENSOR.V_RESMPL_ATTN_K: (
  887. "resampler.attn.in_proj_k", # tensor generated from resampler.attn.in_proj
  888. ),
  889. MODEL_TENSOR.V_RESMPL_ATTN_V: (
  890. "resampler.attn.in_proj_v", # tensor generated from resampler.attn.in_proj
  891. ),
  892. MODEL_TENSOR.V_RESMPL_ATTN_OUT: (
  893. "resampler.attn.out_proj",
  894. ),
  895. MODEL_TENSOR.V_RESMPL_KV: (
  896. "resampler.kv_proj",
  897. ),
  898. MODEL_TENSOR.V_RESMPL_POST_NORM: (
  899. "resampler.ln_post",
  900. ),
  901. MODEL_TENSOR.V_RESMPL_KV_NORM: (
  902. "resampler.ln_kv",
  903. ),
  904. MODEL_TENSOR.V_RESMPL_Q_NORM: (
  905. "resampler.ln_q",
  906. ),
  907. MODEL_TENSOR.V_RESMPL_PROJ: (
  908. "resampler.proj",
  909. ),
  910. MODEL_TENSOR.V_RESMPL_QUERY: (
  911. "resampler.query",
  912. ),
  913. MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK: (
  914. "v.token_embd.img_break", # for pixtral, this is a generated vector
  915. ),
  916. MODEL_TENSOR.V_MM_PATCH_MERGER: (
  917. "multi_modal_projector.patch_merger.merging_layer", # mistral small 3.1
  918. ),
  919. }
  920. # architecture-specific block mappings
  921. arch_block_mappings_cfg: dict[MODEL_ARCH, dict[MODEL_TENSOR, tuple[str, ...]]] = {
  922. MODEL_ARCH.ARCTIC: {
  923. MODEL_TENSOR.FFN_NORM: (
  924. "model.layers.{bid}.residual_layernorm",
  925. ),
  926. MODEL_TENSOR.FFN_NORM_EXP: (
  927. "model.layers.{bid}.post_attention_layernorm",
  928. ),
  929. },
  930. }
  931. mapping: dict[str, tuple[MODEL_TENSOR, str]]
  932. def __init__(self, arch: MODEL_ARCH, n_blocks: int):
  933. self.mapping = {}
  934. for tensor, keys in self.mappings_cfg.items():
  935. if tensor not in MODEL_TENSORS[arch]:
  936. continue
  937. tensor_name = TENSOR_NAMES[tensor]
  938. self.mapping[tensor_name] = (tensor, tensor_name)
  939. for key in keys:
  940. self.mapping[key] = (tensor, tensor_name)
  941. if arch in self.arch_block_mappings_cfg:
  942. self.block_mappings_cfg.update(self.arch_block_mappings_cfg[arch])
  943. for bid in range(n_blocks):
  944. for tensor, keys in self.block_mappings_cfg.items():
  945. if tensor not in MODEL_TENSORS[arch]:
  946. continue
  947. tensor_name = TENSOR_NAMES[tensor].format(bid = bid)
  948. self.mapping[tensor_name] = (tensor, tensor_name)
  949. for key in keys:
  950. key = key.format(bid = bid)
  951. self.mapping[key] = (tensor, tensor_name)
  952. def get_type_and_name(self, key: str, try_suffixes: Sequence[str] = ()) -> tuple[MODEL_TENSOR, str] | None:
  953. result = self.mapping.get(key)
  954. if result is not None:
  955. return result
  956. for suffix in try_suffixes:
  957. if key.endswith(suffix):
  958. result = self.mapping.get(key[:-len(suffix)])
  959. if result is not None:
  960. return result[0], result[1] + suffix
  961. return None
  962. def get_name(self, key: str, try_suffixes: Sequence[str] = ()) -> str | None:
  963. result = self.get_type_and_name(key, try_suffixes = try_suffixes)
  964. if result is None:
  965. return None
  966. return result[1]
  967. def get_type(self, key: str, try_suffixes: Sequence[str] = ()) -> MODEL_TENSOR | None:
  968. result = self.get_type_and_name(key, try_suffixes = try_suffixes)
  969. if result is None:
  970. return None
  971. return result[0]
  972. def __getitem__(self, key: str) -> str:
  973. try:
  974. return self.mapping[key][1]
  975. except KeyError:
  976. raise KeyError(key)
  977. def __contains__(self, key: str) -> bool:
  978. return key in self.mapping
  979. def __repr__(self) -> str:
  980. return repr(self.mapping)
  981. def get_tensor_name_map(arch: MODEL_ARCH, n_blocks: int) -> TensorNameMap:
  982. return TensorNameMap(arch, n_blocks)