From 9d92821cf54720ac232e3065c5d6b7c6c4c40031 Mon Sep 17 00:00:00 2001
From: Meng Zhang <wsxiaoys.lh@gmail.com>
Date: Mon, 27 Mar 2023 11:12:52 +0800
Subject: [PATCH] Add gptj converter (#19)

* Rename deployment-next to development

* Add GPTJ converter
---
 .dockerignore                                 |   4 -
 Dockerfile                                    |   2 +-
 deployment/scripts/triton.sh                  |   4 +
 {deployment-next => development}/.gitignore   |   0
 {deployment-next => development}/Makefile     |   0
 {deployment-next => development}/README.md    |   0
 .../config/repository.toml                    |   0
 .../config/vector.toml                        |   0
 .../docker-compose.dev.yml                    |   0
 .../docker-compose.triton.yml                 |   0
 .../docker-compose.yml                        |   0
 .../scripts/triton.sh                         |   0
 .../converter/huggingface_gptj_convert.py     | 181 ++++++++++++++++++
 .../huggingface_gptneox_convert.py            |   0
 14 files changed, 186 insertions(+), 5 deletions(-)
 rename {deployment-next => development}/.gitignore (100%)
 rename {deployment-next => development}/Makefile (100%)
 rename {deployment-next => development}/README.md (100%)
 rename {deployment-next => development}/config/repository.toml (100%)
 rename {deployment-next => development}/config/vector.toml (100%)
 rename {deployment-next => development}/docker-compose.dev.yml (100%)
 rename {deployment-next => development}/docker-compose.triton.yml (100%)
 rename {deployment-next => development}/docker-compose.yml (100%)
 rename {deployment-next => development}/scripts/triton.sh (100%)
 create mode 100644 tabby/tools/converter/huggingface_gptj_convert.py
 rename tabby/tools/{ => converter}/huggingface_gptneox_convert.py (100%)

diff --git a/.dockerignore b/.dockerignore
index 3cab3d6..ce6b302 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -1,7 +1,3 @@
-data
-deployment
-deployment-next
-
 **/.git
 **/node_modules
 **/__pycache__
diff --git a/Dockerfile b/Dockerfile
index 460b4f8..dcede3d 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -28,4 +28,4 @@ RUN poetry export --without-hashes -o requirements.txt
 
 RUN --mount=type=cache,target=/root/.cache pip install -i $PYPI_INDEX_URL --extra-index-url https://pypi.org/simple --no-dependencies -r requirements.txt
 
-COPY . .
+COPY tabby ./tabby
diff --git a/deployment/scripts/triton.sh b/deployment/scripts/triton.sh
index c86f3eb..bb0b05c 100755
--- a/deployment/scripts/triton.sh
+++ b/deployment/scripts/triton.sh
@@ -1,6 +1,9 @@
 #!/bin/bash
 set -e
 
+if [ -d "$MODEL_NAME" ]; then
+MODEL_DIR="$MODEL_NAME"
+else
 # Get model dir.
 MODEL_DIR=$(python3 <<EOF
 from huggingface_hub import snapshot_download
@@ -8,6 +11,7 @@ from huggingface_hub import snapshot_download
 print(snapshot_download(repo_id='$MODEL_NAME', allow_patterns='triton/**/*', local_files_only=True))
 EOF
 )
+fi
 
 # Set model dir in triton config.
 sed -i 's@${MODEL_DIR}@'$MODEL_DIR'@g' $MODEL_DIR/triton/fastertransformer/config.pbtxt
diff --git a/deployment-next/.gitignore b/development/.gitignore
similarity index 100%
rename from deployment-next/.gitignore
rename to development/.gitignore
diff --git a/deployment-next/Makefile b/development/Makefile
similarity index 100%
rename from deployment-next/Makefile
rename to development/Makefile
diff --git a/deployment-next/README.md b/development/README.md
similarity index 100%
rename from deployment-next/README.md
rename to development/README.md
diff --git a/deployment-next/config/repository.toml b/development/config/repository.toml
similarity index 100%
rename from deployment-next/config/repository.toml
rename to development/config/repository.toml
diff --git a/deployment-next/config/vector.toml b/development/config/vector.toml
similarity index 100%
rename from deployment-next/config/vector.toml
rename to development/config/vector.toml
diff --git a/deployment-next/docker-compose.dev.yml b/development/docker-compose.dev.yml
similarity index 100%
rename from deployment-next/docker-compose.dev.yml
rename to development/docker-compose.dev.yml
diff --git a/deployment-next/docker-compose.triton.yml b/development/docker-compose.triton.yml
similarity index 100%
rename from deployment-next/docker-compose.triton.yml
rename to development/docker-compose.triton.yml
diff --git a/deployment-next/docker-compose.yml b/development/docker-compose.yml
similarity index 100%
rename from deployment-next/docker-compose.yml
rename to development/docker-compose.yml
diff --git a/deployment-next/scripts/triton.sh b/development/scripts/triton.sh
similarity index 100%
rename from deployment-next/scripts/triton.sh
rename to development/scripts/triton.sh
diff --git a/tabby/tools/converter/huggingface_gptj_convert.py b/tabby/tools/converter/huggingface_gptj_convert.py
new file mode 100644
index 0000000..ddd568e
--- /dev/null
+++ b/tabby/tools/converter/huggingface_gptj_convert.py
@@ -0,0 +1,181 @@
+import configparser
+from argparse import ArgumentParser
+from os import makedirs
+from pathlib import Path
+
+import numpy as np
+import torch
+from transformers import PretrainedConfig
+
+torch.set_printoptions(linewidth=130, sci_mode=False)
+np.set_printoptions(linewidth=130, suppress=True)
+
+# This converter is used to convert the huggingface gpt-j-6B model
+# in https://huggingface.co/EleutherAI/gpt-j-6B/blob/main/pytorch_model.bin.
+
+
+def savebin(param, save_path):
+    if isinstance(param, torch.Tensor):
+        param = param.cpu().float().numpy()
+    np.squeeze(param).astype(np.float16).tofile(save_path + ".bin")
+
+
+def param2file(pt_param, layer_id, save_dir, dest_key):
+    base_n = save_dir + "/model.layers." + str(layer_id) + "."
+    save_path = base_n + dest_key
+    savebin(pt_param, save_path)
+
+
+def param2distributed(
+    pt_param,
+    layer_id,
+    save_dir,
+    dest_key,
+    n_inference_gpus,
+    split_axis,
+):
+    np_param = pt_param.cpu().float().numpy()
+    base_n = save_dir + "/model.layers." + str(layer_id) + "."
+    save_path = base_n + dest_key
+    split_param = np.split(np_param, n_inference_gpus, axis=split_axis)
+    for i, p in enumerate(split_param):
+        savebin(p, save_path + f".{i}")
+
+
+def save(w, save_dir, n_inference_gpus, n_layers, layer_id):
+    makedirs(save_dir, exist_ok=True)
+
+    savebin(w["transformer.wte.weight"], save_dir + "/model.wte")
+    l = layer_id
+    print(f"Saving layer {l + 1} / {n_layers}")
+    base_k = "transformer.h." + str(l) + "."
+    param2file(w[base_k + "ln_1.bias"], l, save_dir, "input_layernorm.bias")
+    param2file(w[base_k + "ln_1.weight"], l, save_dir, "input_layernorm.weight")
+    param2distributed(
+        w[base_k + "mlp.fc_in.weight"].T,
+        l,
+        save_dir,
+        "mlp.dense_h_to_4h.weight",
+        n_inference_gpus,
+        split_axis=-1,  # split fast indx
+    )
+    param2distributed(
+        w[base_k + "mlp.fc_in.bias"],
+        l,
+        save_dir,
+        "mlp.dense_h_to_4h.bias",
+        n_inference_gpus,
+        split_axis=-1,  # split fast indx
+    )
+
+    param2distributed(
+        w[base_k + "mlp.fc_out.weight"].T,
+        l,
+        save_dir,
+        "mlp.dense_4h_to_h.weight",
+        n_inference_gpus,
+        split_axis=0,  # split slow indx
+    )
+    param2file(w[base_k + "mlp.fc_out.bias"], l, save_dir, "mlp.dense_4h_to_h.bias")
+    param2distributed(
+        w[base_k + "attn.out_proj.weight"].T,
+        l,
+        save_dir,
+        "attention.dense.weight",
+        n_inference_gpus,
+        split_axis=0,  # split slow indx
+    )
+    QKV_w = torch.stack(
+        [
+            w[base_k + "attn.q_proj.weight"],
+            w[base_k + "attn.k_proj.weight"],
+            w[base_k + "attn.v_proj.weight"],
+        ]
+    )  # [qkv, n_heads * dim_head, latent_space]
+    QKV_w = QKV_w.permute(2, 0, 1)
+    param2distributed(
+        QKV_w,
+        l,
+        save_dir,
+        "attention.query_key_value.weight",
+        n_inference_gpus,
+        split_axis=-1,  # split fast indx
+    )
+    # Other unneeded per-layer params:
+    # attn.attention.masked_bias = torch.tensor(-1e9)
+    # attn.attention.bias = torch.tril(torch.ones(1, 1, 2048, 2048))
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser(
+        description="Convert GPT-J slim checkpoint to FasterTransformer",
+    )
+    parser.add_argument(
+        "--output-dir",
+        help="Folder where binary files are stored",
+    )
+    parser.add_argument(
+        "--ckpt-dir",
+        help="File of GPT-J huggingface checkpoint",
+    )
+    parser.add_argument(
+        "--n-inference-gpus",
+        help="Number of GPUs used for inference runtime",
+        default=1,
+        type=int,
+    )
+    args = parser.parse_args()
+
+    ckpt_file = args.ckpt_dir + "/pytorch_model.bin"
+    checkpoint = torch.load(ckpt_file)
+    print(f"loading from {ckpt_file}")
+
+    out_path = args.output_dir
+    output_dir = out_path + f"/{args.n_inference_gpus}-gpu/"
+    print(f"saving to {output_dir}")
+
+    config_file = args.ckpt_dir + "/config.json"
+    hf_config = PretrainedConfig.from_json_file(config_file).to_dict()
+
+    # NOTE: save parameters to config files (loaded by triton backends)
+    config = configparser.ConfigParser()
+    config["gptj"] = {}
+    try:
+        config["gptj"]["model_name"] = (
+            "gptj" if hf_config["_name_or_path"] == "" else hf_config["_name_or_path"]
+        )
+        config["gptj"]["head_num"] = str(hf_config["n_head"])
+        n_embd = hf_config["n_embd"]
+        config["gptj"]["size_per_head"] = str(n_embd // hf_config["n_head"])
+        config["gptj"]["inter_size"] = str(n_embd * 4)
+        config["gptj"]["num_layer"] = str(hf_config["n_layer"])
+        rotary_dim = (
+            n_embd // hf_config["n_head"]
+            if hf_config["rotary_dim"] is None
+            else hf_config["rotary_dim"]
+        )
+        config["gptj"]["rotary_embedding"] = str(hf_config["rotary_dim"])
+        config["gptj"]["vocab_size"] = str(hf_config["vocab_size"])
+        config["gptj"]["start_id"] = str(hf_config["bos_token_id"])
+        config["gptj"]["end_id"] = str(hf_config["eos_token_id"])
+        config["gptj"]["weight_data_type"] = "fp16"
+        Path(output_dir).mkdir(exist_ok=True, parents=True)
+        with open(output_dir + "/config.ini", "w") as configfile:
+            config.write(configfile)
+    except:
+        print(f"Fail to save the config in config.ini.")
+
+    n_layers = hf_config["n_layer"]
+    for i in range(n_layers):
+        save(checkpoint, output_dir, args.n_inference_gpus, n_layers, i)
+    savebin(
+        checkpoint["transformer.ln_f.weight"],
+        output_dir + "/model.final_layernorm.weight",
+    )
+    savebin(
+        checkpoint["transformer.ln_f.bias"], output_dir + "/model.final_layernorm.bias"
+    )
+    savebin(checkpoint["lm_head.weight"], output_dir + "/model.lm_head.weight")
+    savebin(checkpoint["lm_head.bias"], output_dir + "/model.lm_head.bias")
+
+    print("done")
diff --git a/tabby/tools/huggingface_gptneox_convert.py b/tabby/tools/converter/huggingface_gptneox_convert.py
similarity index 100%
rename from tabby/tools/huggingface_gptneox_convert.py
rename to tabby/tools/converter/huggingface_gptneox_convert.py