Home

Awesome

KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches

Gul'dan_everything

Jiayi Yuan*, Hongyi Liu*, Shaochen (Henry) Zhong*, Yu-Neng Chuang, Songchen Li, Guanchu Wang, Duy Le, Hongye Jin, Vipin Chaudhary, Zhaozhuo Xu, Zirui Liu, Xia Hu

This is the official implementation of our above-mentioned KV cache compression benchmark paper (arXiv). Work corresponds to Shaochen (Henry) Zhong and Zirui Liu.

Changelog and To-do

Overview

Token evictions, KV cache quantizations, hard prompt compression, RNNs, RNN-Transformer hybrids — plenty of efficiency approaches claim they are long context capable, but which can stand up to comprehensive scrutiny, and what are the trade-offs? Our benchmark paper attempts to answer such questions by evaluating many examplar methods among the above schools against various long context tasks.

cover_vis

In this repo, we provide the implementation of all featured methods and evaluations. We intentionally make our codebase minimalistic for easier hacking and reproducing needs, yet we keep it extensible to include alternative or future-coming approaches that are not under our benchmark coverage. We plan to gradually incorporate more interesting methods and datasets into this repo.

Environment Setup

We provide the minimum environment requirements to support the running of our project. This means there can be a slight difference depending on the actual automatic dependency-solving result of different systems. Given the efficiency focus of our benchmark, some of the featured methods require conflicting environments, causing us unable to provide a single unified environment for all featured methods. Thus, we try to control the version of some of the major packages (torch, transformers, etc.) and provide two sets of requirements.

Should one be interested in reproducing a certain method, please look up the corresponding requirement file and install listed packages accordingly.

cd longctx_bench
pip install --upgrade pip
pip install -r requirements/<requirment_file_name>

We note that some packages often require installation later than other packages (e.g., flash-attn or custom kernel supports). In such cases, we commented out such "follow-up" packages in their requirement files so that users could manually install them later. Please carefully inspect the full requirement file and ensure all necessary packages are installed. As an example, suppose one wants to try out KIVI under our codebase; one should:

cd longctx_bench
pip install -r requirements/tf_4.36.txt
pip install flash-attn==2.6.3
cd pipeline/kivi/quant & pip install -e .

Dataset and Access Preparation

Currently, our benchmark features 15 datasets (NarrativeQA, Qasper, MultiFieldQA, HotpotQA, 2WikiMQA, Musique, GovReport, QMSum, MultiNews, TREC, TriviaQA, SAMSum, PassageRetrieval, LCC, RepoBench-P) from LongBench and a Needle-in-a-Haystack/PaulGraham Passkey test with Paul Graham essays as background and passkey retrieval as needle, following the spirit of Arize-ai's NIAH implementation utilized in Google DeepMind's Gemini 1.5 report.

The needle test's dataset are constructed on the fly, where all necessary material is already supplied in this repo. However, the LongBench datasets would require some extra preparations. We supply the following script to process LongBench:

python scrips/dataset_prep/download_longbench.py

Our paper features some models that are gated (e.g., meta-llama/Meta-Llama-3-8B-Instruct). So please supply your HuggingFace access token under config/access_tokens.py. You may consider setting git update-index --assume-unchanged config/access_tokens.py so that Git will no longer track this file to avoid your locally stored token accidentally get synced to upstream.

Experiment Reproduction

We supply all scripts in the scripts folder, with a folder structure that clearly indicate which script is for which experiment. E.g., should one want to verify the needle results of mistral-7b-instruct-v0.2 model without any compression applied, one may do so via:

bash scripts/paulgraham_passkey/20480words_10x10x3_7digits/baseline/mistral-7b-instruct-v0.2.sh <output_dir_root>

Similarly, should one expect the LongBench results of llama3-8b-instruct model with KIVI-powered compression in 2bit, one may achieve such via:

bash scripts/longbench/kivi/2bit/llama3-8b-instruct.sh <output_dir_root>

Given the long context tasks are often time-consuming to evaluate, we additionally supply a set of scripts in the scripts/quick_start folder, which are quick to conclude (just Qasper for LongBench and just two of the longest needle inputs for PaulGraham Passkey). They are solely here as proxies to confirm the code and environment are (likely) running fine before committing to longer evaluations. We recommend trying these scripts first.

Result Digestion

Once an experiment is running, one may monitor real-time printouts in the terminal as well as the exp.log file under <output_folder_dir>. Once an experiment is concluded, the final results can be find in a subfolder under <output_dir_root>. The exact subfolder location varies depending on what experiment is being executed since our provided script would make a subfolder structure that is legible and makes sense for each method. Here is an example for InfLLM on PaulGraham Passkey; the results can be found under <output_folder_dir>.

compress_ratio=$1 # 2x, 4x, 6x, 8x
output_dir_root=$2

task="paulgraham_passkey"
dataset="20480words_10x10x3_7digits"
model="mistral-7b-instruct-v0.2_infinity_model_len"
method="infllm"

python pipeline/inf_stream_llm/main.py \
--exp_desc ${task}_${dataset}_${model}_${method}_${compress_ratio} \ # this argument is purely cosmetic.
--pipeline_config_dir config/pipeline_config/${method}/${task}/${model}/${compress_ratio}.json \
--eval_config_dir config/eval_config/${task}/${dataset}.json \
--output_folder_dir ${output_dir_root}/${task}/${dataset}/${method}/${model}/${compress_ratio}/

Under this <output_folder_dir> folder, one may expect the following components.

For LongBench, we treat each of its subdatasets as an individual eval and loop them over, as demoed here. In this case, we additionally channel the LongBench summarization scripts to summarize all results corresponding to a full LongBench run into one json file at the subdataset level.

Codebase Design and Contribution

Our codebase mainly revolves around the two configs — pipeline_config and eval_config — where the former one defines a method, and the latter one defines an evaluation scheme. Their code implementations can be respectively found under the pipeline and eval folders. We keep the eval implementations "singleton" for each dataset to ensure they are fairly tested among different pipelines. Yet, for pipeline implementation, we intentionally repeat some code components (e.g., the pipeline/<method>/eval_longbench.py and eval_passkey_retrieval.py files are largely the same under different methods) as we want to keep each pipeline/method folder relatively self-contained for better readability; and to make everything loosely coupled to support intervention requests coming at different angles.

Should you want to add a new evaluation, you may consider adding an eval/<new_eval_utils> folder for its eval utilities, and supply corresponding eval_<new_eval>.py files under all pipeline/<method> folders you'd care to evaluate. Should you want to add a new method to our codebase, please consider adding a pipeline/<method> folder and couple with the evaluations you'd care to conduct with necessary eval_<a_certain_eval>.py files.

Should you need to refer to this work or find our codebase useful, please consider citing our work as:

@misc{yuan_liu_zhong_2024_kvcache_comp_benchmnark,
      title={KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches},
      author={Jiayi Yuan and Hongyi Liu and Shaochen Zhong and Yu-Neng Chuang and Songchen Li and Guanchu Wang and Duy Le and Hongye Jin and Vipin Chaudhary and Zhaozhuo Xu and Zirui Liu and Xia Hu},
      year={2024},
      eprint={2407.01527},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2407.01527},
}