SpQR 模型压缩

它伴随的研究论文是“SpQR：用于近无损 LLM 权重压缩的稀疏量化表示”（https://arxiv.org/abs/2306.03078）。

安装

软件包

要使用 falcon 运行 SpQR，请确保您已安装支持 CUDA 的 torch>=2.0.0。

从 requirements.txt 安装软件包：

pip install -r requirements.txt

注意： ArXiv 论文中报告的结果是使用 transformers 的 4.28.dev0 版本获得的，其提交 ID 为 464d420775。

加载/缓存数据集和分词器

该脚本需要下载并本地缓存相关的分词器和数据集。它们将被保存在默认的 Huggingface 数据集目录中，除非通过环境变量指定了其他位置。请参阅 相关数据集文档部分。

模型

目前，此仓库预计可与 LLaMA、Falcon 和 OPT 系列模型配合使用。

数据

对于使用 SpQR 进行量化，建议使用模型训练时所用数据的一个子集。例如，对于 LLaMA 模型的量化，我们推荐使用 RedPajama 的子集；而对于 Falcon 的量化，则推荐使用 RefinedWeb。这两个子集都存储在 data 目录中：

• data/red_pajama_n=1024.pth
• data/refined_web_n=128.pth

注意 这些子集已经过相应模型分词器的处理。如果用于其他模型，可能会导致意外行为。

对于 OPT，根据 GPTQ 论文的建议，我们推荐使用 c4。

W&B 日志记录

为方便起见，可以选择将数据记录到 Weights and Biases 服务（wandb）。运行 pip install wandb 即可启用 W&B 日志记录。在运行实验之前，请指定 $WANDB_ENTITY、$WANDB_PROJECT 和 $WANDB_NAME 环境变量，并使用 --wandb 参数启用日志记录。

启动

GPU 和 RAM 要求

此代码是在配备 80GB 显存的单个 A100 GPU 上开发和测试的。它也可以在具有 32GB 以上显存的 GPU 上成功运行，用于评估高达 LLaMA-65B 和 Falcon-40B 模型的困惑度。使用 --offload activations 选项后，可以在显存较少的机器上评估模型的困惑度：Llama 65B 需要 24GB 以上，Llama 7B 则需要 6GB 以上。

困惑度测试代码还需要足够的内存来容纳未压缩的模型权重（例如，Llama 65B 约需 130GB）以及测试数据集。对于 Language Model Evaluation Harness 评估，需要有足够的内存来将整个模型加载到一台或多台设备上，再加上激活张量。

模型下载

代码要求以 Huggingface 格式下载 LLaMA 模型并本地保存。以下脚本假定 $TRANSFORMERS_CACHE 变量指向 Huggingface Transformers 的缓存文件夹。

困惑度基准测试

此脚本会先压缩模型，然后使用 WikiText2、C4 和 Penn Treebank 数据集测试其困惑度性能。

启动脚本的命令应如下所示：

export MODEL_PATH=<模型目录路径>
export DATASET=<插入数据集名称或自定义数据路径>

python main.py $MODEL_PATH $DATASET \
    --wbits 4 \
    --groupsize 16 \
    --perchannel \
    --qq_scale_bits 3 \
    --qq_zero_bits 3 \
    --qq_groupsize 16 \
    --outlier_threshold=0.2 \
    --permutation_order act_order \
    --percdamp 1e0 \
    --nsamples 128 

上述命令执行了文章中描述的近无损压缩。调整这些参数可以实现更紧密的压缩，但会带来略微更大的损失。

请注意启动参数：

• <模型目录路径> — 包含 config.json 文件的模型文件夹路径。
• [c4, ptb, wikitext2, pajama, refinedweb, none] 中的一个 — 用于压缩的数据集名称，或预处理并分词过的替代数据集路径。
• --wbits 3 — 量化权重表示的位数。
• --groupsize 16 — 用于压缩的一阶分组大小。
• --qq_groupsize 16 — 用于压缩的二阶（量化）分组大小。
• --qq_scale_bits 3 --qq_zero_bits 3 — 用于量化一阶权重尺度和零值的位数。
• --offload activations — 在不使用激活时将其移至 RAM。这会减少显存占用，但会使运行速度降低约 10%。运行 python main.py --help 可获取有关命令行参数的更多详细信息，包括压缩参数。
• --save --load — 用于保存/加载量化模型的路径。

