🦙 剪枝版 LLaMA：通过结构化剪枝加速语言模型预训练

🌟 ArXiv 预印本 | 博客文章

基础模型：Sheared-LLaMA-1.3B | Sheared-LLaMA-2.7B | Sheared-Pythia-160m
未进行持续预训练的剪枝模型：Sheared-LLaMA-1.3B-Pruned, Sheared-LLaMA-2.7B-Pruned
指令微调模型：Sheared-LLaMA-1.3B-ShareGPT | Sheared-LLaMA-2.7B-ShareGPT

感谢您对我们工作的关注！本项目由 Mengzhou Xia、Tianyu Gao、Zhiyuan Zeng 和 Danqi Chen 共同完成。在此，我们提供了 Sheared-LLaMA 的剪枝与持续预训练算法代码库 :) 我们发现，相较于从头开始预训练小型语言模型，对强大的基础模型进行剪枝是一种极具成本效益的方法。下图显示，在已有 Llama-2-7B 模型（使用 2T tokens 预训练）的情况下，对其进行剪枝后得到的模型性能可媲美 OpenLLaMA 模型，而其预训练成本仅为后者的 3%。

更新

• [2023年12月19日] 更新了仓库中的 评估脚本 和剪枝日志。
• [2023年11月22日] 我们发布了指令微调模型 Sheared-LLaMA-1.3B-ShareGPT 和 Sheared-LLaMA-2.7B-ShareGPT。
• [2023年11月19日] 我们发布了早期阶段开发的 Sheared-Pythia-160m 模型。该模型采用相同的剪枝配方和 Pile 数据集生成。
• [2023年11月5日] 我们开源了 LLM-Shearing 的代码——很高兴看到它被应用于更多不同规模的模型。
• [2023年10月10日] 我们发表了 Sheared-LLaMA 论文，并发布了两款 Sheared LLaMA 模型，同时在 Twitter 上进行了分享 🚀！

🔗 快速链接

• 🦙 剪枝版 LLaMA：通过结构化剪枝加速语言模型预训练
  • 🔗 快速链接
  • 简要介绍
  • 安装要求
  • 数据准备
  • 模型准备
  • 剪枝与持续预训练示例脚本
  • 转换剪枝模型
  • 将 Composer 模型转换为 Hugging Face 模型
  • 训练配置
    • 数据配置
    • 基础训练配置
    • 剪枝配置
    • 动态批量加载配置
  • 吞吐量
  • 未来工作
  • 遇到问题或有疑问？
  • 引用

简要介绍

本代码库基于 MosaicML 令人惊叹的 Composer 包，该包专为大规模语言模型预训练设计并进行了优化。整个实现，包括 剪枝 逻辑和 动态批量加载 逻辑，均以回调函数的形式实现，无需修改原生 Composer 训练器。以下是代码库中各文件夹的简要概述：

• shearing.data: 包含示例数据和数据处理脚本。
• shearing.datasets: 实现自定义数据集，以支持动态数据加载。
• shearing.callbacks: 实现动态加载回调函数和剪枝回调函数。
• shearing.models: 实现模型文件。
• shearing.scripts: 包含运行代码的脚本。
• shearing.utils: 包括所有实用工具函数，例如模型转换和剪枝测试。
• train.py: 运行代码的主要入口

安装要求

步骤 1：要开始使用此仓库，您需要按照以下安装步骤操作。在继续之前，请确保已安装 Pytorch 和 Flash Attention。您可以通过以下 pip 命令完成安装：

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.post

请注意，目前不支持 Flash Attention 2 版本，可能需要手动修改模型文件。

步骤 2：然后安装其余所需软件包：

cd llmshearing
pip install -r requirement.txt

步骤 3：最后，以可编辑模式安装 llmshearing 包，以便在您的开发环境中使用：

pip install -e .

数据准备

有关如何使用 Mosaicml 的 Streaming 包准备数据的详细信息，请参阅 llmshearing/data。

模型准备

要将 Hugging Face 转换器模型与 Composer 结合使用，您需要将模型权重转换为 Composer 所期望的密钥格式。以下是将 Hugging Face 模型 'llama2' 的权重转换为 Composer 兼容格式的示例：

