MInference:长上下文 LLM 推理加速,单卡 A100 实现百万 token 推理,速度提升 10 倍

原文标题:单卡A100实现百万token推理,速度快10倍,这是微软官方的大模型推理加速

原文作者:机器之心

冷月清谈:

-**问题:**大型语言模型(LLM)在处理长上下文时面临预填充阶段时间过长的挑战,影响用户体验和 LLM 的广泛应用。 -**解决方案:**微软研究院提出 MInference 方法,一种基于动态稀疏注意力的预填充加速方法。 -**方法:**MInference 通过离线搜索确定最佳动态稀疏模式,并动态逼近推理过程中的稀疏指数,利用高效的 GPU 内核执行动态稀疏注意力计算。 -**效果:**MInference 在单台 A100 机器上将 LLaMA-3-8B-1M 和 GLM-4-1M 等长上下文 LLM 的推理速度提升了 10 倍,同时保持准确性,扩展了 LLM 的实际上下文窗口大小。 -**应用:**MInference 可直接应用于现有 LLM,无需修改预训练设置或额外微调。

怜星夜思:

1、MInference 的动态稀疏模式如何识别和选择?
2、不同注意力模式的性能差异和适用场景有哪些?
3、MInference 是否适用于所有 LLM 模型和下游任务?

原文内容

机器之心报道

编辑:张倩、陈萍

微软的这项研究让开发者可以在单卡机器上以 10 倍的速度处理超过 1M 的输入文本。


大型语言模型 (LLM) 已进入长上下文处理时代,其支持的上下文窗口从先前的 128K 猛增到 10M token 级别。

然而,由于注意力机制的二次复杂度,模型处理输入提示(即预填充阶段)并开始产生第一个 token 可能需要几分钟时间。导致首个 token 生成的时间过长,从而严重影响了用户体验,这也极大地限制了长上下文 LLM 的广泛应用。 

举例来说(如图 2a 所示),在单台装有 A100 的机器上为 LLaMA-3-8B 提供服务时,如果提示有 30 万个 token,模型需要 6 分钟才能完成预填充( pre-filling)阶段,如果提示增加到 100 万个 token,这个数字将增加到 30 分钟。


自注意力计算的开销占到了总预填充延迟的 90% 以上,这使其成为 LLM 处理长上下文时的主要瓶颈。现有的加速预填充方法在应用于长上下文 LLM 时通常无法保持可接受的准确性或效率。

为了解决上述问题,来自微软、萨里大学的研究者提出了一种旨在加速长序列处理预填充的稀疏计算方法:MInference( Milliontokens Inference )。


  • 论文地址:https://arxiv.org/pdf/2407.02490
  • 论文主页:https://hqjiang.com/minference.html
  • 论文标题:MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention

MInference 可以直接应用于现有 LLM,无需对预训练设置进行修改或额外的微调。

通过对各种下游任务(包括 InfiniteBench、RULER、PG-19 和 Needle In A Haystack)以及模型(包括 LLaMA-3-1M、Yi-200K、GLM-4-1M、Phi-3-128K 和 Qwen2-128K)进行评估,实验证明 MInference 可有效将 A100 上的预填充推理延迟降低多达 10 倍,同时保持准确性。


图片

使用 MInference 1.0 ,长上下文 LLM(如 LLaMA-3-8B-1M、GLM-4-1M)在单个 A100 上的推理速度实现了 10 倍提升,并且准确度更高。

方法介绍

作者提出了 MInference,这个名字反映了他们希望在一台 A100 机器上实现百万(million)token 推理的雄心。

MInference 是一种无需训练的高效方法,用于基于动态稀疏注意力的长上下文 LLM 的预填充阶段。

研究者认为注意力,特别是在长上下文中,是稀疏和动态的,即在不同的输入中,稀疏模式有很大的不同。这种动态稀疏性呈现出三种适用于所有输入的独特空间聚合模式:A 形(A-shape)、垂直 - 斜线(Vertical-Slash)和块状 - 稀疏(Block-Sparse)。


MInference 首先使用内核感知稀疏模式搜索算法为每个头部离线确定最佳动态稀疏模式,如算法 1 所示。在推理过程中,它会根据头部的模式动态逼近动态稀疏指数,如算法 2、3 所示。最后,作者使用优化后的 GPU 内核执行高效的动态稀疏注意力计算,大大减少了长上下文 LLM 的预填充阶段延迟。


例如,对于「垂直 - 斜线」模式,作者首先利用最后一个 Q 和 K 之间的注意力计算来估计垂直线和斜线的最佳指数。然后,他们利用动态稀疏编译器 PIT 和 Triton 构建垂直 - 斜线 FlashAttention 内核,加速注意力计算。对于 A 形、垂直 - 斜线和块状 - 稀疏模式,作者首先在注意力计算中使用 Q 和 K 的均值池。利用均值池和 MatMul 的交换属性,可以估算出块状 - 稀疏指数。然后,他们使用 Triton 构建块稀疏 FlashAttention 内核,加速注意力计算。有关内核的详细实现,请参阅附录 C.4 和代码。

在长上下文基准中的评估结果

