AMUSD: Asynchronous Multi-Device Speculative Decoding for LLM Acceleration

📄 arXiv: 2410.17375v1 📥 PDF

作者: Bradley McDanel

分类: cs.CL, cs.DC, cs.LG

发布日期: 2024-10-22

备注: 4 pages, 5 figures, 1 table, 1 algorithm

🔗 代码/项目: GITHUB

💡 一句话要点

提出AMUSD：一种用于LLM加速的异步多设备推测解码方法

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 推测解码 异步计算 多设备并行 LLM加速

📋 核心要点

  1. 现有推测解码方法在草稿和验证阶段交替进行，设备利用率不高，限制了加速效果。
  2. AMUSD通过异步方式解耦草稿和验证过程，使草稿模型和验证模型可以在不同设备上并行工作。
  3. 实验结果表明，AMUSD在多个数据集上相比推测解码平均提升29%，相比自回归解码最高提升1.96倍。

📝 摘要（中文）

大型语言模型通常以自回归方式生成token，将每个token作为下一个token的输入。最近关于推测解码的研究试图通过使用更小、更快的草稿模型来更快地生成候选token，从而加速这一过程。然后，这些候选token由更大的（原始）验证模型并行验证，与单独以自回归方式使用更大的模型相比，从而实现了整体加速。在这项工作中，我们介绍AMUSD（异步多设备推测解码），该系统通过将草稿和验证阶段解耦为连续的异步方法，进一步加速了生成。与传统的推测解码不同，在传统的推测解码中，一次只有一个模型（草稿或验证）执行token生成，而AMUSD使两个模型都能够在单独的设备（例如，GPU）上独立执行预测。我们在多个数据集上评估了我们的方法，结果表明，AMUSD比推测解码平均提高了29%，比传统的自回归解码提高了高达1.96倍的速度，同时实现了相同的输出质量。我们的系统是开源的，可在https://github.com/BradMcDanel/AMUSD/上获得。

🔬 方法详解

问题定义：现有的大型语言模型（LLM）推理过程通常采用自回归方式，即逐个生成token，速度较慢。推测解码旨在通过引入一个小型草稿模型来预测多个候选token，然后由大型验证模型并行验证这些候选token，从而加速推理。然而，传统的推测解码方法中，草稿模型和验证模型是交替工作的，即一个模型在生成token时，另一个模型处于空闲状态，导致设备利用率不高，限制了加速效果。

核心思路：AMUSD的核心思路是将草稿模型的token生成过程和验证模型的token验证过程完全解耦，使其能够异步并行执行。具体来说，草稿模型持续不断地生成候选token，并将这些token发送给验证模型进行验证。同时，验证模型也持续不断地接收来自草稿模型的候选token，并进行验证。通过这种异步并行的方式，可以充分利用计算资源，提高整体的推理速度。

技术框架：AMUSD的整体框架包含两个主要模块：草稿模型和验证模型。这两个模型部署在不同的设备上（例如，不同的GPU）。草稿模型负责快速生成候选token序列，并将这些序列发送给验证模型。验证模型接收到候选token序列后，并行地验证这些token的正确性。验证结果会反馈给草稿模型，用于指导后续的token生成。整个过程是异步进行的，草稿模型和验证模型可以独立地进行计算，无需等待对方完成任务。

关键创新：AMUSD最重要的创新点在于其异步多设备的设计。与传统的推测解码方法相比，AMUSD不再需要草稿模型和验证模型交替工作，而是允许它们并行执行。这种异步并行的方式可以显著提高设备利用率，从而实现更高的推理速度。此外，AMUSD还支持多设备部署，可以将草稿模型和验证模型部署在不同的设备上，进一步提高并行度。

关键设计：AMUSD的关键设计包括：1) 异步通信机制：草稿模型和验证模型之间需要高效的异步通信机制来传递候选token和验证结果。2) 动态批处理：验证模型可以动态地将接收到的候选token进行批处理，以提高验证效率。3) 优先级调度：可以根据token的置信度对验证任务进行优先级调度，优先验证置信度较低的token。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

AMUSD在多个数据集上进行了评估，实验结果表明，AMUSD相比传统的推测解码方法平均提升了29%的推理速度，相比传统的自回归解码方法最高提升了1.96倍的推理速度，同时保持了与原始模型相同的输出质量。这些结果表明，AMUSD是一种有效的LLM加速方法。

🎯 应用场景

AMUSD可广泛应用于需要快速LLM推理的场景，例如实时对话系统、快速文本生成、机器翻译等。通过提高LLM的推理速度，可以显著提升用户体验，并降低部署成本。未来，AMUSD可以进一步扩展到更多类型的LLM和硬件平台，并与其他加速技术相结合，以实现更高的性能。

📄 摘要（原文）

Large language models typically generate tokens autoregressively, using each token as input for the next. Recent work on Speculative Decoding has sought to accelerate this process by employing a smaller, faster draft model to more quickly generate candidate tokens. These candidates are then verified in parallel by the larger (original) verify model, resulting in overall speedup compared to using the larger model by itself in an autoregressive fashion. In this work, we introduce AMUSD (Asynchronous Multi-device Speculative Decoding), a system that further accelerates generation by decoupling the draft and verify phases into a continuous, asynchronous approach. Unlike conventional speculative decoding, where only one model (draft or verify) performs token generation at a time, AMUSD enables both models to perform predictions independently on separate devices (e.g., GPUs). We evaluate our approach over multiple datasets and show that AMUSD achieves an average 29% improvement over speculative decoding and up to 1.96$\times$ speedup over conventional autoregressive decoding, while achieving identical output quality. Our system is open-source and available at https://github.com/BradMcDanel/AMUSD/.