EdgeShard: Efficient LLM Inference via Collaborative Edge Computing(IOT期刊, C刊)

EdgeShard: 基于协同边缘计算的高效大型语言模型推理

传统上，LLM部署在云上，会导致延迟延长、带宽成本高和隐私问题，最近边缘计算被认为有希望解决这些问题，因为边缘设备更加接近于数据源。然而边缘设备因其资源有限，几乎无法承担大语言模型。

现有的工作通过 将繁重的工作负载从边缘卸载到云 或 通过模型量化压缩大型语言模型 来解决这种限制。但是这些方法要么严重依赖于远程云，要么准确度较低。

这项工作是首次在协作边缘计算环境中部署大型语言模型，在该环境中，边缘设备和云服务器共享资源并协作 以高效率和无准确性损失的方式推断大型语言模型。

我们设计了EdgeShard，这是一种能够 将计算密集型大型语言模型 划分 为可负担的分片 并将其部署在分布式设备上的新方法。

考虑到设备异构性、带宽限制和模型复杂性，分区和分布是非常重要的。为此我们提出了一个自适应的联合设备选择和模型划分问题，并设计了一个动态规划算法来优化推理延迟和吞吐量，

实验是对Llama2模型进行的，EdgeShard实现了高达50%的延迟减少和2倍的吞吐量提高

INTRODUCTION

LLM的特点是 规模通常有数千亿个参数，使模型能够捕获语言和上下文中的复杂模式，使他们在生成连贯和上下文适当的响应方面非常有效，这种现象被称作“intelligence emergence”，LLM的杰出能力使他们在广泛的应用中具有价值和良好的表现

然而，目前的LLM严重依赖于云计算 ，有以下问题：响应时间长、带宽成本高和隐私问题。

对云计算的依赖可能会阻碍实时应用中所需的快速推断能力，例如机器人控制、导航等领域，即时的相应是至关重要的。

其次，向云数据中心传输大量数据会消耗大量带宽和网络架构。

第三，基于云的LLM带来了重大隐私问题，处理医院和银行的敏感数据以及手机上的文本输入和照片等个人数据时。

边缘计算可以通过在更靠近数据源的网络边缘上的边缘设备来部署LLM，从而解决上述问题。然而LLM所需资源极高，全精度Llama2-7B模型的推理至少需要28gb内存，这可能超过大多数边缘设备的容量。

有些works利用模型量化来减小模型尺寸，以此来适应资源约束后的边缘设备，但这样对应的准确率不高。

其它的works倾向于使用云边缘协作，将LLM划分为两个子模型，并将部分计算工作负载卸载到云服务器上，但边缘设备与云服务器之间的延迟通常很高且不稳定。

近年来，边缘计算能力得到了不断的增长，在网络边缘部分部署了大量的边缘服务器和边缘云，留下了大量的资源供使用。协同边缘计算（CEC）的提出是为了整合地理分布的边缘设备和云服务器的计算资源，实现高效的资源利用和性能优化，如图1所示，无处不在的分布式边缘设备与云服务器连接，形成共享资源池，协同提供即时的数据处理和AI服务。

CEC不同于现有的边缘计算研究，现有的研究侧重于 云、边缘和终端设备三者之间的垂直协作，而忽略了横向的边缘到边缘协作，存在资源利用不优化、服务覆盖受限、性能不均衡等问题。

在CEC愿景的激励下，我们提出了一个通用的大型语言模型推理框架，名为EdgeShard，以支持分布式边缘设备和云服务器上高效的协作大型语言模型推理。为简单起见，我们使用下面的计算设备来指代边缘设备和云服务器。

在异构计算设备的网络中，EdgeShard将LLM划分为多个分片，并根据异构计算资源和网络资源，以及设备的内存预算，合理地分配给设备。

为了优化性能，我们提出了一个联合的设备选择和模型划分问题，并设计了一个高效的动态规划算法，分别最小化推理延迟和最大化推理吞吐量。

在实际测试平台上进行的大量实验表明，EdgeShard减少了高达50%的延迟，并实现了设备上和垂直云边缘协同推理方法的2倍吞吐量。

该工作与现有工作不同之处在于，将LLM分区并分配给云数据中心的多个gpu，如Gpipe、PipeDream和Splitwise。云服务器通常使用同构gpu，边缘设备本质上具有异构计算能力。用于LLM的现代云gpu通常通过InfiniBand、nvlink等高带宽网络连接，而边缘设备则通过异构低带宽网络连接。针对云数据中心设计的LLM部署解决方案忽略了异构和资源受限的边缘计算环境。

补充：在边缘环境，各设备的能力差异巨大，必须根据每台设备的实际算力、可用内存和网络带宽来自适应地选择参与设备、细粒度地分割模型，并制定最优的分配策略，才能兼顾性能（如延迟、吞吐）与资源利用。

同时，EdgeShard也区别于之前的 将小尺寸DNN划分为分布式边缘设备进行协同推理，因为LLM参数大，推理是自回归的，我们的工作解决了这个问题。

此外，他们只考虑固定数量设备上的模型划分问题，而EdgeShard则更进一步，执行联合设备选择和模型划分。

本文主要贡献内容如下：

1. 提出了一种用于在边缘计算环境中部署LLM的通用LLM推理框架，实现了异构边缘设备和云服务器之间的协同推理。
2. 进一步，我们定量研究了如何选择计算设备和如何划分LLM以优化性能。我们从数学上提出了一个联合设备选择和模型划分问题，并提出了一种动态规划算法来分别优化延迟和吞吐量。
3. 在物理测试台上用最先进的Llama2系列模型评估了EdgeShard的性能。实验表明，EdgeShard性能明显优于各种baseline。
4. 此外，概述了几个悬而未决的问题，如激励机制、批量大小感知优化和隐私保护技术，以指导边缘协作大型语言模型推理的进一步发展。

PRELIMINARIES AND MOTIVATIONS

**生成式LLM推理：**LLMs通常指具有数十亿个参数的基于解码器的transformer模型。与BERT等基于编码器的体系结构的推理过程为单阶段不同，LLM推理过程是迭代的，通常包括两个阶段：1) 提示处理阶段和 2) 自回归生成阶段。提示处理阶段也称为预填充

在提示处理阶段，模型接受用户初始的token作为输入，通过计算概率来生成一个新的token，也就是对应接跟着的一个新的单词

在自回归生成阶段，模型基于初始输入和迄今为止生成的tokens，一次生成一个token。此阶段依次为多个迭代生成tokens，直到满足停止条件（即在生成序列结束(EOS)的 token 或达到用户指定或LLM约束的最大tokens数量时）。

基于解码器的LLM模型通常由嵌入层、重复的线性连接的 transformer 层和输出层组成。如图2所示，假设LLM模型有N层，它将采用一系列输入tokens，并运行所有层以逐个生成token。

在Prefill阶段，模型接受输入(“Today is a”)，并且第一个生成的token是“good”。在Autoregressive Grneration阶段，模型首先将(“Today is a good”)作为输入，并生成下一个token(“day”)。然后，它将(“Today is a good day”)作为输入，并生成下一个token(“EOS”)，这表明这一代的结束。

由于生成的token是由序列中之前的所有token决定的，因此LLMs使用键值缓存(KV Caching)来避免重复计算。每个transformer层存储过去的计算以加快响应，从而减少计算工作量并改善响应时间。

由于Prefill阶段需要计算所有输入tokens的KV缓存作为初始化，因此在Prefill阶段生成token的时间要比自回归阶段高得多(通常为10倍)。

**LLMs消耗内存：**单个边缘设备可能没有足够的内存来容纳LLM模型。以最流行的LLM模型之一，即Llama2为例。如表1所示，Llama2有三个不同的版本，即7B、13B和70B。从表中我们可以看出，Llama2-7B的全精度推理至少需要28gb内存，但智能手机通常只有6 - 12gb内存，Jetson Orin NX有8 - 16gb内存。它们无法负担设备上的LLM推理。

一些作品尝试使用低精度量化，例如8位和4位。但是，它仍然可能超过边缘设备的内存容量。例如，Llama2-70B的4位推理至少需要35gb内存，这在大多数边缘设备上是无法容纳的。此外，低精度推理会导致性能下降。

在我们的work中，我们利用了CEC，这是一种使得 地理分布式的边缘设备 和 云服务器 协作 来 执行计算任务 的 计算范式。

基于这个想法，我们提出了EdgeShard，一个通用的LLM推理框架，允许自适应设备选择和分布式计算设备上的LLM分区，以解决高内存需求并利用异构资源来优化LLM推理。