# 定义 Hugging Face 模型名称和输出路径
HF_MODEL_NAME=meta-llama/Llama-2-7b-hf
OUTPUT_PATH=models/Llama-2-7b-composer/state_dict.pt

# 如果目录不存在，则创建
mkdir -p $(dirname $OUTPUT_PATH)

# 将 Hugging Face 模型转换为 Composer 密钥格式
python3 -m llmshearing.utils.composer_to_hf save_hf_to_composer $HF_MODEL_NAME $OUTPUT_PATH

此外，您可以使用以下实用函数来测试 Hugging Face 模型与转换后的 Composer 模型之间的等效性：

MODEL_SIZE=7B
python3 -m llmshearing.utils.test_composer_hf_eq $HF_MODEL_NAME $OUTPUT_PATH $MODEL_SIZE

这些函数仅适用于 LLaMA/LLaMA2 模型。不过，将其适配到其他模型（如 Mistral-7B）应该并不困难。

剪枝与继续预训练的示例脚本

对于剪枝，您可以参考位于 llmshearing/scripts/pruning.sh 的示例脚本。在该脚本中，您需要进行调整以纳入 数据配置、基础训练配置、剪枝配置 和 动态批量加载配置。

由于剪枝相比继续预训练的计算成本相对较高，我们会在特定步数后（通常在所有实验中为 3200 步）停止使用剪枝目标的训练。随后，我们将对剪枝后的模型进行进一步的预训练。为确保兼容性，有必要将模型的状态字典键转换为符合标准目标模型结构的形式。有关此转换的详细说明，请参阅 转换剪枝模型。

完成模型转换后，您可以继续对剪枝模型进行预训练。这一过程与预训练标准模型类似。为此，您可以参考位于 llmshearing/scripts/continue_pretraining.sh 的示例脚本。在此脚本中，剪枝配置已被移除。

训练完成后，您可以使用转换脚本将 Composer 模型转换为 Transformers 模型。更多详情请参阅 将 Composer 模型转换为 Hugging Face 模型 部分。

转换剪枝模型

在使用 llmshearing/scripts/pruning.sh 完成训练后，保存的模型包含源模型的全部参数，并附带一组掩码。接下来，我们将对掩码变量采取以下操作：1) 移除掩码变量接近于 0 的子结构；2) 通过矩阵-向量乘法将掩码变量合并到模型参数中，从而得到一个更为紧凑的模型。同时，还需要重命名权重键，以便能够无缝加载到目标模型架构中，确保各层名称连续一致。

MODEL_PATH=$MODEL_DIR/latest-rank0.pt
python3 -m llmshearing.utils.post_pruning_processing prune_and_save_model $MODEL_PATH

剪枝后的模型将保存在 $(dirname $MODEL_PATH)/pruned-latest-rank0.pt 中。

将 Composer 模型转换为 Hugging Face 模型

训练完成后，如果您希望使用 Hugging Face 进行推理或微调，可以使用 llmshearing/scripts/composer_to_hf.py 脚本将您的 Composer 模型转换为 Hugging Face 模型。以下是该脚本的使用示例：

MODEL_PATH=$MODEL_DIR/latest-rank0.pt
OUTPUT_PATH=$MODEL_DIR/hf-latest_rank0
MODEL_CLASS=LlamaForCausalLM
HIDDEN_SIZE=2048
NUM_ATTENTION_HEADS=16
NUM_HIDDEN_LAYERS=24
INTERMEDIATE_SIZE=5504
MODEL_NAME=Sheared-Llama-1.3B

python3 -m llmshearing.utils.composer_to_hf save_composer_to_hf $MODEL_PATH $OUTPUT_PATH \
        model_class=${MODEL_CLASS} \
        hidden_size=${HIDDEN_SIZE} \
        num_attention_heads=${NUM_ATTENTION_HEADS} \
        num_hidden_layers=${NUM_HIDDEN_LAYERS} \
        intermediate_size=${INTERMEDIATE_SIZE} \
        num_key_value_heads=${NUM_ATTENTION_HEADS} \
        _name_or_path=${MODEL_NAME}

请注意，此处提到的参数名称是针对 Llama2 的 Hugging Face 配置设计的，在处理其他类型模型时可能会有所不同。

训练配置

