[Fix] Fix the KV Cache insertion logic for quantized OPT

[dbogunowicz]

Feature Description

Fixes the KV Cache logic for quantized text generation models.

In a nutshell, the QuantizeLinear nodes that were created in the process of quantization were messing up our pattern-matching rules for finding Key Matmul and Value Matmul. Additionally the QuantizeLinear nodes are by default inserted into a graph in such a way, that they break the topology of the graph with the kv cache support.

This fix:

• updates the pattern matching for finding Key Matmul and Value MatMul to be robust against the presence of QuantizeLinear nodes
• moves the QuantizeLinear nodes to their appropriate place in the onnx graph.

Manual Testing

Note: For manual testing it is required to use this branch: neuralmagic/deepsparse#1123. It provides the appropriate support for quantized kv cache.

OPT

Tested with the model provided by @natuan nlg-text_generation/c4-opt_1.3b-pruned50_quant/[email protected]@[email protected]@SP0.5@SQ1@PTQ1@ID15513

1. Injecting the kv cache:

python kv_cache_injector.py --input-file /home/ubuntu/damian/ml-experiments/nlg-text_generation/c4-opt_1.3b-pruned50_quant/[email protected]@[email protected]@SP0.5@SQ1@PTQ1@ID15513/deployment/model.onnx --output-file /home/ubuntu/damian/ml-experiments/nlg-text_generation/c4-opt_1.3b-pruned50_quant/[email protected]@[email protected]@SP0.5@SQ1@PTQ1@ID15513/deployment/model_kvcache.onnx

2023-07-18 09:58:01 sparseml.exporters.transforms.kv_cache.configs INFO     Loaded config file /home/ubuntu/damian/ml-experiments/nlg-text_generation/c4-opt_1.3b-pruned50_quant/[email protected]@[email protected]@SP0.5@SQ1@PTQ1@ID15513/deployment/config.json for model: opt
2023-07-18 09:58:01 sparseml.exporters.transforms.kv_cache.configs INFO     Properly configured arguments for KV Cache Transformation
Attempting to validate an in-memory ONNX model that has been loaded without external data. This is currently not supported by the ONNX checker. The validation will be skipped.
2023-07-18 09:58:03 sparseml.exporters.transforms.onnx_transform INFO     [CacheKeysAndValues] Transformed 48 matches
Attempting to validate an in-memory ONNX model that has been loaded without external data. This is currently not supported by the ONNX checker. The validation will be skipped.
Attempting to validate an in-memory ONNX model that has been loaded without external data. This is currently not supported by the ONNX checker. The validation will be skipped.
2023-07-18 09:58:04 sparseml.exporters.transforms.onnx_transform INFO     [PositionsAdjustmentOPT] Transformed 3 matches
Attempting to validate an in-memory ONNX model that has been loaded without external data. This is currently not supported by the ONNX checker. The validation will be skipped.
Modified model saved to: /home/ubuntu/damian/ml-experiments/nlg-text_generation/c4-opt_1.3b-pruned50_quant/[email protected]@[email protected]@SP0.5@SQ1@PTQ1@ID15513/deployment/model_kvcache.onnx

3. Renaming deployment/model_kvcache.onnx to deployment/model.onnx and running it in the pipeline.

opt = Pipeline.create(task="opt", model_path="/home/ubuntu/damian/ml-experiments/nlg-text_generation/c4-opt_1.3b-pruned50_quant/[email protected]@[email protected] [email protected]@SQ1@PTQ1@ID15513/deployment", engine_type=engine_type, use_deepsparse_cache = False, max_generated_tokens=32)

out = opt(sequences="Who is the president of the United States?")
print(out.sequences[0])

Who is the president of the United States?

Who is the president of the United States?

Who is the president of the United States

Note: This PR does not alter the behavior of KV cache injection for non-quantized OPT model.

CodeGen

Tested with the model provided by @shubhra codegen_mono-350m-apps_bigpython_bigquery_thepile-base_quantized/training/

1. Injecting the kv cache:

python kv_cache_injector.py --input-file /home/ubuntu/damian/codegen_mono-350m-apps_bigpython_bigquery_thepile-base_quantized/training/model.onnx --output-file /home/ubuntu/damian/codegen_mono-350m-apps_bigpython_bigquery_thepile-base_quantized/training/model_kvcache.onnx

2023-07-18 11:22:34 sparseml.exporters.transforms.kv_cache.configs INFO     Loaded config file /home/ubuntu/damian/codegen_mono-350m-apps_bigpython_bigquery_thepile-base_quantized/training/config.json for model: codegen
2023-07-18 11:22:34 sparseml.exporters.transforms.kv_cache.configs INFO     Properly configured arguments for KV Cache Transformation
2023-07-18 11:22:35 sparseml.exporters.transforms.onnx_transform INFO     [CacheKeysAndValues] Transformed 40 matches
2023-07-18 11:22:38 sparseml.exporters.transforms.onnx_transform INFO     [PositionsAdjustmentCodeGen] Transformed 3 matches
Modified model saved to: /home/ubuntu/damian/codegen_mono-350m-apps_bigpython_bigquery_thepile-base_quantized/training/model_kvcache.onnx

3. Renaming training/model_kvcache.onnx to training/model.onnx and running it in the pipeline.

opt = Pipeline.create(task="codegen", model_path="/home/ubuntu/damian/codegen_mono-350m-apps_bigpython_bigquery_thepile-base_quantized/training", engine_type=engine_type, use_deepsparse_cache = False, max_generated_tokens=32)

out = opt(sequences="def hello_world():")
print(out.sequences[0])

print("Hello World")

hello_world()

# This is a comment

# This is a comment

# This is

Note: This PR does not alter the behavior of KV cache injection for non-quantized CodeGen model.