LM Evaluation Harness 基准测试

为了进行零样本评估，我们使用经过轻微修改的 Language Model Evaluation Harness 框架。此仓库在 lm-eval-harness 文件夹中包含 2023 年初版本的 LM Evaluation Harness 代码库。

安装

在运行代码之前，请确保已安装 lm-eval-harness 的所有要求和依赖项。运行以下命令进行安装：

pip install -r lm-evaluation-harness/requirements.txt

执行

启动评估程序的主要脚本是 lmeval.py。

注意：当前版本的脚本仅支持 LLaMA/Falcon 的量化。因此，请设置：

• --model=hf-causal
• --model_args pretrained=$MODEL_PATH，其中 $MODEL_PATH 必须是 LLaMA 模型之一。

--quantization_args — 量化器的逗号分隔参数列表。有关详细信息和选项，请参阅 spqr_config.py。

以下是基准测试启动示例：

export MODEL_PATH=<插入模型目录路径>
export DATASET=<插入数据集名称或自定义数据路径>

python lmeval.py \
    --model hf-causal \
    --model_args pretrained=$MODEL_PATH,dtype=float16,use_accelerate=True \
    --quantization_args dataset=$DATASET,wbits=4,groupsize=16,perchannel=True,qq_scale_bits=3,qq_zero_bits=3,qq_groupsize=16,percdamp=1.0,outlier_threshold=0.2,simplified_outliers=False,nsamples=128,offload_activations=True \
    --tasks winogrande,piqa,hellaswag,arc_easy,arc_challenge \
    --batch_size 1

性能和运行注意事项：

• 对于大型模型（LLaMA-30B、LLaMA-65B），请指定 max_memory_per_gpu={value}GIB，以便每块 GPU 保留 15-20GIB 的空闲显存，用于存储校准所需的激活数据。
• offload_activations=True 可略微降低峰值内存消耗。
• 通常，LLaMA-30B 需要 1-2 块 80Gb 显存的 A100 GPU，而 LLaMA-65B 则需要 3 块 80Gb 显存的 A100 GPU。
• 如果有充足的空闲显存，可以提高批量大小以加快评估进程。

推理

该仓库还包含 SpQR 矩阵-向量乘法的高效 CUDA 内核实现。文件 inference_demo.py 包含一个端到端模型推理的演示。以下是运行该脚本的示例。

用法: inference_demo.py [-h] [--pretrained_model_path PRETRAINED_MODEL_PATH] [--compressed_model_path COMPRESSED_MODEL_PATH] --execution_mode {0,1}

选项:
  -h, --help            显示此帮助消息并退出
  --pretrained_model_path PRETRAINED_MODEL_PATH
                        预训练模型的路径
  --compressed_model_path COMPRESSED_MODEL_PATH
                        压缩后的 .pt 模型路径
  --execution_mode {0,1}
                        如果设置为 0，则评估密集的预训练模型。如果设置为 1，则评估 SpQR 量化模型

该脚本还会报告 forward() 步骤的平均和中位数时间，以及总的推理执行时间。

运行转换脚本、测试和基准测试的先决条件

为了运行基准测试和测试套件，您需要构建这些脚本所使用的源代码。可以通过运行以下命令来完成：

/bin/bash scripts/build.sh 

该命令会简单地运行 setup.py 脚本。

从旧版到优化 SPQR 存储格式的转换

在运行 SpQR 并生成以 int8 存储的张量后，为了使用高效的推理内核，必须将 SpQR 生成的张量（旧版张量）转换为 CUDA 内核所使用的优化存储格式。为此，请运行以下脚本：

用法: convert_legacy_model_format.py [-h] --base_model BASE_MODEL --legacy_model_path LEGACY_MODEL_PATH [--sparse_strategy {csr,ptcsr,optimize_latency}] [--save_pt SAVE_PT] [--save_per_layer SAVE_PER_LAYER]

