MagicDec: LLM 长序列、大 Batch 投机采样 2 倍推理加速

最新推荐文章于 2024-10-04 05:48:08 发布

江湖人称麻花滕

最新推荐文章于 2024-10-04 05:48:08 发布

阅读量1.1k

点赞数 22

文章标签： batch 深度学习开发语言人工智能 chatgpt 知识图谱数据库

本文链接：https://blog.csdn.net/m0_59235699/article/details/141727056

版权

一、摘要

LLM 在交互式聊天机器人、文档分析和 Agent 工作流等长上下文应用中变得越来越普遍，但以低延迟和高吞吐服务长上下文请求却很有挑战。投机解码（Speculative Decoding，SD）是一种广泛采用的技术，可以在不牺牲性能的情况下减少延迟，但传统观点认为，它的功效仅限于小 Batch Size（PS：其实大家说的是在常规序列长度下 SD 只适合小 Batch Size）。在 MagicDec 中，作者发现 SD 对于中长序列也能实现加速，甚至基于严格的分析，得出随着 Batch Size 的增加可以实现更好的加速。

首先，确定随着 Batch Size 和序列长度的增加而发生瓶颈转移，并利用这些见解更有效地实施投机解码，以实现高吞吐推理。然后，利用具有稀疏 KV Cache 的 Draft 模型来解决随着序列长度和 Batch Size 而增加的 KV 瓶颈问题。该方法扩展了投机解码在长上下文服务中的广泛适用性，其可以在不影响准确性的情况下提高吞吐量并减少延迟。

对于中长序列，作者在 8 个 A100 上实施 32 到 256 的 Batch 处理时，LLaMA-2-7B-32K 可以加速 2x，LLaMA-3.1-8B 可以加速 1.84x。

二、计算强度 & Roofline Model

2.1 基本概念

AI 模型由许多算子（OP）组成，比如，常见的矩阵乘矩阵、矩阵乘向量、卷积、ReLU 等激活和矩阵加、减法等。对于任何一个算子，其耗时主要由三部分组成：

访存：SM 从 GPU 显存上读取输入数据或写回输出数据，其速度主要受 GPU 显存带宽限制。
计算：根据输入数据通过数学计算获得输出结果，其速度主要受 GPU 中的算力核限制，比如 CUDA Core、Tensor Core 和 SFU 等。
调度：比如 CUDA 编程中常说的 Kernel Launch。（这里不再展开）

访存时间可以表示为，计算时间可以表示为，其中 BW 表示 BandWidth。理想情况下，如果，则可以通过 Overlap 的方式让访存和计算充分折叠，显存带宽和算力都发挥到极致。否则，必将有某一个无法充分发挥：

，访存时间更短，此时算力成为瓶颈。
，计算时间更短，此时访存成为瓶颈。

由于和都是由硬件决定的，因此可以进一步转换为：

，计算成为瓶颈
，访存成为瓶颈

进一步，通常使用算术强度（Arithmetic Intensity）的方式来表示，如下：

则可以得出，对于任何一个 OP：

，则为计算瓶颈
，则为访存瓶颈

不同 GPU、不同数据类型对应的、和，如下所示：

2.2 常见算子分析

2.2.1 矩阵乘矩阵

以最常见的矩阵乘法为例，，假设数据类型 FP16，也就是都是 2 Byte，如下图所示，如果采用最常规的方式（每个计算核都计算 C 中的一个元素，C 中的每个元素都由 A 的一行和 B 的一列计算内积得到），则对应的访存和计算为：

访存：，读取 A、B，写回 C。
计算：，每个元素 K 个乘法和 K 个加法。

则如下所示，对应的 Arithmetic Intensity 约为 1/2，明显小于，此时为明显的访存瓶颈：

为了避免访存出现明显瓶颈，最常见的方案就是采用分块矩阵乘法，此时一个 SM 可以负责处理 C 中的一个块，也就是，理想情况下，只需读取一次 A、一次 B，写回一次 C，中间值全部在 SM 上的 Shared Memory 或者命中 Cache，则对应的访存和计算为：

访存：，读取 A、B，写回 C。
计算：，计算量不变。

则对应的 Arithmetic Intensity 如下所示，此时是计算瓶颈还是访存瓶颈与 M、N 和 K 的大小有关，但是可以有个近似结论，如果 M、N 和 K 的任何一个值小于，则会存在访存瓶颈：

