Hugging Face's logo

Alpha-VLLM
/
MoE-Mixtral-7B-8Expert

Model card Files Community
Cxxs commited on
Commit
99ed461
•
1 Parent(s): 07550da

Upload split_sparse.py

Browse files
Files changed (1) hide show
  1. converted_sparse/split_sparse.py +153 -0
converted_sparse/split_sparse.py ADDED
@@ -0,0 +1,153 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ weight_parallel_dim = {"llma.tok_embeddings.weight": 1, "llma.layers.0.attention.wq.weight": 0,
2
+ "llma.layers.0.attention.wq.bias": 0, "llma.layers.0.attention.wk.weight": 0,
3
+ "llma.layers.0.attention.wk.bias": 0, "llma.layers.0.attention.wv.weight": 0,
4
+ "llma.layers.0.attention.wv.bias": 0, "llma.layers.0.attention.wo.weight": 1,
5
+ "llma.layers.1.attention.wq.weight": 0, "llma.layers.1.attention.wq.bias": 0,
6
+ "llma.layers.1.attention.wk.weight": 0, "llma.layers.1.attention.wk.bias": 0,
7
+ "llma.layers.1.attention.wv.weight": 0, "llma.layers.1.attention.wv.bias": 0,
8
+ "llma.layers.1.attention.wo.weight": 1, "llma.layers.2.attention.wq.weight": 0,
9
+ "llma.layers.2.attention.wq.bias": 0, "llma.layers.2.attention.wk.weight": 0,
10
+ "llma.layers.2.attention.wk.bias": 0, "llma.layers.2.attention.wv.weight": 0,
11
+ "llma.layers.2.attention.wv.bias": 0, "llma.layers.2.attention.wo.weight": 1,
12
+ "llma.layers.3.attention.wq.weight": 0, "llma.layers.3.attention.wq.bias": 0,
13
+ "llma.layers.3.attention.wk.weight": 0, "llma.layers.3.attention.wk.bias": 0,
14
+ "llma.layers.3.attention.wv.weight": 0, "llma.layers.3.attention.wv.bias": 0,
15
+ "llma.layers.3.attention.wo.weight": 1, "llma.layers.4.attention.wq.weight": 0,
16
+ "llma.layers.4.attention.wq.bias": 0, "llma.layers.4.attention.wk.weight": 0,
17
+ "llma.layers.4.attention.wk.bias": 0, "llma.layers.4.attention.wv.weight": 0,
18
+ "llma.layers.4.attention.wv.bias": 0, "llma.layers.4.attention.wo.weight": 1,
19
+ "llma.layers.5.attention.wq.weight": 0, "llma.layers.5.attention.wq.bias": 0,
20
+ "llma.layers.5.attention.wk.weight": 0, "llma.layers.5.attention.wk.bias": 0,
21
+ "llma.layers.5.attention.wv.weight": 0, "llma.layers.5.attention.wv.bias": 0,
22
+ "llma.layers.5.attention.wo.weight": 1, "llma.layers.6.attention.wq.weight": 0,
23
+ "llma.layers.6.attention.wq.bias": 0, "llma.layers.6.attention.wk.weight": 0,
24
+ "llma.layers.6.attention.wk.bias": 0, "llma.layers.6.attention.wv.weight": 0,
25
+ "llma.layers.6.attention.wv.bias": 0, "llma.layers.6.attention.wo.weight": 1,
26
+ "llma.layers.7.attention.wq.weight": 0, "llma.layers.7.attention.wq.bias": 0,
27
+ "llma.layers.7.attention.wk.weight": 0, "llma.layers.7.attention.wk.bias": 0,
28
+ "llma.layers.7.attention.wv.weight": 0, "llma.layers.7.attention.wv.bias": 0,
29
+ "llma.layers.7.attention.wo.weight": 1, "llma.layers.8.attention.wq.weight": 0,
30
+ "llma.layers.8.attention.wq.bias": 0, "llma.layers.8.attention.wk.weight": 0,
31
+ "llma.layers.8.attention.wk.bias": 0, "llma.layers.8.attention.wv.weight": 0,
32
+ "llma.layers.8.attention.wv.bias": 0, "llma.layers.8.attention.wo.weight": 1,
33
+ "llma.layers.9.attention.wq.weight": 0, "llma.layers.9.attention.wq.bias": 0,
34
+ "llma.layers.9.attention.wk.weight": 0, "llma.layers.9.attention.wk.bias": 0,
35
+ "llma.layers.9.attention.wv.weight": 0, "llma.layers.9.attention.wv.bias": 0,
36
+ "llma.layers.9.attention.wo.weight": 1, "llma.layers.10.attention.wq.weight": 0,
37
+ "llma.layers.10.attention.wq.bias": 0, "llma.layers.10.attention.wk.weight": 0,
38
+ "llma.layers.10.attention.wk.bias": 0, "llma.layers.10.attention.wv.weight": 0,
39
+ "llma.layers.10.attention.wv.bias": 0, "llma.layers.10.attention.wo.weight": 1,
40
+ "llma.layers.11.attention.wq.weight": 0, "llma.layers.11.attention.wq.bias": 0,
41
+ "llma.layers.11.attention.wk.weight": 0, "llma.layers.11.attention.wk.bias": 0,
42
+ "llma.layers.11.attention.wv.weight": 0, "llma.layers.11.attention.wv.bias": 0,
43
+ "llma.layers.11.attention.wo.weight": 1, "llma.layers.12.attention.wq.weight": 0,
44
+ "llma.layers.12.attention.wq.bias": 0, "llma.layers.12.attention.wk.weight": 0,
45
+ "llma.layers.12.attention.wk.bias": 0, "llma.layers.12.attention.wv.weight": 0,
46
+ "llma.layers.12.attention.wv.bias": 0, "llma.layers.12.attention.wo.weight": 1,
47
+ "llma.layers.13.attention.wq.weight": 0, "llma.layers.13.attention.wq.bias": 0,
48
+ "llma.layers.13.attention.wk.weight": 0, "llma.layers.13.attention.wk.bias": 0,
49
+ "llma.layers.13.attention.wv.weight": 0, "llma.layers.13.attention.wv.bias": 0,
50
+ "llma.layers.13.attention.wo.weight": 1, "llma.layers.14.attention.wq.weight": 0,
51
+ "llma.layers.14.attention.wq.bias": 0, "llma.layers.14.attention.wk.weight": 0,
52
+ "llma.layers.14.attention.wk.bias": 0, "llma.layers.14.attention.wv.weight": 0,
53
+ "llma.layers.14.attention.wv.bias": 0, "llma.layers.14.attention.wo.weight": 1,
54
+ "llma.layers.15.attention.wq.weight": 0, "llma.layers.15.attention.wq.bias": 0,
55
+ "llma.layers.15.attention.wk.weight": 0, "llma.layers.15.attention.wk.bias": 0,
56
+ "llma.layers.15.attention.wv.weight": 0, "llma.layers.15.attention.wv.bias": 0,
57
+ "llma.layers.15.attention.wo.weight": 1, "llma.layers.16.attention.wq.weight": 0,
58
+ "llma.layers.16.attention.wq.bias": 0, "llma.layers.16.attention.wk.weight": 0,
59
+ "llma.layers.16.attention.wk.bias": 0, "llma.layers.16.attention.wv.weight": 0,
60
+ "llma.layers.16.attention.wv.bias": 0, "llma.layers.16.attention.wo.weight": 1,
61
+ "llma.layers.17.attention.wq.weight": 0, "llma.layers.17.attention.wq.bias": 0,
62
+ "llma.layers.17.attention.wk.weight": 0, "llma.layers.17.attention.wk.bias": 0,
63
+ "llma.layers.17.attention.wv.weight": 0, "llma.layers.17.attention.wv.bias": 0,
64
+ "llma.layers.17.attention.wo.weight": 1, "llma.layers.18.attention.wq.weight": 0,
65
+ "llma.layers.18.attention.wq.bias": 0, "llma.layers.18.attention.wk.weight": 0,
66
+ "llma.layers.18.attention.wk.bias": 0, "llma.layers.18.attention.wv.weight": 0,
67
+ "llma.layers.18.attention.wv.bias": 0, "llma.layers.18.attention.wo.weight": 1,
68
+ "llma.layers.19.attention.wq.weight": 0, "llma.layers.19.attention.wq.bias": 0,
69
+ "llma.layers.19.attention.wk.weight": 0, "llma.layers.19.attention.wk.bias": 0,
70
+ "llma.layers.19.attention.wv.weight": 0, "llma.layers.19.attention.wv.bias": 0,
71
+ "llma.layers.19.attention.wo.weight": 1, "llma.layers.20.attention.wq.weight": 0,
72
+ "llma.layers.20.attention.wq.bias": 0, "llma.layers.20.attention.wk.weight": 0,
73
+ "llma.layers.20.attention.wk.bias": 0, "llma.layers.20.attention.wv.weight": 0,
74
+ "llma.layers.20.attention.wv.bias": 0, "llma.layers.20.attention.wo.weight": 1,
75
+ "llma.layers.21.attention.wq.weight": 0, "llma.layers.21.attention.wq.bias": 0,
76
+ "llma.layers.21.attention.wk.weight": 0, "llma.layers.21.attention.wk.bias": 0,
77
+ "llma.layers.21.attention.wv.weight": 0, "llma.layers.21.attention.wv.bias": 0,
78
+ "llma.layers.21.attention.wo.weight": 1, "llma.layers.22.attention.wq.weight": 0,
79
+ "llma.layers.22.attention.wq.bias": 0, "llma.layers.22.attention.wk.weight": 0,
80
+ "llma.layers.22.attention.wk.bias": 0, "llma.layers.22.attention.wv.weight": 0,
81
+ "llma.layers.22.attention.wv.bias": 0, "llma.layers.22.attention.wo.weight": 1,
82
+ "llma.layers.23.attention.wq.weight": 0, "llma.layers.23.attention.wq.bias": 0,
83
+ "llma.layers.23.attention.wk.weight": 0, "llma.layers.23.attention.wk.bias": 0,
84
+ "llma.layers.23.attention.wv.weight": 0, "llma.layers.23.attention.wv.bias": 0,
85
+ "llma.layers.23.attention.wo.weight": 1, "llma.layers.24.attention.wq.weight": 0,
86
+ "llma.layers.24.attention.wq.bias": 0, "llma.layers.24.attention.wk.weight": 0,
87
+ "llma.layers.24.attention.wk.bias": 0, "llma.layers.24.attention.wv.weight": 0,
88
+ "llma.layers.24.attention.wv.bias": 0, "llma.layers.24.attention.wo.weight": 1,
89
+ "llma.layers.25.attention.wq.weight": 0, "llma.layers.25.attention.wq.bias": 0,
90
+ "llma.layers.25.attention.wk.weight": 0, "llma.layers.25.attention.wk.bias": 0,
91
+ "llma.layers.25.attention.wv.weight": 0, "llma.layers.25.attention.wv.bias": 0,
92
+ "llma.layers.25.attention.wo.weight": 1, "llma.layers.26.attention.wq.weight": 0,
93
+ "llma.layers.26.attention.wq.bias": 0, "llma.layers.26.attention.wk.weight": 0,
94
+ "llma.layers.26.attention.wk.bias": 0, "llma.layers.26.attention.wv.weight": 0,
95
+ "llma.layers.26.attention.wv.bias": 0, "llma.layers.26.attention.wo.weight": 1,
96
+ "llma.layers.27.attention.wq.weight": 0, "llma.layers.27.attention.wq.bias": 0,
97
+ "llma.layers.27.attention.wk.weight": 0, "llma.layers.27.attention.wk.bias": 0,
98
+ "llma.layers.27.attention.wv.weight": 0, "llma.layers.27.attention.wv.bias": 0,
99
+ "llma.layers.27.attention.wo.weight": 1, "llma.layers.28.attention.wq.weight": 0,
100
+ "llma.layers.28.attention.wq.bias": 0, "llma.layers.28.attention.wk.weight": 0,
101
+ "llma.layers.28.attention.wk.bias": 0, "llma.layers.28.attention.wv.weight": 0,
102
+ "llma.layers.28.attention.wv.bias": 0, "llma.layers.28.attention.wo.weight": 1,
103
+ "llma.layers.29.attention.wq.weight": 0, "llma.layers.29.attention.wq.bias": 0,
104
+ "llma.layers.29.attention.wk.weight": 0, "llma.layers.29.attention.wk.bias": 0,
105
+ "llma.layers.29.attention.wv.weight": 0, "llma.layers.29.attention.wv.bias": 0,
106
+ "llma.layers.29.attention.wo.weight": 1, "llma.layers.30.attention.wq.weight": 0,
107
+ "llma.layers.30.attention.wq.bias": 0, "llma.layers.30.attention.wk.weight": 0,
108
+ "llma.layers.30.attention.wk.bias": 0, "llma.layers.30.attention.wv.weight": 0,
109
+ "llma.layers.30.attention.wv.bias": 0, "llma.layers.30.attention.wo.weight": 1,
110
+ "llma.layers.31.attention.wq.weight": 0, "llma.layers.31.attention.wq.bias": 0,
111
+ "llma.layers.31.attention.wk.weight": 0, "llma.layers.31.attention.wk.bias": 0,
112
+ "llma.layers.31.attention.wv.weight": 0, "llma.layers.31.attention.wv.bias": 0,
113
+ "llma.layers.31.attention.wo.weight": 1, "llma.output.weight": 0, "llma.output.bias": 0}
114
+
115
+ import torch
116
+ from pathlib import Path
117
+
118
+ Path("./converted_sparse").mkdir(exist_ok=True)
119
+
120
+ ori = torch.load("consolidated.00.pth", map_location="cpu")
121
+ ori = {"llma." + key: val for key, val in ori.items()}
122
+
123
+
124
+ def func(rank=0):
125
+ shard_split_to = 8
126
+ split_ckpt = {}
127
+ for key, ori_param in ori.items():
128
+ if key in weight_parallel_dim:
129
+ split_ckpt[key] = torch.chunk(ori_param, shard_split_to, weight_parallel_dim[key])[
130
+ rank % shard_split_to].clone()
131
+ if rank == 0:
132
+ print(f"chunk {key}")
133
+ else:
134
+ if "experts.0." in key:
135
+ weight_all_experts = [ori[key.replace("experts.0.", f"experts.{i}.")] for i in range(8)]
136
+ if "w2" in key:
137
+ weight_all_experts = [torch.transpose(_, 0, 1) for _ in weight_all_experts]
138
+ weight_this_rank = [torch.chunk(_, 8, dim=0)[rank] for _ in weight_all_experts]
139
+ weight_this_rank = torch.cat(weight_this_rank, dim=0).clone()
140
+ key = key.replace("experts.0.", "").replace(".weight", "")
141
+ split_ckpt[key] = weight_this_rank
142
+ print("expert key")
143
+ elif "experts" in key:
144
+ continue
145
+ else:
146
+ split_ckpt[key] = ori_param
147
+ if rank == 0:
148
+ print(f"inherit {key}")
149
+ torch.save({"model": split_ckpt}, f"converted_sparse/consolidated.{rank:02d}-of-08.model.pth")
150
+
151
+
152
+ for r in range(8):
153
+ func(r)