APEX：一种可扩展且具备动态感知能力的模拟器，用于大型语言模型（LLM）服务中的自动化并行执行

《Journal of Parallel and Distributed Computing》：APEX: An Extensible and Dynamism-Aware Simulator for Automated Parallel Execution in LLM Serving

　　我们提出了APEX，这是一个用于优化系统性能的LLM服务模拟器。APEX能够识别出能够最小化延迟和提高能效的系统配置。APEX可以扩展到万亿级别的LLM，并且其模拟速度比现有的模拟器更快。

引言
大型语言模型（LLMs）已成功应用于各种场景，如代码生成[1]、问答[2]等。服务LLM既需要大量的内存也需要强大的计算能力，因此通常需要多设备集群来达到高性能。为了并行化LLM，已经提出了一些技术，包括数据并行性（DP）[4]、流水线并行性（PP）[5]、张量并行性（TP）[6]等。每种并行性都有其优缺点：TP在内存使用上更高效，因为它不像DP那样会产生模型副本，但会增加通信开销[6]；PP可以降低通信开销，但可能会导致工作负载不平衡[5]。混合并行性可以有效地平衡这些权衡以提高性能，但由于涉及的因素众多（包括计算工作负载、网络流量、内存效率等），确定最优的并行执行计划并不容易。此外，最优的并行执行计划还取决于输入请求的特性。计算密集型任务（例如需要较长提示的任务，如摘要生成）和内存密集型任务（例如需要长时间生成结果的任务，如代码编写）适合不同的并行策略。一种常见的做法是采用启发式执行计划，例如在同一节点内使用TP，在不同节点之间使用PP[7]、[8]。然而，最近的研究表明，启发式计划的执行速度可能比最优配置慢2倍[9]，这强调了服务提供者需要寻找并采用最优的并行执行计划，而不是仅仅依赖启发式方法。

直接找到最优计划的方法是彻底地部署和评估所有可能的并行方案；然而，这种方法成本非常高，因为可能需要数千小时的GPU运行时间[10]、[11]。另一种方法是利用基于建模的方法[9]，例如模拟器[12]、[13]，来估计并行执行计划的性能；但是，为深度神经网络[12]、[13]或LLM训练[9]、[14]开发的现有解决方案无法直接应用于LLM服务，因为服务系统引入了独特的挑战：（1）迭代级分批处理的动态性：与使用静态分批处理的传统ML系统不同，在静态系统中必须完成一个批次中的所有请求后才能添加新的请求，而LLM服务系统[10]、[15]、[16]采用迭代级分批处理[17]。在这种方法中，随着内存的释放，新请求会不断被添加到处理批次中，导致批次大小动态变化，这很难建模。此外，批次中的某些请求可能处于预填充阶段，而其他请求则处于解码阶段。这种阶段的交错进一步复杂化了建模过程（详见第2.1节）。（2）设计空间的指数级增长：随着模型大小和集群中设备数量的增加，设计空间会呈指数级增长。由于LLM服务通常涉及大型模型和集群，因此在合理的时间内开发出能够搜索到最优并行执行计划的解决方案并不容易。（3）适应不断演变的系统：LLM服务领域是一个新兴领域，模型架构、硬件平台以及系统优化（如量化[18]、[19]和并行技术[20]）在不断演变。因此，性能模型或模拟器可能会很快过时，不再适用于最新的LLM服务系统。因此，跟上这些系统的快速发展既必要又具有挑战性。

为此，我们提出了APEX，这是一个可扩展且具有动态感知能力的LLM服务自动并行执行模拟器。给定一个LLM、一个设备集群以及具有不同上下文/生成长度的输入请求，APEX通过评估多种候选的并行执行计划来识别出高性能的方案。APEX首先生成多种并行执行计划，每种计划都代表了通过组合不同的并行性（即DP、PP、TP）来并行化模型的独特方式。然后，APEX通过模拟来估计每种计划的执行时间，从而进行动态感知的仿真。仿真后，APEX会为每种执行计划提供全面的评估结果，包括每输出一个令牌所需的时间（TPOT）、到达第一个令牌的时间（TTFT）和能耗等关键指标[21]。这种详细的评估使服务提供者能够选择针对不同目标（例如最小化延迟或能耗）进行优化的执行计划，同时满足服务级别目标（SLOs）。

