From 9e187c39603da5dd1cece0a1637c24893328d6bd Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Tue, 14 Nov 2023 18:26:25 -0500 Subject: [PATCH 01/15] update split/join --- src/deepsparse/v2/pipeline.py | 79 ++++++++++++------- src/deepsparse/v2/text_generation/pipeline.py | 2 +- 2 files changed, 53 insertions(+), 28 deletions(-) diff --git a/src/deepsparse/v2/pipeline.py b/src/deepsparse/v2/pipeline.py index f56680d2b9..6e78f2af10 100644 --- a/src/deepsparse/v2/pipeline.py +++ b/src/deepsparse/v2/pipeline.py @@ -58,6 +58,7 @@ def __init__( self.schedulers = schedulers self.pipeline_state = pipeline_state self.validate() + self.temp_operator = OperatorScheduler() self._scheduler_group = SchedulerGroup(self.schedulers) @@ -67,22 +68,15 @@ def _run_sequential( inference_state: InferenceState, pipeline_state: PipelineState, start: str, - end: 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, - ) - next_step, operator_output, state_update = outputs - if state_update: - inference_state.update_state(state_update) - inp = operator_output - return inp + return self._run_next_step( + func=self._scheduler_group.submit, + operator=self.ops[start], + next_step=start, + input=inp, + pipeline_state=pipeline_state, + inference_state=inference_state, + ) def _apply_split(self, inp: Any, inference_state: InferenceState): """ @@ -96,18 +90,41 @@ def _apply_split(self, inp: Any, inference_state: InferenceState): 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)) + + inf_list = [copy.deepcopy(inference_state) for x in range(len(batches))] + step = [self.router.route[self.router.SPLIT_ROUTE] for x in range(len(batches))] + current_outputs = [ + run_with_state(inp=batches[i], inference_state=inf_list[i], start=step[i]) + for i 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]) + completed_outptus = [] + while True: + for i in range(len(batches)): + if isinstance(current_outputs[i], Future) and current_outputs[i].done(): + operator_output = current_outputs[i].result() + + if isinstance(operator_output, tuple): + state_update = operator_output[-1] + operator_output = operator_output[0] + inf_list[i].update_state(state_update) + + next_step = self.router.next(step[i], self.ops, operator_output) + step[i] = next_step + if next_step == self.router.JOIN_ROUTE: + current_outputs[i] = operator_output + else: + current_outputs[i] = run_with_state( + inp=operator_output, + inference_state=inf_list[i], + start=next_step, + ) + break + + if not any(isinstance(x, Future) for x in current_outputs): + break + + return self.condense_inputs(current_outputs) def _run_next_step( self, @@ -129,7 +146,9 @@ def _run_next_step( ) else: operator_output = func(*args, **kwargs) + return operator_output + """ if isinstance(operator_output, Future): operator_output = operator_output.result() @@ -140,6 +159,7 @@ def _run_next_step( next_step = self.router.next(next_step, self.ops, operator_output) return next_step, operator_output, state_update + """ def run( self, @@ -185,9 +205,14 @@ def run( pipeline_state=pipeline_state, ) - next_step, operator_output, state_update = outputs - if state_update: + operator_output = outputs.result() + + if isinstance(operator_output, tuple): + state_update = operator_output[-1] + operator_output = operator_output[0] inference_state.update_state(state_update) + + next_step = self.router.next(next_step, self.ops, operator_output) return operator_output def __call__(self, *args, **kwargs): diff --git a/src/deepsparse/v2/text_generation/pipeline.py b/src/deepsparse/v2/text_generation/pipeline.py index 240da04907..933d5d0fc5 100644 --- a/src/deepsparse/v2/text_generation/pipeline.py +++ b/src/deepsparse/v2/text_generation/pipeline.py @@ -183,7 +183,7 @@ def __init__( router = GraphRouter( end_route="STOP", start_route="process_input", route=routes ) - scheduler = [OperatorScheduler()] + scheduler = [OperatorScheduler(max_workers=4)] super().__init__( ops=ops, router=router, schedulers=scheduler, pipeline_state=pipeline_state ) From b729991d2e2c6f9ae2455f03aed49fd148818dfb Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Tue, 14 Nov 2023 18:38:46 -0500 Subject: [PATCH 02/15] use map --- src/deepsparse/v2/pipeline.py | 17 ++++++++++------- src/deepsparse/v2/text_generation/pipeline.py | 2 +- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/src/deepsparse/v2/pipeline.py b/src/deepsparse/v2/pipeline.py index 6e78f2af10..834b31ff4e 100644 --- a/src/deepsparse/v2/pipeline.py +++ b/src/deepsparse/v2/pipeline.py @@ -66,8 +66,8 @@ def _run_sequential( self, inp: Any, inference_state: InferenceState, - pipeline_state: PipelineState, start: str, + pipeline_state: PipelineState, ): return self._run_next_step( func=self._scheduler_group.submit, @@ -94,11 +94,14 @@ def _apply_split(self, inp: Any, inference_state: InferenceState): inf_list = [copy.deepcopy(inference_state) for x in range(len(batches))] step = [self.router.route[self.router.SPLIT_ROUTE] for x in range(len(batches))] - current_outputs = [ - run_with_state(inp=batches[i], inference_state=inf_list[i], start=step[i]) - for i in range(len(batches)) - ] - completed_outptus = [] + current_outputs = list( + map( + run_with_state, + batches, + inf_list, + step + ) + ) while True: for i in range(len(batches)): if isinstance(current_outputs[i], Future) and current_outputs[i].done(): @@ -120,7 +123,7 @@ def _apply_split(self, inp: Any, inference_state: InferenceState): start=next_step, ) break - + if not any(isinstance(x, Future) for x in current_outputs): break diff --git a/src/deepsparse/v2/text_generation/pipeline.py b/src/deepsparse/v2/text_generation/pipeline.py index 933d5d0fc5..240da04907 100644 --- a/src/deepsparse/v2/text_generation/pipeline.py +++ b/src/deepsparse/v2/text_generation/pipeline.py @@ -183,7 +183,7 @@ def __init__( router = GraphRouter( end_route="STOP", start_route="process_input", route=routes ) - scheduler = [OperatorScheduler(max_workers=4)] + scheduler = [OperatorScheduler()] super().__init__( ops=ops, router=router, schedulers=scheduler, pipeline_state=pipeline_state ) From 58ef2b9b4f466b4d01e28cef2fffd099f160fe2e Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Wed, 15 Nov 2023 11:51:42 -0500 Subject: [PATCH 03/15] update --- .../v2/operators/engine_operator.py | 5 +- src/deepsparse/v2/pipeline.py | 35 ++++---- .../continuous_batching_scheduler.py | 4 +- .../utils/continuous_batching_executor.py | 2 +- .../v2/text_generation/compile_logits.py | 12 +-- .../v2/text_generation/generate_new_token.py | 10 ++- .../v2/text_generation/nl_engine_operator.py | 56 +++++++++++-- src/deepsparse/v2/text_generation/pipeline.py | 9 +- .../v2/text_generation/run_async.py | 35 ++++++++ .../v2/text_generation/run_text_gen.py | 84 +++++++++++++++++++ 10 files changed, 210 insertions(+), 42 deletions(-) create mode 100644 src/deepsparse/v2/text_generation/run_async.py create mode 100644 src/deepsparse/v2/text_generation/run_text_gen.py diff --git a/src/deepsparse/v2/operators/engine_operator.py b/src/deepsparse/v2/operators/engine_operator.py index 9ee8d734c5..65a7d6d948 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 model_to_path, join_engine_outputs, split_engine_inputs from deepsparse.v2.operators import Operator @@ -29,7 +29,7 @@ SUPPORTED_PIPELINE_ENGINES = [DEEPSPARSE_ENGINE, ORT_ENGINE] -__all__ = ["EngineOperator"] +__all__ = ["EngineOperator", "EngineOperatorInputs", "EngineOperatorOutputs"] class EngineOperatorInputs(BaseModel): @@ -145,6 +145,7 @@ def create_engine( onnx_file_path = self.model_path engine_args = deepcopy(self._engine_args) engine_args.update(kwargs) + engine_type = self._engine_type.lower() if engine_type == DEEPSPARSE_ENGINE: diff --git a/src/deepsparse/v2/pipeline.py b/src/deepsparse/v2/pipeline.py index 834b31ff4e..86136a3673 100644 --- a/src/deepsparse/v2/pipeline.py +++ b/src/deepsparse/v2/pipeline.py @@ -18,9 +18,9 @@ from functools import partial from typing import Any, Callable, Dict, List, Union -from deepsparse.v2.operators import Operator +from deepsparse.v2.operators import Operator, EngineOperator from deepsparse.v2.routers import Router -from deepsparse.v2.schedulers import OperatorScheduler, SchedulerGroup +from deepsparse.v2.schedulers import OperatorScheduler, SchedulerGroup, ContinuousBatchingScheduler from deepsparse.v2.utils import InferenceState, PipelineState @@ -50,6 +50,7 @@ def __init__( ops: Union[Dict[str, Operator], List[Operator]], router: Router, schedulers: List[OperatorScheduler], + continuous_batching_scheduler: ContinuousBatchingScheduler, pipeline_state: PipelineState = None, ): @@ -57,9 +58,9 @@ def __init__( self.router = router self.schedulers = schedulers self.pipeline_state = pipeline_state + self._continuous_batching_scheduler = continuous_batching_scheduler self.validate() - self.temp_operator = OperatorScheduler() - + self._scheduler_group = SchedulerGroup(self.schedulers) def _run_sequential( @@ -69,8 +70,16 @@ def _run_sequential( start: str, pipeline_state: PipelineState, ): + """ + if isinstance(self.ops[start], EngineOperator): + func = self._continuous_batching_scheduler.submit + inp = self.ops[start].input_schema(**inp) + else: + """ + func = self._scheduler_group.submit + return self._run_next_step( - func=self._scheduler_group.submit, + func=func, operator=self.ops[start], next_step=start, input=inp, @@ -138,8 +147,7 @@ def _run_next_step( **kwargs, ): """ - Generic function to run a given func, process the output and determine the next - step. + Generic function to run a given func. """ if input: operator_output = ( @@ -151,19 +159,6 @@ def _run_next_step( operator_output = func(*args, **kwargs) return operator_output - """ - 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] - - next_step = self.router.next(next_step, self.ops, operator_output) - return next_step, operator_output, state_update - """ - def run( self, *args, diff --git a/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py b/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py index 669c5922a0..97bebfaaab 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__(...) ``` @@ -82,6 +82,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 diff --git a/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py b/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py index 86afdf309c..e610956e7e 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, pipeline_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_logits.py b/src/deepsparse/v2/text_generation/compile_logits.py index 21bd50e03e..237df110c5 100644 --- a/src/deepsparse/v2/text_generation/compile_logits.py +++ b/src/deepsparse/v2/text_generation/compile_logits.py @@ -16,6 +16,7 @@ from deepsparse.v2.operators import Operator from deepsparse.v2.utils import InferenceState +from deepsparse.v2.text_generation.nl_engine_operator import NLEngineOutputs __all__ = ["CompilePromptLogits"] @@ -28,12 +29,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 +46,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..4dac5a32dc 100644 --- a/src/deepsparse/v2/text_generation/generate_new_token.py +++ b/src/deepsparse/v2/text_generation/generate_new_token.py @@ -18,6 +18,7 @@ from deepsparse.transformers.pipelines.text_generation import FinishReason from deepsparse.v2.operators import Operator from deepsparse.v2.utils import InferenceState +from deepsparse.v2.text_generation.nl_engine_operator import NLEngineOutputs __all__ = ["GenerateNewTokenOperator"] @@ -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..331ef2903a 100644 --- a/src/deepsparse/v2/text_generation/nl_engine_operator.py +++ b/src/deepsparse/v2/text_generation/nl_engine_operator.py @@ -14,9 +14,10 @@ import copy import os -from typing import Any, List, Tuple +from typing import Any, List, Tuple, Optional from pydantic import BaseModel, Field +from deepsparse import Engine from deepsparse.utils.onnx import ( CACHE_INPUT_PREFIX, @@ -26,14 +27,52 @@ DEEPSPARSE_ENGINE, EngineOperator, EngineOperatorInputs, + EngineOperatorOutputs ) +from deepsparse.utils import model_to_path, join_engine_outputs, split_engine_inputs -__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: Any = Field(description="in_generation", default=None) + engine: Optional[Engine] = 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 = [engine_input.engine_inputs for engine_input in inputs] + all_kv_cache = [engine_input.kv_cache for engine_input in inputs] + all_tokens = [engine_input.tokens for engine_input in inputs] + all_generation = [engine_input.in_generation for engine_input in inputs] + + 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: bool = Field(description="in_generation", default=None) @@ -48,8 +87,8 @@ class NLEngineOperator(EngineOperator): multi-token case. """ - input_schema = NlEngineInput - output_schema = None + input_schema = NLEngineInputs + output_schema = NLEngineOutputs def __init__( self, @@ -86,11 +125,12 @@ def __init__( kwargs["engine_kwargs"] = engine_kwargs kwargs["model_path"] = onnx_file_path + super().__init__(**kwargs) self.input_ids_length = input_ids_length - def run(self, inp: NlEngineInput, **kwargs) -> Any: + def run(self, inp: NLEngineInputs, **kwargs) -> NLEngineOutputs: engine_input = inp.engine_inputs kv_cache = inp.kv_cache @@ -121,7 +161,7 @@ def run(self, inp: NlEngineInput, **kwargs) -> Any: ) output = { - "logits": logits, + "engine_outputs": logits, "kv_cache": kv_cache, "tokens": inp.tokens, "in_generation": inp.in_generation, diff --git a/src/deepsparse/v2/text_generation/pipeline.py b/src/deepsparse/v2/text_generation/pipeline.py index 240da04907..ee17e4769e 100644 --- a/src/deepsparse/v2/text_generation/pipeline.py +++ b/src/deepsparse/v2/text_generation/pipeline.py @@ -18,7 +18,7 @@ from deepsparse.utils import split_engine_inputs from deepsparse.v2.pipeline import Pipeline from deepsparse.v2.routers import GraphRouter -from deepsparse.v2.schedulers import OperatorScheduler +from deepsparse.v2.schedulers import OperatorScheduler, ContinuousBatchingScheduler from deepsparse.v2.text_generation import ( AutoRegressiveOperatorPreprocess, CompileGeneratedTokens, @@ -135,6 +135,11 @@ def __init__( compile_generated_tokens = CompileGeneratedTokens() join_output = JoinOutput(tokenizer=self.tokenizer) + + continuous_batching_scheduler = ContinuousBatchingScheduler.get_instance() + continuous_batching_scheduler.add_engine_operator(single_engine_operator, [1]) + continuous_batching_scheduler.add_engine_operator(multi_engine_operator, [1]) + ops = { "process_input": process_inputs, "single_engine": single_engine_operator, @@ -185,7 +190,7 @@ 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): diff --git a/src/deepsparse/v2/text_generation/run_async.py b/src/deepsparse/v2/text_generation/run_async.py new file mode 100644 index 0000000000..b59b61ad2c --- /dev/null +++ b/src/deepsparse/v2/text_generation/run_async.py @@ -0,0 +1,35 @@ +# 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. + +import asyncio + +from deepsparse.transformers.pipelines.text_generation import TextGenerationInput +from deepsparse.v2.text_generation.pipeline import TextGenerationPipeline +from deepsparse.v2.utils import InferenceState + + +model_path = "hf:mgoin/TinyStories-1M-deepsparse" +pipeline = TextGenerationPipeline(model_path, prompt_sequence_length=3) + +prompts = [["Hello there!", "The sun shined bright"], ["The dog barked"]] + + +async def func(index): + print("Hello World", index) + inference_state = InferenceState() + inference_state.create_state({}) + pipeline_state = pipeline.pipeline_state + + input_value = TextGenerationInput( + diff --git a/src/deepsparse/v2/text_generation/run_text_gen.py b/src/deepsparse/v2/text_generation/run_text_gen.py new file mode 100644 index 0000000000..c889be08c5 --- /dev/null +++ b/src/deepsparse/v2/text_generation/run_text_gen.py @@ -0,0 +1,84 @@ +# 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. +import asyncio +import time + +import numpy as np +from transformers import AutoModelForCausalLM, AutoTokenizer + +import torch +from deepsparse.transformers.pipelines.text_generation import TextGenerationInput +from deepsparse.v2.text_generation.pipeline import TextGenerationPipeline +from deepsparse.v2.utils.state import InferenceState + + +def create_tokenizer(model_name): + tokenizer = AutoTokenizer.from_pretrained(model_name) + tokenizer.padding_side = "left" + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + return tokenizer + + +def get_ground_truth(prompt): + model = AutoModelForCausalLM.from_pretrained("roneneldan/TinyStories-1M") + tokenizer = create_tokenizer("roneneldan/TinyStories-1M") + + input_ids = tokenizer.encode(prompt, return_tensors="pt") + out = model(input_ids=input_ids) + prompt_logits = out.logits.detach().numpy() + return prompt_logits + + +from transformers import GenerationConfig + + +# output = pipeline(input_value) +# print(out) +""" +async def func(): + inference_state = InferenceState() + inference_state.create_state({}) + pipeline_state = pipeline.pipeline_state + + output = await pipeline.run_async(input_value, pipeline_state=pipeline_state, inference_state=inference_state) + return output +print(asyncio.run(func())) +""" + +model_path = "hf:mgoin/TinyStories-1M-deepsparse" +pipeline = TextGenerationPipeline(model_path, prompt_sequence_length=3, engine_kwargs={"engine_type": "onnxruntime"}) + + +def run_requests(): + prompts = [["Hello there!", "How are you?"]] + outputs = [] + for i in range(len(prompts)): + input_value = TextGenerationInput( + prompt=prompts[i], + generation_kwargs={ + "do_sample": False, + "max_length": 20, + }, + ) + output = pipeline(input_value) + yield output + + +output = run_requests() +for x in output: + for g in x.generations: + print("\n") + print(g) From e04a97183b690f5524c102f9e8ac8fce4a4b9b3f Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Wed, 15 Nov 2023 20:18:29 -0500 Subject: [PATCH 04/15] run end-to-end --- .../v2/operators/engine_operator.py | 6 +- src/deepsparse/v2/pipeline.py | 25 ++--- .../continuous_batching_scheduler.py | 2 +- .../utils/continuous_batching_executor.py | 4 +- .../v2/text_generation/compile_logits.py | 2 +- .../v2/text_generation/generate_new_token.py | 2 +- .../v2/text_generation/nl_engine_operator.py | 104 +++++++++++++----- src/deepsparse/v2/text_generation/pipeline.py | 11 +- .../v2/text_generation/run_async.py | 35 ------ .../v2/text_generation/run_text_gen.py | 84 -------------- 10 files changed, 100 insertions(+), 175 deletions(-) delete mode 100644 src/deepsparse/v2/text_generation/run_async.py delete mode 100644 src/deepsparse/v2/text_generation/run_text_gen.py diff --git a/src/deepsparse/v2/operators/engine_operator.py b/src/deepsparse/v2/operators/engine_operator.py index 65a7d6d948..4620008cc1 100644 --- a/src/deepsparse/v2/operators/engine_operator.py +++ b/src/deepsparse/v2/operators/engine_operator.py @@ -13,14 +13,14 @@ # limitations under the License. from copy import deepcopy -from typing import Dict, List, Optional, Union +from typing import Any, Dict, List, Optional, Union from pydantic import BaseModel, Field from deepsparse import Context as EngineContext from deepsparse import Engine, MultiModelEngine, Scheduler from deepsparse.benchmark import ORTEngine -from deepsparse.utils import model_to_path, join_engine_outputs, split_engine_inputs +from deepsparse.utils import join_engine_outputs, model_to_path, split_engine_inputs from deepsparse.v2.operators import Operator @@ -34,7 +34,7 @@ class EngineOperatorInputs(BaseModel): engine_inputs: List = Field(description="engine_inputs") - engine: Optional[Engine] = Field( + engine: Optional[Any] = Field( description="override the engine to run forward pass with", default=None, ) diff --git a/src/deepsparse/v2/pipeline.py b/src/deepsparse/v2/pipeline.py index 86136a3673..16e4b0f781 100644 --- a/src/deepsparse/v2/pipeline.py +++ b/src/deepsparse/v2/pipeline.py @@ -18,9 +18,13 @@ from functools import partial from typing import Any, Callable, Dict, List, Union -from deepsparse.v2.operators import Operator, EngineOperator +from deepsparse.v2.operators import EngineOperator, Operator from deepsparse.v2.routers import Router -from deepsparse.v2.schedulers import OperatorScheduler, SchedulerGroup, ContinuousBatchingScheduler +from deepsparse.v2.schedulers import ( + ContinuousBatchingScheduler, + OperatorScheduler, + SchedulerGroup, +) from deepsparse.v2.utils import InferenceState, PipelineState @@ -60,7 +64,7 @@ def __init__( self.pipeline_state = pipeline_state self._continuous_batching_scheduler = continuous_batching_scheduler self.validate() - + self._scheduler_group = SchedulerGroup(self.schedulers) def _run_sequential( @@ -70,13 +74,11 @@ def _run_sequential( start: str, pipeline_state: PipelineState, ): - """ if isinstance(self.ops[start], EngineOperator): func = self._continuous_batching_scheduler.submit inp = self.ops[start].input_schema(**inp) else: - """ - func = self._scheduler_group.submit + func = self._scheduler_group.submit return self._run_next_step( func=func, @@ -103,14 +105,7 @@ def _apply_split(self, inp: Any, inference_state: InferenceState): inf_list = [copy.deepcopy(inference_state) for x in range(len(batches))] step = [self.router.route[self.router.SPLIT_ROUTE] for x in range(len(batches))] - current_outputs = list( - map( - run_with_state, - batches, - inf_list, - step - ) - ) + current_outputs = list(map(run_with_state, batches, inf_list, step)) while True: for i in range(len(batches)): if isinstance(current_outputs[i], Future) and current_outputs[i].done(): @@ -132,7 +127,7 @@ def _apply_split(self, inp: Any, inference_state: InferenceState): start=next_step, ) break - + if not any(isinstance(x, Future) for x in current_outputs): break diff --git a/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py b/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py index 97bebfaaab..aad885157a 100644 --- a/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py +++ b/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py @@ -163,7 +163,7 @@ 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( + operator_engines[batch_size] = engine_operator.create_engine( batch_size=batch_size ) diff --git a/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py b/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py index e610956e7e..57151e64c3 100644 --- a/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py +++ b/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py @@ -71,7 +71,9 @@ def _working_loop(self): ] # run the engine operator with the given engine at the joined batch size - joined_outputs = engine_operator(joined_inputs, inference_state=None, pipeline_state=None) + joined_outputs = engine_operator( + joined_inputs, inference_state=None, pipeline_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_logits.py b/src/deepsparse/v2/text_generation/compile_logits.py index 237df110c5..1c00261cb6 100644 --- a/src/deepsparse/v2/text_generation/compile_logits.py +++ b/src/deepsparse/v2/text_generation/compile_logits.py @@ -15,8 +15,8 @@ from typing import Any from deepsparse.v2.operators import Operator -from deepsparse.v2.utils import InferenceState from deepsparse.v2.text_generation.nl_engine_operator import NLEngineOutputs +from deepsparse.v2.utils import InferenceState __all__ = ["CompilePromptLogits"] diff --git a/src/deepsparse/v2/text_generation/generate_new_token.py b/src/deepsparse/v2/text_generation/generate_new_token.py index 4dac5a32dc..c7e0836c78 100644 --- a/src/deepsparse/v2/text_generation/generate_new_token.py +++ b/src/deepsparse/v2/text_generation/generate_new_token.py @@ -17,8 +17,8 @@ from deepsparse.transformers.pipelines.text_generation import FinishReason from deepsparse.v2.operators import Operator -from deepsparse.v2.utils import InferenceState from deepsparse.v2.text_generation.nl_engine_operator import NLEngineOutputs +from deepsparse.v2.utils import InferenceState __all__ = ["GenerateNewTokenOperator"] diff --git a/src/deepsparse/v2/text_generation/nl_engine_operator.py b/src/deepsparse/v2/text_generation/nl_engine_operator.py index 331ef2903a..2863e6829a 100644 --- a/src/deepsparse/v2/text_generation/nl_engine_operator.py +++ b/src/deepsparse/v2/text_generation/nl_engine_operator.py @@ -14,11 +14,13 @@ import copy import os -from typing import Any, List, Tuple, Optional +from typing import Any, List, Optional, Tuple +import numpy from pydantic import BaseModel, Field -from deepsparse import Engine +from deepsparse import Engine +from deepsparse.utils import join_engine_outputs, model_to_path, split_engine_inputs from deepsparse.utils.onnx import ( CACHE_INPUT_PREFIX, overwrite_onnx_model_inputs_for_kv_cache_models, @@ -27,9 +29,8 @@ DEEPSPARSE_ENGINE, EngineOperator, EngineOperatorInputs, - EngineOperatorOutputs + EngineOperatorOutputs, ) -from deepsparse.utils import model_to_path, join_engine_outputs, split_engine_inputs __all__ = ["NLEngineOperator", "NLEngineInputs"] @@ -40,23 +41,27 @@ class NLEngineInputs(BaseModel): kv_cache: Any = Field(description="kv_cache object") tokens: List = Field(description="tokens") in_generation: Any = Field(description="in_generation", default=None) - engine: Optional[Engine] = Field( + 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 = [] - all_engine_inputs = [engine_input.engine_inputs for engine_input in inputs] - all_kv_cache = [engine_input.kv_cache for engine_input in inputs] - all_tokens = [engine_input.tokens for engine_input in inputs] - all_generation = [engine_input.in_generation for engine_input in inputs] + 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: @@ -64,8 +69,12 @@ def join(cls, inputs: List["NLEngineInputs"]) -> "NLEngineInputs": "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) + 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 @@ -75,7 +84,25 @@ 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: bool = Field(description="in_generation", default=None) + 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): @@ -134,37 +161,57 @@ 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): + if not isinstance(kv_cache, list): + 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 + print(len(inputs)) + 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 of dim 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: + self._update_kv_cache(kv_cache=kv_cache[0]) output = { "engine_outputs": logits, "kv_cache": kv_cache, - "tokens": inp.tokens, - "in_generation": inp.in_generation, + "tokens": [inp.tokens] if not isinstance(inp.tokens, list) else inp.tokens, + "in_generation": [inp.in_generation] + if not isinstance(inp.in_generation, list) + else inp.in_generation, } return output @@ -177,9 +224,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 = { @@ -187,10 +234,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 ee17e4769e..dbfc3d8464 100644 --- a/src/deepsparse/v2/text_generation/pipeline.py +++ b/src/deepsparse/v2/text_generation/pipeline.py @@ -18,7 +18,7 @@ from deepsparse.utils import split_engine_inputs from deepsparse.v2.pipeline import Pipeline from deepsparse.v2.routers import GraphRouter -from deepsparse.v2.schedulers import OperatorScheduler, ContinuousBatchingScheduler +from deepsparse.v2.schedulers import ContinuousBatchingScheduler, OperatorScheduler from deepsparse.v2.text_generation import ( AutoRegressiveOperatorPreprocess, CompileGeneratedTokens, @@ -135,10 +135,9 @@ def __init__( compile_generated_tokens = CompileGeneratedTokens() join_output = JoinOutput(tokenizer=self.tokenizer) - continuous_batching_scheduler = ContinuousBatchingScheduler.get_instance() continuous_batching_scheduler.add_engine_operator(single_engine_operator, [1]) - continuous_batching_scheduler.add_engine_operator(multi_engine_operator, [1]) + continuous_batching_scheduler.add_engine_operator(multi_engine_operator, [2, 4]) ops = { "process_input": process_inputs, @@ -190,7 +189,11 @@ def __init__( ) scheduler = [OperatorScheduler()] super().__init__( - ops=ops, router=router, schedulers=scheduler, pipeline_state=pipeline_state, continuous_batching_scheduler=continuous_batching_scheduler + ops=ops, + router=router, + schedulers=scheduler, + pipeline_state=pipeline_state, + continuous_batching_scheduler=continuous_batching_scheduler, ) def expand_inputs(self, items, batch_size): diff --git a/src/deepsparse/v2/text_generation/run_async.py b/src/deepsparse/v2/text_generation/run_async.py deleted file mode 100644 index b59b61ad2c..0000000000 --- a/src/deepsparse/v2/text_generation/run_async.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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. - -import asyncio - -from deepsparse.transformers.pipelines.text_generation import TextGenerationInput -from deepsparse.v2.text_generation.pipeline import TextGenerationPipeline -from deepsparse.v2.utils import InferenceState - - -model_path = "hf:mgoin/TinyStories-1M-deepsparse" -pipeline = TextGenerationPipeline(model_path, prompt_sequence_length=3) - -prompts = [["Hello there!", "The sun shined bright"], ["The dog barked"]] - - -async def func(index): - print("Hello World", index) - inference_state = InferenceState() - inference_state.create_state({}) - pipeline_state = pipeline.pipeline_state - - input_value = TextGenerationInput( - diff --git a/src/deepsparse/v2/text_generation/run_text_gen.py b/src/deepsparse/v2/text_generation/run_text_gen.py deleted file mode 100644 index c889be08c5..0000000000 --- a/src/deepsparse/v2/text_generation/run_text_gen.py +++ /dev/null @@ -1,84 +0,0 @@ -# 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. -import asyncio -import time - -import numpy as np -from transformers import AutoModelForCausalLM, AutoTokenizer - -import torch -from deepsparse.transformers.pipelines.text_generation import TextGenerationInput -from deepsparse.v2.text_generation.pipeline import TextGenerationPipeline -from deepsparse.v2.utils.state import InferenceState - - -def create_tokenizer(model_name): - tokenizer = AutoTokenizer.from_pretrained(model_name) - tokenizer.padding_side = "left" - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - - return tokenizer - - -def get_ground_truth(prompt): - model = AutoModelForCausalLM.from_pretrained("roneneldan/TinyStories-1M") - tokenizer = create_tokenizer("roneneldan/TinyStories-1M") - - input_ids = tokenizer.encode(prompt, return_tensors="pt") - out = model(input_ids=input_ids) - prompt_logits = out.logits.detach().numpy() - return prompt_logits - - -from transformers import GenerationConfig - - -# output = pipeline(input_value) -# print(out) -""" -async def func(): - inference_state = InferenceState() - inference_state.create_state({}) - pipeline_state = pipeline.pipeline_state - - output = await pipeline.run_async(input_value, pipeline_state=pipeline_state, inference_state=inference_state) - return output -print(asyncio.run(func())) -""" - -model_path = "hf:mgoin/TinyStories-1M-deepsparse" -pipeline = TextGenerationPipeline(model_path, prompt_sequence_length=3, engine_kwargs={"engine_type": "onnxruntime"}) - - -def run_requests(): - prompts = [["Hello there!", "How are you?"]] - outputs = [] - for i in range(len(prompts)): - input_value = TextGenerationInput( - prompt=prompts[i], - generation_kwargs={ - "do_sample": False, - "max_length": 20, - }, - ) - output = pipeline(input_value) - yield output - - -output = run_requests() -for x in output: - for g in x.generations: - print("\n") - print(g) From ff01185ef0a58d83c1984e991cf7a309a13c4deb Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Wed, 15 Nov 2023 22:24:55 -0500 Subject: [PATCH 05/15] clean-up --- src/deepsparse/v2/operators/operator.py | 5 +- src/deepsparse/v2/pipeline.py | 134 +++++++++++------- src/deepsparse/v2/routers/router.py | 2 - .../compile_generated_tokens.py | 2 +- .../v2/text_generation/compile_logits.py | 2 - .../v2/text_generation/generate_new_token.py | 2 +- .../v2/text_generation/nl_engine_operator.py | 8 +- 7 files changed, 86 insertions(+), 69 deletions(-) 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 16e4b0f781..7d9a041db5 100644 --- a/src/deepsparse/v2/pipeline.py +++ b/src/deepsparse/v2/pipeline.py @@ -31,6 +31,23 @@ __all__ = ["Pipeline"] +class SubGraph: + def __init__(self, step, inf): + self.step = step + self.inf = inf + self.output = None + + def update_next_step(self): + if isinstance(self.output, tuple): + state_update = self.output[-1] + operator_output = self.output[0] + self.inf.update_state(state_update) + else: + operator_output = self.output + + return operator_output + + class Pipeline(Operator): """ Pipeline accepts a series of operators, schedulers, and a router. Calling a pipeline @@ -67,25 +84,23 @@ def __init__( self._scheduler_group = SchedulerGroup(self.schedulers) - def _run_sequential( + def _run_next( self, inp: Any, inference_state: InferenceState, - start: str, - pipeline_state: PipelineState, + next_step: str, ): - if isinstance(self.ops[start], EngineOperator): + if isinstance(self.ops[next_step], EngineOperator): func = self._continuous_batching_scheduler.submit - inp = self.ops[start].input_schema(**inp) + inp = self.ops[next_step].input_schema(**inp) else: func = self._scheduler_group.submit - return self._run_next_step( + return self._run_func( func=func, - operator=self.ops[start], - next_step=start, - input=inp, - pipeline_state=pipeline_state, + operator=self.ops[next_step], + inp=inp, + pipeline_state=self.pipeline_state, inference_state=inference_state, ) @@ -98,57 +113,68 @@ 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, - ) - inf_list = [copy.deepcopy(inference_state) for x in range(len(batches))] - step = [self.router.route[self.router.SPLIT_ROUTE] for x in range(len(batches))] - current_outputs = list(map(run_with_state, batches, inf_list, step)) + sub_graphs = [ + SubGraph( + inf=copy.deepcopy(inference_state), + step=self.router.route[self.router.SPLIT_ROUTE], + ) + for i in range(len(batches)) + ] + + for i in range(len(sub_graphs)): + sub_graphs[i].output = self._run_next( + batches[i], sub_graphs[i].inf, sub_graphs[i].step + ) + while True: - for i in range(len(batches)): - if isinstance(current_outputs[i], Future) and current_outputs[i].done(): - operator_output = current_outputs[i].result() + for i in range(len(sub_graphs)): + current_subgraph = sub_graphs[i] - if isinstance(operator_output, tuple): - state_update = operator_output[-1] - operator_output = operator_output[0] - inf_list[i].update_state(state_update) + if ( + isinstance(current_subgraph.output, Future) + and current_subgraph.output.done() + ): + operator_output = current_subgraph.output.result() + current_subgraph.output = operator_output + + operator_output = current_subgraph.update_next_step() + + next_step = self.router.next( + current_subgraph.step, self.ops, operator_output + ) + current_subgraph.step = next_step - next_step = self.router.next(step[i], self.ops, operator_output) - step[i] = next_step if next_step == self.router.JOIN_ROUTE: - current_outputs[i] = operator_output + current_subgraph.output = operator_output else: - current_outputs[i] = run_with_state( + current_subgraph.output = self._run_next( inp=operator_output, - inference_state=inf_list[i], - start=next_step, + inference_state=current_subgraph.inf, + next_step=next_step, ) break - if not any(isinstance(x, Future) for x in current_outputs): + if not any(isinstance(x.output, Future) for x in sub_graphs): break - return self.condense_inputs(current_outputs) + return self.condense_inputs([x.output for x in sub_graphs]) - def _run_next_step( + def _run_func( self, *args, func: Callable, - next_step: Union[str, int], - input: Any = None, + inp: Any = None, **kwargs, ): """ Generic function to run a given func. """ - if input: + if inp: operator_output = ( - func(*args, **kwargs, **input) - if isinstance(input, dict) - else func(input, *args, **kwargs) + func(*args, **kwargs, **inp) + if isinstance(inp, dict) + else func(inp, *args, **kwargs) ) else: operator_output = func(*args, **kwargs) @@ -158,7 +184,6 @@ def run( self, *args, inference_state: InferenceState, - pipeline_state: PipelineState, **kwargs, ): """ @@ -171,31 +196,34 @@ def run( """ next_step = self.router.START_ROUTE operator_output = None - while next_step != self.router.END_ROUTE: # 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( + outputs = self._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, ) else: - outputs = self._run_next_step( - func=self._scheduler_group.submit, - input=operator_output, + outputs = self._run_next( + inp=operator_output, next_step=next_step, inference_state=inference_state, - operator=self.ops[next_step], - pipeline_state=pipeline_state, ) operator_output = outputs.result() @@ -221,12 +249,8 @@ def __call__(self, *args, **kwargs): else: 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/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 1c00261cb6..48a7158f66 100644 --- a/src/deepsparse/v2/text_generation/compile_logits.py +++ b/src/deepsparse/v2/text_generation/compile_logits.py @@ -12,8 +12,6 @@ # 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 diff --git a/src/deepsparse/v2/text_generation/generate_new_token.py b/src/deepsparse/v2/text_generation/generate_new_token.py index c7e0836c78..5bf48bbdbc 100644 --- a/src/deepsparse/v2/text_generation/generate_new_token.py +++ b/src/deepsparse/v2/text_generation/generate_new_token.py @@ -11,7 +11,7 @@ # 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 diff --git a/src/deepsparse/v2/text_generation/nl_engine_operator.py b/src/deepsparse/v2/text_generation/nl_engine_operator.py index 2863e6829a..94cadd8fc6 100644 --- a/src/deepsparse/v2/text_generation/nl_engine_operator.py +++ b/src/deepsparse/v2/text_generation/nl_engine_operator.py @@ -19,8 +19,7 @@ import numpy from pydantic import BaseModel, Field -from deepsparse import Engine -from deepsparse.utils import join_engine_outputs, model_to_path, split_engine_inputs +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,7 +28,6 @@ DEEPSPARSE_ENGINE, EngineOperator, EngineOperatorInputs, - EngineOperatorOutputs, ) @@ -193,7 +191,8 @@ def run(self, inp: NLEngineInputs, **kwargs) -> NLEngineOutputs: .get("engine_outputs") ) - # logits should be stacked along batch dim, kv_cache_state should be a list of dim batch size + # logits should be stacked along batch dim + # kv_cache_state should be a list of dim batch size logits, *kv_cache_state = out kv_cache_state, _ = split_engine_inputs(kv_cache_state, 1) @@ -203,6 +202,7 @@ def run(self, inp: NLEngineInputs, **kwargs) -> NLEngineOutputs: 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 = { From 3dda4ba26b4db145b9d4af1059a9990e5f5fc4d8 Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Thu, 16 Nov 2023 12:31:08 -0500 Subject: [PATCH 06/15] fix bug with batch size, introduce SplitRoute dataclass --- .../v2/operators/engine_operator.py | 10 +++- src/deepsparse/v2/pipeline.py | 60 +++++++++---------- .../continuous_batching_scheduler.py | 12 +++- .../v2/text_generation/nl_engine_operator.py | 35 +++++++---- src/deepsparse/v2/text_generation/pipeline.py | 4 +- src/deepsparse/v2/utils/__init__.py | 3 + src/deepsparse/v2/utils/data.py | 39 ++++++++++++ 7 files changed, 113 insertions(+), 50 deletions(-) create mode 100644 src/deepsparse/v2/utils/data.py diff --git a/src/deepsparse/v2/operators/engine_operator.py b/src/deepsparse/v2/operators/engine_operator.py index 4620008cc1..679c5efb19 100644 --- a/src/deepsparse/v2/operators/engine_operator.py +++ b/src/deepsparse/v2/operators/engine_operator.py @@ -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,10 +143,10 @@ 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() if engine_type == DEEPSPARSE_ENGINE: @@ -181,3 +182,6 @@ def run(self, inp: EngineOperatorInputs, **kwargs) -> Dict: engine_outputs = self.engine(inp.engine_inputs) return {"engine_outputs": engine_outputs} + + def override_model(*args, **kwargs): + raise NotImplementedError diff --git a/src/deepsparse/v2/pipeline.py b/src/deepsparse/v2/pipeline.py index 7d9a041db5..cf96d467fa 100644 --- a/src/deepsparse/v2/pipeline.py +++ b/src/deepsparse/v2/pipeline.py @@ -15,7 +15,6 @@ import copy from concurrent.futures import Future -from functools import partial from typing import Any, Callable, Dict, List, Union from deepsparse.v2.operators import EngineOperator, Operator @@ -25,29 +24,12 @@ OperatorScheduler, SchedulerGroup, ) -from deepsparse.v2.utils import InferenceState, PipelineState +from deepsparse.v2.utils import InferenceState, PipelineState, SplitRoute __all__ = ["Pipeline"] -class SubGraph: - def __init__(self, step, inf): - self.step = step - self.inf = inf - self.output = None - - def update_next_step(self): - if isinstance(self.output, tuple): - state_update = self.output[-1] - operator_output = self.output[0] - self.inf.update_state(state_update) - else: - operator_output = self.output - - return operator_output - - class Pipeline(Operator): """ Pipeline accepts a series of operators, schedulers, and a router. Calling a pipeline @@ -114,37 +96,48 @@ def _apply_split(self, inp: Any, inference_state: InferenceState): batches, orig_batch_size = self.expand_inputs(inp, 1) - sub_graphs = [ - SubGraph( + # 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_routes = [ + SplitRoute( inf=copy.deepcopy(inference_state), step=self.router.route[self.router.SPLIT_ROUTE], ) for i in range(len(batches)) ] - for i in range(len(sub_graphs)): - sub_graphs[i].output = self._run_next( - batches[i], sub_graphs[i].inf, sub_graphs[i].step + # Start each SPLIT_ROUTE; store the Future for each. + for i in range(len(split_routes)): + split_routes[i].output = self._run_next( + batches[i], split_routes[i].inf, split_routes[i].step ) + # Execute all split batches until all split batches have been completed. while True: - for i in range(len(sub_graphs)): - current_subgraph = sub_graphs[i] + for i in range(len(split_routes)): + current_subgraph = split_routes[i] if ( isinstance(current_subgraph.output, Future) and current_subgraph.output.done() ): + # get the result for the completed operator; resolve its output operator_output = current_subgraph.output.result() - current_subgraph.output = operator_output - - operator_output = current_subgraph.update_next_step() + operator_output = current_subgraph.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( current_subgraph.step, self.ops, operator_output ) + # update the step current_subgraph.step = next_step + # store the output for the next step. If the next step is + # JOIN_ROUTE note, this particular route has completed. Simply + # update the output value if next_step == self.router.JOIN_ROUTE: current_subgraph.output = operator_output else: @@ -155,10 +148,11 @@ def _apply_split(self, inp: Any, inference_state: InferenceState): ) break - if not any(isinstance(x.output, Future) for x in sub_graphs): + # keep running until all split routes have completed. + if not any(isinstance(x.output, Future) for x in split_routes): break - return self.condense_inputs([x.output for x in sub_graphs]) + return self.condense_inputs([x.output for x in split_routes]) def _run_func( self, @@ -168,7 +162,7 @@ def _run_func( **kwargs, ): """ - Generic function to run a given func. + Generic function to run a given Callable. """ if inp: operator_output = ( @@ -249,7 +243,7 @@ def __call__(self, *args, **kwargs): else: inference_state = InferenceState() inference_state.create_state({}) - + kwargs["inference_state"] = inference_state return self.run(*args, **kwargs) diff --git a/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py b/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py index aad885157a..c7a175e5b6 100644 --- a/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py +++ b/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py @@ -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 @@ -163,8 +165,16 @@ def add_engine_operator( for batch_size in batch_sizes: if batch_size == 1: continue # already added + + override_model_path = None + if hasattr(engine_operator, "override_model"): + override_model_path, _, _ = engine_operator.override_model( + 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 + batch_size=batch_size, model_path=override_model_path ) self._operators_to_engines[engine_operator] = operator_engines diff --git a/src/deepsparse/v2/text_generation/nl_engine_operator.py b/src/deepsparse/v2/text_generation/nl_engine_operator.py index 94cadd8fc6..15e2eafaa2 100644 --- a/src/deepsparse/v2/text_generation/nl_engine_operator.py +++ b/src/deepsparse/v2/text_generation/nl_engine_operator.py @@ -14,7 +14,8 @@ import copy import os -from typing import Any, List, Optional, Tuple +from pathlib import Path +from typing import Any, List, Optional, Tuple, Union import numpy from pydantic import BaseModel, Field @@ -123,17 +124,15 @@ def __init__( **kwargs, ): + self.sequence_length = sequence_length + self.input_ids_length = input_ids_length self.kv_cache_data_type = None + self.model_path = kwargs.get("model_path") ( 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.override_model(self.model_path, batch_size=1) engine_kwargs = kwargs.get("engine_kwargs", {}) if kwargs.get("engine_type", DEEPSPARSE_ENGINE) == DEEPSPARSE_ENGINE: @@ -153,7 +152,22 @@ def __init__( super().__init__(**kwargs) - self.input_ids_length = input_ids_length + def override_model(self, model_path: Union[str, Path], batch_size: int): + """ + Override the model based on the provided batch_size, sequence_length, + and input_ids_length. + """ + ( + 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, + ) + return onnx_file_path, output_indices_to_be_cached, kv_cache_data_type def run(self, inp: NLEngineInputs, **kwargs) -> NLEngineOutputs: engine_input = inp.engine_inputs @@ -179,9 +193,8 @@ def run(self, inp: NLEngineInputs, **kwargs) -> NLEngineOutputs: ) else: # run the engine without the LIB.kv_cache object - print(len(inputs)) + # stack multiple batch inputs along the batch dimension inputs = join_engine_outputs(inputs, len(inputs)) - out = ( super() .run( @@ -192,7 +205,7 @@ def run(self, inp: NLEngineInputs, **kwargs) -> NLEngineOutputs: ) # logits should be stacked along batch dim - # kv_cache_state should be a list of dim batch size + # kv_cache_state should be a list where each dim 0 is batch_size logits, *kv_cache_state = out kv_cache_state, _ = split_engine_inputs(kv_cache_state, 1) diff --git a/src/deepsparse/v2/text_generation/pipeline.py b/src/deepsparse/v2/text_generation/pipeline.py index dbfc3d8464..2673c868a6 100644 --- a/src/deepsparse/v2/text_generation/pipeline.py +++ b/src/deepsparse/v2/text_generation/pipeline.py @@ -136,8 +136,8 @@ def __init__( join_output = JoinOutput(tokenizer=self.tokenizer) continuous_batching_scheduler = ContinuousBatchingScheduler.get_instance() - continuous_batching_scheduler.add_engine_operator(single_engine_operator, [1]) - continuous_batching_scheduler.add_engine_operator(multi_engine_operator, [2, 4]) + continuous_batching_scheduler.add_engine_operator(single_engine_operator, [4]) + continuous_batching_scheduler.add_engine_operator(multi_engine_operator, [4]) ops = { "process_input": process_inputs, diff --git a/src/deepsparse/v2/utils/__init__.py b/src/deepsparse/v2/utils/__init__.py index 358405d7af..50d152ad1a 100644 --- a/src/deepsparse/v2/utils/__init__.py +++ b/src/deepsparse/v2/utils/__init__.py @@ -15,3 +15,6 @@ # limitations under the License. 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..dfa7041312 --- /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 + +from deepsparse.v2.utils import InferenceState + + +__all__ = ["SplitRoute"] + + +@dataclass +class SplitRoute: + """ + Helper dataclass to store information about each running split route. + """ + + step: int + inf: InferenceState + output: Any = None + + def parse_output(self, operator_output: Any): + if isinstance(operator_output, tuple): + state_update = operator_output[-1] + operator_output = operator_output[0] + self.inf.update_state(state_update) + return operator_output From 2726d03b18e2bd0c391fd2fe1928462da938a07c Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Thu, 16 Nov 2023 12:57:00 -0500 Subject: [PATCH 07/15] update tests to use new inputs/outputs --- src/deepsparse/v2/text_generation/pipeline.py | 1 - .../deepsparse/v2/unit/text_generation/test_misc.py | 13 ++++++++++--- .../text_generation/test_single_token_engine.py | 6 +++--- .../unit/text_generation/test_token_generation.py | 10 +++++++++- 4 files changed, 22 insertions(+), 8 deletions(-) diff --git a/src/deepsparse/v2/text_generation/pipeline.py b/src/deepsparse/v2/text_generation/pipeline.py index 2673c868a6..6b7f1e7fa6 100644 --- a/src/deepsparse/v2/text_generation/pipeline.py +++ b/src/deepsparse/v2/text_generation/pipeline.py @@ -53,7 +53,6 @@ def __init__( pipeline_state = PipelineState() pipeline_state_vals = {} - # TODO: The code below will be replaced with a transformers set-up Operator. self.tokenizer = None model_path = self.setup_onnx_file_path(model_path, sequence_length) self.tokenizer.padding_side = "left" 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_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 From b595e03476652f357c92ec84d3b4cece12a25833 Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Thu, 16 Nov 2023 15:03:49 -0500 Subject: [PATCH 08/15] use the normal scheduler for internal kv_cache --- src/deepsparse/v2/pipeline.py | 28 +++++++++---------- .../v2/text_generation/nl_engine_operator.py | 11 ++++---- src/deepsparse/v2/text_generation/pipeline.py | 18 ++++++++++-- 3 files changed, 34 insertions(+), 23 deletions(-) diff --git a/src/deepsparse/v2/pipeline.py b/src/deepsparse/v2/pipeline.py index cf96d467fa..f67efda012 100644 --- a/src/deepsparse/v2/pipeline.py +++ b/src/deepsparse/v2/pipeline.py @@ -72,7 +72,10 @@ def _run_next( inference_state: InferenceState, next_step: str, ): - if isinstance(self.ops[next_step], EngineOperator): + 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: @@ -116,34 +119,29 @@ def _apply_split(self, inp: Any, inference_state: InferenceState): # Execute all split batches until all split batches have been completed. while True: - for i in range(len(split_routes)): - current_subgraph = split_routes[i] - - if ( - isinstance(current_subgraph.output, Future) - and current_subgraph.output.done() - ): + for split_route in split_routes: + if isinstance(split_route.output, Future) and split_route.output.done(): # get the result for the completed operator; resolve its output - operator_output = current_subgraph.output.result() - operator_output = current_subgraph.parse_output(operator_output) + operator_output = split_route.output.result() + operator_output = split_route.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( - current_subgraph.step, self.ops, operator_output + split_route.step, self.ops, operator_output ) # update the step - current_subgraph.step = next_step + split_route.step = next_step # store the output for the next step. If the next step is # JOIN_ROUTE note, this particular route has completed. Simply # update the output value if next_step == self.router.JOIN_ROUTE: - current_subgraph.output = operator_output + split_route.output = operator_output else: - current_subgraph.output = self._run_next( + split_route.output = self._run_next( inp=operator_output, - inference_state=current_subgraph.inf, + inference_state=split_route.inf, next_step=next_step, ) break diff --git a/src/deepsparse/v2/text_generation/nl_engine_operator.py b/src/deepsparse/v2/text_generation/nl_engine_operator.py index 15e2eafaa2..038b9409d8 100644 --- a/src/deepsparse/v2/text_generation/nl_engine_operator.py +++ b/src/deepsparse/v2/text_generation/nl_engine_operator.py @@ -127,6 +127,7 @@ def __init__( self.sequence_length = sequence_length self.input_ids_length = input_ids_length self.kv_cache_data_type = None + self.internal_kv_cache = internal_kv_cache self.model_path = kwargs.get("model_path") ( onnx_file_path, @@ -173,7 +174,9 @@ def run(self, inp: NLEngineInputs, **kwargs) -> NLEngineOutputs: engine_input = inp.engine_inputs kv_cache = inp.kv_cache + split = True if not isinstance(kv_cache, list): + split = False kv_cache = [kv_cache] engine_input = [engine_input] @@ -220,11 +223,9 @@ def run(self, inp: NLEngineInputs, **kwargs) -> NLEngineOutputs: output = { "engine_outputs": logits, - "kv_cache": kv_cache, - "tokens": [inp.tokens] if not isinstance(inp.tokens, list) else inp.tokens, - "in_generation": [inp.in_generation] - if not isinstance(inp.in_generation, list) - else inp.in_generation, + "kv_cache": kv_cache if split else kv_cache[0], + "tokens": inp.tokens, + "in_generation": inp.in_generation, } return output diff --git a/src/deepsparse/v2/text_generation/pipeline.py b/src/deepsparse/v2/text_generation/pipeline.py index 6b7f1e7fa6..30ba6e5a85 100644 --- a/src/deepsparse/v2/text_generation/pipeline.py +++ b/src/deepsparse/v2/text_generation/pipeline.py @@ -134,9 +134,21 @@ def __init__( compile_generated_tokens = CompileGeneratedTokens() join_output = JoinOutput(tokenizer=self.tokenizer) - continuous_batching_scheduler = ContinuousBatchingScheduler.get_instance() - continuous_batching_scheduler.add_engine_operator(single_engine_operator, [4]) - continuous_batching_scheduler.add_engine_operator(multi_engine_operator, [4]) + + # TODO: Do we always want to use continuous batching by default, but allow + # the user to turn it off if needed? + + # Temporary; use input param to turn this off + if not internal_kv_cache: + continuous_batching_scheduler = ContinuousBatchingScheduler.get_instance() + continuous_batching_scheduler.add_engine_operator( + single_engine_operator, [4] + ) + continuous_batching_scheduler.add_engine_operator( + multi_engine_operator, [4] + ) + else: + continuous_batching_scheduler = None ops = { "process_input": process_inputs, From a8948998ce2af67fd5894a52f04164fc62d48402 Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Thu, 16 Nov 2023 15:35:48 -0500 Subject: [PATCH 09/15] add pipeline inpuits --- .../utils/continuous_batching_executor.py | 4 +- src/deepsparse/v2/text_generation/pipeline.py | 39 ++++++++++++------- 2 files changed, 25 insertions(+), 18 deletions(-) diff --git a/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py b/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py index 57151e64c3..40ff00ca4f 100644 --- a/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py +++ b/src/deepsparse/v2/schedulers/utils/continuous_batching_executor.py @@ -71,9 +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, inference_state=None, pipeline_state=None - ) + 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/pipeline.py b/src/deepsparse/v2/text_generation/pipeline.py index 30ba6e5a85..c9b2f8978e 100644 --- a/src/deepsparse/v2/text_generation/pipeline.py +++ b/src/deepsparse/v2/text_generation/pipeline.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging from typing import Dict from deepsparse.transformers.utils.helpers import process_generation_config @@ -38,6 +39,9 @@ from deepsparse.v2.utils import PipelineState +_LOGGER = logging.getLogger(__name__) + + class TextGenerationPipeline(Pipeline): def __init__( self, @@ -47,6 +51,7 @@ def __init__( internal_kv_cache: bool = True, force_max_tokens: bool = False, generation_config=None, + continuous_batch_sizes: list = None, engine_kwargs: Dict = None, ): @@ -134,21 +139,25 @@ def __init__( compile_generated_tokens = CompileGeneratedTokens() join_output = JoinOutput(tokenizer=self.tokenizer) - - # TODO: Do we always want to use continuous batching by default, but allow - # the user to turn it off if needed? - - # Temporary; use input param to turn this off - if not internal_kv_cache: - continuous_batching_scheduler = ContinuousBatchingScheduler.get_instance() - continuous_batching_scheduler.add_engine_operator( - single_engine_operator, [4] - ) - continuous_batching_scheduler.add_engine_operator( - multi_engine_operator, [4] - ) - else: - continuous_batching_scheduler = None + # TODO: do we want to support lists for different engines? + continuous_batching_scheduler = None + if continuous_batch_sizes: + if not internal_kv_cache: + continuous_batching_scheduler = ( + ContinuousBatchingScheduler.get_instance() + ) + continuous_batching_scheduler.add_engine_operator( + single_engine_operator, continuous_batch_sizes + ) + continuous_batching_scheduler.add_engine_operator( + multi_engine_operator, continuous_batch_sizes + ) + else: + # Temporary + _LOGGER.warn( + "internal kv_cache is currently not supported with continuous ", + "batching", + ) ops = { "process_input": process_inputs, From aa1357dbf14f51852c4614336ea46b012c5524b8 Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Thu, 16 Nov 2023 15:43:24 -0500 Subject: [PATCH 10/15] clean-up --- src/deepsparse/v2/operators/engine_operator.py | 3 --- src/deepsparse/v2/schedulers/continuous_batching_scheduler.py | 2 ++ 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/deepsparse/v2/operators/engine_operator.py b/src/deepsparse/v2/operators/engine_operator.py index 679c5efb19..8a2400b939 100644 --- a/src/deepsparse/v2/operators/engine_operator.py +++ b/src/deepsparse/v2/operators/engine_operator.py @@ -182,6 +182,3 @@ def run(self, inp: EngineOperatorInputs, **kwargs) -> Dict: engine_outputs = self.engine(inp.engine_inputs) return {"engine_outputs": engine_outputs} - - def override_model(*args, **kwargs): - raise NotImplementedError diff --git a/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py b/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py index c7a175e5b6..8117bd2869 100644 --- a/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py +++ b/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py @@ -167,6 +167,8 @@ def add_engine_operator( continue # already added override_model_path = None + # text generation/NLEngineOperator specific; could add generic method + # for all engine_operators, if desired if hasattr(engine_operator, "override_model"): override_model_path, _, _ = engine_operator.override_model( model_path=engine_operator.model_path, batch_size=batch_size From cea50f99fd2c77ce04c6c56f99cd9a019b95884f Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Mon, 20 Nov 2023 20:33:32 +0000 Subject: [PATCH 11/15] change engine type, update docstrings, update override function to be more generic --- .../v2/operators/engine_operator.py | 4 +-- .../continuous_batching_scheduler.py | 4 +-- .../v2/text_generation/nl_engine_operator.py | 30 ++++++++++++++----- 3 files changed, 27 insertions(+), 11 deletions(-) diff --git a/src/deepsparse/v2/operators/engine_operator.py b/src/deepsparse/v2/operators/engine_operator.py index 8a2400b939..630de2d5bd 100644 --- a/src/deepsparse/v2/operators/engine_operator.py +++ b/src/deepsparse/v2/operators/engine_operator.py @@ -13,7 +13,7 @@ # limitations under the License. from copy import deepcopy -from typing import Any, Dict, List, Optional, Union +from typing import Dict, List, Optional, Union from pydantic import BaseModel, Field @@ -34,7 +34,7 @@ class EngineOperatorInputs(BaseModel): engine_inputs: List = Field(description="engine_inputs") - engine: Optional[Any] = Field( + engine: Optional[Union[ORTEngine, Engine]] = Field( description="override the engine to run forward pass with", default=None, ) diff --git a/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py b/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py index 8117bd2869..cc74ac0996 100644 --- a/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py +++ b/src/deepsparse/v2/schedulers/continuous_batching_scheduler.py @@ -169,8 +169,8 @@ def add_engine_operator( override_model_path = None # text generation/NLEngineOperator specific; could add generic method # for all engine_operators, if desired - if hasattr(engine_operator, "override_model"): - override_model_path, _, _ = engine_operator.override_model( + 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 ) diff --git a/src/deepsparse/v2/text_generation/nl_engine_operator.py b/src/deepsparse/v2/text_generation/nl_engine_operator.py index 038b9409d8..13dfab68d7 100644 --- a/src/deepsparse/v2/text_generation/nl_engine_operator.py +++ b/src/deepsparse/v2/text_generation/nl_engine_operator.py @@ -129,11 +129,12 @@ def __init__( self.kv_cache_data_type = None self.internal_kv_cache = internal_kv_cache self.model_path = kwargs.get("model_path") - ( - onnx_file_path, - output_indices_to_be_cached, - kv_cache_data_type, - ) = self.override_model(self.model_path, batch_size=1) + (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: @@ -153,10 +154,20 @@ def __init__( super().__init__(**kwargs) - def override_model(self, model_path: Union[str, Path], batch_size: int): + 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 particular engine """ ( onnx_file_path, @@ -168,7 +179,12 @@ def override_model(self, model_path: Union[str, Path], batch_size: int): sequence_length=self.sequence_length, input_ids_length=self.input_ids_length, ) - return onnx_file_path, output_indices_to_be_cached, kv_cache_data_type + 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: NLEngineInputs, **kwargs) -> NLEngineOutputs: engine_input = inp.engine_inputs From 89d6558c04951069a25fd05b51f555989261c976 Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Mon, 20 Nov 2023 21:56:48 +0000 Subject: [PATCH 12/15] move subgraph functionality to its own function; clean-up cont batching in text gen pipeline --- src/deepsparse/v2/pipeline.py | 124 +++++++++--------- .../v2/text_generation/nl_engine_operator.py | 2 +- src/deepsparse/v2/text_generation/pipeline.py | 41 ++++-- src/deepsparse/v2/utils/__init__.py | 1 + src/deepsparse/v2/utils/data.py | 7 +- src/deepsparse/v2/utils/helpers.py | 37 ++++++ 6 files changed, 135 insertions(+), 77 deletions(-) create mode 100644 src/deepsparse/v2/utils/helpers.py diff --git a/src/deepsparse/v2/pipeline.py b/src/deepsparse/v2/pipeline.py index f67efda012..8a833c2608 100644 --- a/src/deepsparse/v2/pipeline.py +++ b/src/deepsparse/v2/pipeline.py @@ -15,7 +15,7 @@ import copy from concurrent.futures import Future -from typing import Any, Callable, Dict, List, Union +from typing import Any, Dict, List, Union from deepsparse.v2.operators import EngineOperator, Operator from deepsparse.v2.routers import Router @@ -24,7 +24,9 @@ OperatorScheduler, SchedulerGroup, ) -from deepsparse.v2.utils import InferenceState, PipelineState, SplitRoute +from deepsparse.v2.utils import InferenceState, PipelineState +from deepsparse.v2.utils.data import SubGraph +from deepsparse.v2.utils.helpers import run_func __all__ = ["Pipeline"] @@ -81,7 +83,7 @@ def _run_next( else: func = self._scheduler_group.submit - return self._run_func( + return run_func( func=func, operator=self.ops[next_step], inp=inp, @@ -89,88 +91,92 @@ def _run_next( inference_state=inference_state, ) - def _apply_split(self, inp: Any, inference_state: InferenceState): + def _run_sub_graphs( + self, sub_graph_inputs: List[Any], sub_graphs: List[SubGraph] + ) -> List[Any]: """ - Split inputs using the pipeline's expand_inputs function. Inputs are split - into a batch size of one when a SPLIT_ROUTE node is found in a given pipeline's - provided router. The split batches are run asynchronously and then joined when - a JOIN_ROUTE node is found, using the pipeline's condense_inputs function. + 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. """ - - batches, orig_batch_size = self.expand_inputs(inp, 1) - - # 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_routes = [ - SplitRoute( - inf=copy.deepcopy(inference_state), - step=self.router.route[self.router.SPLIT_ROUTE], - ) - for i in range(len(batches)) - ] - - # Start each SPLIT_ROUTE; store the Future for each. - for i in range(len(split_routes)): - split_routes[i].output = self._run_next( - batches[i], split_routes[i].inf, split_routes[i].step + 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 ) - # Execute all split batches until all split batches have been completed. + # Execute all sub graphs until all graphs have been completed. while True: - for split_route in split_routes: - if isinstance(split_route.output, Future) and split_route.output.done(): + 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 = split_route.output.result() - operator_output = split_route.parse_output(operator_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( - split_route.step, self.ops, operator_output + sub_graph.step, self.ops, operator_output ) # update the step - split_route.step = next_step + sub_graph.step = next_step # store the output for the next step. If the next step is - # JOIN_ROUTE note, this particular route has completed. Simply + # end step, this particular route has completed. Simply # update the output value - if next_step == self.router.JOIN_ROUTE: - split_route.output = operator_output + if next_step == sub_graph.end: + sub_graph.output = operator_output else: - split_route.output = self._run_next( + sub_graph.output = self._run_next( inp=operator_output, - inference_state=split_route.inf, + inference_state=sub_graph.inf, next_step=next_step, ) break - # keep running until all split routes have completed. - if not any(isinstance(x.output, Future) for x in split_routes): + # keep running until all sub graphs have completed. + if not any(isinstance(x.output, Future) for x in sub_graphs): break - return self.condense_inputs([x.output for x in split_routes]) + return [x.output for x in sub_graphs] - def _run_func( - self, - *args, - func: Callable, - inp: Any = None, - **kwargs, - ): + def _apply_split(self, inp: Any, inference_state: InferenceState): """ - Generic function to run a given Callable. + Split inputs using the pipeline's expand_inputs function. Inputs are split + into a batch size of one when a SPLIT_ROUTE node is found in a given pipeline's + provided router. The split batches are run asynchronously and then joined when + a JOIN_ROUTE node is found, using the pipeline's condense_inputs function. """ - if inp: - operator_output = ( - func(*args, **kwargs, **inp) - if isinstance(inp, dict) - else func(inp, *args, **kwargs) + + batches, orig_batch_size = self.expand_inputs(inp, 1) + + # 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) - return operator_output + for i in range(len(batches)) + ] + + outputs = self._run_sub_graphs( + sub_graph_inputs=batches, sub_graphs=split_graphs + ) + return self.condense_inputs(outputs) def run( self, @@ -203,7 +209,7 @@ def run( return operator_output if next_step == self.router.START_ROUTE: - outputs = self._run_func( + outputs = run_func( *args, func=self._scheduler_group.submit, operator=self.ops[next_step], diff --git a/src/deepsparse/v2/text_generation/nl_engine_operator.py b/src/deepsparse/v2/text_generation/nl_engine_operator.py index 13dfab68d7..d8c80bbaee 100644 --- a/src/deepsparse/v2/text_generation/nl_engine_operator.py +++ b/src/deepsparse/v2/text_generation/nl_engine_operator.py @@ -167,7 +167,7 @@ def override_model_inputs( :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 particular engine + specific to the NLDecoder engine """ ( onnx_file_path, diff --git a/src/deepsparse/v2/text_generation/pipeline.py b/src/deepsparse/v2/text_generation/pipeline.py index c9b2f8978e..30ada1cb3b 100644 --- a/src/deepsparse/v2/text_generation/pipeline.py +++ b/src/deepsparse/v2/text_generation/pipeline.py @@ -13,10 +13,11 @@ # limitations under the License. import logging -from typing import Dict +from typing import Dict, List, Optional 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 ContinuousBatchingScheduler, OperatorScheduler @@ -51,7 +52,7 @@ def __init__( internal_kv_cache: bool = True, force_max_tokens: bool = False, generation_config=None, - continuous_batch_sizes: list = None, + continuous_batch_sizes: Optional[List[int]] = None, engine_kwargs: Dict = None, ): @@ -142,22 +143,16 @@ def __init__( # TODO: do we want to support lists for different engines? continuous_batching_scheduler = None if continuous_batch_sizes: - if not internal_kv_cache: - continuous_batching_scheduler = ( - ContinuousBatchingScheduler.get_instance() - ) - continuous_batching_scheduler.add_engine_operator( - single_engine_operator, continuous_batch_sizes - ) - continuous_batching_scheduler.add_engine_operator( - multi_engine_operator, continuous_batch_sizes - ) - else: - # Temporary + 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, @@ -225,6 +220,24 @@ 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 + # TODO: Move to be part of a generic transformers set-up Operator. def setup_onnx_file_path(self, model_path, sequence_length) -> str: import logging diff --git a/src/deepsparse/v2/utils/__init__.py b/src/deepsparse/v2/utils/__init__.py index 50d152ad1a..75935a9729 100644 --- a/src/deepsparse/v2/utils/__init__.py +++ b/src/deepsparse/v2/utils/__init__.py @@ -13,6 +13,7 @@ # 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 * diff --git a/src/deepsparse/v2/utils/data.py b/src/deepsparse/v2/utils/data.py index dfa7041312..8622f37f1c 100644 --- a/src/deepsparse/v2/utils/data.py +++ b/src/deepsparse/v2/utils/data.py @@ -18,17 +18,18 @@ from deepsparse.v2.utils import InferenceState -__all__ = ["SplitRoute"] +__all__ = ["SubGraph"] @dataclass -class SplitRoute: +class SubGraph: """ - Helper dataclass to store information about each running split route. + Helper dataclass to store information about each running sub graph. """ step: int inf: InferenceState + end: str output: Any = None def parse_output(self, operator_output: Any): 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 From ce74ac6735cccba72cb73f3b52a873215ce98c7d Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Mon, 20 Nov 2023 22:36:31 +0000 Subject: [PATCH 13/15] update linear pathway to also use subgraph execution --- src/deepsparse/v2/pipeline.py | 40 +++++++++++++++++++++------------ src/deepsparse/v2/utils/data.py | 7 +++--- 2 files changed, 29 insertions(+), 18 deletions(-) diff --git a/src/deepsparse/v2/pipeline.py b/src/deepsparse/v2/pipeline.py index 8a833c2608..78d112a2b3 100644 --- a/src/deepsparse/v2/pipeline.py +++ b/src/deepsparse/v2/pipeline.py @@ -134,7 +134,7 @@ def _run_sub_graphs( # 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 == sub_graph.end: + if next_step in sub_graph.end: sub_graph.output = operator_output else: sub_graph.output = self._run_next( @@ -168,7 +168,7 @@ def _apply_split(self, inp: Any, inference_state: InferenceState): SubGraph( inf=copy.deepcopy(inference_state), step=self.router.route[self.router.SPLIT_ROUTE], - end=self.router.JOIN_ROUTE, + end=[self.router.JOIN_ROUTE], ) for i in range(len(batches)) ] @@ -195,6 +195,8 @@ 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: @@ -209,29 +211,39 @@ def run( return operator_output if next_step == self.router.START_ROUTE: - outputs = run_func( + operator_output = run_func( *args, func=self._scheduler_group.submit, operator=self.ops[next_step], 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( - inp=operator_output, - next_step=next_step, - inference_state=inference_state, + # Single graph execution + graph = SubGraph( + inf=copy.deepcopy(inference_state), + step=next_step, + end=[self.router.SPLIT_ROUTE, self.router.END_ROUTE], ) - operator_output = outputs.result() + operator_output = self._run_sub_graphs( + sub_graph_inputs=[operator_output], sub_graphs=[graph] + )[0] - if isinstance(operator_output, tuple): - state_update = operator_output[-1] - operator_output = operator_output[0] - inference_state.update_state(state_update) + inference_state = graph.inf + next_step = graph.step - next_step = self.router.next(next_step, self.ops, operator_output) return operator_output def __call__(self, *args, **kwargs): diff --git a/src/deepsparse/v2/utils/data.py b/src/deepsparse/v2/utils/data.py index 8622f37f1c..40402734cf 100644 --- a/src/deepsparse/v2/utils/data.py +++ b/src/deepsparse/v2/utils/data.py @@ -13,7 +13,7 @@ # limitations under the License. from dataclasses import dataclass -from typing import Any +from typing import Any, List from deepsparse.v2.utils import InferenceState @@ -29,12 +29,11 @@ class SubGraph: step: int inf: InferenceState - end: str + end: List[str] output: Any = None def parse_output(self, operator_output: Any): if isinstance(operator_output, tuple): - state_update = operator_output[-1] - operator_output = operator_output[0] + operator_output, state_update = operator_output[0], operator_output[-1] self.inf.update_state(state_update) return operator_output From 2afb461ef3892049800aa2c7a3703818810d88a5 Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Mon, 20 Nov 2023 22:43:24 +0000 Subject: [PATCH 14/15] rebase fix --- src/deepsparse/v2/text_generation/pipeline.py | 1 - tests/deepsparse/v2/unit/text_generation/conftest.py | 11 +++++------ .../v2/unit/text_generation/test_process_inputs.py | 6 ++---- 3 files changed, 7 insertions(+), 11 deletions(-) diff --git a/src/deepsparse/v2/text_generation/pipeline.py b/src/deepsparse/v2/text_generation/pipeline.py index 4622835757..ae7334cffd 100644 --- a/src/deepsparse/v2/text_generation/pipeline.py +++ b/src/deepsparse/v2/text_generation/pipeline.py @@ -219,7 +219,6 @@ 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: 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_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( From 9af60d6d1ae5136eb2b0bf1ad4bfdedf0c877aaa Mon Sep 17 00:00:00 2001 From: Dipika Sikka Date: Mon, 20 Nov 2023 23:02:03 +0000 Subject: [PATCH 15/15] fix tests --- tests/deepsparse/v2/integration_tests/test_llms.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 "