davda54 commited on
Commit
2215488
1 Parent(s): d0a30cf

Delete convert_weight.py

Browse files
Files changed (1) hide show
  1. 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")