本节将深入介绍如何在 YAML 配置文件中设置训练参数。这些配置涵盖多个关键方面，包括数据设置、基础训练设置、剪枝设置以及动态数据加载配置。

数据配置

• data_local: 包含数据的本地目录。
• eval_loader.dataset.split: 对于评估，需提供一个包含所有领域数据的组合分割名称。
• train_loader.dataset.split: 当动态加载配置中的 dynamic=True（请参阅 动态加载部分）时，无需设置此值。然而，若 dynamic=False，则必须指定一个训练分割。

基础训练配置

基础训练配置主要遵循原始的 Composer 包。有关这些配置的详细信息，请参阅 Composer 官方文档。以下是一些需要注意的关键训练参数：

• max_duration: 此参数定义了最大训练时长，可按步数（例如 3200ba）或轮次（例如 1ep）指定。在我们的实验中，剪枝阶段的训练时长设为 3200ba，而继续预训练阶段则设为 48000ba。
• save_interval: 此参数决定了模型状态保存的频率。我们在剪枝和继续预训练阶段均将其设置为 3200ba。
• t_warmup: 此参数指定了学习率调度器的学习率预热时长。在剪枝阶段，其值设为 320ba（占总训练时间的 10%），而在继续预训练阶段则设为 1440ba（占总训练时间的 3%）。
• optimizer.lr: 此参数定义了主模型参数的学习率，默认值为 1e-4。
• max_seq_len: 按照 Llama 2 的训练方法，我们支持的最大序列长度为 4096。
• device_train_microbatch_size: 此参数决定了每台设备在训练过程中的批次大小。在剪枝阶段，我们将其设置为 4，而在继续预训练阶段则提高至 16。
• global_train_batch_size: 此参数指定了训练过程中所有 GPU 上的全局批次大小。在剪枝阶段，其值设为 32，而在继续预训练阶段则增加至 256。
• autoresume: 此参数可通过将其设置为 true 来启用断点续跑功能。然而，需要注意的是，尽管我们在继续预训练阶段成功使用过此功能，但尚不能保证其在剪枝阶段的兼容性。

由于计算资源的限制，我们并未进行全面的超参数搜索，因此可能存在更优的超参数以提升性能。

剪枝配置

剪枝过程允许将源模型裁剪为特定的目标结构，脚本中包含一些关键参数，例如：

• from_model：该参数指定源模型的大小，并与一个配置文件相对应。
• to_model：该参数定义目标模型的大小，源模型将被剪枝以匹配目标配置。
• optimizer.lag_lr：该参数指定在剪枝过程中用于学习掩码变量和拉格朗日乘子的学习率。默认值为 $1.0$。

剪枝相关的参数都归类在 model.l0_module 下：

• model.l0_module.lagrangian_warmup_steps： 在初始预热阶段，剪枝率会从 0 逐步增加到期望的目标值。具体的目标值由目标模型的预定义结构决定。需要注意的是，这个值可能与学习率的预热步数不同。通常，我们会将总步数的约 20% 用于剪枝的预热过程。
• model.l0_module.pruning_modules：默认情况下，此设置会剪枝模型的各个部分，包括头部、中间维度、隐藏维度和层。
• model.l0_module.eval_target_model：当设置为 true 时，评估过程会针对与目标模型结构完全一致的子模型进行评估。如果设置为 false，则评估当前模型，并考虑掩码值。由于掩码可能需要一段时间才能收敛为目标模型形状，因此在训练期间我们基于当前模型形状而非目标结构进行评估。
• model.l0_module.target_model.d_model：指定目标模型的隐藏维度。
• model.l0_module.target_model.n_heads：指定目标模型的注意力头数。
• model.l0_module.target_model.n_layers：指定目标模型的层数。
• model.l0_module.target_model.intermediate_size：指定目标模型的中间维度数量。

这些参数使您可以根据具体需求配置和控制剪枝过程。

动态批量加载配置

我们在 datasets/streaming_dataset.py 中扩展了 Steaming 的 StreamingDataset，以支持动态加载数据。配置动态批量加载的参数主要定义在 DynamicLoadingCallback 中。以下大多数配置都可以在 YAML 配置文件的 callbacks.data_loading 部分中指定。以下是每个参数的说明：

