llm-adapters

1
import json
2
import os
3

4
import torch
5
import transformers
6
from peft import PeftModel
7
from transformers import LlamaForCausalLM, LlamaTokenizer  # noqa: E402
8

9
BASE_MODEL = os.environ.get("BASE_MODEL", None)
10
assert (
11
    BASE_MODEL
12
), "Please specify a value for BASE_MODEL environment variable, e.g. `export BASE_MODEL=decapoda-research/llama-7b-hf`"  # noqa: E501
13

14
tokenizer = LlamaTokenizer.from_pretrained(BASE_MODEL)
15

16
base_model = LlamaForCausalLM.from_pretrained(
17
    BASE_MODEL,
18
    load_in_8bit=False,
19
    torch_dtype=torch.float16,
20
    device_map={"": "cpu"},
21
)
22

23
lora_model = PeftModel.from_pretrained(
24
    base_model,
25
    "tloen/alpaca-lora-7b",
26
    device_map={"": "cpu"},
27
    torch_dtype=torch.float16,
28
)
29

30
# merge weights
31
for layer in lora_model.base_model.model.model.layers:
32
    layer.self_attn.q_proj.merge_weights = True
33
    layer.self_attn.v_proj.merge_weights = True
34

35
lora_model.train(False)
36

37
lora_model_sd = lora_model.state_dict()
38

39
params = {
40
    "dim": 4096,
41
    "multiple_of": 256,
42
    "n_heads": 32,
43
    "n_layers": 32,
44
    "norm_eps": 1e-06,
45
    "vocab_size": -1,
46
}
47
n_layers = params["n_layers"]
48
n_heads = params["n_heads"]
49
dim = params["dim"]
50
dims_per_head = dim // n_heads
51
base = 10000.0
52
inv_freq = 1.0 / (
53
    base ** (torch.arange(0, dims_per_head, 2).float() / dims_per_head)
54
)
55

56

57
def permute(w):
58
    return (
59
        w.view(n_heads, dim // n_heads // 2, 2, dim)
60
        .transpose(1, 2)
61
        .reshape(dim, dim)
62
    )
63

64

65
def unpermute(w):
66
    return (
67
        w.view(n_heads, 2, dim // n_heads // 2, dim)
68
        .transpose(1, 2)
69
        .reshape(dim, dim)
70
    )
71

72

73
def translate_state_dict_key(k):  # noqa: C901
74
    k = k.replace("base_model.model.", "")
75
    if k == "model.embed_tokens.weight":
76
        return "tok_embeddings.weight"
77
    elif k == "model.norm.weight":
78
        return "norm.weight"
79
    elif k == "lm_head.weight":
80
        return "output.weight"
81
    elif k.startswith("model.layers."):
82
        layer = k.split(".")[2]
83
        if k.endswith(".self_attn.q_proj.weight"):
84
            return f"layers.{layer}.attention.wq.weight"
85
        elif k.endswith(".self_attn.k_proj.weight"):
86
            return f"layers.{layer}.attention.wk.weight"
87
        elif k.endswith(".self_attn.v_proj.weight"):
88
            return f"layers.{layer}.attention.wv.weight"
89
        elif k.endswith(".self_attn.o_proj.weight"):
90
            return f"layers.{layer}.attention.wo.weight"
91
        elif k.endswith(".mlp.gate_proj.weight"):
92
            return f"layers.{layer}.feed_forward.w1.weight"
93
        elif k.endswith(".mlp.down_proj.weight"):
94
            return f"layers.{layer}.feed_forward.w2.weight"
95
        elif k.endswith(".mlp.up_proj.weight"):
96
            return f"layers.{layer}.feed_forward.w3.weight"
97
        elif k.endswith(".input_layernorm.weight"):
98
            return f"layers.{layer}.attention_norm.weight"
99
        elif k.endswith(".post_attention_layernorm.weight"):
100
            return f"layers.{layer}.ffn_norm.weight"
101
        elif k.endswith("rotary_emb.inv_freq") or "lora" in k:
102
            return None
103
        else:
104
            print(layer, k)
105
            raise NotImplementedError
106
    else:
107
        print(k)
108
        raise NotImplementedError
109

110

111
new_state_dict = {}
112
for k, v in lora_model_sd.items():
113
    new_k = translate_state_dict_key(k)
114
    if new_k is not None:
115
        if "wq" in new_k or "wk" in new_k:
116
            new_state_dict[new_k] = unpermute(v)
117
        else:
118
            new_state_dict[new_k] = v
119

120
os.makedirs("./ckpt", exist_ok=True)
121

122
torch.save(new_state_dict, "./ckpt/consolidated.00.pth")
123

124
with open("./ckpt/params.json", "w") as f:
125
    json.dump(params, f)
126