跳层

此代码库是 LayerSkip 的实现：启用提前退出推理和自推测解码。

入门

克隆仓库：

$ git clone [email protected]:facebookresearch/LayerSkip.git
$ cd LayerSkip

设置环境：

$ conda create --name layer_skip python=3.10
$ conda activate layer_skip

$ pip install -r requirements.txt

访问模型：为了观察加速情况，您需要访问使用 LayerSkip 配方训练的 LLM。我们在使用 LayerSkip 配方持续预训练的不同 Llama 模型的 HuggingFace 上提供了 6 个检查点：
- facebook/layerskip-llama2-7B
- facebook/layerskip-llama2-13B
- facebook/layerskip-codellama-7B
- facebook/layerskip-codellama-34B
- facebook/layerskip-llama3-8B
- facebook/layerskip-llama3.2-1B

为了访问每个模型：

访问上面模特对应的链接，确保您已使用您的帐户登录 HuggingFace 网站。
填写申请表并提交。批准可能需要一段时间，您应该会收到一封电子邮件通知，通知您模型的权限已被授予。
请按照此处的步骤获取用户访问令牌。
在命令行中运行huggingface-cli login ，系统将提示您提供在步骤 3 中获取的令牌。

运行这些步骤后，下面运行 LayerSkip 检查点的命令应该可以工作。

产生

要使用常规自回归解码在交互模式下运行我们的模型之一：

$ torchrun generate.py --model facebook/layerskip-llama2-7B 
    --sample True 
    --max_steps 512

为了观察加速，您需要使用自推测解码来生成令牌，并指定--exit_layer ，草稿阶段退出的层，以及--num_speculations ，草稿令牌的数量：

$ torchrun generate.py --model facebook/layerskip-llama2-7B 
    --sample True 
    --max_steps 512 
    --generation_strategy self_speculative 
    --exit_layer 8 
    --num_speculations 6

尖端：

您可以将--model更改为任何 HuggingFace 模型，但为了观察自我推测解码的加速情况，请使用使用 LayerSkip 配方训练的模型，例如我们在 HuggingFace 上开源的模型。
默认情况下我们启用采样。您可以使用--sample 、 --temperature 、 --top_p和--top_k参数更改采样行为。
您可以运行python generate.py --help以获取有关不同命令行参数的详细信息。

基准

对数据集进行基准测试：

$ torchrun benchmark.py --model facebook/layerskip-llama2-7B 
    --dataset cnn_dm_summarization 
    --num_samples 100 
    --generation_strategy self_speculative 
    --exit_layer 8 
    --num_speculations 6 
    --output_dir ./logs

尖端：

您可以通过修改--dataset参数来指定不同的任务：
- cnn_dm_summarization ：CNN/DM 总结
- xsum_summarization : XSUM 总结
- cnn_dm_lm ：CNN/DM 语言建模（给定文章的前几个单词，生成剩余的文章）
- human_eval : HumanEval 编码
默认情况下，任务以 0-shot 方式运行。您可以通过指定--n_shot参数更改为任何指定的n -shot。
默认情况下我们启用采样，而论文中报告的结果是不采样的贪婪解码。您可以使用--sample 、 --temperature 、 --top_p和--top_k参数更改采样行为。
您可以运行python benchmark.py --help以获取有关不同命令行参数的详细信息。

评价

我们已将生成脚本与 Eleuther 语言模型评估工具集成，以支持大量任务并正确后处理生成的文本。

$ torchrun eval.py --model facebook/layerskip-llama2-7B 
    --tasks gsm8k 
    --limit 10 
    --generation_strategy self_speculative 
    --exit_layer 8 
    --num_speculations 6 
    --output_dir ./logs

尖端：

