连续 Batching:LLM 推理吞吐和延迟的双重优化
关键推理优化技术,显著提升 LLM 服务吞吐和降低响应延迟,生产环境必读。
关键推理优化技术,显著提升 LLM 服务吞吐和降低响应延迟,生产环境必读。
在本文中,我们将介绍大语言模型(LLM)推理的基础知识,并指出传统批处理策略中的低效之处。我们将介绍连续批处理,并讨论 HuggingFace 的 text-generation-inference 和 vLLM 等现有批处理系统的基准测试结果。借助 vLLM,用户可以将 LLM 推理吞吐量提升至 23 倍,同时降低 p50 延迟。由于 LLM 占用大量 GPU 显存且计算成本高昂,对于大多数实际应用而言,服务阶段占据了主要的计算成本。机器学习工程师通常将 LLM 视为“黑箱”,认为只能通过量化和自定义 CUDA 内核等内部改动来优化它们。然而,事实并非完全如此。由于 LLM 以迭代方式生成输出,并且 LLM 推理通常受内存而非计算能力限制,因此存在一些出人意料的系统级批处理优化,能够在实际工作负载中带来 10 倍乃至更大的性能差异。
最近提出的一种此类优化是连续批处理,也称为动态批处理或采用迭代级调度的批处理。我们希望了解这种优化的实际表现。下文将深入介绍相关细节,包括我们如何模拟生产工作负载;不过,先总结一下我们的发现:
使用连续批处理以及针对连续批处理的显存优化(通过 vLLM),吞吐量最高可提升 23 倍。
使用连续批处理以及针对连续批处理的显存优化(通过 vLLM),吞吐量最高可提升 23 倍。
使用连续批处理(在 Ray Serve 和 Hugging Face 的 text-generation-inference 上均如此),吞吐量可达到朴素批处理的 8 倍。
使用连续批处理(在 Ray Serve 和 Hugging Face 的 text-generation-inference 上均如此),吞吐量可达到朴素批处理的 8 倍。
使用经过优化的模型实现(NVIDIA 的 FasterTransformer),吞吐量可达到朴素批处理的 4 倍。
使用经过优化的模型实现(NVIDIA 的 FasterTransformer),吞吐量可达到朴素批处理的 4 倍。
你现在就可以尝试连续批处理:参阅这个示例,在 Ray Serve 上运行 vLLM。
本文其余部分的结构如下:
我们将介绍 LLM 推理的基本工作原理,并指出传统的、基于请求的动态批处理策略中的低效之处。
我们将介绍 LLM 推理的基本工作原理,并指出传统的、基于请求的动态批处理策略中的低效之处。
我们将介绍连续批处理,以及它如何解决基于请求的动态批处理中存在的许多低效问题。
我们将介绍连续批处理,以及它如何解决基于请求的动态批处理中存在的许多低效问题。
接下来,我们将讨论基准测试结果,以及这些结果对于如何经济高效地提供 LLM 模型服务有何影响。
接下来,我们将讨论基准测试结果,以及这些结果对于如何经济高效地提供 LLM 模型服务有何影响。
关于 LLM 推理有很多内容值得了解,如需更多细节,建议读者参阅《单 GPU 上的高效推理》和《优化案例:Bloom 推理》。不过,从较高层面来看,LLM 推理其实相当直观。
首先,从一段 token 序列开始(称为“前缀”或“提示词”)。
首先,从一段 token 序列开始(称为“前缀”或“提示词”)。
LLM 会生成一段补全 token 序列,只有在生成停止 token 或达到最大序列长度后才会停止。
LLM 会生成一段补全 token 序列,只有在生成停止 token 或达到最大序列长度后才会停止。
这是一个迭代过程。模型每进行一次新的前向传播,就会额外生成一个补全 token。例如,假设你的提示词是句子“What is the capital of California: ”,那么需要进行十次前向传播迭代,才能得到完整响应 ["S", "a", "c", "r", “a”, "m", "e", "n", "t", "o"]。这个示例做了一些简化,因为实际上 token 与 ASCII 字符并非一一对应(一种流行的 token 编码技术是字节对编码,即 Byte-Pair Encoding,但这超出了本文的讨论范围)。不过,无论以何种方式对序列进行 token 化,生成过程的迭代性质都是相同的。
现在,我们已经理解了这一迭代过程的简洁原理。接下来,让我们进一步了解一些你可能不知道的 LLM 推理知识:
提示词“What is the capital of California: ”的初始摄入(“预填充”,prefill)所需的时间,与后续生成每个 token 所需的时间大致相同。这是因为预填充阶段会预先计算注意力机制中的一些输入,而这些输入在整个生成过程中保持不变。由于这些输入可以彼此独立地计算,因此预填充阶段能够高效利用 GPU 的并行计算能力。
提示词“What is the capital of California: ”的初始摄入(“预填充”,prefill)所需的时间,与后续生成每个 token 所需的时间大致相同。这是因为预填充阶段会预先计算注意力机制中的一些输入,而这些输入在整个生成过程中保持不变。由于这些输入可以彼此独立地计算,因此预填充阶段能够高效利用 GPU 的并行计算能力。
LLM 推理受内存 I/O 限制,而非受计算能力限制。换句话说,目前将 1MB 数据加载到 GPU 计算核心所需的时间,比这些计算核心对 1MB 数据执行 LLM 计算所需的时间更长。这意味着,LLM 推理吞吐量主要取决于高带宽 GPU 显存中能容纳多大的批次。更多细节请参阅 NVIDIA 文档中的相关页面。
LLM 推理受内存 I/O 限制,而非受计算能力限制。换句话说,目前将 1MB 数据加载到 GPU 计算核心所需的时间,比这些计算核心对 1MB 数据执行 LLM 计算所需的时间更长。这意味着,LLM 推理吞吐量主要取决于高带宽 GPU 显存中能容纳多大的批次。更多细节请参阅 NVIDIA 文档中的相关页面。
GPU 显存消耗量会随着基础模型大小与 token 序列长度之和的增加而增长。《每位 LLM 开发者都应该知道的数字》估算,一个拥有 130 亿参数的模型,会为序列中的每个 token 消耗近 1MB 的状态空间。在配备 40GB 显存的高端 A100 GPU 上,粗略估算可知:存储 26GB 模型参数后还剩 14GB,因此同一时间大约可在显存中容纳 1.4 万个 token。这个数字看似很高,实际限制却相当大:如果将序列长度限制为 512,那么一个批次最多只能处理约 28 个序列。对于更长的序列,问题会更加严重;序列长度为 2048 时,批次大小将被限制为 7 个序列。请注意,这只是一个上限,因为其中没有为中间计算结果预留空间。
GPU 显存消耗量会随着基础模型大小与 token 序列长度之和的增加而增长。《每位 LLM 开发者都应该知道的数字》估算,一个拥有 130 亿参数的模型,会为序列中的每个 token 消耗近 1MB 的状态空间。在配备 40GB 显存的高端 A100 GPU 上,粗略估算可知:存储 26GB 模型参数后还剩 14GB,因此同一时间大约可在显存中容纳 1.4 万个 token。这个数字看似很高,实际限制却相当大:如果将序列长度限制为 512,那么一个批次最多只能处理约 28 个序列。对于更长的序列,问题会更加严重;序列长度为 2048 时,批次大小将被限制为 7 个序列。请注意,这只是一个上限,因为其中没有为中间计算结果预留空间。
这一切意味着,如果能够优化显存使用方式,就还有巨大的优化空间。这正是 AutoGPTQ 等模型量化策略可能如此强大的原因:如果从 16 位表示切换到 8 位表示,将显存使用量减半,就能把可用于更大批次的空间扩大一倍。不过,并非所有策略都需要修改模型权重。例如,FlashAttention 通过重新组织注意力计算、减少内存 I/O,显著提升了吞吐量。
连续批处理是另一种不需要修改模型的显存优化技术。接下来,我们将解释朴素批处理的工作方式及其低效之处,并说明连续批处理如何提高 LLM 生成过程的显存使用效率。
GPU 是大规模并行计算架构,其计算速率(以每秒浮点运算次数,即 flops 衡量)可达到万亿次浮点运算级别(A100),甚至千万亿次浮点运算级别(H100)。尽管拥有如此惊人的计算能力,LLM 仍难以让 GPU 达到饱和利用率,因为芯片的大量内存带宽都被用于加载模型参数。
批处理是改善这种情况的一种方式;与其每次收到输入序列时都重新加载模型参数,不如只加载一次模型参数,然后用它们处理多个输入序列。这样能够更高效地利用芯片的内存带宽,从而提高计算利用率和吞吐量,并降低 LLM 推理成本。
我们将这种传统的批处理方法称为静态批处理,因为在推理完成之前,批次大小始终保持不变。下面是静态批处理在 LLM 推理场景中的示意图:
与传统深度学习模型不同,由于 LLM 推理具有迭代性质,对其进行批处理可能会比较棘手。直观地说,这是因为批次中的某些请求可能会提前“完成”,但释放这些请求占用的资源,并向批次中加入完成状态可能不同的新请求并不容易。这意味着,当一个批次中不同序列的生成长度小于该批次的最大生成长度时,GPU 就会处于利用不足的状态。在上图右侧,这体现为序列 1、3 和 4 的序列结束 token 之后出现的白色方块。
静态批处理多频繁会导致 GPU 利用不足?这取决于一个批次中各序列的生成长度。例如,可以使用 LLM 推理仅输出一个 token,将其作为分类任务(有更好的实现方法,但这里暂且以此为例)。在这种情况下,每个输出序列的长度都相同(1 个 token)。如果输入序列的长度也相同(比如 512 个 token),那么每个静态批次都能达到最佳的 GPU 利用率。
另一方面,由 LLM 驱动的聊天机器人服务既不能假设输入序列长度固定,也不能假设输出序列长度固定。在撰写本文时,专有模型提供的最大上下文长度已经超过 8K 个 token。在静态批处理下,生成输出长度的差异可能导致 GPU 严重利用不足。难怪 OpenAI CEO Sam Altman 将计算成本形容为高得惊人。
如果不对用户输入和模型输出施加限制性假设,未经优化的生产级 LLM 系统根本无法在不造成 GPU 利用不足和不必要高额成本的情况下承载流量。为了让更多人能够广泛使用 LLM 的能力,我们需要优化 LLM 的服务方式。
业界意识到了这种低效,并提出了更好的方法。《Orca: A Distributed Serving System for Transformer-Based Generative Models》是一篇发表于 OSDI ’22 的论文,据我们所知,它首次解决了这个问题。Orca 不会等到批次中的每个序列都完成生成,而是实现了迭代级调度,在每次迭代时确定批次大小。其结果是,一旦批次中的某个序列完成生成,就可以在其位置插入一个新序列,从而获得比静态批处理更高的 GPU 利用率。
现实情况比这个简化模型稍微复杂一些:由于预填充阶段会消耗计算资源,而且其计算模式与生成阶段不同,因此很难将预填充与 token 生成放在一起进行批处理。目前,连续批处理框架通过一个超参数来管理这种情况:waiting_served_ratio,即等待预填充的请求数量与等待序列结束 token 的请求数量之比。
说到框架,Hugging Face 已在其基于 Rust 和 Python 的 text-generation-inference LLM 推理服务器中将连续批处理投入生产。我们使用其实现来了解连续批处理在下文基准测试中的性能特征。
注意:连续批处理、动态批处理和迭代级调度的含义非常接近,因此其中任何一个术语都可以用来描述这种批处理算法。我们选择使用“连续批处理”。“动态批处理”也很贴切,但它可能与请求级批处理混淆;在请求级批处理中,LLM 推理服务器使用静态批次,并在当前批次完全完成生成后才确定下一个批次的大小。我们认为“迭代级调度”很好地描述了调度机制,但不能完整描述整个过程。
在这篇博客文章中,我们希望展示静态批处理与连续批处理之间的差异。事实证明,通过改进 Orca 的设计,连续批处理可以解锁静态批处理无法实现的内存优化。
PagedAttention 是一种在 vLLM(GitHub)中实现的新型注意力机制。它的灵感来自分页和虚拟内存等传统操作系统概念。通过以固定大小的“页”或块来分配内存,它们允许 KV 缓存(即在上文提到的“预填充”阶段中计算的内容)在内存中非连续存放。随后,可以重写注意力机制,使其在按块对齐的输入上运行,从而能够对非连续的内存范围执行注意力计算。
这意味着缓冲区可以按需分配,而不必提前分配:开始新的生成任务时,框架不需要分配一个大小为 maximum_context_length 的连续缓冲区。在每次迭代中,调度器都可以判断某个生成任务是否需要更多空间,并即时进行分配,同时不会对 PagedAttention 的性能造成任何影响。这并不能保证内存得到完美利用(他们的博客称,目前浪费被限制在 4% 以下,并且只发生在最后一个块中),但相比目前业界广泛使用的提前分配方案所造成的浪费,已经有了显著改善。
总体而言,PagedAttention + vLLM 能够大幅节省内存,因为大多数序列不会占满整个上下文窗口。这些内存节省可以直接转化为更大的批次大小,进而意味着更高的吞吐量和更低的服务成本。我们在下文的基准测试中也纳入了 vLLM。
我们将先介绍实验设置,然后深入分析基准测试结果。
我们的目标是观察在模拟真实世界在线推理工作负载的情况下,连续批处理相较于静态批处理的表现。从根本上说,我们关心的是成本。由于成本直接取决于系统在给定延迟下提供服务的效率,因此我们将其拆解为吞吐量和延迟两个方面。
处理包含 1000 个请求的队列所需的时间,其中每个请求包含 512 个输入 token,生成长度从指数分布中采样。
100 个请求各自的请求延迟,这些请求具有不同的输入长度、输出长度和到达时间,并以固定的平均速率到达。
我们将在相应的结果章节中介绍数据集和实验的其他细节。
我们在 Anyscale 提供的单块 NVIDIA A100 GPU 上对吞吐量和延迟进行了基准测试。我们的 A100 配备了 40GB GPU RAM。我们选择 Meta 的 OPT-13B 模型,是因为受测的每个框架都提供了现成的该模型集成。我们选择 13B 版本,是因为它无需张量并行即可装入我们的 GPU,同时规模又足够大,能够带来内存效率方面的挑战。为了保持实验简单,我们选择不使用张量并行——即将每个 Transformer 块拆分到多个 GPU 上——尽管静态批处理和连续批处理都支持张量并行。
我们测试了两个静态批处理框架和三个连续批处理框架。静态批处理框架包括:
Hugging Face 的 Pipelines。这是最简单的推理解决方案。它通过易用的 API 提供静态批处理,适用于任何模型,并且支持的任务不局限于简单的文本生成。我们将其用作基线。
Hugging Face 的 Pipelines。这是最简单的推理解决方案。它通过易用的 API 提供静态批处理,适用于任何模型,并且支持的任务不局限于简单的文本生成。我们将其用作基线。
NVIDIA 的 FasterTransformer。这是一个提供多种 Transformer 模型优化实现的库。目前它只提供静态批处理(Triton 推理服务器提供请求级动态批处理,但尚不支持连续批处理)。借助它,我们可以了解,在静态批处理模式下,对模型进行高度优化的实现究竟能达到怎样的水平——与 Hugging Face Hub 上相对缺乏优化的 OPT-13B 实现相比,它提供了一个更具竞争力的基线。
NVIDIA 的 FasterTransformer。这是一个提供多种 Transformer 模型优化实现的库。目前它只提供静态批处理(Triton 推理服务器提供请求级动态批处理,但尚不支持连续批处理)。借助它,我们可以了解,在静态批处理模式下,对模型进行高度优化的实现究竟能达到怎样的水平——与 Hugging Face Hub 上相对缺乏优化的 OPT-13B 实现相比,它提供了一个更具竞争力的基线。
我们的连续批处理框架包括:
Hugging Face的text-generation-inference。这是Hugging Face用来支持其LLM实时推理API的推理服务器。它实现了连续批处理。
Ray Serve上的连续批处理。Ray Serve利用Ray的无服务器功能来提供无缝自动扩展、高可用性和对复杂DAG的支持。我们想了解连续批处理的工作原理,因此在Ray Serve上用纯Python重新实现了text-generation-inference的核心连续批处理逻辑。正如您将在我们的结果中看到的那样,我们的实现与text-generation-inference实现了相同的性能,这验证了我们的理解。
vLLM。这是UC Berkeley研究人员最近发布的开源项目(GitHub)。它建立在Orca的连续批处理设计基础上,通过对动态内存分配的完全控制,能够显著减少各种形式的GPU内存碎片。我们测试这个框架是因为它展示了通过迭代级调度和连续批处理实现的进一步优化的影响。
基于我们对静态批处理的理解,当每个批次中序列长度的方差较高时,我们预期连续批处理的性能会显著更好。为了展示这一点,我们为每个框架运行四次吞吐量基准测试,每次都在具有更高序列长度方差的数据集上进行。
为此,我们创建了一个包含1000个序列的数据集,每个序列具有512个输入令牌。我们配置模型通过忽略序列结束令牌并配置max_tokens来始终发出每个序列的生成长度。然后我们生成1000个生成长度,每个请求一个,从平均值为128个令牌的指数分布中采样。我们使用指数分布,因为它很好地近似了在服务类似ChatGPT的应用程序时可能遇到的生成长度。为了改变每次运行的方差,我们只选择来自指数分布中小于或等于32、128、512和1536的样本。因此,总输出序列长度最多分别为512+32=544、512+128=640、512+512=1024和512+1536=2048(我们模型的最大序列长度)。
然后我们使用一个简单的asyncio Python基准测试脚本向我们的模型服务器提交HTTP请求。基准测试脚本以突发方式提交所有请求,以便计算饱和。
结果如下:
如预期所示,对于较低方差的生成长度,静态批处理器和朴素连续批处理器的性能大致相同。然而,随着方差增加,朴素静态批处理的性能下降到81个令牌/秒。FasterTransformers相对于朴素静态批处理有显著改进,直到生成长度限制为1536时,几乎与朴素连续批处理器并驾齐驱。Ray Serve上的连续批处理和text-generation-inference实现了大约相同的性能,这是我们预期的,因为它们使用相同的批处理算法。
最令人印象深刻的是vLLM。对于每个数据集,与朴素连续批处理相比,vLLM的性能提高了一倍多。我们还没有分析哪个优化对vLLM性能贡献最大,但我们怀疑vLLM动态预留空间而不是提前预留的能力允许vLLM戏剧性地增加批大小。
我们绘制了这些相对于朴素静态批处理的性能结果:
重要的是要注意FasterTransformer的4倍改进是多么令人印象深刻;当NVIDIA实现它时,我们对FasterTransformers加上连续批处理进行基准测试非常感兴趣。然而,即使使用优化的模型,连续批处理也清楚地比静态批处理是一个重大改进。当您包含由连续批处理和迭代级调度启用的进一步内存优化(如vLLM所做的那样)时,性能差距变得巨大。
实时推理端点经常面临延迟-吞吐量权衡,必须根据用户需求进行优化。我们在现实工作负载上对延迟进行基准测试,并测量延迟的累积分布函数如何随每个框架变化。
与吞吐量基准测试类似,我们配置模型以始终发出每个请求指定的令牌数量。我们通过从1个令牌到512个令牌之间的均匀分布中采样长度来准备100个随机生成的提示。我们从均值为128、最大大小为1536的上限指数分布中采样100个输出长度。这些数字之所以被选择,是因为它们相当现实,并允许生成利用我们模型的完整上下文长度(512+1536=2048)。
与吞吐量基准测试中一次性提交所有请求不同,我们将每个请求延迟预定的秒数。我们采样泊松分布来确定每个请求在前一个请求提交后等待多长时间。泊松分布由λ参数化,即预期速率,在我们的情况下是每秒有多少查询(QPS)命中我们的模型端点。我们测量QPS=1和QPS=4时的延迟,以查看延迟分布如何随负载变化而变化。
我们看到,虽然改进了吞吐量,但连续批处理系