Vidur: A Large-Scale Simulation Framework For LLM Inference

发表时间: 2024-05 · MLSys 2024 · arXiv:2405.05465 (Georgia Tech, Microsoft Research India)

原文: https://arxiv.org/abs/2405.05465

Amey Agrawal 1 2 Nitin Kedia 3 Jayashree Mohan 3 Ashish Panwar 3 Nipun Kwatra 3 Bhargav S. Gulavani 3 Ramachandran Ramjee 3 Alexey Tumanov 1
1 Georgia Institute of Technology, 2 Microsoft Research, 3 Microsoft

速读

一句话结论 本文提出了一个名为 Vidur 的大规模、高保真大语言模型推理性能模拟框架，并配套了配置搜索工具 Vidur-Search，能够自动且极低成本地为不同模型和工作负载找到满足性能约束的最优部署配置。

要解决什么问题 原有做法在寻找大语言模型推理的最佳部署配置时，需要将不同的并行策略、调度策略、批处理技术和硬件组合起来在真实环境中进行穷举测试。由于最佳配置是模型和具体工作负载的共同函数，其搜索空间的复杂度达到了 $O(|M| \cdot |T|)$，随着模型和应用数量的增加，这种基于物理部署的探索成本高得令人难以承受，且一旦配置错误可能导致高达两倍的成本差异。而现有的深度学习模拟框架主要针对训练任务设计，无法直接迁移到推理场景，主要卡在三个机制上：首先是时间尺度，推理对延迟极度敏感，单次迭代仅需几毫秒，需要比训练更细粒度的运行时间预测；其次是迭代时间波动极大，推理包含预填充和解码两个阶段，且请求的序列长度和动态批大小不断变化；最后是误差级联问题，推理请求是动态到达的，单次批处理运行时间的微小预测误差会改变后续的批处理模式，随着时间推移导致聚合误差急剧放大。

怎么做的 核心思路是结合极简的算子级实验分析与机器学习预测模型，构建一个事件驱动的端到端推理模拟器。为了绕开庞大配置空间带来的测试成本，Vidur 利用了绝大多数大语言模型架构高度相似的特性，通过通用的声明式规范提取计算算子，并将其严格划分为三类以定制预测策略。第一类是标记级算子（如线性层和激活函数），其运行时间仅取决于批次中的总标记数；第二类是序列级算子（如注意力机制），其运行时间受上下文长度影响。针对预填充阶段的注意力计算，假设批次中有 $P$ 个长度为 $p_i$ 的请求，其计算成本与长度的平方和成正比，因此系统会预测一个等效单请求长度 $\sqrt{\Sigma_{i=1}^P p_i^2}$ 的运行时间，而解码阶段的注意力计算则被建模为受限于从 KV-Cache 获取数据量的访存密集型操作；第三类是通信算子，仅依赖传输数据量，可脱离具体模型独立分析。Vidur 的关键设计由三个部件构成。首先是分析器，它自动识别张量分片配置，在单张显卡上收集最少的算子运行时间特征。其次是运行时间估计器，它使用收集到的稀疏数据训练随机森林回归模型，从而在端到端模拟中对未见过的输入维度进行高保真的运行时间插值预测。最后是三层分层调度器：全局调度器负责请求路由；副本调度器负责内存规划和批处理，内置了 vLLM、Orca、FasterTransformers、SarathiServe 和 LightLLM 等主流批处理策略的内存管理逻辑；副本阶段调度器则处理流水线并行中的微批次调度。在此基础上，配套的 Vidur-Search 工具通过二分查找算法在模拟器中快速遍历搜索空间，以最大化每美元的每秒查询数为优化目标，自动筛选出满足首次标记时间和标记间时间约束的最佳硬件与系统参数组合。