COLLABORATIVE EDGE COMPUTING FOR LLMS

该框架有三个阶段，包括离线预填充、任务调度优化和协作推理。工作流程如图3所示。

Profiling

预填充是一个离线步骤，优化步骤所需的必要运行时的跟踪信息，只需要执行一次，这些跟踪信息包括：

1. 各层在不同器件上的执行时间
2. LLM模型各层的激活大小和内存消耗
3. 每个设备的可用内存和设备间的带宽

对于每一层的执行时间，我们分别分析了在预填充阶段和自回归阶段生成token的时间，并取其平均值。对于那些可能没有有效内存来保存执行分析的完整模型的设备，我们利用动态模型加载技术，其中模型层被连续加载以适应受限的内存。

更具体地说，我们首先根据模型体系结构计算每层的内存消耗。根据层的内存消耗和给定设备的内存预算，我们可以估计设备可以容纳的最大层，在此基础上我们将模型划分为几个分片，并依次将分片加载到设备中。此外，在动态加载期间，存储前一个分片的输出/激活，并作为下一个分片的输入。然后，分析信息将用于支持智能任务调度策略。

Scheduling Optimization

在任务调度优化阶段，调度程序通过确定参与推理的设备，如何明智地对LLM层进行分区以及应该将模型分片分配给哪个设备来生成部署策略。

该策略综合考虑了异构资源、设备的内存预算和隐私约束，可以稍后应用于选定的设备以进行有效的LLM推理。

Collaborative Inference

在得到LLM模型的划分和分配后，选择设备进行协同推理，其中的一个关键问题就是KV-cache管理，因为KV缓存大小会随着生成tokens长度而不断地增加。如果以前tokens的KV缓存被新令牌覆盖，则可能严重影响模型的输出。

为解决这个问题，我们在每个参与设备上为KV缓存预先分配内存空间，考虑到每个transformer层都有自己的KV缓存，层是分区策略的基本单位，每个设备都可以保持其分配的transformer层的KV缓存，我们采用贪心原则，将每个设备都预先分配空间以适应KV缓存的最大内存消耗，假设分配的transformer层数为 $N_a$，输入tokens的批大小是 $B_b$，输入tokens的最大长度为 $I_m$，生成tokens的最大长度为 $G_m$，则预分配的KV缓存空间大小可计算为 $N_a B_b (I_m + G_m)$。

对于协同推理，我们分别考虑顺序推理和流水线并行推理：

在顺序推理中，设备轮流使用分配到的模型分片执行计算，如图4(a)所示，其中transformer含有30层，将模型划分为3个分片，分别分配给设备1-3，分别为1-16层、17-26层和27-32层。设备1将首先处理输入的tokens，然后将激活/输出发送到设备2，设备2将处理数据，然后发送到设备3。设备3将把生成的token发送回设备1，然后设备1将继续处理新的输入tokens。此循环将重复多次迭代，直到生成EOS token或达到生成tokens的最大数量。

然而从系统角度来看，顺序推理并不是资源高效的，当设备1执行计算时，设备2和3处于空闲状态，因此我们采用流水线并行性来提高资源利用率，对于前面对云服务器的工作中所做的流水线并行推理，输入的数据将首先被分割成微批，然后馈送到系统中。

如图4(b)所示，设备1首先处理数据B1，然后将中间数据传输给设备2。处理完数据B1后，设备1立即去处理数据B2。在这种流水线方式下，每个设备都很忙，系统资源利用率很高。

OPTIMIZE LLM INFERENCE

我们考虑了一个具有异构设备和带宽连接的通用协同边缘网络。更具体地说，给定一组连接异构带宽的异构设备，EdgeShard旨在选择设备子集并将LLM划分为分片，这些分片将分配给所选设备，以最小化推理延迟或最大化吞吐量。

系统模型：LLMs 通常具有分层架构，由嵌入层、多个解码器层和输出层组成。参数和激活的大小在各层之间是不同的，我们假设模型有$N$ 层，$O_i$ 表示第 i 层的激活大小，满足 0 ≤ i ≤ N−1。层 i 的内存消耗用 $Req_i$ 表示。

我们考虑一个由 M 个边缘设备和云服务器组成的网络。这些设备具有异构计算和内存能力，并且云服务器在计算能力方面比边缘设备强大得多。设备 j 的内存预算为 $Mem_j$。计算设备相互连接。设备 k 和设备 j 之间的带宽为 $B_{k,j}$, 0 ≤ k ≤ M−1,0 ≤ j ≤ M−1。输入tokens驻留在一个源节点中。在不丧失一般性的前提下，我们将源节点设为节点0。本文中使用的主要符号如表2所示。

优化LLM推理延迟

一个层只能分配给一个节点，每个节点的层都有各自的计算时间，从一个节点到另一个节点传输某层的激活对应一个通信时间，数据的传输时间由一层的输出大小和两个节点之间的带宽来决定，如果两层在同一个节点上，那么在我们的work中假设传输时间为0，否则传输时间等于（激活大小除以带宽）

总的推断时间可以被计算为：所有层各自的计算时间 + 每两层之间的传输时间

LLM模型的第一层应该始终分配给节点0，节点0被设置为具有输入tokens的源节点。在这种情况下，原始输入数据驻留在源节点上，避免在计算设备之间传输。由式(5)可知，分配给节点 j 的所有层的内存需求不能超过节点 j 的内存预算。由式(6)可知，一层可以且只能分配给一个节点

解决方案：为了最小化推理延迟，我们设计了一个动态规划算法。直观地说，第一个 i 层的最小执行时间由第一个 i −1 层决定，这意味着可以从子问题的最优结果构造最优解。它具有最优子问题的性质，这促使我们使用动态规划。

令DP(i, j)表示将第 i 层分配给节点 j 后，前 i 层的最小总执行时间，状态转移方程如下：

将 i 层分配给节点 j 后，前 i 层的最小总执行时间 = min（将i−1层分配给设备k时，前 i−1层的最小执行时间 + 将 i 层分配给设备 j 的计算时间 + i - 1 层激活从设备k到设备 j 的传输时间 + 将 i 层激活从设备 j 到设备 0 的传输时间（其中的1(i)为指标函数，只有当i = N - 1 时其值为1，否则为0））

由该式可知，DP(i, j)是通过遍历前一层所有可能的节点，选择前 i 层执行时间最少的节点来确定的，由于 LLM 的自回归特性，生成的token需要发送回源节点进行下一次迭代，因此对于最后一层 N - 1，通信时间处理从N - 2层传输数据过来的时间，还包括到源节点的传输时间

基于上图的式子来遍历每一层和每个节点，我们可以填入动态规划表DP(i, j)来跟踪到达每一层的最小总执行时间。最后，最后一层的最小总执行时间可由下式计算。然后，我们可以通过回溯DP(i, j)得到每层的最优节点分配。

这种方法简单有效。通过动态规划，可以快速遍历解空间，找到最佳的LLM分区和分配策略。算法1给出了寻找最小化推理延迟的最优LLM分区和分配策略的算法。

我们从第1层开始遍历大型语言模型的各层。对于每一层，我们遍历具有足够内存的所有计算设备并计算推理时间(第3-19行)。在填满DP表后，我们可以找到最小DP(N−1,j)，它表示执行LLM模型所需的最小时间，以及到达主机层N−1的最终节点。最后，通过回溯选项(i, j)，我们得到模型分区和分配策略R(第20-28行)。

优化LLM推理吞吐量

在本文中为了优化吞吐量，采用了流水线并行的方式，流水线的并行推理中，计算时间和通信时间是可以重叠的，以最大限度来提高吞吐量，对于推理任务而言，设备 j 的最大延迟可以表示为：

最大化吞吐量等同于最小化最慢设备的延迟，与最小化推理类似，吞吐量最大化的问题也可以用动态规划的方式来求解，通过求解第一个 i - 1 层的分配问题，推导出第一个 i 层的吞吐量使得最大化

设 g(i, S, k) 表示用已用设备集 S 处理前 i 层的最小时间，设备 k 为最后的节点；用 g(m, S, j) 表示用已用设备集S处理前m层的下一个状态，设备 j 为最后要使用的节点

此外，我们在执行状态转换时有约束。它们分别是(13)所示的内存约束和隐私约束

