[to #42322933] GPT-3 model supports batch input

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11322820

modelscope/models/nlp/gpt3/backbone.py

@@ -23,8 +23,10 @@ from torch import nn
 from torch.nn import functional as F
 from transformers.modeling_utils import PreTrainedModel
 
+from modelscope.outputs import TokenGeneratorOutput
 from modelscope.utils.constant import ModelFile
 from .configuration import GPT3Config
+from .distributed_gpt3 import sample
 
 
 class GPT3SelfAttention(nn.Module):
@@ -351,5 +353,63 @@ class GPT3Model(PreTrainedModel):
         model.load_state_dict(state_dict)
         return model
 
-    def prepare_inputs_for_generation(self, input_ids, *args, **kwargs):
-        return {'input_ids': input_ids}
+    def generate(self, tokens, temperature=1.0, **kwargs):
+
+        batch_size = tokens.size(0)
+        lengths = kwargs.pop(
+            'prompt_length',
+            torch.tensor([tokens.size(1)], device=tokens.device))
+
+        min_prompt_length = lengths.min().item()
+        max_sequence_length = tokens.size(1)
+        max_sequence_length = min(max_sequence_length,
+                                  self.config.max_position_embeddings)
+
+        # If the context is too big, this happens
+        if min_prompt_length >= max_sequence_length:
+            raise ValueError('context length + tokens_to_generate too large')
+
+        # Added termination_id to support the case that we want to terminate the
+        # generation once that id is generated.
+        termination_id = self.config.eod_id
+
+        # Whether we have reached a termination id.
+        is_generation_done = torch.zeros(
+            batch_size, dtype=torch.uint8, device=tokens.device)
+
+        with torch.no_grad():
+            for context_length in range(min_prompt_length,
+                                        max_sequence_length):
+
+                # Pick the slice that we need to pass through the network.
+                tokens2use = tokens[:, :context_length]
+
+                # logits will be meanigful only in the last pipeline stage.
+                logits = self(tokens2use).logits
+
+                # Sample.
+                last_token_logits = logits[:, -1, :]
+                new_sample = sample(
+                    last_token_logits,
+                    top_k=self.config.top_k,
+                    top_p=self.config.top_p,
+                    temperature=temperature,
+                    vocab_size=self.config.vocab_size)
+
+                # If a prompt length is smaller or equal th current context
+                # length, it means we have started generating tokens
+                started = lengths <= context_length
+                # Update the tokens.
+                tokens[started, context_length] = new_sample[started]
+
+                done_token = (new_sample == termination_id).byte() & \
+                    started.byte()
+
+                is_generation_done = is_generation_done | done_token
+                done = torch.all(is_generation_done)
+
+                if done:
+                    break
+
+        tokens = tokens[:, :(context_length + 1)]
+        return TokenGeneratorOutput(sequences=tokens)

modelscope/models/nlp/gpt3/distributed_gpt3.py

@@ -851,8 +851,6 @@ def sample(logits, top_k=0, top_p=0.0, temperature=1.0, vocab_size=None):
 
     # Check logits for consistency.
     assert logits.ndim == 2, 'expected the logits to be of [b, v] shape.'
-    assert logits.type() == 'torch.cuda.FloatTensor', \
-        'input logits should be floats.'
 
     # Greedy is just simple argmax.
     if top_k == 1:

modelscope/models/nlp/gpt3/text_generation.py

@@ -1,7 +1,7 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
-import os
 from typing import Dict
 
+import torch
 from transformers import BertTokenizer
 
 from modelscope.metainfo import Models
@@ -49,27 +49,17 @@ class GPT3ForTextGeneration(TorchModel):
         """
         return self.model(**input)
 
-    def generate(self, input: Dict[str, Tensor]) -> Dict[str, Tensor]:
+    def generate(self, inputs: Dict[str, Tensor]) -> Dict[str, Tensor]:
         if not isinstance(self.model, GPT3Model):
-            return self.model.generate(**input)
+            return self.model.generate(**inputs)
 
-        assert 'input_ids' in input, "generate function must accept 'input_ids' key"
-        input_ids = input['input_ids']
-        if 'attention_mask' in input:
-            attention_mask = input['attention_mask']
-            input_ids = input_ids[0][attention_mask[0].nonzero()] \
-                .squeeze().unsqueeze(0)
-        # remove sep token at the end of tokenizer output
-        input_ids = input_ids[:, :-1]
+        tokens = inputs['input_ids']
+        lengths = self._get_length(inputs['attention_mask'])
+        return self.model.generate(tokens, prompt_length=lengths)
 
-        gen_params = dict()
-        gen_params['inputs'] = input_ids
-        gen_params['do_sample'] = input.pop('do_sample', True)
-        gen_params['max_length'] = input.pop('max_length', 128)
-        gen_params['top_k'] = input.pop('top_k', 10)
-        gen_params['top_p'] = input.pop('top_p', None)
-        sample_output = self.model.generate(**gen_params)
-        return {'sequences': sample_output[0]}
+    @staticmethod
+    def _get_length(attention_mask: torch.Tensor) -> Tensor:
+        return attention_mask.sum(-1) - 1
 
     def save_pretrained(self, *args, **kwargs):
         if not isinstance(self.model, GPT3Model):

modelscope/pipelines/base.py

@@ -275,17 +275,15 @@ class Pipeline(ABC):
                                                    **forward_params)
                 else:
                     batched_out = self.forward(batched_input, **forward_params)
-            if real_batch_size == 1:
-                output_list.append(batched_out)
-            else:
-                for batch_idx in range(real_batch_size):
-                    out = {}
-                    for k, element in batched_out.items():
-                        if element is not None:
-                            out[k] = element[batch_idx]
-                    out = self.postprocess(out, **postprocess_params)
-                    self._check_output(out)
-                    output_list.append(out)
+
+            for batch_idx in range(real_batch_size):
+                out = {}
+                for k, element in batched_out.items():
+                    if element is not None:
+                        out[k] = element[batch_idx]
+                out = self.postprocess(out, **postprocess_params)
+                self._check_output(out)
+                output_list.append(out)
 
         return output_list
 

modelscope/preprocessors/nlp/text_generation_preprocessor.py

@@ -161,12 +161,7 @@ class TextGenerationTransformersPreprocessor(TextGenerationPreprocessorBase):
         output = self.nlp_tokenizer(sequence1, **kwargs)
         if self.mode != ModeKeys.INFERENCE:
             if sequence2 is not None:
-                self.nlp_tokenizer.tokenize_kwargs[
-                    'max_length'] = self.tgt_length
-                labels = self.nlp_tokenizer(sequence2)['input_ids']
-                self.nlp_tokenizer.tokenize_kwargs[
-                    'max_length'] = self.src_length
-
+                labels = self._get_labels_from_tgt(sequence2)
                 src_input_ids = output['input_ids']
                 src_attention_mask = output['attention_mask']
             else:
@@ -181,6 +176,12 @@ class TextGenerationTransformersPreprocessor(TextGenerationPreprocessorBase):
             }
         return output
 
+    def _get_labels_from_tgt(self, sequence: str) -> torch.Tensor:
+        self.nlp_tokenizer.tokenize_kwargs['max_length'] = self.tgt_length
+        labels = self.nlp_tokenizer(sequence)['input_ids']
+        self.nlp_tokenizer.tokenize_kwargs['max_length'] = self.src_length
+        return labels
+
 
 @PREPROCESSORS.register_module(
     Fields.nlp, module_name=Preprocessors.text_gen_jieba_tokenizer)

tests/pipelines/test_text_generation.py

@@ -95,6 +95,19 @@ class TextGenerationTest(unittest.TestCase, DemoCompatibilityCheck):
         self.run_pipeline_with_model_id(self.gpt3_base_model_id,
                                         self.gpt3_input)
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_gpt_base_with_model_name_batch(self):
+        self.run_pipeline_with_model_id(
+            self.gpt3_base_model_id,
+            [self.gpt3_input, self.gpt3_input[:10], self.gpt3_input[10:]],
+            run_kwargs={'batch_size': 2})
+
+    @unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
+    def test_gpt_base_with_model_name_batch_iter(self):
+        self.run_pipeline_with_model_id(
+            self.gpt3_base_model_id,
+            [self.gpt3_input, self.gpt3_input[:10], self.gpt3_input[10:]])
+
     @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_gpt_large_with_model_name(self):
         self.run_pipeline_with_model_id(self.gpt3_large_model_id,