比如，当 M=1024，K=1024, N=8，则对应算术强度如下，远小于 T4，V100、A100 相应 FP16 的，因此为访存瓶颈：

比如，当 M=1024，K=1024, N=512，则对应计算强度如下，大于 T4，V100、A100 相应 FP16 的，因此为计算瓶颈：

以 GPT-3 系列模型中涉及的矩阵乘进行评估，令 M=N=dmodel，令 K 等于 batch size。在 T4 GPU 上基于 cublas 实现 FP32 类型矩阵乘法，如下图所示为对应的 Arithmetic Intensity 与 K 的关系，可以看出，随着 K 的增加 Arithmetic Intensity 逐渐增加：

如下图所示为 GPU 的 GFLOPs（可以等效于算力利用率，T4 的 FP32 峰值算力为 8100 GFLOPs），可以看出，当 K 比较大时才能充分发挥 GPU 的算力（K 较小时存在访存瓶颈，并且 K 越小访存瓶颈越明显）：

2.2.2 矩阵乘向量

矩阵乘向量其实就是矩阵乘矩阵的特例，也就是 M 或 N 为 1 的情况（假设此时 K 和另一个值远大于 1），此时对应的 Arithmetic Intensity 如下所示约为 1，存在明显访存瓶颈，对于 V100 FP16 的为 138.9，则相当于只能发挥 1/138.9 左右的算力，存在极大算力浪费：

2.2.3 总结

如下所示为常见 OP 的 Arithmetic Intensity，可以看出，除了矩阵乘矩阵之外都是访存瓶颈：

OP	Arithmetic Intensity	瓶颈
矩阵乘矩阵(4096, 1024, 512)	315 FLOPS/B	计算
矩阵乘向量(4096, 1024, 1)	1 FLOPS/B	访存
MaxPooling(3x3 window)	2.25 FLOPS/B	访存
ReLU	0.25 FLOPS/B	访存
Layer Normalization	<10 FLOPS/B	访存

2.3 Roofline Model

Roofline Model 是一种直观的可视化性能模型，基于给定的硬件约束为我们提供潜在的优化建议和优先级。如下图所示：

左侧顶部红线表示硬件的峰值内存带宽（GPU 显存带宽）。
右侧顶部绿线表示硬件的峰值算力（GPU FP16 算力）。
中间的虚线表示基于以上两个硬件指标确定的分割线。
红色区域表示实现的函数比较高效，已经接近显存带宽峰值，但因为计算强度比较低，对应的浮点算力比较低，此时为访存瓶颈。
绿色区域表示实现的函数比较高效，已经接近峰值 FP16 算力，但因为计算强度比较高，对应的显存带宽并没有充分发挥，此时为计算瓶颈。

假设我们实现的 OP 对应图中不同颜色的点：

黄色：表示接近显存带宽峰值，优化的目标是减少访存，向右上移动。
蓝色：表示接近 FP16 算力峰值，如果无法减少计算量则没有进一步优化的空间。
紫色：计算强度比较小，但离显存带宽峰值还有距离，可以优先提升访存效率，但优化到极限也无法达到 FP16 算力极限，会受峰值显存带宽限制。
红色：计算强度比较大，但离 FP16 算力峰值还有距离，可以优先优化计算效率。

三、Transformer Attention 分析

3.1 MHA Attention 计算

如下图所示为标准的 LLM Decoding 阶段的 Multi-Head Attention（MHA）计算，其中的 D 表示 hidden size，H 表示 Head 个数，L 表示当前是在序列的第 L 个 Token。可以看出：

当 Batch Size 为 1 时，图中红色、绿色、蓝色处的矩阵乘法全部为矩阵乘向量，是明显的 Memory Bound，算术强度不到 1。
当 Batch Size 大于 1 时（比如 Continuous Batching）：
红色和蓝色部分：因为是 Weight 乘以 Activation，所以不同的 Request 之间可以共享 Weight。这里变成矩阵乘矩阵，并且 Batch Size 越大，算术强度越大，也就越趋近于 Compute Bound（FFN 层也类似，这里就不再赘述）。
绿色部分：这里 Q、K 和 V 的 Attention 计算，是 Activation 乘以 Activation，所以不同的 Request 之间没有任何相关性。即使 Batching，这里也是 Batched 矩阵乘向量，并且因为序列长度可能不同，这里不同 Request 的矩阵乘向量是不规则的。也就是说，这里算术强度始终不到 1，是明显的 Memory Bound。

