sess_megatron.py

importsys
importos
importpathlib
importtorch
importtraceback
importnumpyasnp
fromtypingimportList, Tuple
frommegatron_miniimportget_args
frommegatron_mini.initializeimportinitialize_megatron
frommegatron_mini.modelimportLLaMAModel
frommegatron_mini.utilsimportget_model_for_infer, Tokenizer


defprint_rank_0(message):
"""If distributed is initialized, print only on rank 0."""
iftorch.distributed.is_initialized():
iftorch.distributed.get_rank() ==0:
print(message, flush=True, file=sys.stderr)
else:
print(message, flush=True, file=sys.stderr)


defadd_code_generation_args(parser):
"""Code generation arguments."""
group=parser.add_argument_group(title="code generation")

group.add_argument(
"--padded_vocab_size",
type=int,
default=40000,
help="Start id for whitespace encoding",
 )
group.add_argument("--model_dir", type=str, default="")
group.add_argument("--model_name", type=str, default="aix3-7b-base")

returnparser

classPredictor(object):

def__init__(self, args):

self.args=args
self.checkpoint_head_hash: str=""
self.np_rand=np.random.RandomState(seed=1414)

# build predictor
self.tokenizer=self.create_tokenizer()

self.dtype=torch.float32
ifself.args.bf16:
self.dtype=torch.bfloat16
elifself.args.fp16:
self.dtype=torch.half

self.predictor=self.create_predictor()

iftorch.distributed.is_initialized():
torch.distributed.barrier()

@staticmethod
defmodel_provider(pre_process=True, post_process=True):
"""Build the model."""

print_rank_0("Building Codemodel ...")
model=LLaMAModel(parallel_output=False)
returnmodel

@staticmethod
defpad_batch(tokens_id, max_seq_len=2048):
"""
 pad_batch was used by syncing token_ids
 """
tokens_id=np.reshape(tokens_id, [1, -1])
context_length=tokens_id.shape[-1]

assertcontext_length<=max_seq_len, f"{context_length}, {max_seq_len}"

ifcontext_length<max_seq_len:
tokens_id=np.concatenate([tokens_id, np.zeros(shape=[1, max_seq_len-context_length], dtype=tokens_id.dtype)], axis=-1)
returntokens_id.astype(np.int64), np.array([context_length], dtype=np.int64)


@staticmethod
defsync_type_info(sess_id: int) ->int:
input_info=np.array([sess_id], dtype=np.int64)
input_info_tensor=torch.tensor(input_info, dtype=torch.int64, device='cuda')
torch.distributed.broadcast(
input_info_tensor,
0,
 )
sess_id=input_info_tensor[0].item()
returnsess_id

@staticmethod
defsync_obj_info(model_dir: str) ->str:

tmp_list= [model_dir]
torch.distributed.broadcast_object_list(
tmp_list,
0,
 )
returntmp_list[0]

defcreate_predictor(self):

model_dir=self.args.model_dir

assertself.args.num_attention_heads%self.args.tensor_model_parallel_size==0
assertself.args.hidden_size%self.args.num_attention_heads==0

model=get_model_for_infer(self.model_provider)
print_rank_0("Loading state dict ...")
_=self.load_checkpoint(model, model_dir)

assertlen(model) ==1, "Above condition should have caught this"

model=model[0]
model.eval()

ifself.args.bf16orself.args.fp16 :
print_rank_0(f" > converting model to {'bf16'ifself.args.bf16else'fp16'} ...")
model.to(self.dtype)

print_rank_0(f" > moving model to GPU ...")
model.cuda(torch.cuda.current_device())

returnmodel

defcreate_tokenizer(self):
assertos.path.exists(os.path.join(self.args.model_dir, "tokenizer.model"))
tokenizer=Tokenizer(model_path=os.path.join(self.args.model_dir, "tokenizer.model"))
returntokenizer

defload_checkpoint(self, model: List[LLaMAModel], path):

assertisinstance(model, list)

ifnot (pathisnotNoneandos.path.exists(path)):
raiseValueError

iteration=0
ifself.args.tensor_model_parallel_size==1andself.args.rank<self.args.tensor_model_parallel_size:
checkpoint_name=os.path.join(path, f"{self.args.model_name}.pt")
assertos.path.isfile(checkpoint_name)
elifself.args.rank<self.args.tensor_model_parallel_size:
checkpoints=sorted(pathlib.Path(path).glob(f"{self.args.model_name}_states_*.pt"))
assertlen(checkpoints) ==self.args.tensor_model_parallel_size
checkpoint_name=checkpoints[self.args.rank]
else:
raiseValueError