算法2描述了寻找Topt最优解的伪代码以及相应的模型划分和分配策略。在算法2中，我们首先初始化动态规划表g(m, S, j)和选择表choice(m, S, j)，并将t0,0分配给g(1,1,0)(第1行和第2行)。然后我们从第1层遍历大型语言模型的各层。对于每一层，我们遍历具有足够内存的所有计算设备，并计算最大延迟(第3-23行)。在填满DP表之后，我们可以找到最大延迟，然后基于此回溯选择表，最终获得模型分区和分配策略(第24-32行)。算法2的计算复杂度为0 (N2×2M× M2)，其中N为LLM模型的层数，M为网络中设备的数量。

上述问题基于理想化情况，即任何时候都没有空闲的设备，设备处理一批数据后无需等待，然而对于现实世界的LLM来说这是不现实的。

如图5(a)所示，与那些单阶段计算应用程序不同，基于解码器的LLM应用程序具有自回归的性质，其中会有多次迭代，每次迭代生成一个token。当前token的计算依赖于所有以前的tokens。因此，迭代不能开始，直到前一个迭代结束，它得到之前生成的token。对于流水线并行推理，这种方案会产生气泡，代表设备的空闲状态。x

如图5(a)所示，假设有4个微批次，P1-P4分别表示这些微批次的第一次迭代(预填充阶段)。G1A和G1B分别表示第二次和第三次迭代。它们也是生成阶段的第一次和第二次迭代。在这种情况下，直到第一次迭代结束，即P4的结束，G1A才能启动，这导致有设备处于空闲状态，并没有发挥极致，其中设备1-3处于空闲状态。

为了接近理想情况，提高吞吐量，我们尽可能减少pipeline中的旗袍，我们提出了EdgeShard-No-Bubbles版本，它优化了跨批的token生成顺序，允许立即生成对应token，而无需等待迭代中所有批的结束，如图5(b)所示，在第一批预填充阶段P1结束后，device-1立即执行G1A所示的第一批微批token生成。同样，当G1A结束时，device-1进入第一个微批的token生成的下一次迭代，由G1B表示。EdgeShard-No-Bubbles通过减少设备空闲时间来减少气泡，有望提高吞吐量。

EXPERIMENTAL EVALUATION

实验设置

测试平台:我们使用各种边缘设备和云服务器作为CEC中的异构计算设备。设备规格如表三所示。我们使用了15台设备，包括12台Jetson AGX Orin, 2台Jetson Orin NX和一台云服务器来配置协作边缘网络。

物理试验台如图6所示。这些设备通过路由和交换机连接。任意两台设备之间的带宽为1000mb /s。我们使用 Linux TC工具[24]来改变设备之间的网络带宽和通信延迟。

我们使用Llama2型号测试EdgeShard的性能，包括Llama2- 7b、Llama213B和Llama2- 70b。

Llama2由Meta于2023年7月发布，是最受欢迎和最强大的开源LLMs之一，代表了人工智能和自然语言处理领域的突破性飞跃。对于模型推理，我们采用文本生成任务来测试性能。

我们使用HuggingFace的WikiText-2数据集[25]。我们提取一个输入tokens长度为32的样本子集，并生成96个tokens。在接下来的所有实验中，我们都使用了全精度推理。我们将EdgeShard的延迟和吞吐量与各种baseline算法进行比较。

1. edge - solo: LLMs本地部署在边缘设备上，没有模型分区。
2. Cloud-Edge-Even。在这种情况下，LLMs被均匀地分成两部分。一个分配给边缘设备，另一个分配给云服务器。
3. Cloud-Edge-Opt:在这种情况下，LLMs被划分为两个分片。一个分配给边缘设备，另一个分配给云服务器。对于LLMs的划分策略，我们也使用了所提出的动态规划算法。不同之处在于只有两个设备作为算法输入。

总体评估

我们设置AGX Orin作为源节点，源节点与云服务器之间的带宽为1mb /s。其他计算设备之间的带宽设置为50mb /s，差异为20%。为了测试吞吐量，我们将批大小设置为参与设备可以支持的最大批大小。LLM推理的延迟和吞吐量如表4所示。

EdgeShard对于大型语言模型的部署是潜在的和有益的。对于Llama2-70B型号，内存需求约为280gb，远远超过单独边缘部署和云边缘协同部署的内存容量。它们会出现内存不足问题(OOM)。然而，EdgeShard通过 将大型模型分割成碎片 并将其分配给多个设备 来解决这一挑战，从而实现协作模型推理。

源节点和云服务器之间的带宽非常有限，即1 Mb/s。与Edge-Solo一样，cloud-edgecollaboration的最优部署策略是本地执行。

带宽的影响

我们将源节点设置为AGX Orin，并将云服务器和源节点之间的带宽从1 Mb/s到50 Mb/s不等。LLM推理的延迟和吞吐量性能分别如图7和图8所示。

在延迟方面，除了Edge-Solo之外，其他三种方法的延迟都随着带宽的增加而减小。这是因为这三种方法都是基于协作的，延迟受数据传输时间的影响。带宽的增加导致通信时间的减少。我们还可以看到，对于协作方法，当云源带宽从1 Mb/s变化到10 Mb/s时，延迟显著减少，而从10 Mb/s变化到50 Mb/s的变化较小。这是因为此时带宽逐渐饱和，计算时间成为瓶颈。

当带宽大于10时，云边缘写作方法优于Edge-Solo方法，因为云边缘协作方法引入了强大的云服务器来进行计算加速，但是，当带宽为1mb /s时，Cloud-Edge-Even的性能不如EdgeSolo。这是因为这种情况下的数据传输成本很高。当带宽大于10mb /s时，Cloud-Edge-Opt和EdgeShard的延迟几乎相同。我们发现EdgeShard生成与Cloud-Edge-Opt方法相同的模型分区和分配策略。

对于Llama2-70B, EdgeShard的性能优于其变体EdgeShard- even，因为云服务器和边缘设备之间存在资源异构，并且EdgeShard在计算设备之间自适应分区LLMs。然而，性能改进并不明显，因为有11个具有相同功能的AGX和只有一个RTX 3090。

源节点的影响

我们还测试了源节点对推理延迟和吞吐量的影响，因为源节点可能具有不同的计算和内存容量，并且EdgeShard强制驻留在源节点上的第一层LLM模型以避免原始数据传输。我们将源节点分别设置为AGX Orin和Orin NX，并比较它们的性能。我们将源节点和云服务器之间的带宽设置为1mb /s。Llama2-7B推断结果如图9所示。

当源节点为Orin NX时，EdgeSolo和Cloud-Edge-Even方法都会遇到OOM错误。这是由于Orin NX的内存相对较低，无法容纳Llama2-7B车型的一半。Cloud-Edge-Opt方法下两种情况的差异要比EdgeShard方法明显得多。对于Cloud-Edge-Opt，大约有60毫秒的差距，对于EdgeShard，大约有5毫秒的差距。这是因为在Cloud-Edge-Opt的情况下只有两个设备，并且它倾向于在源节点上放置更多的层。

然而，AGX Orin在计算能力上比Orin NX强大得多。EdgeShard倾向于涉及更多的设备，在源节点上放置更少的模型层，这可以填补源节点之间的计算能力差距。在吞吐量方面也观察到类似的现象，在Cloud-Edge-Opt方法下，AGX Orin的吞吐量比Orin Nx高6倍，而在EdgeShard方法下，吞吐量只比Orin Nx高2倍。结果表明，EdgeShard可以充分利用网络中的计算资源。

预填充策略的影响

Edgeshard 的第一步是进行性能分析，这需要每个设备运行 LLM 并记录执行时间。然而，对于 LLM 来说，预填充阶段和生成阶段的执行时间变化很大。在本研究中，我们将预填充阶段和生成阶段的平均值作为性能分析步骤中的稳定值。我们还研究了 EdgeShard 在不同性能分析策略下的性能。

我们以LLama2-13B为例。EdgeShard在不同profiling方法下对LLama2-13B推理的性能如图10(a)和(b)所示，其中prefill-profile、generation-profile和average- profile分别表示将prefill阶段的执行时间、generation阶段的执行时间和average作为稳定值。我们发现三种方法的性能差异较小，波动较小。这是因为虽然预填充和生成阶段之间存在10倍的差距，但不同计算设备的相对计算能力保持不变，即RTX3090(最强大)在预填充和生成阶段的推理时间都远低于AGX Orin(最弱)。计算能力的异构性对设备选择和模型划分策略有很大的影响。

流水线执行策略的影响

我们评估了两种流水线执行策略。我们将云服务器和源节点之间的带宽设置为1mb /s。结果如图11所示。