尽管进行了复杂的仿真，但由于利用了Transformer模型的重复结构，APEX仍然具有很高的可扩展性，显著减少了设计空间，并能够对多节点集群中的万亿级别模型进行仿真。为了模拟最先进的LLM服务系统，APEX支持多种LLM、并行性、量化格式和设备集群。考虑到LLM服务的快速发展，APEX被设计为模块化且可扩展的，可以轻松集成新的模型和硬件而无需进行大量工作。此外，APEX利用操作级性能分析来估计系统执行时间，从而能够根据底层平台收集的性能数据来适应不同的集群。我们使用多种LLM、设备集群和不同的工作负载对APEX进行了评估。APEX始终能够识别出性能优异的执行计划，其性能优于启发式计划，同时保持了较低的仿真开销。我们的主要贡献包括：
- 自动并行执行：我们提出了APEX，这是一个能够自动为给定的LLM、设备集群和输入请求找到高性能并行执行计划的模拟器。
- 动态感知：APEX能够模拟迭代级分批处理，捕捉实际的LLM服务行为。
- 准确且高效：APEX提供了准确的加速预测，平均相对误差为10.7%。APEX识别的执行计划比启发式计划快高达3.37倍，并且节能计划最多可减少45%的能耗。
- 广泛的兼容性和可扩展性：APEX支持多种LLM、并行性、量化格式和设备集群，并且易于扩展到新的配置。
- 成本效益高且可扩展：APEX在CPU上仅需15分钟即可识别出最优计划，比基于云的GPU部署快71倍，成本效益高1234倍；同时，与现有的LLM服务模拟器相比，速度快8.4倍。此外，当从十亿级别模型扩展到万亿级别模型时，APEX的仿真开销保持不变。

**LLM推理**
大型语言模型（LLM）推理以用户提供的提示作为输入，并逐个令牌生成响应。这种顺序的令牌生成过程也称为自回归生成[22]，它利用先前生成的令牌作为输入来预测下一个令牌。LLM的推理过程包含两个阶段：
1. 预填充阶段：在预填充阶段，LLM处理输入请求以设置用于预测第一个令牌的中间状态（键和值）。

**系统概述**
我们在图2中展示了APEX的工作流程。用户提供的输入数据以黄色框表示，内置系统数据用红色圆柱体表示。给定一个LLM和一个设备集群，APEX使用内置的并行模板生成多个并行执行计划。每个生成的执行计划代表了将LLM部署到集群中的不同策略。分批模块将这些生成的计划和请求跟踪作为输入，动态生成执行方案。

**实验设置**
为了评估APEX，我们选择了多种不同工作负载的LLM和数据集。
- 模型：我们选择了三种最先进的LLM：Llama-3.1-70B-Instruct [34]、Mistral-Large-Instruct (123B) [37] 和 Mixtral 8x22B-Instruct [38]；这些模型涵盖了不同规模的LLM，还包括一个混合专家（MoE）模型。虽然这些模型默认使用半精度（FP16），但我们也测试了将模型量化为FP8格式的情况。
- 数据集：我们准备了三组请求跟踪数据。

**讨论**
**局限性**
APEX通过建模LLM服务中的主要计算、通信和内存成本（包括注意力机制、GEMM、集体通信、模型权重和KV缓存）来进行粗粒度并行执行规划。这种抽象方法对于评估并行性之间的权衡非常有效。然而，APEX并没有明确地模拟Transformer中的每个内核。尽管像RMSNorm和嵌入这样的内核单独来看可能很小，但它们的累积开销可能会影响整体性能。

**相关工作**
已经提出了几种LLM服务系统[10]、[15]、[16]、[17]、[24]、[29]、[49]。每个系统都引入了创新来解决LLM服务中的关键问题。Orca [17]引入了迭代级分批处理，大大提高了服务吞吐量。DistServe和Splitwise [16]、[29]在不同的设备上分离了预填充和解码阶段，实现了特定阶段的优化。MegaScale-Infer [31]将这一想法进一步细化，实现了更细粒度的分离。

**结论**
我们提出了APEX，这是一个具有动态感知能力的LLM服务模拟器。我们在多种LLM和不同的工作负载下评估了APEX的性能。APEX在其评估的范围内识别出了性能最佳的延迟计划，与启发式计划相比，速度提高了多达3.37倍，并识别出了节能计划，与仅优化延迟的计划相比，能耗最多减少了45%。APEX的速度提高了71倍，成本效益比实际GPU部署高1234倍。

**关于作者声明**
在准备本文的过程中，作者们使用了ChatGPT来润色文本并检查语法错误。使用该工具/服务后，作者根据需要对内容进行了审查和编辑，并对发表文章的内容承担全部责任。

**作者贡献声明**
- 林一坚（Yi-Chien Lin）：负责写作 – 审稿与编辑、撰写原始草稿、软件和 methodology。
- 罗纳德·皮内达（Ronald Pineda）：负责写作 – 审稿与编辑、验证和 methodology。
- 范妮·尼娜·帕拉韦西诺（Fanny Nina Paravecino）：负责写作 – 审稿与编辑、软件、资源管理、methodology以及概念化。

**利益冲突声明**
作者声明了以下可能被视为潜在利益冲突的财务关系/个人关系：
- 林一坚与Meta公司存在雇佣关系。
- 罗纳德·皮内达与微软公司存在雇佣关系。
- 范妮·尼娜·帕拉韦西诺与微软公司存在雇佣关系。
林一坚曾任职于南加州大学。
罗纳德·皮内达
林一坚
范妮·尼娜·帕拉韦西诺