File size: 4,411 Bytes
57c9ebd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
import glob
import re
import shutil
import sys
import accelerate
import torch
from safetensors import safe_open
from configuration_afmoe_scm import AfmoeSCMConfig
from modeling_afmoe import AfmoeForCausalLM
from configuration_afmoe import AfmoeConfig
input_model = sys.argv[1]
output_model_path = sys.argv[2]
auto_map = {
"AutoConfig": "configuration_afmoe.AfmoeConfig",
"AutoModel": "modeling_afmoe.AfmoeModel",
"AutoModelForCausalLM": "modeling_afmoe.AfmoeForCausalLM"
}
cfg_shared_moe = AfmoeSCMConfig.from_pretrained(input_model)
cfg_standard_moe = AfmoeConfig(
auto_map=auto_map,
layer_types=cfg_shared_moe.layer_types,
global_attn_every_n_layers=cfg_shared_moe.global_attn_every_n_layers,
load_balance_coeff=cfg_shared_moe.load_balance_coeff,
mup_enabled=cfg_shared_moe.mup_enabled,
num_dense_layers=cfg_shared_moe.num_dense_layers,
num_expert_groups=cfg_shared_moe.num_expert_groups,
num_limited_groups=cfg_shared_moe.num_limited_groups,
route_norm=cfg_shared_moe.route_norm,
route_scale=cfg_shared_moe.route_scale,
score_func=cfg_shared_moe.score_func,
topk_group=cfg_shared_moe.topk_group,
num_shared_experts=cfg_shared_moe.num_shared_experts,
vocab_size=cfg_shared_moe.vocab_size,
hidden_size=cfg_shared_moe.hidden_size,
intermediate_size=cfg_shared_moe.intermediate_size,
num_hidden_layers=cfg_shared_moe.num_hidden_layers,
num_attention_heads=cfg_shared_moe.num_attention_heads,
num_key_value_heads=cfg_shared_moe.num_key_value_heads,
hidden_act=cfg_shared_moe.hidden_act,
max_position_embeddings=cfg_shared_moe.max_position_embeddings,
initializer_range=cfg_shared_moe.initializer_range,
rms_norm_eps=cfg_shared_moe.rms_norm_eps,
use_cache=cfg_shared_moe.use_cache,
tie_word_embeddings=cfg_shared_moe.tie_word_embeddings,
rope_theta=cfg_shared_moe.rope_theta,
rope_scaling=cfg_shared_moe.rope_scaling,
sliding_window=cfg_shared_moe.sliding_window,
attention_dropout=cfg_shared_moe.attention_dropout,
moe_intermediate_size=cfg_shared_moe.moe_intermediate_size,
num_experts_per_tok=cfg_shared_moe.num_experts_per_tok,
num_experts=cfg_shared_moe.num_experts,
head_dim=cfg_shared_moe.head_dim,
eos_token_id=cfg_shared_moe.eos_token_id,
pad_token_id=cfg_shared_moe.pad_token_id,
torch_dtype=cfg_shared_moe.torch_dtype,
)
num_experts = cfg_shared_moe.num_experts
with accelerate.init_empty_weights():
model_standard_moe = AfmoeForCausalLM(cfg_standard_moe)
model_standard_moe = model_standard_moe.to(torch.bfloat16)
new_state_dict = {}
pattern = f"{input_model}/model-*-of-*.safetensors"
files = sorted(glob.glob(pattern))
if len(files) == 0:
raise FileNotFoundError
tensors = {}
for file_path in files:
print(f"processing {file_path}")
with safe_open(file_path, framework="pt", device="cpu") as f:
for key in f.keys():
tensor = f.get_tensor(key)
tensors[key] = tensor
for key in tensors:
if "moe_mlp" not in key:
new_state_dict[key] = tensors[key]
elif "moe_mlp.output_experts" in key:
layer_num = int(re.search(r"\d+", key).group())
for i, tensor in enumerate(torch.unbind(tensors[key])):
new_state_dict[
f"model.layers.{layer_num}.mlp.experts.{i}.down_proj.weight"
] = tensor.contiguous()
elif "moe_mlp.experts" in key:
layer_num = int(re.search(r"\d+", key).group())
for i, tensor in enumerate(torch.unbind(tensors[key])):
(
new_state_dict[
f"model.layers.{layer_num}.mlp.experts.{i}.up_proj.weight"
],
new_state_dict[
f"model.layers.{layer_num}.mlp.experts.{i}.gate_proj.weight"
],
) = torch.chunk(tensor, 2, dim=0)
model_standard_moe.load_state_dict(new_state_dict, strict=True, assign=True)
model_standard_moe.save_pretrained(output_model_path)
cfg_standard_moe.save_pretrained(output_model_path)
shutil.copy(
"modeling_afmoe.py",
output_model_path + "/" + "modeling_afmoe.py",
)
shutil.copy(
"configuration_afmoe.py",
output_model_path + "/" + "configuration_afmoe.py",
)
for i in ["special_tokens_map.json", "tokenizer_config.json", "tokenizer.json", "chat_template.jinja"]:
shutil.copy(input_model + "/" + i, output_model_path + "/" + i)
|