constants.py 81 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252
  1. from __future__ import annotations
  2. from enum import Enum, IntEnum, auto
  3. from typing import Any
  4. #
  5. # constants
  6. #
  7. GGUF_MAGIC = 0x46554747 # "GGUF"
  8. GGUF_VERSION = 3
  9. GGUF_DEFAULT_ALIGNMENT = 32
  10. GGML_QUANT_VERSION = 2 # GGML_QNT_VERSION from ggml.h
  11. #
  12. # metadata keys
  13. #
  14. class Keys:
  15. class General:
  16. TYPE = "general.type"
  17. ARCHITECTURE = "general.architecture"
  18. QUANTIZATION_VERSION = "general.quantization_version"
  19. ALIGNMENT = "general.alignment"
  20. FILE_TYPE = "general.file_type"
  21. # Authorship Metadata
  22. NAME = "general.name"
  23. AUTHOR = "general.author"
  24. VERSION = "general.version"
  25. ORGANIZATION = "general.organization"
  26. FINETUNE = "general.finetune"
  27. BASENAME = "general.basename"
  28. DESCRIPTION = "general.description"
  29. QUANTIZED_BY = "general.quantized_by"
  30. SIZE_LABEL = "general.size_label"
  31. # Licensing details
  32. LICENSE = "general.license"
  33. LICENSE_NAME = "general.license.name"
  34. LICENSE_LINK = "general.license.link"
  35. # Typically represents the converted GGUF repo (Unless native)
  36. URL = "general.url" # Model Website/Paper
  37. DOI = "general.doi"
  38. UUID = "general.uuid"
  39. REPO_URL = "general.repo_url" # Model Source Repository (git/svn/etc...)
  40. # Model Source during conversion
  41. SOURCE_URL = "general.source.url" # Model Website/Paper
  42. SOURCE_DOI = "general.source.doi"
  43. SOURCE_UUID = "general.source.uuid"
  44. SOURCE_REPO_URL = "general.source.repo_url" # Model Source Repository (git/svn/etc...)
  45. # Base Model Source. There can be more than one source if it's a merged
  46. # model like with 'Mistral-7B-Merge-14-v0.1'. This will assist in
  47. # tracing linage of models as it is finetuned or merged over time.
  48. BASE_MODEL_COUNT = "general.base_model.count"
  49. BASE_MODEL_NAME = "general.base_model.{id}.name"
  50. BASE_MODEL_AUTHOR = "general.base_model.{id}.author"
  51. BASE_MODEL_VERSION = "general.base_model.{id}.version"
  52. BASE_MODEL_ORGANIZATION = "general.base_model.{id}.organization"
  53. BASE_MODEL_DESCRIPTION = "general.base_model.{id}.description"
  54. BASE_MODEL_URL = "general.base_model.{id}.url" # Model Website/Paper
  55. BASE_MODEL_DOI = "general.base_model.{id}.doi"
  56. BASE_MODEL_UUID = "general.base_model.{id}.uuid"
  57. BASE_MODEL_REPO_URL = "general.base_model.{id}.repo_url" # Model Source Repository (git/svn/etc...)
  58. # Dataset Source
  59. DATASET_COUNT = "general.dataset.count"
  60. DATASET_NAME = "general.dataset.{id}.name"
  61. DATASET_AUTHOR = "general.dataset.{id}.author"
  62. DATASET_VERSION = "general.dataset.{id}.version"
  63. DATASET_ORGANIZATION = "general.dataset.{id}.organization"
  64. DATASET_DESCRIPTION = "general.dataset.{id}.description"
  65. DATASET_URL = "general.dataset.{id}.url" # Model Website/Paper
  66. DATASET_DOI = "general.dataset.{id}.doi"
  67. DATASET_UUID = "general.dataset.{id}.uuid"
  68. DATASET_REPO_URL = "general.dataset.{id}.repo_url" # Model Source Repository (git/svn/etc...)
  69. # Array based KV stores
  70. TAGS = "general.tags"
  71. LANGUAGES = "general.languages"
  72. class LLM:
  73. VOCAB_SIZE = "{arch}.vocab_size"
  74. CONTEXT_LENGTH = "{arch}.context_length"
  75. EMBEDDING_LENGTH = "{arch}.embedding_length"
  76. FEATURES_LENGTH = "{arch}.features_length"
  77. BLOCK_COUNT = "{arch}.block_count"
  78. LEADING_DENSE_BLOCK_COUNT = "{arch}.leading_dense_block_count"
  79. FEED_FORWARD_LENGTH = "{arch}.feed_forward_length"
  80. EXPERT_FEED_FORWARD_LENGTH = "{arch}.expert_feed_forward_length"
  81. EXPERT_SHARED_FEED_FORWARD_LENGTH = "{arch}.expert_shared_feed_forward_length"
  82. USE_PARALLEL_RESIDUAL = "{arch}.use_parallel_residual"
  83. TENSOR_DATA_LAYOUT = "{arch}.tensor_data_layout"
  84. EXPERT_COUNT = "{arch}.expert_count"
  85. EXPERT_USED_COUNT = "{arch}.expert_used_count"
  86. EXPERT_SHARED_COUNT = "{arch}.expert_shared_count"
  87. EXPERT_WEIGHTS_SCALE = "{arch}.expert_weights_scale"
  88. EXPERT_WEIGHTS_NORM = "{arch}.expert_weights_norm"
  89. EXPERT_GATING_FUNC = "{arch}.expert_gating_func"
  90. POOLING_TYPE = "{arch}.pooling_type"
  91. LOGIT_SCALE = "{arch}.logit_scale"
  92. DECODER_START_TOKEN_ID = "{arch}.decoder_start_token_id"
  93. ATTN_LOGIT_SOFTCAPPING = "{arch}.attn_logit_softcapping"
  94. FINAL_LOGIT_SOFTCAPPING = "{arch}.final_logit_softcapping"
  95. SWIN_NORM = "{arch}.swin_norm"
  96. RESCALE_EVERY_N_LAYERS = "{arch}.rescale_every_n_layers"
  97. TIME_MIX_EXTRA_DIM = "{arch}.time_mix_extra_dim"
  98. TIME_DECAY_EXTRA_DIM = "{arch}.time_decay_extra_dim"
  99. RESIDUAL_SCALE = "{arch}.residual_scale"
  100. EMBEDDING_SCALE = "{arch}.embedding_scale"
  101. TOKEN_SHIFT_COUNT = "{arch}.token_shift_count"
  102. INTERLEAVE_MOE_LAYER_STEP = "{arch}.interleave_moe_layer_step"
  103. class Attention:
  104. HEAD_COUNT = "{arch}.attention.head_count"
  105. HEAD_COUNT_KV = "{arch}.attention.head_count_kv"
  106. MAX_ALIBI_BIAS = "{arch}.attention.max_alibi_bias"
  107. CLAMP_KQV = "{arch}.attention.clamp_kqv"
  108. KEY_LENGTH = "{arch}.attention.key_length"
  109. VALUE_LENGTH = "{arch}.attention.value_length"
  110. LAYERNORM_EPS = "{arch}.attention.layer_norm_epsilon"
  111. LAYERNORM_RMS_EPS = "{arch}.attention.layer_norm_rms_epsilon"
  112. GROUPNORM_EPS = "{arch}.attention.group_norm_epsilon"
  113. GROUPNORM_GROUPS = "{arch}.attention.group_norm_groups"
  114. CAUSAL = "{arch}.attention.causal"
  115. Q_LORA_RANK = "{arch}.attention.q_lora_rank"
  116. KV_LORA_RANK = "{arch}.attention.kv_lora_rank"
  117. DECAY_LORA_RANK = "{arch}.attention.decay_lora_rank"
  118. ICLR_LORA_RANK = "{arch}.attention.iclr_lora_rank"
  119. VALUE_RESIDUAL_MIX_LORA_RANK = "{arch}.attention.value_residual_mix_lora_rank"
  120. GATE_LORA_RANK = "{arch}.attention.gate_lora_rank"
  121. REL_BUCKETS_COUNT = "{arch}.attention.relative_buckets_count"
  122. SLIDING_WINDOW = "{arch}.attention.sliding_window"
  123. SCALE = "{arch}.attention.scale"
  124. KEY_LENGTH_MLA = "{arch}.attention.key_length_mla"
  125. VALUE_LENGTH_MLA = "{arch}.attention.value_length_mla"
  126. class Rope:
  127. DIMENSION_COUNT = "{arch}.rope.dimension_count"
  128. DIMENSION_SECTIONS = "{arch}.rope.dimension_sections"
  129. FREQ_BASE = "{arch}.rope.freq_base"
  130. SCALING_TYPE = "{arch}.rope.scaling.type"
  131. SCALING_FACTOR = "{arch}.rope.scaling.factor"
  132. SCALING_ATTN_FACTOR = "{arch}.rope.scaling.attn_factor"
  133. SCALING_ORIG_CTX_LEN = "{arch}.rope.scaling.original_context_length"
  134. SCALING_FINETUNED = "{arch}.rope.scaling.finetuned"
  135. SCALING_YARN_LOG_MUL = "{arch}.rope.scaling.yarn_log_multiplier"
  136. class Split:
  137. LLM_KV_SPLIT_NO = "split.no"
  138. LLM_KV_SPLIT_COUNT = "split.count"
  139. LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count"
  140. class SSM:
  141. CONV_KERNEL = "{arch}.ssm.conv_kernel"
  142. INNER_SIZE = "{arch}.ssm.inner_size"
  143. STATE_SIZE = "{arch}.ssm.state_size"
  144. TIME_STEP_RANK = "{arch}.ssm.time_step_rank"
  145. DT_B_C_RMS = "{arch}.ssm.dt_b_c_rms"
  146. class WKV:
  147. HEAD_SIZE = "{arch}.wkv.head_size"
  148. class PosNet:
  149. EMBEDDING_LENGTH = "{arch}.posnet.embedding_length"
  150. BLOCK_COUNT = "{arch}.posnet.block_count"
  151. class ConvNext:
  152. EMBEDDING_LENGTH = "{arch}.convnext.embedding_length"
  153. BLOCK_COUNT = "{arch}.convnext.block_count"
  154. class Tokenizer:
  155. MODEL = "tokenizer.ggml.model"
  156. PRE = "tokenizer.ggml.pre"
  157. LIST = "tokenizer.ggml.tokens"
  158. TOKEN_TYPE = "tokenizer.ggml.token_type"
  159. TOKEN_TYPE_COUNT = "tokenizer.ggml.token_type_count" # for BERT-style token types
  160. SCORES = "tokenizer.ggml.scores"
  161. MERGES = "tokenizer.ggml.merges"
  162. BOS_ID = "tokenizer.ggml.bos_token_id"
  163. EOS_ID = "tokenizer.ggml.eos_token_id"
  164. EOT_ID = "tokenizer.ggml.eot_token_id"
  165. EOM_ID = "tokenizer.ggml.eom_token_id"
  166. UNK_ID = "tokenizer.ggml.unknown_token_id"
  167. SEP_ID = "tokenizer.ggml.seperator_token_id"
  168. PAD_ID = "tokenizer.ggml.padding_token_id"
  169. MASK_ID = "tokenizer.ggml.mask_token_id"
  170. ADD_BOS = "tokenizer.ggml.add_bos_token"
  171. ADD_EOS = "tokenizer.ggml.add_eos_token"
  172. ADD_PREFIX = "tokenizer.ggml.add_space_prefix"
  173. REMOVE_EXTRA_WS = "tokenizer.ggml.remove_extra_whitespaces"
  174. PRECOMPILED_CHARSMAP = "tokenizer.ggml.precompiled_charsmap"
  175. HF_JSON = "tokenizer.huggingface.json"
  176. RWKV = "tokenizer.rwkv.world"
  177. CHAT_TEMPLATE = "tokenizer.chat_template"
  178. CHAT_TEMPLATE_N = "tokenizer.chat_template.{name}"
  179. CHAT_TEMPLATES = "tokenizer.chat_templates"
  180. # FIM/Infill special tokens constants
  181. FIM_PRE_ID = "tokenizer.ggml.fim_pre_token_id"
  182. FIM_SUF_ID = "tokenizer.ggml.fim_suf_token_id"
  183. FIM_MID_ID = "tokenizer.ggml.fim_mid_token_id"
  184. FIM_PAD_ID = "tokenizer.ggml.fim_pad_token_id"
  185. FIM_REP_ID = "tokenizer.ggml.fim_rep_token_id"
  186. FIM_SEP_ID = "tokenizer.ggml.fim_sep_token_id"
  187. # deprecated:
  188. PREFIX_ID = "tokenizer.ggml.prefix_token_id"
  189. SUFFIX_ID = "tokenizer.ggml.suffix_token_id"
  190. MIDDLE_ID = "tokenizer.ggml.middle_token_id"
  191. class Adapter:
  192. TYPE = "adapter.type"
  193. LORA_ALPHA = "adapter.lora.alpha"
  194. class ClipVision:
  195. PROJECTOR_TYPE = "clip.projector_type"
  196. HAS_VISION_ENCODER = "clip.has_vision_encoder"
  197. HAS_LLAVA_PROJECTOR = "clip.has_llava_projector"
  198. IMAGE_SIZE = "clip.vision.image_size"
  199. PATCH_SIZE = "clip.vision.patch_size"
  200. EMBEDDING_LENGTH = "clip.vision.embedding_length"
  201. FEED_FORWARD_LENGTH = "clip.vision.feed_forward_length"
  202. PROJECTION_DIM = "clip.vision.projection_dim"
  203. BLOCK_COUNT = "clip.vision.block_count"
  204. IMAGE_MEAN = "clip.vision.image_mean"
  205. IMAGE_STD = "clip.vision.image_std"
  206. USE_GELU = "clip.use_gelu"
  207. USE_SILU = "clip.use_silu"
  208. class Attention:
  209. HEAD_COUNT = "clip.vision.attention.head_count"
  210. LAYERNORM_EPS = "clip.vision.attention.layer_norm_epsilon"
  211. class Projector:
  212. SCALE_FACTOR = "clip.vision.projector.scale_factor"
  213. #
  214. # recommended mapping of model tensor names for storage in gguf
  215. #
  216. class GGUFType:
  217. MODEL = "model"
  218. ADAPTER = "adapter"
  219. CLIP_VISION = "clip-vision"
  220. class MODEL_ARCH(IntEnum):
  221. CLIP_VISION = auto() # dummy arch for clip.cpp
  222. LLAMA = auto()
  223. LLAMA4 = auto()
  224. DECI = auto()
  225. FALCON = auto()
  226. BAICHUAN = auto()
  227. GROK = auto()
  228. GPT2 = auto()
  229. GPTJ = auto()
  230. GPTNEOX = auto()
  231. MPT = auto()
  232. STARCODER = auto()
  233. REFACT = auto()
  234. BERT = auto()
  235. NOMIC_BERT = auto()
  236. JINA_BERT_V2 = auto()
  237. BLOOM = auto()
  238. STABLELM = auto()
  239. QWEN = auto()
  240. QWEN2 = auto()
  241. QWEN2MOE = auto()
  242. QWEN2VL = auto()
  243. QWEN3 = auto()
  244. QWEN3MOE = auto()
  245. PHI2 = auto()
  246. PHI3 = auto()
  247. PHIMOE = auto()
  248. PLAMO = auto()
  249. CODESHELL = auto()
  250. ORION = auto()
  251. INTERNLM2 = auto()
  252. MINICPM = auto()
  253. MINICPM3 = auto()
  254. GEMMA = auto()
  255. GEMMA2 = auto()
  256. GEMMA3 = auto()
  257. STARCODER2 = auto()
  258. RWKV6 = auto()
  259. RWKV6QWEN2 = auto()
  260. RWKV7 = auto()
  261. ARWKV7 = auto()
  262. MAMBA = auto()
  263. XVERSE = auto()
  264. COMMAND_R = auto()
  265. COHERE2 = auto()
  266. DBRX = auto()
  267. OLMO = auto()
  268. OLMO2 = auto()
  269. OLMOE = auto()
  270. OPENELM = auto()
  271. ARCTIC = auto()
  272. DEEPSEEK = auto()
  273. DEEPSEEK2 = auto()
  274. CHATGLM = auto()
  275. GLM4 = auto()
  276. BITNET = auto()
  277. T5 = auto()
  278. T5ENCODER = auto()
  279. JAIS = auto()
  280. NEMOTRON = auto()
  281. EXAONE = auto()
  282. GRANITE = auto()
  283. GRANITE_MOE = auto()
  284. CHAMELEON = auto()
  285. WAVTOKENIZER_DEC = auto()
  286. PLM = auto()
  287. BAILINGMOE = auto()
  288. class VISION_PROJECTOR_TYPE(IntEnum):
  289. MLP = auto()
  290. LDP = auto()
  291. LDPV2 = auto()
  292. RESAMPLER = auto()
  293. GLM_EDGE = auto()
  294. MERGER = auto()
  295. GEMMA3 = auto()
  296. class MODEL_TENSOR(IntEnum):
  297. TOKEN_EMBD = auto()
  298. TOKEN_EMBD_NORM = auto()
  299. TOKEN_TYPES = auto()
  300. POS_EMBD = auto()
  301. OUTPUT = auto()
  302. OUTPUT_NORM = auto()
  303. ROPE_FREQS = auto()
  304. ROPE_FACTORS_LONG = auto()
  305. ROPE_FACTORS_SHORT = auto()
  306. ATTN_Q = auto()
  307. ATTN_K = auto()
  308. ATTN_V = auto()
  309. ATTN_QKV = auto()
  310. ATTN_OUT = auto()
  311. ATTN_NORM = auto()
  312. ATTN_NORM_2 = auto()
  313. ATTN_OUT_NORM = auto()
  314. ATTN_POST_NORM = auto()
  315. ATTN_ROT_EMBD = auto()
  316. FFN_GATE_INP = auto()
  317. FFN_GATE_INP_SHEXP = auto()
  318. FFN_NORM = auto()
  319. FFN_PRE_NORM = auto()
  320. FFN_POST_NORM = auto()
  321. FFN_GATE = auto()
  322. FFN_DOWN = auto()
  323. FFN_UP = auto()
  324. FFN_ACT = auto()
  325. FFN_NORM_EXP = auto()
  326. FFN_GATE_EXP = auto()
  327. FFN_DOWN_EXP = auto()
  328. FFN_UP_EXP = auto()
  329. FFN_GATE_SHEXP = auto()
  330. FFN_DOWN_SHEXP = auto()
  331. FFN_UP_SHEXP = auto()
  332. FFN_EXP_PROBS_B = auto()
  333. ATTN_Q_NORM = auto()
  334. ATTN_K_NORM = auto()
  335. LAYER_OUT_NORM = auto()
  336. SSM_IN = auto()
  337. SSM_CONV1D = auto()
  338. SSM_X = auto()
  339. SSM_DT = auto()
  340. SSM_A = auto()
  341. SSM_D = auto()
  342. SSM_OUT = auto()
  343. TIME_MIX_W0 = auto()
  344. TIME_MIX_W1 = auto()
  345. TIME_MIX_W2 = auto()
  346. TIME_MIX_A0 = auto()
  347. TIME_MIX_A1 = auto()
  348. TIME_MIX_A2 = auto()
  349. TIME_MIX_V0 = auto()
  350. TIME_MIX_V1 = auto()
  351. TIME_MIX_V2 = auto()
  352. TIME_MIX_G1 = auto()
  353. TIME_MIX_G2 = auto()
  354. TIME_MIX_K_K = auto()
  355. TIME_MIX_K_A = auto()
  356. TIME_MIX_R_K = auto()
  357. TIME_MIX_LERP_X = auto()
  358. TIME_MIX_LERP_K = auto()
  359. TIME_MIX_LERP_V = auto()
  360. TIME_MIX_LERP_R = auto()
  361. TIME_MIX_LERP_G = auto()
  362. TIME_MIX_LERP_FUSED = auto()
  363. TIME_MIX_LERP_W = auto()
  364. TIME_MIX_FIRST = auto()
  365. TIME_MIX_DECAY = auto()
  366. TIME_MIX_DECAY_W1 = auto()
  367. TIME_MIX_DECAY_W2 = auto()
  368. TIME_MIX_KEY = auto()
  369. TIME_MIX_VALUE = auto()
  370. TIME_MIX_RECEPTANCE = auto()
  371. TIME_MIX_GATE = auto()
  372. TIME_MIX_LN = auto()
  373. TIME_MIX_OUTPUT = auto()
  374. CHANNEL_MIX_LERP_K = auto()
  375. CHANNEL_MIX_LERP_R = auto()
  376. CHANNEL_MIX_KEY = auto()
  377. CHANNEL_MIX_RECEPTANCE = auto()
  378. CHANNEL_MIX_VALUE = auto()
  379. ATTN_Q_A = auto()
  380. ATTN_Q_B = auto()
  381. ATTN_KV_A_MQA = auto()
  382. ATTN_KV_B = auto()
  383. ATTN_K_B = auto()
  384. ATTN_V_B = auto()
  385. ATTN_Q_A_NORM = auto()
  386. ATTN_KV_A_NORM = auto()
  387. FFN_SUB_NORM = auto()
  388. ATTN_SUB_NORM = auto()
  389. DEC_ATTN_NORM = auto()
  390. DEC_ATTN_Q = auto()
  391. DEC_ATTN_K = auto()
  392. DEC_ATTN_V = auto()
  393. DEC_ATTN_OUT = auto()
  394. DEC_ATTN_REL_B = auto()
  395. DEC_CROSS_ATTN_NORM = auto()
  396. DEC_CROSS_ATTN_Q = auto()
  397. DEC_CROSS_ATTN_K = auto()
  398. DEC_CROSS_ATTN_V = auto()
  399. DEC_CROSS_ATTN_OUT = auto()
  400. DEC_CROSS_ATTN_REL_B = auto()
  401. DEC_FFN_NORM = auto()
  402. DEC_FFN_GATE = auto()
  403. DEC_FFN_DOWN = auto()
  404. DEC_FFN_UP = auto()
  405. DEC_OUTPUT_NORM = auto()
  406. ENC_ATTN_NORM = auto()
  407. ENC_ATTN_Q = auto()
  408. ENC_ATTN_K = auto()
  409. ENC_ATTN_V = auto()
  410. ENC_ATTN_OUT = auto()
  411. ENC_ATTN_REL_B = auto()
  412. ENC_FFN_NORM = auto()
  413. ENC_FFN_GATE = auto()
  414. ENC_FFN_DOWN = auto()
  415. ENC_FFN_UP = auto()
  416. ENC_OUTPUT_NORM = auto()
  417. CLS = auto() # classifier
  418. CLS_OUT = auto() # classifier output projection
  419. CONV1D = auto()
  420. CONVNEXT_DW = auto()
  421. CONVNEXT_NORM = auto()
  422. CONVNEXT_PW1 = auto()
  423. CONVNEXT_PW2 = auto()
  424. CONVNEXT_GAMMA = auto()
  425. POSNET_CONV1 = auto()
  426. POSNET_CONV2 = auto()
  427. POSNET_NORM = auto()
  428. POSNET_NORM1 = auto()
  429. POSNET_NORM2 = auto()
  430. POSNET_ATTN_NORM = auto()
  431. POSNET_ATTN_Q = auto()
  432. POSNET_ATTN_K = auto()
  433. POSNET_ATTN_V = auto()
  434. POSNET_ATTN_OUT = auto()
  435. # vision
  436. V_MMPROJ = auto()
  437. V_MMPROJ_FC = auto()
  438. V_MMPROJ_MLP = auto()
  439. V_MMPROJ_PEG = auto()
  440. V_ENC_EMBD_CLS = auto()
  441. V_ENC_EMBD_PATCH = auto()
  442. V_ENC_EMBD_POS = auto()
  443. V_ENC_ATTN_Q = auto()
  444. V_ENC_ATTN_K = auto()
  445. V_ENC_ATTN_V = auto()
  446. V_ENC_INPUT_NORM = auto()
  447. V_ENC_OUTPUT = auto()
  448. V_ENC_OUTPUT_NORM = auto()
  449. V_ENC_FFN_UP = auto()
  450. V_ENC_FFN_GATE = auto()
  451. V_ENC_FFN_DOWN = auto()
  452. V_PRE_NORM = auto()
  453. V_POST_NORM = auto()
  454. V_MM_INP_PROJ = auto() # gemma3
  455. V_MM_SOFT_EMB_NORM = auto() # gemma3
  456. V_RESMPL_POS_EMBD_K = auto() # minicpmv
  457. V_RESMPL_ATTN_Q = auto() # minicpmv
  458. V_RESMPL_ATTN_K = auto() # minicpmv
  459. V_RESMPL_ATTN_V = auto() # minicpmv
  460. V_RESMPL_ATTN_OUT = auto() # minicpmv
  461. V_RESMPL_KV = auto() # minicpmv
  462. V_RESMPL_KV_NORM = auto() # minicpmv
  463. V_RESMPL_POST_NORM = auto() # minicpmv
  464. V_RESMPL_Q_NORM = auto() # minicpmv
  465. V_RESMPL_PROJ = auto() # minicpmv
  466. V_RESMPL_QUERY = auto() # minicpmv
  467. V_TOK_EMBD_IMG_BREAK = auto() # pixtral
  468. MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
  469. MODEL_ARCH.CLIP_VISION: "clip", # dummy arch for clip.cpp
  470. MODEL_ARCH.LLAMA: "llama",
  471. MODEL_ARCH.LLAMA4: "llama4",
  472. MODEL_ARCH.DECI: "deci",
  473. MODEL_ARCH.FALCON: "falcon",
  474. MODEL_ARCH.BAICHUAN: "baichuan",
  475. MODEL_ARCH.GROK: "grok",
  476. MODEL_ARCH.GPT2: "gpt2",
  477. MODEL_ARCH.GPTJ: "gptj",
  478. MODEL_ARCH.GPTNEOX: "gptneox",
  479. MODEL_ARCH.MPT: "mpt",
  480. MODEL_ARCH.STARCODER: "starcoder",
  481. MODEL_ARCH.REFACT: "refact",
  482. MODEL_ARCH.BERT: "bert",
  483. MODEL_ARCH.NOMIC_BERT: "nomic-bert",
  484. MODEL_ARCH.JINA_BERT_V2: "jina-bert-v2",
  485. MODEL_ARCH.BLOOM: "bloom",
  486. MODEL_ARCH.STABLELM: "stablelm",
  487. MODEL_ARCH.QWEN: "qwen",
  488. MODEL_ARCH.QWEN2: "qwen2",
  489. MODEL_ARCH.QWEN2MOE: "qwen2moe",
  490. MODEL_ARCH.QWEN2VL: "qwen2vl",
  491. MODEL_ARCH.QWEN3: "qwen3",
  492. MODEL_ARCH.QWEN3MOE: "qwen3moe",
  493. MODEL_ARCH.PHI2: "phi2",
  494. MODEL_ARCH.PHI3: "phi3",
  495. MODEL_ARCH.PHIMOE: "phimoe",
  496. MODEL_ARCH.PLAMO: "plamo",
  497. MODEL_ARCH.CODESHELL: "codeshell",
  498. MODEL_ARCH.ORION: "orion",
  499. MODEL_ARCH.INTERNLM2: "internlm2",
  500. MODEL_ARCH.MINICPM: "minicpm",
  501. MODEL_ARCH.MINICPM3: "minicpm3",
  502. MODEL_ARCH.GEMMA: "gemma",
  503. MODEL_ARCH.GEMMA2: "gemma2",
  504. MODEL_ARCH.GEMMA3: "gemma3",
  505. MODEL_ARCH.STARCODER2: "starcoder2",
  506. MODEL_ARCH.RWKV6: "rwkv6",
  507. MODEL_ARCH.RWKV6QWEN2: "rwkv6qwen2",
  508. MODEL_ARCH.RWKV7: "rwkv7",
  509. MODEL_ARCH.ARWKV7: "arwkv7",
  510. MODEL_ARCH.MAMBA: "mamba",
  511. MODEL_ARCH.XVERSE: "xverse",
  512. MODEL_ARCH.COMMAND_R: "command-r",
  513. MODEL_ARCH.COHERE2: "cohere2",
  514. MODEL_ARCH.DBRX: "dbrx",
  515. MODEL_ARCH.OLMO: "olmo",
  516. MODEL_ARCH.OLMO2: "olmo2",
  517. MODEL_ARCH.OLMOE: "olmoe",
  518. MODEL_ARCH.OPENELM: "openelm",
  519. MODEL_ARCH.ARCTIC: "arctic",
  520. MODEL_ARCH.DEEPSEEK: "deepseek",
  521. MODEL_ARCH.DEEPSEEK2: "deepseek2",
  522. MODEL_ARCH.CHATGLM: "chatglm",
  523. MODEL_ARCH.GLM4: "glm4",
  524. MODEL_ARCH.BITNET: "bitnet",
  525. MODEL_ARCH.T5: "t5",
  526. MODEL_ARCH.T5ENCODER: "t5encoder",
  527. MODEL_ARCH.JAIS: "jais",
  528. MODEL_ARCH.NEMOTRON: "nemotron",
  529. MODEL_ARCH.EXAONE: "exaone",
  530. MODEL_ARCH.GRANITE: "granite",
  531. MODEL_ARCH.GRANITE_MOE: "granitemoe",
  532. MODEL_ARCH.CHAMELEON: "chameleon",
  533. MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec",
  534. MODEL_ARCH.PLM: "plm",
  535. MODEL_ARCH.BAILINGMOE: "bailingmoe",
  536. }
  537. VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = {
  538. VISION_PROJECTOR_TYPE.MLP: "mlp",
  539. VISION_PROJECTOR_TYPE.LDP: "ldp",
  540. VISION_PROJECTOR_TYPE.LDPV2: "ldpv2",
  541. VISION_PROJECTOR_TYPE.RESAMPLER: "resampler",
  542. VISION_PROJECTOR_TYPE.GLM_EDGE: "adapter",
  543. VISION_PROJECTOR_TYPE.MERGER: "qwen2vl_merger",
  544. VISION_PROJECTOR_TYPE.GEMMA3: "gemma3",
  545. }
  546. TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
  547. MODEL_TENSOR.TOKEN_EMBD: "token_embd",
  548. MODEL_TENSOR.TOKEN_EMBD_NORM: "token_embd_norm",
  549. MODEL_TENSOR.TOKEN_TYPES: "token_types",
  550. MODEL_TENSOR.POS_EMBD: "position_embd",
  551. MODEL_TENSOR.OUTPUT_NORM: "output_norm",
  552. MODEL_TENSOR.OUTPUT: "output",
  553. MODEL_TENSOR.ROPE_FREQS: "rope_freqs",
  554. MODEL_TENSOR.ROPE_FACTORS_LONG: "rope_factors_long",
  555. MODEL_TENSOR.ROPE_FACTORS_SHORT: "rope_factors_short",
  556. MODEL_TENSOR.ATTN_NORM: "blk.{bid}.attn_norm",
  557. MODEL_TENSOR.ATTN_NORM_2: "blk.{bid}.attn_norm_2",
  558. MODEL_TENSOR.ATTN_QKV: "blk.{bid}.attn_qkv",
  559. MODEL_TENSOR.ATTN_Q: "blk.{bid}.attn_q",
  560. MODEL_TENSOR.ATTN_K: "blk.{bid}.attn_k",
  561. MODEL_TENSOR.ATTN_V: "blk.{bid}.attn_v",
  562. MODEL_TENSOR.ATTN_OUT: "blk.{bid}.attn_output",
  563. MODEL_TENSOR.ATTN_ROT_EMBD: "blk.{bid}.attn_rot_embd",
  564. MODEL_TENSOR.ATTN_Q_NORM: "blk.{bid}.attn_q_norm",
  565. MODEL_TENSOR.ATTN_K_NORM: "blk.{bid}.attn_k_norm",
  566. MODEL_TENSOR.ATTN_OUT_NORM: "blk.{bid}.attn_output_norm",
  567. MODEL_TENSOR.ATTN_POST_NORM: "blk.{bid}.post_attention_norm",
  568. MODEL_TENSOR.FFN_GATE_INP: "blk.{bid}.ffn_gate_inp",
  569. MODEL_TENSOR.FFN_GATE_INP_SHEXP: "blk.{bid}.ffn_gate_inp_shexp",
  570. MODEL_TENSOR.FFN_NORM: "blk.{bid}.ffn_norm",
  571. MODEL_TENSOR.FFN_PRE_NORM: "blk.{bid}.ffn_norm",
  572. MODEL_TENSOR.FFN_POST_NORM: "blk.{bid}.post_ffw_norm",
  573. MODEL_TENSOR.FFN_GATE: "blk.{bid}.ffn_gate",
  574. MODEL_TENSOR.FFN_DOWN: "blk.{bid}.ffn_down",
  575. MODEL_TENSOR.FFN_UP: "blk.{bid}.ffn_up",
  576. MODEL_TENSOR.FFN_GATE_SHEXP: "blk.{bid}.ffn_gate_shexp",
  577. MODEL_TENSOR.FFN_DOWN_SHEXP: "blk.{bid}.ffn_down_shexp",
  578. MODEL_TENSOR.FFN_UP_SHEXP: "blk.{bid}.ffn_up_shexp",
  579. MODEL_TENSOR.FFN_ACT: "blk.{bid}.ffn",
  580. MODEL_TENSOR.FFN_NORM_EXP: "blk.{bid}.ffn_norm_exps",
  581. MODEL_TENSOR.FFN_GATE_EXP: "blk.{bid}.ffn_gate_exps",
  582. MODEL_TENSOR.FFN_DOWN_EXP: "blk.{bid}.ffn_down_exps",
  583. MODEL_TENSOR.FFN_UP_EXP: "blk.{bid}.ffn_up_exps",
  584. MODEL_TENSOR.FFN_EXP_PROBS_B: "blk.{bid}.exp_probs_b",
  585. MODEL_TENSOR.LAYER_OUT_NORM: "blk.{bid}.layer_output_norm",
  586. MODEL_TENSOR.SSM_IN: "blk.{bid}.ssm_in",
  587. MODEL_TENSOR.SSM_CONV1D: "blk.{bid}.ssm_conv1d",
  588. MODEL_TENSOR.SSM_X: "blk.{bid}.ssm_x",
  589. MODEL_TENSOR.SSM_DT: "blk.{bid}.ssm_dt",
  590. MODEL_TENSOR.SSM_A: "blk.{bid}.ssm_a",
  591. MODEL_TENSOR.SSM_D: "blk.{bid}.ssm_d",
  592. MODEL_TENSOR.SSM_OUT: "blk.{bid}.ssm_out",
  593. MODEL_TENSOR.TIME_MIX_W0: "blk.{bid}.time_mix_w0",
  594. MODEL_TENSOR.TIME_MIX_W1: "blk.{bid}.time_mix_w1",
  595. MODEL_TENSOR.TIME_MIX_W2: "blk.{bid}.time_mix_w2",
  596. MODEL_TENSOR.TIME_MIX_A0: "blk.{bid}.time_mix_a0",
  597. MODEL_TENSOR.TIME_MIX_A1: "blk.{bid}.time_mix_a1",
  598. MODEL_TENSOR.TIME_MIX_A2: "blk.{bid}.time_mix_a2",
  599. MODEL_TENSOR.TIME_MIX_V0: "blk.{bid}.time_mix_v0",
  600. MODEL_TENSOR.TIME_MIX_V1: "blk.{bid}.time_mix_v1",
  601. MODEL_TENSOR.TIME_MIX_V2: "blk.{bid}.time_mix_v2",
  602. MODEL_TENSOR.TIME_MIX_G1: "blk.{bid}.time_mix_g1",
  603. MODEL_TENSOR.TIME_MIX_G2: "blk.{bid}.time_mix_g2",
  604. MODEL_TENSOR.TIME_MIX_K_K: "blk.{bid}.time_mix_k_k",
  605. MODEL_TENSOR.TIME_MIX_K_A: "blk.{bid}.time_mix_k_a",
  606. MODEL_TENSOR.TIME_MIX_R_K: "blk.{bid}.time_mix_r_k",
  607. MODEL_TENSOR.TIME_MIX_LERP_X: "blk.{bid}.time_mix_lerp_x",
  608. MODEL_TENSOR.TIME_MIX_LERP_K: "blk.{bid}.time_mix_lerp_k",
  609. MODEL_TENSOR.TIME_MIX_LERP_V: "blk.{bid}.time_mix_lerp_v",
  610. MODEL_TENSOR.TIME_MIX_LERP_R: "blk.{bid}.time_mix_lerp_r",
  611. MODEL_TENSOR.TIME_MIX_LERP_G: "blk.{bid}.time_mix_lerp_g",
  612. MODEL_TENSOR.TIME_MIX_LERP_FUSED: "blk.{bid}.time_mix_lerp_fused",
  613. MODEL_TENSOR.TIME_MIX_LERP_W: "blk.{bid}.time_mix_lerp_w",
  614. MODEL_TENSOR.TIME_MIX_FIRST: "blk.{bid}.time_mix_first",
  615. MODEL_TENSOR.TIME_MIX_DECAY: "blk.{bid}.time_mix_decay",
  616. MODEL_TENSOR.TIME_MIX_DECAY_W1: "blk.{bid}.time_mix_decay_w1",
  617. MODEL_TENSOR.TIME_MIX_DECAY_W2: "blk.{bid}.time_mix_decay_w2",
  618. MODEL_TENSOR.TIME_MIX_KEY: "blk.{bid}.time_mix_key",
  619. MODEL_TENSOR.TIME_MIX_VALUE: "blk.{bid}.time_mix_value",
  620. MODEL_TENSOR.TIME_MIX_RECEPTANCE: "blk.{bid}.time_mix_receptance",
  621. MODEL_TENSOR.TIME_MIX_GATE: "blk.{bid}.time_mix_gate",
  622. MODEL_TENSOR.TIME_MIX_LN: "blk.{bid}.time_mix_ln",
  623. MODEL_TENSOR.TIME_MIX_OUTPUT: "blk.{bid}.time_mix_output",
  624. MODEL_TENSOR.CHANNEL_MIX_LERP_K: "blk.{bid}.channel_mix_lerp_k",
  625. MODEL_TENSOR.CHANNEL_MIX_LERP_R: "blk.{bid}.channel_mix_lerp_r",
  626. MODEL_TENSOR.CHANNEL_MIX_KEY: "blk.{bid}.channel_mix_key",
  627. MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE: "blk.{bid}.channel_mix_receptance",
  628. MODEL_TENSOR.CHANNEL_MIX_VALUE: "blk.{bid}.channel_mix_value",
  629. MODEL_TENSOR.ATTN_Q_A: "blk.{bid}.attn_q_a",
  630. MODEL_TENSOR.ATTN_Q_B: "blk.{bid}.attn_q_b",
  631. MODEL_TENSOR.ATTN_KV_A_MQA: "blk.{bid}.attn_kv_a_mqa",
  632. MODEL_TENSOR.ATTN_KV_B: "blk.{bid}.attn_kv_b",
  633. MODEL_TENSOR.ATTN_K_B: "blk.{bid}.attn_k_b",
  634. MODEL_TENSOR.ATTN_V_B: "blk.{bid}.attn_v_b",
  635. MODEL_TENSOR.ATTN_Q_A_NORM: "blk.{bid}.attn_q_a_norm",
  636. MODEL_TENSOR.ATTN_KV_A_NORM: "blk.{bid}.attn_kv_a_norm",
  637. MODEL_TENSOR.ATTN_SUB_NORM: "blk.{bid}.attn_sub_norm",
  638. MODEL_TENSOR.FFN_SUB_NORM: "blk.{bid}.ffn_sub_norm",
  639. MODEL_TENSOR.DEC_ATTN_NORM: "dec.blk.{bid}.attn_norm",
  640. MODEL_TENSOR.DEC_ATTN_Q: "dec.blk.{bid}.attn_q",
  641. MODEL_TENSOR.DEC_ATTN_K: "dec.blk.{bid}.attn_k",
  642. MODEL_TENSOR.DEC_ATTN_V: "dec.blk.{bid}.attn_v",
  643. MODEL_TENSOR.DEC_ATTN_OUT: "dec.blk.{bid}.attn_o",
  644. MODEL_TENSOR.DEC_ATTN_REL_B: "dec.blk.{bid}.attn_rel_b",
  645. MODEL_TENSOR.DEC_CROSS_ATTN_NORM: "dec.blk.{bid}.cross_attn_norm",
  646. MODEL_TENSOR.DEC_CROSS_ATTN_Q: "dec.blk.{bid}.cross_attn_q",
  647. MODEL_TENSOR.DEC_CROSS_ATTN_K: "dec.blk.{bid}.cross_attn_k",
  648. MODEL_TENSOR.DEC_CROSS_ATTN_V: "dec.blk.{bid}.cross_attn_v",
  649. MODEL_TENSOR.DEC_CROSS_ATTN_OUT: "dec.blk.{bid}.cross_attn_o",
  650. MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: "dec.blk.{bid}.cross_attn_rel_b",
  651. MODEL_TENSOR.DEC_FFN_NORM: "dec.blk.{bid}.ffn_norm",
  652. MODEL_TENSOR.DEC_FFN_GATE: "dec.blk.{bid}.ffn_gate",
  653. MODEL_TENSOR.DEC_FFN_DOWN: "dec.blk.{bid}.ffn_down",
  654. MODEL_TENSOR.DEC_FFN_UP: "dec.blk.{bid}.ffn_up",
  655. MODEL_TENSOR.DEC_OUTPUT_NORM: "dec.output_norm",
  656. MODEL_TENSOR.ENC_ATTN_NORM: "enc.blk.{bid}.attn_norm",
  657. MODEL_TENSOR.ENC_ATTN_Q: "enc.blk.{bid}.attn_q",
  658. MODEL_TENSOR.ENC_ATTN_K: "enc.blk.{bid}.attn_k",
  659. MODEL_TENSOR.ENC_ATTN_V: "enc.blk.{bid}.attn_v",
  660. MODEL_TENSOR.ENC_ATTN_OUT: "enc.blk.{bid}.attn_o",
  661. MODEL_TENSOR.ENC_ATTN_REL_B: "enc.blk.{bid}.attn_rel_b",
  662. MODEL_TENSOR.ENC_FFN_NORM: "enc.blk.{bid}.ffn_norm",
  663. MODEL_TENSOR.ENC_FFN_GATE: "enc.blk.{bid}.ffn_gate",
  664. MODEL_TENSOR.ENC_FFN_DOWN: "enc.blk.{bid}.ffn_down",
  665. MODEL_TENSOR.ENC_FFN_UP: "enc.blk.{bid}.ffn_up",
  666. MODEL_TENSOR.ENC_OUTPUT_NORM: "enc.output_norm",
  667. MODEL_TENSOR.CLS: "cls",
  668. MODEL_TENSOR.CLS_OUT: "cls.output",
  669. MODEL_TENSOR.CONV1D: "conv1d",
  670. MODEL_TENSOR.CONVNEXT_DW: "convnext.{bid}.dw",
  671. MODEL_TENSOR.CONVNEXT_NORM: "convnext.{bid}.norm",
  672. MODEL_TENSOR.CONVNEXT_PW1: "convnext.{bid}.pw1",
  673. MODEL_TENSOR.CONVNEXT_PW2: "convnext.{bid}.pw2",
  674. MODEL_TENSOR.CONVNEXT_GAMMA: "convnext.{bid}.gamma",
  675. MODEL_TENSOR.POSNET_CONV1: "posnet.{bid}.conv1",
  676. MODEL_TENSOR.POSNET_CONV2: "posnet.{bid}.conv2",
  677. MODEL_TENSOR.POSNET_NORM: "posnet.{bid}.norm",
  678. MODEL_TENSOR.POSNET_NORM1: "posnet.{bid}.norm1",
  679. MODEL_TENSOR.POSNET_NORM2: "posnet.{bid}.norm2",
  680. MODEL_TENSOR.POSNET_ATTN_NORM: "posnet.{bid}.attn_norm",
  681. MODEL_TENSOR.POSNET_ATTN_Q: "posnet.{bid}.attn_q",
  682. MODEL_TENSOR.POSNET_ATTN_K: "posnet.{bid}.attn_k",
  683. MODEL_TENSOR.POSNET_ATTN_V: "posnet.{bid}.attn_v",
  684. MODEL_TENSOR.POSNET_ATTN_OUT: "posnet.{bid}.attn_output",
  685. # vision
  686. MODEL_TENSOR.V_MMPROJ: "mm.{bid}",
  687. MODEL_TENSOR.V_MMPROJ_FC: "mm.model.fc",
  688. MODEL_TENSOR.V_MMPROJ_MLP: "mm.model.mlp.{bid}",
  689. MODEL_TENSOR.V_MMPROJ_PEG: "mm.model.peg.{bid}",
  690. MODEL_TENSOR.V_ENC_EMBD_CLS: "v.class_embd",
  691. MODEL_TENSOR.V_ENC_EMBD_PATCH: "v.patch_embd",
  692. MODEL_TENSOR.V_ENC_EMBD_POS: "v.position_embd",
  693. MODEL_TENSOR.V_ENC_ATTN_Q: "v.blk.{bid}.attn_q",
  694. MODEL_TENSOR.V_ENC_ATTN_K: "v.blk.{bid}.attn_k",
  695. MODEL_TENSOR.V_ENC_ATTN_V: "v.blk.{bid}.attn_v",
  696. MODEL_TENSOR.V_ENC_INPUT_NORM: "v.blk.{bid}.ln1",
  697. MODEL_TENSOR.V_ENC_OUTPUT: "v.blk.{bid}.attn_out",
  698. MODEL_TENSOR.V_ENC_OUTPUT_NORM: "v.blk.{bid}.ln2",
  699. MODEL_TENSOR.V_ENC_FFN_UP: "v.blk.{bid}.ffn_up",
  700. MODEL_TENSOR.V_ENC_FFN_GATE: "v.blk.{bid}.ffn_gate",
  701. MODEL_TENSOR.V_ENC_FFN_DOWN: "v.blk.{bid}.ffn_down",
  702. MODEL_TENSOR.V_PRE_NORM: "v.pre_ln",
  703. MODEL_TENSOR.V_POST_NORM: "v.post_ln",
  704. MODEL_TENSOR.V_MM_INP_PROJ: "mm.input_projection",
  705. MODEL_TENSOR.V_MM_SOFT_EMB_NORM: "mm.soft_emb_norm",
  706. MODEL_TENSOR.V_RESMPL_POS_EMBD_K: "resampler.pos_embd_k",
  707. MODEL_TENSOR.V_RESMPL_ATTN_Q: "resampler.attn.q",
  708. MODEL_TENSOR.V_RESMPL_ATTN_K: "resampler.attn.k",
  709. MODEL_TENSOR.V_RESMPL_ATTN_V: "resampler.attn.v",
  710. MODEL_TENSOR.V_RESMPL_ATTN_OUT: "resampler.attn.out",
  711. MODEL_TENSOR.V_RESMPL_KV: "resampler.kv",
  712. MODEL_TENSOR.V_RESMPL_KV_NORM: "resampler.ln_kv",
  713. MODEL_TENSOR.V_RESMPL_POST_NORM: "resampler.ln_post",
  714. MODEL_TENSOR.V_RESMPL_Q_NORM: "resampler.ln_q",
  715. MODEL_TENSOR.V_RESMPL_PROJ: "resampler.proj",
  716. MODEL_TENSOR.V_RESMPL_QUERY: "resampler.query",
  717. MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK: "v.token_embd.img_break", # pixtral
  718. }
  719. MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
  720. MODEL_ARCH.CLIP_VISION: [
  721. MODEL_TENSOR.V_MMPROJ,
  722. MODEL_TENSOR.V_MMPROJ_FC,
  723. MODEL_TENSOR.V_MMPROJ_MLP,
  724. MODEL_TENSOR.V_MMPROJ_PEG,
  725. MODEL_TENSOR.V_ENC_EMBD_CLS,
  726. MODEL_TENSOR.V_ENC_EMBD_PATCH,
  727. MODEL_TENSOR.V_ENC_EMBD_POS,
  728. MODEL_TENSOR.V_ENC_ATTN_Q,
  729. MODEL_TENSOR.V_ENC_ATTN_K,
  730. MODEL_TENSOR.V_ENC_ATTN_V,
  731. MODEL_TENSOR.V_ENC_INPUT_NORM,
  732. MODEL_TENSOR.V_ENC_OUTPUT,
  733. MODEL_TENSOR.V_ENC_OUTPUT_NORM,
  734. MODEL_TENSOR.V_ENC_FFN_UP,
  735. MODEL_TENSOR.V_ENC_FFN_GATE,
  736. MODEL_TENSOR.V_ENC_FFN_DOWN,
  737. MODEL_TENSOR.V_PRE_NORM,
  738. MODEL_TENSOR.V_POST_NORM,
  739. MODEL_TENSOR.V_MM_INP_PROJ,
  740. MODEL_TENSOR.V_MM_SOFT_EMB_NORM,
  741. MODEL_TENSOR.V_RESMPL_POS_EMBD_K,
  742. MODEL_TENSOR.V_RESMPL_ATTN_Q,
  743. MODEL_TENSOR.V_RESMPL_ATTN_K,
  744. MODEL_TENSOR.V_RESMPL_ATTN_V,
  745. MODEL_TENSOR.V_RESMPL_ATTN_OUT,
  746. MODEL_TENSOR.V_RESMPL_KV,
  747. MODEL_TENSOR.V_RESMPL_KV_NORM,
  748. MODEL_TENSOR.V_RESMPL_POST_NORM,
  749. MODEL_TENSOR.V_RESMPL_Q_NORM,
  750. MODEL_TENSOR.V_RESMPL_PROJ,
  751. MODEL_TENSOR.V_RESMPL_QUERY,
  752. MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK,
  753. ],
  754. MODEL_ARCH.LLAMA: [
  755. MODEL_TENSOR.TOKEN_EMBD,
  756. MODEL_TENSOR.OUTPUT_NORM,
  757. MODEL_TENSOR.OUTPUT,
  758. MODEL_TENSOR.ROPE_FREQS,
  759. MODEL_TENSOR.ATTN_NORM,
  760. MODEL_TENSOR.ATTN_Q,
  761. MODEL_TENSOR.ATTN_K,
  762. MODEL_TENSOR.ATTN_V,
  763. MODEL_TENSOR.ATTN_OUT,
  764. MODEL_TENSOR.ATTN_ROT_EMBD,
  765. MODEL_TENSOR.FFN_GATE_INP,
  766. MODEL_TENSOR.FFN_NORM,
  767. MODEL_TENSOR.FFN_GATE,
  768. MODEL_TENSOR.FFN_DOWN,
  769. MODEL_TENSOR.FFN_UP,
  770. MODEL_TENSOR.FFN_GATE_EXP,
  771. MODEL_TENSOR.FFN_DOWN_EXP,
  772. MODEL_TENSOR.FFN_UP_EXP,
  773. ],
  774. MODEL_ARCH.LLAMA4: [
  775. MODEL_TENSOR.TOKEN_EMBD,
  776. MODEL_TENSOR.OUTPUT_NORM,
  777. MODEL_TENSOR.OUTPUT,
  778. MODEL_TENSOR.ROPE_FREQS,
  779. MODEL_TENSOR.ATTN_NORM,
  780. MODEL_TENSOR.ATTN_Q,
  781. MODEL_TENSOR.ATTN_K,
  782. MODEL_TENSOR.ATTN_V,
  783. MODEL_TENSOR.ATTN_OUT,
  784. MODEL_TENSOR.ATTN_ROT_EMBD,
  785. MODEL_TENSOR.FFN_GATE_INP,
  786. MODEL_TENSOR.FFN_NORM,
  787. MODEL_TENSOR.FFN_GATE,
  788. MODEL_TENSOR.FFN_DOWN,
  789. MODEL_TENSOR.FFN_UP,
  790. MODEL_TENSOR.FFN_GATE_EXP,
  791. MODEL_TENSOR.FFN_DOWN_EXP,
  792. MODEL_TENSOR.FFN_UP_EXP,
  793. MODEL_TENSOR.FFN_GATE_SHEXP,
  794. MODEL_TENSOR.FFN_DOWN_SHEXP,
  795. MODEL_TENSOR.FFN_UP_SHEXP,
  796. ],
  797. MODEL_ARCH.DECI: [
  798. MODEL_TENSOR.TOKEN_EMBD,
  799. MODEL_TENSOR.OUTPUT_NORM,
  800. MODEL_TENSOR.OUTPUT,
  801. MODEL_TENSOR.ROPE_FREQS,
  802. MODEL_TENSOR.ATTN_NORM,
  803. MODEL_TENSOR.ATTN_Q,
  804. MODEL_TENSOR.ATTN_K,
  805. MODEL_TENSOR.ATTN_V,
  806. MODEL_TENSOR.ATTN_OUT,
  807. MODEL_TENSOR.ATTN_ROT_EMBD,
  808. MODEL_TENSOR.FFN_GATE_INP,
  809. MODEL_TENSOR.FFN_NORM,
  810. MODEL_TENSOR.FFN_GATE,
  811. MODEL_TENSOR.FFN_DOWN,
  812. MODEL_TENSOR.FFN_UP,
  813. MODEL_TENSOR.FFN_GATE_EXP,
  814. MODEL_TENSOR.FFN_DOWN_EXP,
  815. MODEL_TENSOR.FFN_UP_EXP,
  816. ],
  817. MODEL_ARCH.GROK: [
  818. MODEL_TENSOR.TOKEN_EMBD,
  819. MODEL_TENSOR.OUTPUT_NORM,
  820. MODEL_TENSOR.OUTPUT,
  821. MODEL_TENSOR.ROPE_FREQS,
  822. MODEL_TENSOR.ATTN_NORM,
  823. MODEL_TENSOR.ATTN_Q,
  824. MODEL_TENSOR.ATTN_K,
  825. MODEL_TENSOR.ATTN_V,
  826. MODEL_TENSOR.ATTN_OUT,
  827. MODEL_TENSOR.ATTN_ROT_EMBD,
  828. MODEL_TENSOR.ATTN_OUT_NORM,
  829. MODEL_TENSOR.FFN_GATE_INP,
  830. MODEL_TENSOR.FFN_NORM,
  831. MODEL_TENSOR.FFN_GATE,
  832. MODEL_TENSOR.FFN_DOWN,
  833. MODEL_TENSOR.FFN_UP,
  834. MODEL_TENSOR.FFN_GATE_EXP,
  835. MODEL_TENSOR.FFN_DOWN_EXP,
  836. MODEL_TENSOR.FFN_UP_EXP,
  837. MODEL_TENSOR.LAYER_OUT_NORM,
  838. ],
  839. MODEL_ARCH.GPTNEOX: [
  840. MODEL_TENSOR.TOKEN_EMBD,
  841. MODEL_TENSOR.OUTPUT_NORM,
  842. MODEL_TENSOR.OUTPUT,
  843. MODEL_TENSOR.ATTN_NORM,
  844. MODEL_TENSOR.ATTN_QKV,
  845. MODEL_TENSOR.ATTN_OUT,
  846. MODEL_TENSOR.FFN_NORM,
  847. MODEL_TENSOR.FFN_DOWN,
  848. MODEL_TENSOR.FFN_UP,
  849. ],
  850. MODEL_ARCH.FALCON: [
  851. MODEL_TENSOR.TOKEN_EMBD,
  852. MODEL_TENSOR.OUTPUT_NORM,
  853. MODEL_TENSOR.OUTPUT,
  854. MODEL_TENSOR.ATTN_NORM,
  855. MODEL_TENSOR.ATTN_NORM_2,
  856. MODEL_TENSOR.ATTN_QKV,
  857. MODEL_TENSOR.ATTN_OUT,
  858. MODEL_TENSOR.FFN_DOWN,
  859. MODEL_TENSOR.FFN_UP,
  860. ],
  861. MODEL_ARCH.BAICHUAN: [
  862. MODEL_TENSOR.TOKEN_EMBD,
  863. MODEL_TENSOR.OUTPUT_NORM,
  864. MODEL_TENSOR.OUTPUT,
  865. MODEL_TENSOR.ROPE_FREQS,
  866. MODEL_TENSOR.ATTN_NORM,
  867. MODEL_TENSOR.ATTN_Q,
  868. MODEL_TENSOR.ATTN_K,
  869. MODEL_TENSOR.ATTN_V,
  870. MODEL_TENSOR.ATTN_OUT,
  871. MODEL_TENSOR.ATTN_ROT_EMBD,
  872. MODEL_TENSOR.FFN_NORM,
  873. MODEL_TENSOR.FFN_GATE,
  874. MODEL_TENSOR.FFN_DOWN,
  875. MODEL_TENSOR.FFN_UP,
  876. ],
  877. MODEL_ARCH.STARCODER: [
  878. MODEL_TENSOR.TOKEN_EMBD,
  879. MODEL_TENSOR.POS_EMBD,
  880. MODEL_TENSOR.OUTPUT_NORM,
  881. MODEL_TENSOR.OUTPUT,
  882. MODEL_TENSOR.ATTN_NORM,
  883. MODEL_TENSOR.ATTN_QKV,
  884. MODEL_TENSOR.ATTN_OUT,
  885. MODEL_TENSOR.FFN_NORM,
  886. MODEL_TENSOR.FFN_DOWN,
  887. MODEL_TENSOR.FFN_UP,
  888. ],
  889. MODEL_ARCH.BERT: [
  890. MODEL_TENSOR.TOKEN_EMBD,
  891. MODEL_TENSOR.TOKEN_EMBD_NORM,
  892. MODEL_TENSOR.TOKEN_TYPES,
  893. MODEL_TENSOR.POS_EMBD,
  894. MODEL_TENSOR.OUTPUT_NORM,
  895. MODEL_TENSOR.ATTN_OUT_NORM,
  896. MODEL_TENSOR.ATTN_Q,
  897. MODEL_TENSOR.ATTN_K,
  898. MODEL_TENSOR.ATTN_V,
  899. MODEL_TENSOR.ATTN_OUT,
  900. MODEL_TENSOR.FFN_DOWN,
  901. MODEL_TENSOR.FFN_UP,
  902. MODEL_TENSOR.LAYER_OUT_NORM,
  903. MODEL_TENSOR.CLS,
  904. MODEL_TENSOR.CLS_OUT,
  905. ],
  906. MODEL_ARCH.NOMIC_BERT: [
  907. MODEL_TENSOR.TOKEN_EMBD,
  908. MODEL_TENSOR.TOKEN_EMBD_NORM,
  909. MODEL_TENSOR.TOKEN_TYPES,
  910. MODEL_TENSOR.POS_EMBD,
  911. MODEL_TENSOR.OUTPUT_NORM,
  912. MODEL_TENSOR.ATTN_OUT_NORM,
  913. MODEL_TENSOR.ATTN_QKV,
  914. MODEL_TENSOR.ATTN_OUT,
  915. MODEL_TENSOR.FFN_GATE,
  916. MODEL_TENSOR.FFN_DOWN,
  917. MODEL_TENSOR.FFN_UP,
  918. MODEL_TENSOR.LAYER_OUT_NORM,
  919. ],
  920. MODEL_ARCH.JINA_BERT_V2: [
  921. MODEL_TENSOR.TOKEN_EMBD,
  922. MODEL_TENSOR.TOKEN_EMBD_NORM,
  923. MODEL_TENSOR.TOKEN_TYPES,
  924. MODEL_TENSOR.ATTN_NORM_2,
  925. MODEL_TENSOR.ATTN_OUT_NORM,
  926. MODEL_TENSOR.ATTN_Q,
  927. MODEL_TENSOR.ATTN_Q_NORM,
  928. MODEL_TENSOR.ATTN_K,
  929. MODEL_TENSOR.ATTN_K_NORM,
  930. MODEL_TENSOR.ATTN_V,
  931. MODEL_TENSOR.ATTN_OUT,
  932. MODEL_TENSOR.FFN_UP,
  933. MODEL_TENSOR.FFN_GATE,
  934. MODEL_TENSOR.FFN_DOWN,
  935. MODEL_TENSOR.LAYER_OUT_NORM,
  936. MODEL_TENSOR.CLS,
  937. ],
  938. MODEL_ARCH.MPT: [
  939. MODEL_TENSOR.TOKEN_EMBD,
  940. MODEL_TENSOR.OUTPUT_NORM,
  941. MODEL_TENSOR.OUTPUT,
  942. MODEL_TENSOR.ATTN_NORM,
  943. MODEL_TENSOR.ATTN_QKV,
  944. MODEL_TENSOR.ATTN_OUT,
  945. MODEL_TENSOR.FFN_NORM,
  946. MODEL_TENSOR.FFN_DOWN,
  947. MODEL_TENSOR.FFN_UP,
  948. MODEL_TENSOR.FFN_ACT,
  949. MODEL_TENSOR.ATTN_Q_NORM,
  950. MODEL_TENSOR.ATTN_K_NORM,
  951. MODEL_TENSOR.POS_EMBD,
  952. ],
  953. MODEL_ARCH.GPTJ: [
  954. MODEL_TENSOR.TOKEN_EMBD,
  955. MODEL_TENSOR.OUTPUT_NORM,
  956. MODEL_TENSOR.OUTPUT,
  957. MODEL_TENSOR.ATTN_NORM,
  958. MODEL_TENSOR.ATTN_Q,
  959. MODEL_TENSOR.ATTN_K,
  960. MODEL_TENSOR.ATTN_V,
  961. MODEL_TENSOR.ATTN_OUT,
  962. MODEL_TENSOR.FFN_DOWN,
  963. MODEL_TENSOR.FFN_UP,
  964. ],
  965. MODEL_ARCH.REFACT: [
  966. MODEL_TENSOR.TOKEN_EMBD,
  967. MODEL_TENSOR.OUTPUT_NORM,
  968. MODEL_TENSOR.OUTPUT,
  969. MODEL_TENSOR.ATTN_NORM,
  970. MODEL_TENSOR.ATTN_Q,
  971. MODEL_TENSOR.ATTN_K,
  972. MODEL_TENSOR.ATTN_V,
  973. MODEL_TENSOR.ATTN_OUT,
  974. MODEL_TENSOR.FFN_NORM,
  975. MODEL_TENSOR.FFN_GATE,
  976. MODEL_TENSOR.FFN_DOWN,
  977. MODEL_TENSOR.FFN_UP,
  978. ],
  979. MODEL_ARCH.BLOOM: [
  980. MODEL_TENSOR.TOKEN_EMBD,
  981. MODEL_TENSOR.TOKEN_EMBD_NORM,
  982. MODEL_TENSOR.OUTPUT_NORM,
  983. MODEL_TENSOR.OUTPUT,
  984. MODEL_TENSOR.ATTN_NORM,
  985. MODEL_TENSOR.ATTN_QKV,
  986. MODEL_TENSOR.ATTN_OUT,
  987. MODEL_TENSOR.FFN_NORM,
  988. MODEL_TENSOR.FFN_DOWN,
  989. MODEL_TENSOR.FFN_UP,
  990. ],
  991. MODEL_ARCH.STABLELM: [
  992. MODEL_TENSOR.TOKEN_EMBD,
  993. MODEL_TENSOR.OUTPUT_NORM,
  994. MODEL_TENSOR.OUTPUT,
  995. MODEL_TENSOR.ROPE_FREQS,
  996. MODEL_TENSOR.ATTN_NORM,
  997. MODEL_TENSOR.ATTN_Q,
  998. MODEL_TENSOR.ATTN_K,
  999. MODEL_TENSOR.ATTN_V,
  1000. MODEL_TENSOR.ATTN_OUT,
  1001. MODEL_TENSOR.FFN_NORM,
  1002. MODEL_TENSOR.FFN_GATE,
  1003. MODEL_TENSOR.FFN_DOWN,
  1004. MODEL_TENSOR.FFN_UP,
  1005. MODEL_TENSOR.ATTN_Q_NORM,
  1006. MODEL_TENSOR.ATTN_K_NORM,
  1007. ],
  1008. MODEL_ARCH.QWEN: [
  1009. MODEL_TENSOR.TOKEN_EMBD,
  1010. MODEL_TENSOR.OUTPUT_NORM,
  1011. MODEL_TENSOR.OUTPUT,
  1012. MODEL_TENSOR.ROPE_FREQS,
  1013. MODEL_TENSOR.ATTN_NORM,
  1014. MODEL_TENSOR.ATTN_QKV,
  1015. MODEL_TENSOR.ATTN_OUT,
  1016. MODEL_TENSOR.ATTN_ROT_EMBD,
  1017. MODEL_TENSOR.FFN_NORM,
  1018. MODEL_TENSOR.FFN_GATE,
  1019. MODEL_TENSOR.FFN_DOWN,
  1020. MODEL_TENSOR.FFN_UP,
  1021. ],
  1022. MODEL_ARCH.QWEN2: [
  1023. MODEL_TENSOR.TOKEN_EMBD,
  1024. MODEL_TENSOR.OUTPUT_NORM,
  1025. MODEL_TENSOR.OUTPUT,
  1026. MODEL_TENSOR.ROPE_FREQS,
  1027. MODEL_TENSOR.ATTN_NORM,
  1028. MODEL_TENSOR.ATTN_Q,
  1029. MODEL_TENSOR.ATTN_K,
  1030. MODEL_TENSOR.ATTN_V,
  1031. MODEL_TENSOR.ATTN_OUT,
  1032. MODEL_TENSOR.FFN_NORM,
  1033. MODEL_TENSOR.FFN_GATE,
  1034. MODEL_TENSOR.FFN_DOWN,
  1035. MODEL_TENSOR.FFN_UP,
  1036. ],
  1037. MODEL_ARCH.QWEN2VL: [
  1038. MODEL_TENSOR.TOKEN_EMBD,
  1039. MODEL_TENSOR.OUTPUT_NORM,
  1040. MODEL_TENSOR.OUTPUT,
  1041. MODEL_TENSOR.ATTN_NORM,
  1042. MODEL_TENSOR.ATTN_Q,
  1043. MODEL_TENSOR.ATTN_K,
  1044. MODEL_TENSOR.ATTN_V,
  1045. MODEL_TENSOR.ATTN_OUT,
  1046. MODEL_TENSOR.FFN_NORM,
  1047. MODEL_TENSOR.FFN_GATE,
  1048. MODEL_TENSOR.FFN_DOWN,
  1049. MODEL_TENSOR.FFN_UP,
  1050. ],
  1051. MODEL_ARCH.QWEN2MOE: [
  1052. MODEL_TENSOR.TOKEN_EMBD,
  1053. MODEL_TENSOR.OUTPUT_NORM,
  1054. MODEL_TENSOR.OUTPUT,
  1055. MODEL_TENSOR.ATTN_NORM,
  1056. MODEL_TENSOR.ATTN_Q,
  1057. MODEL_TENSOR.ATTN_K,
  1058. MODEL_TENSOR.ATTN_V,
  1059. MODEL_TENSOR.ATTN_OUT,
  1060. MODEL_TENSOR.FFN_NORM,
  1061. MODEL_TENSOR.FFN_GATE_INP,
  1062. MODEL_TENSOR.FFN_GATE_EXP,
  1063. MODEL_TENSOR.FFN_DOWN_EXP,
  1064. MODEL_TENSOR.FFN_UP_EXP,
  1065. MODEL_TENSOR.FFN_GATE_INP_SHEXP,
  1066. MODEL_TENSOR.FFN_GATE_SHEXP,
  1067. MODEL_TENSOR.FFN_DOWN_SHEXP,
  1068. MODEL_TENSOR.FFN_UP_SHEXP,
  1069. ],
  1070. MODEL_ARCH.QWEN3: [
  1071. MODEL_TENSOR.TOKEN_EMBD,
  1072. MODEL_TENSOR.OUTPUT_NORM,
  1073. MODEL_TENSOR.OUTPUT,
  1074. MODEL_TENSOR.ROPE_FREQS,
  1075. MODEL_TENSOR.ATTN_NORM,
  1076. MODEL_TENSOR.ATTN_Q,
  1077. MODEL_TENSOR.ATTN_Q_NORM,
  1078. MODEL_TENSOR.ATTN_K,
  1079. MODEL_TENSOR.ATTN_K_NORM,
  1080. MODEL_TENSOR.ATTN_V,
  1081. MODEL_TENSOR.ATTN_OUT,
  1082. MODEL_TENSOR.FFN_NORM,
  1083. MODEL_TENSOR.FFN_GATE,
  1084. MODEL_TENSOR.FFN_DOWN,
  1085. MODEL_TENSOR.FFN_UP,
  1086. ],
  1087. MODEL_ARCH.QWEN3MOE: [
  1088. MODEL_TENSOR.TOKEN_EMBD,
  1089. MODEL_TENSOR.OUTPUT_NORM,
  1090. MODEL_TENSOR.OUTPUT,
  1091. MODEL_TENSOR.ATTN_NORM,
  1092. MODEL_TENSOR.ATTN_Q,
  1093. MODEL_TENSOR.ATTN_Q_NORM,
  1094. MODEL_TENSOR.ATTN_K,
  1095. MODEL_TENSOR.ATTN_K_NORM,
  1096. MODEL_TENSOR.ATTN_V,
  1097. MODEL_TENSOR.ATTN_OUT,
  1098. MODEL_TENSOR.FFN_NORM,
  1099. MODEL_TENSOR.FFN_GATE_INP,
  1100. MODEL_TENSOR.FFN_GATE_EXP,
  1101. MODEL_TENSOR.FFN_DOWN_EXP,
  1102. MODEL_TENSOR.FFN_UP_EXP,
  1103. ],
  1104. MODEL_ARCH.PLAMO: [
  1105. MODEL_TENSOR.TOKEN_EMBD,
  1106. MODEL_TENSOR.OUTPUT_NORM,
  1107. MODEL_TENSOR.OUTPUT,
  1108. MODEL_TENSOR.ROPE_FREQS,
  1109. MODEL_TENSOR.ATTN_NORM,
  1110. MODEL_TENSOR.ATTN_Q,
  1111. MODEL_TENSOR.ATTN_K,
  1112. MODEL_TENSOR.ATTN_V,
  1113. MODEL_TENSOR.ATTN_OUT,
  1114. MODEL_TENSOR.ATTN_ROT_EMBD,
  1115. MODEL_TENSOR.FFN_GATE,
  1116. MODEL_TENSOR.FFN_DOWN,
  1117. MODEL_TENSOR.FFN_UP,
  1118. ],
  1119. MODEL_ARCH.GPT2: [
  1120. MODEL_TENSOR.TOKEN_EMBD,
  1121. MODEL_TENSOR.POS_EMBD,
  1122. MODEL_TENSOR.OUTPUT_NORM,
  1123. MODEL_TENSOR.OUTPUT,
  1124. MODEL_TENSOR.ATTN_NORM,
  1125. MODEL_TENSOR.ATTN_QKV,
  1126. MODEL_TENSOR.ATTN_OUT,
  1127. MODEL_TENSOR.FFN_NORM,
  1128. MODEL_TENSOR.FFN_DOWN,
  1129. MODEL_TENSOR.FFN_UP,
  1130. ],
  1131. MODEL_ARCH.PHI2: [
  1132. MODEL_TENSOR.TOKEN_EMBD,
  1133. MODEL_TENSOR.OUTPUT_NORM,
  1134. MODEL_TENSOR.OUTPUT,
  1135. MODEL_TENSOR.ATTN_NORM,
  1136. MODEL_TENSOR.ATTN_QKV,
  1137. MODEL_TENSOR.ATTN_Q,
  1138. MODEL_TENSOR.ATTN_K,
  1139. MODEL_TENSOR.ATTN_V,
  1140. MODEL_TENSOR.ATTN_OUT,
  1141. MODEL_TENSOR.FFN_NORM,
  1142. MODEL_TENSOR.FFN_DOWN,
  1143. MODEL_TENSOR.FFN_UP,
  1144. ],
  1145. MODEL_ARCH.PHI3: [
  1146. MODEL_TENSOR.TOKEN_EMBD,
  1147. MODEL_TENSOR.OUTPUT_NORM,
  1148. MODEL_TENSOR.OUTPUT,
  1149. MODEL_TENSOR.ROPE_FACTORS_LONG,
  1150. MODEL_TENSOR.ROPE_FACTORS_SHORT,
  1151. MODEL_TENSOR.ATTN_NORM,
  1152. MODEL_TENSOR.ATTN_QKV,
  1153. MODEL_TENSOR.ATTN_Q,
  1154. MODEL_TENSOR.ATTN_K,
  1155. MODEL_TENSOR.ATTN_V,
  1156. MODEL_TENSOR.ATTN_OUT,
  1157. MODEL_TENSOR.FFN_NORM,
  1158. MODEL_TENSOR.FFN_DOWN,
  1159. MODEL_TENSOR.FFN_UP,
  1160. ],
  1161. MODEL_ARCH.PHIMOE: [
  1162. MODEL_TENSOR.TOKEN_EMBD,
  1163. MODEL_TENSOR.OUTPUT_NORM,
  1164. MODEL_TENSOR.OUTPUT,
  1165. MODEL_TENSOR.ROPE_FACTORS_LONG,
  1166. MODEL_TENSOR.ROPE_FACTORS_SHORT,
  1167. MODEL_TENSOR.ATTN_NORM,
  1168. MODEL_TENSOR.ATTN_QKV,
  1169. MODEL_TENSOR.ATTN_Q,
  1170. MODEL_TENSOR.ATTN_K,
  1171. MODEL_TENSOR.ATTN_V,
  1172. MODEL_TENSOR.ATTN_OUT,
  1173. MODEL_TENSOR.FFN_NORM,
  1174. MODEL_TENSOR.FFN_GATE_INP,
  1175. MODEL_TENSOR.FFN_GATE_EXP,
  1176. MODEL_TENSOR.FFN_DOWN_EXP,
  1177. MODEL_TENSOR.FFN_UP_EXP,
  1178. ],
  1179. MODEL_ARCH.CODESHELL: [
  1180. MODEL_TENSOR.TOKEN_EMBD,
  1181. MODEL_TENSOR.POS_EMBD,
  1182. MODEL_TENSOR.OUTPUT_NORM,
  1183. MODEL_TENSOR.OUTPUT,
  1184. MODEL_TENSOR.ATTN_NORM,
  1185. MODEL_TENSOR.ATTN_QKV,
  1186. MODEL_TENSOR.ATTN_OUT,
  1187. MODEL_TENSOR.ATTN_ROT_EMBD,
  1188. MODEL_TENSOR.FFN_NORM,
  1189. MODEL_TENSOR.FFN_DOWN,
  1190. MODEL_TENSOR.FFN_UP,
  1191. ],
  1192. MODEL_ARCH.ORION: [
  1193. MODEL_TENSOR.TOKEN_EMBD,
  1194. MODEL_TENSOR.OUTPUT_NORM,
  1195. MODEL_TENSOR.OUTPUT,
  1196. MODEL_TENSOR.ROPE_FREQS,
  1197. MODEL_TENSOR.ATTN_NORM,
  1198. MODEL_TENSOR.ATTN_Q,
  1199. MODEL_TENSOR.ATTN_K,
  1200. MODEL_TENSOR.ATTN_V,
  1201. MODEL_TENSOR.ATTN_OUT,
  1202. MODEL_TENSOR.ATTN_ROT_EMBD,
  1203. MODEL_TENSOR.FFN_NORM,
  1204. MODEL_TENSOR.FFN_GATE,
  1205. MODEL_TENSOR.FFN_DOWN,
  1206. MODEL_TENSOR.FFN_UP,
  1207. ],
  1208. MODEL_ARCH.INTERNLM2: [
  1209. MODEL_TENSOR.TOKEN_EMBD,
  1210. MODEL_TENSOR.OUTPUT_NORM,
  1211. MODEL_TENSOR.OUTPUT,
  1212. MODEL_TENSOR.ATTN_NORM,
  1213. MODEL_TENSOR.ATTN_Q,
  1214. MODEL_TENSOR.ATTN_K,
  1215. MODEL_TENSOR.ATTN_V,
  1216. MODEL_TENSOR.ATTN_OUT,
  1217. MODEL_TENSOR.ATTN_ROT_EMBD,
  1218. MODEL_TENSOR.FFN_NORM,
  1219. MODEL_TENSOR.FFN_GATE,
  1220. MODEL_TENSOR.FFN_DOWN,
  1221. MODEL_TENSOR.FFN_UP,
  1222. ],
  1223. MODEL_ARCH.MINICPM: [
  1224. MODEL_TENSOR.TOKEN_EMBD,
  1225. MODEL_TENSOR.OUTPUT,
  1226. MODEL_TENSOR.OUTPUT_NORM,
  1227. MODEL_TENSOR.ROPE_FREQS,
  1228. MODEL_TENSOR.ROPE_FACTORS_LONG,
  1229. MODEL_TENSOR.ROPE_FACTORS_SHORT,
  1230. MODEL_TENSOR.ATTN_NORM,
  1231. MODEL_TENSOR.ATTN_Q,
  1232. MODEL_TENSOR.ATTN_K,
  1233. MODEL_TENSOR.ATTN_V,
  1234. MODEL_TENSOR.ATTN_OUT,
  1235. MODEL_TENSOR.ATTN_ROT_EMBD,
  1236. MODEL_TENSOR.FFN_GATE_INP,
  1237. MODEL_TENSOR.FFN_NORM,
  1238. MODEL_TENSOR.FFN_GATE,
  1239. MODEL_TENSOR.FFN_DOWN,
  1240. MODEL_TENSOR.FFN_UP,
  1241. MODEL_TENSOR.FFN_GATE_EXP,
  1242. MODEL_TENSOR.FFN_DOWN_EXP,
  1243. MODEL_TENSOR.FFN_UP_EXP,
  1244. ],
  1245. MODEL_ARCH.MINICPM3: [
  1246. MODEL_TENSOR.TOKEN_EMBD,
  1247. MODEL_TENSOR.OUTPUT_NORM,
  1248. MODEL_TENSOR.OUTPUT,
  1249. MODEL_TENSOR.ROPE_FACTORS_LONG,
  1250. MODEL_TENSOR.ROPE_FACTORS_SHORT,
  1251. MODEL_TENSOR.ATTN_NORM,
  1252. MODEL_TENSOR.ATTN_Q_A,
  1253. MODEL_TENSOR.ATTN_Q_B,
  1254. MODEL_TENSOR.ATTN_KV_A_MQA,
  1255. MODEL_TENSOR.ATTN_KV_B,
  1256. MODEL_TENSOR.ATTN_Q_A_NORM,
  1257. MODEL_TENSOR.ATTN_KV_A_NORM,
  1258. MODEL_TENSOR.ATTN_OUT,
  1259. MODEL_TENSOR.FFN_NORM,
  1260. MODEL_TENSOR.FFN_GATE,
  1261. MODEL_TENSOR.FFN_DOWN,
  1262. MODEL_TENSOR.FFN_UP,
  1263. ],
  1264. MODEL_ARCH.GEMMA: [
  1265. MODEL_TENSOR.TOKEN_EMBD,
  1266. MODEL_TENSOR.OUTPUT_NORM,
  1267. MODEL_TENSOR.ATTN_NORM,
  1268. MODEL_TENSOR.ATTN_Q,
  1269. MODEL_TENSOR.ATTN_K,
  1270. MODEL_TENSOR.ATTN_V,
  1271. MODEL_TENSOR.ATTN_OUT,
  1272. MODEL_TENSOR.FFN_GATE,
  1273. MODEL_TENSOR.FFN_DOWN,
  1274. MODEL_TENSOR.FFN_UP,
  1275. MODEL_TENSOR.FFN_NORM,
  1276. ],
  1277. MODEL_ARCH.GEMMA2: [
  1278. MODEL_TENSOR.TOKEN_EMBD,
  1279. MODEL_TENSOR.OUTPUT_NORM,
  1280. MODEL_TENSOR.ATTN_Q,
  1281. MODEL_TENSOR.ATTN_K,
  1282. MODEL_TENSOR.ATTN_V,
  1283. MODEL_TENSOR.ATTN_OUT,
  1284. MODEL_TENSOR.FFN_GATE,
  1285. MODEL_TENSOR.FFN_DOWN,
  1286. MODEL_TENSOR.FFN_UP,
  1287. MODEL_TENSOR.ATTN_NORM,
  1288. MODEL_TENSOR.ATTN_POST_NORM,
  1289. MODEL_TENSOR.FFN_PRE_NORM,
  1290. MODEL_TENSOR.FFN_POST_NORM,
  1291. ],
  1292. MODEL_ARCH.GEMMA3: [
  1293. MODEL_TENSOR.TOKEN_EMBD,
  1294. MODEL_TENSOR.OUTPUT,
  1295. MODEL_TENSOR.OUTPUT_NORM,
  1296. MODEL_TENSOR.ATTN_Q,
  1297. MODEL_TENSOR.ATTN_Q_NORM,
  1298. MODEL_TENSOR.ATTN_K,
  1299. MODEL_TENSOR.ATTN_K_NORM,
  1300. MODEL_TENSOR.ATTN_V,
  1301. MODEL_TENSOR.ATTN_OUT,
  1302. MODEL_TENSOR.FFN_GATE,
  1303. MODEL_TENSOR.FFN_DOWN,
  1304. MODEL_TENSOR.FFN_UP,
  1305. MODEL_TENSOR.ATTN_NORM,
  1306. MODEL_TENSOR.ATTN_POST_NORM,
  1307. MODEL_TENSOR.FFN_PRE_NORM,
  1308. MODEL_TENSOR.FFN_POST_NORM,
  1309. ],
  1310. MODEL_ARCH.STARCODER2: [
  1311. MODEL_TENSOR.TOKEN_EMBD,
  1312. MODEL_TENSOR.OUTPUT_NORM,
  1313. MODEL_TENSOR.OUTPUT,
  1314. MODEL_TENSOR.ROPE_FREQS,
  1315. MODEL_TENSOR.ATTN_NORM,
  1316. MODEL_TENSOR.ATTN_Q,
  1317. MODEL_TENSOR.ATTN_K,
  1318. MODEL_TENSOR.ATTN_V,
  1319. MODEL_TENSOR.ATTN_OUT,
  1320. MODEL_TENSOR.ATTN_ROT_EMBD,
  1321. MODEL_TENSOR.FFN_NORM,
  1322. MODEL_TENSOR.FFN_DOWN,
  1323. MODEL_TENSOR.FFN_UP,
  1324. ],
  1325. MODEL_ARCH.RWKV6: [
  1326. MODEL_TENSOR.TOKEN_EMBD,
  1327. MODEL_TENSOR.TOKEN_EMBD_NORM,
  1328. MODEL_TENSOR.OUTPUT_NORM,
  1329. MODEL_TENSOR.OUTPUT,
  1330. MODEL_TENSOR.ATTN_NORM,
  1331. MODEL_TENSOR.ATTN_NORM_2,
  1332. MODEL_TENSOR.TIME_MIX_W1,
  1333. MODEL_TENSOR.TIME_MIX_W2,
  1334. MODEL_TENSOR.TIME_MIX_LERP_X,
  1335. MODEL_TENSOR.TIME_MIX_LERP_K,
  1336. MODEL_TENSOR.TIME_MIX_LERP_V,
  1337. MODEL_TENSOR.TIME_MIX_LERP_R,
  1338. MODEL_TENSOR.TIME_MIX_LERP_G,
  1339. MODEL_TENSOR.TIME_MIX_LERP_W,
  1340. MODEL_TENSOR.TIME_MIX_LERP_FUSED,
  1341. MODEL_TENSOR.TIME_MIX_FIRST,
  1342. MODEL_TENSOR.TIME_MIX_DECAY,
  1343. MODEL_TENSOR.TIME_MIX_DECAY_W1,
  1344. MODEL_TENSOR.TIME_MIX_DECAY_W2,
  1345. MODEL_TENSOR.TIME_MIX_KEY,
  1346. MODEL_TENSOR.TIME_MIX_VALUE,
  1347. MODEL_TENSOR.TIME_MIX_RECEPTANCE,
  1348. MODEL_TENSOR.TIME_MIX_GATE,
  1349. MODEL_TENSOR.TIME_MIX_LN,
  1350. MODEL_TENSOR.TIME_MIX_OUTPUT,
  1351. MODEL_TENSOR.CHANNEL_MIX_LERP_K,
  1352. MODEL_TENSOR.CHANNEL_MIX_LERP_R,
  1353. MODEL_TENSOR.CHANNEL_MIX_KEY,
  1354. MODEL_TENSOR.CHANNEL_MIX_RECEPTANCE,
  1355. MODEL_TENSOR.CHANNEL_MIX_VALUE,
  1356. ],
  1357. MODEL_ARCH.RWKV6QWEN2: [
  1358. MODEL_TENSOR.TOKEN_EMBD,
  1359. MODEL_TENSOR.OUTPUT_NORM,
  1360. MODEL_TENSOR.OUTPUT,
  1361. MODEL_TENSOR.ATTN_NORM,
  1362. MODEL_TENSOR.TIME_MIX_W1,
  1363. MODEL_TENSOR.TIME_MIX_W2,
  1364. MODEL_TENSOR.TIME_MIX_LERP_X,
  1365. MODEL_TENSOR.TIME_MIX_LERP_K,
  1366. MODEL_TENSOR.TIME_MIX_LERP_V,
  1367. MODEL_TENSOR.TIME_MIX_LERP_R,
  1368. MODEL_TENSOR.TIME_MIX_LERP_G,
  1369. MODEL_TENSOR.TIME_MIX_LERP_W,
  1370. MODEL_TENSOR.TIME_MIX_LERP_FUSED,
  1371. MODEL_TENSOR.TIME_MIX_FIRST,
  1372. MODEL_TENSOR.TIME_MIX_DECAY,
  1373. MODEL_TENSOR.TIME_MIX_DECAY_W1,
  1374. MODEL_TENSOR.TIME_MIX_DECAY_W2,
  1375. MODEL_TENSOR.TIME_MIX_KEY,
  1376. MODEL_TENSOR.TIME_MIX_VALUE,
  1377. MODEL_TENSOR.TIME_MIX_RECEPTANCE,
  1378. MODEL_TENSOR.TIME_MIX_GATE,
  1379. MODEL_TENSOR.TIME_MIX_LN,
  1380. MODEL_TENSOR.TIME_MIX_OUTPUT,
  1381. MODEL_TENSOR.FFN_NORM,
  1382. MODEL_TENSOR.FFN_GATE,
  1383. MODEL_TENSOR.FFN_DOWN,
  1384. MODEL_TENSOR.FFN_UP,
  1385. ],
  1386. MODEL_ARCH.RWKV7: [
  1387. MODEL_TENSOR.TOKEN_EMBD,
  1388. MODEL_TENSOR.TOKEN_EMBD_NORM,
  1389. MODEL_TENSOR.OUTPUT_NORM,
  1390. MODEL_TENSOR.OUTPUT,
  1391. MODEL_TENSOR.ATTN_NORM,
  1392. MODEL_TENSOR.ATTN_NORM_2,
  1393. MODEL_TENSOR.TIME_MIX_LERP_FUSED,
  1394. MODEL_TENSOR.TIME_MIX_W0,
  1395. MODEL_TENSOR.TIME_MIX_W1,
  1396. MODEL_TENSOR.TIME_MIX_W2,
  1397. MODEL_TENSOR.TIME_MIX_A0,
  1398. MODEL_TENSOR.TIME_MIX_A1,
  1399. MODEL_TENSOR.TIME_MIX_A2,
  1400. MODEL_TENSOR.TIME_MIX_V0,
  1401. MODEL_TENSOR.TIME_MIX_V1,
  1402. MODEL_TENSOR.TIME_MIX_V2,
  1403. MODEL_TENSOR.TIME_MIX_G1,
  1404. MODEL_TENSOR.TIME_MIX_G2,
  1405. MODEL_TENSOR.TIME_MIX_K_K,
  1406. MODEL_TENSOR.TIME_MIX_K_A,
  1407. MODEL_TENSOR.TIME_MIX_R_K,
  1408. MODEL_TENSOR.TIME_MIX_KEY,
  1409. MODEL_TENSOR.TIME_MIX_VALUE,
  1410. MODEL_TENSOR.TIME_MIX_RECEPTANCE,
  1411. MODEL_TENSOR.TIME_MIX_LN,
  1412. MODEL_TENSOR.TIME_MIX_OUTPUT,
  1413. MODEL_TENSOR.CHANNEL_MIX_LERP_K,
  1414. MODEL_TENSOR.CHANNEL_MIX_KEY,
  1415. MODEL_TENSOR.CHANNEL_MIX_VALUE,
  1416. ],
  1417. MODEL_ARCH.ARWKV7: [
  1418. MODEL_TENSOR.TOKEN_EMBD,
  1419. MODEL_TENSOR.TOKEN_EMBD_NORM,
  1420. MODEL_TENSOR.OUTPUT_NORM,
  1421. MODEL_TENSOR.OUTPUT,
  1422. MODEL_TENSOR.ATTN_NORM,
  1423. MODEL_TENSOR.TIME_MIX_LERP_FUSED,
  1424. MODEL_TENSOR.TIME_MIX_W0,
  1425. MODEL_TENSOR.TIME_MIX_W1,
  1426. MODEL_TENSOR.TIME_MIX_W2,
  1427. MODEL_TENSOR.TIME_MIX_A0,
  1428. MODEL_TENSOR.TIME_MIX_A1,
  1429. MODEL_TENSOR.TIME_MIX_A2,
  1430. MODEL_TENSOR.TIME_MIX_V0,
  1431. MODEL_TENSOR.TIME_MIX_V1,
  1432. MODEL_TENSOR.TIME_MIX_V2,
  1433. MODEL_TENSOR.TIME_MIX_G1,
  1434. MODEL_TENSOR.TIME_MIX_G2,
  1435. MODEL_TENSOR.TIME_MIX_K_K,
  1436. MODEL_TENSOR.TIME_MIX_K_A,
  1437. MODEL_TENSOR.TIME_MIX_R_K,
  1438. MODEL_TENSOR.TIME_MIX_KEY,
  1439. MODEL_TENSOR.TIME_MIX_VALUE,
  1440. MODEL_TENSOR.TIME_MIX_RECEPTANCE,
  1441. MODEL_TENSOR.TIME_MIX_LN,
  1442. MODEL_TENSOR.TIME_MIX_OUTPUT,
  1443. MODEL_TENSOR.FFN_NORM,
  1444. MODEL_TENSOR.FFN_GATE,
  1445. MODEL_TENSOR.FFN_DOWN,
  1446. MODEL_TENSOR.FFN_UP,
  1447. ],
  1448. MODEL_ARCH.MAMBA: [
  1449. MODEL_TENSOR.TOKEN_EMBD,
  1450. MODEL_TENSOR.OUTPUT_NORM,
  1451. MODEL_TENSOR.OUTPUT,
  1452. MODEL_TENSOR.ATTN_NORM,
  1453. MODEL_TENSOR.SSM_IN,
  1454. MODEL_TENSOR.SSM_CONV1D,
  1455. MODEL_TENSOR.SSM_X,
  1456. MODEL_TENSOR.SSM_DT,
  1457. MODEL_TENSOR.SSM_A,
  1458. MODEL_TENSOR.SSM_D,
  1459. MODEL_TENSOR.SSM_OUT,
  1460. ],
  1461. MODEL_ARCH.XVERSE: [
  1462. MODEL_TENSOR.TOKEN_EMBD,
  1463. MODEL_TENSOR.OUTPUT_NORM,
  1464. MODEL_TENSOR.OUTPUT,
  1465. MODEL_TENSOR.ROPE_FREQS,
  1466. MODEL_TENSOR.ATTN_NORM,
  1467. MODEL_TENSOR.ATTN_Q,
  1468. MODEL_TENSOR.ATTN_K,
  1469. MODEL_TENSOR.ATTN_V,
  1470. MODEL_TENSOR.ATTN_OUT,
  1471. MODEL_TENSOR.ATTN_ROT_EMBD,
  1472. MODEL_TENSOR.FFN_NORM,
  1473. MODEL_TENSOR.FFN_GATE,
  1474. MODEL_TENSOR.FFN_DOWN,
  1475. MODEL_TENSOR.FFN_UP,
  1476. ],
  1477. MODEL_ARCH.COMMAND_R: [
  1478. MODEL_TENSOR.TOKEN_EMBD,
  1479. MODEL_TENSOR.OUTPUT_NORM,
  1480. MODEL_TENSOR.ATTN_NORM,
  1481. MODEL_TENSOR.ATTN_Q,
  1482. MODEL_TENSOR.ATTN_K,
  1483. MODEL_TENSOR.ATTN_V,
  1484. MODEL_TENSOR.ATTN_OUT,
  1485. MODEL_TENSOR.FFN_GATE,
  1486. MODEL_TENSOR.FFN_DOWN,
  1487. MODEL_TENSOR.FFN_UP,
  1488. MODEL_TENSOR.ATTN_K_NORM,
  1489. MODEL_TENSOR.ATTN_Q_NORM,
  1490. ],
  1491. MODEL_ARCH.COHERE2: [
  1492. MODEL_TENSOR.TOKEN_EMBD,
  1493. MODEL_TENSOR.OUTPUT_NORM,
  1494. MODEL_TENSOR.ATTN_NORM,
  1495. MODEL_TENSOR.ATTN_Q,
  1496. MODEL_TENSOR.ATTN_K,
  1497. MODEL_TENSOR.ATTN_V,
  1498. MODEL_TENSOR.ATTN_OUT,
  1499. MODEL_TENSOR.FFN_GATE,
  1500. MODEL_TENSOR.FFN_DOWN,
  1501. MODEL_TENSOR.FFN_UP,
  1502. ],
  1503. MODEL_ARCH.DBRX: [
  1504. MODEL_TENSOR.TOKEN_EMBD,
  1505. MODEL_TENSOR.OUTPUT_NORM,
  1506. MODEL_TENSOR.OUTPUT,
  1507. MODEL_TENSOR.ATTN_NORM,
  1508. MODEL_TENSOR.ATTN_QKV,
  1509. MODEL_TENSOR.ATTN_OUT,
  1510. MODEL_TENSOR.ATTN_OUT_NORM,
  1511. MODEL_TENSOR.FFN_GATE_INP,
  1512. MODEL_TENSOR.FFN_GATE_EXP,
  1513. MODEL_TENSOR.FFN_DOWN_EXP,
  1514. MODEL_TENSOR.FFN_UP_EXP,
  1515. ],
  1516. MODEL_ARCH.OLMO: [
  1517. MODEL_TENSOR.TOKEN_EMBD,
  1518. MODEL_TENSOR.OUTPUT,
  1519. MODEL_TENSOR.ATTN_Q,
  1520. MODEL_TENSOR.ATTN_K,
  1521. MODEL_TENSOR.ATTN_V,
  1522. MODEL_TENSOR.ATTN_OUT,
  1523. MODEL_TENSOR.FFN_GATE,
  1524. MODEL_TENSOR.FFN_DOWN,
  1525. MODEL_TENSOR.FFN_UP,
  1526. ],
  1527. MODEL_ARCH.OLMO2: [
  1528. MODEL_TENSOR.TOKEN_EMBD,
  1529. MODEL_TENSOR.OUTPUT_NORM,
  1530. MODEL_TENSOR.OUTPUT,
  1531. MODEL_TENSOR.ATTN_Q,
  1532. MODEL_TENSOR.ATTN_K,
  1533. MODEL_TENSOR.ATTN_V,
  1534. MODEL_TENSOR.ATTN_OUT,
  1535. MODEL_TENSOR.ATTN_POST_NORM,
  1536. MODEL_TENSOR.ATTN_Q_NORM,
  1537. MODEL_TENSOR.ATTN_K_NORM,
  1538. MODEL_TENSOR.FFN_POST_NORM,
  1539. MODEL_TENSOR.FFN_GATE,
  1540. MODEL_TENSOR.FFN_DOWN,
  1541. MODEL_TENSOR.FFN_UP,
  1542. ],
  1543. MODEL_ARCH.OLMOE: [
  1544. MODEL_TENSOR.TOKEN_EMBD,
  1545. MODEL_TENSOR.OUTPUT_NORM,
  1546. MODEL_TENSOR.OUTPUT,
  1547. MODEL_TENSOR.ATTN_OUT,
  1548. MODEL_TENSOR.ATTN_Q,
  1549. MODEL_TENSOR.ATTN_K,
  1550. MODEL_TENSOR.ATTN_V,
  1551. MODEL_TENSOR.ATTN_NORM,
  1552. MODEL_TENSOR.ATTN_Q_NORM,
  1553. MODEL_TENSOR.ATTN_K_NORM,
  1554. MODEL_TENSOR.FFN_NORM,
  1555. MODEL_TENSOR.FFN_GATE_INP,
  1556. MODEL_TENSOR.FFN_GATE_EXP,
  1557. MODEL_TENSOR.FFN_UP_EXP,
  1558. MODEL_TENSOR.FFN_DOWN_EXP,
  1559. ],
  1560. MODEL_ARCH.OPENELM: [
  1561. MODEL_TENSOR.TOKEN_EMBD,
  1562. MODEL_TENSOR.OUTPUT_NORM,
  1563. MODEL_TENSOR.ATTN_NORM,
  1564. MODEL_TENSOR.ATTN_QKV,
  1565. MODEL_TENSOR.ATTN_Q_NORM,
  1566. MODEL_TENSOR.ATTN_K_NORM,
  1567. MODEL_TENSOR.ATTN_OUT,
  1568. MODEL_TENSOR.FFN_NORM,
  1569. MODEL_TENSOR.FFN_GATE,
  1570. MODEL_TENSOR.FFN_DOWN,
  1571. MODEL_TENSOR.FFN_UP,
  1572. ],
  1573. MODEL_ARCH.ARCTIC: [
  1574. MODEL_TENSOR.TOKEN_EMBD,
  1575. MODEL_TENSOR.OUTPUT_NORM,
  1576. MODEL_TENSOR.OUTPUT,
  1577. MODEL_TENSOR.ROPE_FREQS,
  1578. MODEL_TENSOR.ATTN_NORM,
  1579. MODEL_TENSOR.ATTN_Q,
  1580. MODEL_TENSOR.ATTN_K,
  1581. MODEL_TENSOR.ATTN_V,
  1582. MODEL_TENSOR.ATTN_OUT,
  1583. MODEL_TENSOR.ATTN_ROT_EMBD,
  1584. MODEL_TENSOR.FFN_GATE_INP,
  1585. MODEL_TENSOR.FFN_NORM,
  1586. MODEL_TENSOR.FFN_GATE,
  1587. MODEL_TENSOR.FFN_DOWN,
  1588. MODEL_TENSOR.FFN_UP,
  1589. MODEL_TENSOR.FFN_NORM_EXP,
  1590. MODEL_TENSOR.FFN_GATE_EXP,
  1591. MODEL_TENSOR.FFN_DOWN_EXP,
  1592. MODEL_TENSOR.FFN_UP_EXP,
  1593. ],
  1594. MODEL_ARCH.DEEPSEEK: [
  1595. MODEL_TENSOR.TOKEN_EMBD,
  1596. MODEL_TENSOR.OUTPUT_NORM,
  1597. MODEL_TENSOR.OUTPUT,
  1598. MODEL_TENSOR.ROPE_FREQS,
  1599. MODEL_TENSOR.ATTN_NORM,
  1600. MODEL_TENSOR.ATTN_Q,
  1601. MODEL_TENSOR.ATTN_K,
  1602. MODEL_TENSOR.ATTN_V,
  1603. MODEL_TENSOR.ATTN_OUT,
  1604. MODEL_TENSOR.ATTN_ROT_EMBD,
  1605. MODEL_TENSOR.FFN_GATE_INP,
  1606. MODEL_TENSOR.FFN_NORM,
  1607. MODEL_TENSOR.FFN_GATE,
  1608. MODEL_TENSOR.FFN_DOWN,
  1609. MODEL_TENSOR.FFN_UP,
  1610. MODEL_TENSOR.FFN_GATE_EXP,
  1611. MODEL_TENSOR.FFN_DOWN_EXP,
  1612. MODEL_TENSOR.FFN_UP_EXP,
  1613. MODEL_TENSOR.FFN_GATE_SHEXP,
  1614. MODEL_TENSOR.FFN_DOWN_SHEXP,
  1615. MODEL_TENSOR.FFN_UP_SHEXP,
  1616. ],
  1617. MODEL_ARCH.DEEPSEEK2: [
  1618. MODEL_TENSOR.TOKEN_EMBD,
  1619. MODEL_TENSOR.OUTPUT_NORM,
  1620. MODEL_TENSOR.OUTPUT,
  1621. MODEL_TENSOR.ROPE_FREQS,
  1622. MODEL_TENSOR.ATTN_NORM,
  1623. MODEL_TENSOR.ATTN_Q,
  1624. MODEL_TENSOR.ATTN_Q_A,
  1625. MODEL_TENSOR.ATTN_Q_B,
  1626. MODEL_TENSOR.ATTN_KV_A_MQA,
  1627. MODEL_TENSOR.ATTN_KV_B,
  1628. MODEL_TENSOR.ATTN_K_B,
  1629. MODEL_TENSOR.ATTN_V_B,
  1630. MODEL_TENSOR.ATTN_Q_A_NORM,
  1631. MODEL_TENSOR.ATTN_KV_A_NORM,
  1632. MODEL_TENSOR.ATTN_OUT,
  1633. MODEL_TENSOR.ATTN_ROT_EMBD,
  1634. MODEL_TENSOR.FFN_GATE_INP,
  1635. MODEL_TENSOR.FFN_NORM,
  1636. MODEL_TENSOR.FFN_GATE,
  1637. MODEL_TENSOR.FFN_DOWN,
  1638. MODEL_TENSOR.FFN_UP,
  1639. MODEL_TENSOR.FFN_GATE_EXP,
  1640. MODEL_TENSOR.FFN_DOWN_EXP,
  1641. MODEL_TENSOR.FFN_UP_EXP,
  1642. MODEL_TENSOR.FFN_GATE_SHEXP,
  1643. MODEL_TENSOR.FFN_DOWN_SHEXP,
  1644. MODEL_TENSOR.FFN_UP_SHEXP,
  1645. MODEL_TENSOR.FFN_EXP_PROBS_B,
  1646. ],
  1647. MODEL_ARCH.PLM: [
  1648. MODEL_TENSOR.TOKEN_EMBD,
  1649. MODEL_TENSOR.OUTPUT,
  1650. MODEL_TENSOR.OUTPUT_NORM,
  1651. MODEL_TENSOR.ATTN_NORM,
  1652. MODEL_TENSOR.ATTN_Q,
  1653. MODEL_TENSOR.ATTN_KV_A_MQA,
  1654. MODEL_TENSOR.ATTN_KV_A_NORM,
  1655. MODEL_TENSOR.ATTN_KV_B,
  1656. MODEL_TENSOR.ATTN_OUT,
  1657. MODEL_TENSOR.FFN_NORM,
  1658. MODEL_TENSOR.FFN_UP,
  1659. MODEL_TENSOR.FFN_DOWN,
  1660. ],
  1661. MODEL_ARCH.CHATGLM : [
  1662. MODEL_TENSOR.TOKEN_EMBD,
  1663. MODEL_TENSOR.ROPE_FREQS,
  1664. MODEL_TENSOR.OUTPUT_NORM,
  1665. MODEL_TENSOR.OUTPUT,
  1666. MODEL_TENSOR.ATTN_NORM,
  1667. MODEL_TENSOR.ATTN_QKV,
  1668. MODEL_TENSOR.ATTN_Q,
  1669. MODEL_TENSOR.ATTN_K,
  1670. MODEL_TENSOR.ATTN_V,
  1671. MODEL_TENSOR.ATTN_OUT,
  1672. MODEL_TENSOR.FFN_NORM,
  1673. MODEL_TENSOR.FFN_DOWN,
  1674. MODEL_TENSOR.FFN_UP,
  1675. ],
  1676. MODEL_ARCH.GLM4 : [
  1677. MODEL_TENSOR.TOKEN_EMBD,
  1678. MODEL_TENSOR.ROPE_FREQS,
  1679. MODEL_TENSOR.OUTPUT_NORM,
  1680. MODEL_TENSOR.OUTPUT,
  1681. MODEL_TENSOR.ATTN_NORM,
  1682. MODEL_TENSOR.ATTN_QKV,
  1683. MODEL_TENSOR.ATTN_Q,
  1684. MODEL_TENSOR.ATTN_K,
  1685. MODEL_TENSOR.ATTN_V,
  1686. MODEL_TENSOR.ATTN_OUT,
  1687. MODEL_TENSOR.FFN_NORM,
  1688. MODEL_TENSOR.FFN_DOWN,
  1689. MODEL_TENSOR.FFN_UP,
  1690. MODEL_TENSOR.ATTN_POST_NORM,
  1691. MODEL_TENSOR.FFN_POST_NORM,
  1692. ],
  1693. MODEL_ARCH.BITNET: [
  1694. MODEL_TENSOR.ATTN_Q,
  1695. MODEL_TENSOR.ATTN_K,
  1696. MODEL_TENSOR.ATTN_V,
  1697. MODEL_TENSOR.TOKEN_EMBD,
  1698. MODEL_TENSOR.OUTPUT_NORM,
  1699. MODEL_TENSOR.ATTN_NORM,
  1700. MODEL_TENSOR.ATTN_OUT,
  1701. MODEL_TENSOR.FFN_NORM,
  1702. MODEL_TENSOR.FFN_GATE,
  1703. MODEL_TENSOR.FFN_DOWN,
  1704. MODEL_TENSOR.FFN_UP,
  1705. MODEL_TENSOR.ATTN_SUB_NORM,
  1706. MODEL_TENSOR.FFN_SUB_NORM,
  1707. ],
  1708. MODEL_ARCH.T5: [
  1709. MODEL_TENSOR.TOKEN_EMBD,
  1710. MODEL_TENSOR.OUTPUT,
  1711. MODEL_TENSOR.DEC_ATTN_NORM,
  1712. MODEL_TENSOR.DEC_ATTN_Q,
  1713. MODEL_TENSOR.DEC_ATTN_K,
  1714. MODEL_TENSOR.DEC_ATTN_V,
  1715. MODEL_TENSOR.DEC_ATTN_OUT,
  1716. MODEL_TENSOR.DEC_ATTN_REL_B,
  1717. MODEL_TENSOR.DEC_CROSS_ATTN_NORM,
  1718. MODEL_TENSOR.DEC_CROSS_ATTN_Q,
  1719. MODEL_TENSOR.DEC_CROSS_ATTN_K,
  1720. MODEL_TENSOR.DEC_CROSS_ATTN_V,
  1721. MODEL_TENSOR.DEC_CROSS_ATTN_OUT,
  1722. MODEL_TENSOR.DEC_CROSS_ATTN_REL_B,
  1723. MODEL_TENSOR.DEC_FFN_NORM,
  1724. MODEL_TENSOR.DEC_FFN_GATE,
  1725. MODEL_TENSOR.DEC_FFN_DOWN,
  1726. MODEL_TENSOR.DEC_FFN_UP,
  1727. MODEL_TENSOR.DEC_OUTPUT_NORM,
  1728. MODEL_TENSOR.ENC_ATTN_NORM,
  1729. MODEL_TENSOR.ENC_ATTN_Q,
  1730. MODEL_TENSOR.ENC_ATTN_K,
  1731. MODEL_TENSOR.ENC_ATTN_V,
  1732. MODEL_TENSOR.ENC_ATTN_OUT,
  1733. MODEL_TENSOR.ENC_ATTN_REL_B,
  1734. MODEL_TENSOR.ENC_FFN_NORM,
  1735. MODEL_TENSOR.ENC_FFN_GATE,
  1736. MODEL_TENSOR.ENC_FFN_DOWN,
  1737. MODEL_TENSOR.ENC_FFN_UP,
  1738. MODEL_TENSOR.ENC_OUTPUT_NORM,
  1739. ],
  1740. MODEL_ARCH.T5ENCODER: [
  1741. MODEL_TENSOR.TOKEN_EMBD,
  1742. MODEL_TENSOR.OUTPUT,
  1743. MODEL_TENSOR.ENC_ATTN_NORM,
  1744. MODEL_TENSOR.ENC_ATTN_Q,
  1745. MODEL_TENSOR.ENC_ATTN_K,
  1746. MODEL_TENSOR.ENC_ATTN_V,
  1747. MODEL_TENSOR.ENC_ATTN_OUT,
  1748. MODEL_TENSOR.ENC_ATTN_REL_B,
  1749. MODEL_TENSOR.ENC_FFN_NORM,
  1750. MODEL_TENSOR.ENC_FFN_GATE,
  1751. MODEL_TENSOR.ENC_FFN_DOWN,
  1752. MODEL_TENSOR.ENC_FFN_UP,
  1753. MODEL_TENSOR.ENC_OUTPUT_NORM,
  1754. ],
  1755. MODEL_ARCH.JAIS: [
  1756. MODEL_TENSOR.TOKEN_EMBD,
  1757. MODEL_TENSOR.OUTPUT_NORM,
  1758. MODEL_TENSOR.OUTPUT,
  1759. MODEL_TENSOR.ATTN_NORM,
  1760. MODEL_TENSOR.ATTN_QKV,
  1761. MODEL_TENSOR.ATTN_OUT,
  1762. MODEL_TENSOR.FFN_NORM,
  1763. MODEL_TENSOR.FFN_DOWN,
  1764. MODEL_TENSOR.FFN_GATE,
  1765. MODEL_TENSOR.FFN_UP,
  1766. ],
  1767. MODEL_ARCH.NEMOTRON: [
  1768. MODEL_TENSOR.TOKEN_EMBD,
  1769. MODEL_TENSOR.OUTPUT_NORM,
  1770. MODEL_TENSOR.OUTPUT,
  1771. MODEL_TENSOR.ROPE_FREQS,
  1772. MODEL_TENSOR.ATTN_NORM,
  1773. MODEL_TENSOR.ATTN_Q,
  1774. MODEL_TENSOR.ATTN_K,
  1775. MODEL_TENSOR.ATTN_V,
  1776. MODEL_TENSOR.ATTN_OUT,
  1777. MODEL_TENSOR.ATTN_ROT_EMBD,
  1778. MODEL_TENSOR.FFN_NORM,
  1779. MODEL_TENSOR.FFN_DOWN,
  1780. MODEL_TENSOR.FFN_UP,
  1781. ],
  1782. MODEL_ARCH.EXAONE: [
  1783. MODEL_TENSOR.TOKEN_EMBD,
  1784. MODEL_TENSOR.OUTPUT_NORM,
  1785. MODEL_TENSOR.OUTPUT,
  1786. MODEL_TENSOR.ROPE_FREQS,
  1787. MODEL_TENSOR.ATTN_NORM,
  1788. MODEL_TENSOR.ATTN_Q,
  1789. MODEL_TENSOR.ATTN_K,
  1790. MODEL_TENSOR.ATTN_V,
  1791. MODEL_TENSOR.ATTN_OUT,
  1792. MODEL_TENSOR.ATTN_ROT_EMBD,
  1793. MODEL_TENSOR.FFN_NORM,
  1794. MODEL_TENSOR.FFN_GATE,
  1795. MODEL_TENSOR.FFN_DOWN,
  1796. MODEL_TENSOR.FFN_UP,
  1797. ],
  1798. MODEL_ARCH.GRANITE: [
  1799. MODEL_TENSOR.TOKEN_EMBD,
  1800. MODEL_TENSOR.OUTPUT_NORM,
  1801. MODEL_TENSOR.OUTPUT,
  1802. MODEL_TENSOR.ATTN_NORM,
  1803. MODEL_TENSOR.ATTN_Q,
  1804. MODEL_TENSOR.ATTN_K,
  1805. MODEL_TENSOR.ATTN_V,
  1806. MODEL_TENSOR.ATTN_OUT,
  1807. MODEL_TENSOR.FFN_NORM,
  1808. MODEL_TENSOR.FFN_GATE,
  1809. MODEL_TENSOR.FFN_DOWN,
  1810. MODEL_TENSOR.FFN_UP,
  1811. ],
  1812. MODEL_ARCH.GRANITE_MOE: [
  1813. MODEL_TENSOR.TOKEN_EMBD,
  1814. MODEL_TENSOR.OUTPUT_NORM,
  1815. MODEL_TENSOR.OUTPUT,
  1816. MODEL_TENSOR.ATTN_NORM,
  1817. MODEL_TENSOR.ATTN_Q,
  1818. MODEL_TENSOR.ATTN_K,
  1819. MODEL_TENSOR.ATTN_V,
  1820. MODEL_TENSOR.ATTN_OUT,
  1821. MODEL_TENSOR.FFN_NORM,
  1822. MODEL_TENSOR.FFN_GATE_INP,
  1823. MODEL_TENSOR.FFN_GATE_EXP,
  1824. MODEL_TENSOR.FFN_DOWN_EXP,
  1825. MODEL_TENSOR.FFN_UP_EXP,
  1826. ],
  1827. MODEL_ARCH.CHAMELEON: [
  1828. MODEL_TENSOR.TOKEN_EMBD,
  1829. MODEL_TENSOR.OUTPUT_NORM,
  1830. MODEL_TENSOR.OUTPUT,
  1831. MODEL_TENSOR.ATTN_NORM,
  1832. MODEL_TENSOR.ATTN_Q,
  1833. MODEL_TENSOR.ATTN_Q_NORM,
  1834. MODEL_TENSOR.ATTN_K,
  1835. MODEL_TENSOR.ATTN_K_NORM,
  1836. MODEL_TENSOR.ATTN_V,
  1837. MODEL_TENSOR.ATTN_OUT,
  1838. MODEL_TENSOR.FFN_NORM,
  1839. MODEL_TENSOR.FFN_GATE,
  1840. MODEL_TENSOR.FFN_DOWN,
  1841. MODEL_TENSOR.FFN_UP,
  1842. ],
  1843. MODEL_ARCH.WAVTOKENIZER_DEC: [
  1844. MODEL_TENSOR.TOKEN_EMBD,
  1845. MODEL_TENSOR.TOKEN_EMBD_NORM,
  1846. MODEL_TENSOR.CONV1D,
  1847. MODEL_TENSOR.CONVNEXT_DW,
  1848. MODEL_TENSOR.CONVNEXT_NORM,
  1849. MODEL_TENSOR.CONVNEXT_PW1,
  1850. MODEL_TENSOR.CONVNEXT_PW2,
  1851. MODEL_TENSOR.CONVNEXT_GAMMA,
  1852. MODEL_TENSOR.OUTPUT,
  1853. MODEL_TENSOR.OUTPUT_NORM,
  1854. MODEL_TENSOR.POSNET_CONV1,
  1855. MODEL_TENSOR.POSNET_CONV2,
  1856. MODEL_TENSOR.POSNET_NORM,
  1857. MODEL_TENSOR.POSNET_NORM1,
  1858. MODEL_TENSOR.POSNET_NORM2,
  1859. MODEL_TENSOR.POSNET_ATTN_NORM,
  1860. MODEL_TENSOR.POSNET_ATTN_Q,
  1861. MODEL_TENSOR.POSNET_ATTN_K,
  1862. MODEL_TENSOR.POSNET_ATTN_V,
  1863. MODEL_TENSOR.POSNET_ATTN_OUT,
  1864. ],
  1865. MODEL_ARCH.BAILINGMOE: [
  1866. MODEL_TENSOR.TOKEN_EMBD,
  1867. MODEL_TENSOR.OUTPUT_NORM,
  1868. MODEL_TENSOR.OUTPUT,
  1869. MODEL_TENSOR.ROPE_FREQS,
  1870. MODEL_TENSOR.ATTN_NORM,
  1871. MODEL_TENSOR.ATTN_Q,
  1872. MODEL_TENSOR.ATTN_K,
  1873. MODEL_TENSOR.ATTN_V,
  1874. MODEL_TENSOR.ATTN_OUT,
  1875. MODEL_TENSOR.FFN_GATE_INP,
  1876. MODEL_TENSOR.FFN_NORM,
  1877. MODEL_TENSOR.FFN_GATE_EXP,
  1878. MODEL_TENSOR.FFN_DOWN_EXP,
  1879. MODEL_TENSOR.FFN_UP_EXP,
  1880. MODEL_TENSOR.FFN_GATE_SHEXP,
  1881. MODEL_TENSOR.FFN_DOWN_SHEXP,
  1882. MODEL_TENSOR.FFN_UP_SHEXP,
  1883. ],
  1884. # TODO
  1885. }
  1886. # tensors that will not be serialized
  1887. MODEL_TENSOR_SKIP: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
  1888. MODEL_ARCH.LLAMA: [
  1889. MODEL_TENSOR.ROPE_FREQS,
  1890. MODEL_TENSOR.ATTN_ROT_EMBD,
  1891. ],
  1892. MODEL_ARCH.DECI: [
  1893. MODEL_TENSOR.ROPE_FREQS,
  1894. MODEL_TENSOR.ATTN_ROT_EMBD,
  1895. ],
  1896. MODEL_ARCH.BAICHUAN: [
  1897. MODEL_TENSOR.ROPE_FREQS,
  1898. MODEL_TENSOR.ATTN_ROT_EMBD,
  1899. ],
  1900. MODEL_ARCH.QWEN: [
  1901. MODEL_TENSOR.ROPE_FREQS,
  1902. MODEL_TENSOR.ATTN_ROT_EMBD,
  1903. ],
  1904. MODEL_ARCH.CODESHELL: [
  1905. MODEL_TENSOR.ROPE_FREQS,
  1906. MODEL_TENSOR.ATTN_ROT_EMBD,
  1907. ],
  1908. MODEL_ARCH.ORION: [
  1909. MODEL_TENSOR.ROPE_FREQS,
  1910. MODEL_TENSOR.ATTN_ROT_EMBD,
  1911. ],
  1912. MODEL_ARCH.STARCODER2: [
  1913. MODEL_TENSOR.ROPE_FREQS,
  1914. MODEL_TENSOR.ATTN_ROT_EMBD,
  1915. ],
  1916. MODEL_ARCH.XVERSE: [
  1917. MODEL_TENSOR.ROPE_FREQS,
  1918. MODEL_TENSOR.ATTN_ROT_EMBD,
  1919. ],
  1920. MODEL_ARCH.DEEPSEEK: [
  1921. MODEL_TENSOR.ROPE_FREQS,
  1922. MODEL_TENSOR.ATTN_ROT_EMBD,
  1923. ],
  1924. MODEL_ARCH.DEEPSEEK2: [
  1925. MODEL_TENSOR.ROPE_FREQS,
  1926. MODEL_TENSOR.ATTN_ROT_EMBD,
  1927. ],
  1928. MODEL_ARCH.CHATGLM: [
  1929. MODEL_TENSOR.ROPE_FREQS,
  1930. ],
  1931. MODEL_ARCH.NEMOTRON: [
  1932. MODEL_TENSOR.ROPE_FREQS,
  1933. MODEL_TENSOR.ATTN_ROT_EMBD,
  1934. ],
  1935. MODEL_ARCH.BAILINGMOE: [
  1936. MODEL_TENSOR.ROPE_FREQS,
  1937. ],
  1938. }
  1939. #
  1940. # types
  1941. #
  1942. class TokenType(IntEnum):
  1943. NORMAL = 1
  1944. UNKNOWN = 2
  1945. CONTROL = 3
  1946. USER_DEFINED = 4
  1947. UNUSED = 5
  1948. BYTE = 6
  1949. class RopeScalingType(Enum):
  1950. NONE = 'none'
  1951. LINEAR = 'linear'
  1952. YARN = 'yarn'
  1953. LONGROPE = 'longrope'
  1954. class PoolingType(IntEnum):
  1955. NONE = 0
  1956. MEAN = 1
  1957. CLS = 2
  1958. class GGMLQuantizationType(IntEnum):
  1959. F32 = 0
  1960. F16 = 1
  1961. Q4_0 = 2
  1962. Q4_1 = 3
  1963. Q5_0 = 6
  1964. Q5_1 = 7
  1965. Q8_0 = 8
  1966. Q8_1 = 9
  1967. Q2_K = 10
  1968. Q3_K = 11
  1969. Q4_K = 12
  1970. Q5_K = 13
  1971. Q6_K = 14
  1972. Q8_K = 15
  1973. IQ2_XXS = 16
  1974. IQ2_XS = 17
  1975. IQ3_XXS = 18
  1976. IQ1_S = 19
  1977. IQ4_NL = 20
  1978. IQ3_S = 21
  1979. IQ2_S = 22
  1980. IQ4_XS = 23
  1981. I8 = 24
  1982. I16 = 25
  1983. I32 = 26
  1984. I64 = 27
  1985. F64 = 28
  1986. IQ1_M = 29
  1987. BF16 = 30
  1988. TQ1_0 = 34
  1989. TQ2_0 = 35
  1990. class ExpertGatingFuncType(IntEnum):
  1991. SOFTMAX = 1
  1992. SIGMOID = 2
  1993. # TODO: add GGMLFileType from ggml_ftype in ggml.h
  1994. # from llama_ftype in llama.h
  1995. # ALL VALUES SHOULD BE THE SAME HERE AS THEY ARE OVER THERE.
  1996. class LlamaFileType(IntEnum):
  1997. ALL_F32 = 0
  1998. MOSTLY_F16 = 1 # except 1d tensors
  1999. MOSTLY_Q4_0 = 2 # except 1d tensors
  2000. MOSTLY_Q4_1 = 3 # except 1d tensors
  2001. # MOSTLY_Q4_1_SOME_F16 = 4 # tok_embeddings.weight and output.weight are F16
  2002. # MOSTLY_Q4_2 = 5 # support has been removed
  2003. # MOSTLY_Q4_3 = 6 # support has been removed
  2004. MOSTLY_Q8_0 = 7 # except 1d tensors
  2005. MOSTLY_Q5_0 = 8 # except 1d tensors
  2006. MOSTLY_Q5_1 = 9 # except 1d tensors
  2007. MOSTLY_Q2_K = 10 # except 1d tensors
  2008. MOSTLY_Q3_K_S = 11 # except 1d tensors
  2009. MOSTLY_Q3_K_M = 12 # except 1d tensors
  2010. MOSTLY_Q3_K_L = 13 # except 1d tensors
  2011. MOSTLY_Q4_K_S = 14 # except 1d tensors
  2012. MOSTLY_Q4_K_M = 15 # except 1d tensors
  2013. MOSTLY_Q5_K_S = 16 # except 1d tensors
  2014. MOSTLY_Q5_K_M = 17 # except 1d tensors
  2015. MOSTLY_Q6_K = 18 # except 1d tensors
  2016. MOSTLY_IQ2_XXS = 19 # except 1d tensors
  2017. MOSTLY_IQ2_XS = 20 # except 1d tensors
  2018. MOSTLY_Q2_K_S = 21 # except 1d tensors
  2019. MOSTLY_IQ3_XS = 22 # except 1d tensors
  2020. MOSTLY_IQ3_XXS = 23 # except 1d tensors
  2021. MOSTLY_IQ1_S = 24 # except 1d tensors
  2022. MOSTLY_IQ4_NL = 25 # except 1d tensors
  2023. MOSTLY_IQ3_S = 26 # except 1d tensors
  2024. MOSTLY_IQ3_M = 27 # except 1d tensors
  2025. MOSTLY_IQ2_S = 28 # except 1d tensors
  2026. MOSTLY_IQ2_M = 29 # except 1d tensors
  2027. MOSTLY_IQ4_XS = 30 # except 1d tensors
  2028. MOSTLY_IQ1_M = 31 # except 1d tensors
  2029. MOSTLY_BF16 = 32 # except 1d tensors
  2030. # MOSTLY_Q4_0_4_4 = 33 # removed from gguf files, use Q4_0 and runtime repack
  2031. # MOSTLY_Q4_0_4_8 = 34 # removed from gguf files, use Q4_0 and runtime repack
  2032. # MOSTLY_Q4_0_8_8 = 35 # removed from gguf files, use Q4_0 and runtime repack
  2033. MOSTLY_TQ1_0 = 36 # except 1d tensors
  2034. MOSTLY_TQ2_0 = 37 # except 1d tensors
  2035. GUESSED = 1024 # not specified in the model file
  2036. class GGUFEndian(IntEnum):
  2037. LITTLE = 0
  2038. BIG = 1
  2039. class GGUFValueType(IntEnum):
  2040. UINT8 = 0
  2041. INT8 = 1
  2042. UINT16 = 2
  2043. INT16 = 3
  2044. UINT32 = 4
  2045. INT32 = 5
  2046. FLOAT32 = 6
  2047. BOOL = 7
  2048. STRING = 8
  2049. ARRAY = 9
  2050. UINT64 = 10
  2051. INT64 = 11
  2052. FLOAT64 = 12
  2053. @staticmethod
  2054. def get_type(val: Any) -> GGUFValueType:
  2055. if isinstance(val, (str, bytes, bytearray)):
  2056. return GGUFValueType.STRING
  2057. elif isinstance(val, list):
  2058. return GGUFValueType.ARRAY
  2059. elif isinstance(val, float):
  2060. return GGUFValueType.FLOAT32
  2061. elif isinstance(val, bool):
  2062. return GGUFValueType.BOOL
  2063. elif isinstance(val, int):
  2064. return GGUFValueType.INT32
  2065. # TODO: need help with 64-bit types in Python
  2066. else:
  2067. raise ValueError(f"Unknown type: {type(val)}")
  2068. class VisionProjectorType:
  2069. GEMMA3 = "gemma3"
  2070. IDEFICS3 = "idefics3"
  2071. PIXTRAL = "pixtral"
  2072. # Items here are (block size, type size)
  2073. QK_K = 256
  2074. GGML_QUANT_SIZES: dict[GGMLQuantizationType, tuple[int, int]] = {
  2075. GGMLQuantizationType.F32: (1, 4),
  2076. GGMLQuantizationType.F16: (1, 2),
  2077. GGMLQuantizationType.Q4_0: (32, 2 + 16),
  2078. GGMLQuantizationType.Q4_1: (32, 2 + 2 + 16),
  2079. GGMLQuantizationType.Q5_0: (32, 2 + 4 + 16),
  2080. GGMLQuantizationType.Q5_1: (32, 2 + 2 + 4 + 16),
  2081. GGMLQuantizationType.Q8_0: (32, 2 + 32),
  2082. GGMLQuantizationType.Q8_1: (32, 4 + 4 + 32),
  2083. GGMLQuantizationType.Q2_K: (256, 2 + 2 + QK_K // 16 + QK_K // 4),
  2084. GGMLQuantizationType.Q3_K: (256, 2 + QK_K // 4 + QK_K // 8 + 12),
  2085. GGMLQuantizationType.Q4_K: (256, 2 + 2 + QK_K // 2 + 12),
  2086. GGMLQuantizationType.Q5_K: (256, 2 + 2 + QK_K // 2 + QK_K // 8 + 12),
  2087. GGMLQuantizationType.Q6_K: (256, 2 + QK_K // 2 + QK_K // 4 + QK_K // 16),
  2088. GGMLQuantizationType.Q8_K: (256, 4 + QK_K + QK_K // 8),
  2089. GGMLQuantizationType.IQ2_XXS: (256, 2 + QK_K // 4),
  2090. GGMLQuantizationType.IQ2_XS: (256, 2 + QK_K // 4 + QK_K // 32),
  2091. GGMLQuantizationType.IQ3_XXS: (256, 2 + QK_K // 4 + QK_K // 8),
  2092. GGMLQuantizationType.IQ1_S: (256, 2 + QK_K // 8 + QK_K // 16),
  2093. GGMLQuantizationType.IQ4_NL: (32, 2 + 16),
  2094. GGMLQuantizationType.IQ3_S: (256, 2 + QK_K // 4 + QK_K // 8 + QK_K // 32 + 4),
  2095. GGMLQuantizationType.IQ2_S: (256, 2 + QK_K // 4 + QK_K // 16),
  2096. GGMLQuantizationType.IQ4_XS: (256, 2 + 2 + QK_K // 2 + QK_K // 64),
  2097. GGMLQuantizationType.I8: (1, 1),
  2098. GGMLQuantizationType.I16: (1, 2),
  2099. GGMLQuantizationType.I32: (1, 4),
  2100. GGMLQuantizationType.I64: (1, 8),
  2101. GGMLQuantizationType.F64: (1, 8),
  2102. GGMLQuantizationType.IQ1_M: (256, QK_K // 8 + QK_K // 16 + QK_K // 32),
  2103. GGMLQuantizationType.BF16: (1, 2),
  2104. GGMLQuantizationType.TQ1_0: (256, 2 + 4 * 13),
  2105. GGMLQuantizationType.TQ2_0: (256, 2 + 64),
  2106. }
  2107. # Aliases for backward compatibility.
  2108. # general
  2109. KEY_GENERAL_ARCHITECTURE = Keys.General.ARCHITECTURE
  2110. KEY_GENERAL_QUANTIZATION_VERSION = Keys.General.QUANTIZATION_VERSION
  2111. KEY_GENERAL_ALIGNMENT = Keys.General.ALIGNMENT
  2112. KEY_GENERAL_NAME = Keys.General.NAME
  2113. KEY_GENERAL_AUTHOR = Keys.General.AUTHOR
  2114. KEY_GENERAL_URL = Keys.General.URL
  2115. KEY_GENERAL_DESCRIPTION = Keys.General.DESCRIPTION
  2116. KEY_GENERAL_LICENSE = Keys.General.LICENSE
  2117. KEY_GENERAL_SOURCE_URL = Keys.General.SOURCE_URL
  2118. KEY_GENERAL_FILE_TYPE = Keys.General.FILE_TYPE
  2119. # LLM
  2120. KEY_VOCAB_SIZE = Keys.LLM.VOCAB_SIZE
  2121. KEY_CONTEXT_LENGTH = Keys.LLM.CONTEXT_LENGTH
  2122. KEY_EMBEDDING_LENGTH = Keys.LLM.EMBEDDING_LENGTH
  2123. KEY_BLOCK_COUNT = Keys.LLM.BLOCK_COUNT
  2124. KEY_FEED_FORWARD_LENGTH = Keys.LLM.FEED_FORWARD_LENGTH
  2125. KEY_USE_PARALLEL_RESIDUAL = Keys.LLM.USE_PARALLEL_RESIDUAL
  2126. KEY_TENSOR_DATA_LAYOUT = Keys.LLM.TENSOR_DATA_LAYOUT
  2127. # attention
  2128. KEY_ATTENTION_HEAD_COUNT = Keys.Attention.HEAD_COUNT
  2129. KEY_ATTENTION_HEAD_COUNT_KV = Keys.Attention.HEAD_COUNT_KV
  2130. KEY_ATTENTION_MAX_ALIBI_BIAS = Keys.Attention.MAX_ALIBI_BIAS
  2131. KEY_ATTENTION_CLAMP_KQV = Keys.Attention.CLAMP_KQV
  2132. KEY_ATTENTION_LAYERNORM_EPS = Keys.Attention.LAYERNORM_EPS
  2133. KEY_ATTENTION_LAYERNORM_RMS_EPS = Keys.Attention.LAYERNORM_RMS_EPS
  2134. # RoPE
  2135. KEY_ROPE_DIMENSION_COUNT = Keys.Rope.DIMENSION_COUNT
  2136. KEY_ROPE_FREQ_BASE = Keys.Rope.FREQ_BASE
  2137. KEY_ROPE_SCALING_TYPE = Keys.Rope.SCALING_TYPE
  2138. KEY_ROPE_SCALING_FACTOR = Keys.Rope.SCALING_FACTOR
  2139. KEY_ROPE_SCALING_ORIG_CTX_LEN = Keys.Rope.SCALING_ORIG_CTX_LEN
  2140. KEY_ROPE_SCALING_FINETUNED = Keys.Rope.SCALING_FINETUNED
  2141. # SSM
  2142. KEY_SSM_CONV_KERNEL = Keys.SSM.CONV_KERNEL
  2143. KEY_SSM_INNER_SIZE = Keys.SSM.INNER_SIZE
  2144. KEY_SSM_STATE_SIZE = Keys.SSM.STATE_SIZE
  2145. KEY_SSM_TIME_STEP_RANK = Keys.SSM.TIME_STEP_RANK
  2146. KEY_SSM_DT_B_C_RMS = Keys.SSM.DT_B_C_RMS
  2147. # tokenization
  2148. KEY_TOKENIZER_MODEL = Keys.Tokenizer.MODEL
  2149. KEY_TOKENIZER_PRE = Keys.Tokenizer.PRE
  2150. KEY_TOKENIZER_LIST = Keys.Tokenizer.LIST
  2151. KEY_TOKENIZER_TOKEN_TYPE = Keys.Tokenizer.TOKEN_TYPE
  2152. KEY_TOKENIZER_SCORES = Keys.Tokenizer.SCORES
  2153. KEY_TOKENIZER_MERGES = Keys.Tokenizer.MERGES
  2154. KEY_TOKENIZER_BOS_ID = Keys.Tokenizer.BOS_ID
  2155. KEY_TOKENIZER_EOS_ID = Keys.Tokenizer.EOS_ID
  2156. KEY_TOKENIZER_EOT_ID = Keys.Tokenizer.EOT_ID
  2157. KEY_TOKENIZER_EOM_ID = Keys.Tokenizer.EOM_ID
  2158. KEY_TOKENIZER_UNK_ID = Keys.Tokenizer.UNK_ID
  2159. KEY_TOKENIZER_SEP_ID = Keys.Tokenizer.SEP_ID
  2160. KEY_TOKENIZER_PAD_ID = Keys.Tokenizer.PAD_ID
  2161. KEY_TOKENIZER_MASK_ID = Keys.Tokenizer.MASK_ID
  2162. KEY_TOKENIZER_HF_JSON = Keys.Tokenizer.HF_JSON
  2163. KEY_TOKENIZER_RWKV = Keys.Tokenizer.RWKV
  2164. KEY_TOKENIZER_FIM_PRE_ID = Keys.Tokenizer.FIM_PRE_ID
  2165. KEY_TOKENIZER_FIM_SUF_ID = Keys.Tokenizer.FIM_SUF_ID
  2166. KEY_TOKENIZER_FIM_MID_ID = Keys.Tokenizer.FIM_MID_ID
  2167. KEY_TOKENIZER_FIM_PAD_ID = Keys.Tokenizer.FIM_PAD_ID
  2168. KEY_TOKENIZER_FIM_REP_ID = Keys.Tokenizer.FIM_REP_ID
  2169. KEY_TOKENIZER_FIM_SEP_ID = Keys.Tokenizer.FIM_SEP_ID
  2170. # deprecated
  2171. KEY_TOKENIZER_PREFIX_ID = Keys.Tokenizer.PREFIX_ID
  2172. KEY_TOKENIZER_SUFFIX_ID = Keys.Tokenizer.SUFFIX_ID
  2173. KEY_TOKENIZER_MIDDLE_ID = Keys.Tokenizer.MIDDLE_ID