mirror of https://github.com/THUDM/CodeGeeX.git
Refactor
Stanislas0
2 years ago
parent
31c3bf351d
commit
a866d8db20
@ -0,0 +1,145 @@
|
|||||||
|
"""Get model parallel partitions."""
|
||||||
|
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
import random
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch
|
||||||
|
|
||||||
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
|
||||||
|
os.path.pardir)))
|
||||||
|
|
||||||
|
from codegeex.megatron import get_args
|
||||||
|
from codegeex.megatron.model import CodeGeeXModel
|
||||||
|
from codegeex.megatron.initialize import initialize_megatron
|
||||||
|
from codegeex.megatron.checkpointing import ensure_directory_exists
|
||||||
|
|
||||||
|
|
||||||
|
def get_change_ckpt_args(parser):
|
||||||
|
"""Provide extra arguments required for merging."""
|
||||||
|
group = parser.add_argument_group(title='Mindspore to megatron')
|
||||||
|
group.add_argument(
|
||||||
|
'--load-ckpt-path',
|
||||||
|
type=str,
|
||||||
|
required=True,
|
||||||
|
help='path to load ".pt" checkpoint.',
|
||||||
|
)
|
||||||
|
group.add_argument(
|
||||||
|
'--save-ckpt-path',
|
||||||
|
type=str,
|
||||||
|
required=True,
|
||||||
|
help='dir to save converted checkpoints.',
|
||||||
|
)
|
||||||
|
group.add_argument(
|
||||||
|
'--target-tensor-model-parallel-size',
|
||||||
|
type=int,
|
||||||
|
default=2,
|
||||||
|
help='target tensor model parallel size',
|
||||||
|
)
|
||||||
|
|
||||||
|
return parser
|
||||||
|
|
||||||
|
|
||||||
|
def get_element_from_dict_by_path(d, path):
|
||||||
|
"""
|
||||||
|
Get element from dictionary by path. If element is not present, recursively add empty dictionaries.
|
||||||
|
Args:
|
||||||
|
d (dict): the dictionary to get the element from
|
||||||
|
path (list): the path to the element which is delimited by "."
|
||||||
|
"""
|
||||||
|
path = path.split(".")
|
||||||
|
for k in path:
|
||||||
|
if k not in d:
|
||||||
|
d[k] = {}
|
||||||
|
d = d[k]
|
||||||
|
return d
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
os.environ["MASTER_ADDR"] = "localhost"
|
||||||
|
os.environ["MASTER_PORT"] = str(random.randint(10000, 20000))
|
||||||
|
|
||||||
|
initialize_megatron(
|
||||||
|
extra_args_provider=get_change_ckpt_args,
|
||||||
|
args_defaults={
|
||||||
|
"tokenizer_type": "GPT2BPETokenizer",
|
||||||
|
"no_load_rng" : True,
|
||||||
|
"no_load_optim" : True,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
args = get_args()
|
||||||
|
print(f"Load ckpt from {args.load_ckpt_path}...")
|
||||||
|
state_dict = torch.load(args.load_ckpt_path, map_location="cpu")
|
||||||
|
|
||||||
|
print(f"Spliting ckpt into {args.target_tensor_model_parallel_size} parts...")
|
||||||
|
output_state_dict = []
|
||||||
|
for i in range(args.target_tensor_model_parallel_size):
|
||||||
|
output_state_dict.append({})
|
||||||
|
|
||||||
|
print("Converting Embedding layers...")
|
||||||
|
word_embeddings = state_dict['module']['language_model']['embedding']['word_embeddings']['weight']
|
||||||
|
position_embeddings = state_dict['module']['language_model']['embedding']['position_embeddings']['weight']
|
||||||
|
out_word_embeddings = torch.chunk(word_embeddings, args.target_tensor_model_parallel_size, dim=0)
|
||||||
|
|
||||||
|
for i in range(args.target_tensor_model_parallel_size):
|
||||||
|
pos_emb_dict = get_element_from_dict_by_path(
|
||||||
|
output_state_dict[i], "module.language_model.embedding.position_embeddings"
|
||||||
|
)
|
||||||
|
pos_emb_dict["weight"] = position_embeddings
|
||||||
|
|
||||||
|
word_emb_dict = get_element_from_dict_by_path(
|
||||||
|
output_state_dict[i], "module.language_model.embedding.word_embeddings"
|
||||||
|
)
|
||||||
|
word_emb_dict["weight"] = out_word_embeddings[i]
|
||||||
|
|
||||||
|
print("Converting QueryEmbedding layers...")
|
||||||
|
query_embeddings = state_dict['module']['language_model']['topQueryEmbedding']['top_query_embeddings']['weight']
|
||||||
|
out_query_embeddings = torch.chunk(query_embeddings, args.target_tensor_model_parallel_size, dim=0)
|
||||||
|
|
||||||
|
for i in range(args.target_tensor_model_parallel_size):
|
||||||
|
query_emb_dict = get_element_from_dict_by_path(
|
||||||
|
output_state_dict[i], "module.language_model.topQueryEmbedding.top_query_embeddings"
|
||||||
|
)
|
||||||
|
query_emb_dict["weight"] = out_query_embeddings[i]
|
||||||
|
|
||||||
|
print("Converting Transformer layers...")
|
||||||
|
for layer_name in state_dict['module']['language_model']['transformer'].keys():
|
||||||
|
params = state_dict['module']['language_model']['transformer'][layer_name]
|
||||||
|
if "layernorm" in layer_name:
|
||||||
|
pass
|
||||||
|
elif "attention" in layer_name and "weight" in layer_name:
|
||||||
|
if "dense" in layer_name:
|
||||||
|
params = torch.chunk(params, args.target_tensor_model_parallel_size, dim=1)
|
||||||
|
else:
|
||||||
|
params = torch.chunk(params, args.target_tensor_model_parallel_size, dim=0)
|
||||||
|
elif "weight" in layer_name and "dense" in layer_name:
|
||||||
|
if "h_to_4h" in layer_name:
|
||||||
|
params = torch.chunk(params, args.target_tensor_model_parallel_size, dim=0)
|
||||||
|
else:
|
||||||
|
params = torch.chunk(params, args.target_tensor_model_parallel_size, dim=1)
|
||||||
|
elif "bias" in layer_name:
|
||||||
|
if "dense" not in layer_name or "mlp" in layer_name:
|
||||||
|
if "4h_to_h" in layer_name:
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
params = torch.chunk(params, args.target_tensor_model_parallel_size, dim=0)
|
||||||
|
|
||||||
|
for i in range(args.target_tensor_model_parallel_size):
|
||||||
|
params_dict = get_element_from_dict_by_path(output_state_dict[i], "module.language_model.transformer")
|
||||||
|
if type(params) is tuple:
|
||||||
|
params_dict[layer_name] = params[i]
|
||||||
|
else:
|
||||||
|
params_dict[layer_name] = params
|
||||||
|
|
||||||
|
os.makedirs(args.save_ckpt_path, exist_ok=True)
|
||||||
|
for rank in range(args.target_tensor_model_parallel_size):
|
||||||
|
save_ckpt_path = os.path.join(args.save_ckpt_path, f"mp_rank_{rank:02d}_model_states.pt")
|
||||||
|
torch.save(output_state_dict[rank], save_ckpt_path)
|
||||||
|
print(f"Converted checkpoint saved in {save_ckpt_path}.")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
main()
|
||||||
@ -1,201 +0,0 @@
|
|||||||
|
|
||||||
import os
|
|
||||||
import copy
|
|
||||||
import time
|
|
||||||
import torch
|
|
||||||
import random
|
|
||||||
import argparse
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
from codegeex.torch.inference import get_token_stream
|
|
||||||
from codegeex.torch import CodeGeeXModel
|
|
||||||
from codegeex.tokenizer import CodeGeeXTokenizer
|
|
||||||
from codegeex.quantization import quantize
|
|
||||||
|
|
||||||
|
|
||||||
def model_provider(args):
|
|
||||||
"""Build the model."""
|
|
||||||
|
|
||||||
model = CodeGeeXModel(
|
|
||||||
args.hidden_size,
|
|
||||||
args.num_layers,
|
|
||||||
args.num_attention_heads,
|
|
||||||
args.padded_vocab_size,
|
|
||||||
args.max_position_embeddings
|
|
||||||
)
|
|
||||||
|
|
||||||
return model
|
|
||||||
|
|
||||||
|
|
||||||
def add_code_generation_args(parser):
|
|
||||||
group = parser.add_argument_group(title="code generation")
|
|
||||||
group.add_argument(
|
|
||||||
"--num-layers",
|
|
||||||
type=int,
|
|
||||||
default=39,
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--hidden-size",
|
|
||||||
type=int,
|
|
||||||
default=5120,
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--num-attention-heads",
|
|
||||||
type=int,
|
|
||||||
default=40,
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--padded-vocab-size",
|
|
||||||
type=int,
|
|
||||||
default=52224,
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--max-position-embeddings",
|
|
||||||
type=int,
|
|
||||||
default=2048,
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--temperature",
|
|
||||||
type=float,
|
|
||||||
default=1.0,
|
|
||||||
help="Sampling temperature.",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--greedy",
|
|
||||||
action="store_true",
|
|
||||||
default=False,
|
|
||||||
help="Use greedy sampling.",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--top-p",
|
|
||||||
type=float,
|
|
||||||
default=0.0,
|
|
||||||
help="Top p sampling.",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--top-k",
|
|
||||||
type=int,
|
|
||||||
default=0,
|
|
||||||
help="Top k sampling.",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--out-seq-length",
|
|
||||||
type=int,
|
|
||||||
default=2048,
|
|
||||||
help="Size of the output generated text.",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--prompt-file",
|
|
||||||
type=str,
|
|
||||||
default="./test_prompt.txt",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--tokenizer-path",
|
|
||||||
type=str,
|
|
||||||
default="./tokenizer",
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--load",
|
|
||||||
type=str,
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--state-dict-path",
|
|
||||||
type=str,
|
|
||||||
)
|
|
||||||
group.add_argument(
|
|
||||||
"--micro-batch-size",
|
|
||||||
type=int,
|
|
||||||
default=1,
|
|
||||||
)
|
|
||||||
|
|
||||||
return parser
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
parser = argparse.ArgumentParser()
|
|
||||||
parser = add_code_generation_args(parser)
|
|
||||||
args, _ = parser.parse_known_args()
|
|
||||||
|
|
||||||
print("Building CodeGeeX model ...")
|
|
||||||
model = model_provider(args)
|
|
||||||
|
|
||||||
print("Loading tokenizer ...")
|
|
||||||
tokenizer = CodeGeeXTokenizer(
|
|
||||||
tokenizer_path=args.tokenizer_path,
|
|
||||||
mode="codegeex-13b")
|
|
||||||
|
|
||||||
print("Loading state dict ...")
|
|
||||||
state_dict = torch.load(args.load, map_location="cpu")
|
|
||||||
state_dict = state_dict["module"]
|
|
||||||
|
|
||||||
print("Building CodeGeeX model ...")
|
|
||||||
model = model_provider(args)
|
|
||||||
model.load_state_dict(state_dict)
|
|
||||||
model.eval()
|
|
||||||
model.half()
|
|
||||||
model = quantize(model, weight_bit_width=8)
|
|
||||||
model.cuda()
|
|
||||||
|
|
||||||
with open(args.prompt_file, "r") as f:
|
|
||||||
prompt = f.readlines()
|
|
||||||
prompt = "".join(prompt)
|
|
||||||
|
|
||||||
times = {}
|
|
||||||
out_seq_lengths = [args.out_seq_length]
|
|
||||||
micro_batch_size = args.micro_batch_size
|
|
||||||
seq_length = args.max_position_embeddings
|
|
||||||
for out_seq_length in out_seq_lengths:
|
|
||||||
print(f"Generating with out_seq_len {out_seq_length}...")
|
|
||||||
|
|
||||||
times[out_seq_length] = []
|
|
||||||
for prompt in [prompt]:
|
|
||||||
t0 = time.perf_counter()
|
|
||||||
tokens = tokenizer.encode_code(prompt)
|
|
||||||
print(tokens)
|
|
||||||
print("Current prompt:")
|
|
||||||
print(prompt)
|
|
||||||
n_token_prompt = len(tokens)
|
|
||||||
print("N_token_prompt:", n_token_prompt)
|
|
||||||
token_stream = get_token_stream(
|
|
||||||
model,
|
|
||||||
tokenizer,
|
|
||||||
seq_length,
|
|
||||||
out_seq_length,
|
|
||||||
[copy.deepcopy(tokens) for _ in range(micro_batch_size)],
|
|
||||||
micro_batch_size=micro_batch_size,
|
|
||||||
topk=args.top_k,
|
|
||||||
topp=args.top_p,
|
|
||||||
temperature=args.temperature,
|
|
||||||
greedy=args.greedy,
|
|
||||||
)
|
|
||||||
is_finished = [False for _ in range(micro_batch_size)]
|
|
||||||
for i, generated in enumerate(token_stream):
|
|
||||||
generated_tokens = generated[0]
|
|
||||||
for j in range(micro_batch_size):
|
|
||||||
if is_finished[j]:
|
|
||||||
continue
|
|
||||||
if generated_tokens[j].cpu().numpy()[-1] == tokenizer.eos_token_id or len(
|
|
||||||
generated_tokens[j]) >= out_seq_length:
|
|
||||||
is_finished[j] = True
|
|
||||||
generated_tokens_ = generated_tokens[j].cpu().numpy().tolist()
|
|
||||||
generated_code = tokenizer.decode_code(generated_tokens_[n_token_prompt:])
|
|
||||||
generated_code = "".join(generated_code)
|
|
||||||
t1 = time.perf_counter()
|
|
||||||
print("Total generation time:", t1 - t0, "# Tokens:", len(generated_tokens_) - n_token_prompt)
|
|
||||||
print(f"{(t1 - t0) / (len(generated_tokens_) - n_token_prompt)}s/token")
|
|
||||||
times[out_seq_length].append(t1 - t0)
|
|
||||||
print("================================= Generated code:")
|
|
||||||
print(generated_code)
|
|
||||||
|
|
||||||
if all(is_finished):
|
|
||||||
break
|
|
||||||
|
|
||||||
print(times)
|
|
||||||
for out_seq_length in times.keys():
|
|
||||||
print(out_seq_length, np.mean(times[out_seq_length]))
|
|
||||||
|
|
||||||
print("Generation finished.")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
Loading…
Reference in New Issue