SparseGPT

本仓库包含了复现论文SparseGPT: 大型语言模型可以在一次性操作中被精确剪枝中关键结果的代码。

具体来说，它提供了以下脚本和实现：

在原始WikiText2、PTB和C4子集上评估基线和剪枝后的模型。（datautils.py、opt.py、bloom.py）
对OPT和BLOOM模型执行非结构化、n:m和稀疏+量化的SparseGPT压缩。（sparsegpt.py、opt.py、bloom.py）

我们注意到，这个SparseGPT实现是基于我们的开源GPTQ代码。

依赖

torch：在v1.10.1+cu111上测试通过
transformers：在v4.21.2上测试通过
datasets：在v1.17.0上测试通过

使用方法

以下是在OPT模型上运行基线和稀疏化，然后在原始WikiText2、PTB和C4上进行困惑度评估的一些示例命令。另请参阅CMD参数文档。

# 运行密集基线
python opt.py facebook/opt-125m c4

# 运行幅度基线
python opt.py facebook/opt-125m c4 --sparsity .5 --gmp

# 使用SparseGPT剪枝至50%均匀稀疏度
python opt.py facebook/opt-125m c4 --sparsity .5

# 使用SparseGPT剪枝至完整2:4稀疏度
python opt.py facebook/opt-125m c4 --prunen 2 --prunem 4

# 使用SparseGPT剪枝至50% + 4位
python opt.py facebook/opt-125m c4 --sparsity .5 --wbits 4

要在其他OPT模型上运行，请将"facebook/opt-125m"替换为相应模型的HuggingFace名称。对于175B模型，必须首先向Meta请求访问权限并将检查点转换为HuggingFace格式，然后可以将其位置作为名称传递给此脚本。

BLOOM脚本bloom.py的接口非常相似，但目前有些功能仅适用于OPT，例如：

# 使用SparseGPT对BLOOM-176B进行稀疏化
python bloom.py bigscience/bloom c4 --sparsity .5

我们还提供了具有相同接口的LLaMA剪枝脚本：

# 使用SparseGPT对LLaMa进行稀疏化
python llama.py LLAMA_HF_WEIGHTS_LOCATION c4 --sparsity 0.5

如果想保存稀疏化后的模型，请通过--save标志指定保存检查点的路径。

可以选择使用--log_wandb将评估结果记录到W&B。

演示

可以通过colab演示demo.ipynb尝试SparseGPT。

引用

如果您觉得这项工作有用，请考虑引用：

@article{frantar-sparsegpt,
  title={{SparseGPT}: Massive Language Models Can Be Accurately Pruned in One-Shot}, 
  author={Elias Frantar and Dan Alistarh},
  year={2023},
  journal={arXiv preprint arXiv:2301.00774}
}