adntaha · March 31, 2023 02:54 · Mar 25, 2023 · Mar 23, 2023
diff --git a/export_state_dict_checkpoint.py b/export_state_dict_checkpoint.py
@@ -1,4 +1,3 @@
-from transformers import LlamaTokenizer, LlamaForCausalLM
 import os
 import json
 
@@ -9,6 +8,7 @@
 assert (
     "LlamaTokenizer" in transformers._import_structure["models.llama"]
 ), "LLaMA is now in HuggingFace's main branch.\nPlease reinstall it: pip uninstall transformers && pip install git+https://github.com/huggingface/transformers.git"
+from transformers import LlamaTokenizer, LlamaForCausalLM
 
 tokenizer = LlamaTokenizer.from_pretrained("./result")
 
@@ -41,21 +41,18 @@
 dim = params["dim"]
 dims_per_head = dim // n_heads
 base = 10000.0
-inv_freq = 1.0 / \
-    (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head))
+inv_freq = 1.0 / (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head))
 
 
 def permute(w):
     return (
-        w.view(n_heads, dim // n_heads // 2, 2,
-               dim).transpose(1, 2).reshape(dim, dim)
+        w.view(n_heads, dim // n_heads // 2, 2, dim).transpose(1, 2).reshape(dim, dim)
     )
 
 
 def unpermute(w):
     return (
-        w.view(n_heads, 2, dim // n_heads // 2,
-               dim).transpose(1, 2).reshape(dim, dim)
+        w.view(n_heads, 2, dim // n_heads // 2, dim).transpose(1, 2).reshape(dim, dim)
     )
 
 

diff --git a/export_state_dict_checkpoint.py b/export_state_dict_checkpoint.py
@@ -0,0 +1,114 @@
+from transformers import LlamaTokenizer, LlamaForCausalLM
+import os
+import json
+
+import torch
+
+import transformers
+
+assert (
+    "LlamaTokenizer" in transformers._import_structure["models.llama"]
+), "LLaMA is now in HuggingFace's main branch.\nPlease reinstall it: pip uninstall transformers && pip install git+https://github.com/huggingface/transformers.git"
+
+tokenizer = LlamaTokenizer.from_pretrained("./result")
+
+base_model = LlamaForCausalLM.from_pretrained(
+    "./result",
+    load_in_8bit=False,
+    torch_dtype=torch.float16,
+    device_map={"": "cpu"},
+)
+
+# merge weights
+for layer in base_model.model.layers:
+    layer.self_attn.q_proj.merge_weights = True
+    layer.self_attn.v_proj.merge_weights = True
+
+base_model.train(False)
+
+base_model_sd = base_model.state_dict()
+
+params = {
+    "dim": 4096,
+    "multiple_of": 256,
+    "n_heads": 32,
+    "n_layers": 32,
+    "norm_eps": 1e-06,
+    "vocab_size": -1,
+}
+n_layers = params["n_layers"]
+n_heads = params["n_heads"]
+dim = params["dim"]
+dims_per_head = dim // n_heads
+base = 10000.0
+inv_freq = 1.0 / \
+    (base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head))
+
+
+def permute(w):
+    return (
+        w.view(n_heads, dim // n_heads // 2, 2,
+               dim).transpose(1, 2).reshape(dim, dim)
+    )
+
+
+def unpermute(w):
+    return (
+        w.view(n_heads, 2, dim // n_heads // 2,
+               dim).transpose(1, 2).reshape(dim, dim)
+    )
+
+
+def translate_state_dict_key(k):
+    k = k.replace("base_model.model.", "")
+    if k == "model.embed_tokens.weight":
+        return "tok_embeddings.weight"
+    elif k == "model.norm.weight":
+        return "norm.weight"
+    elif k == "lm_head.weight":
+        return "output.weight"
+    elif k.startswith("model.layers."):
+        layer = k.split(".")[2]
+        if k.endswith(".self_attn.q_proj.weight"):
+            return f"layers.{layer}.attention.wq.weight"
+        elif k.endswith(".self_attn.k_proj.weight"):
+            return f"layers.{layer}.attention.wk.weight"
+        elif k.endswith(".self_attn.v_proj.weight"):
+            return f"layers.{layer}.attention.wv.weight"
+        elif k.endswith(".self_attn.o_proj.weight"):
+            return f"layers.{layer}.attention.wo.weight"
+        elif k.endswith(".mlp.gate_proj.weight"):
+            return f"layers.{layer}.feed_forward.w1.weight"
+        elif k.endswith(".mlp.down_proj.weight"):
+            return f"layers.{layer}.feed_forward.w2.weight"
+        elif k.endswith(".mlp.up_proj.weight"):
+            return f"layers.{layer}.feed_forward.w3.weight"
+        elif k.endswith(".input_layernorm.weight"):
+            return f"layers.{layer}.attention_norm.weight"
+        elif k.endswith(".post_attention_layernorm.weight"):
+            return f"layers.{layer}.ffn_norm.weight"
+        elif k.endswith("rotary_emb.inv_freq") or "lora" in k:
+            return None
+        else:
+            print(layer, k)
+            raise NotImplementedError
+    else:
+        print(k)
+        raise NotImplementedError
+
+
+new_state_dict = {}
+for k, v in base_model_sd.items():
+    new_k = translate_state_dict_key(k)
+    if new_k is not None:
+        if "wq" in new_k or "wk" in new_k:
+            new_state_dict[new_k] = unpermute(v)
+        else:
+            new_state_dict[new_k] = v
+
+os.makedirs("./ckpt", exist_ok=True)
+
+torch.save(new_state_dict, "./ckpt/consolidated.00.pth")
+
+with open("./ckpt/params.json", "w") as f:
+    json.dump(params, f)
No results found