From 871b345e79a557f10fbc260db50abecefc2f8024 Mon Sep 17 00:00:00 2001 From: "hemu.zp" Date: Mon, 9 Jan 2023 09:31:44 +0800 Subject: [PATCH] [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 | 64 ++++++++++++++++++- .../models/nlp/gpt3/distributed_gpt3.py | 2 - modelscope/models/nlp/gpt3/text_generation.py | 28 +++----- modelscope/pipelines/base.py | 20 +++--- .../nlp/text_generation_preprocessor.py | 13 ++-- tests/pipelines/test_text_generation.py | 13 ++++ 6 files changed, 100 insertions(+), 40 deletions(-) diff --git a/modelscope/models/nlp/gpt3/backbone.py b/modelscope/models/nlp/gpt3/backbone.py index 4647428e..a86f01e4 100644 --- a/modelscope/models/nlp/gpt3/backbone.py +++ b/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) diff --git a/modelscope/models/nlp/gpt3/distributed_gpt3.py b/modelscope/models/nlp/gpt3/distributed_gpt3.py index 0fb3843b..ceb8c218 100644 --- a/modelscope/models/nlp/gpt3/distributed_gpt3.py +++ b/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: diff --git a/modelscope/models/nlp/gpt3/text_generation.py b/modelscope/models/nlp/gpt3/text_generation.py index 0d6f33b5..9361f0a2 100644 --- a/modelscope/models/nlp/gpt3/text_generation.py +++ b/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): diff --git a/modelscope/pipelines/base.py b/modelscope/pipelines/base.py index d6e23f0c..d20549f2 100644 --- a/modelscope/pipelines/base.py +++ b/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 diff --git a/modelscope/preprocessors/nlp/text_generation_preprocessor.py b/modelscope/preprocessors/nlp/text_generation_preprocessor.py index c1c32bf0..93d2f16c 100644 --- a/modelscope/preprocessors/nlp/text_generation_preprocessor.py +++ b/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) diff --git a/tests/pipelines/test_text_generation.py b/tests/pipelines/test_text_generation.py index 7977d3ee..cbb1b29b 100644 --- a/tests/pipelines/test_text_generation.py +++ b/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,