效果如何 实验在 4 张 80GB A100 或 H100 显卡组成的集群上进行，底层软件基于 vLLM 优化版，测试了 LLaMA2 7B 与 70B、InternLM-20B 和 Qwen-72B 模型，工作负载涵盖了自然语言对话、论文摘要和双语平行语料。对比基线是传统的基于真实物理部署的穷举搜索路线。量化结果显示，在静态工作负载下，Vidur 对请求执行时间的预测误差（即使是 95 分位数的尾部延迟）最高仅为 3.33%；在动态工作负载且请求速率达到系统最大承载能力 85% 的高压场景下，预测误差依然保持在 5% 以内。在 LLaMA2-70B 的部署优化任务中，传统的物理部署探索路线需要耗费 42000 个 GPU 小时，成本高达 218000 美元；而 Vidur-Search 仅需在一台普通 CPU 机器上运行 1 小时即可找到最优配置。在更大规模的假设分析中，执行 35565 次模拟运行仅需 12.5 小时和 125 美元。作者也承认了该方法的局限性：当系统处于极限负载（达到最大容量的 95%）且运行 LLaMA2-7B 这种较小模型时，由于 CPU 的调度开销占比变大，预测误差会发生级联放大，最大误差会上升至 12.65%。

1 主要贡献

优化大型语言模型（LLMs）的部署目前非常昂贵，因为这需要通过实验在大量的配置空间（由并行化策略、批处理技术、调度策略等系统参数构成）中运行应用程序工作负载来寻找最佳配置。为了解决这一挑战，本文提出了Vidur，一个用于LLM推理性能的大规模、高保真、易于扩展的模拟框架。Vidur结合了实验分析和预测建模来模拟LLM算子的性能，并通过估计延迟和吞吐量等关键指标来评估不同工作负载的端到端推理性能。我们在多个LLMs上验证了Vidur的保真度，结果表明它在各个范围内的推理延迟估计误差小于9%。此外，我们还提出了Vidur-Search，这是一个配置搜索工具，可帮助优化LLM的部署。Vidur-Search利用Vidur自动识别在满足应用性能约束条件下最具成本效益的部署配置。例如，Vidur-Search在一个CPU机器上仅需一小时就能找到LLaMA2-70B的最佳部署配置，而基于部署的探索则需要耗费42,000个GPU小时，成本高达218,000美元。Vidur的源代码已在 https://github.com/microsoft/vidur 开源。

2 背景知识/关键Observation/设计原则

LLM概述
LLMs利用基于自注意力机制的Transformer架构作为其核心构建块。自注意力机制帮助语言模型学习输入序列中不同元素之间的关系，并随后生成输出序列。LLM由两个主要的子模块组成：自注意力（self-attention）和多层感知机（MLP）。近年来，利用这些模块的变体开发了各种LLMs（如GPTs、LLaMAs、Falcons）。主要地，这些模型仅在嵌入大小、Transformer块的数量以及模型使用的注意力机制方面有所不同。

LLM推理效率优化
LLM推理请求处理包括两个不同的阶段：预填充（prefill）和解码（decode）。预填充阶段处理整个用户输入提示并生成第一个输出标记。随后，以自回归方式一次生成一个输出标记。在此解码阶段，前一步骤生成的标记通过模型传递以生成下一个标记，直到生成特殊的序列结束标记，此时请求处理完成。解码过程需要访问先前处理标记的键（key）和值（value）激活以执行注意力操作。为了避免重复计算，当代LLM推理系统将它们存储在KV-Cache中。鉴于LLM推理的巨大成本，LLM推理效率已成为系统研究的活跃领域。为此，最近提出了多种优化机制。每种技术都做出不同的权衡。对于具有成本效益的推理，应根据具体的应用需求组合适当的优化集。例如，张量并行（TP）是并行化LLM推理的常见策略。TP通过在参与的GPU上平均拆分模型权重和KV-Cache，在这些GPU上分片每一层。另一方面，流水线并行（PP）是另一种并行化策略，其中模型被划分为连续Transformer块的阶段。最近，研究人员确定了LLM推理调度器设计中固有的权衡，并提出了将现有LLM推理调度器分为两类：预填充优先和解码优先。SarathiServe试图通过利用解码阶段的计算松弛来缓解这种权衡。另一组最近的工作，Splitwise和DistServe，通过在不同的设备上分离预填充和解码的计算来解决这种延迟-吞吐量权衡。