# Load the checkpoint.
print(f"rank_{self.args.rank} load: {checkpoint_name}", flush=True, file=sys.stderr)
state_dict=torch.load(checkpoint_name, map_location="cpu")

# Set iteration.
iteration=state_dict.get("iteration", 0)

if"model"instate_dict:
state_dict=state_dict["model"]
if"module"instate_dict:
state_dict=state_dict["module"]

# Model.
model[0].load_state_dict(state_dict, strict=True)

print_rank_0(
f"successfully loaded checkpoint from {path} "
f"at iteration {iteration}"
 )

returniteration

defpredict_batch(self, data):

common_len=int(data[1].item())

withtorch.no_grad():
tokens_ids=data[0].clone().detach().cuda()

logits=self.predictor(
tokens=tokens_ids, # shape: [bsz, 1024]
start_pos=common_len,
 )
logits=logits[:, -1].view(1, -1).contiguous()
probs=torch.softmax(logits, dim=-1).cpu().numpy()

return [np.squeeze(probs)]

defpredict(self, token_ids: List[int], common_len: int) ->Tuple[List[int], List[float]]:
iftorch.distributed.is_initialized():
torch.distributed.barrier()
try:
common_len_nda=np.array([common_len]).astype("int64")
token_ids_nda=np.array([token_ids], dtype=np.int64)

max_pad_len=max(token_ids_nda.shape[-1], 128)
max_pad_len=self.sync_type_info(max_pad_len)
token_ids_nda, tokens_id_len=self.pad_batch(token_ids_nda, max_seq_len=max_pad_len)

context_tensor=torch.tensor(token_ids_nda, dtype=torch.int64, device='cuda')
context_tensor_length=torch.tensor(tokens_id_len, dtype=torch.int64, device='cuda')
context_common_len=torch.tensor(common_len_nda, dtype=torch.int64, device='cuda')

torch.distributed.broadcast(
context_tensor,
0,
 )
torch.distributed.broadcast(
context_tensor_length,
0,
 )
torch.distributed.broadcast(
context_common_len,
0,
 )

tokens_id_len=context_tensor_length.min().item()
batch= [context_tensor[:, :tokens_id_len], context_common_len]

out=self.predict_batch(batch)

# shape: [bsz, vocab_size] => [vocab_size]
out=out[0]

predict_id=np.argmax(out)
return [int(predict_id)], [out[predict_id]]
exceptExceptionase:
traceback.print_exc(file=sys.stderr)
raiseRuntimeError(e)


classTestInference:
def__init__(self) ->None:
aix_config= {
"num_layers": 32, "hidden_size": 4096, "num_attention_heads": 32,
"max_position_embeddings": 32768, "fp16": False, "bf16": True, 
"rope_theta": 256000, "inner_hidden_dim": 14464, "padded_vocab_size": 49152, 
"seq_length": 4096, "micro_batch_size": 1, "use_flash_attn": True,
"use_cpu_initialization": True, "attention_head_type": "groupedquery"
 }
initialize_megatron(
extra_args_provider=add_code_generation_args,
aix_config=aix_config
 )
args=get_args()

self.sess=Predictor(args=args)
self.end_token_set=self.sess.tokenizer.end_token_set

defrun_infer(self, code_string: str, max_new_tokens: int=256, later_code: str="", file_path: str="") ->None:
tokens=self.sess.tokenizer.encode(
code_string=code_string, later_code=later_code, file_path=file_path
 )

iflen(tokens) ==0:
returnself.sess.sync_obj_info("")

predict_list= []
common_len=0
whileTrue:

iftorch.distributed.get_rank() ==0:

output_vals=self.sess.predict(
np.array([tokens], dtype='int32'),
np.array([common_len], dtype='int32')
 )

predict_list.append(output_vals[0][0])

iflen(predict_list) >=max_new_tokensorpredict_list[-1] inself.end_token_set:
terminate_runs=1
else:
terminate_runs=0

common_len+=len(tokens)
tokens=predict_list[-1:]

else:
tokens= [0] *4
output_vals=self.sess.predict([], [], input_vals=[
np.array([tokens], dtype='int32'),
np.array([0], dtype='int32')
 ])
predict_list.append(0)
terminate_runs=0


ifself.sess.sync_type_info(terminate_runs) >0:
break
returnself.sess.sync_obj_info(self.sess.tokenizer.decode(predict_list))

if__name__=="__main__":
infer=TestInference()
res=infer.run_infer(
code_string="""# 快速排序算法""",
later_code="\n",
file_path="test.py"
 )

print(res)