[论文解读]FlexLLM: A System for Co-Serving Large Language Model Inference and Finetuning

FlexLLM: A System for Co-Serving Large Language Model Inference and Parameter-Efficient Finetuning

这篇论文是CMU Zhihao Jia课题组在2024年发表于archive，对于研究inference和finetuning的协同调度非常有参考意义。但是该论文是建立在OSDI’22 Unity这篇论文基础上的，需要很强的前置知识，很多地方还没有看懂，在此先记录一下。

摘要

PEFT是一种被广泛应用于不同任务的LLM技术。云服务提供商通常给用户创建分离的系统运行PEFT模型的finetuning和inference任务。这是因为现有的系统不能很好地处理finetuning和inference请求混合的负载。因此，共享的GPU资源通常利用率低下。为了解决这个问题，作者提出了FlexLLM，这是第一个在同一个iteration中服务inference和finetuning请求的系统。我们的系统利用这两个任务的互补性，并利用共享的GPU资源共同运行它们，使用一种称为co-serving的方法。为了实现这一点，FlexLLM引入了一种新的token级 finetuning机制，它将sequence的微调计算分解为更小的token级计算，并使用依赖并行化和图剪枝，这两种静态编译优化，以最小化co-serving的内存开销和延迟。与现有系统相比，FlexLLM的co-serving方法将激活GPU内存开销减少了高达8倍，并将finetuning的端到端GPU显存需求降低了高达36%，同时保持低推理延迟并提高finetuning吞吐量。例如，在重推理工作负载下，FlexLLM仍然可以保持超过80%的峰值finetuning吞吐量，而现有系统则无法在finetuning上取得任何进展。FlexLLM源代码 https://github.com/flexflow/FlexFlow/。

2.背景与挑战

2.1 PEFT

PEFT可以用于训练特定的LLM，但是它也需要很大的计算开销。现有的方法集中解决减少训练参数量的同时优化模型的性能，但是，减少训练参数量并不能减少相应的显存占用。

2.2 PEFT-based LLM inference

现有的大型语言模型（LLM）推理系统中的一个关键挑战是内存密集型的自回归解码机制，因为生成每个词符都需要访问模型的所有参数。这导致了对计算资源的高需求，尤其是在需要快速生成大量文本时，这种机制可能会成为性能瓶颈。为了克服内存带宽限制，已经提出了几种技术，包括量化、并行化、批处理和内核优化。但是，显存贷款的问题仍然存在，特别是当latency和response质量需要同时保证的时候。
近期的研究工作引入了多种系统，用于同时服务多个基于PEFT的大型语言模型（LLM）。PetS是第一个统一的多任务参数高效变换器（PETs）服务框架，它将计算图中的线性层解耦为任务无关的共享操作和任务特定的PET操作，然后使用一个高性能的推理引擎和创新的操作调度策略来提高并发性和吞吐量。虽然PetS仅支持编码器模型，但近期的工作也集中在对生成性推理模型的优化上。PUNICA 是一个在共享GPU集群中服务多租户LoRA的系统。它提出了一个新的CUDA内核，称为分段收集矩阵-向量乘法，它允许批处理GPU操作以并发执行多个不同的LoRA模型。S-LoRA在PUNICA引入的内核基础上进行了优化。它提出了统一分页、异构批处理和一种新颖的张量并行算法，以在单台机器上服务数千个LoRA组件。然而，所有这些系统也像传统的LLM服务系统一样受到内存访问的限制。

3. Co-serving LLM inference and PEFT

已经有很多研究提出不同的方法来调度异构的任务（in this case, inference and PEFT）来共享GPU。这些方法可以被分为三类：资源隔离，时分共享，和空分共享。该章节主要讨论这些方法应用到inference和PEFT上时的limitations，以及重点突出FlexLLM与现有工作的差异。Fig3 描述了不同调度方法在执行inference 和finetuning时的GPU执行时间线。

资源隔离。 不同的任务分别在不同的GPU上执行，如图3(a)。资源隔离满足了finetuning任务不会影响到inference任务，而且用户可以一次提交所有需要训练的数据集，这样可以使用更大的batch size来训练。另一方面，用户可以独立且随机地地提交inference请求。但是这种方式的缺点是GPU资源利用率低。Finetuning和inference请求不能batch到一起运行。而且，当GPU资源有限时，inference请求并不一定能meet SLO。

