什么是托管大型语言模型 (LLMs) 用于文本生成生成 AI 应用的最佳框架和配置？尽管有很多用于服务 LLMs 的选项，但由于模型的大小、不同的模型架构、应用的性能要求等原因，这个问题很难回答。Amazon SageMaker 的大型模型推断 (LMI) 容器通过将多种不同的框架和技术整合到一起，优化了 LLMs 的部署，使其变得简单易用。LMI 容器具有一个称为 DJL serving 的强大服务堆栈，它对底层的 LLM 是不可知的。它提供了可以针对给定的 LLM 调整的系统级配置参数，以提取托管基础设施的最佳性能。它还支持最近的优化，如连续批处理，也称为迭代批处理或滚动批处理，这在吞吐量方面提供了显著的改进。
在之前的一篇文章中，我们展示了如何使用 LMI 容器在 SageMaker 上部署 Falcon 系列模型。在本文中，我们演示了如何通过连续批处理等技术来提高 Falcon-40B 的吞吐量和延迟。我们还提供了 SageMaker LMI 容器提供的配置参数的直观理解，可以帮助您找到适合实际应用的最佳配置。
LLM 文本生成推断的基础知识
首先，让我们来看一下有关如何执行 LLM 文本生成的一些基础知识。
前向传播、激活和 KV 缓存
给定一系列标记的输入序列，它们在 LLM (如 Falcon) 的所有层上进行前向传播，以生成下一个标记。前向传播是指将输入数据通过神经网络以产生输出的过程。在文本生成的情况下，前向传播涉及将初始种子或上下文输入到语言模型中，并生成序列中的下一个字符或标记。要生成一系列文本，通常是迭代进行的，这意味着它会针对输出序列的每个步骤或位置重复进行。在每次迭代中，模型生成下一个字符或标记，并将其作为生成文本的一部分。此过程持续进行，直到生成所需长度的文本。
像 Falcon 或 GPT 这样的语言模型的文本生成是自回归的。这意味着模型一次生成一个标记，同时根据先前生成的标记进行调整。换句话说，每次迭代，模型都会将先前生成的文本作为输入，并根据该上下文预测下一个标记。如在 vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention 中提到的，在这个自回归解码过程中，LLM 的所有输入标记都会生成它们的注意力键和值张量，并将这些张量保留在 GPU 内存中以生成下一个标记。这些缓存的键和值张量通常被称为 KV 缓存。
填充和解码阶段
在自回归解码过程中，例如在使用 Falcon 等语言模型进行文本生成时，通常有两个主要阶段：填充阶段和解码阶段。这些阶段对于生成连贯和语境相关的文本至关重要。
填充阶段包括以下内容：
- 初始上下文 – 填充阶段从用户提供的初始上下文或种子文本开始。这个初始上下文可以是一个句子、一个短语，甚至只是一个单词。它设置了文本生成的起点，并为接下来的内容提供了上下文。
- 模型调整 – 使用提供的上下文来调整语言模型。模型将此上下文作为输入，并根据对上下文的理解生成序列中的下一个标记（单词或字符）。
- 标记生成 – 模型逐个生成一个标记，预测文本中接下来应该出现的内容。该标记会附加到上下文中，从而将其扩展。
- 迭代过程 – 生成标记的过程是迭代进行的。在每个步骤中，模型都会生成一个标记，同时考虑到更新的上下文，其中现在包括在之前步骤中生成的标记。
预填充阶段持续进行直到满足预定的停止条件。该条件可以是生成文本的最大长度、标志文本结束的特定标记，或用户或应用程序设置的任何其他标准。
解码阶段包括以下内容：
- 完成 – 在预填充阶段之后，您有一个部分生成的文本，可能是不完整的或在某个位置被截断的。解码阶段负责将文本完成，使其连贯且符合语法。
- 从上一标记继续 – 在解码阶段，模型从预填充阶段生成的最后一个标记开始。它将此标记作为初始上下文，并生成下一个标记以继续文本。
- 迭代完成 – 与预填充阶段一样，生成标记的过程再次是迭代的。模型一次生成一个标记，基于序列中的前一个标记进行条件化。
- 停止条件 – 解码阶段也具有停止条件，可能与预填充阶段相同，例如达到最大长度或遇到文本结束标记。满足此条件时，生成过程停止。
预填充和解码阶段的组合使得自回归模型能够生成以初始上下文为基础并产生连贯、有关上下文和前后一致的文本序列。
详细说明，请参考基于Transformer的生成模型的分布式服务系统。
使用动态批处理优化吞吐量
到目前为止，我们只讨论了单个输入。在实践中，我们预计会随机地处理来自应用程序客户端的多个请求，并行或交错地进行推理。传统方法中，可以使用基本批处理来增加GPU的吞吐量和计算资源的利用率。批处理实际上是将多个请求的数值表示组合成一个批次，并进行自回归的前向传递的并行运行。此智能批处理是在服务端完成的。SageMaker LMI的DJLServing服务器可以通过在serving.properties中设置以下参数来将多个请求批处理在一起以并行处理：
- max_batch_delay = 100 – 批处理聚合的最大延迟（毫秒）。默认值为100毫秒。
- batch_size = 32 – 动态批处理大小。默认值为1。
这基本上显示了DJLServing将每次排队等待100毫秒的请求或者如果排队请求的数量达到指定的batch_size，批次将被调度到后端进行推理。这被称为动态批处理。它是动态的，因为批处理大小可能会根据在该时间段内添加的请求数而改变。然而，由于请求可能具有不同的特征（例如，某些请求的输入形状为20个标记，输出为500个标记，而其他请求可能相反，输入为500个标记，但只有20个输出），同一批次中的某些请求可能处理速度更快。尽管队列中还有其他请求等待处理，但这可能导致GPU的利用率不足，因为等待批处理中的所有正在处理的请求完成解码阶段，例如，等待所有正在推理的请求完成。下图说明了这个过程。
动态批处理可视化 – 注意请求2和3的空闲窗口
使用连续批处理优化吞吐量
通过使用连续批处理，也称为迭代或滚动批处理，我们利用预加载和解码阶段之间的差异。为了激活连续批处理，DJServing根据serving.properties提供以下附加配置：
- engine=MPI – 我们建议您使用MPI引擎进行连续批处理。
- option.rolling_batch=auto或lmi-dist – 我们建议使用auto，因为它将自动选择最合适的滚动批处理算法，并在未来进行其他优化。
- option.max_rolling_batch_size=32 – 这限制了并发请求的数量。默认值为32。
使用连续批处理，服务堆栈（DJLServing）不会等待批处理中所有正在进行中的请求完成解码阶段。相反，在逻辑中断处（解码阶段的一次迭代结束时），它会拉取等待队列中的其他请求，而当前批处理仍在处理中（因此称为滚动批处理）。它在每个解码阶段迭代结束时检查待处理请求的情况。请记住，对于每个请求，我们需要运行预加载阶段，然后是顺序解码阶段。因为我们可以并行处理请求的初始提示中的所有标记，所以每当有新请求被拉取时，我们会临时暂停批处理中正在进行的请求的解码阶段 – 我们会暂时保存其内存中的KV缓存和激活状态，并运行新请求的预加载阶段。
此缓存的大小可以使用以下选项进行配置：
- option.max_rolling_batch_prefill_tokens=1024 – 限制滚动批处理（在解码和预加载阶段之间）中缓存中保存的同时预加载标记的数量
当预加载完成后，我们将新请求和旧的暂停的请求合并到一个新的滚动批处理中，这些请求可以并行进行解码阶段。请注意，旧的暂停请求可以继续进行它们中断的解码阶段，而新请求将从它们的第一个新标记开始。
连续或迭代批处理可视化 – 注意空闲时间被后续请求替代
您可能已经意识到，连续批处理几乎是一种我们在日常生活中自然并行化任务的方法。我们有随机时间到达的消息、电子邮件、电话通知（潜在的新请求）（类似于随机交错方式到达多个请求以供GPU处理）。当我们继续进行正在进行的任务 – 撰写电子邮件、编码、参加会议（与GPU中当前正在处理的任务类似）时，这一切都在同时进行。在逻辑中断点处，我们暂停正在进行的任务，并检查通知，以确定是否需要我们采取某些行动，如果需要，我们将其添加到正在进行的任务中（现实生活中的滚动批处理），或将其放入待办事项列表（队列）。
把一切都放在一起：如何思考GPU的内存利用
建议对模型进行负载测试，以确定哪种配置对您的业务用例最具成本效益。为了建立理解，让我们将GPU的内存占用作为模型加载时和连续请求在滚动批处理中处理时的内存占用进行可视化。在本文中，我们假设我们正在将Falcon-40B模型加载到安装有24 GB内存的NVIDIA A10G GPU的G5实例类型中的一个实例上。请注意，类似的理解适用于配备V100、A100和H100 GPU系列的p3、p4和p5实例类型。
以下是获取为Falcon-40B提供服务所需的总内存的近似值的概述：
- 模型大小 = 模型参数数量（Falcon-40B有400亿个参数）x 每个参数的字节数（FP32为4个字节）= 160 GB
- 加载Falcon-40B用于推理所需的总内存的近似值 = 模型大小（=160 GB）+ KV Cache（Attention Cache）（=*20 GB）+ 机器学习框架的其他内存开销（约2 GB）
内存可视化 – 了解已加载的Falcon-40B模型的内存占用
对于Falcon-40B，如果我们通过量化模型到 bfloat16（2字节）数据类型来压缩模型，模型大小约为80 GB。正如您所见，这仍然大于单个加速器设备支持的内存，因此我们需要采用模型分区（分片）技术与特殊的张量并行（TP）方法，并将模型划分到多个加速器设备上。假设我们选择了g5.24xlarge，它具有4个A10G GPU设备。如果我们将DJLServing（serving.properties）配置为以下内容，可以期望80 GB的模型权重将平均分配给所有4个GPU：
- option.tensor_parallel_degree%0Aoption.-,tensor_parallel_degree,-%3D2%0A%23%20specify%20per) = 4或8，或使用max（检测到的最大GPU数）
将tensor_parallel_degree设置为4，即使在处理单个请求之前，24 GB GPU内存的约80%（约20 GB）已被使用。剩余的16%的GPU将用于接收请求的KV缓存。对于您的业务场景及其延迟和吞吐量要求，剩余内存中的2-3 GB可能已经足够。如果不够，您可以将实例大小增加到g5.48xlarge，它具有8个GPU，并将tensor_parallel_degree设置为8。在这种情况下，每个GPU的可用24 GB内存中仅使用约10 GB用于模型权重，我们还剩下约60%的GPU用于激活和KV缓存。直观地说，我们可以看到这种配置可能允许我们实现更高的吞吐量。此外，由于我们现在有一个更大的缓冲区，我们可以增加max_rolling_batch_prefill_tokens和max_rolling_batch_size参数以进一步优化吞吐量。这两个参数将控制模型的激活预填充和KV缓存的预分配。这两个参数设置较大的值将与较大的吞吐量相关联，前提是您在GPU内存中有足够的缓冲区用于KV缓存。
使用PagedAttention进行连续批处理
PagedAttention是由加州大学伯克利分校开发的一种新的优化算法，通过允许注意力缓存 (KV缓存) 在内存中非连续分布并分配固定大小的页面或块来改进连续批处理过程。这种方法受到了操作系统使用的虚拟内存和分页概念的启发。
根据vLLM论文，每个令牌序列的注意力缓存被分成块，并通过一个块表映射到物理块。在计算注意力时，PagedAttention内核可以使用块表从物理内存高效地提取块。这减少了内存浪费，允许更大的批量大小、增加GPU利用率和提高吞吐量。
性能比较
为确保对部署配置进行有效的负载测试，建议首先考虑业务场景，并清楚定义基于LLM的应用程序的输入和输出特性。例如，如果您正在处理一个呼叫中心摘要的使用案例，输入可能包含较大的文本，如客户服务代理和客户之间的500个令牌的聊天记录，但输出可能相对较小，大约为100个令牌，表示对聊天记录的摘要。另一方面，如果您正在处理代码生成场景，输入可能只有15个令牌，如“编写一个高效的Python实现，用于描述包括分页在内的所有EC2资源”，但输出可能更大，达到500个令牌。此外，还要考虑在您的特定场景中，是否追求更低的延迟还是最大化吞吐量。
在全面了解业务场景之后，您可以分析并确定您的托管环境的最佳配置。在这个上下文中，托管环境包括各种关键元素，包括实例类型和其他配置参数，如tensor_parallel_degree、max_rolling_batch_size、max_rolling_batch_prefill_tokens等等。我们的目标是确定最有效的设置来支持响应时间、吞吐量和模型输出质量的需求。
在我们的分析中，我们通过动态批处理和迭代批处理，使用SageMaker上的LMI容器中的serving.properties中详细列出的配置，对性能进行了基准测试，展示了连续批处理相对传统动态批处理的优势。
| 动态批处理 | 连续批处理 | 使用PagedAttention的连续批处理 |
|
engine=Python option.model_id=tiiuae/falcon-40b option.tensor_parallel_degree=8 option.dtype=fp16 batch_size=4 max_batch_delay=100 option.trust_remote_code=true |
engine=MPI option.model_id={{s3_url}} option.trust_remote_code=true option.tensor_parallel_degree=8 option.max_rolling_batch_size=32 option.rolling_batch=auto option.dtype=fp16 option.max_rolling_batch_prefill_tokens=1024 option.paged_attention=False |
engine=MPI option.model_id={{s3_url}} option.trust_remote_code=true option.tensor_parallel_degree=8 option.max_rolling_batch_size=32 option.rolling_batch=auto option.dtype=fp16 option.max_rolling_batch_prefill_tokens=1024 option.paged_attention=True |
使用FP16数据类型在ml.g5.48xlarge上部署的Falcon-40B的两种配置在几个不同的代表实际应用的场景中进行了基准测试：
- 大量生成的输入令牌和生成的令牌数量较少的情况 – 在该场景下，输入令牌的数量固定为32个，并生成了128个新令牌
| 批处理策略 | 吞吐量（令牌/秒） | 延迟 p90（秒） |
| 动态批处理 | 5.53 | 58.34 |
| 连续批处理 | 56.04 | 4.74 |
| 连续批处理 with PagedAttention | 59.18 | 4.76 |
- 大量输入和生成令牌数量较少的情况 – 在这里，我们将输入令牌的数量固定为256个，并提示LLM将输入总结为32个令牌
| 批处理策略 | 吞吐量（令牌/秒） | 延迟 p90（秒） |
| 动态批处理 | 19.96 | 59.31 |
| 连续批处理 | 46.69 | 3.88 |
| 连续批处理 with PagedAttention | 44.75 | 2.67 |
我们可以看到，在使用SageMaker上的LMI容器时，连续批处理与PagedAttention相比，场景1的吞吐量提高了10倍，场景2提高了2.3倍。
结论
在本文中，我们介绍了LLMs如何使用内存，并解释了连续批处理如何通过在SageMaker上使用LMI容器提高吞吐量。我们通过显示基准测试结果，展示了对使用LMI容器在SageMaker上的Falcon-40B进行连续批处理的好处。您可以在GitHub存储库中找到代码。