REST: 基于检索的推测解码

如果训练让你焦头烂额，休息一下，速度翻倍。

新闻

🎉 2024-3-14: REST 被 NAACL 2024 接收！

简介

REST 是一种基于检索的推测解码方法，旨在提高大语言模型的生成速度。与依赖草稿语言模型的推测解码不同，REST 利用数据存储来检索和使用草稿标记。此外，REST 与分块并行解码和 Medusa 不同，它不需要额外的训练步骤。它作为一种即插即用的解决方案，能够加速任何现有的语言模型。

<div align="center"> <picture> <img src="https://yellow-cdn.veclightyear.com/835a84d5/f8873971-45a8-4c2f-8f0b-198bcf66f82d.png" width="85%"> </picture> <br> <div align="left" width="80%"> <em>REST 概述。在推理过程中，输入上下文用作查询，从数据存储中检索与输入最长后缀匹配的文档。使用检索到的文档的延续构建 Trie，并对低频分支进行修剪。修剪后子树中的候选项将进一步输入到 LLM 中，并使用树注意力掩码进行验证。从开始的所有正确标记将被接受，第一个错误之后的草稿标记将被拒绝。</em> </div> <br> </div> <div align="center"> <picture> <img src="https://yellow-cdn.veclightyear.com/835a84d5/1599ab88-65e1-4721-9f88-b0b11720ab54.png" width="85%"> </picture> <br> <div align="left" width="40%"> <em>HumanEval 和 MT-Bench 上使用标准自回归生成和 REST 的速度。在 HumanEval 中，核采样的温度设置为 0.8，top-p 设置为 0.95。对于 MT-Bench，温度设置为 0.7，top-p 设置为 0.8。所有实验都在单个 NVIDIA A6000 GPU 和 96 个 CPU 核心上进行，批量大小为 1。</em> </div> <br> </div>

安装

conda create -n rest python=3.9
conda activate rest
pip3 install -r requirements.txt # 注意 Pytorch CUDA 版本
pip3 install DraftRetriever/wheels/draftretriever-0.1.0-cp39-cp39-manylinux_2_34_x86_64.whl

构建数据存储

构建小型数据存储

使用 ShareGPT 的数据在 10 分钟内构建聊天数据存储（需要 465MB 磁盘存储）

cd datastore
python3 get_datastore_chat.py --model-path lmsys/vicuna-7b-v1.5 # 在此文件夹中获取 datastore_chat_small.idx

从 The Stack 在 20 分钟内构建 Python 代码生成数据存储（需要 924MB 磁盘存储）

cd datastore
python3 get_datastore_code.py --model-path codellama/CodeLlama-7b-instruct-hf # 在此文件夹中获取 datastore_stack_small.idx

构建大型数据存储

（可选）使用 UltraChat 的数据构建聊天数据存储（需要 12GB 磁盘存储）

cd datastore
python3 get_datastore_chat.py --model-path lmsys/vicuna-7b-v1.5 --large-datastore True # 在此文件夹中获取 datastore_chat_large.idx

（可选）从 The Stack 构建 Python 代码生成数据存储（需要 27GB 磁盘存储）

cd datastore
python3 get_datastore_code.py --model-path codellama/CodeLlama-7b-instruct-hf --large-datastore True # 在此文件夹中获取 datastore_stack_large.idx

推理

在 MT-Bench 上推理

cd llm_judge
RAYON_NUM_THREADS=6 CUDA_VISIBLE_DEVICES=0 python3 gen_model_answer_rest.py --model-path lmsys/vicuna-7b-v1.5 --model-id vicuna-7b-v1.5 --datastore-path ../datastore/datastore_chat_small.idx

在 HumanEval 上推理

cd human_eval
RAYON_NUM_THREADS=6 CUDA_VISIBLE_DEVICES=0 python3 rest_test.py --model-path codellama/CodeLlama-7b-instruct-hf --datastore-path ../datastore/datastore_stack_small.idx

自由聊天

RAYON_NUM_THREADS=6 CUDA_VISIBLE_DEVICES=0 python3 -m rest.inference.cli --datastore-path datastore/datastore_chat_small.idx --base-model lmsys/vicuna-7b-v1.5

请注意，RAYON_NUM_THREADS 环境变量控制检索的最大线程数。您可以根据机器情况进行调整。

其他模型和数据存储

在上面的示例中，我们默认使用 Vicuna 和 CodeLlama。但实际上，您可以通过简单地更改 "--model-path" 参数来使用任何您喜欢的基于 LLaMA 的模型。您还可以从任何您喜欢的数据构建数据存储。如果您想使用 LLaMA 以外的架构，您也可以修改文件 model/modeling_llama_kv.py 以匹配相应的模型。

引用

@misc{he2023rest,
      title={REST: Retrieval-Based Speculative Decoding}, 
      author={Zhenyu He and Zexuan Zhong and Tianle Cai and Jason D Lee and Di He},
      year={2023},
      eprint={2311.08252},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}