api.py 9.41 KB
Newer Older
Jared Casper's avatar
Jared Casper committed
1
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
mshoeybi's avatar
working  
mshoeybi committed
2
3
4
5
6
7

"""Inference API."""


import torch

8
from megatron.core import mpu
mshoeybi's avatar
working  
mshoeybi committed
9
from .communication import broadcast_float_list
10
11
from .generation import (
        generate_tokens_probs_and_return_on_first_stage,
rprenger's avatar
rprenger committed
12
13
        score_and_return_on_first_stage,
        beam_search_and_return_on_first_stage)
mshoeybi's avatar
mshoeybi committed
14
15
16
from .tokenization import (
    tokenize_prompts,
    detokenize_generations)
liangjing's avatar
liangjing committed
17
from .forward_step import ForwardStep
mshoeybi's avatar
mshoeybi committed
18
19

def generate_and_post_process(model,
liangjing's avatar
liangjing committed
20
                              forward_step=ForwardStep,
mshoeybi's avatar
mshoeybi committed
21
22
23
                              prompts=None,
                              tokens_to_generate=0,
                              return_output_log_probs=False,
mshoeybi's avatar
mshoeybi committed
24
25
                              top_k_sampling=0,
                              top_p_sampling=0.0,
26
27
                              top_p_decay=0.0,
                              top_p_bound=0.0,
mshoeybi's avatar
mshoeybi committed
28
                              temperature=1.0,
mshoeybi's avatar
mshoeybi committed
29
                              add_BOS=False,
30
31
                              use_eod_token_for_early_termination=True,
                              stop_on_double_eol=False,
32
                              stop_on_eol=False,
Peng Xu's avatar
Peng Xu committed
33
                              prevent_newline_after_colon=False,
liangjing's avatar
liangjing committed
34
35
36
                              random_seed=-1,
                              return_logits=False,
                              detokenize_segments=True):
mshoeybi's avatar
mshoeybi committed
37
    """Run inference and post-process outputs, i.e., detokenize,
mshoeybi's avatar
mshoeybi committed
38
    move to cpu and convert to list."""
mshoeybi's avatar
mshoeybi committed
39
40

    # Main inference.
liangjing's avatar
liangjing committed
41
    tokens, lengths, output_log_probs, logits = generate(
mshoeybi's avatar
mshoeybi committed
42
        model,
liangjing's avatar
liangjing committed
43
        forward_step=forward_step,
mshoeybi's avatar
mshoeybi committed
44
45
46
        prompts=prompts,
        tokens_to_generate=tokens_to_generate,
        return_output_log_probs=return_output_log_probs,
mshoeybi's avatar
mshoeybi committed
47
48
        top_k_sampling=top_k_sampling,
        top_p_sampling=top_p_sampling,
49
50
        top_p_decay=top_p_decay,
        top_p_bound=top_p_bound,
mshoeybi's avatar
mshoeybi committed
51
        temperature=temperature,
mshoeybi's avatar
mshoeybi committed
52
        add_BOS=add_BOS,
53
54
        use_eod_token_for_early_termination=use_eod_token_for_early_termination,
        stop_on_double_eol=stop_on_double_eol,
55
        stop_on_eol=stop_on_eol,
Peng Xu's avatar
Peng Xu committed
56
        prevent_newline_after_colon=prevent_newline_after_colon,
57
        random_seed=random_seed)
mshoeybi's avatar
mshoeybi committed
58
59
60
61

    # Only post-process on first stage.
    if mpu.is_pipeline_first_stage():
        tokens, prompts_plus_generations, prompts_plus_generations_segments = \
liangjing's avatar
liangjing committed
62
            detokenize_generations(tokens, lengths, detokenize_segments)
mshoeybi's avatar
mshoeybi committed
63
64
65

        if return_output_log_probs:
            output_log_probs = output_log_probs.cpu().numpy().tolist()
66
67
            for i, (prob, seg) in enumerate(zip(output_log_probs, prompts_plus_generations_segments)):
                output_log_probs[i] = prob[:len(seg)-1]
mshoeybi's avatar
mshoeybi committed
68

liangjing's avatar
liangjing committed
69
70
71
72
73
74
75
        if return_logits:
            assert(tokens_to_generate == 0)
            assert(mpu.get_pipeline_model_parallel_world_size() == 1)
            return prompts_plus_generations, prompts_plus_generations_segments, \
            output_log_probs, tokens, logits
        else:
            return prompts_plus_generations, prompts_plus_generations_segments, \
