通过 Amazon SageMaker 提升 Falcon 模型的性能 机器学习博客

使用 Amazon SageMaker 提升 Falcon 模型性能

关键要点：

在这篇文章中，我们讨论了如何通过 Amazon SageMaker 和 LMI 容器来提升 Falcon 模型的性能。我们介绍了动态批处理和连续批处理也称为滚动批处理的技术，以优化推理的吞吐量和延迟。此外，我们还通过基准测试结果对比了不同批处理的效果。

随着生成式 AI 应用程序需求的增加，如何使用最优框架和配置来托管大型语言模型LLMs以提高文本生成性能，变得尤为重要。虽然有多个选项用于服务 LLM，但由于模型规模、架构和性能要求的多样性，要找到最佳方案并不容易。Amazon SageMaker 的大型模型推理LMI容器通过整合多种不同的框架和优化技术，使得部署 LLM 变得简单。LMI 容器具有强大的服务堆栈 DJL Serving，能够根据特定 LLM 进行系统级配置，提取最佳性能，并支持连续批处理这样的新优化，显著提高了吞吐量。

在之前的一篇文章中，我们展示了如何使用 LMI 容器在 SageMaker 上部署 Falcon 系列模型。本篇文章将进一步介绍如何通过连续批处理等技术提高 Falcon40B 的吞吐量和延迟，同时解释 SageMaker LMI 容器提供的配置参数，以帮助您找到适合实际应用的最佳配置。

LLM 文本生成推理的基本原理

首先，让我们来看一下如何进行 LLM 文本生成推理的一些基本原理。

前向传播、激活与KV缓存

给定一个输入的 tokens 序列，它们会通过 LLM如 Falcon的所有层进行 前向传播，以生成下一个 token。前向传播 指的是将输入数据传递通过神经网络以产生输出的过程。在文本生成的情况下，前向传播涉及将初始的种子或上下文输入语言模型，并生成序列中的下一个字符或 token。为了生成文本序列，这一过程通常是迭代进行的，也就是说，针对输出序列中的每一步或位置重复此过程。每次迭代，模型生成下一个字符或 token，该 token 成为生成文本的一部分，直到达到所需的文本长度。

例如，像 Falcon 或 GPT 这类语言模型的文本生成是 自回归 的。这意味着模型一次生成一个 token，并基于之前生成的 tokens 进行条件预测。在自回归解码过程中，所有输入 tokens 会生成它们的注意力 Key 和 Value 张量，这些张量会保存在 GPU 内存中以生成下一个 tokens。这些缓存的 Key 和 Value 张量通常被称为 KV 缓存。

填充与解码阶段

在自回归解码过程中，通常有两个主要阶段：填充 阶段和 解码 阶段。这两个阶段对生成连贯的、上下文相关的文本至关重要。

填充阶段包括以下内容：

填充阶段持续进行，直到满足预设的停止条件。此条件可以是生成文本的最大长度、特定 token标记文本结束的信号或用户或应用程序设定的其他标准。

解码阶段包括以下内容：

填充与解码阶段的结合使自回归模型能够生成建立在初始上下文上的文本，生成连贯、上下文相关且一致的文本序列。

有关该过程的详细说明，可以参考 分布式 Transformer 基于生成模型的服务系统。

使用动态批处理优化吞吐量

迄今为止，我们只讨论了单个输入。在实践中，我们预计会遇到来自应用客户端的多个请求随机地并发或错开到达。传统方法中，基本批处理可以用于增加吞吐量和提升 GPU 计算资源的利用率。批处理将多个请求的数值表示组合在一起并进行自回归前向传播的并行运行。智能批处理在服务端进行。SageMaker LMI 的 DJL Serving 服务器可以通过在 servingproperties 中设置以下参数来配置多个请求的批处理：

这表明 DJL Serving 会在一个时间点上排队请求100毫秒，或者当排队的请求数量达到指定的 batchsize 时，批处理会被调度到后端进行推理。这被称为 动态批处理。它是动态的，因为批处理大小可以根据在该时间段内添加的请求数变化。然而，由于请求可能具有不同的特征例如，有些请求可能是 20 个输入 token 和 500 个输出 token，而其他请求可能正好相反，500 个输入 token 只有 20 个输出 token，批处理中的某些请求可能会比其他请求更快完成处理。这可能导致在等待同一批中的所有任务完成解码阶段时 GPU 的利用效率降低，即使在队列中还有其他请求等待处理。下图说明了这一过程。

使用连续批处理优化吞吐量

