[KV cache] Base, generic KV Cache injection support

[bfineran]

Feature Preview:

import onnx
from sparseml.exporters.kv_cache_injector import KeyValueCacheInjector

model = onnx.load("deployment/model.onnx")
model = KeyValueCacheInjector(model_type="opt").apply(model)
onnx.save(model, "deployment/model_kvcache.onnx")

The implementation above guarantees full "onnx checker safety".
However, to allow the user to run the transform faster / in insolation, the following
path is also enabled:

import onnx
from sparseml.exporters.kv_cache_injector import KeyValueCacheInjector
model = onnx.load("deployment/model.onnx", load_external_data=False) # we operate only on the model graph
model = KeyValueCacheInjector(model_type="opt").apply(model)
onnx.save(model, "deployment/model_kvcache.onnx")

Note: this will raise multiple warnings, making the user conscious of the fact, that the models with load_external_data=True, cannot be properly validated.

Additional changes:

• makes MatMulAddToMatMulIntegerAddCastMul a bit more generic by making the Add portion optional - should look into renaming this transform...

Testing

This functionality has been tested multiple times for OPT models dense/sparse/quantize:

model = onnx.load("/network/damian/sparsegpt_webinar_fp32/sparsegpt_1.3b/model.onnx")
model = KeyValueCacheInjector(model_type = "opt").apply(model)
onnx.save(model, "temp.onnx")
onnx.checker.check_model("temp.onnx")

2023-05-29 11:26:27 sparseml.exporters.transforms.onnx_transform INFO     [CacheKeysAndValues] Transformed 48 matches
2023-05-29 11:26:29 sparseml.exporters.transforms.onnx_transform INFO     [PositionEmbeddingsAdjustment] Transformed 5 matches

With quantized models, we have observed that the kv cache export leads to the presence of the "dead MatMuls":

However, the current head of the branch does not produce "dead MatMuls" anymore. Does that mean that the problem has been solved along the way?

Also, the quantized models run fine in the pipeline (tested with ORT)

[bfineran]

status so far - confirmed the 'generic' pattern matching works for OPT, next step is adding the injection

[bfineran]

@dbogunowicz initial implementation of KV cache concats + OPT Cache length adjustment completed

for some reason the exporter I wrote was cleared in my environment, will get that up at a later date
Sample code to try:

import onnx
from sparseml.exporters.transforms.kv_cache import *
model = onnx.load("/home/benjamin/tmp-models/small_decoder_opt.onnx", load_external_data=False)
model = CacheKeysAndValues().transform(model)
model = OPTCacheLengthAdjustment().transform(model)

Example Cached MatMul:

Example Cache Length adjustment:

[bfineran]

@dbogunowicz quick comments, pushing up adjustment to reshape

[bfineran]

why did we take this out?

[dbogunowicz]

for now it can be hard-coded as an argument to the function - no changes between CodeGen and OPT

[KSGulin]

This looks great. Building KV cache injection in ONNX feels like the ML equivalent of Chris Sawyer creating RollerCoaster Tycoon in assembly, but this was clean and easy to follow. Left a couple non-blocking comments

[KSGulin]

Is there a reason to check if cache_input_idx == 1 here instead of setting to 3 regardless?