从上可以看出，通过 Continuous Batching 可以很好地将 Memory Bound 问题转变为 Compute Bound，但 Q、K 和 V 的 Attention 计算的算术强度却始终小于 1。根据 Amdahl 法则，如果系统中有一部分无法优化，即使把其他部分优化到可以忽略，不可优化的部分也会决定整个系统的性能上限。不幸的是，Sequence Length 越长，这里的计算量就越不可忽略。

根据模型配置信息可以估算出模型中 Q、K 和 V 的 Attention 计算与其他矩阵计算的比例大约为 (L+D)/(12*D)（PS：准确值需要根据具体的模型参数计算）。也就是说，当序列长度 L 等于 11 倍的 hidden size 时，两部分的计算量相当，即使其他矩阵计算优化到 0，加速比也只有 2x。比如 LLaMA 2 7B 的 hidden size 为 4K，当序列长度达到 44K 时，两部分的计算量相当，要优化的重点也会很不一样，这也是很多长序列相关工作会在 Attention 部分采用稀疏 Attention 的一个重要原因。

此外，计算量与计算时间并不等价，实际计算时由于 Attention 部分始终是 Memory Bound，并且算术强度小于 1，而其他部分通过 Continuous Batching 的方式可以将算术强度近似提升到对应的 Batch Size，也就是说，当 Batch Size 比较大时，可能在序列长度为 4K 时，Attention 部分的计算时间已经超过非 Attention 的计算时间。

3.2 GQA Attention 计算

早期通常只有比较大的模型才会采用 GQA（[2305.13245] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints），比如 LLaMA -2 70B，而 LLaMA-2 7B/13B 都没有采用 GQA。然而，LLaMA-3 8B 中也用上了 GQA，甚至其他更小的模型也在将 MHA 替换为 GQA。

使用 GQA 有个非常大的好处：在推理阶段可以显著降低 KV Cache 的大小。比如，相比 32 个 KV Head 的 MHA，32 个 Query Head，8 个 KV Head 的 GQA 的 KV Cache 大小可以降低到 MHA 的 8/32=1/4，这也为更大的 Batch Size 提供了空间，可以进一步提升吞吐。
除此之外，还有一个比较大的好处：可以明显提升 Q、K 和 V 的 Attention 计算的算术强度。此时虽然不同的 Request 之间同样不能共享，但是同一个 Request 中的不同 Head 可以共享，比如 8 个 Query Head 共享 1 个 KV Head，则算术强度会接近于 8，计算效率明显提升，也可以更充分发挥 Tensor Core 的算力。

使用 MHA 时，Q、K 和 V 的 Attention 计算可以使用 CUDA Core 也可以使用 Tensor Core。由于 Tensor Core 要求矩阵的 Shape 是 8 的整数倍，如果不满足就只能 Padding：

对于 MHA 而言，其是矩阵乘向量，则有 7/8 的计算是冗余的。
对于 GQA 而言，如果 4 个 Query Head 共享 1 个 KV Head，则 Attention 计算有 4/8 的计算是冗余的，如果 8 个 Query Head 共享 1 个 KV Head，则没有计算的冗余。很多框架已经做了相关优化，比如 LMDeploy，TRT-LLM 的 XQA 等。

PS：对于 GQA 而言，理论上也可以期望 GPU 的 L2 Cache 能够缓存到共享的 Key 和 Value Cache，从而缓解 IO Bound 问题，然而实际上无法人为控制，不一定能达到理想的效果。所以一般是在算法实现上直接优化。

3.3 投机采样

我们在之前的文章中已经详细介绍过各种投机采样技术，其核心思路如下图所示，首先以低成本的方式快速生成多个候选 Token（小模型，多头，检索等方式），然后通过一次并行验证阶段快速验证多个 Token，进而减少大模型的 Decoding Step，实现加速的目的：

