constants.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890
  1. # ruff: noqa
  2. from __future__ import annotations
  3. import sys
  4. from enum import Enum, IntEnum, auto
  5. from typing import Any
  6. #
  7. # constants
  8. #
  9. GGUF_MAGIC = 0x46554747 # "GGUF"
  10. GGUF_VERSION = 3
  11. GGUF_DEFAULT_ALIGNMENT = 32
  12. #
  13. # metadata keys
  14. #
  15. class Keys:
  16. class General:
  17. ARCHITECTURE = "general.architecture"
  18. QUANTIZATION_VERSION = "general.quantization_version"
  19. ALIGNMENT = "general.alignment"
  20. NAME = "general.name"
  21. AUTHOR = "general.author"
  22. VERSION = "general.version"
  23. URL = "general.url"
  24. DESCRIPTION = "general.description"
  25. LICENSE = "general.license"
  26. SOURCE_URL = "general.source.url"
  27. SOURCE_HF_REPO = "general.source.huggingface.repository"
  28. FILE_TYPE = "general.file_type"
  29. class LLM:
  30. VOCAB_SIZE = "{arch}.vocab_size"
  31. CONTEXT_LENGTH = "{arch}.context_length"
  32. EMBEDDING_LENGTH = "{arch}.embedding_length"
  33. BLOCK_COUNT = "{arch}.block_count"
  34. FEED_FORWARD_LENGTH = "{arch}.feed_forward_length"
  35. USE_PARALLEL_RESIDUAL = "{arch}.use_parallel_residual"
  36. TENSOR_DATA_LAYOUT = "{arch}.tensor_data_layout"
  37. EXPERT_COUNT = "{arch}.expert_count"
  38. EXPERT_USED_COUNT = "{arch}.expert_used_count"
  39. POOLING_TYPE = "{arch}.pooling_type"
  40. LOGIT_SCALE = "{arch}.logit_scale"
  41. class Attention:
  42. HEAD_COUNT = "{arch}.attention.head_count"
  43. HEAD_COUNT_KV = "{arch}.attention.head_count_kv"
  44. MAX_ALIBI_BIAS = "{arch}.attention.max_alibi_bias"
  45. CLAMP_KQV = "{arch}.attention.clamp_kqv"
  46. KEY_LENGTH = "{arch}.attention.key_length"
  47. VALUE_LENGTH = "{arch}.attention.value_length"
  48. LAYERNORM_EPS = "{arch}.attention.layer_norm_epsilon"
  49. LAYERNORM_RMS_EPS = "{arch}.attention.layer_norm_rms_epsilon"
  50. CAUSAL = "{arch}.attention.causal"
  51. class Rope:
  52. DIMENSION_COUNT = "{arch}.rope.dimension_count"
  53. FREQ_BASE = "{arch}.rope.freq_base"
  54. SCALING_TYPE = "{arch}.rope.scaling.type"
  55. SCALING_FACTOR = "{arch}.rope.scaling.factor"
  56. SCALING_ORIG_CTX_LEN = "{arch}.rope.scaling.original_context_length"
  57. SCALING_FINETUNED = "{arch}.rope.scaling.finetuned"
  58. class SSM:
  59. CONV_KERNEL = "{arch}.ssm.conv_kernel"
  60. INNER_SIZE = "{arch}.ssm.inner_size"
  61. STATE_SIZE = "{arch}.ssm.state_size"
  62. TIME_STEP_RANK = "{arch}.ssm.time_step_rank"
  63. class Tokenizer:
  64. MODEL = "tokenizer.ggml.model"
  65. LIST = "tokenizer.ggml.tokens"
  66. TOKEN_TYPE = "tokenizer.ggml.token_type"
  67. TOKEN_TYPE_COUNT = "tokenizer.ggml.token_type_count" # for BERT-style token types
  68. SCORES = "tokenizer.ggml.scores"
  69. MERGES = "tokenizer.ggml.merges"
  70. BOS_ID = "tokenizer.ggml.bos_token_id"
  71. EOS_ID = "tokenizer.ggml.eos_token_id"
  72. UNK_ID = "tokenizer.ggml.unknown_token_id"
  73. SEP_ID = "tokenizer.ggml.seperator_token_id"
  74. PAD_ID = "tokenizer.ggml.padding_token_id"
  75. CLS_ID = "tokenizer.ggml.cls_token_id"
  76. MASK_ID = "tokenizer.ggml.mask_token_id"
  77. ADD_BOS = "tokenizer.ggml.add_bos_token"
  78. ADD_EOS = "tokenizer.ggml.add_eos_token"
  79. ADD_PREFIX = "tokenizer.ggml.add_space_prefix"
  80. HF_JSON = "tokenizer.huggingface.json"
  81. RWKV = "tokenizer.rwkv.world"
  82. CHAT_TEMPLATE = "tokenizer.chat_template"
  83. #
  84. # recommended mapping of model tensor names for storage in gguf
  85. #
  86. class MODEL_ARCH(IntEnum):
  87. LLAMA = auto()
  88. FALCON = auto()
  89. BAICHUAN = auto()
  90. GROK = auto()
  91. GPT2 = auto()
  92. GPTJ = auto()
  93. GPTNEOX = auto()
  94. MPT = auto()
  95. STARCODER = auto()
  96. PERSIMMON = auto()
  97. REFACT = auto()
  98. BERT = auto()
  99. NOMIC_BERT = auto()
  100. BLOOM = auto()
  101. STABLELM = auto()
  102. QWEN = auto()
  103. QWEN2 = auto()
  104. PHI2 = auto()
  105. PLAMO = auto()
  106. CODESHELL = auto()
  107. ORION = auto()
  108. INTERNLM2 = auto()
  109. MINICPM = auto()
  110. GEMMA = auto()
  111. STARCODER2 = auto()
  112. MAMBA = auto()
  113. XVERSE = auto()
  114. COMMAND_R = auto()
  115. DBRX = auto()
  116. class MODEL_TENSOR(IntEnum):
  117. TOKEN_EMBD = auto()
  118. TOKEN_EMBD_NORM = auto()
  119. TOKEN_TYPES = auto()
  120. POS_EMBD = auto()
  121. OUTPUT = auto()
  122. OUTPUT_NORM = auto()
  123. ROPE_FREQS = auto()
  124. ATTN_Q = auto()
  125. ATTN_K = auto()
  126. ATTN_V = auto()
  127. ATTN_QKV = auto()
  128. ATTN_OUT = auto()
  129. ATTN_NORM = auto()
  130. ATTN_NORM_2 = auto()
  131. ATTN_OUT_NORM = auto()
  132. ATTN_ROT_EMBD = auto()
  133. FFN_GATE_INP = auto()
  134. FFN_NORM = auto()
  135. FFN_GATE = auto()
  136. FFN_DOWN = auto()
  137. FFN_UP = auto()
  138. FFN_ACT = auto()
  139. FFN_GATE_EXP = auto()
  140. FFN_DOWN_EXP = auto()
  141. FFN_UP_EXP = auto()
  142. ATTN_Q_NORM = auto()
  143. ATTN_K_NORM = auto()
  144. LAYER_OUT_NORM = auto()
  145. SSM_IN = auto()
  146. SSM_CONV1D = auto()
  147. SSM_X = auto()
  148. SSM_DT = auto()
  149. SSM_A = auto()
  150. SSM_D = auto()
  151. SSM_OUT = auto()
  152. MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
  153. MODEL_ARCH.LLAMA: "llama",
  154. MODEL_ARCH.FALCON: "falcon",
  155. MODEL_ARCH.BAICHUAN: "baichuan",
  156. MODEL_ARCH.GROK: "grok",
  157. MODEL_ARCH.GPT2: "gpt2",
  158. MODEL_ARCH.GPTJ: "gptj",
  159. MODEL_ARCH.GPTNEOX: "gptneox",
  160. MODEL_ARCH.MPT: "mpt",
  161. MODEL_ARCH.STARCODER: "starcoder",
  162. MODEL_ARCH.PERSIMMON: "persimmon",
  163. MODEL_ARCH.REFACT: "refact",
  164. MODEL_ARCH.BERT: "bert",
  165. MODEL_ARCH.NOMIC_BERT: "nomic-bert",
  166. MODEL_ARCH.BLOOM: "bloom",
  167. MODEL_ARCH.STABLELM: "stablelm",
  168. MODEL_ARCH.QWEN: "qwen",
  169. MODEL_ARCH.QWEN2: "qwen2",
  170. MODEL_ARCH.PHI2: "phi2",
  171. MODEL_ARCH.PLAMO: "plamo",
  172. MODEL_ARCH.CODESHELL: "codeshell",
  173. MODEL_ARCH.ORION: "orion",
  174. MODEL_ARCH.INTERNLM2: "internlm2",
  175. MODEL_ARCH.MINICPM: "minicpm",
  176. MODEL_ARCH.GEMMA: "gemma",
  177. MODEL_ARCH.STARCODER2: "starcoder2",
  178. MODEL_ARCH.MAMBA: "mamba",
  179. MODEL_ARCH.XVERSE: "xverse",
  180. MODEL_ARCH.COMMAND_R: "command-r",
  181. MODEL_ARCH.DBRX: "dbrx",
  182. }
  183. TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
  184. MODEL_TENSOR.TOKEN_EMBD: "token_embd",
  185. MODEL_TENSOR.TOKEN_EMBD_NORM: "token_embd_norm",
  186. MODEL_TENSOR.TOKEN_TYPES: "token_types",
  187. MODEL_TENSOR.POS_EMBD: "position_embd",
  188. MODEL_TENSOR.OUTPUT_NORM: "output_norm",
  189. MODEL_TENSOR.OUTPUT: "output",
  190. MODEL_TENSOR.ROPE_FREQS: "rope_freqs",
  191. MODEL_TENSOR.ATTN_NORM: "blk.{bid}.attn_norm",
  192. MODEL_TENSOR.ATTN_NORM_2: "blk.{bid}.attn_norm_2",
  193. MODEL_TENSOR.ATTN_QKV: "blk.{bid}.attn_qkv",
  194. MODEL_TENSOR.ATTN_Q: "blk.{bid}.attn_q",
  195. MODEL_TENSOR.ATTN_K: "blk.{bid}.attn_k",
  196. MODEL_TENSOR.ATTN_V: "blk.{bid}.attn_v",
  197. MODEL_TENSOR.ATTN_OUT: "blk.{bid}.attn_output",
  198. MODEL_TENSOR.ATTN_ROT_EMBD: "blk.{bid}.attn_rot_embd",
  199. MODEL_TENSOR.ATTN_Q_NORM: "blk.{bid}.attn_q_norm",
  200. MODEL_TENSOR.ATTN_K_NORM: "blk.{bid}.attn_k_norm",
  201. MODEL_TENSOR.ATTN_OUT_NORM: "blk.{bid}.attn_output_norm",
  202. MODEL_TENSOR.FFN_GATE_INP: "blk.{bid}.ffn_gate_inp",
  203. MODEL_TENSOR.FFN_NORM: "blk.{bid}.ffn_norm",
  204. MODEL_TENSOR.FFN_GATE: "blk.{bid}.ffn_gate",
  205. MODEL_TENSOR.FFN_DOWN: "blk.{bid}.ffn_down",
  206. MODEL_TENSOR.FFN_UP: "blk.{bid}.ffn_up",
  207. MODEL_TENSOR.FFN_ACT: "blk.{bid}.ffn",
  208. MODEL_TENSOR.FFN_GATE_EXP: "blk.{bid}.ffn_gate_exps",
  209. MODEL_TENSOR.FFN_DOWN_EXP: "blk.{bid}.ffn_down_exps",
  210. MODEL_TENSOR.FFN_UP_EXP: "blk.{bid}.ffn_up_exps",
  211. MODEL_TENSOR.LAYER_OUT_NORM: "blk.{bid}.layer_output_norm",
  212. MODEL_TENSOR.SSM_IN: "blk.{bid}.ssm_in",
  213. MODEL_TENSOR.SSM_CONV1D: "blk.{bid}.ssm_conv1d",
  214. MODEL_TENSOR.SSM_X: "blk.{bid}.ssm_x",
  215. MODEL_TENSOR.SSM_DT: "blk.{bid}.ssm_dt",
  216. MODEL_TENSOR.SSM_A: "blk.{bid}.ssm_a",
  217. MODEL_TENSOR.SSM_D: "blk.{bid}.ssm_d",
  218. MODEL_TENSOR.SSM_OUT: "blk.{bid}.ssm_out",
  219. }
  220. MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
  221. MODEL_ARCH.LLAMA: [
  222. MODEL_TENSOR.TOKEN_EMBD,
  223. MODEL_TENSOR.OUTPUT_NORM,
  224. MODEL_TENSOR.OUTPUT,
  225. MODEL_TENSOR.ROPE_FREQS,
  226. MODEL_TENSOR.ATTN_NORM,
  227. MODEL_TENSOR.ATTN_Q,
  228. MODEL_TENSOR.ATTN_K,
  229. MODEL_TENSOR.ATTN_V,
  230. MODEL_TENSOR.ATTN_OUT,
  231. MODEL_TENSOR.ATTN_ROT_EMBD,
  232. MODEL_TENSOR.FFN_GATE_INP,
  233. MODEL_TENSOR.FFN_NORM,
  234. MODEL_TENSOR.FFN_GATE,
  235. MODEL_TENSOR.FFN_DOWN,
  236. MODEL_TENSOR.FFN_UP,
  237. MODEL_TENSOR.FFN_GATE_EXP,
  238. MODEL_TENSOR.FFN_DOWN_EXP,
  239. MODEL_TENSOR.FFN_UP_EXP,
  240. ],
  241. MODEL_ARCH.GROK: [
  242. MODEL_TENSOR.TOKEN_EMBD,
  243. MODEL_TENSOR.OUTPUT_NORM,
  244. MODEL_TENSOR.OUTPUT,
  245. MODEL_TENSOR.ROPE_FREQS,
  246. MODEL_TENSOR.ATTN_NORM,
  247. MODEL_TENSOR.ATTN_Q,
  248. MODEL_TENSOR.ATTN_K,
  249. MODEL_TENSOR.ATTN_V,
  250. MODEL_TENSOR.ATTN_OUT,
  251. MODEL_TENSOR.ATTN_ROT_EMBD,
  252. MODEL_TENSOR.ATTN_OUT_NORM,
  253. MODEL_TENSOR.FFN_GATE_INP,
  254. MODEL_TENSOR.FFN_NORM,
  255. MODEL_TENSOR.FFN_GATE,
  256. MODEL_TENSOR.FFN_DOWN,
  257. MODEL_TENSOR.FFN_UP,
  258. MODEL_TENSOR.FFN_GATE_EXP,
  259. MODEL_TENSOR.FFN_DOWN_EXP,
  260. MODEL_TENSOR.FFN_UP_EXP,
  261. MODEL_TENSOR.LAYER_OUT_NORM,
  262. ],
  263. MODEL_ARCH.GPTNEOX: [
  264. MODEL_TENSOR.TOKEN_EMBD,
  265. MODEL_TENSOR.OUTPUT_NORM,
  266. MODEL_TENSOR.OUTPUT,
  267. MODEL_TENSOR.ATTN_NORM,
  268. MODEL_TENSOR.ATTN_QKV,
  269. MODEL_TENSOR.ATTN_OUT,
  270. MODEL_TENSOR.FFN_NORM,
  271. MODEL_TENSOR.FFN_DOWN,
  272. MODEL_TENSOR.FFN_UP,
  273. ],
  274. MODEL_ARCH.FALCON: [
  275. MODEL_TENSOR.TOKEN_EMBD,
  276. MODEL_TENSOR.OUTPUT_NORM,
  277. MODEL_TENSOR.OUTPUT,
  278. MODEL_TENSOR.ATTN_NORM,
  279. MODEL_TENSOR.ATTN_NORM_2,
  280. MODEL_TENSOR.ATTN_QKV,
  281. MODEL_TENSOR.ATTN_OUT,
  282. MODEL_TENSOR.FFN_DOWN,
  283. MODEL_TENSOR.FFN_UP,
  284. ],
  285. MODEL_ARCH.BAICHUAN: [
  286. MODEL_TENSOR.TOKEN_EMBD,
  287. MODEL_TENSOR.OUTPUT_NORM,
  288. MODEL_TENSOR.OUTPUT,
  289. MODEL_TENSOR.ROPE_FREQS,
  290. MODEL_TENSOR.ATTN_NORM,
  291. MODEL_TENSOR.ATTN_Q,
  292. MODEL_TENSOR.ATTN_K,
  293. MODEL_TENSOR.ATTN_V,
  294. MODEL_TENSOR.ATTN_OUT,
  295. MODEL_TENSOR.ATTN_ROT_EMBD,
  296. MODEL_TENSOR.FFN_NORM,
  297. MODEL_TENSOR.FFN_GATE,
  298. MODEL_TENSOR.FFN_DOWN,
  299. MODEL_TENSOR.FFN_UP,
  300. ],
  301. MODEL_ARCH.STARCODER: [
  302. MODEL_TENSOR.TOKEN_EMBD,
  303. MODEL_TENSOR.POS_EMBD,
  304. MODEL_TENSOR.OUTPUT_NORM,
  305. MODEL_TENSOR.OUTPUT,
  306. MODEL_TENSOR.ATTN_NORM,
  307. MODEL_TENSOR.ATTN_QKV,
  308. MODEL_TENSOR.ATTN_OUT,
  309. MODEL_TENSOR.FFN_NORM,
  310. MODEL_TENSOR.FFN_DOWN,
  311. MODEL_TENSOR.FFN_UP,
  312. ],
  313. MODEL_ARCH.BERT: [
  314. MODEL_TENSOR.TOKEN_EMBD,
  315. MODEL_TENSOR.TOKEN_EMBD_NORM,
  316. MODEL_TENSOR.TOKEN_TYPES,
  317. MODEL_TENSOR.POS_EMBD,
  318. MODEL_TENSOR.OUTPUT_NORM,
  319. MODEL_TENSOR.ATTN_OUT_NORM,
  320. MODEL_TENSOR.ATTN_Q,
  321. MODEL_TENSOR.ATTN_K,
  322. MODEL_TENSOR.ATTN_V,
  323. MODEL_TENSOR.ATTN_OUT,
  324. MODEL_TENSOR.FFN_DOWN,
  325. MODEL_TENSOR.FFN_UP,
  326. MODEL_TENSOR.LAYER_OUT_NORM,
  327. ],
  328. MODEL_ARCH.NOMIC_BERT: [
  329. MODEL_TENSOR.TOKEN_EMBD,
  330. MODEL_TENSOR.TOKEN_EMBD_NORM,
  331. MODEL_TENSOR.TOKEN_TYPES,
  332. MODEL_TENSOR.POS_EMBD,
  333. MODEL_TENSOR.OUTPUT_NORM,
  334. MODEL_TENSOR.ATTN_OUT_NORM,
  335. MODEL_TENSOR.ATTN_QKV,
  336. MODEL_TENSOR.ATTN_OUT,
  337. MODEL_TENSOR.FFN_GATE,
  338. MODEL_TENSOR.FFN_DOWN,
  339. MODEL_TENSOR.FFN_UP,
  340. MODEL_TENSOR.LAYER_OUT_NORM,
  341. ],
  342. MODEL_ARCH.MPT: [
  343. MODEL_TENSOR.TOKEN_EMBD,
  344. MODEL_TENSOR.OUTPUT_NORM,
  345. MODEL_TENSOR.OUTPUT,
  346. MODEL_TENSOR.ATTN_NORM,
  347. MODEL_TENSOR.ATTN_QKV,
  348. MODEL_TENSOR.ATTN_OUT,
  349. MODEL_TENSOR.FFN_NORM,
  350. MODEL_TENSOR.FFN_DOWN,
  351. MODEL_TENSOR.FFN_UP,
  352. MODEL_TENSOR.FFN_ACT,
  353. MODEL_TENSOR.ATTN_Q_NORM,
  354. MODEL_TENSOR.ATTN_K_NORM,
  355. MODEL_TENSOR.POS_EMBD,
  356. ],
  357. MODEL_ARCH.GPTJ: [
  358. MODEL_TENSOR.TOKEN_EMBD,
  359. MODEL_TENSOR.OUTPUT_NORM,
  360. MODEL_TENSOR.OUTPUT,
  361. MODEL_TENSOR.ATTN_NORM,
  362. MODEL_TENSOR.ATTN_Q,
  363. MODEL_TENSOR.ATTN_K,
  364. MODEL_TENSOR.ATTN_V,
  365. MODEL_TENSOR.ATTN_OUT,
  366. MODEL_TENSOR.FFN_DOWN,
  367. MODEL_TENSOR.FFN_UP,
  368. ],
  369. MODEL_ARCH.PERSIMMON: [
  370. MODEL_TENSOR.TOKEN_EMBD,
  371. MODEL_TENSOR.OUTPUT,
  372. MODEL_TENSOR.OUTPUT_NORM,
  373. MODEL_TENSOR.ATTN_NORM,
  374. MODEL_TENSOR.ATTN_QKV,
  375. MODEL_TENSOR.ATTN_OUT,
  376. MODEL_TENSOR.FFN_NORM,
  377. MODEL_TENSOR.FFN_DOWN,
  378. MODEL_TENSOR.FFN_UP,
  379. MODEL_TENSOR.ATTN_Q_NORM,
  380. MODEL_TENSOR.ATTN_K_NORM,
  381. MODEL_TENSOR.ATTN_ROT_EMBD,
  382. ],
  383. MODEL_ARCH.REFACT: [
  384. MODEL_TENSOR.TOKEN_EMBD,
  385. MODEL_TENSOR.OUTPUT_NORM,
  386. MODEL_TENSOR.OUTPUT,
  387. MODEL_TENSOR.ATTN_NORM,
  388. MODEL_TENSOR.ATTN_Q,
  389. MODEL_TENSOR.ATTN_K,
  390. MODEL_TENSOR.ATTN_V,
  391. MODEL_TENSOR.ATTN_OUT,
  392. MODEL_TENSOR.FFN_NORM,
  393. MODEL_TENSOR.FFN_GATE,
  394. MODEL_TENSOR.FFN_DOWN,
  395. MODEL_TENSOR.FFN_UP,
  396. ],
  397. MODEL_ARCH.BLOOM: [
  398. MODEL_TENSOR.TOKEN_EMBD,
  399. MODEL_TENSOR.TOKEN_EMBD_NORM,
  400. MODEL_TENSOR.OUTPUT_NORM,
  401. MODEL_TENSOR.OUTPUT,
  402. MODEL_TENSOR.ATTN_NORM,
  403. MODEL_TENSOR.ATTN_QKV,
  404. MODEL_TENSOR.ATTN_OUT,
  405. MODEL_TENSOR.FFN_NORM,
  406. MODEL_TENSOR.FFN_DOWN,
  407. MODEL_TENSOR.FFN_UP,
  408. ],
  409. MODEL_ARCH.STABLELM: [
  410. MODEL_TENSOR.TOKEN_EMBD,
  411. MODEL_TENSOR.OUTPUT_NORM,
  412. MODEL_TENSOR.OUTPUT,
  413. MODEL_TENSOR.ROPE_FREQS,
  414. MODEL_TENSOR.ATTN_NORM,
  415. MODEL_TENSOR.ATTN_Q,
  416. MODEL_TENSOR.ATTN_K,
  417. MODEL_TENSOR.ATTN_V,
  418. MODEL_TENSOR.ATTN_OUT,
  419. MODEL_TENSOR.FFN_NORM,
  420. MODEL_TENSOR.FFN_GATE,
  421. MODEL_TENSOR.FFN_DOWN,
  422. MODEL_TENSOR.FFN_UP,
  423. ],
  424. MODEL_ARCH.QWEN: [
  425. MODEL_TENSOR.TOKEN_EMBD,
  426. MODEL_TENSOR.OUTPUT_NORM,
  427. MODEL_TENSOR.OUTPUT,
  428. MODEL_TENSOR.ROPE_FREQS,
  429. MODEL_TENSOR.ATTN_NORM,
  430. MODEL_TENSOR.ATTN_QKV,
  431. MODEL_TENSOR.ATTN_OUT,
  432. MODEL_TENSOR.ATTN_ROT_EMBD,
  433. MODEL_TENSOR.FFN_NORM,
  434. MODEL_TENSOR.FFN_GATE,
  435. MODEL_TENSOR.FFN_DOWN,
  436. MODEL_TENSOR.FFN_UP,
  437. ],
  438. MODEL_ARCH.QWEN2: [
  439. MODEL_TENSOR.TOKEN_EMBD,
  440. MODEL_TENSOR.OUTPUT_NORM,
  441. MODEL_TENSOR.OUTPUT,
  442. MODEL_TENSOR.ATTN_NORM,
  443. MODEL_TENSOR.ATTN_Q,
  444. MODEL_TENSOR.ATTN_K,
  445. MODEL_TENSOR.ATTN_V,
  446. MODEL_TENSOR.ATTN_OUT,
  447. MODEL_TENSOR.FFN_NORM,
  448. MODEL_TENSOR.FFN_GATE,
  449. MODEL_TENSOR.FFN_DOWN,
  450. MODEL_TENSOR.FFN_UP,
  451. ],
  452. MODEL_ARCH.PLAMO: [
  453. MODEL_TENSOR.TOKEN_EMBD,
  454. MODEL_TENSOR.OUTPUT_NORM,
  455. MODEL_TENSOR.OUTPUT,
  456. MODEL_TENSOR.ROPE_FREQS,
  457. MODEL_TENSOR.ATTN_NORM,
  458. MODEL_TENSOR.ATTN_Q,
  459. MODEL_TENSOR.ATTN_K,
  460. MODEL_TENSOR.ATTN_V,
  461. MODEL_TENSOR.ATTN_OUT,
  462. MODEL_TENSOR.ATTN_ROT_EMBD,
  463. MODEL_TENSOR.FFN_GATE,
  464. MODEL_TENSOR.FFN_DOWN,
  465. MODEL_TENSOR.FFN_UP,
  466. ],
  467. MODEL_ARCH.GPT2: [
  468. MODEL_TENSOR.TOKEN_EMBD,
  469. MODEL_TENSOR.POS_EMBD,
  470. MODEL_TENSOR.OUTPUT_NORM,
  471. MODEL_TENSOR.OUTPUT,
  472. MODEL_TENSOR.ATTN_NORM,
  473. MODEL_TENSOR.ATTN_QKV,
  474. MODEL_TENSOR.ATTN_OUT,
  475. MODEL_TENSOR.FFN_NORM,
  476. MODEL_TENSOR.FFN_DOWN,
  477. MODEL_TENSOR.FFN_UP,
  478. ],
  479. MODEL_ARCH.PHI2: [
  480. MODEL_TENSOR.TOKEN_EMBD,
  481. MODEL_TENSOR.OUTPUT_NORM,
  482. MODEL_TENSOR.OUTPUT,
  483. MODEL_TENSOR.ATTN_NORM,
  484. MODEL_TENSOR.ATTN_QKV,
  485. MODEL_TENSOR.ATTN_Q,
  486. MODEL_TENSOR.ATTN_K,
  487. MODEL_TENSOR.ATTN_V,
  488. MODEL_TENSOR.ATTN_OUT,
  489. MODEL_TENSOR.FFN_NORM,
  490. MODEL_TENSOR.FFN_DOWN,
  491. MODEL_TENSOR.FFN_UP,
  492. ],
  493. MODEL_ARCH.CODESHELL: [
  494. MODEL_TENSOR.TOKEN_EMBD,
  495. MODEL_TENSOR.POS_EMBD,
  496. MODEL_TENSOR.OUTPUT_NORM,
  497. MODEL_TENSOR.OUTPUT,
  498. MODEL_TENSOR.ATTN_NORM,
  499. MODEL_TENSOR.ATTN_QKV,
  500. MODEL_TENSOR.ATTN_OUT,
  501. MODEL_TENSOR.ATTN_ROT_EMBD,
  502. MODEL_TENSOR.FFN_NORM,
  503. MODEL_TENSOR.FFN_DOWN,
  504. MODEL_TENSOR.FFN_UP,
  505. ],
  506. MODEL_ARCH.ORION: [
  507. MODEL_TENSOR.TOKEN_EMBD,
  508. MODEL_TENSOR.OUTPUT_NORM,
  509. MODEL_TENSOR.OUTPUT,
  510. MODEL_TENSOR.ROPE_FREQS,
  511. MODEL_TENSOR.ATTN_NORM,
  512. MODEL_TENSOR.ATTN_Q,
  513. MODEL_TENSOR.ATTN_K,
  514. MODEL_TENSOR.ATTN_V,
  515. MODEL_TENSOR.ATTN_OUT,
  516. MODEL_TENSOR.ATTN_ROT_EMBD,
  517. MODEL_TENSOR.FFN_NORM,
  518. MODEL_TENSOR.FFN_GATE,
  519. MODEL_TENSOR.FFN_DOWN,
  520. MODEL_TENSOR.FFN_UP,
  521. ],
  522. MODEL_ARCH.INTERNLM2: [
  523. MODEL_TENSOR.TOKEN_EMBD,
  524. MODEL_TENSOR.OUTPUT_NORM,
  525. MODEL_TENSOR.OUTPUT,
  526. MODEL_TENSOR.ATTN_NORM,
  527. MODEL_TENSOR.ATTN_Q,
  528. MODEL_TENSOR.ATTN_K,
  529. MODEL_TENSOR.ATTN_V,
  530. MODEL_TENSOR.ATTN_OUT,
  531. MODEL_TENSOR.ATTN_ROT_EMBD,
  532. MODEL_TENSOR.FFN_NORM,
  533. MODEL_TENSOR.FFN_GATE,
  534. MODEL_TENSOR.FFN_DOWN,
  535. MODEL_TENSOR.FFN_UP,
  536. ],
  537. MODEL_ARCH.MINICPM: [
  538. MODEL_TENSOR.TOKEN_EMBD,
  539. MODEL_TENSOR.OUTPUT_NORM,
  540. MODEL_TENSOR.ROPE_FREQS,
  541. MODEL_TENSOR.ATTN_NORM,
  542. MODEL_TENSOR.ATTN_Q,
  543. MODEL_TENSOR.ATTN_K,
  544. MODEL_TENSOR.ATTN_V,
  545. MODEL_TENSOR.ATTN_OUT,
  546. MODEL_TENSOR.ATTN_ROT_EMBD,
  547. MODEL_TENSOR.FFN_GATE_INP,
  548. MODEL_TENSOR.FFN_NORM,
  549. MODEL_TENSOR.FFN_GATE,
  550. MODEL_TENSOR.FFN_DOWN,
  551. MODEL_TENSOR.FFN_UP,
  552. MODEL_TENSOR.FFN_GATE_EXP,
  553. MODEL_TENSOR.FFN_DOWN_EXP,
  554. MODEL_TENSOR.FFN_UP_EXP,
  555. ],
  556. MODEL_ARCH.GEMMA: [
  557. MODEL_TENSOR.TOKEN_EMBD,
  558. MODEL_TENSOR.OUTPUT_NORM,
  559. MODEL_TENSOR.ATTN_NORM,
  560. MODEL_TENSOR.ATTN_Q,
  561. MODEL_TENSOR.ATTN_K,
  562. MODEL_TENSOR.ATTN_V,
  563. MODEL_TENSOR.ATTN_OUT,
  564. MODEL_TENSOR.FFN_GATE,
  565. MODEL_TENSOR.FFN_DOWN,
  566. MODEL_TENSOR.FFN_UP,
  567. MODEL_TENSOR.FFN_NORM,
  568. ],
  569. MODEL_ARCH.STARCODER2: [
  570. MODEL_TENSOR.TOKEN_EMBD,
  571. MODEL_TENSOR.OUTPUT_NORM,
  572. MODEL_TENSOR.OUTPUT,
  573. MODEL_TENSOR.ROPE_FREQS,
  574. MODEL_TENSOR.ATTN_NORM,
  575. MODEL_TENSOR.ATTN_Q,
  576. MODEL_TENSOR.ATTN_K,
  577. MODEL_TENSOR.ATTN_V,
  578. MODEL_TENSOR.ATTN_OUT,
  579. MODEL_TENSOR.ATTN_ROT_EMBD,
  580. MODEL_TENSOR.FFN_NORM,
  581. MODEL_TENSOR.FFN_DOWN,
  582. MODEL_TENSOR.FFN_UP,
  583. ],
  584. MODEL_ARCH.MAMBA: [
  585. MODEL_TENSOR.TOKEN_EMBD,
  586. MODEL_TENSOR.OUTPUT_NORM,
  587. MODEL_TENSOR.OUTPUT,
  588. MODEL_TENSOR.ATTN_NORM,
  589. MODEL_TENSOR.SSM_IN,
  590. MODEL_TENSOR.SSM_CONV1D,
  591. MODEL_TENSOR.SSM_X,
  592. MODEL_TENSOR.SSM_DT,
  593. MODEL_TENSOR.SSM_A,
  594. MODEL_TENSOR.SSM_D,
  595. MODEL_TENSOR.SSM_OUT,
  596. ],
  597. MODEL_ARCH.XVERSE: [
  598. MODEL_TENSOR.TOKEN_EMBD,
  599. MODEL_TENSOR.OUTPUT_NORM,
  600. MODEL_TENSOR.OUTPUT,
  601. MODEL_TENSOR.ROPE_FREQS,
  602. MODEL_TENSOR.ATTN_NORM,
  603. MODEL_TENSOR.ATTN_Q,
  604. MODEL_TENSOR.ATTN_K,
  605. MODEL_TENSOR.ATTN_V,
  606. MODEL_TENSOR.ATTN_OUT,
  607. MODEL_TENSOR.ATTN_ROT_EMBD,
  608. MODEL_TENSOR.FFN_NORM,
  609. MODEL_TENSOR.FFN_GATE,
  610. MODEL_TENSOR.FFN_DOWN,
  611. MODEL_TENSOR.FFN_UP,
  612. ],
  613. MODEL_ARCH.COMMAND_R: [
  614. MODEL_TENSOR.TOKEN_EMBD,
  615. MODEL_TENSOR.OUTPUT_NORM,
  616. MODEL_TENSOR.ATTN_NORM,
  617. MODEL_TENSOR.ATTN_Q,
  618. MODEL_TENSOR.ATTN_K,
  619. MODEL_TENSOR.ATTN_V,
  620. MODEL_TENSOR.ATTN_OUT,
  621. MODEL_TENSOR.FFN_GATE,
  622. MODEL_TENSOR.FFN_DOWN,
  623. MODEL_TENSOR.FFN_UP,
  624. MODEL_TENSOR.ATTN_K_NORM,
  625. MODEL_TENSOR.ATTN_Q_NORM,
  626. ],
  627. MODEL_ARCH.DBRX: [
  628. MODEL_TENSOR.TOKEN_EMBD,
  629. MODEL_TENSOR.OUTPUT_NORM,
  630. MODEL_TENSOR.OUTPUT,
  631. MODEL_TENSOR.ATTN_NORM,
  632. MODEL_TENSOR.ATTN_QKV,
  633. MODEL_TENSOR.ATTN_OUT,
  634. MODEL_TENSOR.ATTN_OUT_NORM,
  635. MODEL_TENSOR.FFN_GATE_INP,
  636. MODEL_TENSOR.FFN_GATE_EXP,
  637. MODEL_TENSOR.FFN_DOWN_EXP,
  638. MODEL_TENSOR.FFN_UP_EXP,
  639. ],
  640. # TODO
  641. }
  642. # tensors that will not be serialized
  643. MODEL_TENSOR_SKIP: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
  644. MODEL_ARCH.LLAMA: [
  645. MODEL_TENSOR.ROPE_FREQS,
  646. MODEL_TENSOR.ATTN_ROT_EMBD,
  647. ],
  648. MODEL_ARCH.BAICHUAN: [
  649. MODEL_TENSOR.ROPE_FREQS,
  650. MODEL_TENSOR.ATTN_ROT_EMBD,
  651. ],
  652. MODEL_ARCH.PERSIMMON: [
  653. MODEL_TENSOR.ROPE_FREQS,
  654. ],
  655. MODEL_ARCH.QWEN: [
  656. MODEL_TENSOR.ROPE_FREQS,
  657. MODEL_TENSOR.ATTN_ROT_EMBD,
  658. ],
  659. MODEL_ARCH.CODESHELL: [
  660. MODEL_TENSOR.ROPE_FREQS,
  661. MODEL_TENSOR.ATTN_ROT_EMBD,
  662. ],
  663. MODEL_ARCH.ORION: [
  664. MODEL_TENSOR.ROPE_FREQS,
  665. MODEL_TENSOR.ATTN_ROT_EMBD,
  666. ],
  667. MODEL_ARCH.STARCODER2: [
  668. MODEL_TENSOR.ROPE_FREQS,
  669. MODEL_TENSOR.ATTN_ROT_EMBD,
  670. ],
  671. MODEL_ARCH.XVERSE: [
  672. MODEL_TENSOR.ROPE_FREQS,
  673. MODEL_TENSOR.ATTN_ROT_EMBD,
  674. ],
  675. }
  676. #
  677. # types
  678. #
  679. class TokenType(IntEnum):
  680. NORMAL = 1
  681. UNKNOWN = 2
  682. CONTROL = 3
  683. USER_DEFINED = 4
  684. UNUSED = 5
  685. BYTE = 6
  686. class RopeScalingType(Enum):
  687. NONE = 'none'
  688. LINEAR = 'linear'
  689. YARN = 'yarn'
  690. class PoolingType(IntEnum):
  691. NONE = 0
  692. MEAN = 1
  693. CLS = 2
  694. class GGMLQuantizationType(IntEnum):
  695. F32 = 0
  696. F16 = 1
  697. Q4_0 = 2
  698. Q4_1 = 3
  699. Q5_0 = 6
  700. Q5_1 = 7
  701. Q8_0 = 8
  702. Q8_1 = 9
  703. Q2_K = 10
  704. Q3_K = 11
  705. Q4_K = 12
  706. Q5_K = 13
  707. Q6_K = 14
  708. Q8_K = 15
  709. IQ2_XXS = 16
  710. IQ2_XS = 17
  711. IQ3_XXS = 18
  712. IQ1_S = 19
  713. IQ4_NL = 20
  714. IQ3_S = 21
  715. IQ2_S = 22
  716. IQ4_XS = 23
  717. I8 = 24
  718. I16 = 25
  719. I32 = 26
  720. I64 = 27
  721. F64 = 28
  722. IQ1_M = 29
  723. class GGUFEndian(IntEnum):
  724. LITTLE = 0
  725. BIG = 1
  726. class GGUFValueType(IntEnum):
  727. UINT8 = 0
  728. INT8 = 1
  729. UINT16 = 2
  730. INT16 = 3
  731. UINT32 = 4
  732. INT32 = 5
  733. FLOAT32 = 6
  734. BOOL = 7
  735. STRING = 8
  736. ARRAY = 9
  737. UINT64 = 10
  738. INT64 = 11
  739. FLOAT64 = 12
  740. @staticmethod
  741. def get_type(val: Any) -> GGUFValueType:
  742. if isinstance(val, (str, bytes, bytearray)):
  743. return GGUFValueType.STRING
  744. elif isinstance(val, list):
  745. return GGUFValueType.ARRAY
  746. elif isinstance(val, float):
  747. return GGUFValueType.FLOAT32
  748. elif isinstance(val, bool):
  749. return GGUFValueType.BOOL
  750. elif isinstance(val, int):
  751. return GGUFValueType.INT32
  752. # TODO: need help with 64-bit types in Python
  753. else:
  754. print("Unknown type:", type(val))
  755. sys.exit()
  756. # Note: Does not support GGML_QKK_64
  757. QK_K = 256
  758. # Items here are (block size, type size)
  759. GGML_QUANT_SIZES = {
  760. GGMLQuantizationType.F32: (1, 4),
  761. GGMLQuantizationType.F16: (1, 2),
  762. GGMLQuantizationType.Q4_0: (32, 2 + 16),
  763. GGMLQuantizationType.Q4_1: (32, 2 + 2 + 16),
  764. GGMLQuantizationType.Q5_0: (32, 2 + 4 + 16),
  765. GGMLQuantizationType.Q5_1: (32, 2 + 2 + 4 + 16),
  766. GGMLQuantizationType.Q8_0: (32, 2 + 32),
  767. GGMLQuantizationType.Q8_1: (32, 4 + 4 + 32),
  768. GGMLQuantizationType.Q2_K: (256, 2 + 2 + QK_K // 16 + QK_K // 4),
  769. GGMLQuantizationType.Q3_K: (256, 2 + QK_K // 4 + QK_K // 8 + 12),
  770. GGMLQuantizationType.Q4_K: (256, 2 + 2 + QK_K // 2 + 12),
  771. GGMLQuantizationType.Q5_K: (256, 2 + 2 + QK_K // 2 + QK_K // 8 + 12),
  772. GGMLQuantizationType.Q6_K: (256, 2 + QK_K // 2 + QK_K // 4 + QK_K // 16),
  773. GGMLQuantizationType.Q8_K: (256, 4 + QK_K + QK_K // 8),
  774. GGMLQuantizationType.IQ2_XXS: (256, 2 + QK_K // 4),
  775. GGMLQuantizationType.IQ2_XS: (256, 2 + QK_K // 4 + QK_K // 32),
  776. GGMLQuantizationType.IQ3_XXS: (256, 2 + QK_K // 4 + QK_K // 8),
  777. GGMLQuantizationType.IQ1_S: (256, 2 + QK_K // 8 + QK_K // 16),
  778. GGMLQuantizationType.IQ4_NL: (32, 2 + 16),
  779. GGMLQuantizationType.IQ3_S:
  780. (256, 2 + QK_K // 4 + QK_K // 8 + QK_K // 32 + 4),
  781. GGMLQuantizationType.IQ2_S: (256, 2 + QK_K // 4 + QK_K // 16),
  782. GGMLQuantizationType.IQ4_XS: (256, 2 + 2 + QK_K // 2 + QK_K // 64),
  783. GGMLQuantizationType.I8: (1, 1),
  784. GGMLQuantizationType.I16: (1, 2),
  785. GGMLQuantizationType.I32: (1, 4),
  786. GGMLQuantizationType.I64: (1, 8),
  787. GGMLQuantizationType.F64: (1, 8),
  788. }
  789. # Aliases for backward compatibility.
  790. # general
  791. KEY_GENERAL_ARCHITECTURE = Keys.General.ARCHITECTURE
  792. KEY_GENERAL_QUANTIZATION_VERSION = Keys.General.QUANTIZATION_VERSION
  793. KEY_GENERAL_ALIGNMENT = Keys.General.ALIGNMENT
  794. KEY_GENERAL_NAME = Keys.General.NAME
  795. KEY_GENERAL_AUTHOR = Keys.General.AUTHOR
  796. KEY_GENERAL_URL = Keys.General.URL
  797. KEY_GENERAL_DESCRIPTION = Keys.General.DESCRIPTION
  798. KEY_GENERAL_LICENSE = Keys.General.LICENSE
  799. KEY_GENERAL_SOURCE_URL = Keys.General.SOURCE_URL
  800. KEY_GENERAL_SOURCE_HF_REPO = Keys.General.SOURCE_HF_REPO
  801. KEY_GENERAL_FILE_TYPE = Keys.General.FILE_TYPE
  802. # LLM
  803. KEY_VOCAB_SIZE = Keys.LLM.VOCAB_SIZE
  804. KEY_CONTEXT_LENGTH = Keys.LLM.CONTEXT_LENGTH
  805. KEY_EMBEDDING_LENGTH = Keys.LLM.EMBEDDING_LENGTH
  806. KEY_BLOCK_COUNT = Keys.LLM.BLOCK_COUNT
  807. KEY_FEED_FORWARD_LENGTH = Keys.LLM.FEED_FORWARD_LENGTH
  808. KEY_USE_PARALLEL_RESIDUAL = Keys.LLM.USE_PARALLEL_RESIDUAL
  809. KEY_TENSOR_DATA_LAYOUT = Keys.LLM.TENSOR_DATA_LAYOUT
  810. # attention
  811. KEY_ATTENTION_HEAD_COUNT = Keys.Attention.HEAD_COUNT
  812. KEY_ATTENTION_HEAD_COUNT_KV = Keys.Attention.HEAD_COUNT_KV
  813. KEY_ATTENTION_MAX_ALIBI_BIAS = Keys.Attention.MAX_ALIBI_BIAS
  814. KEY_ATTENTION_CLAMP_KQV = Keys.Attention.CLAMP_KQV
  815. KEY_ATTENTION_LAYERNORM_EPS = Keys.Attention.LAYERNORM_EPS
  816. KEY_ATTENTION_LAYERNORM_RMS_EPS = Keys.Attention.LAYERNORM_RMS_EPS
  817. # RoPE
  818. KEY_ROPE_DIMENSION_COUNT = Keys.Rope.DIMENSION_COUNT
  819. KEY_ROPE_FREQ_BASE = Keys.Rope.FREQ_BASE
  820. KEY_ROPE_SCALING_TYPE = Keys.Rope.SCALING_TYPE
  821. KEY_ROPE_SCALING_FACTOR = Keys.Rope.SCALING_FACTOR
  822. KEY_ROPE_SCALING_ORIG_CTX_LEN = Keys.Rope.SCALING_ORIG_CTX_LEN
  823. KEY_ROPE_SCALING_FINETUNED = Keys.Rope.SCALING_FINETUNED
  824. # SSM
  825. KEY_SSM_CONV_KERNEL = Keys.SSM.CONV_KERNEL
  826. KEY_SSM_INNER_SIZE = Keys.SSM.INNER_SIZE
  827. KEY_SSM_STATE_SIZE = Keys.SSM.STATE_SIZE
  828. KEY_SSM_TIME_STEP_RANK = Keys.SSM.TIME_STEP_RANK
  829. # tokenization
  830. KEY_TOKENIZER_MODEL = Keys.Tokenizer.MODEL
  831. KEY_TOKENIZER_LIST = Keys.Tokenizer.LIST
  832. KEY_TOKENIZER_TOKEN_TYPE = Keys.Tokenizer.TOKEN_TYPE
  833. KEY_TOKENIZER_SCORES = Keys.Tokenizer.SCORES
  834. KEY_TOKENIZER_MERGES = Keys.Tokenizer.MERGES
  835. KEY_TOKENIZER_BOS_ID = Keys.Tokenizer.BOS_ID
  836. KEY_TOKENIZER_EOS_ID = Keys.Tokenizer.EOS_ID
  837. KEY_TOKENIZER_UNK_ID = Keys.Tokenizer.UNK_ID
  838. KEY_TOKENIZER_SEP_ID = Keys.Tokenizer.SEP_ID
  839. KEY_TOKENIZER_PAD_ID = Keys.Tokenizer.PAD_ID
  840. KEY_TOKENIZER_CLS_ID = Keys.Tokenizer.CLS_ID
  841. KEY_TOKENIZER_MASK_ID = Keys.Tokenizer.MASK_ID
  842. KEY_TOKENIZER_HF_JSON = Keys.Tokenizer.HF_JSON
  843. KEY_TOKENIZER_RWKV = Keys.Tokenizer.RWKV