通过使用 连续批处理也称为 迭代 或 滚动 批处理，我们可以利用填充与解码阶段之间的区别。为了启用连续批处理，DJ Serving 根据 servingproperties 提供以下额外配置：

使用连续批处理时，服务堆栈DJL Serving不会等待批处理中的所有请求完成解码阶段。相反，在解码阶段的逻辑中断即结束一个迭代时，它会拉入排队的额外请求，而当前批处理仍在处理。因此，我们在每次迭代的解码阶段结束时检查是否有待处理请求。请记住，对于每个请求，我们需要先执行填充阶段，然后再顺序进行解码阶段。因为我们可以并行处理请求的初始提示中的所有 tokens，所以当有新请求被拉入时，我们会暂时暂停正在进行的批处理解码阶段我们会暂时保存其 KV 缓存和激活，并执行新请求的填充阶段。

该缓存的大小可以使用以下选项配置：

填充完成后，我们会将新请求和旧暂停请求组合成一个新的滚动批处理，它们可以并行进行解码阶段。请注意，旧的暂停请求可以继续其解码阶段，而新请求将从其首个新的 token 开始。

您可能已经意识到，连续批处理是一种非常类似于我们在日常生活中自然并行化任务的方法。我们时常会在随机的时间接收到信息、电子邮件或电话通知可以看作是新请求。同时，我们依旧在完成当前的任务例如撰写电子邮件、编码、参与会议类似于当前 GPU 正在处理的任务。在逻辑上中断时，我们暂停正在进行的任务并检查通知，以决定是否需要采取行动，如果需要，我们将其添加到正在进行的任务中现实生活中的滚动批处理，或者将其列入待办事项列表队列。

综上所述：如何思考 GPU 内存的利用

建议您对模型进行负载测试，以了解哪种配置对您的业务用例最具成本效益。为了建立这种理解，让我们对 GPU 的内存占用情况有个可视化的概念，观察模型加载及其在滚动批处理中的后续请求处理。本文假设我们正在将 Falcon40B 模型加载到安装了 NVIDIA A10G GPU每个 GPU 具有 24 GB 内存的某些 G5 实例类型上。请注意，对 p3、p4 和 p5 实例类型的类似理解同样适用，因为它们分别配备 V100、A100 和 H100 GPU 系列。

以下是估算服务 Falcon40B 所需总内存的概述：

对于 Falcon40B，如果我们将模型量化压缩为 bfloat162 字节数据类型，则模型大小约为 80 GB。如您所见，这仍然超过单个加速器设备所支持的内存，因此我们需要采用模型分割分片技术，以及特殊的 张量并行性TP方法，将模型分散到多个加速器设备上。假设我们选择 g524xlarge，它配备 4 个 A10G GPU。如果我们在 DJL Servingservingproperties中配置以下内容，预计 80 GB 的模型权重将均匀分配到所有 4 个 GPU 上：

将 tensorparalleldegree 设置为 4，约 20 GB 的 24 GB GPU 内存约 84在处理单个请求前已被占用。剩余的 16 GPU 将用于后续请求的 KV 缓存。有可能，在您的业务场景以及延迟和吞吐量要求下，剩余内存的 23 GB 已绰绰有余。如果不足，则可以将实例大小提升至 g548xlarge，该实例具有 8 个 GPU，并将 tensorparalleldegree 设置为 8。在这种情况下，每个 GPU 大约只会使用 10 GB 的 24 GB 内存用于模型权重，剩下约 60 的 GPU 内存用于激活和 KV 缓存。直观上我们可以看到，这种配置可能会使我们具有更高的吞吐量。此外，由于我们现在有了更大的缓冲区，我们可以增加 maxrollingbatchprefilltokens 和 maxrollingbatchsize 参数，以进一步优化吞吐量。这两个参数一起控制激活填充和 KV 缓存的预分配。对于这两个参数的更大数值，假设您在 GPU 内存中有足够的 KV 缓存缓冲，通常会与更高的吞吐量相关联。

使用 PagedAttention 进行连续批处理

PagedAttention 是 UC Berkeley 开发的一项新优化算法，通过在固定大小的页面或块中分配内存，使注意力缓存KV 缓存不再是连续的，从而改进了连续批处理过程。这也受到操作系统使用的虚拟内存和分页概念的启发。

根据 vLLM 论文，每个 token 序列的注意力缓存被分成块，并通过块表映射到物理块。在计算注意力时，PagedAttention 内核可以使用块表有效地从物理内存中获取块。这显著减少了内存浪费，并使得批量大小增大、GPU 利用率提升以及吞吐量倍增。

性能比较