我们可以看到，对于所有方法，EdgeShard-No-Bubble都优于EdgeShard-Bubble。

具体来说，对于Llama2-7b, EdgeShard- no - bubble比Cloud-Edge-Even和EdgeShard的EdgeShard- bubble每秒分别提高了约0.36和6.96个tokens。对于Cloud-Edge-Opt方法，它在本例中选择本地执行。没有流水线执行，因此这两个方法的吞吐量是相同的。

对于llama1 -13b, EdgeShard- nobubble比Cloud-EdgeEven、Cloud-Edge和EdgeShard的EdgeShard- bubble每秒分别提高了1.69、1.89和5.21个tokens。与EdgeShard-Bubble相比，EdgeShard-No-bubble不需要在一次迭代中等待所有微批的完成，可以有效地减少设备的空闲时间，从而提高吞吐量。

Galaxy: A Resource-Efficient Collaborative Edge AI System for In-situ Transformer Inference(INFOCOM, A会)

Galaxy: 一种资源高效的协同边缘人工智能系统，用于现场 transformer 推理

基于transformer的模型已经在边缘解锁了强大的智能应用程序，传统的部署方法是直接将推理工作负载卸载到远程云服务器上，这将对骨干网络造成巨大的压力，并引起用户的隐私问题。

为解决这个问题，本地推理最近被认可为边缘智能，但它仍然具有密集工作负载的问题和有限设备计算资源冲突的问题。

本文中，我们利用我们的观察，即许多边缘环境通常包含一组丰富的具有空闲资源的伴随可信边缘设备，并提出了Galaxy：一种协作边缘人工智能系统，可以打破异构边缘设备之间的资源墙，以实现高效的Transformer推理加速。Galaxy引入一种新的混合并行模型来协调协同推理，以及一种异构感知的并行规划，以充分利用资源潜力。

此外，Galaxy 设计了一种基于瓦片的细粒度通信与计算重叠机制，以减轻带宽受限的边缘环境中张量同步对推理延迟的影响。基于原型实现的大量评估表明，在各种边缘环境设置下，Galaxy 显著优于最先进的方法，端到端延迟最多可降低 2.5 倍。

INTRODUCTION

目前，大多数基于transformer的智能应用严重依赖于云服务，大规模基于transformer的模型的实际推断发生在云中。在边缘，只部署一个代理守护进程来转发用户请求，然而云辅助方法存在以下问题：

1. 由于边缘设备和远程云[7]之间的广域网(WAN)连接不可靠和容易延迟，服务质量可能会受到影响。
2. 来自众多边缘客户端的推理请求会对骨干网和数据中心造成巨大压力。
3. 智能应用中的传感数据可能包含高度敏感或私人信息。

为了解决这个问题，无需远程协助的边缘设备上的本地推理，能够将数据保存在本地并避免网络传输，已经被认为是一个有前途的范例，然而Transformer推理的计算密集型和资源消耗性对资源受限的边缘设备提出了重大挑战.

在现成的边缘设备中对Bert-L模型[12]的推断施加了几乎700MB的最小可用内存空间，而延迟比数据中心GPU长121倍。这些结果表明了密集的Transformer推理工作量与有限的机载资源之间的根本矛盾。

为了应对这些挑战，现有技术探索设计复杂的调度机制，以利用边缘设备的资源潜力，但仍然受到单个设备有限的板载资源的瓶颈。

我们观察到智能家居等边缘设备通常物理上较近，促使我们将邻近可用的边缘设备视作资源增幅，并以分布式方式与它们进行协作，以在边缘呈现快速的推理。如图1所示，我们可以利用智能家居中的分布式计算资源来加速基于Transformer的语音助手，然而这种模式带来几个挑战：

1. 如何在多个边缘设备之间并行化单次Transformer推理工作负载
2. 如何根据异构边缘设备的资源预算确定工作负载划分策略
3. 如何在带宽有限的边缘环境下降低分布式推理延迟

为了应对这些挑战，我们提出了Galaxy，这是一个协作边缘人工智能系统，它打破了跨异构边缘设备的资源墙，实现低延迟Transformer推理，从而实现实时的原位边缘智能服务。

Galaxy的贡献不仅仅是利用分布式边缘设备来部署Transformer推理，而是在三个层面上解决了上述挑战。

首先，为了最大限度地利用资源协调异构辅助设备，促进协同推理，引入了一种新的混合并行模型(HMP)，它结合了张量并行(TP)和序列并行(SP)的优势，作为一种新的并行架构来管理分布式推理工作流。

其次，为了最大限度地提高边缘设备间HMP的资源利用率，提出了一种 综合考虑设备资源异构性和内存预算 的 工作负载规划算法。

第三，为了在带宽有限的边缘环境中实现低延迟的协同推理，我们通过将 连续计算 和 通信操作 分解成细粒度的块来细致地解耦它们之间的紧密数据依赖，从而实现同步的有效重叠

在实际测试平台上的广泛评估表明，与最先进的协同推理方法相比，Galaxy实现了高达2.5倍的加速。与单设备情况相比，Galaxy的4路并行推理可以实现86%的缩放效率。据我们所知，Galaxy是第一个将混合模型并行性应用于边缘协作Transformer推理场景的作品。

综上所述，本文做出了以下贡献。

• 通过对设备上和并行推理方法的广泛测量研究，我们引入了一种新的HMP架构，用于与可信边缘设备协作进行原位单次Transformer推理加速。
• 我们设计了一个异构和内存预算敏感的工作负载规划算法，以促进资源高效的边缘协作推理。
• 我们提出了一种基于tile的细粒度优化，利用通信和计算重叠的概念来减轻同步开销。
• 我们实现了Galaxy并在实际的边缘测试平台上对其进行了评估。实验结果表明，与最先进的方法相比，延迟减少了2.5倍。

BACKGROUND AND MOTIVATION

基于Transformer的模型

当前与语言相关的应用程序倾向于使用基于Transformer的模型，这些模型由Transformer层堆叠而成，因为它们具有优越的准确性，原始Transformer的组成包括一个编码器和一个解码器，在本文中我们主要关注最近的LLM，如Bert和GPT-2，它们只使用编码器或解码器组件，图2展示了我们在本文中考虑的Transformer层的模型架构

在Transformer层中，主要组件是多头注意力(MHA)模块和多层感知器(MLP)模块，这些组件通过诸如丢弃、残差加法和层归一化等逐元素操作相连接。在MHA块中，第一个线性层为每个注意头生成查询(Q)、键(K)和值(V)矩阵。每个头部独立进行自注意力机制，它们的输出通过最后的线性层连接并进一步处理得到输出。MLP块涉及两个线性操作，将隐藏大小从h增加到4h，然后将其减少到h。

资源受限边缘设备的Transformer推理

本地推理可以利用边缘环境中的空闲资源，同时充分保护用户的数据隐私，使其成为隐私敏感边缘应用中广泛使用的范例。

然而，Transformer推理的资源密集型性质对资源有限的边缘设备提出了重大挑战。

我们进行实验来分析有限的计算资源如何影响设备上的Transformer推理。实验设置在§IV-A中描述，结果如表1所示。

具体而言，我们使用长度为30的输入序列在现成的边缘设备和Nvidia GPU平台上对五种典型的基于transformer的模型进行设备上推理。我们观察到A100和Nano- m之间的推理延迟表现出巨大的差距，例如Jetson Nano与Bert-L上的A100相比延迟了121倍。内存预算是Transformer推理中的另一个关键因素。半精度浮点格式的GPT2-L在推理期间占用1.6GB内存，超过了单个Nano-M的1.5GB预算。

为了减轻资源限制，我们利用我们的观察，即边缘环境通常由物理接近的多个可信边缘设备组成。这使得这些边缘设备之间的计算资源共享是相互信任的。

多设备协同Transformer推理

在跨边缘设备的协同变压器推理中，关键问题是并行策略的选择。我们在图3中说明了不同的并行计划。

(1) 数据和流水线并行：

**数据并行(DP) 和 流水线并行(PP)**是并行执行基于Transformer的模型的常用方法。

DP沿着样本维度划分工作负载，允许每个设备独立执行推断。在边缘智能服务中，频繁提出单次推理请求(例如，向智能助手发送单个语音命令)，由于缺乏数据批次，DP不适用。

PP将模型沿层维水平划分为连续的阶段，每个阶段映射到一个不同的设备。然而，在单次推断的情况下，PP在同时利用多个边缘设备方面仍然不足，因为阶段间的数据依赖迫使每个设备等待前一个设备的完成。