LLM推理配置空间
诸如并行策略、调度器选择、块大小、批大小、SKU等控制旋钮为LLM部署引发了巨大的配置空间。此外，我们观察到最佳配置（定义为每个控制旋钮的特定选择组合）不仅仅是特定模型的函数。相反，最佳配置是模型$m$和在该模型上评估的轨迹$t$的函数。因此，配置搜索的复杂性为$O(|M| \cdot |T|)$，其中$M$是所有感兴趣模型的集合，$T$是工作负载的集合。随着模型和下游应用数量的迅速增加，寻找最佳配置的成本根本无法扩展。然而，配置错误是极其昂贵的。例如，图1b显示，在一个轨迹上使用另一个轨迹的最佳配置可能会有高达$2\times$的成本差异。

模拟LLM推理的挑战
现有的DNN模拟框架主要关注训练作业。构建大规模推理模拟器，特别是针对LLMs，涉及现有模拟器未解决的多个挑战：
- 时间尺度：传统的DNN训练工作负载通常是计算密集型工作负载，每次迭代执行数百毫秒。相比之下，LLM推理是一项对延迟更敏感的任务，其迭代时间可能短得多（每次几毫秒）。因此，模拟LLM推理需要在更细的粒度上预测迭代时间。
- 变化的迭代时间：与传统的深度学习工作负载相比，在LLM推理期间，不同迭代的延迟可能会有很大差异。推理运行时间的变化有多个来源：推理包含不同的阶段（预填充和解码），请求的序列长度差异很大，且在线推理期间批大小不断变化。
- 级联错误：在训练工作负载中，批处理组成在所有批次中都是统一的，并且每个批次的执行是独立的。然而，在推理期间，请求动态到达系统，如果任何批次的运行时间预测有重大错误，这可能会改变批处理模式。因此，单个批次预测中的小错误会随时间推移而级联，导致聚合错误。

3 方法细节

4 VIDUR

Vidur利用领域知识提供LLM推理的高保真性能估计。它模拟了推理堆栈所有层的行为，包括模型执行和各种级别的请求调度，在副本以及集群级别。

关键见解
- LLMs共享关键架构属性：绝大多数LLMs共享基本相似的架构，在激活函数、归一化层、残差连接等的选择上存在细微差别。这使我们能够使用通用的声明式模型规范格式来捕获各种模型的基本架构选择。这种架构统一性的另一个结果是，Vidur只需要对所有模型家族共享的少量计算算子进行建模。
- 用于运行时间预测的操作分类：在运行的批次中，每个请求可能与不同数量的KV-Cache和查询标记相关联，从而导致巨大的组合输入空间。因此，分析每种可能的组合以预测操作运行时间是不可行的。相反，我们观察到LLM算子可以分为不同的类别。例如，某些操作的执行时间取决于批次中所有请求的总上下文长度，而对于其他操作，则仅取决于当前迭代中的标记数量。这种分类使我们能够为每种算子类型设计定制的运行时间预测策略。
- 并行化策略的自动分析：每个模型并行配置具有不同的内存、计算和网络通信特征。朴素的分析和重放方法将需要为每个并行配置进行单独的分析运行，这可能是昂贵的。相比之下，Vidur结合了有关LLM并行化策略的领域知识，这使其能够识别在每个设备上执行的计算子集。在分析阶段，我们自动从模型的声明式规范中识别每个算子的张量分片配置。因此，Vidur可以在单个GPU上执行最少分析的情况下模拟各种并行化方案。

系统概述
Vidur主要有两个处理阶段。首先是模型加载阶段，其中使用模型规范生成一组要分析的计算算子。Vidur分析器收集已识别算子的运行时间特征，并将它们提供给运行时间估计器。为了最大限度地降低向系统添加新模型的成本障碍，我们在分析阶段收集最少的数据，然后训练小型机器学习模型，以在模拟期间可能触发这些操作的大范围参数上生成预测。此阶段由Vidur的运行时间估计器处理，它生成可稍后在模拟期间使用的按操作分类的运行时间查找表。一旦模型加载完毕，用户就可以在Vidur-Bench支持的各种工作负载上，使用各种调度策略和并行化策略执行模拟。在我们的事件驱动模拟器的核心是一个可插拔的分层调度器，它支持几种流行的批处理策略以及内存规划和管理功能。模拟器提供了详细的指标，捕获请求（归一化延迟、首次标记时间、标记间时间等）和集群（模型FLOPs利用率、KV-Cache利用率等）性能指标。