时分共享。 与资源隔离相比，时间共享方法可以通过允许推理和微调任务根据各种参数通过不同长度的时间片共享GPU资源来增加GPU的利用率。这种技术允许推理和微调请求轮流访问集群中的所有GPU，如图3(b)所示。使用所有GPU来并行化计算（例如，采用数据并行或张量模型并行）可以帮助减少每个请求的延迟，这反过来又可以帮助满足推理SLO的限制（例如，在图3(a)中推理请求r3无法完成，但在图3(b)中可以通过利用更多资源来完成）。然而，时间片的安排也可能因为推理请求的不可预测到达时间而变得有害，这可能导致无法满足的严格要求的延迟（例如，在图3(b)中，由于预先安排的微调请求，r4已被延迟）。

空分共享。
影响资源隔离和时间共享的一个问题是它们没有考虑到计算密集型微调任务和内存密集型推理任务的不同特征。这导致两种工作负载都严重浪费资源。一个替代解决方案是空间共享，它允许同时运行不同的任务，每个任务占用GPU资源的一部分（即，流处理器）。例如，可以通过编写多流代码或利用多进程服务（MPS）[4]，在NVIDIA GPU上进行空间共享。此外，在最新的Hopper和Ampere架构上，空间共享可以利用多实例GPU（MIG）[3]来动态调整分区配置。然而，这种调整并不足够灵活或高效[2, 19, 69]，以适应动态的推理工作负载。图3©和(d)展示了两个不同的情况，其中将更大的GPU资源份额分配给微调或推理。尽管为推理任务过度配置可能看起来能提供更好的服务等级目标（SLO）保证，但在遇到像r3这样的意外和紧急请求时仍然会失败。此外，过度配置还会浪费内存带宽资源，并显著减慢微调过程。因此，空间共享不能依赖于在同时服务来自在线工作负载的推理请求时，以令人满意的性能运行微调请求。

FlexLLM’s co-serving方法。 如图3(e)所示，FlexLLM采用co-serving方法，使用细粒度的调度机制来适应性地结合推理和微调请求。在每次co-serving迭代中，FlexLLM通过以下方式联合服务推理和微调请求：(1) 在预填充阶段计算推理请求的键和值，(2) 在解码阶段为推理请求生成一个新词符，以及(3) 对每个微调请求的一组词符执行前向和后向传递。FlexLLM将推理kernel与微调前向kernel融合，以最小化kernel启动开销，并减少访问GPU设备内存以检索模型权重的次数。在每次迭代中，FlexLLM动态调整推理和微调tokens的分配，策略性地平衡满足推理请求的服务等级目标（SLO）的需求与最大化GPU利用率的目标。这种自适应方法确保FlexLLM遵守推理请求的要求，同时优化可用GPU资源的使用。

4.系统概述

图4展示了FlexLLM的概览，它为LLM co-serving PEFT和inference请求。用户可以选择PEFT model hub的模型执行微调或推理服务，该hub存储了backbone LLM和所有微调后的模型。推理和微调请求的编程接口都通过PEFT即服务（PEFTaaS）接口统一。FlexLLM包含三个主要系统组件。首先，对于在多个GPU上并行执行的微调请求，静态编译模块（第5节）生成一个并行计算图。该图指定了在分布式环境中执行PEFT模型，并通过对内存节省进行不必要的张量剪枝来优化图。其次，动态调度模块（第6节）采用了token级微调机制。它采用混合调度策略混合推理和微调tokens。最后，计算图和调度计划由FlexLLM的分布式共同服务运行时（第7节）执行。

5.静态编译

当一个PEFT模型注册后，FlexLLM将它编译为一个并行计算图PCG。为了优化在多个GPU上的模型的finetuning，该系统通过应用图转换到PCG生成最佳的并行策略。两个关键的静态优化是：依赖并行化和图剪枝。依赖并行化允许多个来自于同一个backbone LLM架构的PEFT模型共享backbone模型参数。该系统基于backbone LLM模型的并行化对每个PEFT都产生一个最优的并行规划。图剪枝是另一个优化技术，减少那些需要在前向传播过程保持到后向传播完成的中间激活tensors。这个技术帮助进一步减少PEFTmodel的显存需求。

