tensor_mapping.py 64 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477
  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 plamo2 granite-hybrid
  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. "encoder", # neobert
  31. "model.transformer.wte", # llada
  32. "embed_tokens", # qwen3-embedding
  33. ),
  34. # Token type embeddings
  35. MODEL_TENSOR.TOKEN_TYPES: (
  36. "embeddings.token_type_embeddings", # bert nomic-bert
  37. ),
  38. # Normalization of token embeddings
  39. MODEL_TENSOR.TOKEN_EMBD_NORM: (
  40. "word_embeddings_layernorm", # bloom
  41. "embeddings.LayerNorm", # bert
  42. "emb_ln", # nomic-bert
  43. "transformer.norm", # openelm
  44. "rwkv.blocks.0.pre_ln", # rwkv
  45. "rwkv.blocks.0.pre_ln", # rwkv6
  46. "model.pre_ln", # rwkv7
  47. "model.layers.0.pre_norm", # rwkv7
  48. "backbone.norm", # wavtokenizer
  49. "model.embedding_norm", # lfm2
  50. ),
  51. # Position embeddings
  52. MODEL_TENSOR.POS_EMBD: (
  53. "transformer.wpe", # gpt2
  54. "embeddings.position_embeddings", # bert
  55. "wpe", # gpt2
  56. ),
  57. # Output
  58. MODEL_TENSOR.OUTPUT: (
  59. "embed_out", # gptneox
  60. "lm_head", # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais nemotron exaone olmoe olmo2 phimoe plamo2
  61. "output", # llama-pth bloom internlm2
  62. "word_embeddings_for_head", # persimmon
  63. "lm_head.linear", # phi2
  64. "output_layer", # chatglm
  65. "head", # rwkv
  66. "head.out", # wavtokenizer
  67. "lm_head", # llama4
  68. "model.transformer.ff_out", # llada
  69. ),
  70. # Output norm
  71. MODEL_TENSOR.OUTPUT_NORM: (
  72. "gpt_neox.final_layer_norm", # gptneox
  73. "transformer.ln_f", # gpt2 gpt-j falcon jais exaone
  74. "model.norm", # llama-hf baichuan internlm2 olmoe olmo2 phimoe plamo2
  75. "norm", # llama-pth
  76. "transformer.norm_f", # mpt dbrx
  77. "ln_f", # refact bloom qwen gpt2
  78. "language_model.encoder.final_layernorm", # persimmon
  79. "model.final_layernorm", # persimmon
  80. "lm_head.ln", # phi2
  81. "model.norm_f", # mamba-qbert
  82. "backbone.norm_f", # mamba
  83. "transformer.rms_norm", # Grok
  84. "encoder.final_layernorm", # chatglm
  85. "transformer.norm", # openelm
  86. "model.norm", # nemotron
  87. "rwkv.ln_out", # rwkv6
  88. "model.ln_out", # rwkv7
  89. "backbone.final_layer_norm", # wavtokenizer
  90. "model.norm", # llama4
  91. "model.transformer.ln_f", # llada
  92. ),
  93. # Rope frequencies
  94. MODEL_TENSOR.ROPE_FREQS: (
  95. "rope.freqs", # llama-pth
  96. "rotary_pos_emb.inv_freq", # chatglm
  97. ),
  98. MODEL_TENSOR.ROPE_FACTORS_LONG: (),
  99. MODEL_TENSOR.ROPE_FACTORS_SHORT: (),
  100. MODEL_TENSOR.CONV1D: (
  101. "backbone.embed", # roberta
  102. ),
  103. }
  104. block_mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = {
  105. # Attention norm
  106. MODEL_TENSOR.ATTN_NORM: (
  107. "gpt_neox.layers.{bid}.input_layernorm", # gptneox
  108. "transformer.h.{bid}.ln_1", # gpt2 gpt-j refact qwen jais exaone
  109. "transformer.blocks.{bid}.norm_1", # mpt
  110. "transformer.h.{bid}.input_layernorm", # falcon7b
  111. "h.{bid}.input_layernorm", # bloom
  112. "transformer.h.{bid}.ln_mlp", # falcon40b
  113. "model.layers.{bid}.input_layernorm", # llama-hf nemotron olmoe phimoe granite-hybrid
  114. "layers.{bid}.attention_norm", # llama-pth
  115. "language_model.encoder.layers.{bid}.input_layernorm", # persimmon
  116. "model.layers.{bid}.ln1", # yi
  117. "h.{bid}.ln_1", # gpt2
  118. "transformer.h.{bid}.ln", # phi2
  119. "model.layers.layers.{bid}.norm", # plamo
  120. "model.layers.layers.{bid}.pre_mixer_norm", # plamo2
  121. "model.layers.{bid}.attention_norm", # internlm2
  122. "model.layers.{bid}.norm", # mamba-qbert
  123. "backbone.layers.{bid}.norm", # mamba
  124. "transformer.decoder_layer.{bid}.rms_norm", # Grok
  125. "transformer.blocks.{bid}.norm_attn_norm.norm_1", # dbrx
  126. "encoder.layers.{bid}.input_layernorm", # chatglm
  127. "transformer.layers.{bid}.attn_norm", # openelm
  128. "rwkv.blocks.{bid}.ln1", # rwkv6
  129. "model.layers.{bid}.ln1", # rwkv7
  130. "model.layers.{bid}.input_layernorm", # llama4
  131. "transformer_encoder.{bid}.attention_norm", # neobert
  132. "model.layers.{bid}.operator_norm", # lfm2
  133. "model.transformer.blocks.{bid}.attn_norm", # llada
  134. "layers.{bid}.input_layernorm", # qwen3-embedding
  135. ),
  136. # Attention norm 2
  137. MODEL_TENSOR.ATTN_NORM_2: (
  138. "transformer.h.{bid}.ln_attn", # falcon40b
  139. "encoder.layer.{bid}.layer_norm_1", # jina-v2-code
  140. "rwkv.blocks.{bid}.ln2", # rwkv6
  141. "model.layers.{bid}.ln2", # rwkv7
  142. ),
  143. # Attention query-key-value
  144. MODEL_TENSOR.ATTN_QKV: (
  145. "gpt_neox.layers.{bid}.attention.query_key_value", # gptneox
  146. "transformer.h.{bid}.attn.c_attn", # gpt2 qwen jais
  147. "transformer.blocks.{bid}.attn.Wqkv", # mpt
  148. "transformer.blocks.{bid}.norm_attn_norm.attn.Wqkv", # dbrx
  149. "transformer.h.{bid}.self_attention.query_key_value", # falcon
  150. "h.{bid}.self_attention.query_key_value", # bloom
  151. "language_model.encoder.layers.{bid}.self_attention.query_key_value", # persimmon
  152. "model.layers.{bid}.self_attn.query_key_value", # persimmon
  153. "h.{bid}.attn.c_attn", # gpt2
  154. "transformer.h.{bid}.mixer.Wqkv", # phi2
  155. "encoder.layers.{bid}.attn.Wqkv", # nomic-bert
  156. "encoder.layers.{bid}.mixer.Wqkv", # jina
  157. "model.layers.{bid}.self_attn.qkv_proj", # phi3
  158. "model.layers.layers.{bid}.mixer.qkv_proj", # plamo2
  159. "encoder.layers.{bid}.self_attention.query_key_value", # chatglm
  160. "transformer.layers.{bid}.attn.qkv_proj", # openelm
  161. "transformer_encoder.{bid}.qkv", # neobert
  162. ),
  163. # Attention query
  164. MODEL_TENSOR.ATTN_Q: (
  165. "model.layers.{bid}.self_attn.q_proj", # llama-hf nemotron olmoe olmo2 phimoe
  166. "model.layers.{bid}.self_attn.q_proj_no_perm", # llama-custom
  167. "layers.{bid}.attention.wq", # llama-pth
  168. "encoder.layer.{bid}.attention.self.query", # bert
  169. "transformer.layer.{bid}.attention.q_lin", # distillbert
  170. "transformer.h.{bid}.attn.q_proj", # gpt-j
  171. "model.layers.layers.{bid}.self_attn.q_proj", # plamo
  172. "model.layers.{bid}.attention.wq", # internlm2
  173. "transformer.decoder_layer.{bid}.multi_head_attention.query",# Grok
  174. "transformer.h.{bid}.attn.attention.q_proj", # exaone
  175. "model.layers.{bid}.self_attn.q_proj", # llama4
  176. "model.transformer.blocks.{bid}.q_proj", # llada
  177. "layers.{bid}.self_attn.q_proj", # qwen3-embedding
  178. ),
  179. # Attention key
  180. MODEL_TENSOR.ATTN_K: (
  181. "model.layers.{bid}.self_attn.k_proj", # llama-hf nemotron olmoe olmo2 phimoe
  182. "model.layers.{bid}.self_attn.k_proj_no_perm", # llama-custom
  183. "layers.{bid}.attention.wk", # llama-pth
  184. "encoder.layer.{bid}.attention.self.key", # bert
  185. "transformer.layer.{bid}.attention.k_lin", # distillbert
  186. "transformer.h.{bid}.attn.k_proj", # gpt-j
  187. "transformer.h.{bid}.attn.k", # refact
  188. "model.layers.layers.{bid}.self_attn.k_proj", # plamo
  189. "model.layers.{bid}.attention.wk", # internlm2
  190. "transformer.decoder_layer.{bid}.multi_head_attention.key",# Grok
  191. "transformer.h.{bid}.attn.attention.k_proj", # exaone
  192. "model.layers.{bid}.self_attn.k_proj", # llama4
  193. "model.transformer.blocks.{bid}.k_proj", # llada
  194. "layers.{bid}.self_attn.k_proj", # qwen3-embedding
  195. ),
  196. # Attention value
  197. MODEL_TENSOR.ATTN_V: (
  198. "model.layers.{bid}.self_attn.v_proj", # llama-hf nemotron olmoe olmo2 phimoe
  199. "layers.{bid}.attention.wv", # llama-pth
  200. "encoder.layer.{bid}.attention.self.value", # bert
  201. "transformer.layer.{bid}.attention.v_lin", # distillbert
  202. "transformer.h.{bid}.attn.v_proj", # gpt-j
  203. "transformer.h.{bid}.attn.v", # refact
  204. "model.layers.layers.{bid}.self_attn.v_proj", # plamo
  205. "model.layers.{bid}.attention.wv", # internlm2
  206. "transformer.decoder_layer.{bid}.multi_head_attention.value",# Grok
  207. "transformer.h.{bid}.attn.attention.v_proj", # exaone
  208. "model.layers.{bid}.self_attn.v_proj", # llama4
  209. "model.transformer.blocks.{bid}.v_proj", # llada
  210. "layers.{bid}.self_attn.v_proj", # qwen3-embedding
  211. ),
  212. # Attention output
  213. MODEL_TENSOR.ATTN_OUT: (
  214. "gpt_neox.layers.{bid}.attention.dense", # gptneox
  215. "transformer.h.{bid}.attn.c_proj", # gpt2 refact qwen jais
  216. "transformer.blocks.{bid}.attn.out_proj", # mpt
  217. "transformer.h.{bid}.self_attention.dense", # falcon
  218. "h.{bid}.self_attention.dense", # bloom
  219. "model.layers.{bid}.self_attn.o_proj", # llama-hf nemotron olmoe olmo2 phimoe
  220. "model.layers.{bid}.self_attn.out_proj", # lfm2
  221. "model.layers.{bid}.self_attn.linear_attn", # deci
  222. "layers.{bid}.attention.wo", # llama-pth
  223. "encoder.layer.{bid}.attention.output.dense", # bert
  224. "transformer.layer.{bid}.attention.out_lin", # distillbert
  225. "transformer.h.{bid}.attn.out_proj", # gpt-j
  226. "language_model.encoder.layers.{bid}.self_attention.dense", # persimmon
  227. "model.layers.{bid}.self_attn.dense", # persimmon
  228. "h.{bid}.attn.c_proj", # gpt2
  229. "transformer.h.{bid}.mixer.out_proj", # phi2
  230. "model.layers.layers.{bid}.self_attn.o_proj", # plamo
  231. "model.layers.layers.{bid}.mixer.o_proj", # plamo2
  232. "model.layers.{bid}.attention.wo", # internlm2
  233. "encoder.layers.{bid}.attn.out_proj", # nomic-bert
  234. "encoder.layers.{bid}.mixer.out_proj", # jina
  235. "transformer.decoder_layer.{bid}.multi_head_attention.linear", # Grok
  236. "transformer.blocks.{bid}.norm_attn_norm.attn.out_proj", # dbrx
  237. "encoder.layers.{bid}.self_attention.dense", # chatglm
  238. "transformer.layers.{bid}.attn.out_proj", # openelm
  239. "transformer.h.{bid}.attn.attention.out_proj", # exaone
  240. "model.layers.{bid}.self_attn.o_proj", # llama4
  241. "transformer_encoder.{bid}.wo", # neobert
  242. "model.transformer.blocks.{bid}.attn_out", # llada
  243. "layers.{bid}.self_attn.o_proj", # qwen3-embedding
  244. ),
  245. # Attention output norm
  246. MODEL_TENSOR.ATTN_OUT_NORM: (
  247. "encoder.layer.{bid}.attention.output.LayerNorm", # bert
  248. "transformer.layer.{bid}.sa_layer_norm", # distillbert
  249. "encoder.layers.{bid}.norm1", # nomic-bert
  250. "transformer.decoder_layer.{bid}.rms_norm_1", # Grok
  251. "transformer.blocks.{bid}.norm_attn_norm.norm_2", # dbrx
  252. ),
  253. MODEL_TENSOR.ATTN_POST_NORM: (
  254. "model.layers.{bid}.post_attention_layernorm", # gemma2 olmo2 # ge
  255. "model.layers.{bid}.post_self_attn_layernorm", # glm-4-0414
  256. "model.layers.layers.{bid}.post_mixer_norm.weight", # plamo2
  257. ),
  258. # Rotary embeddings
  259. MODEL_TENSOR.ATTN_ROT_EMBD: (
  260. "model.layers.{bid}.self_attn.rotary_emb.inv_freq", # llama-hf
  261. "layers.{bid}.attention.inner_attention.rope.freqs", # llama-pth
  262. "model.layers.layers.{bid}.self_attn.rotary_emb.inv_freq", # plamo
  263. "transformer.h.{bid}.attn.rotary_emb.inv_freq", # codeshell
  264. ),
  265. MODEL_TENSOR.ATTN_SINKS: (
  266. "model.layers.{bid}.self_attn.sinks", # openai-moe
  267. ),
  268. # Feed-forward norm
  269. MODEL_TENSOR.FFN_NORM: (
  270. "gpt_neox.layers.{bid}.post_attention_layernorm", # gptneox
  271. "transformer.h.{bid}.ln_2", # gpt2 refact qwen jais exaone
  272. "h.{bid}.post_attention_layernorm", # bloom
  273. "transformer.blocks.{bid}.norm_2", # mpt
  274. "model.layers.{bid}.post_attention_layernorm", # llama-hf nemotron olmoe phimoe
  275. "layers.{bid}.ffn_norm", # llama-pth
  276. "language_model.encoder.layers.{bid}.post_attention_layernorm", # persimmon
  277. "model.layers.{bid}.ln2", # yi
  278. "h.{bid}.ln_2", # gpt2
  279. "model.layers.{bid}.ffn_norm", # internlm2
  280. "transformer.decoder_layer.{bid}.rms_norm_2", # Grok
  281. "encoder.layers.{bid}.post_attention_layernorm", # chatglm
  282. "transformer.layers.{bid}.ffn_norm", # openelm
  283. "model.layers.{bid}.pre_ff_layernorm", # jamba granite-hybrid
  284. "model.layers.{bid}.pre_moe_layernorm", # mini-jamba
  285. "model.layers.{bid}.post_attention_layernorm", # llama4
  286. "transformer_encoder.{bid}.ffn_norm", # neobert
  287. "model.layers.layers.{bid}.pre_mlp_norm", # plamo2
  288. "model.transformer.blocks.{bid}.ff_norm", # llada
  289. "layers.{bid}.post_attention_layernorm", # qwen3-embedding
  290. ),
  291. # Post feed-forward norm
  292. MODEL_TENSOR.FFN_PRE_NORM: (
  293. "model.layers.{bid}.pre_feedforward_layernorm", # gemma2
  294. "model.layers.{bid}.pre_ff_layernorm.weight",
  295. ),
  296. # Post feed-forward norm
  297. MODEL_TENSOR.FFN_POST_NORM: (
  298. "model.layers.{bid}.post_feedforward_layernorm", # gemma2 olmo2
  299. "model.layers.{bid}.post_mlp_layernorm", # glm-4-0414
  300. "model.layers.layers.{bid}.post_mlp_norm.weight", # plamo2
  301. "model.layers.{bid}.feed_forward.up_proj",
  302. ),
  303. MODEL_TENSOR.FFN_GATE_INP: (
  304. "layers.{bid}.feed_forward.gate", # mixtral
  305. "model.layers.{bid}.block_sparse_moe.gate", # mixtral phimoe
  306. "model.layers.{bid}.mlp.gate", # qwen2moe olmoe
  307. "transformer.decoder_layer.{bid}.router", # Grok
  308. "transformer.blocks.{bid}.ffn.router.layer", # dbrx
  309. "model.layers.{bid}.block_sparse_moe.router.layer", # granitemoe
  310. "model.layers.{bid}.feed_forward.router", # llama4 jamba
  311. "encoder.layers.{bid}.mlp.router.layer", # nomic-bert-moe
  312. "model.layers.{bid}.mlp.router", # openai-moe
  313. "model.layers.{bid}.mlp.gate.wg", # hunyuan
  314. "model.layers.{bid}.block_sparse_moe.primary_router", # smallthinker
  315. ),
  316. MODEL_TENSOR.FFN_GATE_INP_SHEXP: (
  317. "model.layers.{bid}.mlp.shared_expert_gate", # qwen2moe
  318. ),
  319. MODEL_TENSOR.FFN_EXP_PROBS_B: (
  320. "model.layers.{bid}.mlp.gate.e_score_correction", # deepseek-v3 dots1
  321. "model.layers.{bid}.mlp.moe_statics.e_score_correction", # ernie4.5-moe
  322. ),
  323. # Feed-forward up
  324. MODEL_TENSOR.FFN_UP: (
  325. "gpt_neox.layers.{bid}.mlp.dense_h_to_4h", # gptneox
  326. "transformer.h.{bid}.mlp.c_fc", # gpt2 jais
  327. "transformer.blocks.{bid}.ffn.up_proj", # mpt
  328. "transformer.h.{bid}.mlp.dense_h_to_4h", # falcon
  329. "h.{bid}.mlp.dense_h_to_4h", # bloom
  330. "model.layers.{bid}.mlp.up_proj", # llama-hf refact nemotron olmo2
  331. "layers.{bid}.feed_forward.w3", # llama-pth
  332. "encoder.layer.{bid}.intermediate.dense", # bert
  333. "transformer.layer.{bid}.ffn.lin1", # distillbert
  334. "transformer.h.{bid}.mlp.fc_in", # gpt-j
  335. "transformer.h.{bid}.mlp.linear_3", # refact
  336. "language_model.encoder.layers.{bid}.mlp.dense_h_to_4h", # persimmon
  337. "model.layers.{bid}.mlp.dense_h_to_4h", # persimmon
  338. "transformer.h.{bid}.mlp.w1", # qwen
  339. "h.{bid}.mlp.c_fc", # gpt2
  340. "transformer.h.{bid}.mlp.fc1", # phi2
  341. "model.layers.{bid}.mlp.fc1", # phi2
  342. "model.layers.{bid}.mlp.gate_up_proj", # phi3 glm-4-0414
  343. "model.layers.layers.{bid}.mlp.up_proj", # plamo
  344. "model.layers.layers.{bid}.mlp.gate_up_proj", # plamo2
  345. "model.layers.{bid}.feed_forward.w3", # internlm2
  346. "encoder.layers.{bid}.mlp.fc11", # nomic-bert
  347. "encoder.layers.{bid}.mlp.fc1", # nomic-bert-moe
  348. "model.layers.{bid}.mlp.c_fc", # starcoder2
  349. "encoder.layer.{bid}.mlp.gated_layers_v", # jina-bert-v2 (split up/gate, no longer used)
  350. "encoder.layer.{bid}.mlp.gated_layers", # jina-bert-v2 (GEGLU)
  351. "encoder.layer.{bid}.mlp.up_gated_layer", # jina-v2-code (GEGLU)
  352. "model.layers.{bid}.residual_mlp.w3", # arctic
  353. "encoder.layers.{bid}.mlp.dense_h_to_4h", # chatglm
  354. "transformer.h.{bid}.mlp.c_fc_1", # exaone
  355. "model.layers.{bid}.feed_forward.up_proj", # llama4 jamba granite-hybrid
  356. "transformer_encoder.{bid}.ffn.w12", # neobert
  357. "model.layers.{bid}.block_sparse_moe.up", # smallthinker
  358. "model.transformer.blocks.{bid}.up_proj", # llada
  359. "layers.{bid}.mlp.up_proj", # qwen3-embedding
  360. ),
  361. MODEL_TENSOR.FFN_UP_EXP: (
  362. "layers.{bid}.feed_forward.experts.w3", # mixtral (merged)
  363. "transformer.decoder_layer.{bid}.moe.linear_v", # Grok (merged)
  364. "transformer.blocks.{bid}.ffn.experts.mlp.v1", # dbrx
  365. "model.layers.{bid}.mlp.experts.up_proj", # qwen2moe olmoe (merged) ernie4.5-moe
  366. "model.layers.{bid}.block_sparse_moe.experts.w3", # phimoe (merged)
  367. "model.layers.{bid}.feed_forward.experts.up_proj", # llama4
  368. "encoder.layers.{bid}.mlp.experts.mlp.w1", # nomic-bert-moe
  369. "model.layers.{bid}.block_sparse_moe.experts.up", # smallthinker
  370. ),
  371. MODEL_TENSOR.FFN_UP_SHEXP: (
  372. "model.layers.{bid}.mlp.shared_expert.up_proj", # qwen2moe
  373. "model.layers.{bid}.mlp.shared_experts.up_proj", # deepseek deepseek2
  374. "model.layers.{bid}.feed_forward.shared_expert.up_proj", # llama4
  375. "model.layers.{bid}.feed_forward.down_proj",
  376. "model.layers.{bid}.mlp.shared_mlp.up_proj", # hunyuan
  377. ),
  378. # AWQ-activation gate
  379. MODEL_TENSOR.FFN_ACT: (
  380. "transformer.blocks.{bid}.ffn.act", # mpt
  381. ),
  382. # Feed-forward gate
  383. MODEL_TENSOR.FFN_GATE: (
  384. "model.layers.{bid}.mlp.gate_proj", # llama-hf refact olmo2
  385. "layers.{bid}.feed_forward.w1", # llama-pth
  386. "transformer.h.{bid}.mlp.w2", # qwen
  387. "transformer.h.{bid}.mlp.c_fc2", # jais
  388. "model.layers.layers.{bid}.mlp.gate_proj", # plamo
  389. "model.layers.{bid}.feed_forward.w1", # internlm2
  390. "encoder.layers.{bid}.mlp.fc12", # nomic-bert
  391. "encoder.layer.{bid}.mlp.gated_layers_w", # jina-bert-v2 (split up/gate, no longer used)
  392. "transformer.h.{bid}.mlp.linear_1", # refact
  393. "model.layers.{bid}.residual_mlp.w1", # arctic
  394. "transformer.h.{bid}.mlp.c_fc_0", # exaone
  395. "model.layers.{bid}.feed_forward.gate_proj", # llama4 jamba granite-hybrid
  396. "model.layers.{bid}.block_sparse_moe.gate", # smallthinker
  397. "model.transformer.blocks.{bid}.ff_proj", # llada
  398. "layers.{bid}.mlp.gate_proj", # qwen3-embedding
  399. ),
  400. MODEL_TENSOR.FFN_GATE_EXP: (
  401. "layers.{bid}.feed_forward.experts.w1", # mixtral (merged)
  402. "transformer.decoder_layer.{bid}.moe.linear", # Grok (merged)
  403. "transformer.blocks.{bid}.ffn.experts.mlp.w1", # dbrx
  404. "model.layers.{bid}.mlp.experts.gate_proj", # qwen2moe olmoe (merged) ernie4.5-moe
  405. "model.layers.{bid}.block_sparse_moe.experts.w1", # phimoe (merged)
  406. "model.layers.{bid}.feed_forward.experts.gate_proj", # llama4
  407. "model.layers.{bid}.block_sparse_moe.experts.gate", # smallthinker
  408. ),
  409. MODEL_TENSOR.FFN_GATE_SHEXP: (
  410. "model.layers.{bid}.mlp.shared_expert.gate_proj", # qwen2moe
  411. "model.layers.{bid}.mlp.shared_experts.gate_proj", # deepseek deepseek2
  412. "model.layers.{bid}.feed_forward.shared_expert.gate_proj", # llama4
  413. "model.layers.{bid}.mlp.shared_mlp.gate_proj", # hunyuan
  414. ),
  415. # Feed-forward down
  416. MODEL_TENSOR.FFN_DOWN: (
  417. "gpt_neox.layers.{bid}.mlp.dense_4h_to_h", # gptneox
  418. "transformer.h.{bid}.mlp.c_proj", # gpt2 refact qwen jais
  419. "transformer.blocks.{bid}.ffn.down_proj", # mpt
  420. "transformer.h.{bid}.mlp.dense_4h_to_h", # falcon
  421. "h.{bid}.mlp.dense_4h_to_h", # bloom
  422. "model.layers.{bid}.mlp.down_proj", # llama-hf nemotron olmo2
  423. "layers.{bid}.feed_forward.w2", # llama-pth
  424. "encoder.layer.{bid}.output.dense", # bert
  425. "transformer.layer.{bid}.ffn.lin2", # distillbert
  426. "transformer.h.{bid}.mlp.fc_out", # gpt-j
  427. "language_model.encoder.layers.{bid}.mlp.dense_4h_to_h", # persimmon
  428. "model.layers.{bid}.mlp.dense_4h_to_h", # persimmon
  429. "h.{bid}.mlp.c_proj", # gpt2
  430. "transformer.h.{bid}.mlp.fc2", # phi2
  431. "model.layers.{bid}.mlp.fc2", # phi2
  432. "model.layers.layers.{bid}.mlp.down_proj", # plamo
  433. "model.layers.{bid}.feed_forward.w2", # internlm2
  434. "encoder.layers.{bid}.mlp.fc2", # nomic-bert
  435. "model.layers.{bid}.mlp.c_proj", # starcoder2
  436. "encoder.layer.{bid}.mlp.wo", # jina-bert-v2
  437. "transformer.layers.{bid}.ffn.proj_2", # openelm
  438. "model.layers.{bid}.residual_mlp.w2", # arctic
  439. "encoder.layer.{bid}.mlp.down_layer", # jina-bert-v2
  440. "encoder.layers.{bid}.mlp.dense_4h_to_h", # chatglm
  441. "model.layers.h.{bid}.mlp.c_proj", # exaone
  442. "model.layers.{bid}.feed_forward.down_proj", # llama4 jamba granite-hybrid
  443. "transformer_encoder.{bid}.ffn.w3", # neobert
  444. "model.layers.{bid}.block_sparse_moe.down", # smallthinker
  445. "model.transformer.blocks.{bid}.ff_out", # llada
  446. "layers.{bid}.mlp.down_proj", # qwen3-embedding
  447. ),
  448. MODEL_TENSOR.FFN_DOWN_EXP: (
  449. "layers.{bid}.feed_forward.experts.w2", # mixtral (merged)
  450. "transformer.decoder_layer.{bid}.moe.linear_1", # Grok (merged)
  451. "transformer.blocks.{bid}.ffn.experts.mlp.w2", # dbrx
  452. "model.layers.{bid}.mlp.experts.down_proj", # qwen2moe olmoe (merged) ernie4.5-moe
  453. "model.layers.{bid}.block_sparse_moe.output_linear", # granitemoe
  454. "model.layers.{bid}.block_sparse_moe.experts.w2", # phimoe (merged)
  455. "model.layers.{bid}.feed_forward.experts.down_proj", # llama4
  456. "encoder.layers.{bid}.mlp.experts.mlp.w2", # nomic-bert-moe
  457. "model.layers.{bid}.block_sparse_moe.experts.down", # smallthinker
  458. ),
  459. MODEL_TENSOR.FFN_DOWN_SHEXP: (
  460. "model.layers.{bid}.mlp.shared_expert.down_proj", # qwen2moe
  461. "model.layers.{bid}.mlp.shared_experts.down_proj", # deepseek deepseek2
  462. "model.layers.{bid}.feed_forward.shared_expert.down_proj", # llama4
  463. "model.layers.{bid}.shared_mlp.output_linear", # granitemoe
  464. "model.layers.{bid}.mlp.shared_mlp.down_proj", # hunyuan
  465. ),
  466. MODEL_TENSOR.ATTN_Q_NORM: (
  467. "language_model.encoder.layers.{bid}.self_attention.q_layernorm",
  468. "model.layers.{bid}.self_attn.q_layernorm", # persimmon
  469. "model.layers.{bid}.self_attn.query_layernorm", # hunyuan
  470. "model.layers.{bid}.self_attn.q_norm", # cohere olmoe chameleon olmo2
  471. "transformer.blocks.{bid}.attn.q_ln", # sea-lion
  472. "encoder.layer.{bid}.attention.self.layer_norm_q", # jina-bert-v2
  473. "transformer.layers.{bid}.attn.q_norm", # openelm
  474. "model.layers.layers.{bid}.mixer.q", # plamo2
  475. "layers.{bid}.self_attn.q_norm", # qwen3-embedding
  476. ),
  477. MODEL_TENSOR.ATTN_K_NORM: (
  478. "language_model.encoder.layers.{bid}.self_attention.k_layernorm",
  479. "model.layers.{bid}.self_attn.k_layernorm", # persimmon
  480. "model.layers.{bid}.self_attn.key_layernorm", # hunyuan
  481. "model.layers.{bid}.self_attn.k_norm", # cohere olmoe chameleon olmo2
  482. "transformer.blocks.{bid}.attn.k_ln", # sea-lion
  483. "encoder.layer.{bid}.attention.self.layer_norm_k", # jina-bert-v2
  484. "transformer.layers.{bid}.attn.k_norm", # openelm
  485. "model.layers.layers.{bid}.mixer.k", # plamo2
  486. "layers.{bid}.self_attn.k_norm", # qwen3-embedding
  487. ),
  488. MODEL_TENSOR.ROPE_FREQS: (
  489. "language_model.encoder.layers.{bid}.self_attention.rotary_emb.inv_freq", # persimmon
  490. ),
  491. MODEL_TENSOR.LAYER_OUT_NORM: (
  492. "encoder.layer.{bid}.output.LayerNorm", # bert
  493. "transformer.layer.{bid}.output_layer_norm", # distillbert
  494. "encoder.layers.{bid}.norm2", # nomic-bert
  495. "transformer.decoder_layer.{bid}.rms_norm_3", # Grok
  496. "encoder.layer.{bid}.mlp.layernorm", # jina-bert-v2
  497. "encoder.layer.{bid}.layer_norm_2", # jina-v2-code
  498. ),
  499. MODEL_TENSOR.PER_LAYER_TOKEN_EMBD: (
  500. "model.embed_tokens_per_layer", # gemma3n
  501. ),
  502. MODEL_TENSOR.PER_LAYER_MODEL_PROJ: (
  503. "model.per_layer_model_projection", # gemma3n
  504. ),
  505. MODEL_TENSOR.PER_LAYER_PROJ_NORM: (
  506. "model.per_layer_projection_norm", # gemma3n
  507. ),
  508. MODEL_TENSOR.ALTUP_PROJ: (
  509. "model.altup_projections", # gemma3n
  510. ),
  511. MODEL_TENSOR.ALTUP_UNEMBD_PROJ: (
  512. "model.altup_unembed_projections", # gemma3n
  513. ),
  514. MODEL_TENSOR.PER_LAYER_INP_GATE: (
  515. "model.layers.{bid}.per_layer_input_gate", # gemma3n
  516. ),
  517. MODEL_TENSOR.PER_LAYER_PROJ: (
  518. "model.layers.{bid}.per_layer_projection", # gemma3n
  519. ),
  520. MODEL_TENSOR.PER_LAYER_POST_NORM: (
  521. "model.layers.{bid}.post_per_layer_input_norm", # gemma3n
  522. ),
  523. MODEL_TENSOR.ALTUP_CORRECT_COEF: (
  524. "model.layers.{bid}.altup.correction_coefs", # gemma3n
  525. ),
  526. MODEL_TENSOR.ALTUP_CORRECT_SCALE: (
  527. "model.layers.{bid}.altup.correct_output_scale", # gemma3n
  528. ),
  529. MODEL_TENSOR.ALTUP_PREDICT_COEF: (
  530. "model.layers.{bid}.altup.prediction_coefs", # gemma3n
  531. ),
  532. MODEL_TENSOR.ALTUP_ROUTER: (
  533. "model.layers.{bid}.altup.modality_router", # gemma3n
  534. ),
  535. MODEL_TENSOR.ALTUP_ROUTER_NORM: (
  536. "model.layers.{bid}.altup.router_norm", # gemma3n
  537. ),
  538. MODEL_TENSOR.LAUREL_L: (
  539. "model.layers.{bid}.laurel.linear_left", # gemma3n
  540. ),
  541. MODEL_TENSOR.LAUREL_R: (
  542. "model.layers.{bid}.laurel.linear_right", # gemma3n
  543. ),
  544. MODEL_TENSOR.LAUREL_POST_NORM: (
  545. "model.layers.{bid}.laurel.post_laurel_norm", # gemma3n
  546. ),
  547. MODEL_TENSOR.SSM_IN: (
  548. "model.layers.{bid}.in_proj", # mamba-hf
  549. "backbone.layers.{bid}.mixer.in_proj", # mamba
  550. "model.layers.{bid}.mamba.in_proj", # jamba falcon-h1 granite-hybrid
  551. "model.layers.layers.{bid}.mixer.in_proj", # plamo2
  552. ),
  553. MODEL_TENSOR.SSM_CONV1D: (
  554. "model.layers.{bid}.conv1d", # mamba-hf
  555. "backbone.layers.{bid}.mixer.conv1d", # mamba
  556. "model.layers.{bid}.mamba.conv1d", # jamba falcon-h1 granite-hybrid
  557. "model.layers.layers.{bid}.mixer.conv1d", # plamo2
  558. ),
  559. MODEL_TENSOR.SSM_X: (
  560. "model.layers.{bid}.x_proj", # mamba-hf
  561. "backbone.layers.{bid}.mixer.x_proj", # mamba
  562. "model.layers.{bid}.mamba.x_proj", # jamba
  563. "model.layers.layers.{bid}.mixer.bcdt_proj", # plamo2
  564. ),
  565. MODEL_TENSOR.SSM_DT: (
  566. "model.layers.{bid}.dt_proj", # mamba-hf
  567. "backbone.layers.{bid}.mixer.dt_proj", # mamba
  568. "model.layers.{bid}.mamba.dt_proj", # jamba falcon-h1 granite-hybrid
  569. "model.layers.layers.{bid}.mixer.dt_proj", # plamo2
  570. ),
  571. MODEL_TENSOR.SSM_DT_NORM: (
  572. "model.layers.layers.{bid}.mixer.dt_norm.weight", # plamo2
  573. "model.layers.{bid}.mamba.dt_layernorm", # jamba
  574. ),
  575. MODEL_TENSOR.SSM_A: (
  576. "model.layers.{bid}.A_log", # mamba-hf
  577. "backbone.layers.{bid}.mixer.A_log", # mamba
  578. "model.layers.{bid}.mamba.A_log", # jamba falcon-h1 granite-hybrid
  579. "model.layers.layers.{bid}.mixer.A_log", # plamo2
  580. ),
  581. MODEL_TENSOR.SSM_B_NORM: (
  582. "model.layers.{bid}.mamba.b_layernorm", # jamba
  583. "model.layers.{bid}.mamba.B_layernorm", # mini-jamba
  584. "model.layers.layers.{bid}.mixer.B_norm.weight", # plamo2
  585. ),
  586. MODEL_TENSOR.SSM_C_NORM: (
  587. "model.layers.{bid}.mamba.c_layernorm", # jamba
  588. "model.layers.{bid}.mamba.C_layernorm", # mini-jamba
  589. "model.layers.layers.{bid}.mixer.C_norm.weight", # plamo2
  590. ),
  591. MODEL_TENSOR.SSM_D: (
  592. "model.layers.{bid}.D", # mamba-hf
  593. "backbone.layers.{bid}.mixer.D", # mamba
  594. "model.layers.{bid}.mamba.D", # jamba falcon-h1 granite-hybrid
  595. "model.layers.layers.{bid}.mixer.D", # plamo2
  596. ),
  597. MODEL_TENSOR.SSM_NORM: (
  598. "model.layers.{bid}.mamba.norm", # falcon-h1 granite-hybrid
  599. "backbone.layers.{bid}.mixer.norm", # mamba2
  600. ),
  601. MODEL_TENSOR.SSM_OUT: (
  602. "model.layers.{bid}.out_proj", # mamba-hf
  603. "backbone.layers.{bid}.mixer.out_proj", # mamba
  604. "model.layers.{bid}.mamba.out_proj", # jamba falcon-h1 granite-hybrid
  605. "model.layers.layers.{bid}.mixer.out_proj", # plamo2
  606. ),
  607. MODEL_TENSOR.TIME_MIX_W0: (
  608. "model.layers.{bid}.attention.w0", # rwkv7
  609. ),
  610. MODEL_TENSOR.TIME_MIX_W1: (
  611. "rwkv.blocks.{bid}.attention.time_maa_w1", # rwkv6
  612. "model.layers.{bid}.self_attn.time_maa_w1", # rwkv6qwen2
  613. "model.layers.{bid}.attention.w1", # rwkv7
  614. ),
  615. MODEL_TENSOR.TIME_MIX_W2: (
  616. "rwkv.blocks.{bid}.attention.time_maa_w2", # rwkv6
  617. "model.layers.{bid}.self_attn.time_maa_w2", # rwkv6qwen2
  618. "model.layers.{bid}.attention.w2", # rwkv7
  619. ),
  620. MODEL_TENSOR.TIME_MIX_A0: (
  621. "model.layers.{bid}.attention.a0", # rwkv7
  622. ),
  623. MODEL_TENSOR.TIME_MIX_A1: (
  624. "model.layers.{bid}.attention.a1", # rwkv7
  625. ),
  626. MODEL_TENSOR.TIME_MIX_A2: (
  627. "model.layers.{bid}.attention.a2", # rwkv7
  628. ),
  629. MODEL_TENSOR.TIME_MIX_V0: (
  630. "model.layers.{bid}.attention.v0", # rwkv7
  631. ),
  632. MODEL_TENSOR.TIME_MIX_V1: (
  633. "model.layers.{bid}.attention.v1", # rwkv7
  634. ),
  635. MODEL_TENSOR.TIME_MIX_V2: (
  636. "model.layers.{bid}.attention.v2", # rwkv7
  637. ),
  638. MODEL_TENSOR.TIME_MIX_G1: (
  639. "model.layers.{bid}.attention.g1", # rwkv7
  640. ),
  641. MODEL_TENSOR.TIME_MIX_G2: (
  642. "model.layers.{bid}.attention.g2", # rwkv7
  643. ),
  644. MODEL_TENSOR.TIME_MIX_K_K: (
  645. "model.layers.{bid}.attention.k_k", # rwkv7
  646. ),
  647. MODEL_TENSOR.TIME_MIX_K_A: (
  648. "model.layers.{bid}.attention.k_a", # rwkv7
  649. ),
  650. MODEL_TENSOR.TIME_MIX_R_K: (
  651. "model.layers.{bid}.attention.r_k", # rwkv7
  652. ),
  653. MODEL_TENSOR.TIME_MIX_LERP_X: (
  654. "rwkv.blocks.{bid}.attention.time_maa_x", # rwkv6
  655. "model.layers.{bid}.self_attn.time_maa_x", # rwkv6qwen2
  656. ),
  657. MODEL_TENSOR.TIME_MIX_LERP_K: (
  658. "rwkv.blocks.{bid}.attention.time_maa_k", # rwkv6
  659. "model.layers.{bid}.self_attn.time_maa_k", # rwkv6qwen2
  660. ),
  661. MODEL_TENSOR.TIME_MIX_LERP_V: (
  662. "rwkv.blocks.{bid}.attention.time_maa_v", # rwkv6
  663. "model.layers.{bid}.self_attn.time_maa_v", # rwkv6qwen2
  664. ),
  665. MODEL_TENSOR.TIME_MIX_LERP_R: (
  666. "rwkv.blocks.{bid}.attention.time_maa_r", # rwkv6
  667. "model.layers.{bid}.self_attn.time_maa_r", # rwkv6qwen2
  668. ),
  669. MODEL_TENSOR.TIME_MIX_LERP_G: (
  670. "rwkv.blocks.{bid}.attention.time_maa_g", # rwkv6
  671. "model.layers.{bid}.self_attn.time_maa_g", # rwkv6qwen2
  672. ),
  673. MODEL_TENSOR.TIME_MIX_LERP_W: (
  674. "rwkv.blocks.{bid}.attention.time_maa_w", # rwkv6
  675. "model.layers.{bid}.self_attn.time_maa_w", # rwkv6qwen2
  676. ),
  677. MODEL_TENSOR.TIME_MIX_FIRST: (
  678. "rwkv.blocks.{bid}.attention.time_faaaa", # rwkv6
  679. ),
  680. MODEL_TENSOR.TIME_MIX_DECAY: (
  681. "rwkv.blocks.{bid}.attention.time_decay", # rwkv6
  682. "model.layers.{bid}.self_attn.time_decay", # rwkv6qwen2
  683. ),
  684. MODEL_TENSOR.TIME_MIX_DECAY_W1: (
  685. "rwkv.blocks.{bid}.attention.time_decay_w1", # rwkv6
  686. "model.layers.{bid}.self_attn.time_decay_w1", # rwkv6qwen2
  687. ),
  688. MODEL_TENSOR.TIME_MIX_DECAY_W2: (
  689. "rwkv.blocks.{bid}.attention.time_decay_w2", # rwkv6
  690. "model.layers.{bid}.self_attn.time_decay_w2", # rwkv6qwen2
  691. ),
  692. MODEL_TENSOR.TIME_MIX_KEY: (
  693. "rwkv.blocks.{bid}.attention.key", # rwkv6
  694. "model.layers.{bid}.self_attn.k_proj", # rwkv6qwen2
  695. "model.layers.{bid}.attention.key", # rwkv7
  696. "model.layers.{bid}.attention.k_proj", # rwkv7
  697. ),
  698. MODEL_TENSOR.TIME_MIX_VALUE: (
  699. "rwkv.blocks.{bid}.attention.value", # rwkv6
  700. "model.layers.{bid}.self_attn.v_proj", # rwkv6qwen2
  701. "model.layers.{bid}.attention.value", # rwkv7
  702. "model.layers.{bid}.attention.v_proj", # rwkv7
  703. ),
  704. MODEL_TENSOR.TIME_MIX_RECEPTANCE: (
  705. "rwkv.blocks.{bid}.attention.receptance", # rwkv6
  706. "model.layers.{bid}.self_attn.q_proj", # rwkv6qwen2
  707. "model.layers.{bid}.attention.receptance", # rwkv7
  708. "model.layers.{bid}.attention.r_proj", # rwkv7
  709. ),
  710. MODEL_TENSOR.TIME_MIX_GATE: (
  711. "rwkv.blocks.{bid}.attention.gate", # rwkv6
  712. "model.layers.{bid}.self_attn.gate", # rwkv6qwen2
  713. ),
  714. MODEL_TENSOR.TIME_MIX_LN: (
  715. "rwkv.blocks.{bid}.attention.ln_x", # rwkv6
  716. "model.layers.{bid}.attention.ln_x" # rwkv7
  717. ),
  718. MODEL_TENSOR.TIME_MIX_OUTPUT: (
  719. "rwkv.blocks.{bid}.attention.output", # rwkv6
  720. "model.layers.{bid}.self_attn.o_proj", # rwkv6qwen2
  721. "model.layers.{bid}.attention.output", # rwkv7
  722. "model.layers.{bid}.attention.o_proj", # rwkv7
  723. ),
  724. MODEL_TENSOR.CHANNEL_MIX_LERP_K: (
  725. "rwkv.blocks.{bid}.feed_forward.time_maa_k", # rwkv6
  726. "model.layers.{bid}.feed_forward.x_k", # rwkv7
  727. ),
  728. MODEL_TENSOR.CHANNEL_MIX_LERP_R: (
  729. "rwkv.blocks.{bid}.feed_forward.time_maa_r", # rwkv6
  730. ),
  731. MODEL_TENSOR.CHANNEL_MIX_KEY: (
  732. "rwkv.blocks.{bid}.feed_forward.key", # rwkv6
  733. "model.layers.{bid}.feed_forward.key", # rwkv7
  734. ),
  735. MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE: (
  736. "rwkv.blocks.{bid}.feed_forward.receptance", # rwkv6
  737. ),
  738. MODEL_TENSOR.CHANNEL_MIX_VALUE: (
  739. "rwkv.blocks.{bid}.feed_forward.value", # rwkv6
  740. "model.layers.{bid}.feed_forward.value", # rwkv7
  741. ),
  742. MODEL_TENSOR.ATTN_Q_A: (
  743. "model.layers.{bid}.self_attn.q_a_proj", # deepseek2
  744. ),
  745. MODEL_TENSOR.ATTN_Q_B: (
  746. "model.layers.{bid}.self_attn.q_b_proj", # deepseek2
  747. ),
  748. MODEL_TENSOR.ATTN_KV_A_MQA: (
  749. "model.layers.{bid}.self_attn.kv_a_proj_with_mqa", # deepseek2
  750. ),
  751. MODEL_TENSOR.ATTN_KV_B: (
  752. "model.layers.{bid}.self_attn.kv_b_proj", # deepseek2
  753. ),
  754. MODEL_TENSOR.ATTN_K_B: (
  755. "model.layers.{bid}.self_attn.k_b_proj", # deepseek2
  756. ),
  757. MODEL_TENSOR.ATTN_V_B: (
  758. "model.layers.{bid}.self_attn.v_b_proj", # deepseek2
  759. ),
  760. MODEL_TENSOR.ATTN_Q_A_NORM: (
  761. "model.layers.{bid}.self_attn.q_a_layernorm", # deepseek2
  762. ),
  763. MODEL_TENSOR.ATTN_KV_A_NORM: (
  764. "model.layers.{bid}.self_attn.kv_a_layernorm", # deepseek2
  765. ),
  766. MODEL_TENSOR.ATTN_SUB_NORM: (
  767. "model.layers.{bid}.self_attn.inner_attn_ln", # bitnet
  768. ),
  769. MODEL_TENSOR.FFN_SUB_NORM: (
  770. "model.layers.{bid}.mlp.ffn_layernorm", # bitnet
  771. ),
  772. MODEL_TENSOR.DEC_ATTN_NORM: (
  773. "decoder.block.{bid}.layer.0.layer_norm", # t5
  774. ),
  775. MODEL_TENSOR.DEC_ATTN_Q: (
  776. "decoder.block.{bid}.layer.0.SelfAttention.q", # t5
  777. ),
  778. MODEL_TENSOR.DEC_ATTN_K: (
  779. "decoder.block.{bid}.layer.0.SelfAttention.k", # t5
  780. ),
  781. MODEL_TENSOR.DEC_ATTN_V: (
  782. "decoder.block.{bid}.layer.0.SelfAttention.v", # t5
  783. ),
  784. MODEL_TENSOR.DEC_ATTN_OUT: (
  785. "decoder.block.{bid}.layer.0.SelfAttention.o", # t5
  786. ),
  787. MODEL_TENSOR.DEC_ATTN_REL_B: (
  788. "decoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5
  789. ),
  790. MODEL_TENSOR.DEC_CROSS_ATTN_NORM: (
  791. "decoder.block.{bid}.layer.1.layer_norm", # t5
  792. ),
  793. MODEL_TENSOR.DEC_CROSS_ATTN_Q: (
  794. "decoder.block.{bid}.layer.1.EncDecAttention.q", # t5
  795. ),
  796. MODEL_TENSOR.DEC_CROSS_ATTN_K: (
  797. "decoder.block.{bid}.layer.1.EncDecAttention.k", # t5
  798. ),
  799. MODEL_TENSOR.DEC_CROSS_ATTN_V: (
  800. "decoder.block.{bid}.layer.1.EncDecAttention.v", # t5
  801. ),
  802. MODEL_TENSOR.DEC_CROSS_ATTN_OUT: (
  803. "decoder.block.{bid}.layer.1.EncDecAttention.o", # t5
  804. ),
  805. MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: (
  806. "decoder.block.{bid}.layer.1.EncDecAttention.relative_attention_bias", # t5
  807. ),
  808. MODEL_TENSOR.DEC_FFN_NORM: (
  809. "decoder.block.{bid}.layer.2.layer_norm", # t5
  810. ),
  811. MODEL_TENSOR.DEC_FFN_GATE: (
  812. "decoder.block.{bid}.layer.2.DenseReluDense.wi_0", # flan-t5
  813. ),
  814. MODEL_TENSOR.DEC_FFN_UP: (
  815. "decoder.block.{bid}.layer.2.DenseReluDense.wi", # t5
  816. "decoder.block.{bid}.layer.2.DenseReluDense.wi_1", # flan-t5
  817. ),
  818. MODEL_TENSOR.DEC_FFN_DOWN: (
  819. "decoder.block.{bid}.layer.2.DenseReluDense.wo", # t5
  820. ),
  821. MODEL_TENSOR.DEC_OUTPUT_NORM: (
  822. "decoder.final_layer_norm", # t5
  823. ),
  824. MODEL_TENSOR.ENC_ATTN_NORM: (
  825. "encoder.block.{bid}.layer.0.layer_norm", # t5
  826. ),
  827. MODEL_TENSOR.ENC_ATTN_Q: (
  828. "encoder.block.{bid}.layer.0.SelfAttention.q", # t5
  829. ),
  830. MODEL_TENSOR.ENC_ATTN_K: (
  831. "encoder.block.{bid}.layer.0.SelfAttention.k", # t5
  832. ),
  833. MODEL_TENSOR.ENC_ATTN_V: (
  834. "encoder.block.{bid}.layer.0.SelfAttention.v", # t5
  835. ),
  836. MODEL_TENSOR.ENC_ATTN_OUT: (
  837. "encoder.block.{bid}.layer.0.SelfAttention.o", # t5
  838. ),
  839. MODEL_TENSOR.ENC_ATTN_REL_B: (
  840. "encoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5
  841. ),
  842. MODEL_TENSOR.ENC_FFN_NORM: (
  843. "encoder.block.{bid}.layer.1.layer_norm", # t5
  844. ),
  845. MODEL_TENSOR.ENC_FFN_GATE: (
  846. "encoder.block.{bid}.layer.1.DenseReluDense.wi_0", # flan-t5
  847. ),
  848. MODEL_TENSOR.ENC_FFN_UP: (
  849. "encoder.block.{bid}.layer.1.DenseReluDense.wi", # t5
  850. "encoder.block.{bid}.layer.1.DenseReluDense.wi_1", # flan-t5
  851. ),
  852. MODEL_TENSOR.ENC_FFN_DOWN: (
  853. "encoder.block.{bid}.layer.1.DenseReluDense.wo", # t5
  854. ),
  855. ############################################################################
  856. # TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg
  857. MODEL_TENSOR.ENC_OUTPUT_NORM: (
  858. "encoder.final_layer_norm", # t5
  859. "layer_norm", # neobert
  860. ),
  861. MODEL_TENSOR.CLS: (
  862. "classifier", # jina
  863. "classifier.dense", # roberta
  864. "pre_classifier", # distillbert
  865. "dense", # neobert
  866. ),
  867. MODEL_TENSOR.CLS_OUT: (
  868. "classifier.out_proj", # roberta
  869. ),
  870. #############################################################################
  871. MODEL_TENSOR.CONVNEXT_DW: (
  872. "backbone.convnext.{bid}.dwconv", # wavtokenizer
  873. ),
  874. MODEL_TENSOR.CONVNEXT_NORM: (
  875. "backbone.convnext.{bid}.norm", # wavtokenizer
  876. ),
  877. MODEL_TENSOR.CONVNEXT_PW1: (
  878. "backbone.convnext.{bid}.pwconv1", # wavtokenizer
  879. ),
  880. MODEL_TENSOR.CONVNEXT_PW2: (
  881. "backbone.convnext.{bid}.pwconv2", # wavtokenizer
  882. ),
  883. MODEL_TENSOR.CONVNEXT_GAMMA: (
  884. "backbone.convnext.{bid}.gamma", # wavtokenizer
  885. ),
  886. MODEL_TENSOR.POSNET_CONV1: (
  887. "backbone.posnet.{bid}.conv1", # wavtokenizer
  888. ),
  889. MODEL_TENSOR.POSNET_CONV2: (
  890. "backbone.posnet.{bid}.conv2", # wavtokenizer
  891. ),
  892. MODEL_TENSOR.POSNET_NORM: (
  893. "backbone.posnet.{bid}.norm", # wavtokenizer
  894. ),
  895. MODEL_TENSOR.POSNET_NORM1: (
  896. "backbone.posnet.{bid}.norm1", # wavtokenizer
  897. ),
  898. MODEL_TENSOR.POSNET_NORM2: (
  899. "backbone.posnet.{bid}.norm2", # wavtokenizer
  900. ),
  901. MODEL_TENSOR.POSNET_ATTN_NORM: (
  902. "backbone.posnet.{bid}.norm", # wavtokenizer
  903. ),
  904. MODEL_TENSOR.POSNET_ATTN_Q: (
  905. "backbone.posnet.{bid}.q", # wavtokenizer
  906. ),
  907. MODEL_TENSOR.POSNET_ATTN_K: (
  908. "backbone.posnet.{bid}.k", # wavtokenizer
  909. ),
  910. MODEL_TENSOR.POSNET_ATTN_V: (
  911. "backbone.posnet.{bid}.v", # wavtokenizer
  912. ),
  913. MODEL_TENSOR.POSNET_ATTN_OUT: (
  914. "backbone.posnet.{bid}.proj_out", # wavtokenizer
  915. ),
  916. MODEL_TENSOR.SHORTCONV_CONV: (
  917. "model.layers.{bid}.conv.conv",
  918. ),
  919. MODEL_TENSOR.SHORTCONV_INPROJ: (
  920. "model.layers.{bid}.conv.in_proj",
  921. ),
  922. MODEL_TENSOR.SHORTCONV_OUTPROJ: (
  923. "model.layers.{bid}.conv.out_proj",
  924. ),
  925. #############################################################################
  926. ## Vision encoder
  927. MODEL_TENSOR.V_MMPROJ: (
  928. "multi_modal_projector.linear_{bid}",
  929. "visual.merger.mlp.{bid}", # qwen2vl
  930. ),
  931. MODEL_TENSOR.V_MMPROJ_FC: (
  932. "model.connector.modality_projection.proj", # SmolVLM
  933. ),
  934. MODEL_TENSOR.V_MMPROJ_MLP: (
  935. "model.mm_projector.mlp.mlp.{bid}",
  936. "vision_model.vision_adapter.mlp.fc{bid}", # llama 4
  937. "mlp1.{bid}", # InternVL
  938. ),
  939. MODEL_TENSOR.V_MMPROJ_PEG: (
  940. "model.mm_projector.peg.peg.{bid}",
  941. ),
  942. MODEL_TENSOR.V_ENC_EMBD_CLS: (
  943. "vision_tower.vision_model.embeddings.class_embedding",
  944. "vision_model.class_embedding", # llama 4
  945. ),
  946. MODEL_TENSOR.V_ENC_EMBD_PATCH: (
  947. "vision_tower.vision_model.embeddings.patch_embedding",
  948. "vpm.embeddings.patch_embedding",
  949. "model.vision_model.embeddings.patch_embedding", # SmolVLM
  950. "vision_tower.patch_conv", # pixtral
  951. "vision_model.patch_embedding.linear", # llama 4
  952. "visual.patch_embed.proj", # qwen2vl
  953. ),
  954. MODEL_TENSOR.V_ENC_EMBD_POS: (
  955. "vision_tower.vision_model.embeddings.position_embedding",
  956. "vpm.embeddings.position_embedding",
  957. "model.vision_model.embeddings.position_embedding", # SmolVLM
  958. "vision_model.positional_embedding_vlm", # llama 4
  959. ),
  960. MODEL_TENSOR.V_ENC_ATTN_Q: (
  961. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.q_proj",
  962. "vpm.encoder.layers.{bid}.self_attn.q_proj",
  963. "model.vision_model.encoder.layers.{bid}.self_attn.q_proj", # SmolVLM
  964. "vision_model.model.layers.{bid}.self_attn.q_proj", # llama4
  965. "vision_tower.transformer.layers.{bid}.attention.q_proj", # pixtral
  966. "visual.blocks.{bid}.attn.q", # qwen2vl, generated
  967. ),
  968. MODEL_TENSOR.V_ENC_ATTN_Q_NORM: (
  969. "vision_tower.vision_model.encoder.layers.{bid}.attn.q_norm", # InternVL
  970. ),
  971. MODEL_TENSOR.V_ENC_ATTN_K: (
  972. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.k_proj",
  973. "vpm.encoder.layers.{bid}.self_attn.k_proj",
  974. "model.vision_model.encoder.layers.{bid}.self_attn.k_proj", # SmolVLM
  975. "vision_model.model.layers.{bid}.self_attn.k_proj", # llama4
  976. "vision_tower.transformer.layers.{bid}.attention.k_proj", # pixtral
  977. "visual.blocks.{bid}.attn.k", # qwen2vl, generated
  978. ),
  979. MODEL_TENSOR.V_ENC_ATTN_K_NORM: (
  980. "vision_tower.vision_model.encoder.layers.{bid}.attn.k_norm", # InternVL
  981. ),
  982. MODEL_TENSOR.V_ENC_ATTN_V: (
  983. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.v_proj",
  984. "vpm.encoder.layers.{bid}.self_attn.v_proj",
  985. "model.vision_model.encoder.layers.{bid}.self_attn.v_proj", # SmolVLM
  986. "vision_model.model.layers.{bid}.self_attn.v_proj", # llama4
  987. "vision_tower.transformer.layers.{bid}.attention.v_proj", # pixtral
  988. "visual.blocks.{bid}.attn.v", # qwen2vl, generated
  989. ),
  990. MODEL_TENSOR.V_ENC_INPUT_NORM: (
  991. "vision_tower.vision_model.encoder.layers.{bid}.layer_norm1",
  992. "vision_tower.vision_model.encoder.layers.{bid}.norm1", # InternVL
  993. "vpm.encoder.layers.{bid}.layer_norm1",
  994. "model.vision_model.encoder.layers.{bid}.layer_norm1", # SmolVLM
  995. "vision_tower.transformer.layers.{bid}.attention_norm", # pixtral
  996. "vision_model.model.layers.{bid}.input_layernorm", # llama4
  997. "visual.blocks.{bid}.norm1", # qwen2vl
  998. ),
  999. MODEL_TENSOR.V_ENC_ATTN_O: (
  1000. "vision_tower.vision_model.encoder.layers.{bid}.self_attn.out_proj",
  1001. "vision_tower.vision_model.encoder.layers.{bid}.attn.proj", # InternVL
  1002. "vpm.encoder.layers.{bid}.self_attn.out_proj",
  1003. "model.vision_model.encoder.layers.{bid}.self_attn.out_proj", # SmolVLM
  1004. "vision_model.model.layers.{bid}.self_attn.o_proj", # llama4
  1005. "vision_tower.transformer.layers.{bid}.attention.o_proj", # pixtral
  1006. "visual.blocks.{bid}.attn.proj", # qwen2vl
  1007. ),
  1008. MODEL_TENSOR.V_ENC_POST_ATTN_NORM: (
  1009. "vision_tower.vision_model.encoder.layers.{bid}.layer_norm2",
  1010. "vision_tower.vision_model.encoder.layers.{bid}.norm2", # InternVL
  1011. "vpm.encoder.layers.{bid}.layer_norm2",
  1012. "model.vision_model.encoder.layers.{bid}.layer_norm2", # SmolVLM
  1013. "vision_model.model.layers.{bid}.post_attention_layernorm", # llama4
  1014. "vision_tower.transformer.layers.{bid}.ffn_norm", # pixtral
  1015. "visual.blocks.{bid}.norm2", # qwen2vl
  1016. ),
  1017. MODEL_TENSOR.V_ENC_FFN_UP: (
  1018. "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc1",
  1019. "vpm.encoder.layers.{bid}.mlp.fc1",
  1020. "model.vision_model.encoder.layers.{bid}.mlp.fc1", # SmolVLM, gemma3
  1021. "vision_tower.transformer.layers.{bid}.feed_forward.up_proj", # pixtral
  1022. "vision_model.model.layers.{bid}.mlp.fc1", # llama4
  1023. "visual.blocks.{bid}.mlp.fc1", # qwen2vl
  1024. "visual.blocks.{bid}.mlp.up_proj", # qwen2.5vl
  1025. ),
  1026. MODEL_TENSOR.V_ENC_FFN_GATE: (
  1027. "vision_tower.transformer.layers.{bid}.feed_forward.gate_proj", # pixtral
  1028. "visual.blocks.{bid}.mlp.gate_proj", # qwen2.5vl
  1029. ),
  1030. MODEL_TENSOR.V_ENC_FFN_DOWN: (
  1031. "vision_tower.vision_model.encoder.layers.{bid}.mlp.fc2",
  1032. "vpm.encoder.layers.{bid}.mlp.fc2",
  1033. "model.vision_model.encoder.layers.{bid}.mlp.fc2", # SmolVLM, gemma3
  1034. "vision_tower.transformer.layers.{bid}.feed_forward.down_proj", # pixtral
  1035. "vision_model.model.layers.{bid}.mlp.fc2", # llama4
  1036. "visual.blocks.{bid}.mlp.fc2", # qwen2vl
  1037. "visual.blocks.{bid}.mlp.down_proj", # qwen2.5vl
  1038. ),
  1039. MODEL_TENSOR.V_LAYER_SCALE_1: (
  1040. "vision_tower.vision_model.encoder.layers.{bid}.ls1", # InternVL
  1041. ),
  1042. MODEL_TENSOR.V_LAYER_SCALE_2: (
  1043. "vision_tower.vision_model.encoder.layers.{bid}.ls2", # InternVL
  1044. ),
  1045. MODEL_TENSOR.V_PRE_NORM: (
  1046. "vision_tower.vision_model.pre_layrnorm",
  1047. "vision_tower.ln_pre", # pixtral
  1048. "vision_model.layernorm_pre", # llama4
  1049. ),
  1050. MODEL_TENSOR.V_POST_NORM: (
  1051. "vision_tower.vision_model.post_layernorm",
  1052. "model.vision_model.post_layernorm", # SmolVLM
  1053. "vision_model.layernorm_post", # llama4
  1054. "visual.merger.ln_q", # qwen2vl
  1055. ),
  1056. MODEL_TENSOR.V_MM_INP_PROJ: (
  1057. "multi_modal_projector.mm_input_projection",
  1058. ),
  1059. MODEL_TENSOR.V_MM_INP_NORM: (
  1060. "multi_modal_projector.norm",
  1061. ),
  1062. MODEL_TENSOR.V_MM_SOFT_EMB_NORM: (
  1063. "multi_modal_projector.mm_soft_emb_norm",
  1064. ),
  1065. MODEL_TENSOR.V_RESMPL_POS_EMBD_K: (
  1066. "resampler.pos_embed_k",
  1067. ),
  1068. MODEL_TENSOR.V_RESMPL_ATTN_Q: (
  1069. "resampler.attn.in_proj_q", # tensor generated from resampler.attn.in_proj
  1070. ),
  1071. MODEL_TENSOR.V_RESMPL_ATTN_K: (
  1072. "resampler.attn.in_proj_k", # tensor generated from resampler.attn.in_proj
  1073. ),
  1074. MODEL_TENSOR.V_RESMPL_ATTN_V: (
  1075. "resampler.attn.in_proj_v", # tensor generated from resampler.attn.in_proj
  1076. ),
  1077. MODEL_TENSOR.V_RESMPL_ATTN_OUT: (
  1078. "resampler.attn.out_proj",
  1079. ),
  1080. MODEL_TENSOR.V_RESMPL_KV: (
  1081. "resampler.kv_proj",
  1082. ),
  1083. MODEL_TENSOR.V_RESMPL_POST_NORM: (
  1084. "resampler.ln_post",
  1085. ),
  1086. MODEL_TENSOR.V_RESMPL_KV_NORM: (
  1087. "resampler.ln_kv",
  1088. ),
  1089. MODEL_TENSOR.V_RESMPL_Q_NORM: (
  1090. "resampler.ln_q",
  1091. ),
  1092. MODEL_TENSOR.V_RESMPL_PROJ: (
  1093. "resampler.proj",
  1094. ),
  1095. MODEL_TENSOR.V_RESMPL_QUERY: (
  1096. "resampler.query",
  1097. ),
  1098. MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK: (
  1099. "v.token_embd.img_break", # for pixtral, this is a generated vector
  1100. ),
  1101. MODEL_TENSOR.V_MM_PATCH_MERGER: (
  1102. "multi_modal_projector.patch_merger.merging_layer", # mistral small 3.1
  1103. ),
  1104. # audio (mtmd)
  1105. MODEL_TENSOR.A_ENC_EMBD_POS: (
  1106. "audio_tower.embed_positions", # ultravox
  1107. ),
  1108. MODEL_TENSOR.A_ENC_CONV1D: (
  1109. "audio_tower.conv{bid}", # ultravox
  1110. ),
  1111. MODEL_TENSOR.A_PRE_NORM: (),
  1112. MODEL_TENSOR.A_POST_NORM: (
  1113. "audio_tower.layer_norm", # ultravox
  1114. "audio_tower.ln_post", # qwen2omni
  1115. ),
  1116. MODEL_TENSOR.A_ENC_ATTN_Q: (
  1117. "audio_tower.layers.{bid}.self_attn.q_proj", # ultravox
  1118. ),
  1119. MODEL_TENSOR.A_ENC_ATTN_K: (
  1120. "audio_tower.layers.{bid}.self_attn.k_proj", # ultravox
  1121. ),
  1122. MODEL_TENSOR.A_ENC_ATTN_V: (
  1123. "audio_tower.layers.{bid}.self_attn.v_proj", # ultravox
  1124. ),
  1125. MODEL_TENSOR.A_ENC_INPUT_NORM: (
  1126. "audio_tower.layers.{bid}.self_attn_layer_norm", # ultravox
  1127. ),
  1128. MODEL_TENSOR.A_ENC_OUTPUT: (
  1129. "audio_tower.layers.{bid}.self_attn.out_proj", # ultravox
  1130. ),
  1131. MODEL_TENSOR.A_ENC_OUTPUT_NORM: (
  1132. "audio_tower.layers.{bid}.final_layer_norm", # ultravox
  1133. ),
  1134. MODEL_TENSOR.A_ENC_FFN_UP: (
  1135. "audio_tower.layers.{bid}.fc1", # ultravox
  1136. ),
  1137. MODEL_TENSOR.A_ENC_FFN_GATE: (),
  1138. MODEL_TENSOR.A_ENC_FFN_DOWN: (
  1139. "audio_tower.layers.{bid}.fc2", # ultravox
  1140. ),
  1141. # note: some tensors below has "audio." pseudo-prefix, to prevent conflicts with vision tensors
  1142. # this prefix is added in the conversion code in modify_tensors()
  1143. MODEL_TENSOR.A_MMPROJ: (
  1144. "audio.multi_modal_projector.linear_{bid}", # ultravox
  1145. ),
  1146. MODEL_TENSOR.A_MMPROJ_FC: (
  1147. "audio.multi_modal_projector.linear", # qwen2audio
  1148. "audio_tower.proj", # qwen2omni
  1149. ),
  1150. MODEL_TENSOR.A_MM_NORM_PRE: (
  1151. "audio.multi_modal_projector.ln_pre", # ultravox
  1152. ),
  1153. MODEL_TENSOR.A_MM_NORM_MID: (
  1154. "audio.multi_modal_projector.ln_mid", # ultravox
  1155. ),
  1156. # NextN/MTP tensors for GLM4_MOE
  1157. MODEL_TENSOR.NEXTN_EH_PROJ: (
  1158. "model.layers.{bid}.eh_proj",
  1159. ),
  1160. MODEL_TENSOR.NEXTN_EMBED_TOKENS: (
  1161. "model.layers.{bid}.embed_tokens",
  1162. ),
  1163. MODEL_TENSOR.NEXTN_ENORM: (
  1164. "model.layers.{bid}.enorm",
  1165. ),
  1166. MODEL_TENSOR.NEXTN_HNORM: (
  1167. "model.layers.{bid}.hnorm",
  1168. ),
  1169. MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: (
  1170. "model.layers.{bid}.shared_head.head",
  1171. ),
  1172. MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: (
  1173. "model.layers.{bid}.shared_head.norm",
  1174. ),
  1175. }
  1176. # architecture-specific block mappings
  1177. arch_block_mappings_cfg: dict[MODEL_ARCH, dict[MODEL_TENSOR, tuple[str, ...]]] = {
  1178. MODEL_ARCH.ARCTIC: {
  1179. MODEL_TENSOR.FFN_NORM: (
  1180. "model.layers.{bid}.residual_layernorm",
  1181. ),
  1182. MODEL_TENSOR.FFN_NORM_EXP: (
  1183. "model.layers.{bid}.post_attention_layernorm",
  1184. ),
  1185. },
  1186. }
  1187. mapping: dict[str, tuple[MODEL_TENSOR, str]]
  1188. def __init__(self, arch: MODEL_ARCH, n_blocks: int):
  1189. self.mapping = {}
  1190. for tensor, keys in self.mappings_cfg.items():
  1191. if tensor not in MODEL_TENSORS[arch]:
  1192. continue
  1193. tensor_name = TENSOR_NAMES[tensor]
  1194. self.mapping[tensor_name] = (tensor, tensor_name)
  1195. for key in keys:
  1196. self.mapping[key] = (tensor, tensor_name)
  1197. if arch in self.arch_block_mappings_cfg:
  1198. self.block_mappings_cfg.update(self.arch_block_mappings_cfg[arch])
  1199. for bid in range(n_blocks):
  1200. for tensor, keys in self.block_mappings_cfg.items():
  1201. if tensor not in MODEL_TENSORS[arch]:
  1202. continue
  1203. tensor_name = TENSOR_NAMES[tensor].format(bid = bid)
  1204. self.mapping[tensor_name] = (tensor, tensor_name)
  1205. for key in keys:
  1206. key = key.format(bid = bid)
  1207. self.mapping[key] = (tensor, tensor_name)
  1208. def get_type_and_name(self, key: str, try_suffixes: Sequence[str] = ()) -> tuple[MODEL_TENSOR, str] | None:
  1209. result = self.mapping.get(key)
  1210. if result is not None:
  1211. return result
  1212. for suffix in try_suffixes:
  1213. if key.endswith(suffix):
  1214. result = self.mapping.get(key[:-len(suffix)])
  1215. if result is not None:
  1216. return result[0], result[1] + suffix
  1217. return None
  1218. def get_name(self, key: str, try_suffixes: Sequence[str] = ()) -> str | None:
  1219. result = self.get_type_and_name(key, try_suffixes = try_suffixes)
  1220. if result is None:
  1221. return None
  1222. return result[1]
  1223. def get_type(self, key: str, try_suffixes: Sequence[str] = ()) -> MODEL_TENSOR | None:
  1224. result = self.get_type_and_name(key, try_suffixes = try_suffixes)
  1225. if result is None:
  1226. return None
  1227. return result[0]
  1228. def __getitem__(self, key: str) -> str:
  1229. try:
  1230. return self.mapping[key][1]
  1231. except KeyError:
  1232. raise KeyError(key)
  1233. def __contains__(self, key: str) -> bool:
  1234. return key in self.mapping
  1235. def __repr__(self) -> str:
  1236. return repr(self.mapping)
  1237. def get_tensor_name_map(arch: MODEL_ARCH, n_blocks: int) -> TensorNameMap:
  1238. return TensorNameMap(arch, n_blocks)