[Text Generation] Multitoken prefill enablement

[dbogunowicz]

Enable running the pipeline in the mode, where the prompt is processed (prefill scenario) through multiple, consecutive passes through the multitoken engine. The goal is to achieve optimal inference speed with the deepsparse engine.

Manual Testing

from deepsparse import Pipeline



def _test_pipeline(engine_type, prompt_processing_sequence_length):
    opt = Pipeline.create(task="opt",
                          model_path="/home/ubuntu/damian/sparseml/deployment",
                          engine_type=engine_type,
                          use_deepsparse_cache = False,
                          prompt_processing_sequence_length = prompt_processing_sequence_length,
                          max_generated_tokens=32)
    prompt = "def hello_world():" * 20 # long prompt, so it gets processed by multitoken engine
    out = opt(sequences=prompt, return_logits=True)
    print(out.sequences[0])

for prompt_processing_sequence_length in (8, 16, 55, 128):
    _test_pipeline(engine_type ="onnxruntime", prompt_processing_sequence_length=prompt_processing_sequence_length)

Results:

# Results are identical to the pytorch baseline 

2023-07-20 15:58:14 deepsparse.transformers.engines.nl_decoder_engine INFO     Overwriting in-place the input shapes of the transformer model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 15:58:24 deepsparse.utils.onnx INFO     Overwriting in-place the batch size of the model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 15:58:32 deepsparse.transformers.engines.nl_decoder_engine INFO     Overwriting in-place the input shapes of the transformer model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 15:58:39 deepsparse.utils.onnx INFO     Overwriting in-place the batch size of the model at /home/ubuntu/damian/sparseml/deployment/model.onnx
def hello_world():def hello_world():def hello_world():def hello_world():def hello_world():def hello_world():def hello

2023-07-20 15:58:53 deepsparse.transformers.engines.nl_decoder_engine INFO     Overwriting in-place the input shapes of the transformer model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 15:58:59 deepsparse.utils.onnx INFO     Overwriting in-place the batch size of the model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 15:59:12 deepsparse.transformers.engines.nl_decoder_engine INFO     Overwriting in-place the input shapes of the transformer model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 15:59:20 deepsparse.utils.onnx INFO     Overwriting in-place the batch size of the model at /home/ubuntu/damian/sparseml/deployment/model.onnx
def hello_world():def hello_world():def hello_world():def hello_world():def hello_world():def hello_world():def hello

2023-07-20 15:59:33 deepsparse.transformers.engines.nl_decoder_engine INFO     Overwriting in-place the input shapes of the transformer model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 15:59:40 deepsparse.utils.onnx INFO     Overwriting in-place the batch size of the model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 16:01:12 deepsparse.transformers.engines.nl_decoder_engine INFO     Overwriting in-place the input shapes of the transformer model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 16:01:18 deepsparse.utils.onnx INFO     Overwriting in-place the batch size of the model at /home/ubuntu/damian/sparseml/deployment/model.onnx
def hello_world():def hello_world():def hello_world():def hello_world():def hello_world():def hello_world():def hello

2023-07-20 16:02:38 deepsparse.transformers.engines.nl_decoder_engine INFO     Overwriting in-place the input shapes of the transformer model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 16:02:44 deepsparse.utils.onnx INFO     Overwriting in-place the batch size of the model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 16:02:54 deepsparse.transformers.engines.nl_decoder_engine INFO     Overwriting in-place the input shapes of the transformer model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-20 16:03:00 deepsparse.utils.onnx INFO     Overwriting in-place the batch size of the model at /home/ubuntu/damian/sparseml/deployment/model.onnx
def hello_world():def hello_world():def hello_world():def hello_world():def hello_world():def hello_world():def hello

Process finished with exit code 0

[bfineran]

Took a deeper look following our offline conversation and understand why you had to go this route - LGTM, but let's update an existing diagram or add a new one to explain the relationship between decoder engine, cache, state, state transfer, and capacity

[bfineran]

shouldn't this be the total remaining tokens ie something like self.sequence_length - idx * self.prompt_processing_sequence_length or am I missing something?

[dbogunowicz]

We are essentially talking about the same thing, but my logic was way too overcomplicated. Refactored the function, so now hopefully anyone reading should grasp what's going on.