(2) 模型并行：

模型并行(MP)是一种并行计算范式，它在一个模型层内水平划分操作，促进单次推理的并发执行。

应用于Transformer模型的最常见的模型并行技术是 张量并行(TP) 和 序列并行(SP)。

TP分区跨设备建模权重，每个设备托管一个参数子集，但它无法并行化MHA和MLP块之间的一些元素操作。

相比之下，SP沿着序列维度分割输入，促进所有操作的并行性，但要求每个设备存储整个模型参数。

由于层内数据依赖，在MP过程中插入同步点，以保证协同推理结果与本地推理结果的一致性。然而，这些同步点引入了显著的通信延迟，可能成为推理性能的瓶颈，特别是在带宽有限的边缘环境中。

总结上述分析，促使我们设计一种混合模型并行架构，该架构结合了TP和SP的优点，并采用通信优化方法来减轻同步开销。

GALAXY DESIGN

Galaxy Workflow

我们的系统设计旨在同时利用多个异构边缘设备来实现低延迟的本地推理，图4说明了我们提出的Galaxy的工作流程，它具有三个主要阶段：预处理阶段，并行规划阶段和执行阶段。

预处理阶段是在部署前运行一次的脱机过程，Galaxy Profiler使用校准数据作为物理边缘设备上的输入来执行推理过程，以记录并行性规划(步骤1)所需的运行时跟踪。在并行规划阶段，Galaxy采用了一种新的混合模型并行(HMP)架构，该架构结合了TP和SP来编排分布式边缘设备(步骤2)。

Galaxy Planner将来自Galaxy Profiler的分析结果作为输入，以生成并行规划配置(步骤3)。该配置综合考虑了资源异质性和内存预算，并随后应用于执行阶段的目标模型和边缘设备，以实现高效的边缘协同推理(步骤4)。

分布式推理不可避免地涉及张量同步操作。Galaxy集成了基于tile的细粒度通信优化，以减轻额外通信开销带来的性能下降(步骤5)。通过以上模块，Galaxy主要实现以下设计目标：

• 用于跨多个边缘设备的低延迟单次Transformer推理的HMP架构(§III-B)。
• 明智的并行规划器，全面考虑设备异构性和内存预算，旨在以负载均衡的方式分配工作负载，以充分利用边缘设备的计算资源(§III-C)。
• 基于tile的细粒度通信优化将连续计算和通信操作之间的紧密依赖解耦，从而实现它们之间的有效重叠(§III-D)。

Hybrid Model Parallelism

Galaxy采用了创新的HMP（混合模型并行）架构，可在边缘环境中促进高效的并行Transformer推理。在本节中，我们将使用跨两个边缘设备进行协作推理的示例来详细说明我们的HMP架构。如图5所示，TP和SP在Transformer层中交替存在。TP应用于MHA块和MLP块，SP应用于连接MHA（多头注意力模块）和MLP块的操作，即连接块。

MHA块上的张量并行性：

设计一种高效的TP方法的目的是减少跨不同设备的操作员之间的数据依赖性，从而降低张量同步的频率。如图5所示，应用TP的第一个块是MHA块。我们利用了MHA固有的并行性优势:多个注意头的计算是完全独立的。这种头级依赖使我们能够在执行Multi SelfAttention操作期间，在没有任何张量同步的情况下，在边缘设备上拆分每个注意头的操作。考虑到这一点，我们沿着它们的头部维度划分与键(WK)、查询(WQ)和值(WV)相关的权重矩阵。初始的通用矩阵乘法(GEMM)分布到不同的器件上，并沿头部维数(1)并行化。随后，每个注意头对应的自注意力机制在各自的设备上局部进行(2)。输出线性层的最终GEMM沿着其行维并行化，确保与初始GEMM的正向分区对齐(3)。对设备i (i∈{0,1})的操作可以表示为([·|·]为concat操作):

MLP块上的张量并行性：

如图5所示，使用TP的第二个区块是MLP区块，它由两个连续的GEMM组成。为了避免第一次和第二次GEMM操作之间的张量同步，我们利用矩阵平铺的概念来消除数据依赖性。我们沿着第一个GEMM的列维度(4)划分权重矩阵，并沿着其行划分第二个GEMM，以与第一个GEMM的列方向划分(5)对齐。第二个GEMM可以直接将第一个GEMM的输出作为输入，而不需要同步点。对设备i (i∈{0,1})的操作可表述为:

连接块上的序列并行性：

TP加速了每个Transformer层中计算最密集的部分，同时使连接MHA块和MLP块的Dropout, Residual Addition和layer Norm保持不变(6)。尽管这些操作是基于元素的，并且不需要密集的矩阵乘法，但它们需要相当多的内存访问，因此也会产生不可忽略的执行延迟。我们注意到这些元素操作在序列维度上是独立的，这允许我们通过划分输入序列来并行化它们。对设备i (i∈{0,1})的操作可表述为:

张量同步点：

为了保证我们HMP的推理结果与本地推理结果一致，需要在每个TP和SP块的末端设置一个同步点，如图5所示。

为了完成TP块，需要一个ReduceSum操作来聚合多个设备(G←C0 + C1和G←F0 + F1)的计算结果。随后，聚合结果沿着序列维度进行分区，并分散在SP ([G0|G1]←G)的各个边缘设备上。这两个操作可以有效地结合起来，并使用单个ReduceScatter操作实现(7)。

在完成SP块时，每个设备只保留输入序列的一部分。收集所有这些片段，将它们连接起来，并将它们分发到所有设备以进行后续TP (A←[H0|H1]和D←[H0|H1])。因此，我们在每个SP块的末尾执行一个AllGather通信原语(8)。

混合模型并行体系结构的优点：

与直接的TP或SP体系结构相比，采用HMP体系结构具有许多优点。

与TP相比:

(1)HMP架构消除了连接块中的冗余计算，充分利用了Transformer层的并行潜力。

(2) HMP不引入额外的通信开销。乍一看，最先进的TP[24]需要两个AllReduce操作，而HMP需要每个Transformer层两个ReduceScatter和两个AllGather操作。然而，在通信原语的实现中，单个Ring-AllReduce操作的通信量相当于一个Ring-ReduceScatter和一个RingAllGather[27]。

(3) HMP架构将较大的AllReduce操作拆分为两个较小的原语ReduceScatter和AllGather，极大地促进了§III-D中提出的基于平铺的通信重叠。

与SP相比:

SP沿着序列维度划分输入张量，而不划分权重矩阵。这种范式要求每个设备都能容纳全局模型的整体副本。HMP通过跨设备分布模型参数缓解了这个问题，从而打破了单个设备的内存墙，实现了内存资源的可伸缩性。

异构性和内存感知工作负载规划

图5显示，每个TP或SP块完成后需要一个同步点。这些同步点的启动受最慢设备(离散器)完成时间的限制。这样的掉队者会使其他更快的设备挨饿，从而导致资源利用率不足。考虑到设备计算能力的固有异质性，特别是在边缘环境中，采用 异构感知工作负载规划 对于以平衡的方式分配工作负载至关重要。

此外，基于Transformer的模型的推理需要大量的内存。在实际部署中，内存不足(OOM)问题是推理的一个障碍，这对通常在严格内存限制下运行的边缘设备构成了重大挑战。因此，我们的 工作负载规划还应该全面考虑每个设备的内存预算，以防止可用内存的过度消耗。

1)优化目标公式:

如§III-B所述，我们的HMP架构通过沿着三个不同的维度划分来分配工作负载:MHA块的头部维度，MLP块的权重矩阵的行维度，以及连接块的输入张量的序列维度。我们的工作负载规划侧重于确定每个块的分区配置，即:MHA块分区A = {a0, a1，…， aD−1}，MLP阻塞分区B = {b0, b1，…， bD−1}，连接块分区S = {s0, s1，…， sD−1}，其中D为边缘设备个数。我们引入符号L(MHA, Ad, d)， L(MLP, Bd, d)和L(CON, Sd, d)来分别表示设备d上的MHA块，MLP块和连接块的执行延迟，每个块给定其分区配置Ad, Bd和Sd。每个TP或SP块的执行时间L由散列器决定:

除了最小化执行延迟之外，我们的策略还需要在推理期间防止OOM错误。在部署基于transformer的模型时，巨大的内存占用源于MHA和MLP块中包含的大量权重矩阵。因此，我们的工作负载规划明智地划分了MHA和MLP块，从而允许模型的内存需求由多个设备协作处理。

