mirror of
https://github.com/modelscope/modelscope.git
synced 2026-07-10 04:22:33 +02:00
60780769b1
* fix(llm ppl): 1. cache position; 2. stream_gready_search; 3. swift_mapping * fix punkt --------- Co-authored-by: suluyan <suluyan.sly@alibaba-inc.com>
562 lines
26 KiB
Python
562 lines
26 KiB
Python
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import types
|
|
import warnings
|
|
from contextlib import contextmanager
|
|
from typing import Any, Dict, Generator, List, Optional, Union
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
import transformers
|
|
from packaging import version
|
|
from torch import nn
|
|
from transformers import PreTrainedModel
|
|
from transformers.generation import GreedySearchDecoderOnlyOutput # noqa
|
|
from transformers.generation import (GreedySearchEncoderDecoderOutput,
|
|
LogitsProcessorList,
|
|
SampleDecoderOnlyOutput,
|
|
SampleEncoderDecoderOutput,
|
|
StoppingCriteriaList,
|
|
validate_stopping_criteria)
|
|
|
|
from modelscope.pipelines.base import Input
|
|
from modelscope.utils.constant import Frameworks
|
|
from modelscope.utils.device import device_placement
|
|
|
|
|
|
class StreamingOutputMixin:
|
|
|
|
def stream_generate(self, *args, **kwargs) -> Generator:
|
|
"""
|
|
Support the input of Model and Pipeline.
|
|
The output is a `Generator` type,
|
|
which conforms to the output standard of modelscope.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
|
|
class PipelineStreamingOutputMixin(StreamingOutputMixin):
|
|
|
|
def stream_generate(self, input: Union[Input, List[Input]], *args,
|
|
**kwargs) -> Generator:
|
|
"""
|
|
Similar to the `Pipeline.__call__` method.
|
|
it supports the input that the pipeline can accept,
|
|
and also supports batch input.
|
|
|
|
self.model must be a subclass of StreamingOutputMixin
|
|
and implement the stream method.
|
|
"""
|
|
assert isinstance(self.model, StreamingOutputMixin
|
|
), 'pipeline.model must be StreamingOutputMixin!'
|
|
if (self.model or (self.has_multiple_models and self.models[0])):
|
|
if not self._model_prepare:
|
|
self.prepare_model()
|
|
|
|
batch_size = kwargs.pop('batch_size', None)
|
|
preprocess_params, forward_params, postprocess_params = self._sanitize_parameters(
|
|
**kwargs)
|
|
|
|
if isinstance(input, list):
|
|
model_input_list = [
|
|
self._preprocess_with_check(i, preprocess_params)
|
|
for i in input
|
|
]
|
|
|
|
if batch_size is None:
|
|
output = []
|
|
for ele in model_input_list:
|
|
output.append(
|
|
self._stream_single(ele, forward_params,
|
|
postprocess_params))
|
|
else:
|
|
output = self._stream_batch(model_input_list, batch_size,
|
|
forward_params, postprocess_params)
|
|
|
|
else:
|
|
model_input = self._preprocess_with_check(input, preprocess_params)
|
|
output = self._stream_single(model_input, forward_params,
|
|
postprocess_params)
|
|
return output
|
|
|
|
def _preprocess_with_check(
|
|
self, input: Input,
|
|
preprocess_params: Dict[str, Any]) -> Dict[str, Any]:
|
|
self._check_input(input)
|
|
return self.preprocess(input, **preprocess_params)
|
|
|
|
def _stream_single(self, model_input: Dict[str, Any],
|
|
forward_params: Dict[str, Any],
|
|
postprocess_params: Dict[str, Any]) -> Generator:
|
|
|
|
with device_placement(self.framework, self.device_name):
|
|
if self.framework == Frameworks.torch:
|
|
with torch.no_grad():
|
|
if self._auto_collate:
|
|
model_input = self._collate_fn(model_input)
|
|
stream = self.model.stream_generate(
|
|
model_input, **forward_params)
|
|
else:
|
|
stream = self.model.stream_generate(model_input,
|
|
**forward_params)
|
|
|
|
for out in stream:
|
|
out = self.postprocess(out, **postprocess_params)
|
|
self._check_output(out)
|
|
yield out
|
|
|
|
def _stream_batch(self, model_input_list: List[Dict[str, Any]],
|
|
batch_size: int, forward_params: Dict[str, Any],
|
|
postprocess_params: Dict[str, Any]) -> Generator:
|
|
|
|
stream_list = []
|
|
real_batch_sizes = []
|
|
with device_placement(self.framework, self.device_name):
|
|
for i in range(0, len(model_input_list), batch_size):
|
|
end = min(i + batch_size, len(model_input_list))
|
|
real_batch_size = end - i
|
|
real_batch_sizes.append(real_batch_size)
|
|
|
|
batched_out = self._batch(model_input_list[i:end])
|
|
if self.framework == Frameworks.torch:
|
|
with torch.no_grad():
|
|
if self._auto_collate:
|
|
batched_out = self._collate_fn(batched_out)
|
|
stream_list.append(
|
|
self.model.stream_generate(batched_out,
|
|
**forward_params))
|
|
else:
|
|
stream_list.append(
|
|
self.model.stream_generate(batched_out,
|
|
**forward_params))
|
|
|
|
output_list = [None] * len(model_input_list)
|
|
stop_streams = 0
|
|
while stop_streams < len(stream_list):
|
|
stop_streams = 0
|
|
for i, (stream, real_batch_size) in enumerate(
|
|
zip(stream_list, real_batch_sizes)):
|
|
try:
|
|
batched_out = next(stream)
|
|
for batch_idx in range(real_batch_size):
|
|
out = {}
|
|
for k, element in batched_out.items():
|
|
if element is not None:
|
|
if isinstance(element, (tuple, list)):
|
|
if isinstance(element[0],
|
|
torch.Tensor):
|
|
out[k] = type(element)(
|
|
e[batch_idx:batch_idx + 1]
|
|
for e in element)
|
|
else:
|
|
# Compatible with traditional pipelines
|
|
out[k] = element[batch_idx]
|
|
else:
|
|
out[k] = element[batch_idx:batch_idx
|
|
+ 1]
|
|
out = self.postprocess(out, **postprocess_params)
|
|
self._check_output(out)
|
|
output_index = i * batch_size + batch_idx
|
|
output_list[output_index] = out
|
|
except StopIteration:
|
|
stop_streams += 1
|
|
yield output_list
|
|
|
|
return output_list
|
|
|
|
|
|
class PretrainedModelStreamingOutputMixin(StreamingOutputMixin):
|
|
|
|
def stream_generate(self, *args, **kwargs) -> Generator:
|
|
model = self if isinstance(self, PreTrainedModel) else self.model
|
|
assert isinstance(model, PreTrainedModel), \
|
|
'self or self.model must be `PretrainedModel`!'
|
|
with self._replace_generate(model):
|
|
return model.generate(*args, **kwargs)
|
|
|
|
@contextmanager
|
|
def _replace_generate(self, model: PreTrainedModel) -> Generator:
|
|
if version.parse(transformers.__version__) >= version.parse('4.43.0'):
|
|
greedy_search_name = 'stream_greedy_search'
|
|
sample_name = '_sample'
|
|
elif version.parse(
|
|
transformers.__version__) >= version.parse('4.39.0'):
|
|
greedy_search_name = '_greedy_search'
|
|
sample_name = '_sample'
|
|
else:
|
|
greedy_search_name = 'greedy_search'
|
|
sample_name = 'sample'
|
|
origin_greedy_search = getattr(model, greedy_search_name)
|
|
origin_sample = getattr(model, sample_name)
|
|
setattr(model, greedy_search_name,
|
|
types.MethodType(self.stream_greedy_search, model))
|
|
setattr(model, sample_name, types.MethodType(self.stream_sample,
|
|
model))
|
|
yield
|
|
setattr(model, greedy_search_name, origin_greedy_search)
|
|
setattr(model, sample_name, origin_sample)
|
|
|
|
@staticmethod
|
|
def stream_greedy_search(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
logits_processor: Optional[LogitsProcessorList] = None,
|
|
stopping_criteria: Optional[StoppingCriteriaList] = None,
|
|
max_length: Optional[int] = None,
|
|
pad_token_id: Optional[int] = None,
|
|
eos_token_id: Optional[Union[int, List[int]]] = None,
|
|
output_attentions: Optional[bool] = None,
|
|
output_hidden_states: Optional[bool] = None,
|
|
output_scores: Optional[bool] = None,
|
|
return_dict_in_generate: Optional[bool] = None,
|
|
synced_gpus: bool = False,
|
|
**model_kwargs,
|
|
) -> Generator:
|
|
logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList(
|
|
)
|
|
stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList(
|
|
)
|
|
if max_length is not None:
|
|
warnings.warn(
|
|
'`max_length` is deprecated in this function, use'
|
|
' `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.',
|
|
UserWarning,
|
|
)
|
|
stopping_criteria = validate_stopping_criteria(
|
|
stopping_criteria, max_length)
|
|
pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
|
|
eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
|
|
if isinstance(eos_token_id, int):
|
|
eos_token_id = [eos_token_id]
|
|
eos_token_id_tensor = torch.tensor(eos_token_id).to(
|
|
input_ids.device) if eos_token_id is not None else None
|
|
output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
|
|
output_attentions = (
|
|
output_attentions if output_attentions is not None else
|
|
self.generation_config.output_attentions)
|
|
output_hidden_states = (
|
|
output_hidden_states if output_hidden_states is not None else
|
|
self.generation_config.output_hidden_states)
|
|
return_dict_in_generate = (
|
|
return_dict_in_generate if return_dict_in_generate is not None else
|
|
self.generation_config.return_dict_in_generate)
|
|
|
|
# init attention / hidden states / scores tuples
|
|
scores = () if (return_dict_in_generate and output_scores) else None
|
|
decoder_attentions = () if (return_dict_in_generate
|
|
and output_attentions) else None
|
|
cross_attentions = () if (return_dict_in_generate
|
|
and output_attentions) else None
|
|
decoder_hidden_states = () if (return_dict_in_generate
|
|
and output_hidden_states) else None
|
|
|
|
# if model is an encoder-decoder, retrieve encoder attention weights and hidden states
|
|
if return_dict_in_generate and self.config.is_encoder_decoder:
|
|
encoder_attentions = model_kwargs['encoder_outputs'].get(
|
|
'attentions') if output_attentions else None
|
|
encoder_hidden_states = (
|
|
model_kwargs['encoder_outputs'].get('hidden_states')
|
|
if output_hidden_states else None)
|
|
|
|
# keep track of which sequences are already finished
|
|
unfinished_sequences = torch.ones(
|
|
input_ids.shape[0], dtype=torch.long, device=input_ids.device)
|
|
|
|
this_peer_finished = False # used by synced_gpus only
|
|
while True:
|
|
if synced_gpus:
|
|
# Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
|
|
# The following logic allows an early break if all peers finished generating their sequence
|
|
this_peer_finished_flag = torch.tensor(
|
|
0.0 if this_peer_finished else 1.0).to(input_ids.device)
|
|
# send 0.0 if we finished, 1.0 otherwise
|
|
dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
|
|
# did all peers finish? the reduced sum will be 0.0 then
|
|
if this_peer_finished_flag.item() == 0.0:
|
|
break
|
|
|
|
# prepare model inputs
|
|
model_inputs = self.prepare_inputs_for_generation(
|
|
input_ids, **model_kwargs)
|
|
|
|
# forward pass to get next token
|
|
outputs = self(
|
|
**model_inputs,
|
|
return_dict=True,
|
|
output_attentions=output_attentions,
|
|
output_hidden_states=output_hidden_states,
|
|
)
|
|
|
|
if synced_gpus and this_peer_finished:
|
|
continue # don't waste resources running the code we don't need
|
|
|
|
next_token_logits = outputs.logits[:, -1, :]
|
|
|
|
# pre-process distribution
|
|
next_tokens_scores = logits_processor(input_ids, next_token_logits)
|
|
|
|
# Store scores, attentions and hidden_states when required
|
|
if return_dict_in_generate:
|
|
if output_scores:
|
|
scores += (next_tokens_scores, )
|
|
if output_attentions:
|
|
decoder_attentions += ((outputs.decoder_attentions, ) if
|
|
self.config.is_encoder_decoder else
|
|
(outputs.attentions, ))
|
|
if self.config.is_encoder_decoder:
|
|
cross_attentions += (outputs.cross_attentions, )
|
|
|
|
if output_hidden_states:
|
|
decoder_hidden_states += ((outputs.decoder_hidden_states, )
|
|
if self.config.is_encoder_decoder
|
|
else (outputs.hidden_states, ))
|
|
|
|
# argmax
|
|
next_tokens = torch.argmax(next_tokens_scores, dim=-1)
|
|
|
|
# finished sentences should have their next token be a padding token
|
|
if eos_token_id is not None:
|
|
if pad_token_id is None:
|
|
raise ValueError(
|
|
'If `eos_token_id` is defined, make sure that `pad_token_id` is defined.'
|
|
)
|
|
next_tokens = next_tokens * unfinished_sequences + pad_token_id * (
|
|
1 - unfinished_sequences)
|
|
|
|
# update generated ids, model inputs, and length for next step
|
|
input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
|
|
|
|
# return Generator for stream
|
|
if return_dict_in_generate:
|
|
if self.config.is_encoder_decoder:
|
|
yield GreedySearchEncoderDecoderOutput(
|
|
sequences=input_ids,
|
|
scores=scores,
|
|
encoder_attentions=encoder_attentions,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
decoder_attentions=decoder_attentions,
|
|
cross_attentions=cross_attentions,
|
|
decoder_hidden_states=decoder_hidden_states,
|
|
)
|
|
else:
|
|
yield GreedySearchDecoderOnlyOutput(
|
|
sequences=input_ids,
|
|
scores=scores,
|
|
attentions=decoder_attentions,
|
|
hidden_states=decoder_hidden_states,
|
|
)
|
|
else:
|
|
yield input_ids
|
|
|
|
model_kwargs = self._update_model_kwargs_for_generation(
|
|
outputs,
|
|
model_kwargs,
|
|
is_encoder_decoder=self.config.is_encoder_decoder)
|
|
|
|
# if eos_token was found in one sentence, set sentence to finished
|
|
if eos_token_id_tensor is not None:
|
|
unfinished_sequences = unfinished_sequences.mul(
|
|
next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(
|
|
eos_token_id_tensor.unsqueeze(1)).prod(dim=0))
|
|
|
|
# stop when each sentence is finished
|
|
if unfinished_sequences.max() == 0:
|
|
this_peer_finished = True
|
|
|
|
# stop if we exceed the maximum length
|
|
if stopping_criteria(input_ids, scores):
|
|
this_peer_finished = True
|
|
|
|
if this_peer_finished and not synced_gpus:
|
|
break
|
|
|
|
@staticmethod
|
|
def stream_sample(
|
|
self,
|
|
input_ids: torch.LongTensor,
|
|
logits_processor: Optional[LogitsProcessorList] = None,
|
|
stopping_criteria: Optional[StoppingCriteriaList] = None,
|
|
logits_warper: Optional[LogitsProcessorList] = None,
|
|
max_length: Optional[int] = None,
|
|
pad_token_id: Optional[int] = None,
|
|
eos_token_id: Optional[Union[int, List[int]]] = None,
|
|
output_attentions: Optional[bool] = None,
|
|
output_hidden_states: Optional[bool] = None,
|
|
output_scores: Optional[bool] = None,
|
|
return_dict_in_generate: Optional[bool] = None,
|
|
synced_gpus: bool = False,
|
|
**model_kwargs,
|
|
) -> Generator:
|
|
logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList(
|
|
)
|
|
stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList(
|
|
)
|
|
if max_length is not None:
|
|
warnings.warn(
|
|
'`max_length` is deprecated in this function, use'
|
|
' `stopping_criteria=StoppingCriteriaList(MaxLengthCriteria(max_length=max_length))` instead.',
|
|
UserWarning,
|
|
)
|
|
stopping_criteria = validate_stopping_criteria(
|
|
stopping_criteria, max_length)
|
|
logits_warper = logits_warper if logits_warper is not None else LogitsProcessorList(
|
|
)
|
|
pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
|
|
eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
|
|
if isinstance(eos_token_id, int):
|
|
eos_token_id = [eos_token_id]
|
|
eos_token_id_tensor = torch.tensor(eos_token_id).to(
|
|
input_ids.device) if eos_token_id is not None else None
|
|
output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
|
|
output_attentions = (
|
|
output_attentions if output_attentions is not None else
|
|
self.generation_config.output_attentions)
|
|
output_hidden_states = (
|
|
output_hidden_states if output_hidden_states is not None else
|
|
self.generation_config.output_hidden_states)
|
|
return_dict_in_generate = (
|
|
return_dict_in_generate if return_dict_in_generate is not None else
|
|
self.generation_config.return_dict_in_generate)
|
|
|
|
# init attention / hidden states / scores tuples
|
|
scores = () if (return_dict_in_generate and output_scores) else None
|
|
decoder_attentions = () if (return_dict_in_generate
|
|
and output_attentions) else None
|
|
cross_attentions = () if (return_dict_in_generate
|
|
and output_attentions) else None
|
|
decoder_hidden_states = () if (return_dict_in_generate
|
|
and output_hidden_states) else None
|
|
|
|
# if model is an encoder-decoder, retrieve encoder attention weights and hidden states
|
|
if return_dict_in_generate and self.config.is_encoder_decoder:
|
|
encoder_attentions = model_kwargs['encoder_outputs'].get(
|
|
'attentions') if output_attentions else None
|
|
encoder_hidden_states = (
|
|
model_kwargs['encoder_outputs'].get('hidden_states')
|
|
if output_hidden_states else None)
|
|
|
|
# keep track of which sequences are already finished
|
|
unfinished_sequences = torch.ones(
|
|
input_ids.shape[0], dtype=torch.long, device=input_ids.device)
|
|
|
|
this_peer_finished = False # used by synced_gpus only
|
|
# auto-regressive generation
|
|
while True:
|
|
if synced_gpus:
|
|
# Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
|
|
# The following logic allows an early break if all peers finished generating their sequence
|
|
this_peer_finished_flag = torch.tensor(
|
|
0.0 if this_peer_finished else 1.0).to(input_ids.device)
|
|
# send 0.0 if we finished, 1.0 otherwise
|
|
dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
|
|
# did all peers finish? the reduced sum will be 0.0 then
|
|
if this_peer_finished_flag.item() == 0.0:
|
|
break
|
|
|
|
# prepare model inputs
|
|
model_kwargs = self._get_initial_cache_position(
|
|
input_ids, model_kwargs)
|
|
model_inputs = self.prepare_inputs_for_generation(
|
|
input_ids, **model_kwargs)
|
|
|
|
# forward pass to get next token
|
|
outputs = self(
|
|
**model_inputs,
|
|
return_dict=True,
|
|
output_attentions=output_attentions,
|
|
output_hidden_states=output_hidden_states,
|
|
)
|
|
|
|
if synced_gpus and this_peer_finished:
|
|
continue # don't waste resources running the code we don't need
|
|
|
|
next_token_logits = outputs.logits[:, -1, :]
|
|
|
|
# pre-process distribution
|
|
next_token_scores = logits_processor(input_ids, next_token_logits)
|
|
next_token_scores = logits_warper(input_ids, next_token_scores)
|
|
|
|
# Store scores, attentions and hidden_states when required
|
|
if return_dict_in_generate:
|
|
if output_scores:
|
|
scores += (next_token_scores, )
|
|
if output_attentions:
|
|
decoder_attentions += ((outputs.decoder_attentions, ) if
|
|
self.config.is_encoder_decoder else
|
|
(outputs.attentions, ))
|
|
if self.config.is_encoder_decoder:
|
|
cross_attentions += (outputs.cross_attentions, )
|
|
|
|
if output_hidden_states:
|
|
decoder_hidden_states += ((outputs.decoder_hidden_states, )
|
|
if self.config.is_encoder_decoder
|
|
else (outputs.hidden_states, ))
|
|
|
|
# sample
|
|
probs = nn.functional.softmax(next_token_scores, dim=-1)
|
|
next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
|
|
|
|
# finished sentences should have their next token be a padding token
|
|
if eos_token_id is not None:
|
|
if pad_token_id is None:
|
|
raise ValueError(
|
|
'If `eos_token_id` is defined, make sure that `pad_token_id` is defined.'
|
|
)
|
|
next_tokens = next_tokens * unfinished_sequences + pad_token_id * (
|
|
1 - unfinished_sequences)
|
|
|
|
# update generated ids, model inputs, and length for next step
|
|
input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
|
|
|
|
# return Generator for stream
|
|
if return_dict_in_generate:
|
|
if self.config.is_encoder_decoder:
|
|
yield SampleEncoderDecoderOutput(
|
|
sequences=input_ids,
|
|
scores=scores,
|
|
encoder_attentions=encoder_attentions,
|
|
encoder_hidden_states=encoder_hidden_states,
|
|
decoder_attentions=decoder_attentions,
|
|
cross_attentions=cross_attentions,
|
|
decoder_hidden_states=decoder_hidden_states,
|
|
)
|
|
else:
|
|
yield SampleDecoderOnlyOutput(
|
|
sequences=input_ids,
|
|
scores=scores,
|
|
attentions=decoder_attentions,
|
|
hidden_states=decoder_hidden_states,
|
|
)
|
|
else:
|
|
yield input_ids
|
|
|
|
model_kwargs = self._update_model_kwargs_for_generation(
|
|
outputs,
|
|
model_kwargs,
|
|
is_encoder_decoder=self.config.is_encoder_decoder)
|
|
|
|
# if eos_token was found in one sentence, set sentence to finished
|
|
if eos_token_id_tensor is not None:
|
|
unfinished_sequences = unfinished_sequences.mul(
|
|
next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(
|
|
eos_token_id_tensor.unsqueeze(1)).prod(dim=0))
|
|
|
|
# stop when each sentence is finished
|
|
if unfinished_sequences.max() == 0:
|
|
this_peer_finished = True
|
|
|
|
# stop if we exceed the maximum length
|
|
if stopping_criteria(input_ids, scores):
|
|
this_peer_finished = True
|
|
|
|
if this_peer_finished and not synced_gpus:
|
|
break
|
|
|
|
|
|
def add_stream_generate(model: PreTrainedModel):
|
|
pretrained_class = type(model)
|
|
parent_classes = (pretrained_class, PretrainedModelStreamingOutputMixin)
|
|
new_model = type(pretrained_class.__name__, parent_classes, {})(
|
|
model.config)
|
|
new_model.__dict__.update(model.__dict__)
|
|
return new_model
|