• callbacks.data_loading.dynamic：此布尔参数决定是否启用动态数据加载。当设置为 true 时，数据会从不同的领域或流中动态加载。如果设置为 false，则禁用动态数据加载。
• callbacks.data_loading.set_names：指定用于动态数据加载的领域名称或流名称。
• callbacks.data_loading.proportion：该参数定义每个领域或流的初始数据加载比例。所有比例之和必须等于 1，表示初始数据加载配置中各来源的相对权重。
• callbacks.data_loading.update_type：选择在训练过程中调整数据加载比例的更新方式。有两种选项：
  • doremi：在此模式下，数据加载比例会采用指数下降的方式进行更新，类似于 Doremi 中描述的方法。这使得数据加载比例能够随时间自适应地调整。
  • constant：选择此选项后，数据加载比例在整个训练过程中保持不变。相当于禁用了动态数据加载。
• callbacks.data_loading.target_loss：指定训练过程中的目标验证损失。该目标损失值应在训练开始前计算或预先确定。加载比例会根据模型当前损失与目标损失之间的差异进行动态调整，从而引导训练过程达到预期的性能水平。
• eval_interval：确定训练过程中评估的频率。如果 dynamic=True，则每次评估后都会调整数据加载比例。

该代码设计仅适用于本地数据，不支持远程流式数据。此外，目前它仅支持单个工作线程的数据加载器，且不提供预取功能。在我们的测试中，这一限制并未带来额外的计算开销。

吞吐量

以下是使用 A100 80GB GPU 运行剪枝和继续预训练步骤的吞吐量。吞吐量以每秒处理的标记数来衡量。请参考 llm-foundry 的标准吞吐量。

	GPU 数量	每设备吞吐量	总吞吐量
剪枝 7B	8	1844	14750
预训练 3B	16	4957	79306
预训练 1.3B	16	8684	138945

未来工作

源模型：尽管大型模型无疑功能强大，并且在不久的将来还有进一步增强的潜力，但我们认为小规模模型（参数量少于70亿）仍具有未被充分挖掘的潜力。然而，目前鲜有研究致力于提升小模型的能力，而我们的工作正是朝着这一目标迈进。这项工作的自然延伸是将代码库扩展到以下方面：

• 更强大的基础模型，例如 Mistral-7B；
• 领域特定的语言模型，如代码生成模型，包括 CodeLlama 和 DeepSeek-Coder；
• 不同规模的模型。由于计算资源的限制，我们主要使用了70亿参数的模型进行实验。尚不清楚从更大规模的模型中进行剪枝是否会有更大的收益。

为了使代码库适配其他模型，一个关键步骤是确保使用掩码运行模型与运行剪枝后的模型效果等价。我们使用 llmshearing/utils/test_pruning.py 来执行此类测试，以验证模型文件中 prune_params 函数的正确性。

数据来源：请务必注意，最终模型的性能不仅取决于剪枝算法和基础模型，还取决于数据的质量。在我们的实验中，主要使用了 RedPajama v1 数据集。不过，以下是一些可以考虑纳入的额外资源：

• Dolma 数据集，这是一个3TB的预训练数据集，涵盖了 CommonCrawl、C4、peS2o、The Stack、Project Gutenberg 和 Wikipedia 等多个领域；
• proof-pile-2 数据集，包含550亿个标记的数学和科学文献；
• RedPajama-v2 数据集，一个30TB的预训练数据集。

发现错误或有疑问？

如果您对代码或论文有任何疑问，请随时发送邮件至 Mengzhou (mengzhou@princeton.edu)。如果在使用代码时遇到任何问题，或希望报告错误，您可以提交一个问题。请尽可能详细地描述问题，以便我们能够更快更好地帮助您！

引用

如果您在工作中发现本仓库有所帮助，请引用我们的论文：

@article{xia2023sheared,
  title={Sheared llama: 通过结构化剪枝加速语言模型预训练},
  author={Xia, Mengzhou and Gao, Tianyu and Zeng, Zhiyuan and Chen, Danqi},
  journal={arXiv 预印本 arXiv:2310.06694},
  year={2023}
}

LLM-Shearing 快速上手指南