我们将 $M_{att}$ 和 $M_{mlp}$ 分别表示为加载一个MHA块和一个MLP块的内存占用。$Budget_d$ 表示分配给设备d的内存预算，l 表示模型中Transformer层的总数。综上所述，在内存约束下最小化延迟的优化目标如下:

为了方便我们的工作负载规划算法，我们使用了Galaxy Profiler，它使用校准数据集作为物理边缘设备上的输入来进行推理过程，以记录并行规划所需的运行时概要。对于TP和SP块，该分析器仔细地捕获各种分区配置下的计算延迟。同时，Galaxy Profiler也记录模型信息，包括MHA和MLP块中包含的参数数量。

2)工作负载规划算法:

解决上述约束优化问题的稻草人方法包括穷举搜索所有可能的分区组合，然后选择满足内存约束的最优解。然而，这种方法具有指数复杂度，使得它不适用于大规模的Transformer模型。

连接块的执行时间主要取决于内存访问量，而不是SoC的计算能力，我们在SP规划中采用相等分区的策略。在张量同步过程中，相等的分区保持了所有设备之间均匀的通信量，为§III-D中基于tile的通信重叠奠定了有利的基础。

对于TP，我们可以实现与每个设备的计算能力成比例的工作负载分配的块的最佳分区，而不考虑内存预算。这种比例分区确保所有设备几乎同时完成它们的任务，有效地减少了可能导致资源利用率不理想的潜在延迟。

有了这些见解，我们设计了一个两步启发式算法，在算法1中概述。第一步，算法不考虑设备的内存约束，根据设备的计算能力分配相应的工作负载，从而保证工作负载的均衡(第1-8行)。随后，在此初始分布的基础上，第二步对工作负载分配进行微调。

它将超出内存预算的设备的多余工作负载重新分配给具有空闲内存容量的设备。(第9-19行)。考虑到MHA块(头部维度)的分区粒度通常比MLP块(列维度)的分区粒度粗，我们首先为MLP块(第21行)重新分配工作负载，然后是MHA块(第22行)。如果尽管工作负载重新分配，OOM错误仍然存在，这表明协作推理中涉及的边缘设备无法容纳目标模型，从而导致算法失败(第23-24行)。我们将设备的计算能力Vd定义为在设备d上执行MHA块和MLP块所需总时间的倒数。

工作负载规划是在部署前运行一次的脱机过程。算法1的时间复杂度上界为0 (D3)。在我们的实验中，4个异构边缘设备在国内笔记本电脑上的运行时间不到1秒。

基于tile的通信优化

与数据中心中稳定的高带宽网络相比，边缘环境经常与不一致的、带宽有限的连接作斗争。这增加了协同推理期间的同步延迟，成为全局系统性能的一个重要瓶颈。通信与计算重叠是一种有效的优化策略。然而，由于通信和计算之间严格的数据依赖关系，它的实现在Transformer推理中变得复杂。

为了解决这个问题，Galaxy引入了一种基于tile的方法来有效地解耦它们的依赖关系，以实现细粒度的重叠。

我们从图5中观察到，每个TP块都以GEMM操作（矩阵乘法操作）开始和结束。我们设计在进入和退出TP块时将这些GEMM操作与AllGather和ReduceScatter操作重叠。为了说明这一点，下面一节提供了一个跨三个设备的协作推理示例，演示如何在MLP块之前和之后使用同步点重叠GEMM(也适用于MHA块)。

AllGather重叠:

如图5所示，在MLP块中，AllGather与初始矩阵乘法(GEMM1)之间存在严格的数据依赖关系。具体来说，设备i (i∈{0,1,2})上的GEMM1只有在AllGather完成所有子序列的聚合后才能开始:D = AllGather(H0,H1,H2)， Ei = GEMM1(D,WD i)。

为了解耦AllGather和GEMM1之间的严格依赖关系，我们利用矩阵平铺来分解GEMM1。我们发现GEMM1的直接计算可以等价地通过将矩阵D水平分割成块，在每个块上独立执行GEMM1，然后将结果串联起来来实现。

我们采用Ring-AllGather实现，并将其与上述矩阵平铺方法相结合，实现通信和计算的重叠。图6中示出了涉及三个协作设备的重叠过程的示例。

在包含D个设备的基于tile的重叠流程的上下文中，通常需要D个步骤(在本例中为三个步骤)。我们定义(i +1)%3和(i−1)%3表示在3-设备环拓扑中设备i的后继设备和前导设备的索引。

步骤1:设备i在设备上tile Hi和WD i之间进行GEMM操作，并发地将Hi分发给后续设备。同时，设备i接收并存储从其前一个设备发送的磁块H(i−1)%3。步骤2:设备i对磁片H(i−1)%3进行GEMM操作，并发分发给后续设备。同时，设备i接收从其前一个设备发送的磁块H(i−2)%3。步骤3:设备i对磁片H(i−2)%3执行GEMM操作。值得注意的是，最后一步不需要任何通信。三个GEMM操作的结果沿着序列维度连接起来，产生最终结果Ei。

ReduceScatter重叠：

如图5所示，MLP块中最终矩阵乘法(GEMM2)与ReduceScatter操作(i∈{0,1,2})之间存在严格的数据依赖关系:

Fi = GEMM2(Ei,WE i)， Gi = ReduceScatter(F0, F1, F2)。

为了解耦ReduceScatter和GEMM2之间的严格依赖关系，我们镜像了AllGather使用的平铺方法。我们沿着行维将矩阵Ei分成三个大小相等的块Ei,r (r∈{0,1,2})(与连接块的分区配置一致)，并为每个块独立计算GEMM2 (Eq.10)。为了获得最终结果Gr，需要在所有设备上进行额外的ReduceSum操作(见公式11)。

与AllGather类似，我们使用Ring-ReduceScatter实现与矩阵平铺相结合来实现通信和计算重叠。如图6所示，ReduceScatter重叠的过程也包括三个步骤。

步骤1:设备 i 在Ei，(i+2)%3和WE i之间执行GEMM操作，产生结果Oi，(i+2)%3。

步骤2:设备i对磁片Ei，(i+1)%3执行GEMM操作，并产生结果Oi，(i+1)%3。同时，设备i将步骤1中的GEMM结果转发给后续设备。在接收到来自前一个设备的数据后，设备i在它和Oi (i+1)%3之间执行一个ReduceSum操作。

步骤3:设备i对磁片Ei,i进行GEMM操作，得到结果Oi,i。设备i并发送步骤2的ReduceSum结果给后续设备。在从前一个设备接收到的数据块和Oi,i之间执行一个ReduceSum操作，产生最终结果Gi。

我们基于tile的通信优化无缝地将D−1轮环通信与D轮GEMM操作重叠，而不会增加额外的开销，也不会产生与非重叠方法不一致的结果。

在本节中，我们在物理测试平台上评估了五种不同尺寸的基于基于Transformer的模型的Galaxy原型的性能。

模型和数据集。我们使用5个典型的基于Transformer的模型对Galaxy进行评估，范围从6600万到27亿个参数，详见表4。我们从流行GLUE数据集[31]的QNLI语料库中提取了一个平均序列长度为284的样本子集进行评估。

边缘环境设置。我们在各种现实边缘环境中评估Galaxy，包括现成边缘设备(Jetson Nano[32])的同质和异质配置，详见表II和III。在同构环境中，Nano-M的内存预算设置为1.5GB。在异构环境中，nano - ol的内存预算分别设置为1.5GB、Nano-M的1.2GB和Nano-S的0.7GB。我们限制使用板载CPU来模拟资源受限的边缘场景。我们还将在§IV-E中演示Galaxy在GPU环境中的有效性。我们调整了D2D带宽，以模拟现实边缘环境中的各种网络条件。

我们将Galaxy与单设备方法和最先进的并行方法进行比较:

Local Inference (Local):单设备上的推理模型。我们将其与Galaxy进行比较，分析Galaxy的可扩展性性能。

Megatron-LM (M-LM)[24]:一种最先进的TP方法，将权重矩阵拆分为MHA和MLP块，以并行化GEMM算子。每个MHA和MLP块之后都需要进行一次AllReduce同步。

序列并行(SP)[25]:最先进的SP方法沿着其序列维度划分输入，并在工作人员之间并行化推理。每个MHA块之间需要进行两次AllGather同步。

