革命性大型语言模型推理:推测解码与低精度量化

革命性大型语言模型推理:推测解码与低精度量化

随着人工智能(AI)的快速发展,大型语言模型(LLM)已成为自然语言处理(NLP)的基石。这些模型在语言生成和理解方面展现出卓越的能力,使其成为机器翻译、对话系统和文本生成等应用中不可或缺的一部分。然而,随着这些模型规模的持续增长,其推理效率和资源消耗仍然是大规模部署的主要瓶颈。

近年来,研究人员提出了多种优化技术来提升 LLM 推理效率。在我们之前的工作中,我们通过基于超大型语言模型(vLLM)框架的动态 KV-cache 压缩,实现了 1.5 倍的推理加速,为优化 LLM 性能奠定了坚实基础。

本文深入探讨了另外两种优化策略:推测解码和低精度量化。这些方法旨在最大化主流硬件的计算能力,降低推理成本并提高速度。通过采用这些技术,我们力求为提升 LLM 推理效率提供新的视角和实用解决方案,以应对可扩展 AI 部署的挑战。

推测解码

技术原理

推测解码是一种通过引入较小模型生成多个候选 token,再由大模型进行验证,从而实现并行解码以提升速度的推理加速技术。该方法的可行性基于两个关键因素:

内存瓶颈
在现代 GPU 硬件中,内存访问带宽通常慢于推理所需的计算速度,导致该过程显著受限于内存。大型语言模型推理阶段的 GPU 内存流量主要由模型大小决定,而 GPU 往往有剩余的计算能力。通过利用小模型并行推理来充分利用这些余量,可以提升整体效率。

小模型预测准确性
当前主流语言生成模型通常是基于 Transformer 架构的自回归模型。小模型在预测常见语言模式(例如习语搭配或标准表达)时具有较高的准确性。当小模型成功预测下一个 token 时,大模型可以直接复用该结果,从而显著减少计算开销。

实现方法

推测解码方法主要包括以下步骤:

  • 多轮候选生成

    • 使用经过专门训练的小模型进行快速推理,为每个位置生成多个高质量候选 token。
    • 创新地采用动态概率阈值调整机制,自适应控制候选数量。
    • 引入上下文感知缓存机制,提高小模型在连续文本生成过程中的预测准确性。
  • 高效验证策略

    • 设计批量验证机制,将多个候选 token 分组并一次性送入大模型进行评分。
    • 实现提前终止策略,在发现高置信度匹配时快速返回结果。
    • 支持验证结果的异步处理,减少 GPU 空闲等待时间。
  • 智能调度系统

    • 动态平衡小模型预测与大模型验证之间的计算资源分配。
    • 基于历史统计的自适应预测,优化候选 token 生成策略。
    • 实现细粒度的任务划分与调度,最大化硬件利用率。

图1. 动态草稿长度的推测解码

图1. 动态草稿长度的推测解码

实验结果

基于业界领先的 vLLM 推理框架,我们进行了深度定制与开发,取得了显著的性能提升。我们在 H20 GPU 上使用 LLaMA-3.1-70B-Instruct 模型进行了实验,测试了不同上下文长度场景。此外,我们还与无推测解码的 vLLM 进行了性能对比,重点关注不同批量大小下的吞吐量,实验结果如图 2 所示。

  • 性能提升

    • 在保持输出质量的前提下,实现了 1.4 倍的整体推理加速。
    • 完全兼容现有 vLLM 功能,可无缝集成到当前系统中。
    • 支持多种规模的语言模型,对大规模模型的提升尤为明显。

图2. 推测解码加速比图

图2. 推测解码加速比图

  • 技术优势

    • 基于概率的动态采样策略,实现高效的 token 草拟与验证。
    • 稳定的智能调度机制,确保一致的加速效果。
    • 成本效益高,资源开销低。

低精度量化

技术原理

  • 量化概述

量化是一种广泛采用的技术,通过将模型权重和激活从高位宽表示转换为较低位宽表示,从而降低大规模语言模型(LLM)的计算和内存成本。具体而言,许多方法涉及将 FP16 张量量化为低位整数张量,如下所示:

将FP16张量量化为低位整数张量

  • 训练后量化(PTQ)

PTQ 对预训练模型进行量化,无需重新训练,从而避免了高昂的再训练成本。然而,将传统量化方法应用于 LLM 会带来独特的挑战:

  • LLM 在权重和激活中表现出更高频率的异常值以及更广的分布范围。
    • 因此,对 LLM 进行有效量化需要专门的优化技术,以在保持效率的同时最小化精度损失。

实现方法

  • 异常值优化

大型语言模型权重和激活中存在的异常值给量化过程带来了显著挑战。传统方法通常通过裁剪或简单缩放来处理异常值,这可能适用于较小模型,但往往会对 LLM 的性能产生重大影响。

我们的方法通过详细的分析和优化策略改进了异常值处理过程。它不仅细化了异常值的计算,还优化了相关数据结构,确保模型在最小化精度损失的同时保持高效计算能力。例如,我们采用更灵活的动态范围分布策略,对异常值进行分层处理,从而更准确地捕捉模型的特征。

  • 改进算子与层融合

