|
|
|
import os |
|
import json |
|
|
|
import torch |
|
from transformers import LlamaTokenizer, LlamaForCausalLM |
|
|
|
tokenizer = LlamaTokenizer.from_pretrained("../7B-2nd-train") |
|
|
|
base_model = LlamaForCausalLM.from_pretrained( |
|
"../7B-2nd-train", |
|
load_in_8bit=False, |
|
torch_dtype=torch.float16, |
|
device_map={"": "cpu"}, |
|
) |
|
|
|
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 |
|
|
|
torch.save(new_state_dict, "consolidated.00.pth") |
|
|
|
with open("params.json", "w") as f: |
|
json.dump(params, f) |
|
|
