[Text Generation][V2] NonKVCachePipeline

src/deepsparse/transformers/helpers.py

@@ -30,11 +30,7 @@
 from onnx import ModelProto
 
 from deepsparse.log import get_main_logger
-from deepsparse.utils.onnx import (
-    _MODEL_DIR_ONNX_NAME,
-    model_to_path,
-    truncate_onnx_model,
-)
+from deepsparse.utils.onnx import MODEL_ONNX_NAME, model_to_path, truncate_onnx_model
 from sparsezoo.utils import save_onnx
 
 
@@ -55,6 +51,7 @@ def setup_transformers_pipeline(
     sequence_length: int,
     tokenizer_padding_side: str = "left",
     engine_kwargs: Optional[Dict] = None,
+    onnx_model_name: Optional[str] = None,
 ) -> Tuple[
     str, transformers.PretrainedConfig, transformers.PreTrainedTokenizer, Dict[str, Any]
 ]:
@@ -65,10 +62,14 @@ def setup_transformers_pipeline(
     :param sequence_length: The sequence length to use for the model
     :param tokenizer_padding_side: The side to pad on for the tokenizer,
         either "left" or "right"
+    :param onnx_model_name: The name of the onnx model to be loaded.
+        If not specified, defaults are used (see fetch_onnx_file_path)
     :param engine_kwargs: The kwargs to pass to the engine
     :return The model path, config, tokenizer, and engine kwargs
     """
-    model_path, config, tokenizer = fetch_onnx_file_path(model_path, sequence_length)
+    model_path, config, tokenizer = fetch_onnx_file_path(
+        model_path, sequence_length, onnx_model_name
+    )
 
     tokenizer.padding_side = tokenizer_padding_side
     if not tokenizer.pad_token:
@@ -89,6 +90,7 @@ def setup_transformers_pipeline(
 def fetch_onnx_file_path(
     model_path: str,
     sequence_length: int,
+    onnx_model_name: Optional[str] = None,
     task: Optional[str] = None,
 ) -> Tuple[str, transformers.PretrainedConfig, transformers.PreTrainedTokenizer]:
     """
@@ -97,9 +99,13 @@ def fetch_onnx_file_path(
     derived from the `model_path` provided.
     :param model_path: path to the model to be parsed
     :param sequence_length: maximum sequence length of the model
+    :param onnx_model_name: optionally, the precise name of the ONNX model
+        of interest may be specified. If not specified, the default ONNX model
+        name will be used (refer to `get_deployment_path` for details)
+    :param task: task to use for the config. Defaults to None
     :return: file path to the processed ONNX file for the engine to compile
     """
-    deployment_path, onnx_path = get_deployment_path(model_path)
+    deployment_path, onnx_path = get_deployment_path(model_path, onnx_model_name)
 
     hf_logger = logging.getLogger("transformers")
     hf_logger_level = hf_logger.level
@@ -126,7 +132,9 @@ def fetch_onnx_file_path(
     return onnx_path, config, tokenizer
 
 
-def get_deployment_path(model_path: str) -> Tuple[str, str]:
+def get_deployment_path(
+    model_path: str, onnx_model_name: Optional[str] = None
+) -> Tuple[str, str]:
     """
     Returns the path to the deployment directory
     for the given model path and the path to the mandatory
@@ -135,27 +143,33 @@ def get_deployment_path(model_path: str) -> Tuple[str, str]:
     for running the transformers model in the deepsparse pipeline
 
     :param model_path: path to model directory, sparsezoo stub, or ONNX file
+    :param onnx_model_name: name of the ONNX file to look for in the deployment
+        directory. Defaults to MODEL_ONNX_NAME
     :return: path to the deployment directory and path to the ONNX file inside
         the deployment directory
     """
+    onnx_model_name = onnx_model_name or MODEL_ONNX_NAME
+
     if os.path.isfile(model_path):
         # return the parent directory of the ONNX file
         return os.path.dirname(model_path), model_path
 
     if os.path.isdir(model_path):
         model_files = os.listdir(model_path)
 
-        if _MODEL_DIR_ONNX_NAME not in model_files:
+        if onnx_model_name not in model_files:
             raise ValueError(
-                f"{_MODEL_DIR_ONNX_NAME} not found in transformers model directory "
+                f"{onnx_model_name} not found in transformers model directory "
                 f"{model_path}. Be sure that an export of the model is written to "
-                f"{os.path.join(model_path, _MODEL_DIR_ONNX_NAME)}"
+                f"{os.path.join(model_path, onnx_model_name)}"
             )
-        return model_path, os.path.join(model_path, _MODEL_DIR_ONNX_NAME)
+        return model_path, os.path.join(model_path, onnx_model_name)
 
     elif model_path.startswith("zoo:") or model_path.startswith("hf:"):
         onnx_model_path = model_to_path(model_path)
-        return os.path.dirname(onnx_model_path), onnx_model_path
+        return os.path.dirname(onnx_model_path), onnx_model_path.replace(
+            MODEL_ONNX_NAME, onnx_model_name
+        )
     else:
         raise ValueError(
             f"model_path {model_path} is not a valid file, directory, or zoo stub"

src/deepsparse/transformers/pipelines/text_generation/__init__.py

@@ -21,6 +21,7 @@
 from .kv_cache_operator import *
 from .multi_engine_prefill_operator import *
 from .nl_engine_operator import *
+from .nl_engine_operator_no_kv_cache import *
 from .parse_inputs import *
 from .prep_for_prefill import *
 from .process_inputs import *
@@ -31,3 +32,4 @@
 from .prep_for_generation import *  # isort:skip
 
 from .pipeline import *  # isort:skip
+from .pipeline_no_kv_cache import *  # isort:skip

src/deepsparse/transformers/pipelines/text_generation/generate_new_token.py

@@ -49,8 +49,9 @@ def run(self, inp: NLEngineOutputs, inference_state: InferenceState, **kwargs):
         callback = inference_state.current_state.get("callback")
         stop = inference_state.current_state.get("stop")
 
-        if kv_cache.total_num_processed_tokens >= kv_cache.capacity:
-            finish_reason = FinishReason.CAPACITY
+        if kv_cache:
+            if kv_cache.total_num_processed_tokens >= kv_cache.capacity:
+                finish_reason = FinishReason.CAPACITY
 
         if token == self.tokenizer.eos_token_id and not self.force_max_tokens:
             finish_reason = FinishReason.STOP

src/deepsparse/transformers/pipelines/text_generation/join_output.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List
+from typing import Dict, List, Tuple
 
 import numpy
 
@@ -34,7 +34,8 @@ class JoinOutput(Operator):
     def __init__(self, tokenizer):
         self.tokenizer = tokenizer
 
-    def run(self, inp: List[CompileGenerationsOutput], **kwargs):
+    def run(self, inp: Tuple[List[CompileGenerationsOutput], Dict], **kwargs):
+
         batch_outputs = [x for x in inp[0]]
         generated_tokens = [x.generated_tokens for x in batch_outputs]
         generated_logits = [x.generated_logits for x in batch_outputs]

src/deepsparse/transformers/pipelines/text_generation/nl_engine_operator.py

@@ -32,7 +32,7 @@
 )
 
 
-__all__ = ["NLEngineOperator", "NLEngineInputs"]
+__all__ = ["NLEngineOperator", "NLEngineInputs", "NLEngineOutputs"]
 
 
 class NLEngineInputs(BaseModel):

src/deepsparse/transformers/pipelines/text_generation/nl_engine_operator_no_kv_cache.py

@@ -0,0 +1,67 @@
+# 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
+
+import numpy
+from pydantic import BaseModel
+
+from deepsparse.operators.engine_operator import EngineOperator, EngineOperatorInputs
+from deepsparse.transformers.helpers import overwrite_transformer_onnx_model_inputs
+
+
+__all__ = [
+    "NLEngineOperatorNoCache",
+    "NLEngineInputsNoCache",
+]
+
+
+class NLEngineInputsNoCache(BaseModel):
+    input_ids: Any
+    attention_mask: Any
+
+
+class NLEngineOperatorNoCache(EngineOperator):
+    """
+    Operator the Natural Language Engine, that operates without
+    KV Cache. This means that this operator merely maps input_ids
+    and attention_mask to logits
+    """
+
+    input_schema = NLEngineInputsNoCache
+    output_schema = None
+
+    def __init__(self, sequence_length: int, **kwargs):
+        overwrite_transformer_onnx_model_inputs(
+            path=kwargs.get("model_path"),
+            batch_size=kwargs.get("batch_size", 1),
+            max_length=sequence_length,
+        )
+        super().__init__(**kwargs)
+
+    def run(self, inp: NLEngineInputsNoCache, **kwargs) -> Any:
+        engine_inputs = [inp.input_ids, inp.attention_mask]
+        logits = (
+            super()
+            .run(EngineOperatorInputs(engine_inputs=engine_inputs), **kwargs)
+            .get("engine_outputs")
+        )
+
+        # By default, the engine outputs logits for all tokens in the sequence.
+        # Let's filter out the logits for the padding tokens.
+        logits = numpy.compress(inp.attention_mask.flatten(), logits[0], axis=1)
+
+        return {"logits": [logits], "kv_cache": None, "tokens": None}, {
+            "prompt_logits": [logits]
+        }

src/deepsparse/transformers/pipelines/text_generation/pipeline_no_kv_cache.py

@@ -0,0 +1,148 @@
+# 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 logging
+from typing import Dict, Optional
+
+from deepsparse.pipeline import Pipeline
+from deepsparse.routers import GraphRouter
+from deepsparse.schedulers import OperatorScheduler
+from deepsparse.transformers.helpers import setup_transformers_pipeline
+from deepsparse.transformers.pipelines.text_generation import (
+    CompileGenerations,
+    GenerateNewTokenOperator,
+    JoinOutput,
+    NLEngineOperatorNoCache,
+    PrepareGeneration,
+    ProcessInputsTextGeneration,
+    ProcessOutputs,
+    TokenGeneratorOperator,
+)
+from deepsparse.transformers.utils.helpers import process_generation_config
+from deepsparse.utils import split_engine_inputs
+from deepsparse.utils.onnx import default_cached_outputs
+
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class TextGenerationPipelineNoCache(Pipeline):
+    def __init__(
+        self,
+        model_path: str,
+        sequence_length: int = 1024,
+        onnx_model_name: Optional[str] = None,
+        generation_config=None,
+        engine_kwargs: Optional[Dict] = None,
+        **kwargs,
+    ):
+
+        (
+            self.model_path,
+            self.config,
+            self.tokenizer,
+            engine_kwargs,
+        ) = setup_transformers_pipeline(
+            model_path,
+            sequence_length,
+            tokenizer_padding_side="right",
+            onnx_model_name=onnx_model_name,
+            engine_kwargs=engine_kwargs,
+        )
+        self.verify_no_kv_cache_present()
+
+        token_generator = TokenGeneratorOperator()
+
+        process_inputs = ProcessInputsTextGeneration(
+            generation_config=process_generation_config(generation_config),
+            sequence_length=sequence_length,
+            tokenizer=self.tokenizer,
+        )
+        engine_operator = NLEngineOperatorNoCache(
+            sequence_length=sequence_length,
+            **engine_kwargs,
+        )
+        prepare_generation = PrepareGeneration(
+            sequence_length=sequence_length,
+            prompt_sequence_length=1,
+            token_generator=token_generator,
+        )
+        generate_new_token = GenerateNewTokenOperator(
+            tokenizer=self.tokenizer, force_max_tokens=True
+        )
+        compile_generations = CompileGenerations()
+        join_output = JoinOutput(tokenizer=self.tokenizer)
+        process_outputs = ProcessOutputs(tokenizer=self.tokenizer)
+
+        ops = {
+            "process_input": process_inputs,
+            "engine_operator": engine_operator,
+            "prepare_generation": prepare_generation,
+            "generate_new_token": generate_new_token,
+            "compile_generations": compile_generations,
+            "join_output": join_output,
+            "process_outputs": process_outputs,
+        }
+        routes = {
+            "process_input": "SPLIT",
+            "SPLIT": "engine_operator",
+            "engine_operator": "prepare_generation",
+            "prepare_generation": "generate_new_token",
+            "generate_new_token": "compile_generations",
+            "compile_generations": "JOIN",
+            "JOIN": "join_output",
+            "join_output": "process_outputs",
+            "process_outputs": "STOP",
+        }
+
+        # TODO: Using the GraphRouter, but should use
+        # LinearRouter with appropriate split/join support
+        router = GraphRouter(
+            end_route="STOP", start_route="process_input", route=routes
+        )
+        scheduler = [OperatorScheduler()]
+        super().__init__(
+            ops=ops,
+            router=router,
+            schedulers=scheduler,
+        )
+
+    def run(self, *args, **kwargs):
+        # we need to set the fixed_sequences_length flag to True
+        # for the non-kv cache pipeline
+        kwargs.update(dict(fixed_sequences_length=True, max_new_tokens=1))
+        return super().run(*args, **kwargs)
+
+    def condense_inputs(self, *args, **kwargs):
+        return args[0], kwargs
+
+    def expand_inputs(self, items, batch_size):
+        items = [items.get(key) for key in items.keys()]
+        out, orig_batch_size = split_engine_inputs(items, batch_size)
+        combined_batches = [{"input_ids": b[0], "attention_mask": b[1]} for b in out]
+        return combined_batches, orig_batch_size
+
+    def verify_no_kv_cache_present(self) -> bool:
+        """
+        Verifies that the ONNX model does not have
+        KV cache inputs/outputs present.
+        :return: True if compatible, False otherwise
+        """
+        is_kv_cache_present = any(default_cached_outputs(self.model_path))
+        if is_kv_cache_present:
+            raise ValueError(
+                f"The model: {self.model_path} has KV cache inputs/outputs present. "
+                "Please use the TextGenerationPipeline instead."
+            )
+        return not is_kv_cache_present