compressed-tensors support for KV cache Quantization

[dbogunowicz]

Feature description

Implements quantized kv cache support for the transformer models that have been quantized using compressed-tensors.

Introduces:

• CompressedTensorsQuantizedCacheConfig - a config object that stores the static qparams for quantizing/dequantizing KV Cache
• minor improvements to the QuantizedCache interface, to make it more general in the future for other implementations
• CompressedTensorsCache - very lightweight wrapper around QuantizedCache, that allows us to enable kv cache quantization

Manual testing:

from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers.cache_utils import CompressedTensorsQuantizedCacheConfig
import torch
import time 
investigate_mem_consumption = True

model_id = "/root/compressed-tensors/llama1.1b_new_quant_out" # quantized model with kv cache quantization enabled

tokenizer = AutoTokenizer.from_pretrained("Xenova/llama2.c-stories15M") # somehow the tokenizer is missing from the model in `model_id`
tokenizer.pad_token_id = tokenizer.eos_token_id
model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="eager").to("cuda:0")

inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

cache_config = CompressedTensorsQuantizedCacheConfig.from_pretrained(model_id)

out_quant = model.generate(**inputs, cache_implementation="quantized", cache_config=cache_config, min_new_tokens=40, return_dict_in_generate=True)
out = model.generate(**inputs, min_new_tokens=40, return_dict_in_generate=True)

assert out_quant.sequences.allclose(out.sequences) # assert same tokens get generated regardless of cache type
assert out_quant.past_key_values._quantized_key_cache[-1].dtype == torch.int8 # assert that we are actually caching quantized tensors

Note

Compatible with this branch of compressed-tensors: neuralmagic/compressed-tensors#86
Pending missing items: Add tests, add compressed-tensors as transformers dependency, fix ugly import issues (circular imports between transformers and compressed-tensors)

[dbogunowicz]

I will take care of this hack once the PR is in the "landable" state

[dbogunowicz]

Instead of rewriting update() method for CompressedTensorsCache I decided to expand the interface.

[dbogunowicz]

We should potentially add more supported quantization types, to be discussed @bfineran

[dbogunowicz]

@mgoin

from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers.cache_utils import CompressedTensorsQuantizedCacheConfig, CompressedTensorsCache

model_id = "/root/compressed-tensors/llama1.1b_new_quant_out" # quantized model with kv cache quantization enabled

tokenizer = AutoTokenizer.from_pretrained("Xenova/llama2.c-stories15M") # somehow the tokenizer is missing from the model in `model_id`
tokenizer.pad_token_id = tokenizer.eos_token_id
model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="eager").to("cuda:0")

inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

cache_config = CompressedTensorsQuantizedCacheConfig.from_pretrained(model_id)
cache_object = CompressedTensorsCache(cache_config)

out = model(**inputs)
out_quant = model(**inputs, past_key_values=cache_object)

assert (out.logits == out_quant.logits).all()

working fine!