投机采样可以有效减少 Decoding Step 数量，这也是其存在的意义，然而验证的接受率会很大程度上影响最终的加速比，接受率越高，减少的 Decoding Step 数量就越多，因未接受而浪费的计算就越少（实际上只要不是接受率 100%，就一定存在计算的浪费）。

除此之外，当序列比较长时，由于减少 Decoding Step 而减少的对全局 Key、Value Cache 的访问更加可观，相当于在 Memory Bound 的时候用 Compute 换 IO。也就是说，投机采样可以明显提升 Attention 计算的算术强度，比如投机的 Token 长度为 10，则相当于可以将算术强度从 1 近似提升到 10，那么即使 Token 接受率只有 50%，也依然可以等价于将算术强度提升到 5。此外，投机采样的提升与 GQA 是正交的，两者能共同提升算术强度。

举个简单的例子：假设序列长度比较长，Batch Size 比较大，Attention 是明显的 Memory Bound，计算强度小于 1。此时 Attention 计算时间为 80ms，其他部分为 20ms，每次猜测 10 个 Token，平均接受率为 50%，要 Decoding 生成 100 个 Token。

如果不使用投机采样，则总共的生成时间为：(80+20)*100=10s。
如果使用投机采样，则平均只需要 100/(10*50%)=20 个 Decoding Step。每次猜测 10 个 Token，算术强度从 1 增加到 10，其 Latency 并不会明显增加，假设 Attention 计算时间变为 160ms，则总的生成时间为：(160+20)*20=3.6s。
假设 Draft Model 总的 Latency 为 0.4s，则总的加速比为：10/(3.6+0.4)=2.5x。

四、方案

4.1 TriForce

本文的作者有几个也是 [2404.11912] TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding 的作者，其实现思路也很类似，不过是将其扩展到了大 Batch Size 场景：

如下图所示为 TriForce 的核心思路，其主要分为三个过程，涉及两级投机采样：

Speculation Phase 1（此阶段可以执行多次，直到获得满足条件的多个候选 Token）：
使用比较小的 LLaMA-68M 模型作为草稿模型，同时使用 StreamingLLM 的方案快速生成多个候选。
使用 LLaMA2-7B-128K 模型作为 Target 模型，采用稀疏 Key、Value Cache 方案进行并行验证。验证后的 Token 作为新的草稿。因为采用的部分 Key、Value Cache，因此相应的结果可能是有损的。
Speculation Phase 2（此阶段执行一次）：
使用 LLaMA2-7B-128K 模型作为 Target 模型，采用全量 Key、Value Cache 方案进行并行验证。此阶段可以保证生成的结果无损。

4.2 本文方案

本文中的方案对 TriForce 进行了简化，只包含两个阶段：

第一阶段：使用 StreamingLLM 进行 Draft 推理，采用固定的 512 个稀疏 Token。本文针对的是中长序列场景，所以 512 个稀疏 Token 相比整个序列要小得多，因此也就可以使用模型自身作为 Draft Model。
第二阶段：使用原始模型进行并行验证。

对于 Draft Model 的选择，作者也提供了两种方案，一种是使用更小的 [2401.02385] TinyLlama: An Open-Source Small Language Model；另一种是使用 Target Model 自己。

五、实验

如下图 Figure 5 所示，作者对比了可变 Batch Size，序列长度下最优投机长度（r）带来的加速比。上图对应 StreamingLLM 的 Budget Token 为 256，下图为 512，图中的数据表示 r。可以看出，LLaMA-2-7B-32K 模型在 Batch Size 32，Seq 长度 32K 时获得 2x 加速；LLaMA-3.1-8B 在 Batch Size 32，Seq 长度 100K 时获得 1.84x 加速。

如下图 Table 1 可以看出，随着序列长度增加，本文方案的加速比愈加明显，而且在同样序列长度时，Batch Size 越大加速越明显。然而，我们也可以发现一个有趣的现象，LLaMA-3.1-8B 在 32K 序列长度的加速比和 LLaMA-2-7B在 8K 训练长度的加速比相当，也就是说，对于 LLaMA-3.1-8B 的加速效果要更弱一些，其主要原因是因为 LLaMA-3.1-8B 使用了 GQA，在同样序列长度时 Memory Bound 的问题要比使用 MHA 的 LLaMA-2-7B 小得多。

如下图 Table 2 所示，在 L40 和 H100 上获得了类似的加速：

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述