How continuous batching enables 23x throughput in LLM inference while reducing p50 latency

本文是<How continuous batching enables 23x throughput in LLM inference while reducing p50 latency>的翻译简化版,详细的可以看原文

前言

LLM inference

Simplified LLM inference. This toy example shows a hypothetical model which supports a maximum sequence length of 8 tokens (T1, T2, …, T8). Starting from the prompt tokens (yellow), the iterative process generates a single token at a time (blue). Once the model generates an end-of-sequence token (red), the generation loop stops. This example shows a batch of only one input sequence, so the batch size is 1.

现在我们已经了解了迭代过程的简单性,让我们深入探讨一些关于LLM推理的你可能不知道的事情:

  • The initial ingestion (“prefill”) of “What is the capital of California: ”所需的时间与生成每个后续token所需的时间大致相同。这是因为预填充阶段对注意机制的某些输入进行了预计算,这些输入在生成过程中保持不变。由于这些输入可以独立计算,预填充阶段有效地利用了GPU并行计算能力。

  • LLM推理受内存IO限制而非计算限制。换句话说,将1MB数据加载到GPU计算核心上所需的时间比该计算核心执行1MB数据上的LLM运算所需时间更长。这意味着LLM推理吞吐量主要取决于我们能够将多大批次大小适应高带宽GPU内存中。

  • 消耗的GPU内存量随基础模型大小和token序列长度而增加。一个具有130亿参数模型会为序列中每个token消耗近1MB状态空间。在具有40GB RAM 的高端A100 GPU上,粗略计算表明,在存储了26GB模型参数后还剩下14GB,可以一次性将大约14k个令牌保存在内存中。这可能看起来很高,但实际上相当有限;如果我们将序列长度限制为512,则最多可以处理约28个序列的批次。对于更长的序列长度问题更加严重;序列长度为2048意味着批次大小被限制为7个序列。请注意,这只是一个上界估计,因为它没有留出空间来存储中间计算结果。

所有这些意味着如果我们能够优化内存使用情况,就会有相当大的“潜力”。这就是诸如AutoGPTQ等模型量化策略可能非常强大的原因;如果能够通过从16位转换到8位表示减少一半的内存使用量,则可以将可用空间翻倍以容纳更大的批次大小。然而,并不是所有策略都需要修改模型权重。例如,FlashAttention通过重新组织注意力计算以减少内存IO要求找到了显著提高吞吐量的方法。

连续分批是另一种无需修改模型即可进行内存优化的技术。接下来我们将解释简单分批工作方式(及其效率低下之处),以及连续分批如何提高LLM生成的内存效率。

静态batch

Completing four sequences using static batching. On the first iteration (left), each sequence generates one token (blue) from the prompt tokens (yellow). After several iterations (right), the completed sequences each have different sizes because each emits their end-of-sequence-token (red) at different iterations. Even though sequence 3 finished after two iterations, static batching means that the GPU will be underutilized until the last sequence in the batch finishes generation (in this example, sequence 2 after six iterations).

Continuous batching

Orca: A Distributed Serving System for Transformer-Based Generative Models是OSDI '22上提出的一篇论文,是首次解决这个问题的尝试。与其等待批处理中的每个序列都完成生成,Orca实现了iteration-level scheduling,其中批大小是根据每次迭代确定的。结果是一旦批处理中的一个序列完成生成,就可以插入新的序列来替换它,从而获得比静态分批更高的GPU利用率。

Completing seven sequences using continuous batching. Left shows the batch after a single iteration, right shows the batch after several iterations. Once a sequence emits an end-of-sequence token, we insert a new sequence in its place (i.e. sequences S5, S6, and S7). This achieves higher GPU utilization since the GPU does not wait for all sequences to complete before starting a new one.

其实吧,Continuous batching, dynamic batching, and iteration-level scheduling在意义上都挺像的,可以用其中任何一个来描述当前说的这个算法,不过这里为啥用Continuous batching呢?主要是dynamic batching容易和request-level batching搞混了。而iteration-level scheduling也只是对调度机制做了描述,但不能代表整个过程。

request-level batching指在推理服务器中使用固定大小的批次,无论请求的数量如何,每个批次的大小都是预先设定的。这种方法的优点是可以简化资源管理和调度,但缺点是可能无法充分利用系统资源,特别是在请求量较小的情况下,而动态批处理是一种更灵活的批处理方法。它允许推理服务器根据当前的请求量动态地调整批次的大小

PagedAttention and vLLM

continuous batching可以解锁静态批处理无法实现的内存优化。

PagedAttention是在vLLM(GitHub)中实现的一种新的注意力机制。它从传统的操作系统概念中汲取灵感,如分页和虚拟内存。它们允许KV缓存(在上述“预填充”阶段中计算的内容)通过分配固定大小的“页面”或块来实现非连续性。然后,注意力机制可以被重写为对块对齐的输入进行操作,允许对非连续的内存范围进行注意力操作。

这意味着缓冲区分配可以在需要的时候进行,而不是提前进行:当开始新的生成时,框架不需要分配一个大小为maximum_context_length的连续缓冲区。每次迭代,调度器可以决定是否需要为特定的生成提供更多的空间,这个分配过程很轻(on the fly),而不会对PagedAttention的性能产生任何影响。这并不能保证完美的内存利用(他们的博客说,现在的浪费限制在最后一个块的4%以下),但它显著地改善了目前广泛使用的提前分配方案的浪费。

总的来说,PagedAttention + vLLM使得大多数序列不会消耗整个上下文窗口,从而实现了大量的内存节省。这些内存节省直接转化为更高的批处理大小,这意味着更高的吞吐量和更便宜的服务。我们在下面的基准测试中包含了vLLM。

Inflight batching doesn’t always provide an improvement in throughput and/or latency. It really depends on your dataset. Datasets with uniform input and output lengths are less likely to benefit from inflight batching.

参考