如何通过自推测解码加速大模型推理实现高效推理?

摘要:推测解码是一种新颖的文本生成方法,它结合了推测解码 (Speculative Decoding) 的优势和大语言模型 (LLM) 的提前退出 (Early Exit) 机制该方法出自论文 LayerSkip: Enabling Early-
推测解码是一种新颖的文本生成方法,它结合了推测解码 (Speculative Decoding) 的优势和大语言模型 (LLM) 的提前退出 (Early Exit) 机制该方法出自论文 LayerSkip: Enabling Early-Exit Inference and Self-Speculative Decoding它通过使用 同一个模型 的早期层来生成候选词元 (token),并使用后期层进行验证,从而实现高效生成 这项技术不仅加快了文本生成速度,还显著节省了内存并降低了计算延迟为了实现端到端的加速,早期层的输出需要与最终层的输出足够接近正如论文中所述,这可以通过一种训练方法来实现,该方法可以在预训练期间应用,也可以在特定领域进行微调时应用自推测解码对于实际应用特别高效,它可以在较小的 GPU 上部署,并降低 大规模推理 所需的整体硬件资源- 在本博客中,我们将探讨自推测解码的概念、其实现方式以及在 🤗 transformers 库中的实际应用-您将了解到其技术原理,包括 提前退出层 (Early-Exit Layers) 、 反嵌入 (Unembedding) 和 训练修改 (Training Modifications)为了将这些概念付诸实践,我们提供了代码示例、与传统推测解码的基准比较,以及对性能权衡的见解- 您还可以直接查看以下 Hugging Face 资源,了解更多关于该方法的信息并亲自尝试: Hugging Face 论文讨论论坛 LayerSkip 模型集合 展示自推测解码深入工作原理的 Colab 笔记本 推测解码与自推测解码 LayerSkip 演示 GIF 在 facebook/layerskip-llama2-7B 上的 LayerSkip 推理演示 (使用 LayerSkip 方法持续预训练的 Llama2 7B) 传统的推测解码 使用 两个 模型: 一个较小的模型 (草稿模型) 用于生成一系列候选词元,一个较大的模型 (验证模型) 用于验证草稿的准确性较小的模型执行大部分生成工作,而较大的模型则负责改进结果这提高了文本生成速度,因为较大的模型一次性验证完整序列,而不是逐个生成词元- 在自推测解码中,作者在此概念的基础上,使用大模型的早期层来生成草稿词元,然后由模型的更深层进行验证这种推测解码的“自洽”特性需要特定的训练,使模型能够同时执行草稿生成和验证这反过来又比传统的推测解码提高了速度并降低了计算成本 在 transformers 中的使用 为了在 🤗 transformers 库中启用提前退出自推测解码,我们只需在 generate() 函数中添加 assistant_early_exit 参数 以下是一个简单的代码片段,展示了该功能: pip install transformers from transformers import AutoTokenizer, AutoModelForCausalLM early_exit_layer = 4 prompt = "Alice and Bob" checkpoint = "facebook/layerskip-llama2-7B" tokenizer = AutoTokenizer.from_pretrained(checkpoint) inputs = tokenizer(prompt, return_tensors="pt").to("cuda") model = AutoModelForCausalLM.from_pretrained(checkpoint).to("cuda") outputs = model.generate(**inputs, assistant_early_exit=early_exit_layer) 注意: 虽然 assistant_early_exit 参数可以为任何仅解码器的 transformer 启用提前退出自推测解码,但除非模型经过专门训练,否则无法反嵌入 (通过 LM 头进行解码的过程,在博客文章后面有描述) 中间层的 logits只有对检查点进行这样的训练,以提高早期层的准确性,您才能获得加速LayerSkip 论文提出了一种训练方法来实现这一点 (即应用提前退出损失,并逐步增加层丢弃率)这里 提供了使用 LayerSkip 训练方法持续预训练的 Llama2、Llama3 和 Code Llama 检查点的集合 基准测试 我们进行了一系列广泛的基准测试,以衡量 LayerSkip 的自推测解码相对于自回归解码在各种模型上的加速情况我们还将自推测解码 (基于提前退出) 与标准推测解码技术进行了比较要复现这些结果,您可以在 这里 找到代码
阅读全文