LLM-Shearing 是一个通过结构化剪枝（Structured Pruning）加速大语言模型预训练的工具。它能够将大型基座模型（如 LLaMA-2-7B）剪枝为小型模型，并通过少量继续预训练达到与从头训练相当的效果，显著降低计算成本。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
• Python: 3.8+
• GPU: 支持 CUDA 的 NVIDIA GPU
• 前置依赖:
  • PyTorch (需指定版本以兼容 Flash Attention)
  • Flash Attention (注意：目前仅支持 v1 版本，v2 版本需要手动修改代码)

安装步骤

请按照以下顺序执行命令来配置环境。

1. 安装 PyTorch 和 Flash Attention

首先安装特定版本的 PyTorch 和 Flash Attention v1。以下命令基于 CUDA 11.8，请根据您的实际 CUDA 版本调整。

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.post

提示：国内用户若下载缓慢，可尝试使用清华或阿里镜像源替换 --index-url 或直接使用 pip install -i https://pypi.tuna.tsinghua.edu.cn/simple ...（需注意兼容性）。

2. 克隆仓库并安装依赖

克隆项目代码并安装 requirements.txt 中的其他依赖包。

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

3. 安装本地开发包

以可编辑模式安装 llmshearing 包，以便在开发环境中调用。

pip install -e .

基本使用

本工具的核心流程分为三个阶段：模型格式转换 -> 剪枝与继续预训练 -> 模型导出。

1. 模型准备（格式转换）

LLM-Shearing 基于 MosaicML Composer 构建，使用前需将 Hugging Face 格式的模型权重转换为 Composer 格式。以下以 Llama-2-7b 为例：

# 定义变量
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

(可选) 验证转换后的模型与原模型等价性：

MODEL_SIZE=7B
python3 -m llmshearing.utils.test_composer_hf_eq $HF_MODEL_NAME $OUTPUT_PATH $MODEL_SIZE

2. 执行剪枝与继续预训练

项目提供了示例脚本用于执行剪枝和后续的继续预训练。您需要根据实际需求修改脚本中的配置文件（数据路径、剪枝参数等）。

• 剪枝阶段：运行 llmshearing/scripts/pruning.sh。
  • 该阶段通常会运行固定步数（例如 3200 步），目的是移除冗余结构。
• 继续预训练阶段：剪枝完成后，运行 llmshearing/scripts/continue_pretraining.sh 对剪枝后的模型进行进一步训练以恢复性能。

注意：请务必在脚本中正确配置 data_local（数据本地路径）及相关训练超参数。数据准备需参考 MosaicML Streaming 格式。

3. 处理剪枝后的模型

剪枝训练结束后，保存的模型包含原始参数和掩码（masks）。需要运行后处理脚本将掩码应用到位，移除零值子结构，并重整权重键名以适配标准架构。

MODEL_PATH=$MODEL_DIR/latest-rank0.pt
python3 -m llmshearing.utils.post_pruning_processing prune_and_save_model $MODEL_PATH

处理后的模型将保存为 pruned-latest-rank0.pt。

4. 导出为 Hugging Face 格式

为了方便推理或微调，最后将 Composer 格式的模型转换回 Hugging Face Transformers 格式。

MODEL_PATH=$MODEL_DIR/latest-rank0.pt
OUTPUT_PATH=$MODEL_DIR/hf-latest_rank0
MODEL_CLASS=LlamaForCausalLM
HIDDEN_SIZE=2048
NUM_ATTENTION_HEADS=16
NUM_HIDDEN_LAYERS=24
INTERMEDIATE_SIZE=5504
MODEL_NAME=Sheared-Llama-1.3B

python3 -m llmshearing.utils.composer_to_hf save_composer_to_hf $MODEL_PATH $OUTPUT_PATH \
        model_class=${MODEL_CLASS} \
        hidden_size=${HIDDEN_SIZE} \
        num_attention_heads=${NUM_ATTENTION_HEADS} \
        num_hidden_layers=${NUM_HIDDEN_LAYERS} \
        intermediate_size=${INTERMEDIATE_SIZE} \
        num_key_value_heads=${NUM_ATTENTION_HEADS} \
        _name_or_path=${MODEL_NAME}

转换完成后，即可在 $OUTPUT_PATH 目录下使用标准的 Hugging Face transformers 库加载模型进行推理。