LLM-Shearing 开源项目教程

1. 项目介绍

LLM-Shearing 是一个由普林斯顿大学 NLP 团队开发的开源项目，旨在通过结构化剪枝技术加速大型语言模型（LLM）的预训练过程。该项目基于 MosaicML 的 Composer 包，专门优化和设计用于大规模语言模型的预训练。通过 LLM-Shearing，用户可以有效地将现有的 LLM 剪枝为更小、更高效的模型，从而显著减少计算资源和时间成本。

2. 项目快速启动

2.1 环境准备

在开始之前，请确保您的环境中已经安装了 PyTorch 和 Flash Attention。您可以通过以下命令进行安装：

pip install torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
pip install flash-attn==1.0.3

2.2 安装 LLM-Shearing

克隆项目仓库并安装所需的依赖包：

git clone https://github.com/princeton-nlp/LLM-Shearing.git
cd LLM-Shearing
pip install -r requirement.txt
pip install -e .

2.3 数据准备

参考 llmshearing/data 目录中的说明，使用 MosaicML 的 Streaming 包准备数据。

2.4 模型准备

将 Hugging Face 的模型权重转换为 Composer 兼容的格式：

HF_MODEL_NAME=meta-llama/Llama-2-7b-hf
OUTPUT_PATH=models/Llama-2-7b-composer/state_dict.pt
mkdir -p $(dirname $OUTPUT_PATH)
python3 -m llmshearing.utils.composer_to_hf save_hf_to_composer $HF_MODEL_NAME $OUTPUT_PATH

2.5 剪枝和继续预训练

使用提供的脚本进行模型剪枝和继续预训练：

# 剪枝
bash llmshearing/scripts/pruning.sh

# 继续预训练
bash llmshearing/scripts/continue_pretraining.sh

3. 应用案例和最佳实践

3.1 案例一：加速 LLaMA 模型预训练

通过 LLM-Shearing，用户可以将 LLaMA-2-7B 模型剪枝为 1.3B 和 2.7B 参数的模型，显著减少预训练所需的计算资源和时间。剪枝后的模型在下游任务和指令调优任务中表现优异，且计算成本仅为从头开始训练的 3%。

3.2 案例二：动态批量加载

LLM-Shearing 还提供了一种动态批量加载算法，可以根据不同域的损失减少率动态调整数据加载比例，从而提高数据利用效率和模型性能。

4. 典型生态项目

4.1 MosaicML Composer

LLM-Shearing 基于 MosaicML 的 Composer 包，Composer 是一个专门为大规模语言模型预训练设计的优化工具包，提供了丰富的功能和优化策略。

4.2 Hugging Face Transformers

LLM-Shearing 支持将剪枝后的模型转换为 Hugging Face 的 Transformers 格式，方便用户进行推理和微调。

4.3 RedPajama 数据集

LLM-Shearing 使用 RedPajama 数据集进行预训练，RedPajama 是一个高质量的开源数据集，适用于各种语言模型的训练。

通过以上步骤，您可以快速上手 LLM-Shearing 项目，并利用其强大的剪枝和预训练功能，加速您的语言模型开发和应用。