thomasgauthier
commited on
Commit
•
48bedfe
1
Parent(s):
04b7792
Added expert extraction code
Browse files
README.md
CHANGED
@@ -35,4 +35,109 @@ The following named weight correspondance was used:
|
|
35 |
| [Unmixtraled-22B-v0.1-expert-5](https://huggingface.co/thomasgauthier/Unmixtraled-22B-v0.1-expert-5) | Mixtral 8x22B embed, attn, layernorm, lm_head + expert 5 MLPs | 1099.32373046875 |
|
36 |
| [Unmixtraled-22B-v0.1-expert-6](https://huggingface.co/thomasgauthier/Unmixtraled-22B-v0.1-expert-6) | Mixtral 8x22B embed, attn, layernorm, lm_head + expert 6 MLPs | 341.5309753417969 |
|
37 |
| [Unmixtraled-22B-v0.1-expert-7](https://huggingface.co/thomasgauthier/Unmixtraled-22B-v0.1-expert-7) | Mixtral 8x22B embed, attn, layernorm, lm_head + expert 7 MLPs | 2099.63818359375 |
|
38 |
-
| [Unmixtraled-22B-v0.1-lerp](https://huggingface.co/thomasgauthier/Unmixtraled-22B-v0.1-lerp) | Mixtral 8x22B embed, attn, layernorm, lm_head + linear merge of expert 0-7 MLPs | 1873.9874267578125 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
35 |
| [Unmixtraled-22B-v0.1-expert-5](https://huggingface.co/thomasgauthier/Unmixtraled-22B-v0.1-expert-5) | Mixtral 8x22B embed, attn, layernorm, lm_head + expert 5 MLPs | 1099.32373046875 |
|
36 |
| [Unmixtraled-22B-v0.1-expert-6](https://huggingface.co/thomasgauthier/Unmixtraled-22B-v0.1-expert-6) | Mixtral 8x22B embed, attn, layernorm, lm_head + expert 6 MLPs | 341.5309753417969 |
|
37 |
| [Unmixtraled-22B-v0.1-expert-7](https://huggingface.co/thomasgauthier/Unmixtraled-22B-v0.1-expert-7) | Mixtral 8x22B embed, attn, layernorm, lm_head + expert 7 MLPs | 2099.63818359375 |
|
38 |
+
| [Unmixtraled-22B-v0.1-lerp](https://huggingface.co/thomasgauthier/Unmixtraled-22B-v0.1-lerp) | Mixtral 8x22B embed, attn, layernorm, lm_head + linear merge of expert 0-7 MLPs | 1873.9874267578125 |
|
39 |
+
|
40 |
+
# Code
|
41 |
+
|
42 |
+
The following code was used to extract the experts and construct the dense models:
|
43 |
+
|
44 |
+
```python
|
45 |
+
# pip install -U transformers huggingface_hub "git+https://github.com/arcee-ai/mergekit@7467108c05d56ef2bb4b8f33936d437dc448f7dd"
|
46 |
+
|
47 |
+
import fnmatch
|
48 |
+
import json
|
49 |
+
import os
|
50 |
+
import re
|
51 |
+
import shutil
|
52 |
+
|
53 |
+
import torch
|
54 |
+
from huggingface_hub import snapshot_download
|
55 |
+
from mergekit.architecture import get_architecture_info
|
56 |
+
from mergekit.common import ModelReference
|
57 |
+
from mergekit.io import LazyTensorLoader, TensorWriter
|
58 |
+
from tqdm import tqdm
|
59 |
+
|
60 |
+
MIXTRAL_MODEL_ID = "mistral-community/Mixtral-8x22B-v0.1"
|
61 |
+
MIXTRAL_PATH = snapshot_download(repo_id=MIXTRAL_MODEL_ID)
|
62 |
+
print(f"Mixtral downloaded to: {MIXTRAL_PATH}")
|
63 |
+
|
64 |
+
MISTRAL_PATH = snapshot_download(
|
65 |
+
repo_id="mistralai/Mistral-7B-v0.1", allow_patterns=["config.json"]
|
66 |
+
)
|
67 |
+
print(f"Mistral config downloaded to: {MISTRAL_PATH}")
|
68 |
+
|
69 |
+
with open(os.path.join(MISTRAL_PATH, "config.json"), "r") as f:
|
70 |
+
mistral_config = json.load(f)
|
71 |
+
|
72 |
+
with open(os.path.join(MIXTRAL_PATH, "config.json"), "r") as f:
|
73 |
+
mixtral_config = json.load(f)
|
74 |
+
|
75 |
+
combined_config = {
|
76 |
+
key: mixtral_config[key] for key in mistral_config if key in mixtral_config
|
77 |
+
}
|
78 |
+
combined_config["architectures"] = ["MistralForCausalLM"]
|
79 |
+
combined_config["model_type"] = "mistral"
|
80 |
+
|
81 |
+
mixtral_model_ref = ModelReference.parse(MIXTRAL_PATH)
|
82 |
+
mixtral_architecture_info = get_architecture_info(mixtral_model_ref.config())
|
83 |
+
mixtral_loader = LazyTensorLoader(mixtral_model_ref.tensor_index(), lazy_unpickle=True)
|
84 |
+
|
85 |
+
ALLOW_LIST = ["generation_config.json", "tokenizer.model", "tokenizer_config.json"]
|
86 |
+
|
87 |
+
def copy_directory(src, dest, allowed_patterns):
|
88 |
+
os.makedirs(dest, exist_ok=True)
|
89 |
+
for root, dirs, files in os.walk(src):
|
90 |
+
# Only keep directories that match at least one of the allowed patterns
|
91 |
+
dirs[:] = [d for d in dirs if any(fnmatch.fnmatch(d, pattern) for pattern in allowed_patterns)]
|
92 |
+
for file in files:
|
93 |
+
# Only copy files that match at least one of the allowed patterns
|
94 |
+
if any(fnmatch.fnmatch(file, pattern) for pattern in allowed_patterns):
|
95 |
+
src_path = os.path.join(root, file)
|
96 |
+
dest_path = os.path.join(dest, os.path.relpath(src_path, src))
|
97 |
+
os.makedirs(os.path.dirname(dest_path), exist_ok=True)
|
98 |
+
shutil.copy2(src_path, dest_path)
|
99 |
+
|
100 |
+
def get_tensor(layer_num, expert_num, tensor_type):
|
101 |
+
weight_name = f"model.layers.{layer_num}.block_sparse_moe.experts.{expert_num}.{tensor_type}.weight"
|
102 |
+
return mixtral_loader.get_tensor(weight_name)
|
103 |
+
|
104 |
+
|
105 |
+
def extract_layer_number(string):
|
106 |
+
match = re.search(r"layers\.(\d+)\.", string)
|
107 |
+
return int(match.group(1)) if match else None
|
108 |
+
|
109 |
+
|
110 |
+
def save_expert_as_dense(output_path, expert_num):
|
111 |
+
dense_model_ref = ModelReference.parse(output_path)
|
112 |
+
dense_architecture_info = get_architecture_info(dense_model_ref.config())
|
113 |
+
|
114 |
+
writer = TensorWriter(output_path, safe_serialization=True)
|
115 |
+
|
116 |
+
for weight_info in tqdm(dense_architecture_info.all_weights(dense_model_ref.config())):
|
117 |
+
if weight_info.name.endswith(".up_proj.weight"):
|
118 |
+
layer_num = extract_layer_number(weight_info.name)
|
119 |
+
writer.save_tensor(weight_info.name, get_tensor(layer_num, expert_num, "w3"))
|
120 |
+
elif weight_info.name.endswith(".down_proj.weight"):
|
121 |
+
layer_num = extract_layer_number(weight_info.name)
|
122 |
+
writer.save_tensor(weight_info.name, get_tensor(layer_num, expert_num, "w2"))
|
123 |
+
elif weight_info.name.endswith(".gate_proj.weight"):
|
124 |
+
layer_num = extract_layer_number(weight_info.name)
|
125 |
+
writer.save_tensor(weight_info.name, get_tensor(layer_num, expert_num, "w1"))
|
126 |
+
else:
|
127 |
+
writer.save_tensor(weight_info.name, mixtral_loader.get_tensor(weight_info.name))
|
128 |
+
|
129 |
+
writer.finalize()
|
130 |
+
|
131 |
+
|
132 |
+
num_experts = mixtral_config["num_local_experts"]
|
133 |
+
|
134 |
+
for expert_num in range(num_experts):
|
135 |
+
dense_path = f"./dense_expert_{expert_num}"
|
136 |
+
copy_directory(MIXTRAL_PATH, dense_path, ALLOW_LIST)
|
137 |
+
|
138 |
+
with open(os.path.join(dense_path, "config.json"), "w") as f:
|
139 |
+
json.dump(combined_config, f, indent=2)
|
140 |
+
|
141 |
+
save_expert_as_dense(dense_path, expert_num)
|
142 |
+
print(f"Dense model #{expert_num} saved to {os.path.abspath(dense_path)}")
|
143 |
+
```
|