Views
No views yet
1import re
2
3import torch
4
5from transformers import RwkvForCausalLM
6
7def convert_state_dict(state_dict):
8 state_dict_keys = list(state_dict.keys())
9 for name in state_dict_keys:
10 weight = state_dict.pop(name)
11 # emb -> embedding
12 if name.startswith("emb."):
13 name = name.replace("emb.", "embeddings.")
14 # ln_0 -> pre_ln (only present at block 0)
15 if name.startswith("blocks.0.ln0"):
16 name = name.replace("blocks.0.ln0", "blocks.0.pre_ln")
17 # att -> attention
18 name = re.sub(r"blocks\.(\d+)\.att", r"blocks.\1.attention", name)
19 # ffn -> feed_forward
20 name = re.sub(r"blocks\.(\d+)\.ffn", r"blocks.\1.feed_forward", name)
21 # time_mix_k -> time_mix_key and reshape
22 if name.endswith(".time_mix_k"):
23 name = name.replace(".time_mix_k", ".time_mix_key")
24 # time_mix_v -> time_mix_value and reshape
25 if name.endswith(".time_mix_v"):
26 name = name.replace(".time_mix_v", ".time_mix_value")
27 # time_mix_r -> time_mix_key and reshape
28 if name.endswith(".time_mix_r"):
29 name = name.replace(".time_mix_r", ".time_mix_receptance")
30
31 if name != "head.weight":
32 name = "rwkv." + name
33
34 state_dict[name] = weight
35 return state_dict
36
37
38def revert_state_dict(state_dict):
39 state_dict_keys = list(state_dict.keys())
40 for name in state_dict_keys:
41 weight = state_dict.pop(name)
42 name = name.removeprefix("rwkv.")
43
44 # emb -> embedding
45 if name.startswith("embeddings."):
46 name = name.replace("embeddings.", "emb.")
47 # ln_0 -> pre_ln (only present at block 0)
48 if name.startswith("blocks.0.pre_ln"):
49 name = name.replace("blocks.0.pre_ln", "blocks.0.ln0")
50 # att -> attention
51 name = re.sub(r"blocks\.(\d+)\.attention", r"blocks.\1.att", name)
52 # ffn -> feed_forward
53 name = re.sub(r"blocks\.(\d+)\.feed_forward", r"blocks.\1.ffn", name)
54 # time_mix_k -> time_mix_key and reshape
55 if name.endswith(".time_mix_key"):
56 name = name.replace(".time_mix_key", ".time_mix_k")
57 # time_mix_v -> time_mix_value and reshape
58 if name.endswith(".time_mix_value"):
59 name = name.replace(".time_mix_value", ".time_mix_v")
60 # time_mix_r -> time_mix_key and reshape
61 if name.endswith(".time_mix_receptance"):
62 name = name.replace(".time_mix_receptance", ".time_mix_r")
63
64 state_dict[name] = weight
65 return state_dict
66
67
68if __name__ == "__main__":
69 # repo = "beomi/KoRWKV-6B"
70 repo = "beomi/KoAlpaca-KoRWKV-6B"
71 model = RwkvForCausalLM.from_pretrained(repo, torch_dtype=torch.bfloat16)
72
73 state_dict = model.state_dict()
74 converted = revert_state_dict(state_dict)
75 name = repo.split("/")[-1] + ".bf16.pth"
76
77 torch.save(converted, name)