Alibi编码类型
Alibi编码是一种位置编码方法，与RazorAttention结合使用，通过Alibi编码来识别哪些注意力头对位置信息更为敏感，从而决定哪些头可以被压缩。Alibi编码并不直接在网络中加入显式的位置编码，而是通过在query-key注意力分数上施加一个与距离成比例的偏置实现位置信息的建模。
KV Cache的管理需考虑batch, seqlen, num_heads和head_size这四个维度，其中seqlen维度通常是压缩的重点，因为随着序列长度的增加，KV Cache的内存占用会迅速增长。传统的压缩方法可能会忽略不同注意力头（Heads）之间的差异，而RazorAttention加速技术则提供了一种更细粒度的内存压缩方法，针对使用Alibi编码的模型进行优化，可以更有效地识别哪些注意力头对于位置信息更为敏感，并据此调整压缩策略。RazorAttention加速技术支持全量加速和增量加速：
- 全量加速：压缩后的KV Cache可直接用于模型推理，实现全量加速。
- 增量加速：支持只更新和压缩新token对应的KV Cache部分。
目前支持对表1中Alibi编码的大模型进行长序列压缩（包括但不限于）。
前提条件
- 已参考环境准备，完成CANN开发环境的部署、PyTorch 2.1.0及以上版本的安装及Python环境变量的配置。
- 执行命令安装如下依赖。
以下命令若使用非root用户安装，需要在安装命令后加上--user，例如：pip3 install numpy==1.25.2 --user。
pip3 install numpy==1.26.4 pip3 install transformers==4.43.1 pip3 install torch==2.1.0 # 安装CPU版本的PyTorch 2.1.0(不依赖torch_npu)
功能实现流程
图1 压缩接口调用流程
关键步骤说明如下：
- 用户准备原始模型。
- 调用RACompressConfig接口生成压缩配置，新建模型的压缩脚本run.py。
- 执行压缩算法RACompressor启动长序列压缩任务，进行长序列压缩。
- 调用get_compress_heads接口导出压缩窗口，并在指定路径中获取.pt文件，具体请参见MindIE是什么的“MindIE支持模型列表”章节中查看已适配的量化模型。
压缩步骤（以baichuan2-13b为例）
- 用户准备原始模型。用户需要自行准备模型、权重文件。本样例以baichuan2-13b为例，从该网站下载权重文件，并上传至服务器的“baichuan 2-13b”文件夹，目录示例如下：
config.json configuration_baichuan.py cut_utils.py generation_config.json generation_utils.py handler.py model-00001-of-00003.safetensors model-00002-of-00003.safetensors model-00003-of-00003.safetensors modeling_baichuan_cut.py modeling_baichuan.py model.safetensors.index.json pytorch_model.bin.index.json quantizer.py special_tokens_map.json tokenization_baichuan.py tokenizer_config.json tokenizer.model
- 新建模型的压缩脚本run.py，将如下样例代码导入run.py文件，并执行。
from msmodelslim.pytorch.ra_compression import RACompressConfig, RACompressor from transformers import AutoTokenizer, AutoModelForCausalLM config = RACompressConfig(theta=0.00001, alpha=100) # 压缩类的配置，需根据实际情况进行修改 input_model_path = "/data1/models/baichuan/baichuan2-13b/float_path/" # 模型权重文件的保存路径，需根据实际情况进行修改 save_path = "./win.pt" # 生成压缩窗口的路径，需根据实际情况进行修改 tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=input_model_path, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=input_model_path, trust_remote_code=True).float().cpu() # 不支持使用npu方式进行加载 ra = RACompressor(model, config) ra.get_alibi_windows(save_path)
- 执行以下命令，启动长序列压缩任务，并在“baichuan 2-13b”文件夹的路径下获取.pt文件。
python3 run.py
- .pt文件可用于后续的推理部署任务，具体请参见MindIE是什么的“MindIE支持模型列表”章节中查看已适配的量化模型。
父主题： 长序列压缩