请注意，通过推测性解码，我们只能从生成任务（例如gsm8k或cnn_dailymail ）中获得加速，而分类任务，即多项选择题任务（例如piqa 、 social_iqa ）或真/假问题任务（例如boolq ）将获得加速。不会导致加速。
您可以使用--tasks参数指定 Eleuther Evaluation Harness 支持的任意数量的任务。要获取所有可能任务的列表，请检查此链接。
与generate.py和benchmark.py脚本类似，您可以指定不同的模型、数据集和采样参数
您可以运行python benchmark.py --help以获取有关不同命令行参数的详细信息。

扫

我们的推理超参数exit_layer和num_speculations决定推理过程中的加速：

exit_layer ：
- 较小意味着选秀阶段速度更快但准确性较低
- 更大意味着更准确但更慢的选秀阶段
num_speculations ：
- 越小意味着接受率越高，但验证阶段的摊销将减少草稿阶段的摊销
- learger 意味着验证阶段将更好地摊销草稿阶段，但接受率会下降

exit_layer和num_speculations的最佳组合可能会随着模型、数据集和采样参数的变化而变化。因此，我们提供了一个脚本来扫描不同exit_layer和num_speculations的网格：

$ torchrun sweep.py --model facebook/layerskip-llama2-7B 
    --dataset human_eval 
    --generation_strategy self_speculative 
    --num_samples 150 
    --max_steps 256 
    --output_dir ./logs/ 
    --sample False

这将在--outpu_dir参数指定的目录中创建一个 CSV 文件。

尖端：

与generate.py和benchmark.py脚本类似，您可以指定不同的模型、数据集和采样参数
您可以运行python sweep.py --help以获取有关不同命令行参数的详细信息。

正确性

为了验证自推测解码算法生成的标记是否正确，我们创建了一个脚本来比较自回归解码与自推测解码的输出。请注意，我们只能在没有采样时保证输出的等效性（即--sample False ）：

$ torchrun correctness.py --model facebook/layerskip-llama2-7B 
    --dataset human_eval 
    --generation_strategy self_speculative 
    --num_speculations 6 
    --exit_layer 4 
    --num_samples 10 
    --sample False 
    --output_dir ./logs

使用 Docker

Kindy 检查 DOCKER.md 以使用 docker 设置项目

其他实现

我们还有 LayerSkip 推理的其他实现：

gpt-fast：gpt-fast 是一个简单高效的 pytorch 原生转换器文本生成。我们在 gpt-fast 代码库中实现了 LayerSkip，以便能够将其与其他优化（例如torch.compile() 、量化和张量并行性）进行组合。
原生 HuggingFace：在每个 HuggingFace 模型的模型卡中，我们提供了简单的代码片段，这些代码片段利用 HuggingFace 推测解码功能，使用简单的技巧来克隆主模型的早期层，而不克隆其权重。虽然这种实现很简单，不需要实现其他功能或导入其他库，但它在草案和验证阶段之间不共享KV缓存或执行。

训练

我们的培训实施正在进行中。您可以查看此拉取请求以获取详细信息和讨论。

执照

LayerSkip 根据 CC-by-NC 许可证获得许可。请参阅顶级目录中的 LICENSE 文件。

贡献

我们欢迎对 LayerSkip 做出贡献。如果您有兴趣做出贡献，请参阅此文档。

引文

如果您在研究中使用 LayerSkip，请使用以下 BibTex 条目：

 @misc { layerskip ,
    title = { LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding } ,
    author = { Mostafa Elhoushi and Akshat Shrivastava and Diana Liskovich and Basil Hosmer and Bram Wasti and Liangzhen Lai and Anas Mahmoud and Bilge Acun and Saurabh Agarwal and Ahmed Roman and Ahmed A Aly and Beidi Chen and Carole-Jean Wu } ,
    booktitle = " Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) " ,
    month = aug,
    year = " 2024 " ,
    address = " Bangkok, Thailand " ,
    publisher = " Association for Computational Linguistics " ,
    url = " https://aclanthology.org/2024.acl-long.681 " ,
    doi = " 10.18653/v1/2024.acl-long.681 " ,
    pages = " 12622--12642 " ,
}