5.1 依赖并行化

为了最小化内存开销，FlexLLM强制所有微调和推理任务共享基础大型语言模型（LLM）的权重和中间激活，这要求PEFT模型的并行化策略必须与基础LLM的并行化策略兼容，我们称这项任务为依赖并行化。

并行计算图（PCGs）。PCGs利用四种并行化原语（即，partition，conbine，replicate和reduce）共同捕捉与不同并行化策略相关的通信模式、数据移动和计算分布。PCG中的每个张量由一组数据维度表示，每个维度包括两个字段：大小和度数。度数字段指示张量沿该维度被分成的分区数。每个张量还包括一个领先的副本维度，指定张量的副本数量。副本度数大于1表示张量被复制到多个设备上。

并行计算图（PCGs）使用四种并行化原语来表示与不同并行化策略相关的计算和通信成本。首先，分区（partition）和组合（combine）改变了张量的并行度，并且是彼此的“反向操作”。具体来说，分区通过将维度等分地分割成多个分区，从而增加了张量维度的并行度，而组合则通过将多个分区合并成一个，从而减少了张量的并行度。其次，复制（replicate）和归约（reduce）通过复制和求和来更新张量的副本数量。关于PCG原语的更全面介绍可以在Unity论文中找到。

Unity通过迭代地对输入的、未并行化的计算图应用图替换来生成PCGs，这允许Unity通过应用的图替换序列隐式地维护张量数据维度的内部状态。然而，Unity的方法需要从一个完全未并行化的计算图开始，并且不适用于依赖并行化，其中基础模型已经根据用户提供的规范进行了并行化。

为了解决这个问题，FlexLLM通过为张量的每个维度引入一个状态字段来扩展PCG（并行计算图）的表示。每个张量维度可以处于以下四种状态之一，如图6所示。首先，非并行（non-parallel）状态表示张量在该维度上没有并行化（即，并行度为1）。其次，分区（partitioned）和复制（replicated）状态表示张量通过分割和复制原始张量沿给定维度进行分区。最后，预归约（pre-reduce）状态指的是不同副本的张量各自持有部分结果，通过对这些副本进行归约可以恢复张量的最终输出。当使用张量模型并行化（tensor model parallelism）来并行化大型语言模型（LLMs）时，这种情况很常见（参见图5(a)中并行化基础模型的示例）。

依赖并行化问题。给定多个输入和输出并行张量及其数据维度的并行状态，FlexLLM发现一个PCG（并行计算图），通过将输入张量作为输入，产生具有给定并行度和状态的输出张量。与现有自动并行化方法解决的并行化优化问题不同，FlexLLM解决了一个具有额外约束的优化问题。图5(a)显示了使用张量模型并行化来并行化两个线性层的PCG。每个绿色（或灰色）框表示一个张量代数（或并行化）操作符，操作符之间的每条边代表一个并行张量，其维度显示在并行张量旁边。在这个例子中，所有并行张量都有三个维度，其中第一个表示副本维度，其余两个对应于矩阵的行和列。我们考虑将LoRA应用于第一个线性操作符的依赖并行化任务，其输入和输出张量分别处于[-, -, -]和[-,-,|]状态。对于输出张量，最终的分区状态表明输出矩阵的列跨GPU进行了分割。

FlexLLM利用输入和输出张量的并行状态为LoRA网络生成PCG。具体来说，FlexLLM在张量代数操作符之间插入并行化操作符，并利用每个张量数据维度的并行状态来类型检查这些并行化操作符的兼容性。图6显示了并行化操作符如何转换张量的并行状态。大多数PEFT模型只包含逐元素操作符，它们不改变张量的并行状态，以及线性操作符，它们通过遵循图5(b)给出的规则更新输入张量的并行状态。

