Alibi编码类型

Alibi编码是一种位置编码方法，与RazorAttention结合使用，通过Alibi编码来识别哪些注意力头对位置信息更为敏感，从而决定哪些头可以被压缩。Alibi编码并不直接在网络中加入显式的位置编码，而是通过在query-key注意力分数上施加一个与距离成比例的偏置实现位置信息的建模。

KV Cache的管理需考虑batch, seqlen, num_heads和head_size这四个维度，其中seqlen维度通常是压缩的重点，因为随着序列长度的增加，KV Cache的内存占用会迅速增长。传统的压缩方法可能会忽略不同注意力头（Heads）之间的差异，而RazorAttention加速技术则提供了一种更细粒度的内存压缩方法，针对使用Alibi编码的模型进行优化，可以更有效地识别哪些注意力头对于位置信息更为敏感，并据此调整压缩策略。RazorAttention加速技术支持全量加速和增量加速：

目前支持对表1中Alibi编码的大模型进行长序列压缩（包括但不限于）。

表1 已验证模型列表

模型名称

框架

baichuan2-13b

PyTorch

前提条件

功能实现流程

关键步骤说明如下：

  1. 用户准备原始模型。
  2. 调用RACompressConfig接口生成压缩配置，新建模型的压缩脚本run.py。
  3. 执行压缩算法RACompressor启动长序列压缩任务，进行长序列压缩。
  4. 调用get_compress_heads接口导出压缩窗口，并在指定路径中获取.pt文件，具体请参见MindIE的“加速库支持模型列表”章节中已适配量化的模型。

压缩步骤（以baichuan2-13b为例）

  1. 用户准备原始模型。

    用户需要自行准备模型、权重文件。本样例以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

  2. 新建模型的压缩脚本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)

  3. 执行以下命令，启动长序列压缩任务，并在“baichuan 2-13b”文件夹的路径下获取.pt文件。

    python3 run.py

  4. .pt文件可用于后续的推理部署任务，具体请参见MindIE的“加速库支持模型列表”章节中已适配量化的模型。
父主题： 长序列压缩