分析器
为了有效地分析LLMs的运行时间特征，我们利用了绝大多数LLMs共享基本相似架构的见解。
- 算子分类：分析器分析不同的算子以识别它们的输入依赖性。我们发现所有算子都可以放入三个桶之一：
- 标记级算子：线性操作和激活函数等操作的操作数维度取决于模型架构，但它们的运行时间仅取决于批次中正在处理的总标记数（预填充加解码）。
- 序列级算子：注意力操作不仅取决于当前批次中的标记数，还取决于每个请求的上下文长度。
- 通信算子：all-reduce和all-gather等通信操作的运行时间仅取决于要传输的数据量，与模型架构无关。

• 分析标记级算子：标记级算子有两大类 - 矩阵乘法和简单的逐点应用或归约操作，如加法、归一化和激活函数。基于模型规范，我们生成所有不同的张量并行分片配置，并分析每种组合。这种方法使我们能够在单个GPU上分析时获取不同并行配置的轨迹。我们使用标准的PyTorch内核来分析这些操作，并使用CUPTI测量它们的性能。
• 分析序列级算子：批处理序列级算子（如注意力内核）对批次中请求的上下文长度很敏感，从而使要分析的输入状态空间爆炸。我们使用几种技术来解决这个问题。首先，由于它们在计算特征上的差异，我们分别分析预填充和解码阶段的注意力内核。在处理预填充注意力时，我们观察到每个预填充的注意力时间与其长度成二次方。假设我们有一批长度为$p_i$的$P$个预填充，其中$i$从1到$P$变化。因此，整个批次的预填充注意力成本与$\Sigma_{i=1}^P p_i^2$成正比。为了近似此批次的运行时间，我们预测等效批次的单个预填充长度$\sqrt{\Sigma_{i=1}^P p_i^2}$的运行时间。与预填充相反，我们注意到注意力解码操作在很大程度上受内存限制。因此，此操作的运行时间主要取决于需要从KV-Cache中获取的总数据量，而不是批次中不同请求之间上下文长度的精确划分。
• 分析通信算子：LLM推理中经常使用三种集体操作，即all-reduce、all-gather（用于张量并行）和send-recv（用于流水线并行）。由于这些操作不依赖于模型特定的特征，因此我们在提前以与模型无关的方式独立分析不同拓扑的这些内核。

运行时间估计器
为所有算子的每种可能输入组合收集分析数据是极其昂贵的。因此，我们收集一组有限的数据点，并依靠小型机器学习模型来插值运行时间。运行时间估计器首先使用分析数据训练这些模型，然后为它在端到端模拟中遇到的大范围输入张量维度生成运行时间估计。我们发现随机森林（RF）回归模型在数据节俭性和保真度之间实现了适当的平衡。

分层调度器
在Vidur中，我们采用三层分层调度器架构。
- 全局调度器：负责Vidur中的请求路由。除了标准的负载平衡策略（如轮询和最少未完成请求）外，我们还支持有状态的调度策略，其中路由决策可以推迟到以后的时间点。
- 副本调度器：封装了两个关键职责；批处理和内存管理。副本调度器包含一个内存规划器，它使用模型规范和并行配置来计算KV-Cache可用的内存。此信息随后由内存管理器用于提供高级管理API，这些API用于实现自定义批处理策略。Vidur目前支持五种批处理策略：FasterTransformers、Orca、SarathiServe、vLLM和LightLLM。
- 副本阶段调度器：处理流水线阶段内微批次的调度。虽然我们目前仅支持同步流水线并行调度策略，但在未来，我们旨在扩展副本阶段调度器以模拟各种优化，如异步通信、序列并行和推测流水线解码。

5 VIDUR-BENCH

Vidur-Bench是一个用于轻松评估LLM推理系统性能的基准套件。

数据集和工作负载
LLM推理的整体性能对工作负载的类型高度敏感，例如给定查询中的输入和输出标记的数量。Vidur-Bench提供了一组从公开可用的数据集中策划的工作负载（见表1）。这些可用于评估不同请求类型、到达率等的系统性能，或微调服务系统中各种组件的性能敏感参数。