76
            output_log_probs, tokens
mshoeybi's avatar
mshoeybi committed
77
78

    return None
mshoeybi's avatar
working  
mshoeybi committed
79
80

def generate(model,
liangjing's avatar
liangjing committed
81
             forward_step=None,
mshoeybi's avatar
working  
mshoeybi committed
82
83
84
             prompts=None,
             tokens_to_generate=0,
             return_output_log_probs=False,
mshoeybi's avatar
mshoeybi committed
85
86
             top_k_sampling=0,
             top_p_sampling=0.0,
87
88
             top_p_decay=0.0,
             top_p_bound=0.0,
mshoeybi's avatar
mshoeybi committed
89
             temperature=1.0,
mshoeybi's avatar
mshoeybi committed
90
             add_BOS=False,
91
92
             use_eod_token_for_early_termination=True,
             stop_on_double_eol=False,
93
             stop_on_eol=False,
Peng Xu's avatar
Peng Xu committed
94
             prevent_newline_after_colon=False,
95
             random_seed=-1):
mshoeybi's avatar
mshoeybi committed
96
97
98
99
100
101
102
    """Given prompts and input parameters, run inference and return:
       tokens: prompts plus the generated tokens.
       lengths: length of the prompt + generations. Note that we can
           discard tokens in the tokens tensor that are after the
           corresponding length.
       output_log_probs: log probs of the tokens.
    """
mshoeybi's avatar
working  
mshoeybi committed
103
104

    # Make sure input params are avaialble to all ranks.
mshoeybi's avatar
mshoeybi committed
105
    values = [tokens_to_generate,
106
              return_output_log_probs,
107
              top_k_sampling, top_p_sampling, top_p_decay, top_p_bound,
108
109
              temperature, add_BOS, use_eod_token_for_early_termination,
              stop_on_double_eol,
110
              stop_on_eol,
Peng Xu's avatar
Peng Xu committed
111
              prevent_newline_after_colon,
112
              random_seed]
Peng Xu's avatar
Peng Xu committed
113
    values_float_tensor = broadcast_float_list(len(values), float_list=values)
mshoeybi's avatar
working  
mshoeybi committed
114
115
    tokens_to_generate = int(values_float_tensor[0].item())
    return_output_log_probs = bool(values_float_tensor[1].item())
mshoeybi's avatar
mshoeybi committed
116
117
    top_k_sampling = int(values_float_tensor[2].item())
    top_p_sampling = values_float_tensor[3].item()
118
119
    top_p_decay = values_float_tensor[4].item()
    top_p_bound = values_float_tensor[5].item()
120
121
122
123
124
    temperature = values_float_tensor[6].item()
    add_BOS = bool(values_float_tensor[7].item())
    use_eod_token_for_early_termination = bool(values_float_tensor[8].item())
    stop_on_double_eol = bool(values_float_tensor[9].item())
    stop_on_eol = bool(values_float_tensor[10].item())
Peng Xu's avatar
Peng Xu committed
125
126
    prevent_newline_after_colon = bool(values_float_tensor[11].item())
    random_seed = int(values_float_tensor[12].item())
127
128
129

    if random_seed != -1:
        torch.random.manual_seed(random_seed)
mshoeybi's avatar
working  
mshoeybi committed
130
131
132
133
134

    # Tokenize prompts and get the batch.
    # Note that these tensors are broadcaseted to all ranks.
    if torch.distributed.get_rank() == 0:
        assert prompts is not None
liangjing's avatar
liangjing committed
135

mshoeybi's avatar
working  
mshoeybi committed
136
    context_tokens_tensor, context_length_tensor = tokenize_prompts(
mshoeybi's avatar
mshoeybi committed
137
        prompts=prompts, tokens_to_generate=tokens_to_generate, add_BOS=add_BOS)
mshoeybi's avatar
working  
mshoeybi committed
138

139
    if tokens_to_generate == 0:
140
141
        return score_and_return_on_first_stage(
            model, context_tokens_tensor, context_length_tensor)
liangjing's avatar
liangjing committed
142

