执行RACompressor后,可在指定路径下,通过get_alibi_windows()函数生成.pt文件。
RACompressor.get_alibi_windows(save_path)
参数名 |
输入/返回值 |
含义 |
使用限制 |
---|---|---|---|
save_path |
输入 |
长序列压缩时,head压缩窗口参数文件的保存路径。 |
必选。 数据类型:String。 |
from msmodelslim.pytorch.ra_compression import RACompressConfig, RACompressor config = RACompressConfig(theta=0.00001, alpha=100) model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path="baichuan2-13b/float_path/", trust_remote_code=True).float().cpu() # 需根据模型的实际路径配置 ra = RACompressor(model, config) ra.get_alibi_windows(save_path)