mirror of https://github.com/InternLM/InternLM
137 lines
4.7 KiB
Python
137 lines
4.7 KiB
Python
# Copyright (c) InternLM. All rights reserved.
|
|
import argparse
|
|
import json
|
|
import os
|
|
|
|
import torch
|
|
from einops import rearrange
|
|
from tqdm import tqdm
|
|
from transformers import AutoConfig, LlamaConfig, LlamaTokenizer
|
|
|
|
|
|
def save_conifg(config, tgt):
|
|
config_dict = config.to_dict()
|
|
unnecessary_keys = [
|
|
"_name_or_path",
|
|
"auto_map",
|
|
"transformers_version",
|
|
"model_type",
|
|
"architectures",
|
|
"tokenizer_class",
|
|
"attn_implementation",
|
|
]
|
|
for k in unnecessary_keys:
|
|
config_dict.pop(k, None)
|
|
config_dict["attention_bias"] = config_dict.pop("bias")
|
|
config_dict["architectures"] = ["LlamaForCausalLM"]
|
|
llama_config = LlamaConfig(**config_dict)
|
|
llama_config.save_pretrained(tgt)
|
|
|
|
|
|
def convert(src, tgt):
|
|
"""Convert InternLM2 huggingface checkpoints to Llama-style."""
|
|
|
|
print("Convert InternLM2 huggingface checkpoints to Llama...")
|
|
|
|
config = AutoConfig.from_pretrained(src, trust_remote_code=True)
|
|
assert not config.bias, "Cannot convert InternLM Model with bias to LLaMA."
|
|
|
|
head_dim = config.hidden_size // config.num_attention_heads
|
|
num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
|
|
|
|
# load index json file
|
|
index_file = os.path.join(src, "pytorch_model.bin.index.json")
|
|
if os.path.exists(index_file):
|
|
with open(index_file) as fp:
|
|
index_dict = json.load(fp)
|
|
index_dict["weight_map"] = {}
|
|
else:
|
|
index_dict = None
|
|
|
|
os.makedirs(tgt, exist_ok=True)
|
|
for filename in tqdm(os.listdir(src)):
|
|
if not filename.endswith(".bin"):
|
|
continue
|
|
states = torch.load(os.path.join(src, filename))
|
|
llama_states = {}
|
|
for k, v in states.copy().items():
|
|
if "wqkv" in k:
|
|
v = rearrange(
|
|
v,
|
|
"(h gs d) dim -> h gs d dim",
|
|
gs=2 + num_key_value_groups,
|
|
d=head_dim,
|
|
)
|
|
wq, wk, wv = torch.split(v, [num_key_value_groups, 1, 1], dim=1)
|
|
wq = rearrange(wq, "h gs d dim -> (h gs d) dim")
|
|
wk = rearrange(wk, "h gs d dim -> (h gs d) dim")
|
|
wv = rearrange(wv, "h gs d dim -> (h gs d) dim")
|
|
_prefix = k.split("attention")[0]
|
|
wq_key = _prefix + "self_attn.q_proj.weight"
|
|
wk_key = _prefix + "self_attn.k_proj.weight"
|
|
wv_key = _prefix + "self_attn.v_proj.weight"
|
|
llama_states[wq_key] = wq.clone()
|
|
llama_states[wk_key] = wk.clone()
|
|
llama_states[wv_key] = wv.clone()
|
|
|
|
elif "attention.wo" in k:
|
|
new_k = k.replace("attention.wo", "self_attn.o_proj")
|
|
llama_states[new_k] = v
|
|
elif "feed_forward.w1" in k:
|
|
new_k = k.replace("feed_forward.w1", "mlp.gate_proj")
|
|
llama_states[new_k] = v
|
|
elif "feed_forward.w2" in k:
|
|
new_k = k.replace("feed_forward.w2", "mlp.down_proj")
|
|
llama_states[new_k] = v
|
|
elif "feed_forward.w3" in k:
|
|
new_k = k.replace("feed_forward.w3", "mlp.up_proj")
|
|
llama_states[new_k] = v
|
|
elif "attention_norm" in k:
|
|
new_k = k.replace("attention_norm", "input_layernorm")
|
|
llama_states[new_k] = v
|
|
elif "ffn_norm" in k:
|
|
new_k = k.replace("ffn_norm", "post_attention_layernorm")
|
|
llama_states[new_k] = v
|
|
elif "tok_embeddings" in k:
|
|
llama_states["model.embed_tokens.weight"] = v
|
|
elif "output" in k:
|
|
llama_states["lm_head.weight"] = v
|
|
else:
|
|
llama_states[k] = v
|
|
|
|
if index_dict is not None:
|
|
for k in llama_states:
|
|
index_dict["weight_map"][k] = filename
|
|
print(f"Saving to {os.path.join(tgt, filename)}...", flush=True)
|
|
torch.save(llama_states, os.path.join(tgt, filename))
|
|
del states
|
|
|
|
print("Saving config and tokenizer...")
|
|
# index.json
|
|
if index_dict is not None:
|
|
with open(os.path.join(tgt, "pytorch_model.bin.index.json"), "w") as fp:
|
|
json.dump(index_dict, fp, indent=2)
|
|
# tokenizer
|
|
tokenizer = LlamaTokenizer.from_pretrained(src)
|
|
tokenizer.init_kwargs.pop("auto_map", None)
|
|
tokenizer.save_pretrained(tgt)
|
|
# config
|
|
save_conifg(config, tgt)
|
|
print("Done!")
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--src", type=str, help="Input folder")
|
|
parser.add_argument("--tgt", type=str, help="Output folder")
|
|
|
|
args = parser.parse_args()
|
|
|
|
return args
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
|
|
convert(args.src, args.tgt)
|