选项:
  -h, --help            显示此帮助消息并退出
  --base_model BASE_MODEL
                        未量化模型的路径或名称
  --legacy_model_path LEGACY_MODEL_PATH
                        旧版模型的路径
  --sparse_strategy {csr,ptcsr,optimize_latency}
                        稀疏存储策略。选项：csr、ptcsr、auto。CSR 表示压缩稀疏行格式；PTCSR 表示另一种存储格式；optimize_latency 表示使用当前 GPU 自动确定最优存储格式以减少内核延迟。
  --save_pt SAVE_PT     将转换后的量化 .pt 模型保存至此处
  --save_per_layer SAVE_PER_LAYER
                        将转换后的量化 m

Hugging Face 格式转换

要将模型转换为 Hugging Face 兼容格式，请使用 convert_to_hf.py 脚本：

用法: convert_to_hf.py [-h] [--model MODEL] [--config_path CONFIG_PATH] [--in_path_pt IN_PATH_PT] [--out_path OUT_PATH] [--save_safetensors] [--trust_remote_code] [--load_model] [--save_tokenizer]

选项:
  -h, --help            显示此帮助消息并退出
  --model MODEL         用于基于 AutoConfig.from_pretrained() 构建配置的基础模型路径
  --config_path CONFIG_PATH
                        用于基于 AutoConfig.from_pretrained() 构建配置的基础模型路径
  --in_path_pt IN_PATH_PT
                        要转换的检查点路径
  --out_path OUT_PATH   保存 Hugging Face 兼容检查点的路径
  --save_safetensors    是否以 safetensors 格式保存
  --trust_remote_code   是否信任远程代码
  --load_model          是否加载模型
  --save_tokenizer      是否保存分词器

基准测试（matvec 内核）

要运行 matvec 基准测试套件，应执行以下命令：

bench_spqr.py [-h] --tensor_path TENSOR_PATH [--ptcsr_path PTCSR_PATH] [--output_path OUTPUT_PATH]

选项:
  -h, --help            显示此帮助消息并退出
  --tensor_path TENSOR_PATH
                        包含形如 model_path/0/tensor0 tensor1 的张量的文件夹路径
  --ptcsr_path PTCSR_PATH
                        包含形如 model_path/0/tensor0 tensor1 的张量的文件夹路径
  --output_path OUTPUT_PATH
                        结果 *.csv 文件的保存路径。

请确保 <tensor_path> 和可选的 <ptcsr_path> 指向由 convert_legacy_model_format.py 脚本生成的量化矩阵所在的文件夹。使用 <cuda_device_id> 设置基准测试期间的 CUDA 设备。脚本会将结果输出到 <results_output> 中。

测试

要运行单元测试，只需执行以下命令：

python3 tests/test.py

引用