性能指标
Vidur-Bench提供了一套全面的系统级性能指标：
- 算子级指标：这包括每个算子的输入大小和执行时间。
- 请求级指标：这些包括每个请求的指标，如调度延迟、预填充完成时间、首次标记时间（TTFT）和标记间时间（TBT）。
- 副本级指标：这些包括批大小、每次迭代处理的标记数、繁忙和空闲时间以及每个副本的内存和计算利用率。
- 硬件指标：这些捕获集群范围的GPU FLOPs和内存利用率。

6 VIDUR-SEARCH

Vidur-Search是一个帮助系统操作员在满足所需SLO约束的同时找到部署推理系统的最佳成本配置的工具。Vidur-Search利用我们的模拟器以有效的方式计算最佳配置。

Vidur-Search具有以下主要组件：
- 输入：搜索工具的输入包括LLM模型、工作负载、可用的GPU SKU以及副本中的最大GPU数量。
- 约束：对TTFT和TBT等指标的SLO。
- 搜索空间：搜索工具有权配置并行策略、并行度、调度策略、特定于调度器的参数、批大小、GPU SKU的选择等。
- 优化目标：Vidur-Search帮助操作员最大化每美元的QPS。

Vidur-Search通过解决约束优化问题来找到搜索空间中的最佳配置。Vidur-Search首先枚举系统的所有可能部署配置。对于每种配置，我们在指定的QPS下对输入工作负载运行模拟器，并预测TTFT和TBT等指标。我们通过简单的二分搜索来找到系统支持的最大QPS，该搜索寻找不会将调度延迟增加超过阈值的最大QPS。二分搜索的每一步都涉及为相应的配置和QPS运行我们的模拟器。我们在单独的内核上运行每个搜索来并行化这些运行。

4 实验环境

• 数据集：LMSys-Chat-1M（自然语言对话）、Arxiv-Summarization（arXiv论文摘要）、Bilingual-Web-Book（文档级英汉平行数据集）。
• 模型：LLaMA2 7/70B、InternLM-20B、Qwen-72B。
• 硬件配置：Azure Standard NC96ads A100 v4 VMs（4个NVIDIA 80GB A100 GPU，NVLink连接）、H100 VMs（4个NVIDIA 80GB H100 GPU，NVLink连接）。
• 软件配置：基于vLLM的优化版本，支持不同调度策略和CUDA图。

5 实验结果

• 模拟器保真度（静态工作负载）：Vidur预测甚至尾部延迟（P95）的误差率在四种模型和三个数据集上高达3.33%（图3）。
• 模拟器保真度（动态工作负载）：在几乎所有场景中，将请求速率设置为系统容量的85%时，Vidur都实现了高保真度（<5%误差）（图4）。
• 假设分析：工作负载的变化会极大地改变最佳配置；具有相似大小的模型由于架构细节的变化可能具有非常不同的性能特征；每美元的容量遵循预期的趋势（图6）。
• 配置稳定性：为一种工作负载使用最佳配置来服务另一种工作负载可能会导致非常高的开销（图1b）。
• 帕累托前沿分析：在一个指标上最佳的配置可能不满足另一指标上的SLO约束；延迟SLO的微小变化可能导致显著的成本开销（图5）。

6 结论

LLM推理效率取决于大量配置旋钮，如并行类型或程度、调度策略、GPU SKU。在实际硬件上运行所有可能的配置是不切实际的。本文提出了Vidur：一个用于LLM推理的高保真且易于扩展的模拟器，以及一个基准测试和搜索套件。Vidur回答了与部署相关的假设问题，这些问题确定了生产环境的有效部署策略，并有助于以名义成本评估各种系统优化的功效。

7 附录

A.1 动态工作负载请求到达率对保真度的影响

图7展示了Vidur在不同请求到达率下的额外保真度结果。我们发现，对于较大的模型，即使在最大系统容量的95%时，Vidur也能保持高保真度。然而，对于LLaMA2-7B，由于CPU开销导致的误差略高，误差会级联，最大误差达到12.65%。图8还提供了误差趋势。

A.2 假设分析的成本细分

Vidur能够高效且准确地模拟复杂的部署场景，使我们能够以极低的成本探索搜索配置空间。图1a中展示的假设分析需要进行35,565次运行，预计GPU总时长成本为1,139,865美元。同样的搜索在一台96核CPU机器上仅需约12.5小时完成，成本仅为125美元。表2列出了各项任务的细分情况。