为了为旁路网络生成PCG，FlexLLM在所有张量代数操作符之间插入一个并行化操作符，枚举并行化操作符的类型，并检查新PCG中所有张量的兼容性。通过这项检查的PCG被认为是依赖并行化任务的候选PCG。图5©显示了FlexLLM为并行化LoRA网络发现的四个候选PCG。这些候选PCG引入了不同的并行化操作符和通信开销。FlexLLM重用了Unity的基于分析的成本模型[64]，并选择了成本最低的候选PCG。在图5©中的四个候选PCG中，第一个实现了最低成本，因为它在前向或后向过程中没有引入通信。请注意，输入张量后的复制操作符不引入成本，因为它可以与基础模型的融合。其他两个并行化操作符没有引入通信成本，因为它们与输入独立，可以预先处理。

5.2 图剪枝

本节介绍了FlexLLM的图剪枝算法，该算法以PEFT模型及其backbone LLM作为输入，并返回为微调给定旁路网络所必须保留的最小中间激活集合。图剪枝需要推理每个操作符的输入和输出张量来剪枝节点和边，因此FlexLLM为PCG表示引入了以下符号。设G = (N, E)表示一个PEFT模型的PCG，其中每个节点n ∈ N是一个张量代数（或并行化）操作符，每条边e = (n1, n2) ∈ E是操作符之间共享的张量。对于每个节点n，让I(n)和O(n)分别表示其输入和输出张量的集合。因此，如果且仅当O(n1) ∩ I(n2) ≠ ∅（即，n2依赖于n1），则(n1, n2) ∈ E。

FlexLLM的图剪枝设计基于两个重要的观察结果。首先，传统的机器学习（ML）训练过程在前向传播中保留所有中间激活，以便在随后的反向传播中计算梯度。然而，由于大多数操作符的线性代数特性，中间激活被保留用于计算在PEFT期间被冻结的基础大型语言模型（LLM）的可训练参数的梯度。这个属性允许FlexLLM在不使用大多数中间激活的情况下计算PEFT模型的旁路网络的梯度。其次，不同的PEFT模型将旁路网络附加到基础LLM的不同位置，因此需要不同的中间激活来通过网络反向传播（见图8中的示例）。

基于这些观察，FlexLLM在为每个注册的PEFT模型生成PCG（并行计算图）时执行静态图剪枝，并动态调度这些PEFT模型的微调和推理任务并分配内存（见第6节）。算法1展示了FlexLLM的图剪枝算法，它以PEFT模型的PCG G作为输入，并生成一个执行计划，指定了如何在前向传播期间缓存或重新计算中间激活以进行反向传播。首先，REVERSEAUTODIFF使用反向自动微分算法[27]构建反向图G。接下来，由于在微调期间基础LLM的所有权重都被冻结，FlexLLM剪枝这些权重的梯度。剪枝操作符的输出张量可以使得后续剪枝机会。FlexLLM维护一个队列Q来跟踪所有输入和/或输出张量已更新的操作符，并迭代地检查这些操作符以发现额外的剪枝机会。图剪枝后，FlexLLM生成一个张量集合A，其中包含所有需要为反向传播保留的激活。FlexLLM进一步将图剪枝与两种已建立的内存优化技术相结合：重新生成和激活压缩。

重新生成（Rematerialization）。张量重新生成[30]包括在前向传播中选择性地丢弃一些张量，并在反向传播期间重新计算它们。最初设计用于减少深度神经网络（DNN）训练的峰值内存需求，这项技术对于大型语言模型（LLM）的共同服务也是有益的。对于图剪枝后剩余的每个张量t ∈ A，FlexLLM仅当以下条件满足时重新生成t：(1) t的所有输入张量都存储在A中，以及(2) 使用输入张量重新计算t引入的重新生成开销较低。

