Pyramid KV
新闻
-
[2024-06-25] 现在支持大规模LLM的多GPU推理!在LlaMa-3-70B-Instruct上试试PyramidKV吧!
-
[2024-06-10] 现在在Flash Attention v2, Sdpa Attention中支持PyramidKV, SnapKV, H2O和StreamingLLM!如果你的设备(如V100,3090)不支持Flash Attention v2,可以设置attn_implementation=sdpa在Sdpa Attention中尝试PyramidKV!
TODO:
-
支持Stream LLM, H2O和SnapKV的实现
-
支持Mistral模型
-
支持Needle的实现
-
支持在不使用Flash Attention v2(即Sdpa Attention)的情况下进行KV缓存压缩(适用于V100)
-
支持70B LlaMa-3的多GPU推理
-
引入新功能以支持KV缓存预算分配(例如,支持百分比)
-
支持Mixtral
-
支持批量推理
-
支持在解码阶段进行KV缓存压缩
性能
可视化:无效的注意力
Llama模型的注意力图包含3个文档,表示如下:
我们提供了一个名为visualization.ipynb
的笔记本,用于重现给定3个文档的Llama-2-7b-hf模型每一层的可视化结果。
不同层的模型注意力图将存储在./attention
文件夹中
要求
transformers >= 4.41
flash-attn >= 2.4.0.post1
安装
git clone https://github.com/Zefan-Cai/PyramidKV.git
cd PyramidKV
pip install -r requirements.txt .
推理
我们支持在LongBench
上进行推理代码,以重现我们的结果。
请参考scripts/scripts_longBench/eval.sh
根据你的要求修改参数。
我们的代码库支持Flash Attention v2, Sdpa Attention等。本文中的结果基于Flash Attention v2。
export CUDA_VISIBLE_DEVICES=$1
method=$2 # 支持PyramidKV, SnapKV, H2O, StreamingLLM
max_capacity_prompts=64 # 文中为128, 2048
attn_implementation=$3 # 支持"flash_attention_2", "sdpa", "eager"
source_path=$4
model_path=$5
save_dir=${source_path}"results_long_bench" # 结果保存路径
python3 run_longbench.py \
--method ${method} \
--model_path ${model_path} \
--max_capacity_prompts ${max_capacity_prompts} \
--attn_implementation ${attn_implementation} \
--save_dir ${save_dir} \
--use_cache True
- CUDA_VISIBLE_DEVICES: 对于大规模LLM的多GPU推理,只需指定CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7。对于单GPU推理,只需指定CUDA_VISIBLE_DEVICES=0。
- model_path: 模型路径。目前支持"Llama-3-8B-Instruct"。
- method: 支持
PyramidKV
,SnapKV
,StreamingLLM
,H2O
。 - max_capacity_prompts: 每层选择的KV大小。例如,文中为128, 2048。当方法为"PyramidKV"时,由于总KV数量不变,每层的具体KV长度会相应调整。
- save_dir: 保存LongBench结果的目录路径。
修改参数后,运行:
sh scripts/scripts_longBench/eval.sh
Needle in haystack
我们支持在Needle in haystack
上进行推理代码,以重现我们的结果。
请参考scripts/scripts_needle/eval.sh
根据你的需求修改参数。
我们的代码库支持Flash Attention v2, Sdpa Attention等。本文中的结果基于Flash Attention v2。
METHOD='pyramidkv' # ['full', 'pyramidkv', 'snapkv', 'streamingllm', 'h2o']
MAX_CAPACITY_PROMPT=96 # [64, 96, 128, 256, 512, 1024, 2048, ...]
attn_implementation="flash_attention_2" # 支持"flash_attention_2", "sdpa", ""。
TAG=test
# 对于Llama3-8b
(
python -u run_needle_in_haystack.py --s_len 1000 --e_len 8001\
--model_provider LLaMA3 \
--model_name /mnt/workspace/zhiyuanhu/yuliang/models/llama3-8b_raw \
--attn_implementation ${attn_implementation} \
--step 100 \
--method $METHOD \
--max_capacity_prompt $MAX_CAPACITY_PROMPT \
--model_version LlaMA3_${METHOD}_${MAX_CAPACITY_PROMPT}_${TAG}
) 2>&1 | tee results_needle/logs/LlaMA3_${METHOD}_${MAX_CAPACITY_PROMPT}_${TAG}.log
- LLaMA3和Mistral2都支持单GPU推理。
- model_provider: LLaMA3或Mistral2
- model_name: 你的模型路径。目前支持"Llama-3-8B-Instruct"和"Mistral-7B-Instruct-v0.2"。
- step: 上下文长度的增量。
- method: 支持
PyramidKV
,SnapKV
,StreamingLLM
,H2O
。 - max_capacity_prompt: 每层选择的KV大小。例如,文中为128, 2048。当方法为"PyramidKV"时,由于总KV数量不变,每层的具体KV长度会相应调整。
要重现我们的结果,运行
bash scripts/scripts_needle/eval.sh
推理后,运行
python scripts/scripts_needle/visualize.py
来绘制图像,需在visualize.py
中将FOLDER_PATH
更改为你的输出路径(即eval.sh
中的--model_version
参数)。
引用
如果你认为PyramidKV对你的研究和应用有帮助,请使用以下BibTeX进行引用:
@article{zhang2024pyramidkv,
title={PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling},
author={Zhang, Yichi and Gao, Bofei and Liu, Tianyu and Lu, Keming and Xiong, Wayne and Dong, Yue and Chang, Baobao and Hu, Junjie and Xiao, Wen and others},
journal={arXiv preprint arXiv:2406.02069},
year={2024}
}
致谢
感谢**[SnapKV]** SnapKV: LLM Knows What You are Looking for Before Generation提供的开源代码,支持本项目的扩展。