FastDraft: How to Train Your Draft
作者: Ofir Zafrir, Igor Margulis, Dorin Shteyman, Shira Guskin, Guy Boudoukh
分类: cs.CL
发布日期: 2024-11-17 (更新: 2025-06-05)
备注: Accepted at ACL 2025
💡 一句话要点
FastDraft：通过高效预训练和对齐，为大型语言模型快速训练Draft模型。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 推测解码 Draft模型 大型语言模型 模型加速 预训练 模型对齐 高效推理
📋 核心要点
- 推测解码依赖于与目标模型词汇兼容的高效Draft模型，而现有模型往往缺乏这种Draft模型。
- FastDraft通过高效预训练和对齐，为大型语言模型快速训练Draft模型，无需严格的词汇兼容性约束。
- 实验表明，FastDraft在代码补全等任务上实现了高达3倍的加速，并在其他任务上实现了高达2倍的加速。
📝 摘要（中文）
推测解码(Speculative Decoding)作为一种加速大型语言模型自回归推理过程的有效技术，已广受欢迎。然而，推测解码完全依赖于高效Draft模型的可用性，但由于词汇兼容性的严格约束，许多现有语言模型通常缺乏这种模型。本文介绍了一种新颖高效的方法FastDraft，通过结合高效的预训练，然后对目标模型生成的合成数据集进行微调，从而预训练Draft模型并将其与任何大型语言模型对齐。我们通过为流行的Phi-3-mini和Llama-3.1-8B模型训练两个高度参数高效的Draft模型来展示FastDraft。使用FastDraft，我们能够在单个服务器上使用8个Intel® Gaudi® 2加速器，在不到24小时内生成一个大约100亿个token的Draft模型。我们的结果表明，在代码补全方面，该Draft模型在接受率、块效率等关键指标上取得了令人印象深刻的结果，并实现了高达3倍的内存限制加速，在摘要、文本补全和指令任务中实现了高达2倍的加速。我们通过在最新的Intel® Core™ Ultra上进行基准测试验证了我们的理论发现，实现了高达2倍的实际运行时间加速，表明运行时间显著减少。由于其高质量，FastDraft解锁了AI-PC和其他边缘设备上的大型语言模型推理。
🔬 方法详解
问题定义：论文旨在解决大型语言模型推测解码中，由于缺乏高效且词汇兼容的Draft模型而导致的推理速度瓶颈问题。现有方法要么需要耗时的从头训练，要么难以保证Draft模型与目标模型的词汇兼容性，从而限制了推测解码的效率。
核心思路：FastDraft的核心思路是结合高效的预训练和对齐策略，快速生成与目标模型兼容的高质量Draft模型。首先进行高效的预训练，然后利用目标模型生成的合成数据进行微调，从而使Draft模型能够更好地模仿目标模型的行为，提高推测解码的接受率和效率。
技术框架：FastDraft主要包含两个阶段：预训练阶段和对齐阶段。在预训练阶段，使用大规模文本数据对Draft模型进行初始化，使其具备基本的语言建模能力。在对齐阶段，利用目标模型生成合成数据集，并使用这些数据对Draft模型进行微调，使其能够更好地预测目标模型的输出。
关键创新：FastDraft的关键创新在于其高效的预训练和对齐策略。通过高效的预训练，可以快速初始化Draft模型，减少训练时间。通过利用目标模型生成的合成数据进行微调，可以有效地提高Draft模型与目标模型的兼容性，从而提高推测解码的效率。
关键设计：FastDraft的关键设计包括：(1) 使用参数高效的模型结构作为Draft模型，以减少计算和存储开销；(2) 设计合适的损失函数，以促进Draft模型与目标模型的对齐；(3) 精心选择预训练数据和合成数据集，以提高Draft模型的性能。
🖼️ 关键图片
📊 实验亮点
FastDraft在Phi-3-mini和Llama-3.1-8B模型上进行了验证，结果表明，使用FastDraft训练的Draft模型在代码补全任务上实现了高达3倍的内存限制加速，在摘要、文本补全和指令任务中实现了高达2倍的加速。此外，在Intel® Core™ Ultra上的基准测试表明，FastDraft实现了高达2倍的实际运行时间加速。
🎯 应用场景
FastDraft可应用于各种需要加速大型语言模型推理的场景，例如AI-PC、边缘设备和云计算平台。通过使用FastDraft训练的Draft模型，可以显著提高大型语言模型的推理速度，降低延迟，从而改善用户体验，并降低计算成本。该技术还有助于在资源受限的设备上部署大型语言模型。
📄 摘要（原文）
Speculative Decoding has gained popularity as an effective technique for accelerating the auto-regressive inference process of Large Language Models. However, Speculative Decoding entirely relies on the availability of efficient draft models, which are often lacking for many existing language models due to a stringent constraint of vocabulary compatibility. In this work we introduce FastDraft, a novel and efficient approach for pre-training and aligning a draft model to any large language model by incorporating efficient pre-training, followed by fine-tuning over synthetic datasets generated by the target model. We demonstrate FastDraft by training two highly parameter efficient drafts for the popular Phi-3-mini and Llama-3.1-8B models. Using FastDraft, we were able to produce a draft model with approximately 10 billion tokens on a single server with 8 Intel$^\circledR$ Gaudi$^\circledR$ 2 accelerators in under 24 hours. Our results show that the draft model achieves impressive results in key metrics of acceptance rate, block efficiency and up to 3x memory bound speed up when evaluated on code completion and up to 2x in summarization, text completion and instruction tasks. We validate our theoretical findings through benchmarking on the latest Intel$^\circledR$ Core$^{\tiny \text{TM}}$ Ultra, achieving a wall-clock time speedup of up to 2x, indicating a significant reduction in runtime. Due to its high quality, FastDraft unlocks large language models inference on AI-PC and other edge-devices.