@misc{dettmers2023spqr,
      title={SpQR: A Sparse-Quantized Representation for Near-Lossless LLM Weight Compression}, 
      author={Tim Dettmers and Ruslan Svirschevski and Vage Egiazarian and Denis Kuznedelev and Elias Frantar and Saleh Ashkboos and Alexander Borzunov and Torsten Hoefler and Dan Alistarh},
      year={2023},
      eprint={2306.03078},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

SpQR 快速上手指南

SpQR 是一种用于大语言模型（LLM）权重的稀疏量化表示方法，旨在实现近乎无损的模型压缩。本指南帮助中国开发者快速部署并使用该工具。

环境准备

系统要求

• GPU: 推荐单张 NVIDIA A100 (80GB)。
  • 最小显存需求：
    • 困惑度评估：32GB+ (支持 LLaMA-65B/Falcon-40B)。
    • 开启 --offload activations 后：24GB+ (LLaMA-65B) 或 6GB+ (LLaMA-7B)。
• 内存 (RAM): 需足够容纳未压缩的模型权重（例如 LLaMA-65B 约需 130GB）及测试数据集。
• 软件依赖:
  • Python 环境
  • torch >= 2.0.0 (必须支持 CUDA)
  • transformers (论文结果基于版本 4.28.dev0)

前置依赖

确保已安装 CUDA 驱动，并配置好 Hugging Face 缓存目录（可选通过环境变量自定义）。

安装步骤

1. 克隆仓库 (假设已获取源码)

2. 安装 Python 依赖

   pip install -r requirements.txt
   

   提示: 如需使用 W&B 日志记录，请额外运行 pip install wandb。

3. 构建底层组件 (运行基准测试和转换脚本前必需)

   /bin/bash scripts/build.sh
   

4. 准备数据与模型

   • 模型: 下载 Hugging Face 格式的 LLaMA、Falcon 或 OPT 模型到本地。
   • 数据集: 脚本会自动下载并缓存分词器及数据集。
     • LLaMA 推荐使用 RedPajama 子集。
     • Falcon 推荐使用 RefinedWeb 子集。
     • OPT 推荐使用 c4。
     • 注: 仓库 data 目录下已提供预处理好的子集文件，请勿混用不同模型的分词数据。

基本使用

以下示例展示如何对模型进行近无损压缩并评估其困惑度（Perplexity）。

1. 执行模型压缩与困惑度测试

设置模型路径和数据集名称，运行主脚本：

export MODEL_PATH=<PATH_TO_MODEL_DIR>
export DATASET=<INSERT DATASET NAME OR PATH TO CUSTOM DATA>

python main.py $MODEL_PATH $DATASET \
    --wbits 4 \
    --groupsize 16 \
    --perchannel \
    --qq_scale_bits 3 \
    --qq_zero_bits 3 \
    --qq_groupsize 16 \
    --outlier_threshold=0.2 \
    --permutation_order act_order \
    --percdamp 1e0 \
    --nsamples 128

参数说明:

• <PATH_TO_MODEL_DIR>: 包含 config.json 的模型文件夹路径。
• <DATASET>: 数据集名称 (c4, ptb, wikitext2, pajama, refinedweb) 或自定义预处理数据路径。
• --wbits: 量化权重位数 (如 3 或 4)。
• --offload activations: 若显存不足可添加此参数，将激活值移至内存（速度降低约 10%）。
• --save --load: 可用于保存或加载量化后的模型。

2. 模型格式转换 (优化推理)

SpQR 生成的初始张量为 int8 格式，为了使用高效的 CUDA 内核进行推理，需转换为优化存储格式：

python convert_legacy_model_format.py \
    --base_model <PATH_TO_UNQUANTIZED_MODEL> \
    --legacy_model_path <PATH_TO_LEGACY_SPQR_MODEL> \
    --sparse_strategy optimize_latency \
    --save_pt <PATH_TO_SAVE_CONVERTED_PT>

3. 端到端推理演示

使用转换后的模型进行推理测试，对比原始模型与量化模型的性能：

python inference_demo.py \
    --pretrained_model_path <PATH_TO_PRETRAINED_MODEL> \
    --compressed_model_path <PATH_TO_COMPRESSED_PT_MODEL> \
    --execution_mode 1

• --execution_mode 0: 评估原始稠密模型。
• --execution_mode 1: 评估 SpQR 量化模型。

4. (可选) 零样本任务评估

使用修改版的 LM Evaluation Harness 进行基准测试：

export MODEL_PATH=<INSERT PATH_TO_MODEL_DIR>
export DATASET=<INSERT DATASET NAME OR PATH TO CUSTOM DATA>

python lmeval.py \
    --model hf-causal \
    --model_args pretrained=$MODEL_PATH,dtype=float16,use_accelerate=True \
    --quantization_args dataset=$DATASET,wbits=4,groupsize=16,perchannel=True,qq_scale_bits=3,qq_zero_bits=3,qq_groupsize=16,percdamp=1.0,outlier_threshold=0.2,simplified_outliers=False,nsamples=128,offload_activations=True \
    --tasks winogrande,piqa,hellaswag,arc_easy,arc_challenge \
    --batch_size 1

注意：大模型（如 LLaMA-30B/65B）可能需要多卡并行，请根据显存情况调整 max_memory_per_gpu。