Galaxy在异构边缘环境中的卓越性能源于其对设备异构性的考虑，这是M-LM和SP忽略的一个因素，两者都是为配备均匀加速器的数据中心量身定制的。除了设备的异构性外，Galaxy工作负载规划还全面考虑了边缘设备的内存预算，使它们能够协同适应目标模型。相反，M-LM和SP在并行规划期间忽略了内存约束，导致OOM错误。

VLLM 框架

伯克利大学LMSYS组织开源的大语言模型高速推理框架

vLLM是一个快速且易于使用的库，用于 LLM 推理和服务，可以和 HuggingFace 无缝集成

在吞吐量方面，vLLM的性能比 HuggingFace Transformers(HF) 高出 24 倍，文本生成推理（TGI）高出3.5倍。

KV Cache

hugging face框架的generate方法默认会保存前面token的K和V向量，这样可以避免重复计算

这就是KV-Cache，在与模型对话时，随着模型输出逐渐变长，没有感觉到模型输出速度明显变慢，因为每次它只计算当前这个token，前面token的KV向量都已经缓存了

KV-Cache中实际利用率只有20%~40%，大部分KV-Cache里的显存都被浪费了，原因有几下几点：

1. **预分配，但不会用到：**大模型生成时并不知道要生成多少个token，所以总是按照生成参数里设置的最大token数来预分配KV-Cache，比如模型预分配的是1000个token，但输出到100个的时候就输出终止符结束了，那剩下预分配的900个token的KV-Cache就被浪费了
2. **预分配，但尚未用到：**假如一个样本真的可以输出到1000个token，但是在他刚开始输出第一个token的时候，剩下的token都还没有用到，但是显存已经被预分配占用了，这时其它的请求也无法被响应，而未被响应的请求有可能只需要输出10个token就结束了，本来可以和正在输出的样本并行处理
3. **显存之间的间隔碎片，不足以预分配给下一个文本生成：**即使最大生成长度一致，但因为Prompt的长度不同，每次预分配的KV-Cache大小也不同，当一个请求生成完毕释放缓存，但是下一个请求的Prompt长度大于释放的这个请求的Prompt长度，所以无法放入被释放的缓存之中，这种无法被使用的缓存就是碎片

vLLM就是解决了KV-Cache中的浪费问题，可以用更大的batch size来处理请求，从而提高了系统的吞吐量

vLLM中的优化

vllm中做了哪些优化？

Page Attention

第一个优化是Page Attention

Page Attention把显存划分为KV Block，显存按照KV Block来管理KV Cache

vLLM克服了预分配的问题，按需分配，不会提前占用显存，并且按块分配，这样就减少了内存的碎片

虚拟内存是每个请求都有一个逻辑的KV-Cache，在逻辑的KV-Cache里面显存好像是连续的，vLLM框架会在后台维护一个逻辑KV Cache到实际显卡上KV Block的映射表，在进行Page Attention计算时，它会自动找到物理显存上block的KV向量进行计算

每个请求都有自己的逻辑的KV-Cache，其中的Prompt和生成的新的token的KV向量看起来好像都是放在连续的显存上，方便程序操作

vLLM框架内部维护了映射表，在实际计算时，获取显卡上KV Block里的KV向量

虚拟内存的作用就是让程序在使用内存时，感觉自己使用的是连续的内存，但实际操作系统分配时却并不是连续的，这都是通过中间的映射表来实现的

Page Attention的优点总结如下：

Sharing KV Blocks

vLLM的第二个优化：Sharing KV Blocks

当我们在利用大模型进行生成时，有时候会想用一个Prompt生成多个不同的输出

在vLLM的sampling参数里，可以设置N为大于1的一个整数来实现这个功能

此时的同一个Prompt会产生两个不同的序列，它们的Prompt是相同的，这是两个序列，但在显卡的显存里只存放了一份Prompt token的KV Block，每个Block都标记着自己现在被两个序列引用着。只有当引用数为0时，这个Block占用的显存才会被释放。

接着，第一个序列开始生成，生成的第一个token是color，这时候会触发Copy on Write机制：也就是它发现自己要继续写入的block的引用数不是1，它就不能直接写入，必须自己拷贝一份来写入，在自己的拷贝上写入color这个token的KV Cache，然后原来的那个block的引用数就减为1了，自己新写入的这个拷贝引用数变为1

序列2下一个生成的token是matter，它看到自己引用的block2引用数是1，就可以直接写，于是写入自己的下一个token为matter，然后两个序列就各自往下生成，但是它们Prompt的起始部分KV-Block是共享的，这就节省了显存

KV-blocks的共享还可以优化Beam Search的显存占用

中国的历史被3个Beam引用，非常悠久被2个Beam引用，它们之间都共享这些Token的KV-Cache，这就是vLLM的改进

vLLM的代码调用非常简单：

LLAMA模型结构

Meta开发的LLAMA

open和efficient

open：使用了公开渠道可以获取的数据，不包含Meta的客户数据，性能表现媲美GPT3，对研究用途的社区开放，但不能商用

efficient：之前的Scaling Law计算的都是根据训练计算代价花费来评估的，所以会认为在我们到达指定的模型性能时，训练大模型会更加划算，但是Meta认为比起训练计算代价而言，更应该看重推理时的计算代价，因为训练只有一次，推理却是无数次的

Meta发现训练时候应该尽可能多地多砸token数量，这样模型性能就能不断地增长，推理时候的计算代价也就相对较低

Meta在LLAMA系列里，一直都认为：增加训练数据比扩大模型参数更有效

LLAMA1

LLAMA1训练时用的数据：

LLAMA1模型架构：

RMSNorm 归一化公式

silu 激活函数

大体上和 relu 类似；silu 在 0 值附近更为光滑，但因为要计算指数函数，所以计算代价对应较高

LLAMA2

LLAMA2有三个特点：

1. 更加open，可以商用
2. 训练数据更大，贯彻 Meta 的增大数据量比增加模型参数更有效的理念
3. 通过微调训练了 Chat Model

LLAMA2训练情况：

LLAMA在模型架构里引入了GQA(Group Query Attention)，是对**MHA(Multi-Head Attention)**的一种优化，可以减少模型参数量和KV-Cache的大小

MHA中每个token对应的特征都会生成个数相同的K、Q、V的头，其中1号的Q会和其它token以及自己1号的K进行注意力计算

Multi-query中每个token对应的特征还是生成8个Q，但是只生成一个K和V，所有的Q都和其它token以及自己唯一的K进行相似度计算

中间的GQA就是LLAMA2所采用的分组注意力机制，它是左右两者的折中，它会将Q进行分组，每组对应一个K和V

GQA可以节省模型的参数量，这样Q有8个头，K和V只有4个头，实现时为了能够进行矩阵乘法，还是将K和V的头都复制一份，这样就能像MHA那样一一对应了

GQA在LLAMA2中之应用了最大的那个70B模型里，

LLAMA3

LLAMA3实验表现：

LLAMA3模型架构：

LLAMA3训练数据：

LLAMA的下游任务：

1. Clm任务（因果推断）
2. 文本分类任务

DeepSeek加速方法

DeepSeek-MoE原理讲解

MoE（Mixture of Experts）混合专家模型

MoE模型在达到同样效果时，训练花费计算代价会更少，而且最终推理时花费的计算代价也非常少

传统的 Decoder only 的稠密大模型结构：

首先是一个带掩码的多头自注意力机制模块，然后是一个残差链接和层归一化，接下来是一个全链接前向传播模块，然后又是一个残差链接和层归一化

MOE主要是对 FeedForwardLayer（全连接的前向传播模块）进行改造：

专家的选择是对每个token进行的，不是对一个序列

MOE的特点如下：

为了让MOE中的专家负载均衡，办法如下：

设置负载均衡损失可以让模型在训练过程中自己学会负载均衡：

负载均衡损失内容如下：

这个负载均衡损失函数不能直接拿来运用，所以有以下改动：

有了以上的MOE基础知识，我们就可以学习DeepSeek-MOE了

改进点：让专家更加专精（增大专家数量，减小每个专家的网络参数量，也就是更加专精的专家）

对专家进行细分，可以得到更灵活的专家组合

DeepSeek又想到，所有的专家都要学习一些基础的通用的能力，这样就可以将所有专家都要学习的通用基础能力提取出来作为一个共享专家，共享专家保证每次都可以被激活，他负责所有专家都需要的通用能力

这样，路由专家加共享专家的总数不变，和原来选择两个共享专家的原始MOE网络计算代价是相同的

提取共享专家和对专家进行细分，都是有利于模型性能的

DeepSeek-v2-MLA原理讲解

