[LLM Runtime] Support the GGUF format (#40)

neural_speed/application/main_run.cpp

@@ -772,4 +772,4 @@ int main(int argc, char** argv) {  // NOLINT
   model_free(ctx);
 
   return 0;
-}  // NOLINT
+}  // NOLINT

neural_speed/convert/convert_chatglm.py

@@ -17,10 +17,11 @@
 import numpy as np
 from pathlib import Path
 import argparse
-from typing import (IO, TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, TypeVar,
-                    Union)
+from typing import (IO, TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal, Optional, Sequence, Tuple,
+                    TypeVar, Union)
 from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
 from sentencepiece import SentencePieceProcessor  # type: ignore
+import gguf
 
 
 # ref: https://github.com/openai/gpt-2/blob/master/src/encoder.py
@@ -34,7 +35,10 @@ def bytes_to_unicode():
     To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
     And avoids mapping to whitespace/control characters the bpe code barfs on.
     """
-    bs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
+    bs = list(range(ord("!"),
+                    ord("~") + 1)) + list(range(ord("¡"),
+                                                ord("¬") + 1)) + list(range(ord("®"),
+                                                                            ord("ÿ") + 1))
     cs = bs[:]
     n = 0
     for b in range(2**8):
@@ -117,8 +121,7 @@ def load_vocab_for_glm1(path: Path) -> SentencePieceVocab:
         else:
             raise FileNotFoundError(
                 f"Could not find tokenizer.model in {path} or its parent; if it's in another directory, \
-                pass the directory as --vocab-dir"
-            )
+                pass the directory as --vocab-dir")
     added_tokens_path = path.parent / "added_tokens.json"
     print(f"Loading vocab file {path}")
     return SentencePieceVocab(path, added_tokens_path if added_tokens_path.exists() else None)
@@ -139,13 +142,188 @@ def load_vocab_for_glm2(path: Path) -> SentencePieceVocab:
         else:
             raise FileNotFoundError(
                 f"Could not find tokenizer.model in {path} or its parent; if it's in another directory, \
-                pass the directory as --vocab-dir"
-            )
+                pass the directory as --vocab-dir")
     added_tokens_path = path.parent / "added_tokens.json"
     print(f"Loading vocab file {path}")
     return SentencePieceVocab(path, added_tokens_path if added_tokens_path.exists() else None)
 
 
+def chatglm2_convert_gguf(model, tokenizer, dir_model, fname_out, ftype, hparams):
+    print("ChatGLM-2.gguf converting: ")
+    list_vars = model.state_dict()
+    for name in list_vars.keys():
+        print(name, list_vars[name].shape, list_vars[name].dtype)
+
+    print(hparams)
+    fout = open(fname_out, "wb")
+
+    gguf_file = fname_out + '.gguf'
+    gguf_writer = gguf.GGUFWriter(gguf_file, "chatglm2")
+
+    gguf_writer.add_uint32('magic', 0x67676d66)
+    gguf_writer.add_uint32('version', 1)
+    gguf_writer.add_uint32('n_vocab', hparams["padded_vocab_size"])
+    gguf_writer.add_uint32('n_embd', hparams["hidden_size"])
+    gguf_writer.add_uint32('n_mult', 0)
+    gguf_writer.add_uint32('n_head', hparams["num_attention_heads"])
+    gguf_writer.add_uint32('n_head_kv', 0)
+
+    gguf_writer.add_uint32('n_layer', hparams["num_layers"])
+    gguf_writer.add_uint32('n_rot', 0)
+    gguf_writer.add_uint32('ftype', ftype)
+    gguf_writer.add_uint32('max_seq_len', hparams["seq_length"])
+    gguf_writer.add_uint32('alibi_bias_max', 0)
+    gguf_writer.add_uint32('clip_qkv', 0)
+    gguf_writer.add_uint32('par_res', 0)
+
+    gguf_writer.add_uint32('word_embed_proj_dim', 0)
+    gguf_writer.add_uint32('do_layer_norm_before', 0)
+
+    gguf_writer.add_uint32('multi_query_group_num', hparams["multi_query_group_num"])
+    gguf_writer.add_uint32('ffn_hidden_size', hparams["ffn_hidden_size"])
+    gguf_writer.add_uint32('inner_hidden_size', 0)
+
+    gguf_writer.add_int32('bos_token_id', tokenizer.bos_token_id if tokenizer.bos_token_id is not None else -1)
+    gguf_writer.add_int32('eos_token_id', tokenizer.eos_token_id if tokenizer.eos_token_id is not None else -1)
+    gguf_writer.add_int32('pad_token_id', tokenizer.pad_token_id if tokenizer.pad_token_id is not None else -1)
+    gguf_writer.add_int32('sep_token_id', tokenizer.sep_token_id if tokenizer.sep_token_id is not None else -1)
+
+    def write_vocab_gguf(dir_model):
+        print("gguf: get tokenizer metadata")
+
+        tokens: List[bytes] = []
+        scores: List[float] = []
+        toktypes: List[int] = []
+
+        if Path(dir_model + "/tokenizer.model").is_file():
+            # vocab type sentencepiece
+            print("gguf: get sentencepiece tokenizer vocab, scores and token types")
+
+            vocab = load_vocab_for_glm2(Path(dir_model))
+
+            # NOTE: `all_tokens` returns the base vocabulary and added tokens
+            for text, score in vocab.all_tokens():
+                tokens.append(text)
+                scores.append(score)
+
+            if Path(dir_model + "/added_tokens.json").is_file():
+                with open(dir_model + "/added_tokens.json", "r", encoding="utf-8") as f:
+                    addtokens_json = json.load(f)
+
+                    print("gguf: get added tokens")
+
+                    for key in addtokens_json:
+                        tokens.append(key.encode("utf-8"))
+                        scores.append(-1000.0)
+                        toktypes.append(4)  # user-defined token type
+
+            gguf_writer.add_tokenizer_model("chatglm2")
+            gguf_writer.add_token_list(tokens)
+            gguf_writer.add_token_scores(scores)
+
+        print("gguf: get special token ids")
+
+        if Path(dir_model + "/tokenizer.json").is_file():
+            # Look for special tokens in tokenizer.json if it exists
+
+            with open(dir_model + "/tokenizer.json", "r", encoding="utf-8") as f:
+                tokenizer = json.load(f)
+
+            if "added_tokens" in tokenizer and Path(dir_model + "/tokenizer_config.json").is_file():
+
+                with open(dir_model + "/tokenizer_config.json", "r", encoding="utf-8") as f:
+                    tokenizer_config = json.load(f)
+
+                if "bos_token" in tokenizer_config and tokenizer_config["bos_token"] != None:
+                    for key in tokenizer["added_tokens"]:
+                        if key["content"] == tokenizer_config["bos_token"]["content"]:
+                            gguf_writer.add_bos_token_id(key["id"])
+
+                if "eos_token" in tokenizer_config and tokenizer_config["eos_token"] != None:
+                    for key in tokenizer["added_tokens"]:
+                        if key["content"] == tokenizer_config["eos_token"]["content"]:
+                            gguf_writer.add_eos_token_id(key["id"])
+
+                if "unk_token" in tokenizer_config and tokenizer_config["unk_token"] != None:
+                    for key in tokenizer["added_tokens"]:
+                        if key["content"] == tokenizer_config["unk_token"]["content"]:
+                            gguf_writer.add_unk_token_id(key["id"])
+
+                if "sep_token" in tokenizer_config and tokenizer_config["sep_token"] != None:
+                    for key in tokenizer["added_tokens"]:
+                        if key["content"] == tokenizer_config["sep_token"]["content"]:
+                            gguf_writer.add_sep_token_id(key["id"])
+
+                if "pad_token" in tokenizer_config and tokenizer_config["pad_token"] != None:
+                    for key in tokenizer["added_tokens"]:
+                        if key["content"] == tokenizer_config["pad_token"]["content"]:
+                            gguf_writer.add_pad_token_id(key["id"])
+        else:
+            # If no tokenizer.json: Look for special tokens in config.json
+
+            if "bos_token_id" in hparams and hparams["bos_token_id"] != None:
+                gguf_writer.add_bos_token_id(hparams["bos_token_id"])
+
+            if "eos_token_id" in hparams and hparams["eos_token_id"] != None:
+                gguf_writer.add_eos_token_id(hparams["eos_token_id"])
+
+            if "unk_token_id" in hparams and hparams["unk_token_id"] != None:
+                gguf_writer.add_unk_token_id(hparams["unk_token_id"])
+
+            if "sep_token_id" in hparams and hparams["sep_token_id"] != None:
+                gguf_writer.add_sep_token_id(hparams["sep_token_id"])
+
+            if "pad_token_id" in hparams and hparams["pad_token_id"] != None:
+                gguf_writer.add_pad_token_id(hparams["pad_token_id"])
+
+    write_vocab_gguf(dir_model)
+
+    # tensor info
+    print("gguf: get tensor metadata")
+    for name in list_vars.keys():
+        data = list_vars[name].squeeze().numpy()
+
+        print("Processing variable: " + name + " with shape: ", data.shape)
+        if 'inv_freq' in name:
+            continue
+
+        n_dims = len(data.shape)
+
+        # ftype == 0 -> float32, ftype == 1 -> float16
+        ftype_cur = 0
+        if ftype != 0:
+            if name[-7:] == ".weight" and n_dims == 2:
+                print("  Converting to float16")
+                data = data.astype(np.float16)
+                ftype_cur = 1
+            else:
+                print("  Converting to float32")
+                data = data.astype(np.float32)
+                ftype_cur = 0
+        else:
+            if data.dtype != np.float32:
+                print("  Converting to float32")
+                data = data.astype(np.float32)
+                ftype_cur = 0
+
+        # print(f"[{i+1:{padi}d}/{len(model)}]
+        # Writing tensor {name:38s} | size {size:16} | type {lazy_tensor.data_type.name:4}")
+
+        gguf_writer.add_tensor(name, data)
+
+    print("gguf: write header")
+    gguf_writer.write_header_to_file()
+    print("gguf: write metadata")
+    gguf_writer.write_kv_data_to_file()
+    print("gguf: write tensors")
+    gguf_writer.write_tensors_to_file()
+
+    gguf_writer.close()
+
+    print("Done. Output file: " + fname_out)
+    print("")
+
+
 def chatglm2_convert(model, tokenizer, dir_model, fname_out, ftype, hparams):
     print("ChatGLM-2 converting: ")
     list_vars = model.state_dict()
@@ -341,6 +519,11 @@ def main(args_in: Optional[List[str]] = None) -> None:
     parser.add_argument("--outtype", choices=["f32", "f16"], help="output format (default: based on input)")
     parser.add_argument("--outfile", type=Path, help="path to write to; default: based on input")
     parser.add_argument("model", type=Path, help="directory containing model file")
+    parser.add_argument("--format",
+                        type=str,
+                        default="NE",
+                        choices=["NE", "GGUF"],
+                        help="convert to the GGUF or NE format")
     args = parser.parse_args(args_in)
 
     dir_model = args.model.as_posix()
@@ -360,7 +543,10 @@ def main(args_in: Optional[List[str]] = None) -> None:
     model = AutoModel.from_pretrained(dir_model, low_cpu_mem_usage=True, trust_remote_code=True)
 
     if hasattr(model.config, "multi_query_attention"):
-        chatglm2_convert(model, tokenizer, dir_model, fname_out, ftype, hparams)
+        if args.format == "GGUF":
+            chatglm2_convert_gguf(model, tokenizer, dir_model, fname_out, ftype, hparams)
+        else:
+            chatglm2_convert(model, tokenizer, dir_model, fname_out, ftype, hparams)
     else:
         chatglm1_convert(model, tokenizer, dir_model, fname_out, ftype, hparams)