tensor_mapping.py 61 KB

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