diff --git a/src/deepsparse/v2/operators/engine_operator.py b/src/deepsparse/v2/operators/engine_operator.py index 9ee8d734c5..630de2d5bd 100644 --- a/src/deepsparse/v2/operators/engine_operator.py +++ b/src/deepsparse/v2/operators/engine_operator.py @@ -20,7 +20,7 @@ from deepsparse import Context as EngineContext from deepsparse import Engine, MultiModelEngine, Scheduler from deepsparse.benchmark import ORTEngine -from deepsparse.utils import model_to_path +from deepsparse.utils import join_engine_outputs, model_to_path, split_engine_inputs from deepsparse.v2.operators import Operator @@ -29,12 +29,12 @@ SUPPORTED_PIPELINE_ENGINES = [DEEPSPARSE_ENGINE, ORT_ENGINE] -__all__ = ["EngineOperator"] +__all__ = ["EngineOperator", "EngineOperatorInputs", "EngineOperatorOutputs"] class EngineOperatorInputs(BaseModel): engine_inputs: List = Field(description="engine_inputs") - engine: Optional[Engine] = Field( + engine: Optional[Union[ORTEngine, Engine]] = Field( description="override the engine to run forward pass with", default=None, ) @@ -95,8 +95,8 @@ def __init__( engine_kwargs: Dict = None, ): self.model_path = model_to_path(model_path) - self._batch_size = 1 self.engine_context = engine_context + self._batch_size = 1 if self.engine_context is not None: num_cores = num_cores or self.engine_context.num_cores @@ -131,6 +131,7 @@ def batch_size(self) -> int: """ return self._batch_size + # TODO: maybe add a few args to make this less opaque? def create_engine( self, **kwargs, @@ -142,7 +143,8 @@ def create_engine( constructor/compilation :return: inference engine """ - onnx_file_path = self.model_path + + onnx_file_path = kwargs.pop("model_path", self.model_path) engine_args = deepcopy(self._engine_args) engine_args.update(kwargs) engine_type = self._engine_type.lower() diff --git a/src/deepsparse/v2/operators/operator.py b/src/deepsparse/v2/operators/operator.py index 5bb0be841a..2923862b12 100644 --- a/src/deepsparse/v2/operators/operator.py +++ b/src/deepsparse/v2/operators/operator.py @@ -17,7 +17,7 @@ from pydantic import BaseModel -from deepsparse.v2.utils import InferenceState, PipelineState +from deepsparse.v2.utils import InferenceState __all__ = ["Operator"] @@ -57,7 +57,6 @@ def __call__( self, *args, inference_state: InferenceState, - pipeline_state: PipelineState, **kwargs, ) -> Any: """ @@ -90,13 +89,11 @@ def __call__( run_output = self.run( inference_input, inference_state=inference_state, - pipeline_state=pipeline_state, ) else: run_output = self.run( *args, inference_state=inference_state, - pipeline_state=pipeline_state, **kwargs, ) if self.has_output_schema(): diff --git a/src/deepsparse/v2/pipeline.py b/src/deepsparse/v2/pipeline.py index f56680d2b9..78d112a2b3 100644 --- a/src/deepsparse/v2/pipeline.py +++ b/src/deepsparse/v2/pipeline.py @@ -15,13 +15,18 @@ import copy from concurrent.futures import Future -from functools import partial -from typing import Any, Callable, Dict, List, Union +from typing import Any, Dict, List, Union -from deepsparse.v2.operators import Operator +from deepsparse.v2.operators import EngineOperator, Operator from deepsparse.v2.routers import Router -from deepsparse.v2.schedulers import OperatorScheduler, SchedulerGroup +from deepsparse.v2.schedulers import ( + ContinuousBatchingScheduler, + OperatorScheduler, + SchedulerGroup, +) from deepsparse.v2.utils import InferenceState, PipelineState +from deepsparse.v2.utils.data import SubGraph +from deepsparse.v2.utils.helpers import run_func __all__ = ["Pipeline"] @@ -50,6 +55,7 @@ def __init__( ops: Union[Dict[str, Operator], List[Operator]], router: Router, schedulers: List[OperatorScheduler], + continuous_batching_scheduler: ContinuousBatchingScheduler, pipeline_state: PipelineState = None, ): @@ -57,32 +63,92 @@ def __init__( self.router = router self.schedulers = schedulers self.pipeline_state = pipeline_state + self._continuous_batching_scheduler = continuous_batching_scheduler self.validate() self._scheduler_group = SchedulerGroup(self.schedulers) - def _run_sequential( + def _run_next( self, inp: Any, inference_state: InferenceState, - pipeline_state: PipelineState, - start: str, - end: str, + next_step: str, ): - next_step = start - while next_step != end: - outputs = self._run_next_step( - func=self.ops[next_step], - next_step=next_step, - input=inp, - pipeline_state=pipeline_state, - inference_state=inference_state, + if ( + isinstance(self.ops[next_step], EngineOperator) + and self._continuous_batching_scheduler + ): + func = self._continuous_batching_scheduler.submit + inp = self.ops[next_step].input_schema(**inp) + else: + func = self._scheduler_group.submit + + return run_func( + func=func, + operator=self.ops[next_step], + inp=inp, + pipeline_state=self.pipeline_state, + inference_state=inference_state, + ) + + def _run_sub_graphs( + self, sub_graph_inputs: List[Any], sub_graphs: List[SubGraph] + ) -> List[Any]: + """ + Run a list of sub_graphs asynchronously. Polls to identify the sub graph that is + still running but has completed its current step. Schedules the next step + subgraph step. This is repeated until all subgraphs have finished running and + have reached their end step (stored in the Subgraph.end attribute). + + :param sub_graph_inputs: A list of inputs that should be passed to each + subgraph. Each subgraph is given an element of the list as input to its + first node. + :param sub_graphs: A list of Subgraph objects. Each stores the relevant + execution information for the particular subgraph, such as its current step + in the sub graph, inference state, output, and end step. + + :returns: a list of outputs for all the completed Subgraph objects. Returned + in the same order that the subgraphs were passed to the function. + """ + for i in range(len(sub_graphs)): + sub_graphs[i].output = self._run_next( + sub_graph_inputs[i], sub_graphs[i].inf, sub_graphs[i].step ) - next_step, operator_output, state_update = outputs - if state_update: - inference_state.update_state(state_update) - inp = operator_output - return inp + + # Execute all sub graphs until all graphs have been completed. + while True: + for sub_graph in sub_graphs: + if isinstance(sub_graph.output, Future) and sub_graph.output.done(): + # get the result for the completed operator; resolve its output + operator_output = sub_graph.output.result() + operator_output = sub_graph.parse_output(operator_output) + + # determine the next step for the particular operator, using + # its previous output and previously stored step + next_step = self.router.next( + sub_graph.step, self.ops, operator_output + ) + # update the step + sub_graph.step = next_step + + # store the output for the next step. If the next step is + # end step, this particular route has completed. Simply + # update the output value + if next_step in sub_graph.end: + sub_graph.output = operator_output + else: + sub_graph.output = self._run_next( + inp=operator_output, + inference_state=sub_graph.inf, + next_step=next_step, + ) + break + + # keep running until all sub graphs have completed. + if not any(isinstance(x.output, Future) for x in sub_graphs): + break + + return [x.output for x in sub_graphs] def _apply_split(self, inp: Any, inference_state: InferenceState): """ @@ -93,59 +159,29 @@ def _apply_split(self, inp: Any, inference_state: InferenceState): """ batches, orig_batch_size = self.expand_inputs(inp, 1) - run_with_state = partial( - self._run_sequential, - pipeline_state=self.pipeline_state, - start=self.router.route[self.router.SPLIT_ROUTE], - end=self.router.JOIN_ROUTE, - ) - inference_state_list = [ - copy.deepcopy(inference_state) for x in range(len(batches)) - ] - futures = self._scheduler_group.map( - batches, - inference_state_list, - func=run_with_state, - ) - return self.condense_inputs([x.result() for x in futures]) - def _run_next_step( - self, - *args, - func: Callable, - next_step: Union[str, int], - input: Any = None, - **kwargs, - ): - """ - Generic function to run a given func, process the output and determine the next - step. - """ - if input: - operator_output = ( - func(*args, **kwargs, **input) - if isinstance(input, dict) - else func(input, *args, **kwargs) + # Create a list of SplitRoutes, per batch size 1 + # Each SplitRoute object holds information about the particular path it + # follows. All start at the same step defined by SPLIT_ROUTE and start + # with the same inference_state. + split_graphs = [ + SubGraph( + inf=copy.deepcopy(inference_state), + step=self.router.route[self.router.SPLIT_ROUTE], + end=[self.router.JOIN_ROUTE], ) - else: - operator_output = func(*args, **kwargs) - - if isinstance(operator_output, Future): - operator_output = operator_output.result() - - state_update = None - if isinstance(operator_output, tuple): - state_update = operator_output[-1] - operator_output = operator_output[0] + for i in range(len(batches)) + ] - next_step = self.router.next(next_step, self.ops, operator_output) - return next_step, operator_output, state_update + outputs = self._run_sub_graphs( + sub_graph_inputs=batches, sub_graphs=split_graphs + ) + return self.condense_inputs(outputs) def run( self, *args, inference_state: InferenceState, - pipeline_state: PipelineState, **kwargs, ): """ @@ -158,36 +194,56 @@ def run( """ next_step = self.router.START_ROUTE operator_output = None - while next_step != self.router.END_ROUTE: + + # Split Grap Execution (i.e multiple subgraphs) # NOTE: split_route should only appear after the start route node if next_step == self.router.SPLIT_ROUTE: + if operator_output is None: + raise ValueError( + f"{self.router.SPLIT_ROUTE} should appear after " + f"{self.ROUTER.START_ROUTE}" + ) + operator_output = self._apply_split(operator_output, inference_state) next_step = self.router.route[self.router.JOIN_ROUTE] + if next_step == self.router.END_ROUTE: + return operator_output if next_step == self.router.START_ROUTE: - outputs = self._run_next_step( + operator_output = run_func( *args, - next_step=next_step, func=self._scheduler_group.submit, - inference_state=inference_state, operator=self.ops[next_step], - pipeline_state=pipeline_state, + inference_state=inference_state, + pipeline_state=self.pipeline_state, **kwargs, - ) + ).result() + + if isinstance(operator_output, tuple): + operator_output, state_update = ( + operator_output[0], + operator_output[-1], + ) + inference_state.update_state(state_update) + + next_step = self.router.next(next_step, self.ops, operator_output) + else: - outputs = self._run_next_step( - func=self._scheduler_group.submit, - input=operator_output, - next_step=next_step, - inference_state=inference_state, - operator=self.ops[next_step], - pipeline_state=pipeline_state, + # Single graph execution + graph = SubGraph( + inf=copy.deepcopy(inference_state), + step=next_step, + end=[self.router.SPLIT_ROUTE, self.router.END_ROUTE], ) - next_step, operator_output, state_update = outputs - if state_update: - inference_state.update_state(state_update) + operator_output = self._run_sub_graphs( + sub_graph_inputs=[operator_output], sub_graphs=[graph] + )[0] + + inference_state = graph.inf + next_step = graph.step + return operator_output def __call__(self, *args, **kwargs): @@ -204,11 +260,7 @@ def __call__(self, *args, **kwargs): inference_state = InferenceState() inference_state.create_state({}) - if "pipeline_state" in kwargs: - self.pipeline_state = kwargs.get("pipeline_state") - kwargs["inference_state"] = inference_state - kwargs["pipeline_state"] = self.pipeline_state return self.run(*args, **kwargs) diff --git a/src/deepsparse/v2/routers/router.py b/src/deepsparse/v2/routers/router.py index 6b0d851aef..6740f706f1 100644 --- a/src/deepsparse/v2/routers/router.py +++ b/src/deepsparse/v2/routers/router.py @@ -83,8 +83,6 @@ class LinearRouter(Router): def __init__(self, end_route: int, start_route: int = 0): super().__init__(end_route=end_route, start_route=start_route) - self.SPLIT_ROUTE = None - self.JOIN_ROUTE = None _LOGGER.warn("SPLIT and JOIN are not yet supported for the LinearRouter.") def next( diff --git a/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py b/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py index 669c5922a0..cc74ac0996 100644 --- a/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py +++ b/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py @@ -50,7 +50,7 @@ def __init__(self): engine_operator = EngineOperator(...) ... continuous_batching_scheduler = ContinuousBatchingScheduler.get_instance() - continuous_batching_scheduler.add_engine_operator(engine_operator) + continuous_batching_scheduler.add_engine_operator(engine_operator, [1]) super.__init__(...) ``` @@ -58,6 +58,8 @@ def __init__(self): :param max_workers: maximum number of threads to execute at once, default 1 """ + # TODO: If the singleton always returns max_workers 1, should we remove this arg/not + # give the user a choice? def __init__(self, max_workers: int = 1): self._max_workers = max_workers @@ -82,6 +84,8 @@ def get_instance(cls) -> "ContinuousBatchingScheduler": does not exist yet, a scheduler with a single worker thread to schedule all jobs is created and started """ + global _GLOBAL_SCHEDULER + if _GLOBAL_SCHEDULER is not None: return _GLOBAL_SCHEDULER # noqa: F823 @@ -161,8 +165,18 @@ def add_engine_operator( for batch_size in batch_sizes: if batch_size == 1: continue # already added - operator_engines[batch_size] = operator_engines.create_engine( - batch_size=batch_size + + override_model_path = None + # text generation/NLEngineOperator specific; could add generic method + # for all engine_operators, if desired + if hasattr(engine_operator, "override_model_inputs"): + override_model_path = engine_operator.override_model_inputs( + model_path=engine_operator.model_path, batch_size=batch_size + ) + + # will break for internal kv_cache; needs additional argument + operator_engines[batch_size] = engine_operator.create_engine( + batch_size=batch_size, model_path=override_model_path ) self._operators_to_engines[engine_operator] = operator_engines diff --git a/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py b/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py index 86afdf309c..40ff00ca4f 100644 --- a/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py +++ b/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py @@ -71,7 +71,7 @@ def _working_loop(self): ] # run the engine operator with the given engine at the joined batch size - joined_outputs = engine_operator(joined_inputs) + joined_outputs = engine_operator(joined_inputs, inference_state=None) # split outputs and return the results to their respective futures split_outputs = joined_outputs.split() diff --git a/src/deepsparse/v2/text_generation/compile_generated_tokens.py b/src/deepsparse/v2/text_generation/compile_generated_tokens.py index c87436ab3a..630067f8c3 100644 --- a/src/deepsparse/v2/text_generation/compile_generated_tokens.py +++ b/src/deepsparse/v2/text_generation/compile_generated_tokens.py @@ -42,7 +42,7 @@ def run( if finish_reason is not None: in_generation = False - state_update = { # TODO: check if necessary + state_update = { "finished_reason": finished_reason, "generated_tokens": generated_tokens, "generated_logits": generated_logits, diff --git a/src/deepsparse/v2/text_generation/compile_logits.py b/src/deepsparse/v2/text_generation/compile_logits.py index 21bd50e03e..48a7158f66 100644 --- a/src/deepsparse/v2/text_generation/compile_logits.py +++ b/src/deepsparse/v2/text_generation/compile_logits.py @@ -12,9 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any - from deepsparse.v2.operators import Operator +from deepsparse.v2.text_generation.nl_engine_operator import NLEngineOutputs from deepsparse.v2.utils import InferenceState @@ -28,12 +27,13 @@ class CompilePromptLogits(Operator): take prompt logits from each iteration run and update the inference state. """ - def can_operate(self, inp: Any): - if inp.get("in_generation") is None: + def can_operate(self, inp: NLEngineOutputs): + if inp.in_generation is None: return True return False - def run(self, logits, inference_state: InferenceState, **kwargs): + def run(self, inp: NLEngineOutputs, inference_state: InferenceState, **kwargs): + logits = inp.engine_outputs logit_type = "prompt_logits" if inference_state.current_state.get(logit_type) is not None: @@ -44,6 +44,6 @@ def run(self, logits, inference_state: InferenceState, **kwargs): state_update = {logit_type: current_logits} return { - "kv_cache": kwargs.get("kv_cache"), - "tokens": kwargs.get("tokens"), + "kv_cache": inp.kv_cache, + "tokens": inp.tokens, }, state_update diff --git a/src/deepsparse/v2/text_generation/generate_new_token.py b/src/deepsparse/v2/text_generation/generate_new_token.py index 33ab546e39..5bf48bbdbc 100644 --- a/src/deepsparse/v2/text_generation/generate_new_token.py +++ b/src/deepsparse/v2/text_generation/generate_new_token.py @@ -11,12 +11,13 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Sequence, Union +from typing import Sequence, Union import transformers from deepsparse.transformers.pipelines.text_generation import FinishReason from deepsparse.v2.operators import Operator +from deepsparse.v2.text_generation.nl_engine_operator import NLEngineOutputs from deepsparse.v2.utils import InferenceState @@ -30,12 +31,15 @@ def __init__( self.force_max_tokens = force_max_tokens self.tokenizer = tokenizer - def can_operate(self, inp: Any): - if inp.get("in_generation"): + def can_operate(self, inp: NLEngineOutputs): + if inp.in_generation: return True return False - def run(self, logits, kv_cache, inference_state: InferenceState, **kwargs): + def run(self, inp: NLEngineOutputs, inference_state: InferenceState, **kwargs): + logits = inp.engine_outputs + kv_cache = inp.kv_cache + token_generator = inference_state.current_state.get("token_generator") token = token_generator.generate(logits=logits[0, -1, :]) finish_reason = None diff --git a/src/deepsparse/v2/text_generation/nl_engine_operator.py b/src/deepsparse/v2/text_generation/nl_engine_operator.py index 7549f986d9..d8c80bbaee 100644 --- a/src/deepsparse/v2/text_generation/nl_engine_operator.py +++ b/src/deepsparse/v2/text_generation/nl_engine_operator.py @@ -14,10 +14,13 @@ import copy import os -from typing import Any, List, Tuple +from pathlib import Path +from typing import Any, List, Optional, Tuple, Union +import numpy from pydantic import BaseModel, Field +from deepsparse.utils import join_engine_outputs, split_engine_inputs from deepsparse.utils.onnx import ( CACHE_INPUT_PREFIX, overwrite_onnx_model_inputs_for_kv_cache_models, @@ -29,14 +32,76 @@ ) -__all__ = ["NLEngineOperator", "NlEngineInput"] +__all__ = ["NLEngineOperator", "NLEngineInputs"] -class NlEngineInput(BaseModel): - engine_inputs: List = Field(description="engine inputs") +class NLEngineInputs(BaseModel): + engine_inputs: List = Field(description="engine_inputs") kv_cache: Any = Field(description="kv_cache object") tokens: List = Field(description="tokens") - in_generation: bool = Field(description="in_generation", default=None) + in_generation: Any = Field(description="in_generation", default=None) + engine: Optional[Any] = Field( + description="override the engine to run forward pass with", + default=None, + ) + + @classmethod + def join(cls, inputs: List["NLEngineInputs"]) -> "NLEngineInputs": + """ + :param inputs: list of separate EngineOperatorInputs, batch size must be 1 + :return: list of inputs joined into a single input with a multi batch size + """ + all_engine_inputs = [] + all_kv_cache = [] + all_tokens = [] + all_generation = [] + + for engine_input in inputs: + all_engine_inputs.append(engine_input.engine_inputs) + all_kv_cache.append(engine_input.kv_cache) + all_tokens.append(engine_input.tokens) + all_generation.append(engine_input.in_generation) + + for engine_inputs in all_engine_inputs: + if engine_inputs[0].shape[0] != 1: + raise RuntimeError( + "join requires all inputs to have batch size 1, found input with " + f"batch size {engine_inputs[0].shape[0]}" + ) + return cls( + engine_inputs=all_engine_inputs, + tokens=all_tokens, + in_generation=all_generation, + kv_cache=all_kv_cache, + ) + + class Config: + arbitrary_types_allowed = True + + +class NLEngineOutputs(BaseModel): + engine_outputs: Any = Field(description="engine_outputs") + kv_cache: Any = Field(description="kv_cache object") + tokens: List = Field(description="tokens") + in_generation: Any = Field(description="in_generation", default=None) + + def split(self) -> List["NLEngineOutputs"]: + """ + :return: list of the current outputs split to a batch size of 1 each + """ + split_outputs = [ + numpy.expand_dims(self.engine_outputs[i], 0) + for i in range(len(self.engine_outputs)) + ] + return [ + self.__class__( + engine_outputs=split_outputs[i], + kv_cache=self.kv_cache[i], + tokens=self.tokens[i], + in_generation=self.in_generation[i], + ) + for i in range(len(split_outputs)) + ] class NLEngineOperator(EngineOperator): @@ -48,8 +113,8 @@ class NLEngineOperator(EngineOperator): multi-token case. """ - input_schema = NlEngineInput - output_schema = None + input_schema = NLEngineInputs + output_schema = NLEngineOutputs def __init__( self, @@ -59,17 +124,17 @@ def __init__( **kwargs, ): + self.sequence_length = sequence_length + self.input_ids_length = input_ids_length self.kv_cache_data_type = None - ( - onnx_file_path, - output_indices_to_be_cached, - kv_cache_data_type, - ) = overwrite_onnx_model_inputs_for_kv_cache_models( - onnx_file_path=kwargs.get("model_path"), - batch_size=kwargs.get("batch_size", 1), - sequence_length=sequence_length, - input_ids_length=input_ids_length, + self.internal_kv_cache = internal_kv_cache + self.model_path = kwargs.get("model_path") + (onnx_file_path, additional_outputs) = self.override_model_inputs( + self.model_path, batch_size=1, return_additional_outputs=True ) + output_indices_to_be_cached, kv_cache_data_type, = additional_outputs.get( + "output_indices_to_be_cached" + ), additional_outputs.get("kv_cache_data_type") engine_kwargs = kwargs.get("engine_kwargs", {}) if kwargs.get("engine_type", DEEPSPARSE_ENGINE) == DEEPSPARSE_ENGINE: @@ -86,43 +151,95 @@ def __init__( kwargs["engine_kwargs"] = engine_kwargs kwargs["model_path"] = onnx_file_path + super().__init__(**kwargs) - self.input_ids_length = input_ids_length + def override_model_inputs( + self, + model_path: Union[str, Path], + batch_size: int, + return_additional_outputs=False, + ): + """ + Override the model based on the provided batch_size, sequence_length, + and input_ids_length. + + :param model_path: Path to the model + :param batch_size: The batch size to be used for the model + :return: new overwritten model file path. Optionally returns additional outputs + specific to the NLDecoder engine + """ + ( + onnx_file_path, + output_indices_to_be_cached, + kv_cache_data_type, + ) = overwrite_onnx_model_inputs_for_kv_cache_models( + onnx_file_path=model_path, + batch_size=batch_size, + sequence_length=self.sequence_length, + input_ids_length=self.input_ids_length, + ) + if return_additional_outputs: + return onnx_file_path, { + "output_indices_to_be_cached": output_indices_to_be_cached, + "kv_cache_data_type": kv_cache_data_type, + } + return onnx_file_path - def run(self, inp: NlEngineInput, **kwargs) -> Any: + def run(self, inp: NLEngineInputs, **kwargs) -> NLEngineOutputs: engine_input = inp.engine_inputs kv_cache = inp.kv_cache - inputs = self._add_kv_cache_to_input(engine_input, kv_cache) - if bool(kv_cache.engine_internal_cache): + split = True + if not isinstance(kv_cache, list): + split = False + kv_cache = [kv_cache] + engine_input = [engine_input] + + inputs = list(map(self._add_kv_cache_to_input, engine_input, kv_cache)) + + if bool(kv_cache[0].engine_internal_cache): # conventionally, before dispatching # inputs to the engine, we validate them # if val_inp=True. However, in this case # we want to pass the empty kv cache inputs # (batch_size=0) to the engine. Therefore, # we skip the validation + + # Internal kv_cache works for batch_size of 1 atm out = self.engine._eng_net.execute_list_out( - inputs, kv_cache.engine_internal_cache + inputs[0], kv_cache[0].engine_internal_cache ) else: # run the engine without the LIB.kv_cache object + # stack multiple batch inputs along the batch dimension + inputs = join_engine_outputs(inputs, len(inputs)) out = ( super() - .run(EngineOperatorInputs(engine_inputs=inputs), **kwargs) + .run( + EngineOperatorInputs(engine_inputs=inputs, engine=inp.engine), + **kwargs, + ) .get("engine_outputs") ) + # logits should be stacked along batch dim + # kv_cache_state should be a list where each dim 0 is batch_size logits, *kv_cache_state = out - self._update_kv_cache( - kv_cache_state=kv_cache_state, - input_ids_len=self.input_ids_length, - kv_cache=kv_cache, - ) + kv_cache_state, _ = split_engine_inputs(kv_cache_state, 1) + + if len(kv_cache_state) > 0: + for i in range(len(kv_cache)): + self._update_kv_cache( + kv_cache_state=kv_cache_state[i], kv_cache=kv_cache[i] + ) + else: + # internal kv cache case + self._update_kv_cache(kv_cache=kv_cache[0]) output = { - "logits": logits, - "kv_cache": kv_cache, + "engine_outputs": logits, + "kv_cache": kv_cache if split else kv_cache[0], "tokens": inp.tokens, "in_generation": inp.in_generation, } @@ -137,9 +254,9 @@ def _add_kv_cache_to_input(self, engine_input, kv_cache): new_inp = [kv_cache_state[name] for name in self.engine.input_names] return new_inp - def _update_kv_cache(self, kv_cache_state, input_ids_len, kv_cache): + def _update_kv_cache(self, kv_cache, kv_cache_state=None): if bool(kv_cache.engine_internal_cache): - kv_cache.total_num_processed_tokens += input_ids_len + kv_cache.total_num_processed_tokens += self.input_ids_length return kv_cache_state = { @@ -147,10 +264,7 @@ def _update_kv_cache(self, kv_cache_state, input_ids_len, kv_cache): for name, array in zip(self.onnx_input_names_cached, kv_cache_state) } - kv_cache.update( - state=kv_cache_state, - input_ids_len=input_ids_len, - ) + kv_cache.update(state=kv_cache_state, input_ids_len=self.input_ids_length) @property def onnx_input_names_no_cache(self) -> List[str]: diff --git a/src/deepsparse/v2/text_generation/pipeline.py b/src/deepsparse/v2/text_generation/pipeline.py index 5ab73f7a48..ae7334cffd 100644 --- a/src/deepsparse/v2/text_generation/pipeline.py +++ b/src/deepsparse/v2/text_generation/pipeline.py @@ -12,14 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Dict, Optional +import logging +from typing import Dict, List, Optional from deepsparse.transformers.helpers import setup_transformers_pipeline from deepsparse.transformers.utils.helpers import process_generation_config from deepsparse.utils import split_engine_inputs +from deepsparse.v2.operators import EngineOperator from deepsparse.v2.pipeline import Pipeline from deepsparse.v2.routers import GraphRouter -from deepsparse.v2.schedulers import OperatorScheduler +from deepsparse.v2.schedulers import ContinuousBatchingScheduler, OperatorScheduler from deepsparse.v2.text_generation import ( AutoRegressiveOperatorPreprocess, CompileGeneratedTokens, @@ -39,6 +41,9 @@ from deepsparse.v2.utils import PipelineState +_LOGGER = logging.getLogger(__name__) + + class TextGenerationPipeline(Pipeline): def __init__( self, @@ -48,6 +53,7 @@ def __init__( internal_kv_cache: bool = True, force_max_tokens: bool = False, generation_config=None, + continuous_batch_sizes: Optional[List[int]] = None, engine_kwargs: Optional[Dict] = None, ): ( @@ -133,6 +139,20 @@ def __init__( compile_generated_tokens = CompileGeneratedTokens() join_output = JoinOutput(tokenizer=self.tokenizer) + # TODO: do we want to support lists for different engines? + continuous_batching_scheduler = None + if continuous_batch_sizes: + if internal_kv_cache: + _LOGGER.warn( + "internal kv_cache is currently not supported with continuous ", + "batching", + ) + else: + continuous_batching_scheduler = self._get_continuous_batching_scheduler( + batch_sizes=continuous_batch_sizes, + engines=[single_engine_operator, multi_engine_operator], + ) + ops = { "process_input": process_inputs, "single_engine": single_engine_operator, @@ -183,7 +203,11 @@ def __init__( ) scheduler = [OperatorScheduler()] super().__init__( - ops=ops, router=router, schedulers=scheduler, pipeline_state=pipeline_state + ops=ops, + router=router, + schedulers=scheduler, + pipeline_state=pipeline_state, + continuous_batching_scheduler=continuous_batching_scheduler, ) def expand_inputs(self, items, batch_size): @@ -194,3 +218,21 @@ def expand_inputs(self, items, batch_size): def condense_inputs(self, *args, **kwargs): return args[0], kwargs + + def _get_continuous_batching_scheduler( + self, batch_sizes: List[int], engines: List[EngineOperator] + ) -> ContinuousBatchingScheduler: + """ + Fetch the continuous batching scheduler. Requires adding the EngineOperator + that will run through the scheduler. + + :param batch_sizes: List of batch sizes to be used by the models + :param engine: List of EngineOperators which should be scheduled using the + continuous batching scheduler + + :returns: ContinuousBatchingScheduler + """ + continuous_batching_scheduler = ContinuousBatchingScheduler.get_instance() + for op in engines: + continuous_batching_scheduler.add_engine_operator(op, batch_sizes) + return continuous_batching_scheduler diff --git a/src/deepsparse/v2/utils/__init__.py b/src/deepsparse/v2/utils/__init__.py index 358405d7af..75935a9729 100644 --- a/src/deepsparse/v2/utils/__init__.py +++ b/src/deepsparse/v2/utils/__init__.py @@ -13,5 +13,9 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from .helpers import * from .state import * from .types import * + + +from .data import * # isort:skip diff --git a/src/deepsparse/v2/utils/data.py b/src/deepsparse/v2/utils/data.py new file mode 100644 index 0000000000..40402734cf --- /dev/null +++ b/src/deepsparse/v2/utils/data.py @@ -0,0 +1,39 @@ +# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import Any, List + +from deepsparse.v2.utils import InferenceState + + +__all__ = ["SubGraph"] + + +@dataclass +class SubGraph: + """ + Helper dataclass to store information about each running sub graph. + """ + + step: int + inf: InferenceState + end: List[str] + output: Any = None + + def parse_output(self, operator_output: Any): + if isinstance(operator_output, tuple): + operator_output, state_update = operator_output[0], operator_output[-1] + self.inf.update_state(state_update) + return operator_output diff --git a/src/deepsparse/v2/utils/helpers.py b/src/deepsparse/v2/utils/helpers.py new file mode 100644 index 0000000000..1f4bedc6c9 --- /dev/null +++ b/src/deepsparse/v2/utils/helpers.py @@ -0,0 +1,37 @@ +# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from typing import Any, Callable + + +__all__ = ["run_func"] + + +def run_func( + *args, + func: Callable, + inp: Any = None, + **kwargs, +): + """ + Generic function to run a given Callable. + """ + if inp: + output = ( + func(*args, **kwargs, **inp) + if isinstance(inp, dict) + else func(inp, *args, **kwargs) + ) + else: + output = func(*args, **kwargs) + return output diff --git a/tests/deepsparse/v2/integration_tests/test_llms.py b/tests/deepsparse/v2/integration_tests/test_llms.py index 34a8f7a258..c53899f30c 100644 --- a/tests/deepsparse/v2/integration_tests/test_llms.py +++ b/tests/deepsparse/v2/integration_tests/test_llms.py @@ -135,7 +135,7 @@ def test_ort_single_token_prefill(self, setup): pipeline = self.get_pipeline( prompt_sequence_length=1, - engine_type="onnxruntime", + engine_kwargs={"engine_type": "onnxruntime"}, ) output = pipeline( prompt=self.prompt, @@ -163,7 +163,7 @@ def test_ort_multi_token_prefill(self, setup): "Cannot run ORT pipeline with the internal deepsparse cache enabled." ) pipeline = self.get_pipeline( - engine_type="onnxruntime", + engine_kwargs={"engine_type": "onnxruntime"}, ) output = pipeline( prompt=self.prompt, @@ -244,7 +244,7 @@ def test_inference_no_kv_cache_ort(self, setup): def _test_inference_no_kv_cache(self, engine_type): model_path_no_cache = self._get_model_path_no_cache() pipeline = self.get_pipeline( - model_path=model_path_no_cache, engine_type=engine_type + model_path=model_path_no_cache, engine_kwargs={"engine_type": engine_type} ) assert not pipeline.cache_support_enabled, ( "This pipeline test inference using non-kv cache " diff --git a/tests/deepsparse/v2/unit/text_generation/conftest.py b/tests/deepsparse/v2/unit/text_generation/conftest.py index 5d8483e5f6..3840a9bb0a 100644 --- a/tests/deepsparse/v2/unit/text_generation/conftest.py +++ b/tests/deepsparse/v2/unit/text_generation/conftest.py @@ -19,15 +19,14 @@ import pytest from deepsparse.transformers.helpers import get_deployment_path -from deepsparse.transformers.pipelines.text_generation import TextGenerationInput +from deepsparse.transformers.pipelines.text_generation import ( + GenerationDefaults, + TextGenerationInput, +) from deepsparse.transformers.utils import DecoderKVCache from deepsparse.transformers.utils.helpers import initialize_kv_cache_state from deepsparse.v2 import InferenceState, PipelineState -from deepsparse.v2.text_generation import ( - GenerationDefaults, - NLEngineOperator, - TokenGeneratorOperator, -) +from deepsparse.v2.text_generation import NLEngineOperator, TokenGeneratorOperator @pytest.fixture(scope="module") diff --git a/tests/deepsparse/v2/unit/text_generation/test_misc.py b/tests/deepsparse/v2/unit/text_generation/test_misc.py index caa0cc2efd..f215e2aedb 100644 --- a/tests/deepsparse/v2/unit/text_generation/test_misc.py +++ b/tests/deepsparse/v2/unit/text_generation/test_misc.py @@ -13,16 +13,23 @@ # limitations under the License. from deepsparse.v2.text_generation import CompilePromptLogits +from deepsparse.v2.text_generation.nl_engine_operator import NLEngineOutputs -def test_compile_logits(mock_logits, mock_inference_state): +def test_compile_logits(mock_logits, mock_inference_state, mock_tokens, mock_kv_cache): mock_inference_state.update_state({"prompt_logits": [mock_logits]}) compile_prompt_logits = CompilePromptLogits() # Can operate as long as we're not in generation but in prompt_inference. This # can_operate() will check for the `in_generation` flag in the input. - assert compile_prompt_logits.can_operate({}) + inp = NLEngineOutputs( + engine_outputs=mock_logits, + tokens=mock_tokens, + kv_cache=mock_kv_cache, + in_generation=None, + ) + assert compile_prompt_logits.can_operate(inp=inp) output, state = compile_prompt_logits.run( - logits=mock_logits, inference_state=mock_inference_state + inp=inp, inference_state=mock_inference_state ) # The CompilePromptLogits is responsible for updating a list of prompt logits # calculated at each step during prompt inference. After one step of running this diff --git a/tests/deepsparse/v2/unit/text_generation/test_process_inputs.py b/tests/deepsparse/v2/unit/text_generation/test_process_inputs.py index be59db7475..02f4540c44 100644 --- a/tests/deepsparse/v2/unit/text_generation/test_process_inputs.py +++ b/tests/deepsparse/v2/unit/text_generation/test_process_inputs.py @@ -12,10 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from deepsparse.v2.text_generation import ( - GenerationDefaults, - ProcessInputsTextGeneration, -) +from deepsparse.transformers.pipelines.text_generation import GenerationDefaults +from deepsparse.v2.text_generation import ProcessInputsTextGeneration def test_process_inputs( diff --git a/tests/deepsparse/v2/unit/text_generation/test_single_token_engine.py b/tests/deepsparse/v2/unit/text_generation/test_single_token_engine.py index 335a28fbe3..19bb4d1c4a 100644 --- a/tests/deepsparse/v2/unit/text_generation/test_single_token_engine.py +++ b/tests/deepsparse/v2/unit/text_generation/test_single_token_engine.py @@ -16,7 +16,7 @@ from deepsparse.v2.text_generation import ( AutoRegressiveOperatorPreprocess, - NlEngineInput, + NLEngineInputs, ) @@ -89,10 +89,10 @@ def test_run_single_token_engine_once( numpy.array([[0]]), numpy.array([[[[0, 0, 0, 0, 1]]]]), ] - inputs = NlEngineInput( + inputs = NLEngineInputs( engine_inputs=mock_engine_inputs, kv_cache=mock_kv_cache_single_token_engine, tokens=mock_engine_inputs[0].tolist(), ) output = single_token_engine_no_internal_cache.run(inputs) - assert output.get("logits") is not None + assert output.get("engine_outputs") is not None diff --git a/tests/deepsparse/v2/unit/text_generation/test_token_generation.py b/tests/deepsparse/v2/unit/text_generation/test_token_generation.py index fbd9e06778..d04f863171 100644 --- a/tests/deepsparse/v2/unit/text_generation/test_token_generation.py +++ b/tests/deepsparse/v2/unit/text_generation/test_token_generation.py @@ -18,6 +18,7 @@ PrepareGeneration, TokenGeneratorOperator, ) +from deepsparse.v2.text_generation.nl_engine_operator import NLEngineOutputs def test_prep_for_generation( @@ -68,6 +69,7 @@ def test_generate_new_token( mock_kv_cache, mock_inference_state, mock_logits, + mock_tokens, ): """ This test is responsible for testing the GenerateNewTokenOperator, which generates @@ -84,8 +86,14 @@ def test_generate_new_token( "generated_tokens": [mock_token_generator.tokens], } ) + inp = NLEngineOutputs( + engine_outputs=mock_logits, + tokens=mock_tokens, + kv_cache=mock_kv_cache, + in_generation=True, + ) outputs, state = generate_new_token.run( - logits=mock_logits, kv_cache=mock_kv_cache, inference_state=mock_inference_state + inp=inp, inference_state=mock_inference_state ) # The new_token generated/returned by ths operator should match the last token in # token_generator