通过推测解码实现 Transformer 的快速推理

本仓库实现了用于大型语言模型（LLM）解码的推测采样方法。在解码过程中，该方法会同时使用两个模型：目标模型和近似模型。其中，近似模型通常较小，而目标模型则较大。近似模型负责生成候选 token，目标模型则对这些候选进行校验和修正。通过让目标模型并行处理近似模型的输出，这种方法能够显著提升解码效率，相比仅使用目标模型单独解码更为高效。

推测采样由 Google 和 DeepMind 分别独立提出。因此，我实现了两种略有不同的推测采样版本：Google 的 和 DeepMind 的。

更新日志

• 2023年9月21日：添加服务端功能，支持更多模型，如 llama-7B 和 llama-1B。
• 2023年9月19日：为 Google 版本添加 KV 缓存优化。
• 2023年8月16日：首次发布，实现论文中的算法，支持 Bloom-560M 和 Bloomz-7B1。

使用方法

推理

你需要准备一对具有相同嵌入层和词汇表的模型，其中近似模型应小于目标模型。以下是一些经过测试的模型组合。

在示例中，我们以 bloomz-7b1 作为目标模型，bloom-560m 作为近似模型。

python main.py \
    --input "The quick brown fox jumps over the lazy " \
    --target_model_name bigscience/bloomz-7b1 \
    --approx_model_name bigscience/bloom-560m

你还可以使用 -v 参数查看每个 token 是由哪个模型生成的。

我建议将 llama2-7B 和 llama2-70B 分别用作近似模型和目标模型。在这种情况下，我确实观察到了速度提升，如下所示。需要注意的是，近似模型和目标模型的选择对加速效果至关重要。如果两个模型都较小，则由于速度差异不明显，可能无法观察到加速；而如果模型规模差距过大，则会导致更多的拒绝和重采样操作。此外，当前的采样逻辑还不够高效，我发现 Softmax 和 LayerNorm 操作带来了较大的开销，未来我会尝试进一步优化。欢迎提出性能改进的建议。

	llama2-7b	llama2-70b	推测采样
速度（tokens/秒）	1084.86	329.83	427.02

服务端

启动推理服务器：

python serving.py

使用 curl 测试服务端：

curl -X POST -H "Content-Type: application/json" -d '{"prompt": "Who is the president of the USA"}' http://127.0.0.1:5000/predict

参考文献

@inproceedings{leviathan2023fast,
  title={Fast inference from transformers via speculative decoding},
  author={Leviathan, Yaniv and Kalman, Matan and Matias, Yossi},
  booktitle={International Conference on Machine Learning},
  pages={19274--19286},
  year={2023},
  organization={PMLR}
}

@article{chen2023accelerating,
  title={Accelerating large language model decoding with speculative sampling},
  author={Chen, Charlie and Borgeaud, Sebastian and Irving, Geoffrey and Lespiau, Jean-Baptiste and Sifre, Laurent and Jumper, John},
  journal={arXiv preprint arXiv:2302.01318},
  year={2023}
}

局限性

目前，我仅支持批量大小为 1 的请求。由于本仓库主要用于演示目的，其他重要的优化措施，例如批处理和并行化，并未包含在内，而这些措施对于提升效率至关重要。

LLMSpeculativeSampling 快速上手指南

LLMSpeculativeSampling 是一个通过“推测解码”（Speculative Decoding）技术加速大语言模型（LLM）推理的开源工具。它利用一个小模型（近似模型）生成候选 token，再由一个大模型（目标模型）进行并行验证和修正，从而在保持输出质量的同时显著提升推理速度。

环境准备

• 系统要求：Linux 或 macOS（推荐 Linux），具备 CUDA 支持的 NVIDIA GPU（用于加速推理）。
• 前置依赖：
  • Python 3.8+
  • PyTorch (支持 CUDA)
  • Transformers 库
  • Hugging Face Hub 访问权限（需能下载模型权重）
• 模型要求：需要准备一对词表（Vocabulary）和嵌入层（Embedding）完全一致的模型。
  • 目标模型 (Target Model)：较大的模型（如 Llama-2-70B, Bloomz-7B1）。
  • 近似模型 (Approx Model)：较小的模型（如 Llama-2-7B, Bloom-560M）。
  • 注意：如果两个模型规模差异过大或过小，可能无法观察到明显的加速效果。

国内加速建议：若下载 Hugging Face 模型受阻，建议配置镜像源或使用国内加速工具（如 hf-mirror 或在代码中设置 HF_ENDPOINT 环境变量）。

export HF_ENDPOINT=https://hf-mirror.com

安装步骤

克隆仓库并安装必要的 Python 依赖：

git clone https://github.com/hemingkx/LLMSpeculativeSampling.git
cd LLMSpeculativeSampling
pip install -r requirements.txt

(注：若仓库根目录无 requirements.txt，请手动安装核心依赖：pip install torch transformers accelerate)

基本使用

1. 命令行推理示例

以下示例演示如何使用 bloomz-7b1 作为目标模型，bloom-560m 作为近似模型进行文本生成。

python main.py \
    --input "The quick brown fox jumps over the lazy " \
    --target_model_name bigscience/bloomz-7b1 \
    --approx_model_name bigscience/bloom-560m

• 查看详细信息：添加 -v 参数可以查看每个 token 是由哪个模型生成的（有助于调试和分析加速效果）。

  python main.py -v --input "..." --target_model_name ... --approx_model_name ...
  

• 推荐组合：根据作者测试，使用 llama2-7B (近似) 和 llama2-70B (目标) 能获得较好的加速比。

2. 启动推理服务 (Serving)

你可以启动一个本地 HTTP 服务来调用该功能。

启动服务：

python serving.py

发送请求测试：

curl -X POST -H "Content-Type: application/json" -d '{"prompt": "Who is the president of the USA"}' http://127.0.0.1:5000/predict

当前限制：目前版本仅支持 batch size 为 1 的请求，尚未包含批处理（batching）等高级优化。