wehub-resource-sync
173 lines
7.1 KiB
Python
173 lines
7.1 KiB
Python
import argparse
|
|
import os
|
|
from collections import OrderedDict
|
|
from glob import glob
|
|
|
|
import torch
|
|
from safetensors import safe_open
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel, AutoConfig
|
|
from accelerate import init_empty_weights
|
|
from glob import glob
|
|
|
|
|
|
# permute for sliced rotary
|
|
def permute(w, n_heads, dim1, dim2):
|
|
"""
|
|
|
|
"""
|
|
return w.view(n_heads, dim1 // n_heads // 2, 2, dim2).transpose(1, 2).reshape(dim1, dim2)
|
|
|
|
|
|
def split_weights(state_dict: OrderedDict, tp_size: int):
|
|
new_state_dicts = [OrderedDict() for _ in range(tp_size)]
|
|
|
|
split_lists = ["q_proj.weight", "k_proj.weight", "v_proj.weight", "gate_proj.weight", "up_proj.weight",
|
|
"q_proj.bias", "k_proj.bias", "v_proj.bias",
|
|
"embed_tokens.weight", "lm_head.weight"]
|
|
|
|
for k, v in state_dict.items():
|
|
if any(k.endswith(name) for name in split_lists):
|
|
tensor = list(torch.chunk(v, tp_size, dim=0))
|
|
for i, t in enumerate(tensor):
|
|
new_state_dicts[i][k] = t.detach().clone()
|
|
elif k.endswith("o_proj.weight") or k.endswith("down_proj.weight"):
|
|
tensor = list(torch.chunk(v, tp_size, dim=1))
|
|
for i, t in enumerate(tensor):
|
|
new_state_dicts[i][k] = t.detach().clone()
|
|
else:
|
|
for i in range(tp_size):
|
|
new_state_dicts[i][k] = v.detach().clone()
|
|
print(new_state_dicts[0].keys())
|
|
return new_state_dicts
|
|
|
|
|
|
def merge_weights(state_dicts: list, merge_avg: bool = False):
|
|
merged_state_dict = OrderedDict()
|
|
|
|
col_splits = ["q_proj.weight", "k_proj.weight", "v_proj.weight", "gate_proj.weight", "up_proj.weight",
|
|
"q_proj.bias", "k_proj.bias", "v_proj.bias",
|
|
"embed_tokens.weight", "lm_head.weight"]
|
|
row_splits = ["o_proj.weight", "down_proj.weight"]
|
|
|
|
for k in state_dicts[0].keys():
|
|
if any(k.endswith(name) for name in col_splits):
|
|
merged_state_dict[k] = torch.cat([state_dict[k] for state_dict in state_dicts], dim=0)
|
|
elif any(k.endswith(name) for name in row_splits):
|
|
merged_state_dict[k] = torch.cat([state_dict[k] for state_dict in state_dicts], dim=1)
|
|
else:
|
|
merged_state_dict[k] = state_dicts[0][k]
|
|
if merge_avg:
|
|
for i in range(1, len(state_dicts)):
|
|
merged_state_dict[k] += state_dicts[i][k]
|
|
merged_state_dict[k] /= len(state_dicts)
|
|
|
|
# print(merged_state_dict.keys())
|
|
return merged_state_dict
|
|
|
|
|
|
def write_model(input_base_path, tp_size: int):
|
|
# model = LlamaForCausalLM.from_pretrained(input_base_path, torch_dtype="auto", device_map="cpu")
|
|
config = AutoConfig.from_pretrained(input_base_path)
|
|
with init_empty_weights():
|
|
model = AutoModel.from_config(config)
|
|
tokenizer = AutoTokenizer.from_pretrained(input_base_path)
|
|
|
|
weights = OrderedDict()
|
|
files = glob(os.path.join(input_base_path, "*.safetensors"))
|
|
if len(files):
|
|
for weight_path in files:
|
|
with safe_open(weight_path, framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
weights[key] = f.get_tensor(key).clone()
|
|
else:
|
|
for weight_path in glob(os.path.join(input_base_path, "pytorch_model*.bin")):
|
|
weights.update(torch.load(weight_path, map_location="cpu"))
|
|
|
|
new_state_dicts = split_weights(weights, tp_size)
|
|
for i, state_dict in enumerate(new_state_dicts):
|
|
output_folder = os.path.join(input_base_path, f"mp_{i}-of-{tp_size}")
|
|
assert not os.path.exists(output_folder), f"Folder {output_folder} already exists. Please remove it before splitting."
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
model.save_pretrained(output_folder, state_dict=state_dict, safe_serialization=False)
|
|
tokenizer.save_pretrained(output_folder)
|
|
print(f"Model saved to {output_folder}")
|
|
|
|
|
|
def merge_model(input_base_path, tp_size: int, merge_avg: bool = False):
|
|
config = AutoConfig.from_pretrained(os.path.join(input_base_path, f"mp_0-of-{tp_size}"))
|
|
if any(arch in config.architectures for arch in ["LlamaForCausalLMDPO"]):
|
|
config.architectures = ["LlamaForCausalLM"]
|
|
with init_empty_weights():
|
|
model = AutoModelForCausalLM.from_config(config)
|
|
tokenizer = AutoTokenizer.from_pretrained(input_base_path)
|
|
|
|
state_dicts = []
|
|
for i in range(tp_size):
|
|
folder = os.path.join(input_base_path, f"mp_{i}-of-{tp_size}")
|
|
print(folder)
|
|
weights = {}
|
|
files = glob(os.path.join(folder, "*.safetensors"))
|
|
if len(files):
|
|
for weight_path in files:
|
|
with safe_open(weight_path, framework="pt", device="cpu") as f:
|
|
for key in f.keys():
|
|
weights[key] = f.get_tensor(key).clone()
|
|
else:
|
|
for weight_path in glob(os.path.join(folder, "pytorch_model*.bin")):
|
|
weights.update(torch.load(weight_path, map_location="cpu"))
|
|
|
|
for k in weights.keys():
|
|
if weights[k].dtype != config.torch_dtype:
|
|
weights[k] = weights[k].to(config.torch_dtype)
|
|
state_dicts.append(weights)
|
|
|
|
merged_state_dict = merge_weights(state_dicts, merge_avg)
|
|
if merge_avg:
|
|
output_folder = os.path.join(input_base_path, "merged_avg")
|
|
model.save_pretrained(output_folder, state_dict=merged_state_dict, safe_serialization=False)
|
|
tokenizer.save_pretrained(output_folder)
|
|
config.save_pretrained(output_folder)
|
|
print(f"Model saved to {output_folder}")
|
|
else:
|
|
tmp_files = glob(os.path.join(input_base_path, "pytorch_model*.bin"))
|
|
assert len(tmp_files) == 0, "Found existing pytorch_model*.bin files. Please remove them before merging."
|
|
model.save_pretrained(input_base_path, state_dict=merged_state_dict, safe_serialization=False)
|
|
# tokenizer.save_pretrained(input_base_path)
|
|
config.save_pretrained(input_base_path)
|
|
print(f"Model saved to {input_base_path}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--input_dir",
|
|
help="Location of LLaMA weights, which contains tokenizer.model and model folders",
|
|
)
|
|
parser.add_argument("--tp_size", help="Tensor model parallel size.", default=2, type=int)
|
|
parser.add_argument("--do_split", help="Split the model into shards.", action="store_true", default=False)
|
|
parser.add_argument("--merge_avg", help="Merge the model shards using average.", action="store_true", default=False)
|
|
parser.add_argument("--use_ds_weight", default=False, action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
if os.path.exists(args.input_dir):
|
|
input_dirs = [args.input_dir]
|
|
else:
|
|
input_dirs = list(glob(args.input_dir))
|
|
print(input_dirs)
|
|
for _dir in input_dirs:
|
|
if args.do_split:
|
|
write_model(
|
|
input_base_path=_dir,
|
|
tp_size=args.tp_size,
|
|
)
|
|
else:
|
|
merge_model(
|
|
input_base_path=_dir,
|
|
tp_size=args.tp_size,
|
|
merge_avg=args.merge_avg,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|