激活压缩（Activation compression）。此外，FlexLLM还尽可能地应用无损压缩来处理剩余的张量。这是因为对于某些操作符进行反向计算并不需要访问原始输入张量。一个例子是ReLU操作符 y = ReLU(x) = ( x x>0 0 x≤0，其导数是 ∂y/∂x = ( 1 x>0 0 x≤0)。因此，FlexLLM不是存储原始输入张量x，而是保留x的位掩码。注意，这种压缩不会引入任何精度损失。

6. 动态调度

在静态编译生成PEFT模型的可执行广义PCG（并行计算图）之后，FlexLLM需要确定执行时间表来同时处理微调和推理请求。这是一个挑战，因为每种工作负载都有其独特的特点。首先，推理请求以在线方式接收并且动态波动，这需要一个能够快速调整GPU资源以确保低延迟的系统架构。另一方面，微调涉及一个前向传递（类似于推理）和一个后向传递来更新模型参数。第三，生成性LLMs的自回归推理需要一种内存密集型的逐词生成模式，而微调通常采用计算密集型的序列级方法以获得更好的吞吐量。为了应对这些挑战，FlexLLM采用了一种新颖的token级微调机制和混合token调度器。

6.1 token-level finetuning

为了最大化GPU的利用率，现有的大型语言模型（LLM）微调系统通常在序列级别而不是token级别处理微调样本。因此，直接将推理请求与微调请求批量处理会显著增加每次迭代的延迟，并影响推理请求的服务等级目标（SLO）。

我们在算法2中提出了我们的token级微调机制。主要思想是使用由一系列token组成的动态滑动窗口，将微调过程划分为更小的步骤。这个窗口用于前向传递和后向传递，其大小由混合token调度器确定（第4行和第13行），以确保co-serving的推理token不超过SLO约束（第6.2节）。在前向传递过程中（第3-11行），在每一步中，我们将微调请求切片成si个词符的窗口（第5行），并将它们输入模型。模型逐层处理这些tokens（第7-8行），计算这些tokens的生成损失（第10行），然后移动到下一个tokens窗口（第11行）。为了保持与序列级微调相同的计算语义，我们需要缓存键和值（类似于推理的增量解码[70]）以及由后向注意力计算重用的查询，遵循自回归LLMs的causal mask。

在后向传递中（第12-20行），我们每一步获得sj个tokens损失的切片（第14行），并按相反的顺序执行模型层以计算梯度（第15-16行）。然而，与序列级注意力执行不同，我们的token级后向注意力执行更为复杂，因为token依赖于自回归。我们在图9的左侧和右侧分别说明了注意力模块的前向和后向PCG。当我们获得sj个token损失的切片时，生成的查询、键和值的梯度（即∆Q、∆K、∆V）应该具有[sj, h]的形状，其中h是隐藏层大小。∆Q的计算需要访问所有之前lj个词符的键缓存，这与前向执行中的因果causal mask一致。然而，自回归执行还使得生成的∆V和∆K具有[lj, h]的大小，因为窗口内token的损失也对所有之前的词符有所贡献。如图10所示，我们建议在每一步中累积KV梯度，并且只有当序列中所有词符的后向遍历完成时，才应用完全累积的sj个梯度（即大小为[sj, h]）来计算其他模型梯度。

6.2 Hybrid token scheduler

FlexLLM的混合token调度器以两个步骤安排token级的GPU计算。首先，FlexLLM根据特定的调度策略决定可用推理请求的时间表。FlexLLM使用Orca的迭代级调度[70]作为默认调度策略，该策略维持一个固定的最大批处理大小b，并在有可用的新请求时动态地用新请求替换每个已完成的请求。其他调度策略，例如为不同长度的请求分配不同的优先级，或者将长输入提示的预填充阶段分割成多个部分，对FlexLLM来说也是可行的。其次，在每次迭代中确定推理时间表后，为了最大限度地利用资源，FlexLLM尽可能多地向批次添加微调tokens（即滑动窗口大小的s）。可以添加的微调词符数量是通过自动使用公式s = arg max f(c, s) ≤ SLO来确定的，其中f(·, ·)是延迟估计函数，c是当前迭代中正在运行的推理token数量。推理延迟可以通过预先分析LLM执行成本来廉价估计[46]，以便将调度开销与每次迭代的延迟相比变得微不足道。

图11说明了FlexLLM的执行时间线，并展示了在特定推理请求到达模式（即r0至r9）下的微调小批量rA中的一个词符调度计划示例。为了保持相同的微调语义，FlexLLM的调度遵循每个小批量的前向和后向传递之间的执行依赖性。在前向传递中，FlexLLM使用融合kernel来处理推理和微调tokens，避免了额外的kernel启动开销。这是可能的，因为请求遵循相同的token计算逻辑。关于注意力操作符，在预填充阶段，微调请求遵循与推理请求相同的逻辑和causal mask。对于后向传递，FlexLLM采用空间共享方法，使用两个独立的流来处理推理和微调tokens。