[KV Cache Injection] Causal Mask for OPT

[dbogunowicz]

Add causal mask support for OPT models to enable multitoken prefill in the Deepsparse pipeline.

Manual Testing

1. Export the OPT model

python kv_cache_injector.py --input-file deployment/model.onnx --output-file deployment/model_kvcache.onnx

2. Inject the KV Cache

python kv_cache_injector.py --input-file deployment/model.onnx --output-file deployment/model_kvcache.onnx
2023-07-25 13:18:44 sparseml.exporters.transforms.kv_cache.configs INFO     Loaded config file deployment/config.json for model: opt
2023-07-25 13:18:44 sparseml.exporters.transforms.kv_cache.configs INFO     Properly configured arguments for KV Cache Transformation
2023-07-25 13:18:46 sparseml.exporters.transforms.onnx_transform INFO     [CacheKeysAndValues] Transformed 48 matches
2023-07-25 13:18:49 sparseml.exporters.transforms.onnx_transform INFO     [AdditionalTransformsOPT] Transformed 2 matches

3. Run inference (using this branch: [Text Generation] Causal Mask Feature Branch deepsparse#1126)

from deepsparse import Pipeline

def _test_pipeline(engine_type):
    opt = Pipeline.create(task="opt",
                          model_path="/home/ubuntu/damian/sparseml/deployment",
                          engine_type=engine_type,
                          prompt_processing_sequence_length=64,
                          use_deepsparse_cache = False,
                          max_generated_tokens=32)
    print('----------')
    prompt = "Who is the president of the United States?" # the prompt is short, will not be processed by self.multitoken_engine
    out = opt(sequences=prompt, return_logits = True)
    print(out.sequences[0])
    print('---------')
    prompt = "Who is the president of the United States?" * 20 # the prompt is long, will be processed by self.multitoken_engine
    out = opt(sequences=prompt, return_logits=True)
    print(out.sequences[0])

_test_pipeline(engine_type ="onnxruntime")
_test_pipeline(engine_type ="deepsparse")

2023-07-25 13:23:35 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-25 13:23:41 deepsparse.utils.onnx INFO     Overwriting in-place the batch size of the model at /home/ubuntu/damian/sparseml/deployment/model.onnx
2023-07-25 13:23:51 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-25 13:23:57 deepsparse.utils.onnx INFO     Overwriting in-place the batch size of the model at /home/ubuntu/damian/sparseml/deployment/model.onnx
----------


The president of the United States is the head of the executive branch of government. The president is the head of the executive branch of government, and the
---------
Who is the president of the United States?Who is the president of the United States?Who is the president of the United States?Who is the president of
/home/ubuntu/damian/deepsparse/src/deepsparse/transformers/pipelines/text_generation.py:137: UserWarning: AVX512 support not detected, disabling internal management of KV cache which may affect performance. To enable full performance, deploy on an AVX512-compatible system.
  warnings.warn(
2023-07-25 13:24:16 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
DeepSparse, Copyright 2021-present / Neuralmagic, Inc. version: 1.6.0.20230725 COMMUNITY | (f26e1c2e) (release) (optimized) (system=avx2, binary=avx2)
2023-07-25 13:24:31 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
----------


The president of the United States is the head of the executive branch of government. The president is the head of the executive branch of government, and the
---------
Who is the president of the United States?Who is the president of the United States?Who is the president of the United States?Who is the president of

[bfineran]

LGTM pending rebase and comment regarding strong preference on using cast over where

[bfineran]

looks like this needs rebase?

[bfineran]

rebase?

[bfineran]

why are we using a where instead of a cast to bool? can you check with runtime for what would work better if either are fine? additionally, seems like a cast would read better in onnx vs where which involves a condition...