mshoeybi's avatar
working  
mshoeybi committed
143
144
145
    # Main inference function.
    # Note that the outputs are available on the first stage.
    return generate_tokens_probs_and_return_on_first_stage(
liangjing's avatar
liangjing committed
146
        model, forward_step, context_tokens_tensor, context_length_tensor,
mshoeybi's avatar
working  
mshoeybi committed
147
        return_output_log_probs=return_output_log_probs,
mshoeybi's avatar
mshoeybi committed
148
149
        top_k=top_k_sampling,
        top_p=top_p_sampling,
150
151
        top_p_decay=top_p_decay,
        top_p_bound=top_p_bound,
mshoeybi's avatar
mshoeybi committed
152
        temperature=temperature,
153
154
        use_eod_token_for_early_termination=use_eod_token_for_early_termination,
        stop_on_double_eol=stop_on_double_eol,
Peng Xu's avatar
Peng Xu committed
155
156
        stop_on_eol=stop_on_eol,
        prevent_newline_after_colon=prevent_newline_after_colon)
rprenger's avatar
rprenger committed
157
158

def beam_search_and_post_process(model,
liangjing's avatar
liangjing committed
159
                                 forward_step=ForwardStep,
rprenger's avatar
rprenger committed
160
161
162
                                 prompts=None,
                                 tokens_to_generate=0,
                                 beam_size=0,
163
164
                                 add_BOS=False,
                                 stop_token=50256,
165
                                 num_return_gen=1,
Peng Xu's avatar
Peng Xu committed
166
                                 length_penalty=1,
liangjing's avatar
liangjing committed
167
168
                                 prevent_newline_after_colon=False,
                                 detokenize_segments=True):
rprenger's avatar
rprenger committed
169
170
171
172
173
    """Run beam search and post-process outputs, i.e., detokenize,
    move to cpu and convert to list."""

    # Main inference.
    tokens, scores = beam_search(model,
liangjing's avatar
liangjing committed
174
                                 forward_step=forward_step,
rprenger's avatar
rprenger committed
175
176
177
                                 prompts=prompts,
                                 tokens_to_generate=tokens_to_generate,
                                 beam_size=beam_size,
178
179
                                 add_BOS=add_BOS,
                                 stop_token=stop_token,
180
                                 num_return_gen=num_return_gen,
Peng Xu's avatar
Peng Xu committed
181
182
                                 length_penalty=length_penalty,
                                 prevent_newline_after_colon=prevent_newline_after_colon)
rprenger's avatar
rprenger committed
183
184
    # Only post-process on first stage.
    if mpu.is_pipeline_first_stage():
liangjing's avatar
liangjing committed
185
186
        lengths = tokens.size(1)*torch.ones(beam_size, dtype=torch.int64, device=torch.cuda.current_device())
        tokens, prompts_plus_generations, prompts_plus_generations_segments = detokenize_generations(tokens, lengths, detokenize_segments)
rprenger's avatar
rprenger committed
187
188
        scores = scores.cpu().numpy().tolist()
        return prompts_plus_generations, prompts_plus_generations_segments, scores
rprenger's avatar
rprenger committed
189
190
191

    return None

liangjing's avatar
liangjing committed
192
def beam_search(model, forward_step, prompts=None, tokens_to_generate=0, beam_size=0, add_BOS=False, stop_token=50256, num_return_gen=1, length_penalty=1, prevent_newline_after_colon=False):
rprenger's avatar
rprenger committed
193
194
195
    # Make sure input params are avaialble to all ranks.
    values = [tokens_to_generate,
              beam_size,
196
197
198
              add_BOS,
              stop_token,
              num_return_gen,
Peng Xu's avatar
Peng Xu committed
199
200
201
              length_penalty,
              prevent_newline_after_colon]
    values_float_tensor = broadcast_float_list(len(values), float_list=values)
rprenger's avatar
rprenger committed
202
203
204
    tokens_to_generate = int(values_float_tensor[0].item())
    beam_size = int(values_float_tensor[1].item())
    add_BOS = bool(values_float_tensor[2].item())
205
206
207
    stop_token = int(values_float_tensor[3].item())
    num_return_gen = int(values_float_tensor[4].item())
    length_penalty = values_float_tensor[5].item()
Peng Xu's avatar
Peng Xu committed
208
    prevent_newline_after_colon = values_float_tensor[6].item()
rprenger's avatar
rprenger committed
209
210
211

    context_tokens_tensor, context_length_tensor = tokenize_prompts(
        prompts=prompts, tokens_to_generate=tokens_to_generate, add_BOS=add_BOS)
liangjing's avatar
liangjing committed
212
213

    return beam_search_and_return_on_first_stage(model, forward_step, context_tokens_tensor, context_length_tensor,
Peng Xu's avatar
Peng Xu committed
214
215
            beam_size, stop_token=stop_token, num_return_gen=num_return_gen, length_penalty=length_penalty,
            prevent_newline_after_colon=prevent_newline_after_colon)