作者在一系列场景中测试了 MInference,包括 QA、编码、基于检索的任务、multi-hop QA、总结和数学任务。RULER 基准包括几个复杂的 multi-hop 或 multi-needle 任务,有效地反映了 LLM 的实际上下文窗口大小。如表 1 所示,MInference 有效地保留了 LLM 的实际上下文窗口处理能力,甚至将实际上下文窗口大小略微扩展到 32K。


作者还使用平均 token 长度为 214K 的 InfiniteBench 在更广泛的任务中测试了 MInference,如表 2 所示。与 SoTA 基线相比,MInference 在所有任务中都始终保持了良好的性能。值得注意的是,在更具挑战性的检索任务(如 KV 检索任务)中,所有基线都无法做出准确预测,准确率低于 1.2%。但是,MInference 成功地保留了处理动态 KV 对检索的能力。


为了进一步评估不同上下文长度和关键信息在提示中不同位置时的性能,作者使用「大海捞针」任务测试了各种模型和方法。如图 1 所示,MInference 在不同的模型、上下文窗口和提示信息位置下都表现良好,与原始模型相比,其性能保持不变甚至略有提高。在 LLaMA-3-8B 和 GLM-4-9B-1M 的情况下,MInference 在高达 1M 的上下文窗口中实现了完全绿色的性能。相比之下,即使在 70K 上下文窗口中,StreamingLLM 和 InfLLM 在提示的中间段性能也会下降到 20% 以下。


作者还使用 PG-19 在语言模型任务中测试了 MInference,其中包括多达 100k 的 token。如图 2 所示,MInference 有效地保持了 LLaMA-3-8B 和 Yi-9B-200K 的困惑度,而所有基线都出现了不同程度的困惑度下降。此外,与标准的 StreamingLLM 相比,使用膨胀和步长配置的 StreamingLLM 更好地保持了困惑度性能。


延迟和内核中的稀疏模式 

图 3 展示了本文提出的三种注意力模式以及 FlashAttention 的微基准测试结果。可以看出,Vertical-Slash 是三种模式中最慢的,但在 1M 上下文窗口下,相比 FlashAttention 仍然实现了 13 倍的加速。


图 4 展示了 Vertical-Slash 头部内核中的稀疏索引。垂直线通过 PIT FlashAttention 使用 1x64 块计算,而斜线通过块级 FlashAttention 使用 64x64 块计算。


参考链接:https://hqjiang.com/minference.html

留住用户,AIGC如何通过个性化提升转化率?

7月10日,《AIGC体验派》第四期,邀请到火山引擎直播洞察产品负责人刘晨晨和火山引擎销售助手产品负责人梁正,一起聊聊用户留存转化率
  • 提升直播间转化率,AIGC复盘如何留住用户?

  • 聊到客户心坎里,金牌话术生成让小白秒变“老司机”?

  • 销售Copilot,基于大模型的销售助手究竟是噱头还是真香?

识别海报二维码点击阅读原文,立即报名直播。

© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:[email protected]

所谓动态稀疏,就是MInference可以根据不同的情况,自动识别出稀疏模式,这就像人的大脑,可以根据不同的输入自动调整处理方式。估计微软的研究员下了不少功夫来设计这个算法,真厉害!感觉要超越人类智能的时代要提前到来了!

不同模式的性能差异和适用场景主要取决于注意力矩阵的稀疏模式,比如文本中不同 token 之间的关联性。具体到你的应用场景,建议实际测试一下不同模式的效果,选择最优的。

在评估中,MInference 在 QA、编码、检索、多跳 QA、总结和数学等任务上都取得了很好的效果,表明它对各种下游任务具有通用性。

需要注意的是,MInference 专注于加速预填充阶段,对于推理过程中的其他部分,可能需要额外的优化。此外,如果你的 LLM 模型或下游任务有特殊要求,可能需要做一些调整。

垂直 - 斜线模式虽然最慢,但对于 1M 上下文窗口,仍比 FlashAttention 快 13 倍。

另外,MInference 目前只支持单 GPU,如果你需要分布式训练或推理,可能需要进一步扩展。不过,作为一个不断发展的项目,这些功能可能会在未来版本中加入。你想了解更多的技术细节吗?

MInference 可以直接应用于现有 LLM,无需修改预训练设置或额外微调,包括 LLaMA、Yi、GLM、Phi 和 Qwen 等模型。

该算法分析注意力计算中 Q 和 K 之间的关联,估计稀疏模式中垂直线和斜线的最佳指数,然后使用均值池来估计块状稀疏指数。

A 形模式适用于注意力矩阵中存在较多连续非零元素的情况,垂直 - 斜线模式适用于注意力矩阵中存在成组非零元素的情况,块状 - 稀疏模式适用于注意力矩阵中存在较大块状非零元素的情况。

MInference 使用内核感知稀疏模式搜索算法,为每个头部离线确定最佳动态稀疏模式,考虑了注意力在不同输入中的动态稀疏性。

对于特别长的上下文,比如上百万 token 的输入,垂直 - 斜线模式可能更合适,因为它可以有效利用稀疏性,减少计算量。而对于上下文长度较短的情况,A 形或块状 - 稀疏模式可能更合适一些。你想试一试吗?

根据提示内容和上下文的不同,MInference 会选择最合适的稀疏模式,例如 A 形、垂直 - 斜线或块状 - 稀疏。我猜这就像是 LLM 学习的魔法,它可以根据不同的输入自动调整。