run_text_generation_server.py 3.26 KB
Newer Older
Jared Casper's avatar
Jared Casper committed
1
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
Ryan Prenger's avatar
Ryan Prenger committed
2
3
4
5
6
7
8
9
10

"""Sample Generate GPT"""
import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
                                             os.path.pardir)))
import socket
from megatron import get_args
from megatron import print_rank_0
11
from megatron.core import mpu
Ryan Prenger's avatar
Ryan Prenger committed
12
13
14
15
from megatron.checkpointing import load_checkpoint
from megatron.initialize import initialize_megatron
from megatron.model import GPTModel
from megatron.training import get_model
liangjing's avatar
v1  
liangjing committed
16
from megatron.arguments import core_transformer_config_from_args
Ryan Prenger's avatar
Ryan Prenger committed
17
from megatron.text_generation_server import MegatronServer
18
from megatron.text_generation import generate_and_post_process
rprenger's avatar
rprenger committed
19
from megatron.text_generation import beam_search_and_post_process
Ryan Prenger's avatar
Ryan Prenger committed
20
21
22
23
24
import torch

def model_provider(pre_process=True, post_process=True):
    """Build the model."""

liangjing's avatar
v1  
liangjing committed
25
26
    config = core_transformer_config_from_args(get_args())

Ryan Prenger's avatar
Ryan Prenger committed
27
    print_rank_0('building GPT model ...')
liangjing's avatar
v1  
liangjing committed
28
    model = GPTModel(config, num_tokentypes=0, parallel_output=False, pre_process=pre_process, post_process=post_process)
Ryan Prenger's avatar
Ryan Prenger committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42

    return model

def add_text_generate_args(parser):
    group = parser.add_argument_group(title='text generation')

    group.add_argument("--temperature", type=float, default=1.0,
                       help='Sampling temperature.')
    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=1024,
                       help='Size of the output generated text.')
liangjing's avatar
v1  
liangjing committed
43
44
    group.add_argument("--port", type=int, default=5000,
                       help='port for text generation server to run on')
Ryan Prenger's avatar
Ryan Prenger committed
45
46
47
48
49
50
51
52
53
54
55
56
57
    return parser


if __name__ == "__main__":
    initialize_megatron(extra_args_provider=add_text_generate_args,
                        args_defaults={'tokenizer_type': 'GPT2BPETokenizer',
                                       'no_load_rng': True,
                                       'no_load_optim': True})

    args = get_args()
    if args.num_layers_per_virtual_pipeline_stage is not None:
        print("Interleaved pipeline schedule is not yet supported for text generation.")
        exit()
liangjing's avatar
v1  
liangjing committed
58
59
60
    print_rank_0("WARNING: Forcing exit_on_missing_checkpoint to True for text "
                 "generation.")
    args.exit_on_missing_checkpoint = True
Ryan Prenger's avatar
Ryan Prenger committed
61
    # Set up model and load checkpoint
62
    model = get_model(model_provider, wrap_with_ddp=False)
Ryan Prenger's avatar
Ryan Prenger committed
63
64
65
66
67
68
69
70

    if args.load is not None:
        _ = load_checkpoint(model, None, None)

    assert len(model) == 1, "Above condition should have caught this"
    model = model[0]
    if mpu.is_pipeline_first_stage() and mpu.get_tensor_model_parallel_rank() == 0:
        server = MegatronServer(model)
liangjing's avatar
v1  
liangjing committed
71
        server.run("0.0.0.0",port=args.port)
Ryan Prenger's avatar
Ryan Prenger committed
72
73
74

    while True:
        choice = torch.cuda.LongTensor(1)
75
        torch.distributed.broadcast(choice, 0)
Ryan Prenger's avatar
Ryan Prenger committed
76
        if choice[0].item() == 0:
77
78
79
80
            try:
                generate_and_post_process(model)
            except ValueError as ve:
                pass
rprenger's avatar
rprenger committed
81
82
83
84
85
        elif choice[0].item() == 1:
            try:
                beam_search_and_post_process(model)
            except ValueError as ve:
                pass