def main():
parser = ArgumentParser()
parser.add_argument("--model_file",
type=str,
default="",
required=True,
help="Pass path to model's .nemo file")
parser.add_argument("--prompt",
type=str,
default="",
required=True,
help="Prompt for the model (a text to complete)")
parser.add_argument("--tokens_to_generate",
type=int,
default="16",
required=False,
help="How many tokens to add to prompt")
parser.add_argument(
"--tensor_model_parallel_size",
type=int,
default=1,
required=True,
)
args = parser.parse_args()
torch.set_grad_enabled(False)
# trainer required for restoring model parallel models
trainer = Trainer(plugins=NLPDDPPlugin(),
devices=args.tensor_model_parallel_size,
precision=16,
accelerator='gpu')
app_state = AppState()
if args.tensor_model_parallel_size > 1:
app_state.model_parallel_size = args.tensor_model_parallel_size
app_state.model_parallel_rank = compute_model_parallel_rank(
trainer.local_rank, app_state.model_parallel_size)
model = MegatronT5Model.restore_from(restore_path=args.model_file,
trainer=trainer)
model.freeze()
request = {
"prompt": args.prompt,
"tokens_to_generate": args.tokens_to_generate,
}
dataset = T5RequestDataset(request, model.tokenizer)
request_dl = DataLoader(dataset)
response = trainer.predict(model, request_dl)
print("***************************")
print(response)
print("***************************")
def score(self, src: List[str], cand: List[str],
ref: List[str]) -> COMETResult:
data = {"src": src, "mt": cand, "ref": ref}
data = [dict(zip(data, t)) for t in zip(*data.values())]
dataloader = DataLoader(
dataset=data,
batch_size=16,
collate_fn=lambda x: self.model.prepare_sample(x, inference=True),
num_workers=4,
)
cuda = 1 if torch.cuda.is_available() else 0
trainer = Trainer(gpus=cuda, deterministic=True, logger=False)
predictions = trainer.predict(self.model,
dataloaders=dataloader,
return_predictions=True)
scores = torch.cat(predictions, dim=0).tolist()
return COMETResult(
sum(scores) / len(scores), scores, src, cand, ref, self.name,
self.modelname)
def main(cfg) -> None:
# trainer required for restoring model parallel models
trainer = Trainer(plugins=NLPDDPPlugin(), **cfg.trainer)
assert (
cfg.trainer.devices *
cfg.trainer.num_nodes == cfg.tensor_model_parallel_size *
cfg.pipeline_model_parallel_size
), "devices * num_nodes should equal tensor_model_parallel_size * pipeline_model_parallel_size"
# Load prompt tuned model, virtual_prompt_model_file must be provided in config
if cfg.get('virtual_prompt_model_file', None) is not None:
# Update frozen GPT model path in case it has changed
prompt_learning_cfg = MegatronGPTPromptLearningModel.restore_from(
cfg.virtual_prompt_model_file, trainer=trainer, return_config=True)
with open_dict(prompt_learning_cfg):
prompt_learning_cfg.language_model_path = cfg.gpt_model_file
# Now load prompt learning model with frozen gpt model base
model = MegatronGPTPromptLearningModel.restore_from(
restore_path=cfg.virtual_prompt_model_file,
trainer=trainer,
override_config_path=prompt_learning_cfg)
# Or load regular GPT model
elif cfg.gpt_model_file:
model = MegatronGPTModel.restore_from(restore_path=cfg.gpt_model_file,
trainer=trainer)
elif cfg.checkpoint_dir:
app_state = AppState()
if cfg.tensor_model_parallel_size > 1 or cfg.pipeline_model_parallel_size > 1:
app_state.model_parallel_size = cfg.tensor_model_parallel_size * cfg.pipeline_model_parallel_size
(
app_state.tensor_model_parallel_rank,
app_state.pipeline_model_parallel_rank,
app_state.model_parallel_size,
app_state.data_parallel_size,
app_state.pipeline_model_parallel_split_rank,
) = fake_initialize_model_parallel(
world_size=app_state.model_parallel_size,
rank=trainer.global_rank,
tensor_model_parallel_size_=cfg.tensor_model_parallel_size,
pipeline_model_parallel_size_=cfg.pipeline_model_parallel_size,
pipeline_model_parallel_split_rank_=cfg.
pipeline_model_parallel_split_rank,
)
checkpoint_path = inject_model_parallel_rank(
os.path.join(cfg.checkpoint_dir, cfg.checkpoint_name))
model = MegatronGPTModel.load_from_checkpoint(
checkpoint_path, hparams_file=cfg.hparams_file, trainer=trainer)
else:
raise ValueError("need at least a nemo file or checkpoint dir")
model.freeze()
# Have to turn off activations_checkpoint_method for inference
try:
model.model.language_model.encoder.activations_checkpoint_method = None
except AttributeError:
pass
try:
model.frozen_model.language_model.encoder.activations_checkpoint_method = None
except AttributeError:
pass
length_params: LengthParam = {
"max_length": cfg.inference.tokens_to_generate,
"min_length": cfg.inference.min_tokens_to_generate,
}
sampling_params: SamplingParam = {
"use_greedy": cfg.inference.greedy,
"temperature": cfg.inference.temperature,
"top_k": cfg.inference.top_k,
"top_p": cfg.inference.top_p,
"repetition_penalty": cfg.inference.repetition_penalty,
"add_BOS": cfg.inference.add_BOS,
"all_probs": cfg.inference.all_probs,
"compute_logprob": cfg.inference.compute_logprob,
}
# First method of running text generation, call model.generate method
response = model.generate(inputs=OmegaConf.to_container(cfg.prompts),
length_params=length_params,
sampling_params=sampling_params)
print("***************************")
print(response)
print("***************************")
# Second method of running text generation, call trainer.predict
collate_fn = None
if cfg.get('virtual_prompt_model', False):
collate_fn = lambda x: list(x)
ds = RequestDataSet(OmegaConf.to_container(cfg.prompts))
request_dl = DataLoader(dataset=ds, collate_fn=collate_fn, batch_size=2)
config = OmegaConf.to_container(cfg.inference)
model.set_inference_config(config)
response = trainer.predict(model, request_dl)
print("***************************")
print(response)
print("***************************")
# Third method of running text generation, use inference server
if cfg.server:
if parallel_state.is_pipeline_first_stage(
) and parallel_state.get_tensor_model_parallel_rank() == 0:
server = MegatronServer(model.cuda())
server.run("0.0.0.0", port=cfg.port)
while True:
choice = torch.cuda.LongTensor(1)
torch.distributed.broadcast(choice, 0)
if choice[0].item() == 0:
generate(model.cuda())
def main():
parser = ArgumentParser()
parser.add_argument("--use_soft_prompts", action="store_true", help="Use model's existing soft prompts")
parser.add_argument("--model_file", type=str, default="", required=True, help="Pass path to model's .nemo file")
parser.add_argument(
"--path_to_file", type=str, default="", required=False, help="Path to file with prompts (a text to complete)"
)
parser.add_argument(
"--prompt", type=str, default="", required=False, help="Prompt for the model (a text to complete)"
)
parser.add_argument(
"--prompt_tag", type=str, default="", required=False, help="Prompt tag string for task specific soft prompt"
)
parser.add_argument(
"--tokens_to_generate", type=int, default="1", required=False, help="How many tokens to add to prompt"
)
parser.add_argument(
"--stop_after_sentence",
type=bool,
default="True",
required=False,
help="True/False: whether to stop after full sentence has been generated.",
)
parser.add_argument(
"--tensor_model_parallel_size", type=int, default=1, required=False,
)
parser.add_argument("--precision", default=16, help="PyTorch Lightning Trainer precision flag")
parser.add_argument("--batch_size", default=1, required=False, help="Evaluation batch_size")
parser.add_argument(
"--compute_logprobs", type=bool, default=False, required=False, help="Method for logprobs computation"
)
args = parser.parse_args()
# cast precision to int if 32 or 16
if args.precision in ["32", "16"]:
args.precision = int(float(args.precision))
# trainer required for restoring model parallel models
trainer = Trainer(plugins=NLPDDPPlugin(), gpus=args.tensor_model_parallel_size, precision=args.precision)
app_state = AppState()
if args.tensor_model_parallel_size is not None and args.tensor_model_parallel_size > 1:
app_state.model_parallel_size = args.tensor_model_parallel_size
app_state.model_parallel_rank = compute_model_parallel_rank(trainer.local_rank, app_state.model_parallel_size)
model = MegatronGPTModel.restore_from(restore_path=args.model_file, trainer=trainer)
model.freeze()
def pad_collate(batch):
tokens, tokens_to_generate = batch[0]['data'], batch[0]['tokens_to_generate']
compute_logprobs = batch[0]['compute_logprobs']
lens = [len(token) for token in tokens]
tokens_pad = pad_sequence(tokens, batch_first=False, padding_value=50256)
data = []
if 'prompt_tags' in batch[0]:
# Keep track of soft prompt tags
prompt_tags = batch[0]['prompt_tags']
for token, lenn, prompt_tag in zip(tokens_pad.T, lens, prompt_tags):
data.append((token, lenn, tokens_to_generate, compute_logprobs, prompt_tag))
else:
for token, lenn in zip(tokens_pad.T, lens):
data.append((token, lenn, tokens_to_generate, compute_logprobs))
return data
# defining type of request
if args.path_to_file != "":
request = []
prompts = open(args.path_to_file, 'r')
for prompt in prompts.readlines():
prompt = prompt.split('\n')[0]
if args.use_soft_prompts and model.use_soft_prompts:
prompt = json.loads(prompt)
request.append(prompt)
dataset = GPTRequestDataset(request, model.tokenizer, args.tokens_to_generate, args.compute_logprobs)
request_dl = DataLoader(dataset=pad_collate(dataset), batch_size=int(args.batch_size))
else:
if args.use_soft_prompts and model.use_soft_prompts:
request = [{'prompt_tag': args.prompt_tag, 'text': args.prompt}]
else:
request = [args.prompt]
dataset = GPTRequestDataset(request, model.tokenizer, args.tokens_to_generate, args.compute_logprobs)
request_dl = DataLoader(dataset=pad_collate(dataset), batch_size=1)
# For GPT models that have had soft prompt tuning but you don't want to use any soft prompts
if not args.use_soft_prompts and model.use_soft_prompts:
model.use_soft_prompts = False
response = trainer.predict(model, request_dl)
print("***************************")
print(response)
print("***************************")
def main():
parser = ArgumentParser()
parser.add_argument("--model_file",
type=str,
default="",
required=True,
help="Pass path to model's .nemo file")
parser.add_argument("--prompt",
type=str,
default="",
required=True,
help="Prompt for the model (a text to complete)")
parser.add_argument("--tokens_to_generate",
type=int,
default="64",
required=False,
help="How many tokens to add to prompt")
parser.add_argument(
"--stop_after_sentence",
type=bool,
default="True",
required=False,
help=
"True/False: whether to stop after full sentence has been generated.",
)
parser.add_argument(
"--tensor_model_parallel_size",
type=int,
default=1,
required=True,
)
parser.add_argument("--precision",
default=32,
help="PyTorch Lightning Trainer precision flag")
args = parser.parse_args()
# cast precision to int if 32 or 16
if args.precision in ["32", "16"]:
args.precision = int(float(args.precision))
# trainer required for restoring model parallel models
trainer = Trainer(plugins=NLPDDPPlugin(),
gpus=args.tensor_model_parallel_size,
precision=args.precision)
app_state = AppState()
if args.tensor_model_parallel_size is not None and args.tensor_model_parallel_size > 1:
app_state.model_parallel_size = args.tensor_model_parallel_size
app_state.model_parallel_rank = compute_model_parallel_rank(
trainer.local_rank, app_state.model_parallel_size)
model = MegatronGPTModel.restore_from(restore_path=args.model_file,
trainer=trainer)
model.freeze()
request = {
"prompt": args.prompt,
"tokens_to_generate": args.tokens_to_generate,
"stop_after_sentence": args.stop_after_sentence,
}
dataset = GPTRequestDataset(request, model.tokenizer)
request_dl = DataLoader(dataset)
response = trainer.predict(model, request_dl)
print("***************************")
print(response[0]['completion']['text'])
print("***************************")
logging.info(
f"Generation stopped because: {response[0]['completion']['stop reason']}"
)
def main():
parser = ArgumentParser()
parser.add_argument("--model_file",
type=str,
default="",
required=True,
help="Pass path to model's .nemo file")
parser.add_argument("--prompt",
type=str,
default="",
required=True,
help="Prompt for the model (a text to complete)")
parser.add_argument("--tokens_to_generate",
type=int,
default="16",
required=False,
help="How many tokens to add to prompt")
parser.add_argument(
"--tensor_model_parallel_size",
type=int,
default=1,
required=False,
)
parser.add_argument(
"--pipeline_model_parallel_size",
type=int,
default=1,
required=False,
)
parser.add_argument(
"--pipeline_model_parallel_split_rank",
type=int,
default=0,
required=False,
)
parser.add_argument("--precision",
default="16",
type=str,
help="PyTorch Lightning Trainer precision flag")
args = parser.parse_args()
# cast precision to int if 32 or 16
if args.precision in ["32", "16"]:
args.precision = int(float(args.precision))
# trainer required for restoring model parallel models
trainer = Trainer(
plugins=NLPDDPPlugin(),
devices=args.tensor_model_parallel_size *
args.pipeline_model_parallel_size,
accelerator='gpu',
precision=args.precision,
)
app_state = AppState()
if args.tensor_model_parallel_size > 1 or args.pipeline_model_parallel_size > 1:
app_state.model_parallel_size = args.tensor_model_parallel_size * args.pipeline_model_parallel_size
(
app_state.tensor_model_parallel_rank,
app_state.pipeline_model_parallel_rank,
app_state.model_parallel_size,
app_state.data_parallel_size,
app_state.pipeline_model_parallel_split_rank,
) = fake_initialize_model_parallel(
world_size=app_state.model_parallel_size,
rank=trainer.global_rank,
tensor_model_parallel_size_=args.tensor_model_parallel_size,
pipeline_model_parallel_size_=args.pipeline_model_parallel_size,
pipeline_model_parallel_split_rank_=args.
pipeline_model_parallel_split_rank,
)
model = MegatronT5Model.restore_from(restore_path=args.model_file,
trainer=trainer)
model.freeze()
request = {
"prompt": args.prompt,
"tokens_to_generate": args.tokens_to_generate,
}
dataset = T5RequestDataset(request, model.tokenizer)
request_dl = DataLoader(dataset)
response = trainer.predict(model, request_dl)
print("***************************")
print(response)
print("***************************")
def main():
parser = ArgumentParser()
# args for loading the model, either from .nemo file or from PTL checkpoint
parser.add_argument("--model_file",
type=str,
default="",
required=False,
help="Pass path to model's .nemo file")
parser.add_argument(
"--checkpoint_dir",
type=str,
default=None,
required=False,
help=
"If not using a .nemo file. Path to PTL checkpoints saved during training. Ex: /raid/nemo_experiments/megatron_gpt/checkpoints",
)
parser.add_argument(
"--checkpoint_name",
type=str,
default=None,
required=False,
help=
"If not using a .nemo file. Name of checkpoint to be used. Ex: megatron_gpt--val_loss=6.34-step=649-last.ckpt",
)
parser.add_argument(
"--hparams_file",
type=str,
default=None,
required=False,
help=
"If not using a .nemo file. Path to config for restoring. It's created during training and may need to be modified during restore if restore environment is different than training. Ex: /raid/nemo_experiments/megatron_gpt/hparams.yaml",
)
parser.add_argument("--tensor_model_parallel_size",
type=int,
default=1,
required=False,
help="Needed if not using a .nemo file")
parser.add_argument(
"--pipeline_model_parallel_size",
type=int,
default=1,
required=False,
help="Needed if not using a .nemo file",
)
# PTL Trainer args
parser.add_argument("--devices",
default=1,
type=int,
help="PyTorch Lightning Trainer devices flag")
parser.add_argument("--num_nodes",
default=1,
type=int,
help="PyTorch Lightning Trainer num_nodes flag")
parser.add_argument("--precision",
default=16,
help="PyTorch Lightning Trainer precision flag")
# evaluation args
parser.add_argument("--path_to_file",
type=str,
default="",
required=False,
help="Path to file with prompts (a text to complete)")
parser.add_argument("--prompt",
type=str,
default="",
required=False,
help="Prompt for the model (a text to complete)")
parser.add_argument("--use_soft_prompts",
action="store_true",
help="Use model's existing soft prompts")
parser.add_argument("--prompt_tag",
type=str,
default="",
required=False,
help="Prompt tag string for task specific soft prompt")
parser.add_argument("--tokens_to_generate",
type=int,
default="1",
required=False,
help="How many tokens to add to prompt")
parser.add_argument(
"--stop_after_sentence",
type=bool,
default="True",
required=False,
help=
"True/False: whether to stop after full sentence has been generated.",
)
parser.add_argument("--batch_size",
default=1,
type=int,
required=False,
help="Evaluation batch_size")
parser.add_argument("--compute_logprobs",
type=bool,
default=False,
required=False,
help="Method for logprobs computation")
args = parser.parse_args()
assert (
args.devices * args.num_nodes == args.tensor_model_parallel_size *
args.pipeline_model_parallel_size
), "devices * num_nodes should equal tensor_model_parallel_size * pipeline_model_parallel_size"
if args.model_file and args.checkpoint_dir:
raise ValueError(
"Only one of model_file or checkpoint_dir should be used")
# cast precision to int if 32 or 16
if args.precision in ["32", "16"]:
args.precision = int(float(args.precision))
# trainer required for restoring model parallel models
trainer = Trainer(
plugins=[NLPDDPPlugin()],
devices=args.devices,
num_nodes=args.num_nodes,
accelerator='gpu',
precision=args.precision,
)
if args.model_file:
model = MegatronGPTModel.restore_from(restore_path=args.model_file,
trainer=trainer)
elif args.checkpoint_dir:
app_state = AppState()
if args.tensor_model_parallel_size > 1 or args.pipeline_model_parallel_size > 1:
app_state.pipeline_model_parallel_size = args.pipeline_model_parallel_size
app_state.tensor_model_parallel_size = args.tensor_model_parallel_size
app_state.model_parallel_size = args.tensor_model_parallel_size * args.pipeline_model_parallel_size
(
app_state.tensor_model_parallel_rank,
app_state.pipeline_model_parallel_rank,
app_state.model_parallel_size,
_,
) = fake_initialize_model_parallel(
world_size=app_state.model_parallel_size,
rank=trainer.global_rank,
tensor_model_parallel_size_=app_state.
tensor_model_parallel_size,
pipeline_model_parallel_size_=app_state.
pipeline_model_parallel_size,
)
# inject model parallel rank
checkpoint_path = inject_model_parallel_rank(
os.path.join(args.checkpoint_dir, args.checkpoint_name))
model = MegatronGPTModel.load_from_checkpoint(
checkpoint_path, hparams_file=args.hparams_file, trainer=trainer)
model.freeze()
def pad_collate(batch):
tokens, tokens_to_generate = batch[0]['data'], batch[0][
'tokens_to_generate']
compute_logprobs = batch[0]['compute_logprobs']
lens = [len(token) for token in tokens]
tokens_pad = pad_sequence(tokens,
batch_first=False,
padding_value=50256)
data = []
if 'prompt_tags' in batch[0]:
# Keep track of soft prompt tags
prompt_tags = batch[0]['prompt_tags']
for token, lenn, prompt_tag in zip(tokens_pad.T, lens,
prompt_tags):
data.append((token, lenn, tokens_to_generate, compute_logprobs,
prompt_tag))
else:
for token, lenn in zip(tokens_pad.T, lens):
data.append(
(token, lenn, tokens_to_generate, compute_logprobs))
return data
# defining type of request
if args.path_to_file != "":
request = []
prompts = open(args.path_to_file, 'r', encoding='utf-8')
for prompt in prompts.readlines():
prompt = prompt.split('\n')[0]
if args.use_soft_prompts and model.use_soft_prompts:
prompt = json.loads(prompt)
request.append(prompt)
dataset = GPTRequestDataset(request, model.tokenizer,
args.tokens_to_generate,
args.compute_logprobs)
request_dl = DataLoader(dataset=pad_collate(dataset),
batch_size=int(args.batch_size))
else:
if args.use_soft_prompts and model.use_soft_prompts:
request = [{'prompt_tag': args.prompt_tag, 'text': args.prompt}]
else:
request = [args.prompt]
dataset = GPTRequestDataset(request, model.tokenizer,
args.tokens_to_generate,
args.compute_logprobs)
request_dl = DataLoader(dataset=pad_collate(dataset), batch_size=1)
# For GPT models that have had soft prompt tuning but you don't want to use any soft prompts
if not args.use_soft_prompts and model.use_soft_prompts:
model.use_soft_prompts = False
response = trainer.predict(model, request_dl)
print("***************************")
print(response)
print("***************************")
if args.prompt and not args.compute_logprobs:
print(f'Prompt: {args.prompt}\n\nResponse: {response[0][0][0]}')