我们设计了一系列改进算子,能够更高效地处理量化后的数据结构。例如,通过层融合技术,在前向计算时将多个连续算子合并为单个算子,以减少计算开销和中间数据传输。

此外,对于常见的操作(如矩阵乘法与激活函数的组合),我们设计了高效的低位算子来加速模型推理。这一改进不仅显著降低了硬件层面的带宽需求,还进一步提升了整体计算效率。

  • 独特的 FP8 量化方案

业界主流量化方法通常采用从 FP16 到 INT8 的转换,但这种方法可能不适用于需要更低精度的场景。我们的 FP8 量化方案引入了自适应缩放因子,在保留模型数值表达能力的同时进一步降低位宽。在实现中,我们调整了 FP8 中指数和尾数的比例,以更好地适应 LLM 的分布特征,确保模型即使在非常低的位宽下也能保持准确性。

  • KVCache 量化

在支持大规模推理的场景中,KVCache 的存储和访问效率至关重要。传统方法通常对 KVCache 采用标准量化技术,而没有针对其存储结构的具体特征进行优化。

我们提出了一种专门的 KVCache 量化技术,通过对查询和键值对应用特定量化策略,最大程度减少计算资源需求。具体而言,我们优化了 KVCache 的动态更新和存储密度,结合分片量化和稀疏矩阵存储方法,显著提高了推理过程中的缓存利用率。

  • 推理流程优化

模型权重从 FP16 量化为 FP8 格式用于推理。输入 token 经过嵌入层和 Transformer 模块,包括注意力机制和前馈网络。

利用 FP8 TensorCore 和优化的 KVCache,系统高效地存储和更新键值对。

线性变换和 Softmax 层生成输出概率分布,产生最终的 token 输出。

该方法充分利用 FP8 TensorCore 的计算效率,以较低的计算成本实现更快的推理。在图 3 中,原本为 FP16 格式的模型权重被量化为 FP8 进行推理,显著减少了数据访问开销,提高了速度。在 LLM 推理过程中,QKV 投影、flash attention 和前馈层等操作以 FP8 格式进行矩阵乘法,利用 Tensor Core 加速,大幅提升推理速度。为减轻精度损失,我们使用少量额外数据进行 PTQ,确保模型推理精度不受影响。在长文本场景中,KV 数据访问的开销可能拖慢推理速度。为此,KV 数据也被量化为 FP8,进一步优化模型的推理性能。

图3. FP8 精度量化流水线概述

图3. FP8 精度量化流水线概述

实验结果

  • 图 4 展示了将低精度量化应用于 Llama3-8B-BF16 后的加速比,其中输入输出长度分别为(10000-1000,5000-500,2000-200),批量大小(Batch Size)可变。输入输出长度越长,加速效果越显著,尤其是在大批量推理时。在单张 4090 GPU 上运行 Llama3-8B-FP8-KV8,将 ttft(首 token 时间)限制在 2 秒以内时,批量大小约为 4,加速比为 1.4 倍。

图4. FP8 量化加速比

图4. FP8 量化加速比

  • 我们的方法在处理较大输入输出长度时显著降低了异常值的影响,并通过层融合等优化提升了计算和传输效率。传统业界 PTQ 方法在处理大规模模型的长输入输出长度时往往面临困难,而我们的解决方案专门针对这一问题设计,确保模型性能稳定。
  • 实验表明,在 4090 GPU 上,我们的方法在处理长输入输出长度和中等批量大小时展现出优越的加速效果,同时保持推理延迟(ttft)低于 2 秒,满足实时推理需求。

质量评估

要使用 lm-evaluation-harness 工具测试不同模型 API(例如 openrouter.ai)在 mmlu_pro(5-shot)测试集上的表现,可按照以下步骤使用工具的默认配置进行推理。结果如图 5 所示。

图5. Meta-Liama-3.1-8B-Instruct-FP8 与 Meta-Liama-3.1-8B-Instruct 的精确匹配率对比

图5. Meta-Liama-3.1-8B-Instruct-FP8 与 Meta-Liama-3.1-8B-Instruct 的精确匹配率对比

结论

本文介绍了推测解码和低精度量化作为互补技术,旨在提升大型语言模型(LLM)的推理效率。推测解码通过利用较小模型生成候选 token,再由大模型并行验证,加速了推理过程,实现了 1.4 倍的加速。

相比之下,低精度量化将模型权重转换为更低位的表示,同时保持模型性能,从而降低了计算和内存开销。总的来说,这些方法为高效推理提供了新的见解和实用解决方案,对动态 KV-cache 压缩等现有技术形成了补充。

原文发布于 Novita AI

Novita AI 是**一个 AI 云平台,为开发者提供使用简单 API 轻松部署 AI 模型的能力,同时也提供经济实惠且可靠的 GPU 云,用于构建和扩展应用。

推荐阅读

  1. 基于 vLLM 框架的动态 KV Cache 压缩
  2. 如何选择最佳 GPU 进行 LLM 推理:基准测试洞察
  3. KV Sparsity 如何为 vLLM 实现 1.5 倍加速