DeepSeek-MOE的性能几乎接近同等参数的稠密模型，但是训练和推理的计算代价都大幅地缩减

现在我们来看下DeepSeek-V2：

首先我们来看下DeepSeek-V2的模型架构图：

DeepSeek-MoE模块中更改了Feed-Forward Network模块，这个改动在DeepSeek-V2中依然保留，同时DeepSeek-V2也对注意力层进行修改，提出了MLA（Multi-Head Latent Attention）多头潜在注意力机制

标准的多头注意力机制过程：

对于第一个token，它的特征向量为H，通过Query、Key和Value权重矩阵分别得到这个token的Q向量、K向量和V向量

然后K向量和K向量计算点积得到V向量的权重，再与V向量相乘得到输出向量，这个输出向量将进入这个block的DeepSeek-MOE层，最终通过多个transformer block，得到最终这个token的特征向量，接一个分类头可以预测得到下一个token，然后将新生成的token拼接到输入序列，然后用两个token的特征向量经过Query权重、Key权重和Value权重，得到两个token分别的Q、K和V向量

KV-Cache减少了推理时的计算量，加快了推理速度，但是它是以宝贵的显存空间来换取计算量的减少，随着序列越来越长，显存占用的会越来越多

为此提出了分组注意力机制GQA，但还是标准的MHA效果最好，MQA和GQA也有很多的大模型采用，它们在推理时可以减少KV-Cache的大小，但是显著影响了模型的性能，那么有没有既能减少KV-Cache，又能提高模型性能的做法呢？这种方法就是DeepSeek的MLA

三种方法（Cache缓存量对比）

即先将token的特征向量首先通过一个参数矩阵进行压缩转换，缓存时候就只需要缓存压缩后的KV向量，在进行计算时候再从KV压缩向量通过两个减压矩阵转化为原来的维度就可以了

通过实验发现，MLA不仅减少了KV-Cache，而且模型效果也得到了提升！

但是MLA取出缓存时不能直接使用，需要进行减压计算，也就是说推理时引入了额外计算，这和KV-Cache初衷相悖

KV-Cache推理过程（上半部分为MHA、MQA、GQA，下半部分为MLA）：

仔细看MLA的公式，可以通过矩阵相乘的结合律进行简化式子，进行合并矩阵，可以在推理之前提前计算好，这样推理时就不用进行额外的减压计算了

MLA对Q向量也进行了压缩，这样能够减少模型的参数量，也提高了模型性能，但是Q向量不需要进行缓存

然后是旋转位置编码（RoPE，大模型默认的位置编码方式）：对每一层的Q和K向量进行旋转，根据token位置的不同，旋转矩阵的参数也不同

增加了旋转位置编码后，就不能进行矩阵融合！

DeepSeek对此的解决方案是给Q和K向量额外增加一些维度，来表示它们的位置信息

这就是最终的解决方案，得到了既兼容旋转位置编码的压缩KV-Cache的方案，又可以提升模型性能

DeepSeek-v3原理讲解多token预测 MTP

性能高，而且训练成本低，训练时间短，成本不到LAMMA3的十分之一

除了在模型架构上做了一些改动外，也在训练框架方面做了非常多的改进，业界首次使用FP8混合精度训练超大模型

模型架构图如下：

模型架构仍然采用了MOE和MLA

DeepSeek-V3中去掉了针对batch的辅助负载均衡损失函数，而是在训练时在每个step针对每个专家的负载进行监控，如果某个专家负载过高，就减少对应的偏置值 $b_i$，如果负载过低，就增大偏置值，这个偏置值只在专家路由时起作用，在和最终专家的输出相乘时，还是采用原来的权重值，不会加入这个偏置值

增加了序列负载均衡损失函数，防止在序列内部专家负载不均衡，这样的模型效果更好

下一个大的改动就是：训练时采用了多Token预测（MTP：Multi-Token Prediction）

之前学的LLM都是每次只能预测1个token，现在逐渐往多token预测的方向发展

这是DeepSeek-V3出现之前的多token预测的模型架构：

MTP优势：

多token预测还有一个好处，就是可以加速预测：

它的加速预测不是简单的用多个头来生成接下来的多个token，因为预测最准确的还是预测下一个token的那个头，预测越远的token越不准确，需要进行修正

下面是利用MTP进行预测加速的方法：

用第一个头去预测下一个token，对不确定的token逐个进行确认

注意：这里三个序列是通过一个batch并行进行验证的，因为GPU擅长并行计算，主要的时间开销是显存访问

最后接收最长的序列

可以看到通过预测和验证两步运算，就预测出来了4个token，原来是需要四步运算才能够进行计算出来的，由此可看出提高了模型的效率

更进一步，可以将验证和预测放在一起进行：

验证时候从只让一个头工作 改为 让四个头同时工作，验证时以4个序列为一个batch，则在验证完成后就又可以进入下一步的验证了，这样不断地进行验证加预测的循环就可以了

MTP预测加速只在生成单个序列时才有效，也就是单个用户！

所以只在训练时利用MTP来提升模型的性能，推理部署时一般就丢弃掉其它几个预测多步的头，就只用第一个头来做下一个token的预测，和普通的单头的大模型没有区别

DeepSeek的MTP给每个头传入了额外的信息，帮助多token预测的头能更好地预测出自己接下来的token

DeepSeek-GRPO原理

GPRO算法是针对大预言模型改进的一种强化学习算法

PPO（近端策略优化算法）

强化学习概念：

Environment 环境

Agent 智能体

State 状态

​ Observation(State的一部分，因为有的游戏里Agent视野并不是全视野，只能看到整个游戏状态的一部分)

Action 动作(Agent做出的动作)

Reward 奖励(Agent做出的动作对应的奖励和惩罚)

强化学习的目的：

训练一个 Policy 神经网络 $\pi$，在所有状态 S 下，给出相应的 Action，得到 Return 的期望最大

训练一个 Policy 神经网络 $\pi$，在所有的 Trajectory 中，得到 Return 的期望最大

注意：每一步的Action都是按照概率进行采样的，不是直接选取最大值的！

这样对应有一个问题，就是大部分时间在采集数据，训练比较慢，这就是PPO算法要解决的问题

仔细看公式，不难发现是有改进空间的：

1. 我们是否增大或减少在状态S下做动作A的概率，应该看做了这个动作之后到游戏结束累计的reward，而不是整个trajectory累积的reward，因为一个动作只能影响它后面的reward，不能影响它之前的。
2. 动作是可以对接下来的reward产生影响，但是它有可能只影响几步，而且影响会逐渐衰退，后面reward更多是由当时的动作影响。

针对以上两点进行了改进，

1. 首先对reward求和，而不是对整个trajectory的reward进行求和，而是从当前步 t 到 trajectory 结束的 reward 进行求和

2. 引入衰减因子 γ，γ 小于 1，距离当前步 t 越远，当前动作对 reward 影响越小，呈指数衰减，即距离当前越远的reward受当前步动作影响越小

接下来还有一些概念：

PPO算法（Proximal Policy Optimization）邻近策略优化

下面我们来看一下重要性采样

利用重要性采样来更新我们目标函数的梯度公式，这样可以将On Policy的训练转化为Off Policy的训练

如何给loss函数增加训练的策略不能和参考的策略分布相差太大呢？这里就需要加上KL散度的约束

两个概率密度函数分布完全一致，则KL散度为0；分布越不一致，KL散度越大！通过 β 来调整KL散度约束的大小

ppo2算法还有一种实现方式，就是通过截断函数来替代KL散度！

红色是原始公式，蓝色是截断函数

对于一个神经网络的训练，最重要的有3点：

1. 网络结构
2. loss函数
3. 训练数据

质量微调和预训练没有本质的不同，前两点都一样，最大的不同就是训练数据的不同

Reward模型的训练

reward 模型训练采用偏好数据

大模型能力的极限是由大模型预训练时候决定的，强化学习只是逼近极限

PPO模型的训练

PPO模型的训练需要四个模型(HS: hidden_size, VS: vocab_size)：

基准模型一般是SFT(监督微调训练)后的模型

训练模型和基准模型结构一致，PPO训练的目标就是优化训练模型，模型输出的概率分布不能和基准模型相差太大

奖励模型用于对一个问答序列进行评分，输出一个分数

状态价值模型会对每个状态评估它对应的价值，预测问答序列结束后的期望回报是多少，

只加载一个LLM，加载多个adapter的操作，来大大减少对显存的占用，以此训练PPO算法