Delete convert_weight.py
Browse files- convert_weight.py +0 -81
convert_weight.py
DELETED
@@ -1,81 +0,0 @@
|
|
1 |
-
import torch
|
2 |
-
from tqdm import tqdm
|
3 |
-
|
4 |
-
|
5 |
-
input_dir_path = "/scratch/project_462000086/norwegian_gpt/Megatron-DeepSpeed-fixed/mistral-7b-from-scratch-2nd-run/global_step30000"
|
6 |
-
output_dir_path = "/scratch/project_462000086/norwegian_gpt/Megatron-DeepSpeed-fixed/hf_mistral_from_scratch_60k"
|
7 |
-
|
8 |
-
n_hidden = 4096
|
9 |
-
n_ffn_hidden = 14336
|
10 |
-
n_heads = 32
|
11 |
-
n_kv_heads = 8
|
12 |
-
n_layers = 32
|
13 |
-
n_tp = 2
|
14 |
-
|
15 |
-
|
16 |
-
weights = {}
|
17 |
-
|
18 |
-
# embedding
|
19 |
-
embedding_weights = []
|
20 |
-
for i in range(n_tp):
|
21 |
-
path = f"{input_dir_path}/layer_01-model_0{i}-model_states.pt"
|
22 |
-
checkpoint = torch.load(path)
|
23 |
-
|
24 |
-
embedding_weights.append(checkpoint["word_embeddings.weight"].bfloat16())
|
25 |
-
|
26 |
-
weights[f"model.embed_tokens.weight"] = torch.cat(embedding_weights, dim=0)
|
27 |
-
del embedding_weights
|
28 |
-
|
29 |
-
lm_head_weights = []
|
30 |
-
for i in range(n_tp):
|
31 |
-
path = f"{input_dir_path}/layer_{n_layers + 5}-model_0{i}-model_states.pt"
|
32 |
-
checkpoint = torch.load(path)
|
33 |
-
|
34 |
-
lm_head_weights.append(checkpoint["lm_head.weight"].bfloat16())
|
35 |
-
|
36 |
-
weights[f"lm_head.weight"] = torch.cat(lm_head_weights, dim=0)
|
37 |
-
del lm_head_weights
|
38 |
-
|
39 |
-
|
40 |
-
# transformer layers
|
41 |
-
for layer in tqdm(range(n_layers)):
|
42 |
-
q_weights, k_weights, v_weights, o_weights = [], [], [], []
|
43 |
-
up_weights, gate_weights, down_weights = [], [], []
|
44 |
-
|
45 |
-
for i in range(n_tp):
|
46 |
-
path = f"{input_dir_path}/layer_{layer+3:02d}-model_0{i}-model_states.pt"
|
47 |
-
checkpoint = torch.load(path)
|
48 |
-
|
49 |
-
weights[f"model.layers.{layer}.input_layernorm.weight"] = checkpoint["input_layernorm.weight"].bfloat16()
|
50 |
-
weights[f"model.layers.{layer}.post_attention_layernorm.weight"] = checkpoint["post_attention_layernorm.weight"].bfloat16()
|
51 |
-
|
52 |
-
kv_weight = checkpoint["self_attention.key_value.weight"].bfloat16()
|
53 |
-
k_weight, v_weight = torch.chunk(kv_weight, 2, dim=0)
|
54 |
-
k_weights.append(k_weight)
|
55 |
-
v_weights.append(v_weight)
|
56 |
-
|
57 |
-
q_weights.append(checkpoint["self_attention.query.weight"].bfloat16())
|
58 |
-
o_weights.append(checkpoint["self_attention.dense.weight"].bfloat16())
|
59 |
-
down_weights.append(checkpoint["mlp.dense_4h_to_h.weight"].bfloat16())
|
60 |
-
|
61 |
-
up_gate_weight = checkpoint["mlp.dense_h_to_4h.weight"].bfloat16()
|
62 |
-
up_weight, gate_weight = torch.chunk(up_gate_weight, 2, dim=0)
|
63 |
-
up_weights.append(up_weight)
|
64 |
-
gate_weights.append(gate_weight)
|
65 |
-
|
66 |
-
weights[f"model.layers.{layer}.self_attn.q_proj.weight"] = torch.cat(q_weights, dim=0)
|
67 |
-
weights[f"model.layers.{layer}.self_attn.k_proj.weight"] = torch.cat(k_weights, dim=0)
|
68 |
-
weights[f"model.layers.{layer}.self_attn.v_proj.weight"] = torch.cat(v_weights, dim=0)
|
69 |
-
weights[f"model.layers.{layer}.self_attn.o_proj.weight"] = torch.cat(o_weights, dim=1)
|
70 |
-
weights[f"model.layers.{layer}.mlp.up_proj.weight"] = torch.cat(up_weights, dim=0)
|
71 |
-
weights[f"model.layers.{layer}.mlp.gate_proj.weight"] = torch.cat(gate_weights, dim=0)
|
72 |
-
weights[f"model.layers.{layer}.mlp.down_proj.weight"] = torch.cat(down_weights, dim=1)
|
73 |
-
|
74 |
-
|
75 |
-
# output layer norm
|
76 |
-
path = f"{input_dir_path}/layer_{n_layers + 4}-model_00-model_states.pt"
|
77 |
-
checkpoint = torch.load(path)
|
78 |
-
|
79 |
-
weights[f"model.norm.weight"] = checkpoint["weight"].bfloat16()
|
80 |
-
|
81 |
-
torch.save(weights, f"{output_dir_path}/pytorch_model.bin")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|