RoPE编码类型
RoPE(Rotary Position Embedding)编码是一种高效的位置编码方式,有以下特点:
- 旋转编码:通过旋转操作将位置信息编码到每个token的嵌入向量中。这种旋转操作确保了模型能够捕捉到序列中元素的相对位置信息,而不依赖于绝对位置。
- 维度保持:旋转操作在每个维度上独立进行,有助于模型在不同的特征维度上捕获位置信息。
- 计算效率:不需要额外的参数来编码位置信息,而是通过数学旋转操作来实现,计算效率较高。
通过RoPE编码与RazorAttention结合,可分析注意力头对位置编码的依赖性,来决定哪些头可以被压缩,以优化模型的存储、传输和计算效率,提高模型在实际应用中的可部署性和实用性。
- 利用RoPE编码的位置信息:由于RoPE编码已经有效地将位置信息编码到每个token中,RA算法可以利用这一点来更好地识别哪些注意力头对于位置信息更为敏感,从而更有针对性地进行压缩。
- 优化压缩策略:通过结合RoPE编码和RA算法,可以在保持模型性能的同时,针对不同的注意力头实施更精细的压缩策略。例如,对于依赖位置信息的Retrieval Head,可以保持其KV Cache的完整性,而对于不依赖位置信息的Non-retrieval Head,则进行压缩。
目前支持对表1中RoPE编码的大模型进行长序列压缩(包括但不限于)。
前提条件
- 已参考环境准备,完成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) pip3 install torch_npu-2.1.0.xxx.whl # xxx需要根据实际情况进行选择,具体请参见安装torch_npu插件
功能实现流程
图1 压缩接口调用流程

关键步骤说明如下:
- 用户准备原始模型。
- 调用RARopeCompressConfig接口生成压缩配置,并新建模型的压缩脚本run.py。
- 调用RARopeCompressor启动长序列压缩任务,进行长序列压缩。
- 调用get_compress_heads接口导出需保留的Head信息,并在指定路径获取.pt文件。
- 用户可根据.pt文件进行压缩
- 压缩后的文件可用于后续的推理部署,具体请参见MindIE是什么的“MindIE支持模型列表”章节中查看已适配的量化模型。
压缩步骤(以Qwen2-72b-instruct为例)
- 用户准备原始模型。
- 新建模型的量化脚本run.py,并将如下样例代码导入run.py文件,并执行以下命令。
import torch from msmodelslim.pytorch.ra_compression import RARopeCompressConfig, RARopeCompressor from transformers import AutoTokenizer, AutoModelForCausalLM import torch_npu torch.npu.set_compile_mode(jit_compile=False) config = RARopeCompressConfig(induction_head_ratio=0.14, echo_head_ratio=0.01) save_path = "./win.pt" model_path = "./Qwen2-72B-Instruct/" model = AutoModelForCausalLM.from_pretrained( pretrained_model_name_or_path=model_path, torch_dtype=torch.bfloat16, trust_remote_code=True, device_map="auto", ).eval() tokenizer = AutoTokenizer.from_pretrained( pretrained_model_name_or_path=model_path, pad_token='<|extra_0|>', eos_token='<|endoftext|>', padding_side='left', trust_remote_code=True ) ra = RARopeCompressor(model, tokenizer, config) ra.get_compress_heads(save_path)
- 启动长序列压缩任务,并在“Qwen2-72b-instruct”文件夹的路径下获取需要保留KV Cache的Head信息的.pt文件。
- 用户可根据.pt文件进行压缩。
- 压缩后的文件可用于后续的推理部署,具体请参见MindIE是什么的“MindIE支持模型列表”章节中查看已适配的量化模型。
python3 run.py
父主题: 长序列压缩