#!/usr/bin/env python3 # dummy model generation script based on convert-hf-to-gguf-bitnet.py from __future__ import annotations import sys from pathlib import Path import numpy as np import configparser import logging import argparse import contextlib import json import os import re import sys from abc import ABC, abstractmethod from enum import IntEnum from pathlib import Path from hashlib import sha256 from typing import TYPE_CHECKING, Any, Callable, ContextManager, Iterator, Sequence, TypeVar, cast, Tuple, Iterable # Necessary to load the local gguf package sys.path.insert(0, str(Path(__file__).parent.parent)) from gguf import GGUFWriter, GGUFReader, RopeScalingType, TokenType, GGMLQuantizationType # noqa: E402 if TYPE_CHECKING: from torch import Tensor import torch import gguf logger = logging.getLogger("generate-dummy-bitnet-model") ###### MODEL HPARAMS CONFIGURATION ###### model_config = { "125M": { "hidden_size": 768, "intermediate_size": 3072, "num_hidden_layers": 11, "num_attention_heads": 12 }, "350M": { "hidden_size": 1024, "intermediate_size": 3072, "num_hidden_layers": 24, "num_attention_heads": 16 }, "1B": { "hidden_size": 2048, "intermediate_size": 3584, "num_hidden_layers": 24, "num_attention_heads": 32 }, "1.5B": { "hidden_size": 1536, "intermediate_size": 9216, "num_hidden_layers": 28, "num_attention_heads": 32 }, "2.7B": { "hidden_size": 3072, "intermediate_size": 7680, "num_hidden_layers": 24, "num_attention_heads": 32 }, "3.8B": { "hidden_size": 3840, "intermediate_size": 8192, "num_hidden_layers": 24, "num_attention_heads": 32 }, "7B": { "hidden_size": 4096, "intermediate_size": 12032, "num_hidden_layers": 32, "num_attention_heads": 32 }, "13B": { "hidden_size": 5120, "intermediate_size": 13824, "num_hidden_layers": 40, "num_attention_heads": 40 }, "30B": { "hidden_size": 6656, "intermediate_size": 16384, "num_hidden_layers": 60, "num_attention_heads": 52 }, "70B": { "hidden_size": 8192, "intermediate_size": 24576, "num_hidden_layers": 80, "num_attention_heads": 64 }, "100B": { "hidden_size": 8192, "intermediate_size": 45568, "num_hidden_layers": 72, "num_attention_heads": 64 } } ###### MODEL DEFINITIONS ###### class SentencePieceTokenTypes(IntEnum): NORMAL = 1 UNKNOWN = 2 CONTROL = 3 USER_DEFINED = 4 UNUSED = 5 BYTE = 6 AnyModel = TypeVar("AnyModel", bound="type[Model]") class Model(ABC): _model_classes: dict[str, type[Model]] = {} def __init__(self, dir_model: Path, ftype: int, fname_out: Path, is_big_endian: bool, use_temp_file: bool): self.dir_model = dir_model self.ftype = ftype self.fname_out = fname_out self.is_big_endian = is_big_endian self.endianess = gguf.GGUFEndian.BIG if is_big_endian else gguf.GGUFEndian.LITTLE self.use_temp_file = use_temp_file self.is_safetensors = self._is_model_safetensors() self.num_parts = Model.count_model_parts(self.dir_model, ".safetensors" if self.is_safetensors else ".bin") self.part_names = self._get_part_names() self.hparams = Model.load_hparams(self.dir_model) self.gguf_writer = gguf.GGUFWriter(fname_out, gguf.MODEL_ARCH_NAMES[self.model_arch], endianess=self.endianess, use_temp_file=self.use_temp_file) self.block_count = self.find_hparam(["n_layers", "num_hidden_layers", "n_layer"]) self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) @property @abstractmethod def model_arch(self) -> gguf.MODEL_ARCH: pass def find_hparam(self, keys: Sequence[str], optional: bool = False) -> Any: key = next((k for k in keys if k in self.hparams), None) if key is not None: return self.hparams[key] if optional: return None raise KeyError(f"could not find any of: {keys}") def set_vocab(self): self._set_vocab_gpt2() def get_tensors(self) -> Iterator[tuple[str, Tensor]]: for part_name in self.part_names: logger.info(f"gguf: loading model part '{part_name}'") ctx: ContextManager[Any] if self.is_safetensors: from safetensors import safe_open ctx = cast(ContextManager[Any], safe_open(self.dir_model / part_name, framework="pt", device="cpu")) else: ctx = contextlib.nullcontext(torch.load(str(self.dir_model / part_name), map_location="cpu", mmap=True, weights_only=True)) with ctx as model_part: for name in model_part.keys(): data = model_part.get_tensor(name) if self.is_safetensors else model_part[name] yield name, data def match_model_tensor_name(self, name: str, key: gguf.MODEL_TENSOR, bid: int | None, suffix: str = ".weight") -> bool: if key not in gguf.MODEL_TENSORS[self.model_arch]: return False key_name: str = gguf.TENSOR_NAMES[key] if "{bid}" in key_name: if bid is None: return False key_name = key_name.format(bid=bid) else: if bid is not None: return False return name == (key_name + suffix) def map_tensor_name(self, name: str, try_suffixes: Sequence[str] = (".weight", ".bias")) -> str: new_name = self.tensor_map.get_name(key=name, try_suffixes=try_suffixes) if new_name is None: raise ValueError(f"Can not map tensor {name!r}") return new_name def set_gguf_parameters(self): self.gguf_writer.add_name(self.dir_model.name) self.gguf_writer.add_block_count(self.block_count) if (n_ctx := self.find_hparam(["max_position_embeddings", "n_ctx"], optional=True)) is not None: self.gguf_writer.add_context_length(n_ctx) logger.info(f"gguf: context length = {n_ctx}") n_embd = self.find_hparam(["hidden_size", "n_embd"]) self.gguf_writer.add_embedding_length(n_embd) logger.info(f"gguf: embedding length = {n_embd}") if (n_ff := self.find_hparam(["intermediate_size", "n_inner"], optional=True)) is not None: self.gguf_writer.add_feed_forward_length(n_ff) logger.info(f"gguf: feed forward length = {n_ff}") n_head = self.find_hparam(["num_attention_heads", "n_head"]) self.gguf_writer.add_head_count(n_head) logger.info(f"gguf: head count = {n_head}") if (n_head_kv := self.hparams.get("num_key_value_heads")) is not None: self.gguf_writer.add_head_count_kv(n_head_kv) logger.info(f"gguf: key-value head count = {n_head_kv}") if (rope_theta := self.hparams.get("rope_theta")) is not None: self.gguf_writer.add_rope_freq_base(rope_theta) logger.info(f"gguf: rope theta = {rope_theta}") if (f_rms_eps := self.hparams.get("rms_norm_eps")) is not None: self.gguf_writer.add_layer_norm_rms_eps(f_rms_eps) logger.info(f"gguf: rms norm epsilon = {f_rms_eps}") if (f_norm_eps := self.find_hparam(["layer_norm_eps", "layer_norm_epsilon", "norm_epsilon"], optional=True)) is not None: self.gguf_writer.add_layer_norm_eps(f_norm_eps) logger.info(f"gguf: layer norm epsilon = {f_norm_eps}") if (n_experts := self.hparams.get("num_local_experts")) is not None: self.gguf_writer.add_expert_count(n_experts) logger.info(f"gguf: expert count = {n_experts}") if (n_experts_used := self.hparams.get("num_experts_per_tok")) is not None: self.gguf_writer.add_expert_used_count(n_experts_used) logger.info(f"gguf: experts used count = {n_experts_used}") self.gguf_writer.add_file_type(self.ftype) logger.info(f"gguf: file type = {self.ftype}") def write_tensors(self): block_count = self.hparams.get("n_layers", self.hparams.get("num_hidden_layers", self.hparams.get("n_layer"))) tensor_map = gguf.get_tensor_name_map(self.model_arch, block_count) for name, data_torch in self.get_tensors(): # we don't need these if name.endswith((".attention.masked_bias", ".attention.bias", ".attention.rotary_emb.inv_freq")): continue old_dtype = data_torch.dtype # convert any unsupported data types to float32 if data_torch.dtype not in (torch.float16, torch.float32): data_torch = data_torch.to(torch.float32) data = data_torch.squeeze().numpy() # map tensor names new_name = tensor_map.get_name(name, try_suffixes=(".weight", ".bias")) if new_name is None: raise ValueError(f"Can not map tensor {name!r}") n_dims = len(data.shape) data_dtype = data.dtype # if f32 desired, convert any float16 to float32 if self.ftype == 0 and data_dtype == np.float16: data = data.astype(np.float32) # TODO: Why cant we use these float16 as-is? There should be not reason to store float16 as float32 if self.ftype == 1 and data_dtype == np.float16 and (n_dims == 1 or new_name.endswith("_norm.weight")): data = data.astype(np.float32) # if f16 desired, convert any float32 2-dim weight tensors to float16 if self.ftype == 1 and data_dtype == np.float32 and name.endswith(".weight") and n_dims == 2: data = data.astype(np.float16) logger.info(f"{new_name}, n_dims = {n_dims}, {old_dtype} --> {data.dtype}") self.gguf_writer.add_tensor(new_name, data) def write(self): self.write_tensors() self.gguf_writer.write_header_to_file() self.gguf_writer.write_kv_data_to_file() self.gguf_writer.write_tensors_to_file() self.gguf_writer.close() def write_vocab(self): self.gguf_writer.write_header_to_file() self.gguf_writer.write_kv_data_to_file() self.gguf_writer.close() @staticmethod def count_model_parts(dir_model: Path, prefix: str) -> int: num_parts = 0 for filename in os.listdir(dir_model): if filename.endswith(prefix): num_parts += 1 return num_parts @staticmethod def load_hparams(dir_model): with open(dir_model / "config.json", "r", encoding="utf-8") as f: return json.load(f) @classmethod def register(cls, *names: str) -> Callable[[AnyModel], AnyModel]: assert names def func(modelcls: type[Model]): for name in names: cls._model_classes[name] = modelcls return modelcls return func @classmethod def from_model_architecture(cls, arch): try: return cls._model_classes[arch] except KeyError: raise NotImplementedError(f'Architecture {arch!r} not supported!') from None def _is_model_safetensors(self) -> bool: return Model.count_model_parts(self.dir_model, ".safetensors") > 0 def _get_part_names(self): if self.is_safetensors: if self.num_parts == 1: # there's only one .safetensors file return ("model.safetensors",) return (f"model-{n:05}-of-{self.num_parts:05}.safetensors" for n in range(1, self.num_parts + 1)) if self.num_parts == 1: # there's only one .bin file return ("pytorch_model.bin",) return (f"pytorch_model-{n:05}-of-{self.num_parts:05}.bin" for n in range(1, self.num_parts + 1)) # used for GPT-2 BPE and WordPiece vocabs def get_vocab_base(self) -> tuple[list[str], list[int], str]: tokens: list[str] = [] toktypes: list[int] = [] from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(self.dir_model) vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) assert max(tokenizer.vocab.values()) < vocab_size tokpre = self.get_vocab_base_pre(tokenizer) reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()} added_vocab = tokenizer.get_added_vocab() for i in range(vocab_size): if i not in reverse_vocab: tokens.append(f"[PAD{i}]") toktypes.append(gguf.TokenType.USER_DEFINED) elif reverse_vocab[i] in added_vocab: tokens.append(reverse_vocab[i]) if tokenizer.added_tokens_decoder[i].special: toktypes.append(gguf.TokenType.CONTROL) else: toktypes.append(gguf.TokenType.USER_DEFINED) else: tokens.append(reverse_vocab[i]) toktypes.append(gguf.TokenType.NORMAL) return tokens, toktypes, tokpre # NOTE: this function is generated by convert-hf-to-gguf-update.py # do not modify it manually! # ref: https://github.com/ggerganov/llama.cpp/pull/6920 def get_vocab_base_pre(self, tokenizer) -> str: # encoding this string and hashing the resulting tokens would (hopefully) give us a unique identifier that # is specific for the BPE pre-tokenizer used by the model # we will use this unique identifier to write a "tokenizer.ggml.pre" entry in the GGUF file which we can # use in llama.cpp to implement the same pre-tokenizer chktxt = '\n \n\n \n\n\n \t \t\t \t\n \n \n \n \n🚀 (normal) 😶\u200d🌫️ (multiple emojis concatenated) ✅ 🦙🦙 3 33 333 3333 33333 333333 3333333 33333333 3.3 3..3 3...3 កាន់តែពិសេសអាច😁 ?我想在apple工作1314151天～ ------======= нещо на Български \'\'\'\'\'\'```````""""......!!!!!!?????? I\'ve been \'told he\'s there, \'RE you sure? \'M not sure I\'ll make it, \'D you like some tea? We\'Ve a\'lL' chktok = tokenizer.encode(chktxt) chkhsh = sha256(str(chktok).encode()).hexdigest() logger.debug(f"chktok: {chktok}") logger.debug(f"chkhsh: {chkhsh}") res = None # NOTE: if you get an error here, you need to update the convert-hf-to-gguf-update.py script # or pull the latest version of the model from Huggingface # don't edit the hashes manually! if chkhsh == "0ef9807a4087ebef797fc749390439009c3b9eda9ad1a097abbe738f486c01e5": # ref: https://huggingface.co/meta-llama/Meta-Llama-3-8B res = "llama-bpe" if chkhsh == "049ecf7629871e3041641907f3de7c733e4dbfdc736f57d882ba0b0845599754": # ref: https://huggingface.co/deepseek-ai/deepseek-llm-7b-base res = "deepseek-llm" if chkhsh == "347715f544604f9118bb75ed199f68779f423cabb20db6de6f31b908d04d7821": # ref: https://huggingface.co/deepseek-ai/deepseek-coder-6.7b-base res = "deepseek-coder" if chkhsh == "8aeee3860c56296a157a1fe2fad249ec40aa59b1bb5709f4ade11c4e6fe652ed": # ref: https://huggingface.co/tiiuae/falcon-7b res = "falcon" if chkhsh == "0876d13b50744004aa9aeae05e7b0647eac9d801b5ba4668afc01e709c15e19f": # ref: https://huggingface.co/BAAI/bge-small-en-v1.5 res = "bert-bge" if chkhsh == "b6dc8df998e1cfbdc4eac8243701a65afe638679230920b50d6f17d81c098166": # ref: https://huggingface.co/mosaicml/mpt-7b res = "mpt" if chkhsh == "35d91631860c815f952d711435f48d356ebac988362536bed955d43bfa436e34": # ref: https://huggingface.co/bigcode/starcoder2-3b res = "starcoder" if chkhsh == "3ce83efda5659b07b1ad37ca97ca5797ea4285d9b9ab0dc679e4a720c9da7454": # ref: https://huggingface.co/openai-community/gpt2 res = "gpt-2" if chkhsh == "6221ad2852e85ce96f791f476e0b390cf9b474c9e3d1362f53a24a06dc8220ff": # ref: https://huggingface.co/smallcloudai/Refact-1_6-base res = "refact" if chkhsh == "9c2227e4dd922002fb81bde4fc02b0483ca4f12911410dee2255e4987644e3f8": # ref: https://huggingface.co/CohereForAI/c4ai-command-r-v01 res = "command-r" if res is None: logger.warning("\n") logger.warning("**************************************************************************************") logger.warning("** WARNING: The BPE pre-tokenizer was not recognized!") logger.warning("** There are 2 possible reasons for this:") logger.warning("** - the model has not been added to convert-hf-to-gguf-update.py yet") logger.warning("** - the pre-tokenization config has changed upstream") logger.warning("** Check your model files and convert-hf-to-gguf-update.py and update them accordingly.") logger.warning("** ref: https://github.com/ggerganov/llama.cpp/pull/6920") logger.warning("**") logger.warning(f"** chkhsh: {chkhsh}") logger.warning("**************************************************************************************") logger.warning("\n") raise NotImplementedError("BPE pre-tokenizer was not recognized - update get_vocab_base_pre()") logger.debug(f"tokenizer.ggml.pre: {repr(res)}") logger.debug(f"chkhsh: {chkhsh}") return res def _set_vocab_sentencepiece(self): from sentencepiece import SentencePieceProcessor tokenizer_path = self.dir_model / 'tokenizer.model' tokens: list[bytes] = [] scores: list[float] = [] toktypes: list[int] = [] if not tokenizer_path.is_file(): raise FileNotFoundError(f"File not found: {tokenizer_path}") tokenizer = SentencePieceProcessor(str(tokenizer_path)) vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size()) for token_id in range(tokenizer.vocab_size()): piece = tokenizer.id_to_piece(token_id) text = piece.encode("utf-8") score = tokenizer.get_score(token_id) toktype = SentencePieceTokenTypes.NORMAL if tokenizer.is_unknown(token_id): toktype = SentencePieceTokenTypes.UNKNOWN elif tokenizer.is_control(token_id): toktype = SentencePieceTokenTypes.CONTROL elif tokenizer.is_unused(token_id): toktype = SentencePieceTokenTypes.UNUSED elif tokenizer.is_byte(token_id): toktype = SentencePieceTokenTypes.BYTE tokens.append(text) scores.append(score) toktypes.append(toktype) added_tokens_file = self.dir_model / 'added_tokens.json' if added_tokens_file.is_file(): with open(added_tokens_file, "r", encoding="utf-8") as f: added_tokens_json = json.load(f) for key in added_tokens_json: key = key.encode("utf-8") if key not in tokens: tokens.append(key) scores.append(-1000.0) toktypes.append(SentencePieceTokenTypes.USER_DEFINED) if vocab_size > len(tokens): pad_count = vocab_size - len(tokens) logger.debug(f"Padding vocab with {pad_count} token(s) - [PAD1] through [PAD{pad_count}]") for i in range(1, pad_count + 1): tokens.append(f"[PAD{i}]") scores.append(-1000.0) toktypes.append(SentencePieceTokenTypes.UNUSED) assert len(tokens) == vocab_size self.gguf_writer.add_tokenizer_model("llama") self.gguf_writer.add_tokenizer_pre("default") self.gguf_writer.add_token_list(tokens) self.gguf_writer.add_token_scores(scores) self.gguf_writer.add_token_types(toktypes) special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) special_vocab.add_to_gguf(self.gguf_writer) # TL1 def process_tl1(weight, BM, BY, bm, by, M, K): final_weight = [] # split in row with size of BM (160) outer_BM_weights = np.split(weight, (M // BM), axis=0) for outer_BM_weight in outer_BM_weights: # split in col with size of by (16index * 2 == 32nums) outer_BY_weights = np.split(outer_BM_weight, (K // BY), axis=1) for outer_BY_weight in outer_BY_weights: # split in row with size of bm (32) inner_bm_weights = np.split(outer_BY_weight, (BM // bm), axis=0) for inner_bm_weight in inner_bm_weights: # split in col with size of by (2index * 2 == 4nums) inner_by_weights = np.split(inner_bm_weight, (BY // by), axis=1) for inner_by_weight in inner_by_weights: # 16 * 6 minor minor_bm_weights = np.split(inner_by_weight, (bm // 16), axis=0) for minor_bm_weight in minor_bm_weights: minor_by_weights = np.split(minor_bm_weight, (by // 4), axis=1) for minor in minor_by_weights: minor_weight = np.split(minor, 2, axis=1) hi_weight = minor_weight[0].astype(np.uint8) << 4 lo_weight = minor_weight[1].astype(np.uint8) func_weight = lo_weight + hi_weight final_weight.append(func_weight) weight = np.array(final_weight, dtype=np.uint8) return weight # based on t_mac.utils.preprocess_weights def preprocess_weights_tl1( w: np.ndarray, bits = 2, g = 4, ) -> Tuple[np.ndarray, np.ndarray]: M, K = w.shape weight = w weight = np.where(np.abs(weight) < 1e-6, 0, weight).astype(np.float32) weight = np.sign(weight) weight_num = np.prod(weight.shape) model_size = args.model_size KEMD = model_config[model_size]['hidden_size'] # outer loop BMEMD = 256 BYEMD = 256 # inner loop (32row 32num/16index) bmEMD = 32 byEMD = 8 KGQA = model_config[model_size]['intermediate_size'] BMGQA = 256 BYGQA = 256 bmGQA = 32 byGQA = 8 weight = np.reshape(weight, (weight_num // 2, 2)) hi_weight = np.multiply(np.split(weight, 2, axis=1)[0], 3) lo_weight = np.split(weight, 2, axis=1)[1] weight = np.reshape((hi_weight + lo_weight), weight_num // 2) # row-major index weight = weight + 4 weight = np.reshape(weight, (M, K // 2)).astype(np.uint8) if K == KEMD: weight = process_tl1(weight, BMEMD, BYEMD, bmEMD, byEMD, M, K) elif K == KGQA: weight = process_tl1(weight, BMGQA, BYGQA, bmGQA, byGQA, M, K) else: raise NotImplementedError return weight def preprocess_two_weights_tl2(M, K, weight_num, BM, BY, bm, by, weight, final_weight): weight = np.reshape(weight, (weight_num // 2, 2)) hi_weight = np.multiply(np.split(weight, 2, axis=1)[0], 3) lo_weight = np.split(weight, 2, axis=1)[1] weight = np.reshape((hi_weight + lo_weight), weight_num // 2) # row-major index weight = weight + 4 weight = np.reshape(weight, (M, K // 2)).astype(np.uint8) outer_BM_weights = np.split(weight, (M // BM), axis=0) for outer_BM_weight in outer_BM_weights: # split in col with size of by (32index * 3 == 96nums) outer_BY_weights = np.split(outer_BM_weight, (K // BY), axis=1) for outer_BY_weight in outer_BY_weights: # split in row with size of bm (32) inner_bm_weights = np.split(outer_BY_weight, (BM // bm), axis=0) for inner_bm_weight in inner_bm_weights: # split in col with size of by (2index * 2 == 4nums) inner_by_weights = np.split(inner_bm_weight, (BY // by), axis=1) for inner_by_weight in inner_by_weights: func_weights = np.split(inner_by_weight, 2, axis=1) left_weight = func_weights[0] left_sub_weights = np.split(left_weight, 4, axis=0) new_left_weight = np.reshape( np.concatenate([left_sub_weights[0], left_sub_weights[2], left_sub_weights[1], left_sub_weights[3]], axis=0, dtype=np.uint8), (bm)) right_weight = func_weights[1] right_sub_weights = np.split(right_weight, 4, axis=0) new_right_weight = np.reshape( np.concatenate([right_sub_weights[0], right_sub_weights[2], right_sub_weights[1], right_sub_weights[3]], axis=0, dtype=np.uint8), (bm)) hi_weight = new_left_weight.astype(np.uint8) << 4 lo_weight = new_right_weight func_weight = hi_weight + lo_weight func_weight = np.reshape(func_weight, bm * by // 4) final_weight.append(func_weight) def preprocess_three_weights_tl2(M, K, weight_num, BM, BY, bm, by, weight, final_weight): weight = np.reshape(weight, (weight_num // 3, 3)) split_weights = np.split(weight, 3, axis=1) first_weight = np.multiply(split_weights[0], 9) second_weight = np.multiply(split_weights[1], 3) third_weight = split_weights[2] weight = np.reshape((first_weight + second_weight + third_weight), weight_num // 3) sign_weight = np.sign(weight) + 2 sign_weight = np.where(sign_weight > 1, 0, sign_weight) weight = np.abs(weight) # row-major index weight = np.reshape(weight, (M, K // 3)).astype(np.uint8) sign_weight = np.reshape(sign_weight, (M, K // 3)).astype(np.uint8) # print(weight) # split in row with size of BM (160) outer_BM_weights = np.split(weight, (M // BM), axis=0) for outer_BM_weight in outer_BM_weights: # split in col with size of by (32index * 3 == 96nums) outer_BY_weights = np.split(outer_BM_weight, (K // BY), axis=1) for outer_BY_weight in outer_BY_weights: # split in row with size of bm (32) inner_bm_weights = np.split(outer_BY_weight, (BM // bm), axis=0) for inner_bm_weight in inner_bm_weights: # split in col with size of by (2index * 3 == 6nums) inner_by_weights = np.split(inner_bm_weight, (BY // by), axis=1) for inner_by_weight in inner_by_weights: func_weights = np.split(inner_by_weight, 2, axis=1) left_weight = func_weights[0] left_sub_weights = np.split(left_weight, 4, axis=0) new_left_weight = np.reshape( np.concatenate([left_sub_weights[0], left_sub_weights[2], left_sub_weights[1], left_sub_weights[3]], axis=0, dtype=np.uint8), (bm)) right_weight = func_weights[1] right_sub_weights = np.split(right_weight, 4, axis=0) new_right_weight = np.reshape( np.concatenate([right_sub_weights[0], right_sub_weights[2], right_sub_weights[1], right_sub_weights[3]], axis=0, dtype=np.uint8), (bm)) hi_weight = new_left_weight.astype(np.uint8) << 4 lo_weight = new_right_weight func_weight = hi_weight + lo_weight func_weight = np.reshape(func_weight, bm * by // 6) final_weight.append(func_weight) sign_weight_list = [] sign_outer_BM_weights = np.split(sign_weight, (M // BM), axis=0) for sign_outer_BM_weight in sign_outer_BM_weights: # split in col with size of by (32index * 3 == 96nums) sign_outer_BY_weights = np.split(sign_outer_BM_weight, (K // BY), axis=1) for sign_outer_BY_weight in sign_outer_BY_weights: # split in row with size of bm (32) sign_inner_bm_weights = np.split(sign_outer_BY_weight, (BM // bm), axis=0) for sign_inner_bm_weight in sign_inner_bm_weights: # split in col with size of by (4index * 3 == 12nums) sign_inner_by_weights = np.split(sign_inner_bm_weight, (BY // (by * 4)), axis=1) for sign_inner_by_weight in sign_inner_by_weights: func_weight = np.split(sign_inner_by_weight, 8, axis=1) combine_weight = np.zeros((16, 1), dtype=np.uint16) for i in range(len(func_weight)): min_weight = np.split(func_weight[i], 2) min_top_weight = min_weight[0].astype(np.uint16) << 15 - (2 * i) min_bot_weight = min_weight[1].astype(np.uint16) << 15 - (2 * i + 1) combine_weight += min_top_weight combine_weight += min_bot_weight combine_weight = combine_weight.view(np.uint8) # combine_weight = combine_weight[:, [1, 0]] combine_weight = np.reshape(combine_weight, bm) sign_weight_list.append(combine_weight) final_weight.extend(sign_weight_list) final_weight.extend(sign_weight_list) def preprocess_weights_tl2( w: np.ndarray, bits = 2, g = 4, ) -> Tuple[np.ndarray, np.ndarray]: M, K = w.shape weight = w weight = np.where(np.abs(weight) < 1e-6, 0, weight).astype(np.float32) weight = np.sign(weight) weight_num = np.prod(weight.shape) # for three num 6 bit -> # outer loop KEMD = 1536 BMEMD = 256 BYEMD = 96 KGQA = 4096 BMGQA = 128 BYGQA = 96 # inner loop (32row 32num/16index) bm3 = 32 by3 = 6 if K == KEMD: BM3 = BMEMD BY3 = BYEMD elif K == KGQA: BM3 = BMGQA BY3 = BYGQA else: raise NotImplementedError BM2 = BM3 BY2 = 32 # inner loop (32row 32num/16index) bm2 = 32 by2 = 4 if (weight.shape[1] % BY3 != 0): slice_k_idx = weight.shape[1] - weight.shape[1] % BY3 slice_weights = np.split(weight, [slice_k_idx], axis=1) three_weight = slice_weights[0] two_weight = slice_weights[1] else: three_weight = weight final_weight = [] preprocess_three_weights_tl2(three_weight.shape[0], three_weight.shape[1], three_weight.shape[0] * three_weight.shape[1], BM3, BY3, bm3, by3, three_weight, final_weight) if (weight.shape[1] % BY3 != 0): preprocess_two_weights_tl2( two_weight.shape[0], two_weight.shape[1], two_weight.shape[0] * two_weight.shape[1], BM2, BY2, bm2, by2, two_weight, final_weight) weight = np.array(final_weight, dtype=np.uint8) return weight @Model.register("BitnetForCausalLM") class BitnetModel(Model): model_arch = gguf.MODEL_ARCH.BITNET params: str = "" def set_params(self, params: str): self.params = params hp_config = model_config[self.params] self.hparams["hidden_size"] = hp_config["hidden_size"] self.hparams["intermediate_size"] = hp_config["intermediate_size"] self.hparams["num_hidden_layers"] = hp_config["num_hidden_layers"] self.hparams["num_attention_heads"] = hp_config["num_attention_heads"] self.hparams["num_key_value_heads"] = hp_config["num_attention_heads"] self.block_count = self.find_hparam(["n_layers", "num_hidden_layers", "n_layer"]) self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) def set_vocab(self): self._set_vocab_sentencepiece() def set_gguf_parameters(self): super().set_gguf_parameters() self.gguf_writer.add_vocab_size(self.hparams["vocab_size"]) self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) self.gguf_writer.add_rope_scaling_factor(1.0) def weight_quant(self, weight): dtype = weight.dtype weight = weight.float() s = 1 / weight.abs().mean().clamp(min=1e-5) result = (weight * s).round().clamp(-1, 1) / s return result.type(dtype) def transform_to_tl1(self, x: np.ndarray): scale = np.max(np.abs(x)) # res = np.round(x / scale + 2).astype(np.uint8) res = preprocess_weights_tl1(x) return res, scale def transform_to_tl2(self, x: np.ndarray): scale = np.max(np.abs(x)) # res = np.round(x / scale + 2).astype(np.uint8) res = preprocess_weights_tl2(x) return res, scale # generate dummy model def generate_tensors(self) -> Iterator[tuple[str, np.ndarray]]: hp_config = model_config[self.params] hidden_size = hp_config["hidden_size"] intermediate_size = hp_config["intermediate_size"] num_hidden_layers = hp_config["num_hidden_layers"] num_attention_heads = hp_config["num_attention_heads"] # generate dummy tensors tensor = torch.randn((32002, hidden_size), dtype=torch.float32) yield ("model.embed_tokens.weight", tensor) for i in range(num_hidden_layers): yield f"model.layers.{i}.input_layernorm.weight", torch.randn((hidden_size,), dtype=torch.float32) yield f"model.layers.{i}.mlp.down_proj.weight", torch.randn((hidden_size, intermediate_size), dtype=torch.float32) yield f"model.layers.{i}.mlp.ffn_layernorm.weight", torch.randn((intermediate_size,), dtype=torch.float32) yield f"model.layers.{i}.mlp.gate_proj.weight", torch.randn((intermediate_size, hidden_size), dtype=torch.float32) yield f"model.layers.{i}.mlp.up_proj.weight", torch.randn((intermediate_size, hidden_size), dtype=torch.float32) yield f"model.layers.{i}.post_attention_layernorm.weight", torch.randn((hidden_size), dtype=torch.float32) yield f"model.layers.{i}.self_attn.inner_attn_ln.weight", torch.randn((hidden_size,), dtype=torch.float32) yield f"model.layers.{i}.self_attn.k_proj.weight", torch.randn((hidden_size, hidden_size), dtype=torch.float32) yield f"model.layers.{i}.self_attn.o_proj.weight", torch.randn((hidden_size, hidden_size), dtype=torch.float32) yield f"model.layers.{i}.self_attn.q_proj.weight", torch.randn((hidden_size, hidden_size), dtype=torch.float32) yield f"model.layers.{i}.self_attn.rotary_emb.inv_freq", torch.randn((hidden_size // (num_attention_heads * 2),), dtype=torch.float32) yield f"model.layers.{i}.self_attn.v_proj.weight", torch.randn((hidden_size, hidden_size), dtype=torch.float32) tensor = torch.randn((hidden_size,), dtype=torch.float32) yield("model.norm.weight", tensor) def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: # quant weight to i2 (in fp16) if name.endswith(("q_proj.weight", "k_proj.weight", "v_proj.weight", "down_proj.weight", "up_proj.weight", "gate_proj.weight", "o_proj.weight")): data_torch = self.weight_quant(data_torch) return [(self.map_tensor_name(name), data_torch)] def write_tensors(self): max_name_len = max(len(s) for _, s in self.tensor_map.mapping.values()) + len(".weight,") for name, data_torch in self.generate_tensors(): # we don't need these if name.endswith((".attention.masked_bias", ".attention.bias", ".rotary_emb.inv_freq")): continue old_dtype = data_torch.dtype # convert any unsupported data types to float32 if data_torch.dtype not in (torch.float16, torch.float32): data_torch = data_torch.to(torch.float32) # use the first number-like part of the tensor name as the block id bid = None for part in name.split("."): if part.isdecimal(): bid = int(part) break for new_name, data in ((n, d.squeeze().numpy()) for n, d in self.modify_tensors(data_torch, name, bid)): data: np.ndarray = data # type hint data_shape = data.shape n_dims = len(data.shape) data_dtype = data.dtype data_qtype: gguf.GGMLQuantizationType | None = None # when both are True, f32 should win # extra_f32 = self.extra_f32_tensors(name, new_name, bid, n_dims) # extra_f16 = self.extra_f16_tensors(name, new_name, bid, n_dims) extra_f32 = False extra_f16 = False # Most of the codebase that takes in 1D tensors or norms only handles F32 tensors # Conditions should closely match those in llama_model_quantize_internal in llama.cpp extra_f32 = any(cond for cond in ( extra_f32, n_dims == 1, new_name.endswith("_norm.weight"), )) # Some tensor types are always in float32 extra_f32 = extra_f32 or any(self.match_model_tensor_name(new_name, key, bid) for key in ( gguf.MODEL_TENSOR.FFN_GATE_INP, gguf.MODEL_TENSOR.POS_EMBD, gguf.MODEL_TENSOR.TOKEN_TYPES, # for debug / delete when inference gguf.MODEL_TENSOR.TOKEN_EMBD, )) # if f16 desired, convert any float32 2-dim weight tensors to float16 extra_f16 = any(cond for cond in ( extra_f16, (name.endswith(".weight") and n_dims >= 2), )) suit_i2 = True if name.endswith('embed_tokens.weight') or name.endswith('norm.weight'): suit_i2 = False i2_scale = None if self.ftype != gguf.GGMLQuantizationType.F32 and extra_f16 and not extra_f32: if self.ftype == gguf.GGMLQuantizationType.TL1 and suit_i2: data, i2_scale = self.transform_to_tl1(data) assert data.dtype == np.uint8 assert i2_scale.dtype == np.float32 data_qtype = gguf.GGMLQuantizationType.TL1 elif self.ftype == gguf.GGMLQuantizationType.TL2 and suit_i2: data, i2_scale = self.transform_to_tl2(data) assert data.dtype == np.uint8 assert i2_scale.dtype == np.float32 data_qtype = gguf.GGMLQuantizationType.TL2 else: # default to float16 for quantized tensors if data_dtype != np.float16: data = data.astype(np.float16) data_qtype = gguf.GGMLQuantizationType.F16 if data_qtype is None: # by default, convert to float32 if data_dtype != np.float32: data = data.astype(np.float32) data_qtype = gguf.GGMLQuantizationType.F32 shape = data_shape # shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape # reverse shape to make it similar to the internal ggml dimension order shape_str = f"{{{', '.join(str(n) for n in reversed(shape))}}}" # n_dims is implicit in the shape logger.info(f"{f'%-{max_name_len}s' % f'{new_name},'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") self.gguf_writer.add_tensor(new_name, data, raw_shape=shape, raw_dtype=data_qtype) if i2_scale is not None: self.gguf_writer.add_tensor(new_name + "_scale", i2_scale, raw_dtype=gguf.GGMLQuantizationType.F32) ftype_map = { "f32": gguf.GGMLQuantizationType.F32, "f16": gguf.GGMLQuantizationType.F16, "tl1" : gguf.GGMLQuantizationType.TL1, "tl2" : gguf.GGMLQuantizationType.TL2, } def main() -> None: dir_model = args.model fname_out = args.outfile model_size = args.model_size hparams = Model.load_hparams(dir_model) with torch.inference_mode(): model_class = Model.from_model_architecture(hparams["architectures"][0]) model_instance = model_class(dir_model, ftype_map[args.outtype], fname_out, args.bigendian, args.use_temp_file) model_instance.set_params(model_size) logger.info("Set model parameters") model_instance.set_gguf_parameters() logger.info("Set model tokenizer") model_instance.set_vocab() if args.vocab_only: logger.info(f"Exporting model vocab to '{fname_out}'") model_instance.write_vocab() else: logger.info(f"Exporting model to '{fname_out}'") model_instance.write() logger.info(f"Model successfully exported to '{fname_out}'") def read_gguf_file(gguf_file_path): """ Reads and prints key-value pairs and tensor information from a GGUF file in an improved format. Parameters: - gguf_file_path: Path to the GGUF file. """ reader = GGUFReader(gguf_file_path) # List all key-value pairs in a columnized format print("Key-Value Pairs:") # noqa: NP100 max_key_length = max(len(key) for key in reader.fields.keys()) for key, field in reader.fields.items(): value = field.parts[field.data[0]] print(f"{key:{max_key_length}} : {value}") # noqa: NP100 print("----") # noqa: NP100 # List all tensors print("Tensors:") # noqa: NP100 tensor_info_format = "{:<30} | Shape: {:<15} | Size: {:<12} | Quantization: {}" print(tensor_info_format.format("Tensor Name", "Shape", "Size", "Quantization")) # noqa: NP100 print("-" * 80) # noqa: NP100 for tensor in reader.tensors: shape_str = "x".join(map(str, tensor.shape)) size_str = str(tensor.n_elements) quantization_str = tensor.tensor_type.name print(tensor_info_format.format(tensor.name, shape_str, size_str, quantization_str)) # noqa: NP100 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Generate a dummy bitnet model with GGUF format") parser.add_argument( "--vocab-only", action="store_true", help="extract only the vocab", ) parser.add_argument( "--outfile", type=Path, help="path to write to; default: based on input", ) parser.add_argument( "--outtype", type=str, choices=ftype_map.keys(), default="f16", help="output format - use f32 for float32, f16 for float16", ) parser.add_argument("--bigendian", action="store_true", help="model is executed on big endian machine") parser.add_argument( "model", type=Path, help="directory containing model file", ) parser.add_argument("--use-temp-file", action="store_true", help="use the tempfile library while processing (helpful when running out of memory, process killed)") parser.add_argument("--model-name", type=str, default=None, help="name of the model") parser.add_argument("--model-size", type=str, default="7B", help="size of the model") parser.add_argument("--verbose", action="store_true", help="increase output verbosity") return parser.parse_args() if __name__ == '__main__': args = parse_args() main()